feat(agent): upgrade langchain to v1.0+

app/agent/__init__.py

@@ -1,26 +1,25 @@
 import asyncio
-from typing import Dict, List, Any, Union
-import json
-import tiktoken
+from time import strftime
+from typing import Dict, List
 
-from langchain.agents import AgentExecutor
-from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
-from langchain_community.callbacks import get_openai_callback
-from langchain_core.chat_history import InMemoryChatMessageHistory
-from langchain_core.messages import HumanMessage, AIMessage, ToolCall, ToolMessage, SystemMessage, trim_messages
-from langchain_core.runnables import RunnablePassthrough, RunnableLambda
-from langchain_core.runnables.history import RunnableWithMessageHistory
-from langchain.agents.format_scratchpad.openai_tools import format_to_openai_tool_messages
-from langchain.agents.output_parsers.openai_tools import OpenAIToolsAgentOutputParser
+from langchain.agents import create_agent
+from langchain.agents.middleware import (
+    SummarizationMiddleware,
+    ModelRetryMiddleware,
+    ToolRetryMiddleware,
+)
+from langchain_core.messages import (
+    HumanMessage,
+    BaseMessage,
+)
+from langgraph.checkpoint.memory import InMemorySaver
 
-from app.agent.callback import StreamingCallbackHandler
-from app.agent.memory import conversation_manager
+from app.agent.callback import StreamingHandler, StreamingHandler
+from app.agent.memory import memory_manager
 from app.agent.prompt import prompt_manager
 from app.agent.tools.factory import MoviePilotToolFactory
 from app.chain import ChainBase
-from app.core.config import settings
 from app.helper.llm import LLMHelper
-from app.helper.message import MessageHelper
 from app.log import logger
 from app.schemas import Notification
 
@@ -31,42 +30,32 @@ class AgentChain(ChainBase):
 
 class MoviePilotAgent:
     """
-    MoviePilot AI智能体
+    MoviePilot AI智能体（基于 LangChain v1 + LangGraph）
     """
 
-    def __init__(self, session_id: str, user_id: str = None,
-                 channel: str = None, source: str = None, username: str = None):
+    def __init__(
+            self,
+            session_id: str,
+            user_id: str = None,
+            channel: str = None,
+            source: str = None,
+            username: str = None,
+    ):
         self.session_id = session_id
         self.user_id = user_id
-        self.channel = channel  # 消息渠道
-        self.source = source  # 消息来源
-        self.username = username  # 用户名
+        self.channel = channel
+        self.source = source
+        self.username = username
 
-        # 消息助手
-        self.message_helper = MessageHelper()
+        # 流式token管理
+        self.stream_handler = StreamingHandler()
 
-        # 回调处理器
-        self.callback_handler = StreamingCallbackHandler(
-            session_id=session_id
-        )
-
-        # LLM模型
-        self.llm = self._initialize_llm()
-
-        # 工具
-        self.tools = self._initialize_tools()
-
-        # 提示词模板
-        self.prompt = self._initialize_prompt()
-
-        # Agent执行器
-        self.agent_executor = self._create_agent_executor()
-
-    def _initialize_llm(self):
+    @staticmethod
+    def _initialize_llm():
         """
-        初始化LLM模型
+        初始化 LLM（带流式回调）
         """
-        return LLMHelper.get_llm(streaming=True, callbacks=[self.callback_handler])
+        return LLMHelper.get_llm(streaming=True)
 
     def _initialize_tools(self) -> List:
         """
@@ -78,384 +67,118 @@ class MoviePilotAgent:
             channel=self.channel,
             source=self.source,
             username=self.username,
-            callback_handler=self.callback_handler
+            stream_handler=self.stream_handler,
         )
 
-    @staticmethod
-    def _initialize_session_store() -> Dict[str, InMemoryChatMessageHistory]:
+    def _create_agent(self):
         """
