feat(search): pagination + search for relationships and entities

ChrisCoder9000 · ChrisCoder9000 · commit de2620d3f2cc · 2025-11-16T13:59:04.000+01:00
diff --git a/src/adapters/graph.py b/src/adapters/graph.py
@@ -10,7 +10,14 @@
 
 from typing import Optional, Tuple
 from src.adapters.interfaces.graph import GraphClient
-from src.constants.kg import IdentificationParams, Node, Predicate, Triple
+from src.constants.kg import (
+    IdentificationParams,
+    Node,
+    Predicate,
+    SearchEntitiesResult,
+    SearchRelationshipsResult,
+    Triple,
+)
 
 
 class GraphAdapter:
@@ -190,17 +197,47 @@ def get_connected_nodes(
             node=node, uuids=uuids, limit=limit, with_labels=with_labels
         )
 
-    def search_relationships(self, limit: int = 10, skip: int = 0) -> list[Triple]:
+    def search_relationships(
+        self,
+        limit: int = 10,
+        skip: int = 0,
+        relationship_types: Optional[list[str]] = None,
+        from_node_labels: Optional[list[str]] = None,
+        to_node_labels: Optional[list[str]] = None,
+        query_text: Optional[str] = None,
+        query_search_target: Optional[str] = "all",
+    ) -> SearchRelationshipsResult:
         """
         Search the relationships of the graph.
         """
-        return self.graph.search_relationships(limit, skip)
+        relationship_uuids = []
+        # TODO: semantic search + src/core/agents/tools/kg_agent/KGAgentAddTripletsTool.py:165
+        return self.graph.search_relationships(
+            limit,
+            skip,
+            relationship_types,
+            from_node_labels,
+            to_node_labels,
+            relationship_uuids,
+            query_text,
+            query_search_target,
+        )
 
-    def search_entities(self, limit: int = 10, skip: int = 0) -> list[Node]:
+    def search_entities(
+        self,
+        limit: int = 10,
+        skip: int = 0,
+        node_labels: Optional[list[str]] = None,
+        query_text: Optional[str] = None,
+    ) -> SearchEntitiesResult:
         """
         Search the entities of the graph.
         """
-        return self.graph.search_entities(limit, skip)
+        node_uuids = []
+        # TODO: semantic search
+        return self.graph.search_entities(
+            limit, skip, node_labels, node_uuids, query_text
+        )
 
 
 _graph_adapter = GraphAdapter()
diff --git a/src/adapters/interfaces/graph.py b/src/adapters/interfaces/graph.py
@@ -11,7 +11,13 @@
 from abc import ABC, abstractmethod
 from typing import Optional, Tuple
 
-from src.constants.kg import IdentificationParams, Node, Predicate, Triple
+from src.constants.kg import (
+    IdentificationParams,
+    Node,
+    Predicate,
+    SearchEntitiesResult,
+    SearchRelationshipsResult,
+)
 
 
 class GraphClient(ABC):
@@ -196,14 +202,34 @@ def get_connected_nodes(
         raise NotImplementedError("get_connected_nodes method not implemented")
 
     @abstractmethod
-    def search_relationships(self, limit: int = 10, skip: int = 0) -> list[Triple]:
+    def search_relationships(
+        self,
+        limit: int = 10,
+        skip: int = 0,
+        relationship_types: Optional[list[str]] = None,
+        from_node_labels: Optional[list[str]] = None,
+        to_node_labels: Optional[list[str]] = None,
+        relationship_uuids: Optional[list[str]] = None,
+        query_text: Optional[str] = None,
+        query_search_target: Optional[
+            str
+        ] = "all",  # Search into the relationship desc or node names or relationship desc
+    ) -> SearchRelationshipsResult:
         """
         Search the relationships of the graph.
         """
         raise NotImplementedError("search_relationships method not implemented")
 
     @abstractmethod
-    def search_entities(self, limit: int = 10, skip: int = 0) -> list[Node]:
+    def search_entities(
+        self,
+        limit: int = 10,
+        skip: int = 0,
+        node_labels: Optional[list[str]] = None,
+        node_uuids: Optional[list[str]] = None,
+        query_text: Optional[str] = None,
+    ) -> SearchEntitiesResult:
         """
         Search the entities of the graph.
         """
+        raise NotImplementedError("search_entities method not implemented")
diff --git a/src/constants/kg.py b/src/constants/kg.py
@@ -111,3 +111,21 @@ class IdentificationParams(BaseModel):
     )
 
     model_config = ConfigDict(extra="allow")
