# ！/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 foundation.utils.utils import get_models
from foundation.logger.loggering import server_logger as logger

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

    def __init__(self):
        # 获取部署的模型列表
        llm, chat, embed = get_models()
        self.llm = llm
        self.chat = chat

    async def get_model_generate_invoke(self, trace_id, task_prompt_info: dict):
        """
            模型非流式生成（异步）
        """

        prompt_template = task_prompt_info["task_prompt"]
        # 直接格式化消息，不需要额外的invoke步骤
        messages = prompt_template.format_messages()

        # 使用异步方法调用模型，避免阻塞事件循环
        import asyncio
        loop = asyncio.get_event_loop()
        response = await loop.run_in_executor(None, self.llm.invoke, messages)
        # logger.info(f"[模型生成结果]: {response.content}")
        return response.content

    def get_model_generate_stream(self, trace_id, task_prompt_info: dict):
        """
            模型流式生成（异步）
        """
        prompt_template = task_prompt_info["task_prompt"]
        # 直接格式化消息，不需要额外的invoke步骤  stream
        messages = prompt_template.format_messages()
        response = self.llm.stream(messages)
        for chunk in response:
            yield chunk.content

generate_model_client = GenerateModelClient()