feat : 代码梳理 移除所有敏感密钥 通过环境变量方式配置

app/service/chat_robot/script/tools/sql_tools.py

@@ -2,21 +2,15 @@
 """Tools for interacting with a SQL database."""
 from typing import Any, Dict, Optional, Type
 
-from pydantic import BaseModel, Extra, Field, root_validator
-
-from langchain.base_language import BaseLanguageModel
-from langchain.callbacks.manager import (
-    AsyncCallbackManagerForToolRun,
-    CallbackManagerForToolRun,
-)
-from langchain.chains.llm import LLMChain
-from langchain.prompts import PromptTemplate
+from langchain_community.tools.sql_database.prompt import QUERY_CHECKER
+from langchain_community.tools.sql_database.tool import _QuerySQLCheckerToolInput
 # from langchain.sql_database import SQLDatabase
 from langchain_community.utilities import SQLDatabase
-from langchain.tools.base import BaseTool
-from langchain_community.tools.sql_database.prompt import QUERY_CHECKER
-
-from langchain_community.tools.sql_database.tool import QuerySQLCheckerTool, _QuerySQLCheckerToolInput
+from langchain_core.callbacks import CallbackManagerForToolRun, AsyncCallbackManagerForToolRun
+from langchain_core.language_models import BaseLanguageModel
+from langchain_core.prompts import PromptTemplate
+from langchain_core.tools import BaseTool
+from pydantic import BaseModel, Extra, Field, root_validator
 
 
 class BaseSQLDatabaseTool(BaseModel):
@@ -62,7 +56,7 @@ class QuerySQLDataBaseTool(BaseSQLDatabaseTool, BaseTool):
         "LIMIT 1'"
         "Example Input 2: 'SELECT img_name FROM top WHERE sleeve_length = 'Long' AND type = 'Blouse' "
         "order by rand() LIMIT 2'"
-                        )
+    )
 
     def _run(
             self,
@@ -95,9 +89,9 @@ class InfoSQLDatabaseTool(BaseSQLDatabaseTool, BaseTool):
         "Input to this tool is a comma-separated list of tables, output is the schema and sample rows for those tables."
         "There are eight tables covering eight fashion categories: female_top, female_pants, female_dress,"
         "female_skirt, female_outwear, male_bottom, male_top, and male_outwear."
-        
+
         "Example Input: 'female_outwear, male_top'"
-        )
+    )
 
     def _run(
             self,
@@ -183,11 +177,11 @@ class QuerySQLCheckerTool(BaseSQLDatabaseTool, BaseTool):
     args_schema: Type[BaseModel] = _QuerySQLCheckerToolInput
 
     @root_validator(pre=True)
-    def initialize_llm_chain(cls, values: Dict[str, Any]) -> Dict[str, Any]:
+    def initialize_llm_chain(self, values: Dict[str, Any]) -> Dict[str, Any]:
         if "llm_chain" not in values:
             # from langchain.chains.llm import LLMChain
 
-            llm = values.get("llm")   # type: ignore[arg-type]
+            llm = values.get("llm")  # type: ignore[arg-type]
             prompt = PromptTemplate(
                 template=QUERY_CHECKER, input_variables=["dialect", "query"]
             )