-        初始化内存存储
-        """
-        return {}
-
-    def get_session_history(self, session_id: str) -> InMemoryChatMessageHistory:
-        """
-        获取会话历史
-        """
-        chat_history = InMemoryChatMessageHistory()
-        messages: List[dict] = conversation_manager.get_recent_messages_for_agent(
-            session_id=session_id,
-            user_id=self.user_id
-        )
-        if messages:
-            for msg in messages:
-                if msg.get("role") == "user":
-                    chat_history.add_message(HumanMessage(content=msg.get("content", "")))
-                elif msg.get("role") == "agent":
-                    chat_history.add_message(AIMessage(content=msg.get("content", "")))
-                elif msg.get("role") == "tool_call":
-                    metadata = msg.get("metadata", {})
-                    chat_history.add_message(
-                        AIMessage(
-                            content=msg.get("content", ""),
-                            tool_calls=[
-                                ToolCall(
-                                    id=metadata.get("call_id"),
-                                    name=metadata.get("tool_name"),
-                                    args=metadata.get("parameters"),
-                                )
-                            ]
-                        )
-                    )
-                elif msg.get("role") == "tool_result":
-                    metadata = msg.get("metadata", {})
-                    chat_history.add_message(ToolMessage(
-                        content=msg.get("content", ""),
-                        tool_call_id=metadata.get("call_id", "unknown")
-                    ))
-                elif msg.get("role") == "system":
-                    chat_history.add_message(SystemMessage(content=msg.get("content", "")))
-        
-        return chat_history
-
-    @staticmethod
-    def _initialize_prompt() -> ChatPromptTemplate:
-        """
-        初始化提示词模板
+        创建 LangGraph Agent（使用 create_agent + SummarizationMiddleware）
         """
         try:
-            prompt_template = ChatPromptTemplate.from_messages([
-                ("system", "{system_prompt}"),
-                MessagesPlaceholder(variable_name="chat_history"),
-                ("user", "{input}"),
-                MessagesPlaceholder(variable_name="agent_scratchpad"),
-            ])
-            logger.info("LangChain提示词模板初始化成功")
-            return prompt_template
+            # 系统提示词
+            system_prompt = prompt_manager.get_agent_prompt(
+                channel=self.channel
+            ).format(
+                current_date=strftime('%Y-%m-%d')
+            )
+
+            # LLM 模型（用于 agent 执行）
+            llm = self._initialize_llm()
+
+            # 工具列表
+            tools = self._initialize_tools()
+
+            # 中间件
+            middlewares = [
+                # 上下文压缩
+                SummarizationMiddleware(
+                    model=llm,
+                    trigger=("fraction", 0.85)
+                ),
+                # 模型调用失败时自动重试
+                ModelRetryMiddleware(max_retries=3),
+                # 工具调用失败时自动重试
+                ToolRetryMiddleware(max_retries=1)
+            ]
+
+            return create_agent(
+                model=llm,
+                tools=tools,
+                system_prompt=system_prompt,
+                middleware=middlewares,
+                checkpointer=InMemorySaver(),
+            )
         except Exception as e:
-            logger.error(f"初始化提示词失败: {e}")
+            logger.error(f"创建 Agent 失败: {e}")
             raise e
 
-    @staticmethod
-    def _token_counter(messages: List[Union[HumanMessage, AIMessage, ToolMessage, SystemMessage]]) -> int:
+    async def process(self, message: str) -> str:
         """
