openai 替换为 通义千问

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

@@ -1,9 +1,13 @@
 """Callback Handler that add on_chain_end function to record Token usage."""
 from typing import Any, Dict
 
-from langchain.callbacks import OpenAICallbackHandler
+from langchain_community.callbacks.openai_info import OpenAICallbackHandler
 from langchain.schema import LLMResult
-from langchain.callbacks.openai_info import standardize_model_name, MODEL_COST_PER_1K_TOKENS, get_openai_token_cost_for_model
+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
 
 
 class OpenAITokenRecordCallbackHandler(OpenAICallbackHandler):

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

@@ -0,0 +1,28 @@
+from typing import Dict
+
+from dashscope.api_entities.dashscope_response import GenerationUsage
+
+
+class QWenCallbackHandler:
+
+    total_tokens: int = 0
+    input_tokens: int = 0
+    output_tokens: int = 0
+    total_cost: float = 0.0
+
+    def on_llm_end(self, response: GenerationUsage) -> None:
+        """Collect token usage."""
+
+
+        self.input_tokens += response.input_tokens
+        self.output_tokens += response.output_tokens
+        self.total_tokens = self.input_tokens + self.output_tokens
+        self.total_cost = 0.04 * self.input_tokens / 1000 + 0.12 * self.output_tokens / 1000
+
+    def on_chain_end(self, outputs: Dict ) -> None:
+        """Write token usage to redis."""
+        outputs["total_tokens"] = self.total_tokens
+        outputs["total_cost"] = self.total_cost
+        outputs["prompt_tokens"] = self.input_tokens
+        outputs["completion_tokens"] = self.output_tokens
+        print("input_tokens : {} \noutput_tokens : {}".format(outputs["prompt_tokens"] , outputs["completion_tokens"]))