move some annotations.

eosphoros-ai · Jan 20, 2025 · 1f0f401 · 1f0f401
1 parent 0896c52
commit 1f0f401
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 4 deletions.
diff --git a/dbgpt/storage/knowledge_graph/graph_retriever/document_graph_retriever.py b/dbgpt/storage/knowledge_graph/graph_retriever/document_graph_retriever.py
@@ -29,9 +29,9 @@ async def retrieve(
         self, input: Union[Graph, List[str], List[List[float]]]
     ) -> Tuple[Graph, None]:
         """Retrieve from document graph."""
+        # If retrieve subgraph from triplet graph successfully
         if isinstance(input, Graph):
-            # If retrieve subgraph from triplet graph successfully
-            # Using the vids to search chunks and doc
+            # Get entities' vids from triplet subgraph
             keywords_for_document_graph = []
             for vertex in input.vertices():
                 keywords_for_document_graph.append(vertex.name)

diff --git a/dbgpt/storage/knowledge_graph/graph_retriever/graph_retriever.py b/dbgpt/storage/knowledge_graph/graph_retriever/graph_retriever.py
@@ -124,6 +124,7 @@ async def retrieve(self, text: str) -> Tuple[Graph, Tuple[Graph, str]]:
 
             # Using subs to transfer keywords or embeddings
             subs: Union[List[str], List[List[float]]]
+
             if self._enable_similarity_search:
                 # Embedding the question
                 vector = await self._text_embedder.embed(text)
@@ -133,18 +134,21 @@ async def retrieve(self, text: str) -> Tuple[Graph, Tuple[Graph, str]]:
                 )
                 # Using the embeddings of keywords and question
                 vectors.append(vector)
+                # Using vectors as subs
                 subs = vectors
                 logger.info(
                     "Search subgraph with the following keywords and question's "
                     f"embedding vector:\n[KEYWORDS]:{keywords}\n[QUESTION]:{text}"
                 )
             else:
+                # Using keywords as subs
                 subs = keywords
                 logger.info(
                     "Search subgraph with the following keywords:\n"
                     f"[KEYWORDS]:{keywords}"
                 )
 
+            # If enable triplet graph
             if self._triplet_graph_enabled:
                 # Retrieve from triplet graph
                 if self._enable_similarity_search:
@@ -153,10 +157,12 @@ async def retrieve(self, text: str) -> Tuple[Graph, Tuple[Graph, str]]:
                 else:
                     # Retrieve from triplet graph with keywords
                     subgraph = await self._keyword_based_graph_retriever.retrieve(subs)
+
+            # If enable document graph
             if self._document_graph_enabled:
                 # Retrieve from document graph
+                # If not enable triplet graph or failed to retrieve subgraph
                 if subgraph.vertex_count == 0 and subgraph.edge_count == 0:
-                    # If not enable triplet graph or failed to retrieve subgraph
                     # Using subs to retrieve from document graph
                     subgraph_for_doc = await self._document_graph_retriever.retrieve(
                         subs
@@ -165,7 +171,7 @@ async def retrieve(self, text: str) -> Tuple[Graph, Tuple[Graph, str]]:
                     # If retrieve subgraph from triplet graph successfully
                     # Using entities in subgraph to search chunks and doc
                     subgraph_for_doc = await self._document_graph_retriever.retrieve(
-                        triplet_graph=subgraph
+                        subgraph
                     )
 
         return subgraph, (subgraph_for_doc, text2gql_query)