union-app-chat-stream 1.0.3

package/app/service/chat_service.py

@@ -0,0 +1,228 @@
+from pathlib import Path
+from typing import Dict, Generator, List
+
+import yaml
+
+from zai import ZhipuAiClient
+
+from app.models.schemas import ChatResponse
+from app.utils.function_utils import ToolContext, call_function
+from loguru import logger
+
+
+def _preview(text: str, limit: int = 300) -> str:
+    return str(text).replace("\n", " ")[:limit]
+
+
+def _load_prompts() -> Dict[str, str]:
+    path = Path(__file__).resolve().parents[2] / "tools" / "prompts.yaml"
+    if not path.exists():
+        return {}
+    with path.open("r", encoding="utf-8") as f:
+        data = yaml.safe_load(f) or {}
+    return data if isinstance(data, dict) else {}
+
+
+_PROMPTS = _load_prompts()
+
+TOOL_ROUTING_PROMPT = _PROMPTS.get(
+    "tool_routing_prompt",
+    "请根据用户问题和可用 tools 选择合适工具；不要调用未出现在 tools 列表中的函数。",
+).strip()
+
+class ChatService:
+    """
+    聊天服务层
+    整合工具调用流式 LLM + 业务过滤
+    """
+
+    def __init__(self, config, rag_service, union_service):
+        self._config = config
+
+        # ---- 过滤配置 ----
+        self._filter_enabled = config["FILTER_ENABLED"]
+        self._allowed_keywords = config["FILTER_ALLOWED_KEYWORDS"]
+        self._rejection_message = config["FILTER_REJECTION_MESSAGE"]
+
+        # ---- LLM 配置 ----
+        self._client = ZhipuAiClient(
+            api_key=config["LLM_KEY"],
+            base_url=config["LLM_URL"],
+        )
+        self._model = config["LLM_MODEL"]
+        self._max_tokens = config["LLM_MAX_TOKENS"]
+        self._temperature = config["LLM_TEMPERATURE"]
+        self._top_p = config["LLM_TOP_P"]
+        self._system_prompt = config["SYSTEM_PROMPT"]
+        self._tools_max_rounds = config["TOOLS_MAX_ROUNDS"]
+        self._rag = rag_service
+        self._union_service = union_service
+
+    # ========== 过滤 ==========
+
+    def _check_question_valid(self, question: str) -> bool:
+        if not self._filter_enabled or not self._allowed_keywords:
+            return True
+        lower_question = question.lower()
+        return any(k.lower() in lower_question for k in self._allowed_keywords)
+
+    # ========== LLM ==========
+
+    def _build_messages(
+        self,
+        history: List[Dict[str, str]],
+        user_question: str,
+    ) -> List[Dict[str, str]]:
+        messages = []
+        if self._system_prompt:
+            messages.append({"role": "system", "content": self._system_prompt})
+        messages.extend(history)
+        messages.append({"role": "user", "content": user_question})
+        return messages
+
+    def _build_tool_messages(self, history: List[Dict[str, str]], user_question: str) -> List[Dict[str, str]]:
+        messages = self._build_messages(history, user_question)
+        insert_at = 1 if messages and messages[0].get("role") == "system" else 0
+        messages.insert(insert_at, {"role": "system", "content": TOOL_ROUTING_PROMPT})
+        return messages
+
+    def tool_call_stream(
+        self,
+        conversation_id: str,
+        question: str,
+        tools,
+        history: List[Dict[str, str]],
+        jsessionid: str,
+    ) -> Generator[ChatResponse, None, None]:
+        """
+        带工具调用的流式对话（支持交错思考与工具调用）
+        - stream=True + tool_stream=True：模型在流式输出中同时返回推理过程、回答内容与工具调用
+        - 工具执行结果回传模型后继续流式生成，循环直至模型不再调用工具或达到最大轮次
+        """
+        if not self._check_question_valid(question):
+            logger.info(f"问题未通过业务过滤。conversation_id={conversation_id} question={_preview(question, 120)}")
+            yield ChatResponse(
+                conversationId=conversation_id,
+                content=self._rejection_message,
+                finish_reason="rejected",
+            )
+            return
+
+        def content_event(content: str) -> ChatResponse:
+            return ChatResponse(conversationId=conversation_id, content=content)
+
+        def reasoning_event(reasoning_content: str) -> ChatResponse:
+            return ChatResponse(conversationId=conversation_id, reasoning_content=reasoning_content)
+
+        def tool_call_event(tool_call: str) -> ChatResponse:
+            return ChatResponse(conversationId=conversation_id, tool_call=tool_call)
+
+        def tool_result_event(tool_result: str) -> ChatResponse:
+            return ChatResponse(conversationId=conversation_id, tool_result=tool_result)
+
+        try:
+            messages = self._build_tool_messages(history, question)
+            max_rounds = self._tools_max_rounds
+            final_answer = ""
+
+            logger.info(f"开始模型流式调用。conversation_id={conversation_id} model={self._model} question={_preview(question, 120)}")
+            for round_idx in range(max_rounds):
+                response = self._client.chat.completions.create(
+                    model=self._model,
+                    messages=messages,
+                    tools=tools,
+                    tool_choice="auto",
+                    stream=True,
+                    tool_stream=True,
+                    thinking={"type": "enabled", "clear_thinking": False},
+                    max_tokens=self._max_tokens,
+                    temperature=self._temperature,
+                    top_p=self._top_p,
+                )
+
+                current_content = ""
+                current_reasoning = ""
+                reasoning_len = 0
+                tool_calls_map: Dict[int, Dict] = {}
+
+                for chunk in response:
+                    if not chunk.choices:
+                        continue
+                    delta = chunk.choices[0].delta
+
+                    reasoning = getattr(delta, "reasoning_content", None)
+                    if reasoning:
+                        current_reasoning += reasoning
+                        reasoning_len += len(reasoning)
+                        yield reasoning_event(reasoning)
+
+                    content = getattr(delta, "content", None)
+                    if content:
+                        current_content += content
+                        yield content_event(content)
+
+                    for tc in getattr(delta, "tool_calls", None) or []:
+                        self._merge_tool_call_delta(tool_calls_map, tc)
+
+                logger.info(f"模型流式返回完成。conversation_id={conversation_id} round={round_idx + 1} content_chars={len(current_content)} reasoning_chars={reasoning_len} tool_calls={len(tool_calls_map)} content_preview={_preview(current_content)}")
+
+                if not tool_calls_map:
+                    final_answer = current_content
+                    break
+
+                assistant_tool_calls = [tool_calls_map[i] for i in sorted(tool_calls_map)]
+                messages.append({
+                    "role": "assistant",
+                    "content": current_content or None,
+                    "reasoning_content": current_reasoning,
+                    "tool_calls": assistant_tool_calls,
+                })
+
+                for tc in assistant_tool_calls:
+                    name = tc["function"]["name"]
+                    args = tc["function"]["arguments"]
+                    logger.info(f"执行工具调用。conversation_id={conversation_id} tool={name} args={_preview(args, 200)}")
+                    yield tool_call_event(f"\n[调用工具: {name}({args})]\n")
+
+                    tool_context = ToolContext(
+                        union_service=self._union_service,
+                        rag_service=self._rag,
+                        jsessionid=jsessionid,
+                    )
+                    result = call_function(name, args, tool_context)
+                    logger.info(f"工具调用完成。conversation_id={conversation_id} tool={name} result_preview={_preview(result, 300)}")
+                    yield tool_result_event(result)
+
+                    messages.append({
+                        "role": "tool",
+                        "content": result,
+                        "tool_call_id": tc["id"],
+                    })
+
+            logger.info(f"对话完成。conversation_id={conversation_id} final_answer_chars={len(final_answer)} final_answer_preview={_preview(final_answer)}")
+            yield ChatResponse(conversationId=conversation_id, finish_reason="stop")
+
+        except Exception as e:
+            logger.exception(f"模型调用异常。conversation_id={conversation_id} question={_preview(question, 120)}")
+            yield ChatResponse(
+                conversationId=conversation_id,
+                content=f"[错误] 模型调用异常: {str(e)}",
+                finish_reason="error",
+            )
+
+    @staticmethod
+    def _merge_tool_call_delta(tool_calls_map: Dict[int, Dict], tc) -> None:
+        """将单个流式 tool_call 增量按 index 合并到累积字典中"""
+        slot = tool_calls_map.setdefault(tc.index, {
+            "id": "",
+            "type": "function",
+            "function": {"name": "", "arguments": ""},
+        })
+        if tc.id:
+            slot["id"] = tc.id
+        fn = getattr(tc, "function", None)
+        if fn is not None:
+            if getattr(fn, "name", None):
+                slot["function"]["name"] += fn.name
+            if getattr(fn, "arguments", None):
+                slot["function"]["arguments"] += fn.arguments

