LQAgentPlatform/foundation/ai/agent/test_agent.py

# ！/usr/bin/python
# -*- coding: utf-8 -*-
'''
@Project    : lq-agent-api
@File       :agent_mcp.py
@IDE        :PyCharm
@Author     :
@Date       :2025/7/21 10:12
'''
import json

from langgraph.prebuilt import create_react_agent
from sqlalchemy.sql.functions import user
from foundation.observability.logger.loggering import server_logger
from foundation.utils.common import handler_err
from foundation.ai.models import get_models
from foundation.utils.yaml_utils import get_system_prompt_config

import threading
import time
from typing import Dict, List, Optional, AsyncGenerator, Any, OrderedDict
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnableConfig
from foundation.ai.agent.base_agent import BaseAgent
from foundation.schemas.test_schemas import TestForm
# from foundation.agent.function.test_funciton import test_funtion


class TestAgentClient(BaseAgent):
    """
    Xiwuzc 智能助手+MCP（带完整会话管理） - 针对场景优化
    添加会话锁定机制，确保同一时间只有一个客户端可以使用特定会话
    """
    # 单例实例和线程锁
    _instance = None
    _singleton_lock = threading.Lock()

    def __new__(cls):
        """线程安全的单例模式实现"""
        if cls._instance is None:
            with cls._singleton_lock:
                if cls._instance is None:
                    cls._instance = super().__new__(cls)
                    cls._instance._initialize()
        return cls._instance

    def _initialize(self):
        """初始化模型和会话管理"""
        llm, chat, embed = get_models()
        self.llm = llm
        self.chat = chat
        self.embed = embed
        self.agent_executor = None
        self.initialized = False
        self.psutil_available = True

        # 固定系统提示词
        self.system_prompt = get_system_prompt_config()["system_prompt"]

        # 清理任务
        self.cleanup_task = None
        server_logger.info(" client initialized")

    async def init_agent(self):
        """初始化agent_executor（只需一次）"""
        if self.initialized:
            return

        # 获取部署的模型列表
        server_logger.info(f"系统提示词 system_prompt：{self.system_prompt}")

        # 创建提示词模板 - 使用固定的系统提示词
        prompt = ChatPromptTemplate.from_messages([
            ("system", self.system_prompt),
            MessagesPlaceholder(variable_name="messages"),
            ("placeholder", "{agent_scratchpad}")
        ])

        # # 创建Agent - 不再使用MemorySaver
        # self.agent_executor = create_react_agent(
        #     self.llm,
        #     tools=[test_funtion.query_info , test_funtion.execute , test_funtion.handle] ,  # 专用工具集 + 私有知识库检索工具
        #     prompt=prompt
        # )
        self.initialized = True
        server_logger.info(" agent initialized")


    async def handle_query(self, trace_id: str, task_prompt_info: dict, input_query, context=None,
                            config_param: TestForm = None):
        try:
            # 确保agent已初始化
            if not self.initialized:
                await self.init_agent()
            
            session_id = config_param.session_id
           

            try:
                # 构建输入消息
                input_message , input_summary_context = self.get_input_context(
                    trace_id=trace_id,
                    task_prompt_info=task_prompt_info,
                    input_query=input_query,
                    context=context
                )
                # 用于模型对话使用
                input_human_message = HumanMessage(content=input_message)
                # 用于对话历史记录摘要 
                input_human_summary_message = HumanMessage(content=input_summary_context)
                # 获取历史消息
                history_messages = []
                # 构造完整的消息列表
                all_messages = list(history_messages) + [input_human_message]

                # 配置执行上下文
                config = RunnableConfig(
                    configurable={"thread_id": session_id},
                    runnable_kwargs={"recursion_limit": 15}
                )

                # 执行智能体
                events = self.agent_executor.astream(
                    {"messages": all_messages},
                    config=config,
                    stream_mode="values"
                )

                # 处理结果
                full_response = []
                async for event in events:
                    if isinstance(event["messages"][-1], AIMessage):
                        chunk = event["messages"][-1].content
                        full_response.append(chunk)
                    log_content = self.get_pretty_message_str(event["messages"][-1])
                    server_logger.info("\n" + log_content.strip(), trace_id=trace_id)

                if full_response:
                    full_text = "".join(full_response)
                    server_logger.info(trace_id=trace_id, msg=f"full_response: {full_text}")
                    full_text = self.clean_json_output(full_text)
                    return full_text
            finally:
                # 确保释放会话锁
                pass
        except PermissionError as e:
            # 处理会话被其他设备锁定的情况
            return str(e)
        except Exception as e:
            handler_err(server_logger, trace_id=trace_id, err=e, err_name='agent/chat')
            return f"系统错误: {str(e)}"


    async def handle_query_stream(
            self,
            trace_id: str,
            task_prompt_info: dict,
            input_query: str,
            context: Optional[str] = None,
            header_info: Optional[Dict] = None,
            config_param: TestForm = None,
    ) -> AsyncGenerator[str, None]:
        """流式处理查询（优化缓冲管理）"""
        try:
            # 确保agent已初始化
            if not self.initialized:
                await self.init_agent()
            
            session_id = config_param.session_id
        
            try:
                # 构建输入消息
                input_message , input_summary_context = self.get_input_context(
                    trace_id=trace_id,
                    task_prompt_info=task_prompt_info,
                    input_query=input_query,
                    context=context
                )
                server_logger.info(trace_id=trace_id, msg=f"input_context: {input_message}")
                # 用于模型对话使用
                input_human_message = HumanMessage(content=input_message)
                # 用于对话历史记录摘要 
                input_human_summary_message = HumanMessage(content=input_summary_context)
                 # 获取历史消息
                history_messages = []
                # 构造完整的消息列表
                all_messages = list(history_messages) + [input_human_message]
                # 配置执行上下文
                config = RunnableConfig(
                    configurable={"thread_id": session_id},
                    runnable_kwargs={"recursion_limit": 15}
                )

                # 流式执行
                events = self.agent_executor.astream_events(
                    {"messages": all_messages},
                    config=config,
                    stream_mode="values"
                )

                full_response = []
                buffer = []
                last_flush_time = time.time()

                # 流式处理事件
                async for event in events:
                    # 只在特定事件类型时打印日志
                    self.log_stream_pretty_message(trace_id=trace_id, event=event)
                   
                    if 'chunk' in event['data'] and "on_chat_model_stream" in event['event']:
                        chunk = event['data']['chunk'].content
                        full_response.append(chunk)

                        # 缓冲管理策略
                        buffer.append(chunk)
                        current_time = time.time()

                        # 满足以下任一条件即刷新缓冲区
                        if (len(buffer) >= 3 or  # 达到最小块数
                                (current_time - last_flush_time) > 0.5 or  # 超时
                                any(chunk.endswith((c, f"{c} ")) for c in
                                    ['.', '。', '!', '？', '\n', '；', ';'])):  # 自然断点

                            # 合并并发送缓冲内容
                            combined = ''.join(buffer)
                            yield combined

                            # 重置缓冲
                            buffer.clear()
                            last_flush_time = current_time

                # 处理剩余内容
                if buffer:
                    yield ''.join(buffer)

                # 将完整响应添加到历史并进行压缩
                if full_response:
                    full_text = "".join(full_response)
                    server_logger.info(trace_id=trace_id, msg=f"full_response: {full_text}")
            finally:
                # 确保释放会话锁
                pass

        except PermissionError as e:
            yield json.dumps({"error": str(e)})
        except Exception as e:
            handler_err(server_logger, trace_id=trace_id, err=e, err_name='test_stream')
            yield json.dumps({"error": f"系统错误: {str(e)}"})


test_agent_client = TestAgentClient()