-        通用的Token计数器
+        处理用户消息，流式推理并返回 Agent 回复
         """
         try:
-            # 尝试从模型获取编码集，如果失败则回退到 cl100k_base (大多数现代模型使用的编码)
-            try:
-                encoding = tiktoken.encoding_for_model(settings.LLM_MODEL)
-            except KeyError:
-                encoding = tiktoken.get_encoding("cl100k_base")
+            logger.info(f"Agent推理: session_id={self.session_id}, input={message}")
 
-            num_tokens = 0
-            for message in messages:
-                # 基础开销 (每个消息大约 3 个 token)
-                num_tokens += 3
-                
-                # 1. 处理文本内容 (content)
-                if isinstance(message.content, str):
-                    num_tokens += len(encoding.encode(message.content))
-                elif isinstance(message.content, list):
-                    for part in message.content:
-                        if isinstance(part, dict) and part.get("type") == "text":
-                            num_tokens += len(encoding.encode(part.get("text", "")))
-
-                # 2. 处理工具调用 (仅 AIMessage 包含 tool_calls)
-                if getattr(message, "tool_calls", None):
-                    for tool_call in message.tool_calls:
-                        # 函数名
-                        num_tokens += len(encoding.encode(tool_call.get("name", "")))
-                        # 参数 (转为 JSON 估算)
-                        args_str = json.dumps(tool_call.get("args", {}), ensure_ascii=False)
-                        num_tokens += len(encoding.encode(args_str))
-                        # 额外的结构开销 (ID 等)
-                        num_tokens += 3
-
-                # 3. 处理角色权重
-                num_tokens += 1
-
-            # 加上回复的起始 Token (大约 3 个 token)
-            num_tokens += 3
-            return num_tokens
-        except Exception as e:
-            logger.error(f"Token计数失败: {e}")
-            # 发生错误时返回一个保守的估算值
-            return len(str(messages)) // 4
-
-    def _create_agent_executor(self) -> RunnableWithMessageHistory:
-        """
-        创建Agent执行器
-        """
-        try:
-            # 消息裁剪器，防止上下文超出限制
-            base_trimmer = trim_messages(
-                max_tokens=settings.LLM_MAX_CONTEXT_TOKENS * 1000 * 0.8,
-                strategy="last",
-                token_counter=self._token_counter,
-                include_system=True,
-                allow_partial=False,
-                start_on="human",
-            )
-            
-            # 包装trimmer，在裁剪后验证工具调用的完整性
-            def validated_trimmer(messages):
-                # 如果输入是 PromptValue，转换为消息列表
-                if hasattr(messages, "to_messages"):
-                    messages = messages.to_messages()
-                trimmed = base_trimmer.invoke(messages)
-
-                # 二次校验：确保不出现 broken tool chains
-                # 1. AIMessage with tool_calls 必须紧跟着对应的 ToolMessage
-                # 2. ToolMessage 必须有对应的 AIMessage 前置
-                safe_messages = []
-                i = 0
-                while i < len(trimmed):
-                    msg = trimmed[i]
-
-                    if isinstance(msg, AIMessage) and getattr(msg, "tool_calls", None):
-                        # 检查工具调用序列是否完整
-                        tool_calls = msg.tool_calls
-                        is_valid_sequence = True
-                        tool_results = []
-                        
-                        # 向后查找对应的 ToolMessage
-                        temp_i = i + 1
-                        for tool_call in tool_calls:
-                            if temp_i >= len(trimmed):
-                                is_valid_sequence = False
-                                break
-                            
-                            next_msg = trimmed[temp_i]
-                            if isinstance(next_msg, ToolMessage) and next_msg.tool_call_id == tool_call.get("id"):
-                                tool_results.append(next_msg)
-                                temp_i += 1
-                            else:
-                                is_valid_sequence = False
-                                break
-                        
-                        if is_valid_sequence:
-                            # 序列完整，保留消息
-                            safe_messages.append(msg)
-                            safe_messages.extend(tool_results)
-                            i = temp_i  # 跳过已处理的工具结果
-                        else:
-                            # 序列不完整，丢弃该 AIMessage（后续的孤立 ToolMessage 会在下一次循环被当做 orphaned 处理掉）
-                            logger.warning(f"移除无效的工具调用链: {len(tool_calls)} calls, incomplete results")
-                            i += 1
-                        continue
-
-                    if isinstance(msg, ToolMessage):
-                        # 如果在这里遇到 ToolMessage，说明它没有被上面的逻辑消费，则是孤立的（或者顺序错乱）
-                        logger.warning("移除孤立的 ToolMessage")
-                        i += 1
-                        continue
-
-                    # 其他类型的消息直接保留
-                    safe_messages.append(msg)
-                    i += 1
-
-                if len(safe_messages) < len(messages):
-                    logger.info(f"LangChain消息上下文已裁剪: {len(messages)} -> {len(safe_messages)}")
-                return safe_messages
-            
-            # 创建Agent执行链
-            agent = (
-                RunnablePassthrough.assign(
-                    agent_scratchpad=lambda x: format_to_openai_tool_messages(
-                        x["intermediate_steps"]
-                    )
-                )
-                | self.prompt
-                | RunnableLambda(validated_trimmer)
-                | self.llm.bind_tools(self.tools)
-                | OpenAIToolsAgentOutputParser()
-            )
-            executor = AgentExecutor(
-                agent=agent,
-                tools=self.tools,
-                verbose=settings.LLM_VERBOSE,
-                max_iterations=settings.LLM_MAX_ITERATIONS,
-                return_intermediate_steps=True,
-                handle_parsing_errors=True,
-                early_stopping_method="force"
-            )
-            return RunnableWithMessageHistory(
-                executor,
-                self.get_session_history,
-                input_messages_key="input",
-                history_messages_key="chat_history"
-            )
-        except Exception as e:
-            logger.error(f"创建Agent执行器失败: {e}")
-            raise e
-
-    async def _summarize_history(self):
-        """
-        总结提炼之前的对话和工具执行情况，并把会话总结变成新的系统提示词取代之前的对话
-        """
-        try:
-            # 获取当前历史记录
-            chat_history = self.get_session_history(self.session_id)
-            messages = chat_history.messages
-            if not messages:
-                return
-
-            logger.info(f"会话 {self.session_id} 历史消息长度已超过 90%，开始总结并重置上下文...")
-
-            # 将消息转换为摘要所需的文本格式
-            history_text = ""
-            for msg in messages:
-                if isinstance(msg, HumanMessage):
-                    history_text += f"用户: {msg.content}\n"
-                elif isinstance(msg, AIMessage):
-                    history_text += f"智能体: {msg.content}\n"
-                    if getattr(msg, "tool_calls", None):
-                        for tool_call in msg.tool_calls:
-                            history_text += f"智能体调用工具: {tool_call.get('name')}，参数: {tool_call.get('args')}\n"
-                elif isinstance(msg, ToolMessage):
-                    history_text += f"工具响应: {msg.content}\n"
-                elif isinstance(msg, SystemMessage):
-                    history_text += f"系统: {msg.content}\n"
-
-            # 摘要提示词
-            summary_prompt = (
-                "Please provide a comprehensive and highly informational summary of the preceding conversation and tool executions. "
-                "Your goal is to condense the history while retaining all critical details for future reference. "
-                "Ensure you include:\n"
-                "1. User's core intents, specific requests, and any mentioned preferences.\n"
-                "2. Names of movies, TV shows, or other key entities discussed.\n"
-                "3. A concise log of tool calls made and their specific results/outcomes.\n"
-                "4. The current status of any tasks and any pending actions.\n"
-                "5. Any important context that would be necessary for the agent to continue the conversation seamlessly.\n"
-                "The summary should be dense with information and serve as the primary context for the next stage of the interaction."
-            )
-
-            # 调用 LLM 进行总结 (非流式)
-            summary_llm = LLMHelper.get_llm(streaming=False)
-            response = await summary_llm.ainvoke([
-                SystemMessage(content=summary_prompt),
-                HumanMessage(content=f"Here is the conversation history to summarize:\n{history_text}")
-            ])
-            summary_content = str(response.content)
-
-            if not summary_content:
-                logger.warning("总结生成失败，跳过重置逻辑。")
-                return
-
-            # 清空原有的会话记录并插入新的系统总结
-            await conversation_manager.clear_memory(self.session_id, self.user_id)
-            await conversation_manager.add_conversation(
+            # 获取历史消息
+            messages = memory_manager.get_agent_messages(
                 session_id=self.session_id,
-                user_id=self.user_id,
-                role="system",
-                content=f"<history_summary>\n{summary_content}\n</history_summary>"
-            )
-            logger.info(f"会话 {self.session_id} 历史摘要替换完成。")
-        except Exception as e:
-            logger.error(f"执行会话总结出错: {str(e)}")
-
-    async def process_message(self, message: str) -> str:
-        """
-        处理用户消息
-        """
-        try:
-            # 检查上下文长度是否超过 90%
-            history = self.get_session_history(self.session_id)
-            if self._token_counter(history.messages) > settings.LLM_MAX_CONTEXT_TOKENS * 1000 * 0.9:
-                await self._summarize_history()
-
-            # 添加用户消息到记忆
-            await conversation_manager.add_conversation(
-                self.session_id,
-                user_id=self.user_id,
-                role="user",
-                content=message
+                user_id=self.user_id
             )
 
-            # 构建输入上下文
-            input_context = {
-                "system_prompt": prompt_manager.get_agent_prompt(channel=self.channel),
-                "input": message
-            }
+            # 增加用户消息
+            messages.append(HumanMessage(content=message))
 
-            # 执行Agent
-            logger.info(f"Agent执行推理: session_id={self.session_id}, input={message}")
-
-            result = await self._execute_agent(input_context)
-
-            # 获取Agent回复
-            agent_message = await self.callback_handler.get_message()
-
-            # 发送Agent回复给用户（通过原渠道）
-            if agent_message:
-                # 发送回复
-                await self.send_agent_message(agent_message)
-
-                # 添加Agent回复到记忆
-                await conversation_manager.add_conversation(
-                    session_id=self.session_id,
-                    user_id=self.user_id,
-                    role="agent",
-                    content=agent_message
-                )
-            else:
-                agent_message = result.get("output") or "很抱歉，智能体出错了，未能生成回复内容。"
-                await self.send_agent_message(agent_message)
-
-            return agent_message
+            # 执行推理
+            await self._execute_agent(messages)
 
         except Exception as e:
             error_message = f"处理消息时发生错误: {str(e)}"
             logger.error(error_message)
-            # 发送错误消息给用户（通过原渠道）
             await self.send_agent_message(error_message)
             return error_message
 
-    async def _execute_agent(self, input_context: Dict[str, Any]) -> Dict[str, Any]:
+    async def _execute_agent(self, messages: List[BaseMessage]):
         """
