solve some comments.

eosphoros-ai · Jan 16, 2025 · 3e934c2 · 3e934c2
1 parent e78c030
commit 3e934c2
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 6 deletions.
diff --git a/dbgpt/rag/transformer/intent_interpreter.py b/dbgpt/rag/transformer/intent_interpreter.py
@@ -2,7 +2,7 @@
 import json
 import logging
 import re
-from typing import Dict, List
+from typing import Dict, List, Union
 
 from dbgpt.core import BaseMessage, HumanPromptTemplate, LLMClient
 from dbgpt.rag.transformer.llm_translator import LLMTranslator
@@ -110,7 +110,7 @@ def _parse_response(self, text: str) -> Dict:
         else:
             text = ""
 
-        intention = dict()
+        intention: Dict[str, Union[str, List[str]]] = dict()
         intention = json.loads(text)
         if "category" not in intention:
             intention["category"] = ""

diff --git a/dbgpt/rag/transformer/text2gql.py b/dbgpt/rag/transformer/text2gql.py
@@ -101,7 +101,7 @@ def _format_messages(self, text: str, history: str = None) -> List[BaseMessage]:
 
     def _parse_response(self, text: str) -> Dict:
         """Parse llm response."""
-        translation = {}
+        translation: Dict[str, str] = {}
         query = ""
 
         code_block_pattern = re.compile(r"```cypher(.*?)```", re.S)

diff --git a/dbgpt/storage/knowledge_graph/community_summary.py b/dbgpt/storage/knowledge_graph/community_summary.py
@@ -4,7 +4,7 @@
 import logging
 import os
 import uuid
-from typing import List, Optional, Tuple, Union
+from typing import Dict, List, Optional, Tuple, Union
 
 from dbgpt._private.pydantic import ConfigDict, Field
 from dbgpt.core import Chunk, LLMClient
@@ -542,12 +542,16 @@ async def asimilar_search_with_scores(
 
         # if enable text2gql search, use translated query to retrieve subgraph
         if self._graph_store.enable_text_search:
-            intention = await self._intent_interpreter.translate(text)
+            intention: Dict[
+                str, Union[str, List[str]]
+            ] = await self._intent_interpreter.translate(text)
             schema = json.dumps(
                 json.loads(self._graph_store_apdater.get_schema()), indent=4
             )
             intention["schema"] = schema
-            translation = await self._text2gql.translate(json.dumps(intention))
+            translation: Dict[str, str] = await self._text2gql.translate(
+                json.dumps(intention)
+            )
             try:
                 query = translation["query"]
                 if "LIMIT" not in query: