embedding api improvements

Cargo.toml

@@ -24,7 +24,7 @@ readme = "README.md"
 homepage = "https://github.com/Raphtory/raphtory/"
 keywords = ["graph", "temporal-graph", "temporal"]
 authors = ["Pometry"]
-rust-version = "1.86.0"
+rust-version = "1.88.0"
 edition = "2021"
 
 # debug symbols are using a lot of resources
@@ -87,7 +87,7 @@ parking_lot = { version = "0.12.1", features = [
     "send_guard",
 ] }
 ordered-float = "4.2.0"
-chrono = { version = "0.4.42", features = ["serde"] }
+chrono = { version = "0.4.41", features = ["serde"] }
 tempfile = "3.10.0"
 futures-util = "0.3.30"
 thiserror = "2.0.0"
@@ -153,7 +153,7 @@ pest_derive = "2.7.8"
 minijinja = "2.2.0"
 minijinja-contrib = { version = "2.2.0", features = ["datetime"] }
 datafusion = { version = "50.0.0" }
-arroy = "0.6.1"
+lancedb = "0.22.2" # this is the latest and asks for chrono 0.4.41
 heed = "0.22.0"
 sqlparser = "0.58.0"
 futures = "0.3"
@@ -172,4 +172,3 @@ strsim = { version = "0.11.1" }
 uuid = { version = "1.16.0", features = ["v4"] }
 strum = "0.27.2"
 strum_macros = "0.27.2"
-

milvus/docker-compose.yml

@@ -0,0 +1,26 @@
+version: "3.9"
+
+services:
+  milvus:
+    image: milvusdb/milvus:v2.6.1
+    container_name: milvus
+    command: ["milvus", "run", "standalone"]
+    security_opt:
+      - seccomp:unconfined
+    environment:
+      ETCD_USE_EMBED: "true"
+      COMMON_STORAGETYPE: "local"
+      DEPLOY_MODE: "STANDALONE"
+    ports:
+      - "9091:9091"
+      - "19530:19530"
+
+  attu:
+    image: zilliz/attu:v2.6
+    container_name: attu
+    environment:
+      MILVUS_URL: "http://milvus:19530"
+    ports:
+      - "8000:3000"
+    depends_on:
+      - milvus

python/tests/test_base_install/test_graphql/misc/test_graphql_vectors.py

@@ -1,16 +1,12 @@
 import tempfile
 from raphtory.graphql import GraphServer, RaphtoryClient
 from raphtory import Graph
+from raphtory.vectors import OpenAIEmbeddings, embedding_server
 
 
-def embedding(texts):
-    return [[text.count("a"), text.count("b")] for text in texts]
-
-
-def test_embedding():
-    result = embedding(texts=["aaa", "b", "ab", "ba"])
-    assert result == [[3, 0], [0, 1], [1, 1], [1, 1]]
-
+@embedding_server(address="0.0.0.0:7340")
+def embeddings(text: str):
+    return [text.count("a"), text.count("b")]
 
 def setup_graph(g):
     g.add_node(1, "aab")
@@ -58,51 +54,61 @@ def assert_correct_documents(client):
     }
 
 
-def setup_server(work_dir):
-    server = GraphServer(work_dir)
-    server = server.set_embeddings(
-        cache="/tmp/graph-cache",
-        embedding=embedding,
-        nodes="{{ name }}",
-        edges=False,
-    )
-    return server
-
-
 def test_new_graph():
     print("test_new_graph")
     work_dir = tempfile.TemporaryDirectory()
-    server = setup_server(work_dir.name)
-    with server.start():
-        client = RaphtoryClient("http://localhost:1736")
-        client.new_graph("abb", "EVENT")
-        rg = client.remote_graph("abb")
-        setup_graph(rg)
-        assert_correct_documents(client)
+    server = GraphServer(work_dir.name)
+    with embeddings.start():
+        with server.start():
+            client = RaphtoryClient("http://localhost:1736")
+            client.new_graph("abb", "EVENT")
+            rg = client.remote_graph("abb")
+            setup_graph(rg)
+            client.query("""
+                {
+                    vectoriseGraph(path: "abb", model: { openAI: { model: "whatever", apiBase: "http://localhost:7340" } }, nodes: { custom: "{{ name }}" }, edges: { enabled: false })
+                }
+                """)
+            assert_correct_documents(client)
 
 
 def test_upload_graph():
     print("test_upload_graph")