-        执行LangChain Agent
+        调用 LangGraph Agent，通过 astream_events 流式获取 token，
+        同时用 UsageMetadataCallbackHandler 统计 token 用量。
         """
         try:
-            with get_openai_callback() as cb:
-                result = await self.agent_executor.ainvoke(
-                    input_context,
-                    config={"configurable": {"session_id": self.session_id}},
-                    callbacks=[self.callback_handler]
-                )
-                logger.info(f"LLM调用消耗: \n{cb}")
+            # Agent运行配置
+            agent_config = {
+                ""
+            }
+
+            # 创建智能体
+            agent = self._create_agent()
+
+            # 流式运行智能体
+            async for chunk in agent.astream(
+                    {"messages": messages},
+                    stream_mode="messages",
+                    config=agent_config
+            ):
+                token, metadata = chunk
+                if token:
+                    self.stream_handler.emit(token.content)
+
+            # 发送最终消息给用户
+            await self.send_agent_message(
+                self.stream_handler.take()
+            )
+
+            # 保存消息
+            memory_manager.save_agent_messages(
+                session_id=self.session_id,
+                user_id=self.user_id,
+                messages=agent.get_state(agent_config).values("messages")
+            )
 
-                if cb.total_tokens > 0:
-                    result["token_usage"] = {
-                        "prompt_tokens": cb.prompt_tokens,
-                        "completion_tokens": cb.completion_tokens,
-                        "total_tokens": cb.total_tokens
-                    }
-            return result
         except asyncio.CancelledError:
             logger.info(f"Agent执行被取消: session_id={self.session_id}")
-            return {
-                "output": "任务已取消",
-                "intermediate_steps": [],
-                "token_usage": {}
-            }
+            return "任务已取消", {}
         except Exception as e:
             logger.error(f"Agent执行失败: {e}")
-            return {
-                "output": str(e),
-                "intermediate_steps": [],
-                "token_usage": {}
-            }
+            return str(e), {}
 
     async def send_agent_message(self, message: str, title: str = "MoviePilot助手"):
         """
