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

app/service/chat_robot/script/callbacks/openai_token_record_callback.py

@@ -2,18 +2,16 @@
 from typing import Any, Dict
 
 from langchain_community.callbacks.openai_info import OpenAICallbackHandler
-from langchain.schema import LLMResult
 from langchain_community.callbacks.openai_info import standardize_model_name, MODEL_COST_PER_1K_TOKENS, \
     get_openai_token_cost_for_model
-
-
-# from langchain.callbacks.openai_info import standardize_model_name, MODEL_COST_PER_1K_TOKENS, get_openai_token_cost_for_model
+from langchain_core.outputs import LLMResult
 
 
 class OpenAITokenRecordCallbackHandler(OpenAICallbackHandler):
     need_record: bool = True
     response_type: str = "string"
     """Callback Handler that tracks OpenAI info and write to redis after agent finish"""
+
     def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
         """Collect token usage."""
         if response.llm_output is None:
@@ -22,7 +20,7 @@ class OpenAITokenRecordCallbackHandler(OpenAICallbackHandler):
         if "token_usage" not in response.llm_output:
             return None
         if "function_call" in response.generations[0][0].message.additional_kwargs:
-            if response.generations[0][0].message.additional_kwargs["function_call"]["name"] in ["sql_db_query", "sql_db_schema","tutorial_tool"]:
+            if response.generations[0][0].message.additional_kwargs["function_call"]["name"] in ["sql_db_query", "sql_db_schema", "tutorial_tool"]:
                 self.need_record = False
             if response.generations[0][0].message.additional_kwargs["function_call"]["name"] == "sql_db_query":
                 self.response_type = "image"
@@ -39,6 +37,7 @@ class OpenAITokenRecordCallbackHandler(OpenAICallbackHandler):
         self.total_tokens += token_usage.get("total_tokens", 0)
         self.prompt_tokens += prompt_tokens
         self.completion_tokens += completion_tokens
+        return None
 
     def on_chain_end(self, outputs: Dict, **kwargs: Any) -> None:
         """Write token usage to redis."""