Add MetadataFilters to neo4j_property_graph (#14362)

docs/docs/examples/property_graph/property_graph_advanced.ipynb

@@ -47,7 +47,25 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "--2024-06-26 11:12:16--  https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt\n",
+      "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.110.133, 185.199.109.133, 185.199.111.133, ...\n",
+      "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.110.133|:443... connected.\n",
+      "HTTP request sent, awaiting response... 200 OK\n",
+      "Length: 75042 (73K) [text/plain]\n",
+      "Saving to: ‘data/paul_graham/paul_graham_essay.txt’\n",
+      "\n",
+      "data/paul_graham/pa 100%[===================>]  73.28K  --.-KB/s    in 0.007s  \n",
+      "\n",
+      "2024-06-26 11:12:16 (10.4 MB/s) - ‘data/paul_graham/paul_graham_essay.txt’ saved [75042/75042]\n",
+      "\n"
+     ]
+    }
+   ],
    "source": [
     "!mkdir -p 'data/paul_graham/'\n",
     "!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'"
@@ -106,6 +124,15 @@
     "    \"PLACE\": [\"HAS\", \"PART_OF\", \"WORKED_AT\"],\n",
     "    \"ORGANIZATION\": [\"HAS\", \"PART_OF\", \"WORKED_WITH\"],\n",
     "}\n",
+    "validation_schema = [\n",
+    "    (\"ORGANIZATION\", \"HAS\", \"PERSON\"),\n",
+    "    (\"PERSON\", \"WORKED_AT\", \"ORGANIZATION\"),\n",
+    "    (\"PERSON\", \"WORKED_WITH\", \"PERSON\"),\n",
+    "    (\"PERSON\", \"WORKED_ON\", \"ORGANIZATION\"),\n",
+    "    (\"PERSON\", \"PART_OF\", \"ORGANIZATION\"),\n",
+    "    (\"ORGANIZATION\", \"PART_OF\", \"ORGANIZATION\"),\n",
+    "    (\"PERSON\", \"WORKED_AT\", \"PLACE\"),\n",
+    "]\n",
     "\n",
     "kg_extractor = SchemaLLMPathExtractor(\n",
     "    llm=Ollama(model=\"llama3\", json_mode=True, request_timeout=3600),\n",

llama-index-core/llama_index/core/indices/property_graph/sub_retrievers/text_to_cypher.py

@@ -72,8 +72,8 @@ def _clean_query_output(self, query_output: Any) -> Any:
             filtered = {}
             for key, value in query_output.items():
                 if (
-                    key in self.allowed_output_fields
-                    or self.allowed_output_fields is None
+                    self.allowed_output_fields is None
+                    or key in self.allowed_output_fields
                 ):
                     filtered[key] = value
                 elif isinstance(value, (dict, list)):
@@ -101,6 +101,7 @@ def retrieve_from_graph(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
             question=question,
         )
 
+        parsed_cypher_query = response
         if self.allowed_output_fields is not None:
             parsed_cypher_query = self._parse_generated_cyher(response)
 
@@ -134,6 +135,7 @@ async def aretrieve_from_graph(
             question=question,
         )
 
+        parsed_cypher_query = response
         if self.allowed_output_fields is not None:
             parsed_cypher_query = self._parse_generated_cyher(response)
 

llama-index-core/llama_index/core/indices/property_graph/sub_retrievers/vector.py

@@ -11,7 +11,11 @@
 )
 from llama_index.core.settings import Settings
 from llama_index.core.schema import BaseNode, NodeWithScore, QueryBundle
-from llama_index.core.vector_stores.types import VectorStoreQuery, VectorStore
+from llama_index.core.vector_stores.types import (
+    VectorStoreQuery,
+    VectorStore,
+    MetadataFilters,
+)
 
 
 class VectorContextRetriever(BasePGRetriever):
@@ -41,13 +45,15 @@ def __init__(
         vector_store: Optional[VectorStore] = None,
         similarity_top_k: int = 4,
         path_depth: int = 1,
+        filters: Optional[MetadataFilters] = None,
         **kwargs: Any
     ) -> None:
         self._retriever_kwargs = kwargs or {}
         self._embed_model = embed_model or Settings.embed_model
         self._similarity_top_k = similarity_top_k
         self._vector_store = vector_store
         self._path_depth = path_depth
+        self._filters = filters
 
         super().__init__(graph_store=graph_store, include_text=include_text, **kwargs)
 
@@ -60,6 +66,7 @@ def _get_vector_store_query(self, query_bundle: QueryBundle) -> VectorStoreQuery
         return VectorStoreQuery(
             query_embedding=query_bundle.embedding,
             similarity_top_k=self._similarity_top_k,
+            filters=self._filters,
             **self._retriever_kwargs,
         )
 
@@ -80,6 +87,7 @@ async def _aget_vector_store_query(
         return VectorStoreQuery(
             query_embedding=query_bundle.embedding,
             similarity_top_k=self._similarity_top_k,
+            filters=self._filters,
             **self._retriever_kwargs,
         )
 

llama-index-integrations/graph_stores/llama-index-graph-stores-neo4j/llama_index/graph_stores/neo4j/neo4j_property_graph.py

@@ -528,14 +528,26 @@ def vector_query(
         self, query: VectorStoreQuery, **kwargs: Any
     ) -> Tuple[List[LabelledNode], List[float]]:
         """Query the graph store with a vector store query."""
+        conditions = None
+        if query.filters:
+            conditions = [
+                f"e.{filter.key} {filter.operator.value} {filter.value}"
+                for filter in query.filters.filters
+            ]
+        filters = (
+            f" {query.filters.condition.value} ".join(conditions).replace("==", "=")
+            if conditions is not None
+            else "1 = 1"
+        )
+
         data = self.structured_query(
-            """MATCH (e:`__Entity__`)
-            WHERE e.embedding IS NOT NULL AND size(e.embedding) = $dimension
+            f"""MATCH (e:`__Entity__`)
+            WHERE e.embedding IS NOT NULL AND size(e.embedding) = $dimension AND ({filters})
             WITH e, vector.similarity.cosine(e.embedding, $embedding) AS score
             ORDER BY score DESC LIMIT toInteger($limit)
             RETURN e.id AS name,
                [l in labels(e) WHERE l <> '__Entity__' | l][0] AS type,
-               e{.* , embedding: Null, name: Null, id: Null} AS properties,
+               e{{.* , embedding: Null, name: Null, id: Null}} AS properties,
                score""",
             param_map={
                 "embedding": query.query_embedding,

llama-index-integrations/graph_stores/llama-index-graph-stores-neo4j/pyproject.toml

@@ -28,7 +28,7 @@ exclude = ["**/BUILD"]
 license = "MIT"
 name = "llama-index-graph-stores-neo4j"
 readme = "README.md"
-version = "0.2.5"
+version = "0.2.6"
 
 [tool.poetry.dependencies]
 python = ">=3.8.1,<4.0"