@@ -468,7 +191,7 @@ class MoviePilotAgent:
                 userid=self.user_id,
                 username=self.username,
                 title=title,
-                text=message
+                text=message,
             )
         )
 
@@ -492,38 +215,44 @@ class AgentManager:
         """
         初始化管理器
         """
-        await conversation_manager.initialize()
+        await memory_manager.initialize()
 
     async def close(self):
         """
         关闭管理器
         """
-        await conversation_manager.close()
-        # 清理所有活跃的智能体
+        await memory_manager.close()
         for agent in self.active_agents.values():
             await agent.cleanup()
         self.active_agents.clear()
 
-    async def process_message(self, session_id: str, user_id: str, message: str,
-                              channel: str = None, source: str = None, username: str = None) -> str:
+    async def process_message(
+            self,
+            session_id: str,
+            user_id: str,
+            message: str,
+            channel: str = None,
+            source: str = None,
+            username: str = None,
+    ) -> str:
         """
         处理用户消息
         """
-        # 获取或创建Agent实例
         if session_id not in self.active_agents:
-            logger.info(f"创建新的AI智能体实例，session_id: {session_id}, user_id: {user_id}")
+            logger.info(
+                f"创建新的AI智能体实例，session_id: {session_id}, user_id: {user_id}"
+            )
             agent = MoviePilotAgent(
                 session_id=session_id,
                 user_id=user_id,
                 channel=channel,
                 source=source,
-                username=username
+                username=username,
             )
             self.active_agents[session_id] = agent
         else:
             agent = self.active_agents[session_id]
-            agent.user_id = user_id  # 确保user_id是最新的
-            # 更新渠道信息
+            agent.user_id = user_id
             if channel:
                 agent.channel = channel
             if source:
@@ -531,8 +260,7 @@ class AgentManager:
             if username:
                 agent.username = username
 
-        # 处理消息
-        return await agent.process_message(message)
+        return await agent.process(message)
 
     async def clear_session(self, session_id: str, user_id: str):
         """
@@ -542,7 +270,7 @@ class AgentManager:
             agent = self.active_agents[session_id]
             await agent.cleanup()
             del self.active_agents[session_id]
-            await conversation_manager.clear_memory(session_id, user_id)
+            await memory_manager.clear_memory(session_id, user_id)
             logger.info(f"会话 {session_id} 的记忆已清空")