+
+
+class SearchRelationshipsResult(BaseModel):
+    """
+    Search relationships result model.
+    """
+
+    results: List[Triple]
+    total: int
+
+
+class SearchEntitiesResult(BaseModel):
+    """
+    Search entities result model.
+    """
+
+    results: List[Node]
+    total: int
diff --git a/src/core/agents/tools/kg_agent/KGAgentAddTripletsTool.py b/src/core/agents/tools/kg_agent/KGAgentAddTripletsTool.py
@@ -162,6 +162,8 @@ def _run(self, *args, **kwargs) -> str:
                 triplet.object,
             )
 
+            # TODO: save natural language relationship/predicate+s+o vector
+
             # TODO: add changelog relationship created
 
         return f"Triplets added successfully: {triplets}"
diff --git a/src/core/search/entities.py b/src/core/search/entities.py
@@ -8,13 +8,19 @@
 -----
 """
 
-from src.constants.kg import Node
+from typing import Optional
+from src.constants.kg import SearchEntitiesResult
 from src.services.kg_agent.main import graph_adapter
 
 
-def search_entities(limit: int = 10, skip: int = 0) -> list[Node]:
+def search_entities(
+    limit: int = 10,
+    skip: int = 0,
+    node_labels: Optional[list[str]] = None,
+    query_text: Optional[str] = None,
+) -> SearchEntitiesResult:
     """
     Search the entities of the graph.
     """
-    result = graph_adapter.search_entities(limit, skip)
+    result = graph_adapter.search_entities(limit, skip, node_labels, query_text)
     return result
diff --git a/src/core/search/relationships.py b/src/core/search/relationships.py
@@ -8,13 +8,30 @@
 -----
 """
 
-from src.constants.kg import Triple
+from typing import Optional
+from src.constants.kg import SearchRelationshipsResult
 from src.services.kg_agent.main import graph_adapter
 
 
-def search_relationships(limit: int = 10, skip: int = 0) -> list[Triple]:
+def search_relationships(
+    limit: int = 10,
+    skip: int = 0,
+    relationship_types: Optional[list[str]] = None,
+    from_node_labels: Optional[list[str]] = None,
+    to_node_labels: Optional[list[str]] = None,
+    query_text: Optional[str] = None,
+    query_search_target: Optional[str] = "all",
+) -> SearchRelationshipsResult:
     """
     Search the relationships of the graph.
     """
-    result = graph_adapter.search_relationships(limit, skip)
+    result = graph_adapter.search_relationships(
+        limit,
+        skip,
+        relationship_types,
+        from_node_labels,
+        to_node_labels,
+        query_text,
+        query_search_target,
+    )
     return result
diff --git a/src/lib/neo4j/client.py b/src/lib/neo4j/client.py
@@ -13,7 +13,14 @@
 from neo4j import GraphDatabase
 from src.adapters.interfaces.graph import GraphClient
 from src.config import config
-from src.constants.kg import IdentificationParams, Node, Predicate, Triple
+from src.constants.kg import (
+    IdentificationParams,
+    Node,
+    Predicate,
+    SearchEntitiesResult,
+    SearchRelationshipsResult,
+    Triple,
+)
 from src.utils.logging import log
 
 
@@ -650,19 +657,79 @@ def get_connected_nodes(
             for record in result.records
         ]
 
-    def search_relationships(self, limit: int = 10, skip: int = 0) -> list[Triple]:
+    def search_relationships(
+        self,
+        limit: int = 10,
+        skip: int = 0,
+        relationship_types: Optional[list[str]] = None,
+        from_node_labels: Optional[list[str]] = None,
+        to_node_labels: Optional[list[str]] = None,
+        relationship_uuids: Optional[list[str]] = None,
+        query_text: Optional[str] = None,
+        query_search_target: Optional[str] = "all",
+    ) -> SearchRelationshipsResult:
         """
         Search the relationships of the graph.
         """
