update:加入测试速度功能

core/utils/llm.py

@@ -2,6 +2,10 @@ import json
 import logging
 import openai
 import requests
+from datetime import datetime
+from core.utils.util import is_segment
+from core.utils.util import get_string_no_punctuation_or_emoji
+from core.utils.util import read_config, get_project_dir
 from abc import ABC, abstractmethod
 
 logger = logging.getLogger(__name__)
@@ -9,7 +13,7 @@ logger = logging.getLogger(__name__)
 
 class LLM(ABC):
     @abstractmethod
-    def response(self, conn, dialogue):
+    def response(self, session_id, dialogue):
         """LLM response generator"""
         pass
 
@@ -21,7 +25,7 @@ class DeepSeekLLM(LLM):
         self.base_url = config.get("url")
         self.client = openai.OpenAI(api_key=self.api_key, base_url=self.base_url)
 
-    def response(self, conn, dialogue):
+    def response(self, session_id, dialogue):
         logger.info(f"Generating response using {dialogue}")
         try:
             responses = self.client.chat.completions.create(
@@ -47,7 +51,7 @@ class ChatGLMLLM(LLM):
         self.base_url = config.get("url")
         self.client = openai.OpenAI(api_key=self.api_key, base_url=self.base_url)
 
-    def response(self, conn, dialogue):
+    def response(self, session_id, dialogue):
         try:
             responses = self.client.chat.completions.create(
                 model=self.model_name,
@@ -64,13 +68,36 @@ class ChatGLMLLM(LLM):
         except Exception as e:
             logger.error(f"Error in response generation: {e}")
 
+class AliLLM(LLM):
+    def __init__(self, config):
+        self.model_name = config.get("model_name")
+        self.api_key = config.get("api_key")
+        self.base_url = config.get("base_url")
+        self.client = openai.OpenAI(api_key=self.api_key, base_url=self.base_url)
+
+    def response(self, session_id, dialogue):
+        try:
+            responses = self.client.chat.completions.create(
+                model=self.model_name,
+                messages=dialogue,
+                stream=True
+            )
+            for chunk in responses:
+                # 检查是否存在有效的choice且content不为空
+                if chunk.choices and len(chunk.choices) > 0:
+                    delta = chunk.choices[0].delta
+                    content = getattr(delta, 'content', '')
+                    if content:  # 仅在content非空时生成
+                        yield content
+        except Exception as e:
+            logger.error(f"Error in response generation: {e}")
 
 class DifyLLM(LLM):
     def __init__(self, config):
         self.api_key = config["api_key"]
         self.base_url = config.get("base_url", "https://api.dify.ai/v1").rstrip('/')
 
-    def response(self, conn,dialogue):
+    def response(self, session_id, dialogue):
         try:
             # 取最后一条用户消息
             last_msg = next(m for m in reversed(dialogue) if m["role"] == "user")
@@ -82,7 +109,7 @@ class DifyLLM(LLM):
                     json={
                         "query": last_msg["content"],
                         "response_mode": "streaming",
-                        "user": conn.session_id,
+                        "user": session_id,
                         "inputs": {}
                     },
                     stream=True
@@ -103,9 +130,45 @@ def create_instance(class_name, *args, **kwargs):
         "DeepSeekLLM": DeepSeekLLM,
         "ChatGLMLLM": ChatGLMLLM,
         "DifyLLM": DifyLLM,
+        "AliLLM": AliLLM,
         # 可扩展其他LLM实现
     }
 
     if cls := cls_map.get(class_name):
         return cls(*args, **kwargs)
     raise ValueError(f"不支持的LLM类型: {class_name}")
+
+
+if __name__ == "__main__":
+    """
+      响应速度测试
+    """
+    config = read_config(get_project_dir() + "config.yaml")
+    llm = create_instance(
+        config["selected_module"]["LLM"],
+        config["LLM"][config["selected_module"]["LLM"]]
+    )
+
+    start_time = datetime.now()
+
+    dialogue = []
+    dialogue.append({"role": "system", "content": config.get("prompt")})
+    dialogue.append({"role": "user", "content": "你好小智"})
+    llm_responses = llm.response("test", dialogue)
+    response_message = []
+    first_text = None
+    start = 0
+
+    for content in llm_responses:
+        response_message.append(content)
+
+        if is_segment(response_message):
+            segment_text = "".join(response_message[start:])
+            segment_text = get_string_no_punctuation_or_emoji(segment_text)
+            if len(segment_text) > 0:
+                if first_text is None:
+                    first_text = segment_text
+                    print("大模型首次返回耗时：" + str(datetime.now() - start_time))
+                start = len(response_message)
+
+    print("大模型返回总耗时：" + str(datetime.now() - start_time))