package/app/service/llm_service.py

@@ -0,0 +1,214 @@
+from __future__ import annotations
+
+import os
+import httpx
+import json
+from typing import Optional, Generator, Dict, Any, List
+from openai import OpenAI
+from loguru import logger
+
+from app.utils import common_utils
+
+
+class LLMService:
+    # 常量定义
+    DEFAULT_TEMPERATURE = 0.7
+    DEFAULT_MAX_TOKENS = 30000
+    STREAM_MAX_TOKENS = 30000
+    JSON_RESPONSE_FORMAT = {"type": "json_object"}
+
+    def __init__(self):
+        """初始化LLM服务实例"""
+        self._client: Optional[OpenAI] = None
+        self._model_name: Optional[str] = None
+
+    def initialize(self, config) -> None:
+        """初始化LLM客户端"""
+        if self._client is not None:
+            logger.debug("LLM客户端已经初始化，跳过重复初始化")
+            return
+
+        llm_url = config["LLM_URL"]
+        llm_key = config["LLM_KEY"]
+        self._model_name = config["LLM_MODEL"]
+
+        if not llm_url:
+            error_msg = "LLM_URL 配置为空，请检查环境变量配置"
+            logger.error(error_msg)
+            raise RuntimeError(error_msg)
+
+        if not llm_key:
+            error_msg = "LLM_KEY 配置为空，请检查环境变量配置"
+            logger.error(error_msg)
+            raise RuntimeError(error_msg)
+
+        try:
+            self._client = self._create_client(llm_url, llm_key)
+            logger.info(f"LLM客户端初始化成功: {llm_url}")
+        except Exception as e:
+            error_msg = f"LLM客户端初始化失败: {str(e)}"
+            logger.error(error_msg)
+            raise RuntimeError(error_msg) from e
+
+    @property
+    def model_name(self) -> str:
+        """获取模型名称"""
+        return self._model_name
+
+    def _create_client(self, llm_url: str, llm_key: str) -> OpenAI:
+        return OpenAI(
+            base_url=llm_url,
+            api_key=llm_key,
+            http_client=httpx.Client(verify=False)
+        )
+
+    def _ensure_client(self) -> None:
+        """
+        确保客户端已初始化
+
+        Raises:
+            RuntimeError: 当客户端未初始化时
+        """
+        if self._client is None:
+            raise RuntimeError("LLM客户端未初始化，请先调用 initialize() 方法")
+
+    def _build_messages(
+        self,
+        system_prompt: str,
+        user_prompt: Optional[str] = None,
+        query_text: Optional[str] = None,
+        data_json: Optional[str] = None
+    ) -> List[Dict[str, str]]:
+        """
+        构建LLM消息列表
+
+        Args:
+            system_prompt: 系统提示词
+            user_prompt: 用户提示词模板（可选）
+            query_text: 查询文本（可选）
+            data_json: 数据JSON字符串（可选）
+
+        Returns:
+            消息列表
+        """
+        messages = [{"role": "system", "content": system_prompt}]
+
+        if user_prompt:
+            # 将用户模板中的占位符替换为实际值
+            content = user_prompt
+            if query_text is not None:
+                content = content.replace("{query_text}", str(query_text))
+            if data_json is not None:
+                content = content.replace("{data_json}", str(data_json))
+            messages.append({"role": "user", "content": content})
+        elif query_text:
+            messages.append({"role": "user", "content": query_text})
+        else:
+            messages.append({"role": "user", "content": system_prompt})
+        return messages
+
+    def execute_llm(
+        self,
+        system_prompt: str,
+        user_prompt: Optional[str] = None,
+        query_text: Optional[str] = None,
+        data_json: Optional[str] = None,
+        temperature: Optional[float] = None,
+        max_tokens: Optional[int] = None,
+        use_json_format: bool = False
+    ) -> Optional[Dict[str, Any] | str]:
+        """
+        通用非流式LLM执行方法
+
+        Args:
+            system_prompt: 系统提示词
+            user_prompt: 用户提示词模板（可选）
+            query_text: 查询文本（可选）
+            data_json: 数据JSON字符串（可选）
+            temperature: 温度参数，默认为 DEFAULT_TEMPERATURE
+            max_tokens: 最大token数，默认为 DEFAULT_MAX_TOKENS
+            use_json_format: 是否使用JSON格式响应
+
+        Returns:
+            解析后的结果字典或原始字符串，失败返回None
+        """
+        self._ensure_client()
+
+        try:
+            messages = self._build_messages(system_prompt, user_prompt, query_text, data_json)
+
+            request_params = {
+                "model": self.model_name,
+                "messages": messages,
+                "temperature": temperature or self.DEFAULT_TEMPERATURE,
+                "max_tokens": max_tokens or self.DEFAULT_MAX_TOKENS
+            }
+
+            if use_json_format:
+                request_params["response_format"] = self.JSON_RESPONSE_FORMAT
+
+            llm_response = self._client.chat.completions.create(**request_params)
+            content = common_utils.remove_think_tag(llm_response.choices[0].message.content)
+
+            if use_json_format:
+                result = json.loads(content)
+                logger.debug(f"LLM JSON响应成功")
+                return result
+            else:
+                logger.debug("LLM文本响应成功")
+                return content
+        except Exception as e:
+            logger.error(f"LLM调用失败: {e.args}")
+            return None
+
+    def execute_llm_stream(
+        self,
+        system_prompt: str,
+        user_prompt: Optional[str] = None,
+        query_text: Optional[str] = None,
+        data_json: Optional[str] = None,
+        temperature: Optional[float] = None,
+        max_tokens: Optional[int] = None
+    ) -> Generator[str, None, None]:
+        """
+        通用流式LLM执行方法
+
+        Args:
+            system_prompt: 系统提示词
+            user_prompt: 用户提示词模板（可选）
+            query_text: 查询文本（可选）
+            data_json: 数据JSON字符串（可选）
+            temperature: 温度参数，默认为 DEFAULT_TEMPERATURE
+            max_tokens: 最大token数，默认为 STREAM_MAX_TOKENS
+
+        Yields:
+            LLM流式响应内容片段
+        """
+        self._ensure_client()
+        logger.debug("开始LLM流式调用")
+
+        try:
+            messages = self._build_messages(system_prompt, user_prompt, query_text, data_json)
+            logger.info(f"LLM请求: {messages[:2000]}...")
+            llm_stream_response = self._client.chat.completions.create(
+                model=self.model_name,
+                messages=messages,
+                temperature=temperature or self.DEFAULT_TEMPERATURE,
+                stream=True,
+                max_tokens=max_tokens or self.STREAM_MAX_TOKENS
+            )
+
+            for chunk in llm_stream_response:
+                content = getattr(chunk.choices[0].delta, 'content', None) or getattr(chunk.choices[0].delta, 'reasoning_content', None) or ''
+                if content:
+                    logger.info(f"LLM流式响应: {content[:50]}...")
+                    yield content
+
+        except Exception as e:
+            error_msg = f"[ERROR] llm 调用失败：{str(e)}"
+            logger.error(error_msg)
+            yield error_msg
+
+
+# 全局单例实例
+llm_service = LLMService()