From 1f0f401c65dfd3c93cd0e60c274900444590edce Mon Sep 17 00:00:00 2001 From: Lyu Songlin Date: Mon, 20 Jan 2025 16:41:27 +0800 Subject: [PATCH] move some annotations. --- .../graph_retriever/document_graph_retriever.py | 4 ++-- .../knowledge_graph/graph_retriever/graph_retriever.py | 10 ++++++++-- 2 files changed, 10 insertions(+), 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 index 13441a286..811686a16 100644 --- 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 index 5ce549f4a..8979b4b67 100644 --- 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)