-    work_dir = tempfile.mkdtemp()
-    temp_dir = tempfile.mkdtemp()
-    server = setup_server(work_dir)
-    with server.start():
-        client = RaphtoryClient("http://localhost:1736")
-        g = Graph()
-        setup_graph(g)
-        g_path = temp_dir + "/abb"
-        g.save_to_zip(g_path)
-        client.upload_graph(path="abb", file_path=g_path, overwrite=True)
-        assert_correct_documents(client)
+    work_dir = tempfile.TemporaryDirectory()
+    temp_dir = tempfile.TemporaryDirectory()
+    server = GraphServer(work_dir.name)
+    with embeddings.start():
+        with server.start():
+            client = RaphtoryClient("http://localhost:1736")
+            g = Graph()
+            setup_graph(g)
+            g_path = temp_dir.name + "/abb"
+            g.save_to_zip(g_path)
+            client.upload_graph(path="abb", file_path=g_path, overwrite=True)
+            client.query("""
+                {
+                vectoriseGraph(path: "abb", model: { openAI: { model: "whatever", apiBase: "http://localhost:7340" } }, nodes: { custom: "{{ name }}" }, edges: { enabled: false })
+                }
+                """)
+            assert_correct_documents(client)
 
+GRAPH_NAME = "abb"
 
 def test_include_graph():
-    work_dir = tempfile.mkdtemp()
-    g_path = work_dir + "/abb"
+    work_dir = tempfile.TemporaryDirectory()
+    g_path = work_dir.name + "/" + GRAPH_NAME
     g = Graph()
     setup_graph(g)
     g.save_to_file(g_path)
-    server = setup_server(work_dir)
-    with server.start():
-        client = RaphtoryClient("http://localhost:1736")
-        assert_correct_documents(client)
+    server = GraphServer(work_dir.name)
+    with embeddings.start():
+        embedding_client = OpenAIEmbeddings(api_base="http://localhost:7340")
+        server.vectorise_graph(
+            name=GRAPH_NAME,
+            embeddings=embedding_client,
+            nodes="{{ name }}",
+            edges=False
+        )
+        with server.start():
+            client = RaphtoryClient("http://localhost:1736")
+            assert_correct_documents(client)

python/tests/test_base_install/test_vectors.py

@@ -1,7 +1,10 @@
+import pytest
+import requests
 from raphtory import Graph
-from raphtory.vectors import VectorisedGraph
+from raphtory.vectors import VectorisedGraph, OpenAIEmbeddings, embedding_server
 
 embedding_map = {
+    "raphtory": [1.0, 0.0, 0.0], # this is now needed,
     "node1": [1.0, 0.0, 0.0],
     "node2": [0.0, 1.0, 0.0],
     "node3": [0.0, 0.0, 1.0],
@@ -11,20 +14,32 @@
     "edge3": [0.0, 1.0, 1.0],
 }
 
-
-def single_embedding(text: str):
-    try:
+@pytest.fixture(autouse=True)
+def test_server():
+    @embedding_server(address="0.0.0.0:7340") # TODO: ask only for PORT!!!
+    def custom_embeddings(text: str):
         return embedding_map[text]
-    except:
-        raise Exception(f"unexpected document content: {text}")
 
+    with custom_embeddings.start():
+        yield
+
+def test_failing_python_embeddings():
+    @embedding_server(address="0.0.0.0:7342")
+    def failing_embeddings(text: str):
+        assert(False)
+
+    with failing_embeddings.start():
+        headers = { "Content-Type": "application/json" }
+        payload = { "model": "whatever", "input": ["Hello world"] }
+        response = requests.post("http://localhost:7342/embeddings", headers=headers, json=payload)
+        assert(response.status_code == 500)
+        response = requests.post("http://localhost:7342/embeddings", headers=headers, json=payload)
+        assert(response.status_code == 500)
 
-def embedding(texts):
-    return [single_embedding(text) for text in texts]
 
 
 def floats_are_equals(float1: float, float2: float) -> bool:
-    return float1 + 0.001 > float2 and float1 - 0.001 < float2
+    return float1 + 0.00001 > float2 and float1 - 0.01 < float2
 
 
 # the graph generated by this function looks like this:
@@ -48,26 +63,34 @@ def create_graph() -> VectorisedGraph:
     g.add_edge(3, "node1", "node3", {"name": "edge2"})
     g.add_edge(4, "node3", "node4", {"name": "edge3"})
 
-    vg = g.vectorise(embedding, nodes="{{ name }}", edges="{{ properties.name }}")
+    embeddings = OpenAIEmbeddings(api_base="http://localhost:7340")
+    vg = g.vectorise(embeddings, nodes="{{ name }}", edges="{{ properties.name }}")
 
     return vg
 
+def test_embedding_sever_context_manager():
+    @embedding_server(address="0.0.0.0:7341")
+    def constant(text: str):
+        return [1.0]
+
+    with constant.start():
+        headers = { "Content-Type": "application/json" }
+        data = {
+            # "model": "whatever",
+            "input": ["The text to vectorise"]
+        }
+        response = requests.post("http://localhost:7341/embeddings", headers=headers, json=data)
+        response.raise_for_status()
+        result = response.json()
+        vector = result['data'][0]['embedding']
+        assert vector == [1.0]
+
 
 def test_selection():
     vg = create_graph()
 
-    ################################
-    selection = vg.empty_selection()
-    nodes_to_select = ["node1", "node2"]
-    edges_to_select = [("node1", "node2"), ("node1", "node3")]
-    selection = vg.empty_selection()
-    selection.add_nodes(nodes_to_select)
-    selection.add_edges(edges_to_select)
-    nodes = selection.nodes()
-    ###########################
-
     assert len(vg.empty_selection().get_documents()) == 0
