Stopped the embedding from being returned with VectorRetriever (neo4j#58)

* Stopped the embedding from being returned with VectorRetriever

* Updated test_neo4j_queries.py

* Removed apoc calls

* Added nodeLabels back in

* Added nodeLabels and id to VectorRetriever metadata

* Updated CHANGELOG

* Updated CHANGELOG

Author: alexthomas93

CHANGELOG.md

@@ -2,6 +2,9 @@
 
 ## Next
 
+### Added
+-   Stopped embeddings from being returned when searching with `VectorRetriever`. Added `nodeLabels` and `id` to the metadata of `VectorRetriever` results.
+
 ## 0.2.0
 
 ### Fixed

src/neo4j_genai/neo4j_queries.py

@@ -147,7 +147,9 @@ def get_search_query(
     else:
         raise ValueError(f"Search type is not supported: {search_type}")
     query_tail = get_query_tail(
-        retrieval_query, return_properties, fallback_return="RETURN node, score"
+        retrieval_query,
+        return_properties,
+        fallback_return=f"RETURN node {{ .*, `{embedding_node_property}`: null }} AS node, labels(node) AS nodeLabels, elementId(node) AS id, score",
     )
     return f"{query} {query_tail}", params
 

src/neo4j_genai/retrievers/vector.py

@@ -108,6 +108,8 @@ def default_format_record(self, record: neo4j.Record) -> RetrieverResultItem:
         """
         metadata = {
             "score": record.get("score"),
+            "nodeLabels": record.get("nodeLabels"),
+            "id": record.get("id"),
         }
         node = record.get("node")
         return RetrieverResultItem(

tests/e2e/test_vector_e2e.py

@@ -32,6 +32,7 @@ def test_vector_retriever_search_text(
     assert isinstance(results, RetrieverResult)
     assert len(results.items) == 5
     for result in results.items:
+        assert f"'{retriever._embedding_node_property}': None" in result.content
         assert isinstance(result, RetrieverResultItem)
 
 
@@ -66,6 +67,7 @@ def test_vector_retriever_search_vector(driver: Driver) -> None:
     assert isinstance(results, RetrieverResult)
     assert len(results.items) == 5
     for result in results.items:
+        assert f"'{retriever._embedding_node_property}': None" in result.content
         assert isinstance(result, RetrieverResultItem)
 
 

tests/unit/retrievers/test_vector.py

@@ -101,7 +101,8 @@ def test_similarity_search_vector_happy_path(
     assert records == RetrieverResult(
         items=[
             RetrieverResultItem(
-                content="{'text': 'dummy-node'}", metadata={"score": 1.0}
+                content="{'text': 'dummy-node'}",
+                metadata={"score": 1.0, "nodeLabels": None, "id": None},
             ),
         ],
         metadata={"__retriever": "VectorRetriever"},
@@ -143,7 +144,10 @@ def test_similarity_search_text_happy_path(
     )
     assert records == RetrieverResult(
         items=[
-            RetrieverResultItem(content="dummy-node", metadata={"score": 1.0}),
+            RetrieverResultItem(
+                content="dummy-node",
+                metadata={"score": 1.0, "nodeLabels": None, "id": None},
+            ),
         ],
         metadata={"__retriever": "VectorRetriever"},
     )
@@ -191,7 +195,10 @@ def test_similarity_search_text_return_properties(
     )
     assert records == RetrieverResult(
         items=[
-            RetrieverResultItem(content="dummy-node", metadata={"score": 1.0}),
+            RetrieverResultItem(
+                content="dummy-node",
+                metadata={"score": 1.0, "nodeLabels": None, "id": None},
+            ),
         ],
         metadata={"__retriever": "VectorRetriever"},
     )

tests/unit/test_neo4j_queries.py

@@ -22,7 +22,7 @@ def test_vector_search_basic() -> None:
     expected = (
         "CALL db.index.vector.queryNodes($vector_index_name, $top_k, $query_vector) "
         "YIELD node, score "
-        "RETURN node, score"
+        "RETURN node { .*, `None`: null } AS node, labels(node) AS nodeLabels, elementId(node) AS id, score"
     )
     result, params = get_search_query(SearchType.VECTOR)
     assert result.strip() == expected.strip()
@@ -42,7 +42,7 @@ def test_hybrid_search_basic() -> None:
         "RETURN n.node AS node, (n.score / max) AS score "
         "} "
         "WITH node, max(score) AS score ORDER BY score DESC LIMIT $top_k "
-        "RETURN node, score"
+        "RETURN node { .*, `None`: null } AS node, labels(node) AS nodeLabels, elementId(node) AS id, score"
     )
     result, _ = get_search_query(SearchType.HYBRID)
     assert result.strip() == expected.strip()
@@ -78,8 +78,8 @@ def test_vector_search_with_filters(_mock: Any) -> None:
         " AND (True) "
         "WITH node, "
         "vector.similarity.cosine(node.`vector`, $query_vector) AS score "
-        "ORDER BY score DESC LIMIT $top_k"
-        " RETURN node, score"
+        "ORDER BY score DESC LIMIT $top_k "
+        "RETURN node { .*, `vector`: null } AS node, labels(node) AS nodeLabels, elementId(node) AS id, score"
     )
     result, params = get_search_query(
         SearchType.VECTOR,
@@ -104,8 +104,8 @@ def test_vector_search_with_params_from_filters(_mock: Any) -> None:
         " AND (True) "
         "WITH node, "
         "vector.similarity.cosine(node.`vector`, $query_vector) AS score "
-        "ORDER BY score DESC LIMIT $top_k"
-        " RETURN node, score"
+        "ORDER BY score DESC LIMIT $top_k "
+        "RETURN node { .*, `vector`: null } AS node, labels(node) AS nodeLabels, elementId(node) AS id, score"
     )
     result, params = get_search_query(
         SearchType.VECTOR,