+        filters = []
+        if relationship_types:
+            filters.append(
+                "type(r) IN [" + ",".join(f"'{t}'" for t in relationship_types) + "]"
+            )
+        if from_node_labels:
+            filters.append(
+                "ANY(lbl IN labels(n) WHERE lbl IN ["
+                + ",".join(
+                    f"'{label}'" for label in self._clean_labels(from_node_labels)
+                )
+                + "])"
+            )
+        if to_node_labels:
+            filters.append(
+                "ANY(lbl IN labels(m) WHERE lbl IN ["
+                + ",".join(f"'{label}'" for label in self._clean_labels(to_node_labels))
+                + "])"
+            )
+        if relationship_uuids:
+            filters.append(
+                "r.uuid IN [" + ",".join(f"'{u}'" for u in relationship_uuids) + "]"
+            )
+        if query_text:
+            if query_search_target == "all":
+                filters.append(
+                    f"(toLower(coalesce(n.name, n.Name, '')) CONTAINS toLower('{query_text}') OR "
+                    f"toLower(coalesce(m.name, m.Name, '')) CONTAINS toLower('{query_text}') OR "
+                    f"toLower(coalesce(r.description, r.Description, '')) CONTAINS toLower('{query_text}'))"
+                )
+            elif query_search_target == "node_name":
+                filters.append(f"toLower(n.name) CONTAINS toLower('{query_text}')")
+            elif query_search_target == "relationship_description":
+                filters.append(
+                    f"toLower(r.description) CONTAINS toLower('{query_text}')"
+                )
+            elif query_search_target == "relationship_name":
+                filters.append(f"toLower(r.name) CONTAINS toLower('{query_text}')")
         cypher_query = f"""
-        MATCH (n)-[r]-(m)
-        RETURN n.uuid as n_uuid, n.name as n_name, labels(n) as n_labels, n.description as n_description, properties(n) as n_properties,
-               r as rel, type(r) as rel_type, r.description as rel_description,
-               m.uuid as m_uuid, m.name as m_name, labels(m) as m_labels, m.description as m_description, properties(m) as m_properties
+        MATCH (n)-[r]->(m)
+        {"WHERE " + " AND ".join(filters) if filters else ""}
+        RETURN n.uuid AS n_uuid, n.name AS n_name, labels(n) AS n_labels,
+            n.description AS n_description, properties(n) AS n_properties,
+            r AS rel, type(r) AS rel_type, r.description AS rel_description,
+            m.uuid AS m_uuid, m.name AS m_name, labels(m) AS m_labels,
+            m.description AS m_description, properties(m) AS m_properties
         SKIP {skip}
         LIMIT {limit}
         """
+        cypher_count = """
+        MATCH ()-[r]-()
+        RETURN count(r) AS total
+        """
         result = self.driver.execute_query(cypher_query)
+        count_result = self.driver.execute_query(cypher_count)
+        total = 0
+        if count_result and count_result.records:
+            total = count_result.records[0].get("total") or 0
+
         triples: list[Triple] = []
         for record in result.records:
             relationship = record.get("rel")
@@ -745,19 +812,49 @@ def search_relationships(self, limit: int = 10, skip: int = 0) -> list[Triple]:
                     ),
                 )
             )
-        return triples
+        return SearchRelationshipsResult(results=triples, total=total)
 
-    def search_entities(self, limit: int = 10, skip: int = 0) -> list[Node]:
+    def search_entities(
+        self,
+        limit: int = 10,
+        skip: int = 0,
+        node_labels: Optional[list[str]] = None,
+        node_uuids: Optional[list[str]] = None,
+        query_text: Optional[str] = None,
+    ) -> SearchEntitiesResult:
         """
         Search the entities of the graph.
         """
+        filters = []
+        if node_labels:
+            filters.append(
+                "ANY(lbl IN labels(n) WHERE lbl IN ["
+                + ",".join(f"'{label}'" for label in self._clean_labels(node_labels))
+                + "])"
+            )
+        if node_uuids:
+            filters.append(f"n.uuid IN [{','.join(node_uuids)}]")
+        if query_text:
+            filters.append(
+                f"(toLower(coalesce(n.name, n.Name, '')) CONTAINS toLower('{query_text}'))"
+            )
         cypher_query = f"""
         MATCH (n)
+        {"WHERE " + " AND ".join(filters) if filters else ""}
         RETURN n.uuid as uuid, n.name as name, labels(n) as labels, n.description as description, properties(n) as properties
         SKIP {skip}
         LIMIT {limit}
         """
+        cypher_count = """
+        MATCH (n)
+        RETURN count(n) AS total
+        """
         result = self.driver.execute_query(cypher_query)
+        count_result = self.driver.execute_query(cypher_count)
+        total = 0
+        if count_result and count_result.records:
+            total = count_result.records[0].get("total") or 0
+
         nodes: list[Node] = []
         for record in result.records:
             properties_record = record.get("properties") or {}
@@ -784,7 +881,7 @@ def search_entities(self, limit: int = 10, skip: int = 0) -> list[Node]:
                     properties=properties,
                 )
             )
-        return nodes
+        return SearchEntitiesResult(results=nodes, total=total)
 
 
 _neo4j_client = Neo4jClient()
diff --git a/src/services/api/controllers/retrieve.py b/src/services/api/controllers/retrieve.py
diff --git a/src/services/api/routes/retrieve.py b/src/services/api/routes/retrieve.py

Original file line number	Diff line number	Diff line change
`@@ -162,6 +162,8 @@ def _run(self, args, *kwargs) -> str:`
`162`	`162`	`triplet.object,`
`163`	`163`	`)`
`164`	`164`
	`165`	`+ # TODO: save natural language relationship/predicate+s+o vector`
	`166`	`+`
`165`	`167`	`# TODO: add changelog relationship created`
`166`	`168`
`167`	`169`	`return f"Triplets added successfully: {triplets}"`