enhance getting chunks by relationship ids

Author: IANTHEREAL

backend/app/rag/indices/knowledge_graph/graph_store/tidb_graph_store.py

@@ -40,6 +40,7 @@
 from app.models import (
     KnowledgeBase,
     EntityType,
+    Document,
 )
 
 logger = logging.getLogger(__name__)
@@ -312,7 +313,7 @@ def create_relationship(
             self._session.flush()
 
     def get_subgraph_by_relationship_ids(
-            self, ids: list[int], **kwargs
+        self, ids: list[int], **kwargs
     ) -> RetrievedKnowledgeGraph:
         stmt = (
             select(self._relationship_model)
@@ -1081,39 +1082,57 @@ def get_chunks_by_relationships(
         }
 
         if not chunk_ids:
-            return []
+            # Query chunks
+            chunks = session.exec(
+                select(self._chunk_model).where(self._chunk_model.id.in_(chunk_ids))
+            ).all()
 
-        # Query chunks
-        chunks = session.exec(
-            select(self._chunk_model).where(self._chunk_model.id.in_(chunk_ids))
+            return [
+                {
+                    "id": chunk.id,
+                    "text": chunk.text,
+                    "document_id": chunk.document_id,
+                    "meta": {
+                        "language": chunk.meta.get("language"),
+                        "product": chunk.meta.get("product"),
+                        "resource": chunk.meta.get("resource"),
+                        "source_uri": chunk.meta.get("source_uri"),
+                        "tidb_version": chunk.meta.get("tidb_version"),
+                    },
+                }
+                for chunk in chunks
+            ]
+
+        document_ids = {
+            rel.meta.get("document_id")
+            for rel in relationships
+            if rel.meta.get("document_id") is not None
+        }
+
+        documents = session.exec(
+            select(Document).where(Document.id.in_(document_ids))
         ).all()
 
         return [
             {
-                "id": chunk.id,
-                "text": chunk.text,
-                "document_id": chunk.document_id,
-                "meta": {
-                    "language": chunk.meta.get("language"),
-                    "product": chunk.meta.get("product"),
-                    "resource": chunk.meta.get("resource"),
-                    "source_uri": chunk.meta.get("source_uri"),
-                    "tidb_version": chunk.meta.get("tidb_version"),
-                },
+                "id": doc.id,
+                "text": doc.content,
+                "document_id": doc.id,
+                "meta": doc.meta,
             }
-            for chunk in chunks
+            for doc in documents
         ]
-    
+
     def get_entire_knowledge_graph(self) -> RetrievedKnowledgeGraph:
         """Retrieve all entities and relationships from the knowledge graph store.
-        
+
         Returns:
             RetrievedKnowledgeGraph containing all entities and relationships
         """
         # Query all entities
         entity_query = select(self._entity_model).order_by(self._entity_model.id)
         db_entities = self._session.exec(entity_query).all()
-        
+
         # Query all relationships with their related entities
         relationship_query = (
             select(self._relationship_model)
@@ -1124,7 +1143,7 @@ def get_entire_knowledge_graph(self) -> RetrievedKnowledgeGraph:
             .order_by(self._relationship_model.id)
         )
         db_relationships = self._session.exec(relationship_query).all()
-        
+
         # Convert entities to RetrievedEntity objects
         entities = []
         for entity in db_entities:
@@ -1138,7 +1157,7 @@ def get_entire_knowledge_graph(self) -> RetrievedKnowledgeGraph:
                     entity_type=entity.entity_type,
                 )
             )
-        
+
         # Convert relationships to RetrievedRelationship objects
         relationships = []
         for rel in db_relationships: