embedding api improvements

Cargo.toml

@@ -82,7 +82,8 @@ parking_lot = { version = "0.12.1", features = [
     "send_guard",
 ] }
 ordered-float = "4.2.0"
-chrono = { version = "=0.4.38", features = ["serde"] }
+# chrono = { version = "=0.4.41", features = ["serde"] } # this works for lancedb version requirements
+chrono = { version = "=0.4.39", features = ["serde"] } # this works for arrow-arith at compilation time
 tempfile = "3.10.0"
 futures-util = "0.3.30"
 thiserror = "2.0.0"
@@ -154,6 +155,8 @@ minijinja = "2.2.0"
 minijinja-contrib = { version = "2.2.0", features = ["datetime"] }
 datafusion = { version = "43.0.0" }
 arroy = "0.6.1"
+# lancedb = "0.22.0" # this is the latest and asks for chrono 0.4.41
+lancedb = "0.19.1" # this uses chrono 0.4.39, which doesn't causes the problem when compiling arrow-arith 53.2.0
 heed = "0.22.0"
 sqlparser = "0.51.0"
 futures = "0.3"
@@ -172,7 +175,11 @@ uuid = { version = "1.16.0", features = ["v4"] }
 
 # Make sure that transitive dependencies stick to disk_graph 50
 [patch.crates-io]
+# if we ever can, please try to sync these versions with those used by lancedb, ask Pedro for more details
 arrow-buffer = { git = "https://github.com/apache/arrow-rs.git", tag = "53.2.0" }
 arrow-schema = { git = "https://github.com/apache/arrow-rs.git", tag = "53.2.0" }
 arrow-data = { git = "https://github.com/apache/arrow-rs.git", tag = "53.2.0" }
 arrow-array = { git = "https://github.com/apache/arrow-rs.git", tag = "53.2.0" }
+# chrono = { version = "0.4.39" }
+# chrono = { version="0.4.41", git = "https://github.com/chronotope/chrono.git", tag = "v0.4.39" }
+chrono = { git = "https://github.com/chronotope/chrono.git", tag = "v0.4.38" } # this forces lancedb to work with chrono 0.4.38

Makefile

@@ -54,7 +54,7 @@ rust-fmt:
 rust-build:
 	cargo build -q
 
-rust-build-docs: 
+rust-build-docs:
 	cargo doc --no-deps -p raphtory -q
 
 run-graphql:

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_vectors.py

@@ -1,7 +1,10 @@
+from time import sleep
+import pytest
 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,16 +14,16 @@
     "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}")
-
 
-def embedding(texts):
-    return [single_embedding(text) for text in texts]
+    custom_embeddings.start()
+    sleep(1)
+    yield
+    custom_embeddings.stop()
 
 
 def floats_are_equals(float1: float, float2: float) -> bool:
@@ -48,26 +51,28 @@ 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 }}")
+# FIXME: I should not need to write /v1?, change that in rust, so the route is /embeddings, not v1/embeddings
+    embeddings = OpenAIEmbeddings(api_base="http://localhost:7340/v1")
+    vg = g.vectorise(embeddings, nodes="{{ name }}", edges="{{ properties.name }}")
 
     return vg
 
 
 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()
-    ###########################
+    # ################################
+    # 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")]

raphtory-benchmark/src/common/vectors.rs

@@ -35,7 +35,7 @@ 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(embedding_model).await.unwrap();
     let template = DocumentTemplate {
         node_template: Some("{{name}}".to_owned()),
         edge_template: None,

raphtory-graphql/schema.graphql

@@ -13,7 +13,6 @@ enum AllPropertySpec {
 	ALL_PROPERTIES
 }
 
-
 """
 Collection of items
 """
@@ -637,7 +636,6 @@ type EdgesWindowSet {
 	list: [Edges!]!
 }
 
-
 """
 Document in a vector graph
 """
@@ -940,8 +938,8 @@ type Graph {
 }
 
 type GraphAlgorithmPlugin {
-	pagerank(iterCount: Int!, threads: Int, tol: Float): [PagerankOutput!]!
 	shortest_path(source: String!, targets: [String!]!, direction: String): [ShortestPathOutput!]!
+	pagerank(iterCount: Int!, threads: Int, tol: Float): [PagerankOutput!]!
 }
 
 type GraphSchema {
@@ -1063,7 +1061,6 @@ type GraphWindowSet {
 	list: [Graph!]!
 }
 
-
 input IndexSpecInput {
 	"""
 	Node properties.
@@ -1086,7 +1083,6 @@ input InputEdge {
 	dst: String!
 }
 
-
 type LayerSchema {
 	"""
 	Returns the name of the layer with this schema
@@ -2299,7 +2295,6 @@ enum SortByTime {
 	EARLIEST
 }
 
-
 type TemporalProperties {
 	"""
 	Get property value matching the specified key.
@@ -2452,9 +2447,21 @@ input WindowDuration @oneOf {
 	epoch: Int
 }
 
+"""
+Directs the executor to include this field or fragment only when the `if` argument is true.
+"""
 directive @include(if: Boolean!) on FIELD | FRAGMENT_SPREAD | INLINE_FRAGMENT
+"""
+Indicates that an Input Object is a OneOf Input Object (and thus requires exactly one of its field be provided)
+"""
 directive @oneOf on INPUT_OBJECT
+"""
+Directs the executor to skip this field or fragment when the `if` argument is true.
+"""
 directive @skip(if: Boolean!) on FIELD | FRAGMENT_SPREAD | INLINE_FRAGMENT
+"""
+Provides a scalar specification URL for specifying the behavior of custom scalar types.
+"""
 directive @specifiedBy(url: String!) on SCALAR
 schema {
 	query: QueryRoot