-    assert len(vg.empty_selection().get_documents_with_scores()) == 0
+    assert len(vg.empty_selection().get_documents_with_distances()) == 0
 
     nodes_to_select = ["node1", "node2"]
     edges_to_select = [("node1", "node2"), ("node1", "node3")]
@@ -77,7 +100,9 @@ def test_selection():
     nodes = selection.nodes()
     node_names_returned = [node.name for node in nodes]
     assert node_names_returned == nodes_to_select
+    print("before get documents")
     docs = [doc.content for doc in selection.get_documents()]
+    print("after get documents")
     assert docs == ["node1", "node2"]
 
     selection = vg.empty_selection()
@@ -113,8 +138,8 @@ def test_search():
     assert edge_names_returned == [("node1", "node2")]
     # TODO: same for edges ?
 
-    [(doc1, score1)] = vg.entities_by_similarity("node1", 1).get_documents_with_scores()
-    assert floats_are_equals(score1, 1.0)
+    [(doc1, distance1)] = vg.entities_by_similarity("node1", 1).get_documents_with_distances()
+    assert floats_are_equals(distance1, 0.0)
     assert (doc1.entity.name, doc1.content) == ("node1", "node1")
 
     # chained search
@@ -205,16 +230,19 @@ def test_filtering_by_entity_type():
     assert contents == ["edge1", "edge2", "edge3"]
 
 
-def constant_embedding(texts):
-    return [[1.0, 0.0, 0.0] for text in texts]
 
+@embedding_server(address="0.0.0.0:7341")
+def constant_embedding(_text):
+    return [1.0, 0.0, 0.0]
 
 def test_default_template():
     g = Graph()
     g.add_node(1, "node1")
     g.add_edge(2, "node1", "node1")
 
-    vg = g.vectorise(constant_embedding)
+    running = constant_embedding.start()
+
+    vg = g.vectorise(OpenAIEmbeddings(api_base="http://localhost:7341"))
 
     node_docs = vg.nodes_by_similarity(query="whatever", limit=10).get_documents()
     assert len(node_docs) == 1
@@ -226,3 +254,5 @@ def test_default_template():
         edge_docs[0].content
         == "There is an edge from node1 to node1 with events at:\n- Jan 1 1970 00:00\n"
     )
+
+    running.stop()

raphtory-benchmark/src/common/vectors.rs

@@ -4,8 +4,9 @@ use rand::{rngs::StdRng, Rng, SeedableRng};
 use raphtory::{
     prelude::{AdditionOps, Graph, NO_PROPS},
     vectors::{
-        cache::VectorCache, embeddings::EmbeddingResult, template::DocumentTemplate,
-        vectorisable::Vectorisable, vectorised_graph::VectorisedGraph, Embedding,
+        cache::VectorCache, embeddings::EmbeddingResult, storage::OpenAIEmbeddings,
+        template::DocumentTemplate, vectorisable::Vectorisable, vectorised_graph::VectorisedGraph,
+        Embedding,
     },
 };
 use tokio::runtime::Runtime;
@@ -35,12 +36,22 @@ pub fn create_graph_for_vector_bench(size: usize) -> Graph {
 }
 
 pub async fn vectorise_graph_for_bench_async(graph: Graph) -> VectorisedGraph<Graph> {
-    let cache = VectorCache::in_memory(embedding_model);
+    let cache = VectorCache::in_memory();
+    let model = cache
+        .openai(OpenAIEmbeddings {
+            model: "whatever".to_owned(),
+            api_base: Some("localhost://1783".to_owned()), // TODO: run embedding server as well on the background so that this works
+            api_key_env: None,
+            project_id: None,
+            org_id: None,
+        })
+        .await
+        .unwrap();
     let template = DocumentTemplate {
         node_template: Some("{{name}}".to_owned()),
         edge_template: None,
     };
-    graph.vectorise(cache, template, None, true).await.unwrap()
+    graph.vectorise(model, template, None, true).await.unwrap()
 }
 
 // TODO: remove this version

raphtory-graphql/schema.graphql

@@ -2546,6 +2546,12 @@ type QueryRoot {
 	"""
 	updateGraph(path: String!): MutableGraph!
 	"""
+	Update graph query, has side effects to update graph state
+
+	Returns:: GqlMutableGraph
+	"""
+	vectoriseGraph(path: String!, model: String, apiBase: String, nodes: Template, edges: Template): Boolean!
+	"""
 	Create vectorised graph in the format used for queries
 
 	Returns:: GqlVectorisedGraph
@@ -2612,6 +2618,17 @@ enum SortByTime {
 	EARLIEST
 }
 
+input Template @oneOf {
+	"""
+	The default template.
+	"""
+	enabled: Boolean
+	"""
+	A custom template.
+	"""
+	custom: String
+}
+
 type TemporalProperties {
 	"""
 	Get property value matching the specified key.