LQAgentPlatform/foundation/ai/agent/generate/model_generate.py

# ！/usr/bin/ python
# -*- coding: utf-8 -*-
'''
@Project    : lq-agent-api
@File       :model_generate.py
@IDE        :PyCharm
@Author     :
@Date       :2025/7/14 14:22
'''

from langchain_core.prompts import ChatPromptTemplate
from langchain_core.messages import BaseMessage, SystemMessage, HumanMessage
from foundation.ai.models.model_handler import model_handler
from foundation.observability.logger.loggering import review_logger as logger
import asyncio
import time
from typing import Optional, Callable, Any, List, Union

class GenerateModelClient:
    """
        主要是生成式模型
    """

    def __init__(self, default_timeout: int = 60, max_retries: int = 3, backoff_factor: float = 1.0):
        # 获取默认模型
        self.llm = model_handler.get_models()
        self.chat = self.llm  # 当前chat和llm使用相同模型

        # 配置参数
        self.default_timeout = default_timeout
        self.max_retries = max_retries
        self.backoff_factor = backoff_factor

        # 保存model_handler引用,用于动态获取模型
        self.model_handler = model_handler

    async def _retry_with_backoff(self, func: Callable, *args, timeout: Optional[int] = None, **kwargs):
        """
        带指数退避的重试机制，每次重试都有独立的超时控制
        """
        current_timeout = timeout or self.default_timeout

        for attempt in range(self.max_retries + 1):
            try:
                # 每次重试都有独立的超时时间
                return await asyncio.wait_for(
                    func(*args, **kwargs),
                    timeout=current_timeout
                )
            except asyncio.TimeoutError:
                if attempt == self.max_retries:
                    logger.error(f"[模型调用] 达到最大重试次数 {self.max_retries}，最终超时")
                    raise TimeoutError(f"模型调用在 {self.max_retries} 次重试后均超时")

                wait_time = self.backoff_factor * (2 ** attempt)
                logger.warning(f"[模型调用] 第 {attempt + 1} 次超时, {wait_time}秒后重试...")
                await asyncio.sleep(wait_time)
            except Exception as e:
                if attempt == self.max_retries:
                    logger.error(f"[模型调用] 达到最大重试次数 {self.max_retries}，最终失败: {str(e)}")
                    raise

                wait_time = self.backoff_factor * (2 ** attempt)
                logger.warning(f"[模型调用] 第 {attempt + 1} 次尝试失败: {str(e)}, {wait_time}秒后重试...")
                await asyncio.sleep(wait_time)

    async def get_model_generate_invoke(
        self,
        trace_id: str,
        task_prompt_info: Optional[dict] = None,
        messages: Optional[List[BaseMessage]] = None,
        system_prompt: Optional[str] = None,
        user_prompt: Optional[str] = None,
        prompt: Optional[str] = None,
        timeout: Optional[int] = None,
        model_name: Optional[str] = None
    ) -> str:
        """模型非流式生成（异步）

        支持多种调用方式（优先级从高到低）：
        1. messages: 直接传入 LangChain Message 对象列表
        2. system_prompt + user_prompt: 分别传入系统和用户提示词
        3. prompt: 传入单条用户提示词字符串
        4. task_prompt_info: 传入包含 ChatPromptTemplate 的字典（兼容旧接口）

        Args:
            trace_id: 追踪ID
            task_prompt_info: 任务提示词信息（兼容旧接口），需包含 format_messages() 方法
            messages: LangChain Message 对象列表（如 [SystemMessage, HumanMessage]）
            system_prompt: 系统提示词字符串
            user_prompt: 用户提示词字符串
            prompt: 单条用户提示词字符串（无系统提示时使用）
            timeout: 超时时间（秒），默认使用构造时的 default_timeout
            model_name: 模型名称（可选），支持 doubao/qwen/deepseek/gemini 等

        Returns:
            str: 模型生成的文本内容

        Raises:
            ValueError: 参数组合错误
            TimeoutError: 调用超时
            Exception: 模型调用异常

        Examples:
            # 方式1: 使用 Message 列表（推荐）
            messages = [SystemMessage(content="你是专家"), HumanMessage(content="请分析...")]
            result = await client.get_model_generate_invoke("trace-001", messages=messages)

            # 方式2: 分别传入系统和用户提示词
            result = await client.get_model_generate_invoke(
                "trace-001",
                system_prompt="你是专家",
                user_prompt="请分析..."
            )

            # 方式3: 传入单条提示词
            result = await client.get_model_generate_invoke("trace-001", prompt="请分析...")

            # 方式4: 兼容旧接口（使用 PromptLoader）
            task_prompt_info = {"task_prompt": chat_template}
            result = await client.get_model_generate_invoke("trace-001", task_prompt_info=task_prompt_info)
        """
        start_time = time.time()
        current_timeout = timeout or self.default_timeout

        try:
            # 选择模型
            llm_to_use = self.model_handler.get_model_by_name(model_name) if model_name else self.llm
            logger.info(f"[模型调用] 使用{'指定' if model_name else '默认'}模型: {model_name or 'default'}, trace_id: {trace_id}")

            # 构建消息列表（按优先级）
            final_messages = self._build_messages(
                messages=messages,
                system_prompt=system_prompt,
                user_prompt=user_prompt,
                prompt=prompt,
                task_prompt_info=task_prompt_info
            )

            # 定义模型调用函数，使用原生 ainvoke
            async def _invoke():
                return await llm_to_use.ainvoke(final_messages)

            # 调用带重试机制
            response = await self._retry_with_backoff(_invoke, timeout=current_timeout)

            elapsed_time = time.time() - start_time
            logger.info(f"[模型调用] 成功 trace_id: {trace_id}, 耗时: {elapsed_time:.2f}s")
            return response.content

        except asyncio.TimeoutError:
            elapsed_time = time.time() - start_time
            logger.error(f"[模型调用] 超时 trace_id: {trace_id}, 耗时: {elapsed_time:.2f}s, 超时阈值: {current_timeout}s")
            raise TimeoutError(f"模型调用超时，trace_id: {trace_id}")

        except Exception as e:
            elapsed_time = time.time() - start_time
            logger.error(f"[模型调用] 异常 trace_id: {trace_id}, 耗时: {elapsed_time:.2f}s, 错误: {type(e).__name__}: {str(e)}")
            raise

    def _build_messages(
        self,
        messages: Optional[List[BaseMessage]] = None,
        system_prompt: Optional[str] = None,
        user_prompt: Optional[str] = None,
        prompt: Optional[str] = None,
        task_prompt_info: Optional[dict] = None
    ) -> List[BaseMessage]:
        """构建消息列表（内部方法）

        优先级：messages > system_prompt+user_prompt > prompt > task_prompt_info
        """
        # 方式1: 直接使用传入的 Message 列表
        if messages is not None:
            if not isinstance(messages, list):
                raise ValueError("messages 必须是列表")
            if len(messages) == 0:
                raise ValueError("messages 不能为空列表")
            logger.debug(f"使用传入的 messages 列表，共 {len(messages)} 条消息")
            return messages

        # 方式2: system_prompt + user_prompt
        if system_prompt is not None and user_prompt is not None:
            logger.debug("使用 system_prompt + user_prompt 构建消息")
            return [SystemMessage(content=system_prompt), HumanMessage(content=user_prompt)]

        # 方式3: 单独 system_prompt（可能是特殊情况）
        if system_prompt is not None:
            logger.debug("使用单独的 system_prompt 构建消息")
            return [SystemMessage(content=system_prompt)]

        # 方式4: 单条 prompt 字符串
        if prompt is not None:
            logger.debug("使用单条 prompt 字符串构建消息")
            return [HumanMessage(content=prompt)]

        # 方式5: 兼容旧接口 task_prompt_info
        if task_prompt_info is not None:
            if "task_prompt" not in task_prompt_info:
                raise ValueError("task_prompt_info 必须包含 'task_prompt' 键")
            task_prompt = task_prompt_info["task_prompt"]
            if hasattr(task_prompt, 'format_messages'):
                logger.debug("使用 task_prompt_info 中的 ChatPromptTemplate 构建消息")
                return task_prompt.format_messages()
            elif isinstance(task_prompt, str):
                logger.debug("使用 task_prompt_info 中的字符串构建消息")
                return [HumanMessage(content=task_prompt)]
            else:
                raise ValueError(f"task_prompt 类型不支持: {type(task_prompt)}")

        # 没有提供任何有效参数
        raise ValueError(
            "必须提供以下参数之一: "
            "messages, system_prompt+user_prompt, prompt, 或 task_prompt_info"
        )

    def get_model_generate_stream(
        self,
        trace_id: str,
        task_prompt_info: Optional[dict] = None,
        messages: Optional[List[BaseMessage]] = None,
        system_prompt: Optional[str] = None,
        user_prompt: Optional[str] = None,
        prompt: Optional[str] = None,
        timeout: Optional[int] = None
    ):
        """模型流式生成（同步生成器）

        支持多种调用方式（优先级从高到低）：
        1. messages: 直接传入 LangChain Message 对象列表
        2. system_prompt + user_prompt: 分别传入系统和用户提示词
        3. prompt: 传入单条用户提示词字符串
        4. task_prompt_info: 传入包含 ChatPromptTemplate 的字典（兼容旧接口）

        Args:
            trace_id: 追踪ID
            task_prompt_info: 任务提示词信息（兼容旧接口）
            messages: LangChain Message 对象列表
            system_prompt: 系统提示词字符串
            user_prompt: 用户提示词字符串
            prompt: 单条用户提示词字符串
            timeout: 超时时间（秒）

        Yields:
            str: 生成的文本块

        Raises:
            ValueError: 参数组合错误
        """
        start_time = time.time()
        current_timeout = timeout or self.default_timeout

        try:
            logger.info(f"[模型流式调用] 开始处理 trace_id: {trace_id}, 超时配置: {current_timeout}s")

            # 构建消息列表
            final_messages = self._build_messages(
                messages=messages,
                system_prompt=system_prompt,
                user_prompt=user_prompt,
                prompt=prompt,
                task_prompt_info=task_prompt_info
            )

            response = self.llm.stream(final_messages)

            chunk_count = 0
            for chunk in response:
                chunk_count += 1
                if hasattr(chunk, 'content') and chunk.content:
                    yield chunk.content
                elif chunk:
                    yield chunk

            elapsed_time = time.time() - start_time
            logger.info(f"[模型流式调用] 成功 trace_id: {trace_id}, 生成块数: {chunk_count}, 耗时: {elapsed_time:.2f}s")

        except Exception as e:
            elapsed_time = time.time() - start_time
            logger.error(f"[模型流式调用] 异常 trace_id: {trace_id}, 耗时: {elapsed_time:.2f}s, 错误: {type(e).__name__}: {str(e)}")
            raise

generate_model_client = GenerateModelClient(default_timeout=15, max_retries=2, backoff_factor=0.5)