mini-openclaw/backend/graph/agent.py

"""AgentManager — Core Agent using LangChain create_agent API with DeepSeek."""

import os
from pathlib import Path
from typing import Any, AsyncGenerator

from langchain_core.messages import HumanMessage, AIMessage

from config import get_rag_mode
from graph.prompt_builder import build_system_prompt
from graph.session_manager import session_manager
from graph.chat_model_options import chat_model_options
from tools import get_all_tools


class AgentManager:
    """Manages the Agent lifecycle: initialization, streaming, invocation."""

    def __init__(self) -> None:
        self._base_dir: Path | None = None
        self._tools: list = []
        self._llm = None
        self._raw_logs_dir: Path | None = None

    def initialize(self, base_dir: Path) -> None:
        """Initialize LLM (DeepSeek) and tools. Called once at startup."""
        self._base_dir = base_dir
        self._tools = get_all_tools(base_dir)

        # 初始化原始报文日志目录 (base_dir 已是 backend/ 目录)
        self._raw_logs_dir = base_dir / "raw_logs"
        self._raw_logs_dir.mkdir(exist_ok=True, parents=True)

        # Use langchain-deepseek for native DeepSeek support
        from langchain_deepseek import ChatDeepSeek

        self._llm = ChatDeepSeek(
            model=os.getenv("DEEPSEEK_MODEL", "deepseek-chat"),
            api_key=os.getenv("DEEPSEEK_API_KEY"),
            api_base=os.getenv("DEEPSEEK_BASE_URL", "https://api.deepseek.com"),
            temperature=0.7,
            streaming=True,
            **chat_model_options(),
        )

        session_manager.initialize(base_dir)
        print(f"🤖 Agent initialized with {len(self._tools)} tools (DeepSeek: {os.getenv('DEEPSEEK_MODEL', 'deepseek-chat')})")

    def _build_agent(self):
        """Build a fresh agent with current system prompt (re-reads files each time)."""
        assert self._base_dir is not None
        assert self._llm is not None

        rag_mode = get_rag_mode()
        system_prompt = build_system_prompt(self._base_dir, rag_mode=rag_mode)

        return self.build_task_agent(self._llm, self._tools, system_prompt)

    def build_task_agent(self, model, tools: list, system_prompt: str):
        """Shared core factory. Scoped tasks supply explicit model, tools and prompt."""
        from langchain.agents import create_agent
        return create_agent(model=model, tools=tools, system_prompt=system_prompt)

    def invoke_scoped(self, message: str, *, model, tools: list, system_prompt: str) -> dict:
        """Run an isolated task through the same core without shared chat memory.

        No mutation of the default model/tools: concurrent projects cannot share state.
        The caller stores the trace under its own authorization boundary.
        """
        from langchain_core.messages import messages_to_dict
        agent = self.build_task_agent(model, tools, system_prompt)
        result = agent.invoke(
            {"messages": self._build_messages(message, [])},
            config={"recursion_limit": 16, "callbacks": []},
        )
        messages = result.get("messages", [])
        if not messages or messages[-1].type != "ai" or getattr(messages[-1], "tool_calls", None):
            raise ValueError("Agent 未完成工具调用循环")
        content = messages[-1].content
        if not isinstance(content, str):
            content = "\n".join(b.get("text", "") for b in content if isinstance(b, dict) and b.get("type") == "text")
        if not content.strip():
            raise ValueError("Agent 未返回有效建议")
        return {"text": content, "messages": messages_to_dict(messages), "engine": "mini-openclaw/AgentManager"}

    def _build_messages(self, user_message: str, history: list[dict[str, Any]]) -> list:
        """Convert session history + new message into LangChain messages."""
        messages = []
        for msg in history:
            role = msg.get("role", "")
            content = msg.get("content", "")
            if role == "user":
                messages.append(HumanMessage(content=content))
            elif role == "assistant":
                messages.append(AIMessage(content=content))
        messages.append(HumanMessage(content=user_message))
        return messages

    async def astream(
        self, message: str, history: list[dict[str, Any]], session_id: str = "default"
    ) -> AsyncGenerator[dict[str, Any], None]:
        """Stream agent response with token-level and node-level events.

        Yields events:
          {"type": "retrieval", "query": "...", "results": [...]}  (RAG mode only)
          {"type": "token", "content": "..."}
          {"type": "tool_start", "tool": "...", "input": "..."}
          {"type": "tool_end", "tool": "...", "output": "..."}
          {"type": "done", "content": "..."}

        Also persists complete message flow to {session}.msgs.jsonl for trajectory.
        """
        # RAG retrieval: inject memory context if enabled
        rag_mode = get_rag_mode()
        rag_context = ""
        if rag_mode and self._base_dir:
            from graph.memory_indexer import get_memory_indexer

            indexer = get_memory_indexer(self._base_dir)
            results = indexer.retrieve(message)
            if results:
                yield {
                    "type": "retrieval",
                    "query": message,
                    "results": results,
                }
                snippets = "\n\n".join(
                    f"[片段 {i+1}] (score: {r['score']})\n{r['text']}"
                    for i, r in enumerate(results)
                )
                rag_context = f"[记忆检索结果]\n{snippets}"

        # 为当前 session 创建原始报文回调
        from graph.raw_trace import RawTraceCallback

        raw_callback = RawTraceCallback(session_id, self._raw_logs_dir)

        # 消息流日志文件（记录完整的 human/assistant/tool 消息序列）
        import json
        import time
        msgs_log_file = self._raw_logs_dir / f"{session_id}.msgs.jsonl"

        # 记录 human 消息
        def _append_msg(role: str, content: str, **kwargs):
            """追加一条消息到 msgs.jsonl"""
            msg_record = {
                "role": role,
                "content": content,
                "timestamp": time.time(),
                **kwargs
            }
            with open(msgs_log_file, "a", encoding="utf-8") as f:
                f.write(json.dumps(msg_record, ensure_ascii=False) + "\n")

        _append_msg("human", message)

        # 构建 agent 时传入 callback
        agent = self._build_agent()

        # Build messages with optional RAG context appended to history
        augmented_history = list(history)
        if rag_context:
            augmented_history.append(
                {"role": "assistant", "content": rag_context}
            )
        messages = self._build_messages(message, augmented_history)

        full_response = ""
        tools_just_finished = False
        current_assistant_content = ""  # 当前 assistant 消息累积
        current_tool_calls = []  # 当前 assistant 的 tool_calls

        async for event in agent.astream(
            {"messages": messages},
            stream_mode=["messages", "updates"],
            config={"callbacks": [raw_callback]},
        ):
            # event is a tuple of (stream_mode, data) when using multiple modes
            if isinstance(event, tuple):
                mode, data = event
            else:
                mode = "messages"
                data = event

            if mode == "messages":
                # Token-level streaming from LLM
                msg, metadata = data
                if hasattr(msg, "content") and msg.content:
                    if msg.type == "AIMessageChunk" or msg.type == "ai":
                        if msg.content and not getattr(msg, "tool_calls", None):
                            # If tools just finished, signal a new response segment
                            if tools_just_finished:
                                yield {"type": "new_response"}
                                tools_just_finished = False
                            full_response += msg.content
                            current_assistant_content += msg.content
                            yield {"type": "token", "content": msg.content}

            elif mode == "updates":
                if isinstance(data, dict):
                    for node_name, node_data in data.items():
                        if node_name == "model" and "messages" in node_data:
                            # 模型节点：记录 assistant 消息和 tool_calls
                            for agent_msg in node_data["messages"]:
                                if hasattr(agent_msg, "tool_calls") and agent_msg.tool_calls:
                                    # 记录 assistant 消息（带 tool_calls）
                                    tool_calls_data = []
                                    for tc in agent_msg.tool_calls:
                                        tool_call_info = {
                                            "name": tc["name"],
                                            "args": tc.get("args", {})
                                        }
                                        tool_calls_data.append(tool_call_info)
                                        current_tool_calls.append(tool_call_info)
                                        yield {
                                            "type": "tool_start",
                                            "tool": tc["name"],
                                            "input": str(tc.get("args", ""))[:1000],
                                        }

                                    # 记录带 tool_calls 的 assistant 消息
                                    _append_msg(
                                        "assistant",
                                        agent_msg.content if hasattr(agent_msg, "content") else "",
                                        tool_calls=tool_calls_data
                                    )

                        elif node_name == "tools" and "messages" in node_data:
                            # 工具节点：记录每个 ToolMessage 的返回内容
                            for tool_msg in node_data["messages"]:
                                if hasattr(tool_msg, "name"):
                                    tool_output = str(tool_msg.content)
                                    # 记录 tool 消息（完整 content）
                                    _append_msg(
                                        "tool",
                                        tool_output,
                                        name=tool_msg.name,
                                        tool_call_id=getattr(tool_msg, "tool_call_id", None)
                                    )
                                    yield {
                                        "type": "tool_end",
                                        "tool": tool_msg.name,
                                        "output": tool_output[:2000],  # yield 时截断，但完整内容已记录
                                    }
                            # After all tool results, mark that tools finished
                            tools_just_finished = True

        # 记录最终的 assistant 消息（纯回复，无 tool_calls）
        if current_assistant_content and not current_tool_calls:
            _append_msg("assistant", current_assistant_content)
        elif current_assistant_content and current_tool_calls:
            # 如果有 tool_calls，记录最后一轮的纯文本回复
            _append_msg("assistant", current_assistant_content)

        yield {"type": "done", "content": full_response}

    async def ainvoke(self, message: str, session_id: str) -> str:
        """Non-streaming invocation (fallback)."""
        history = session_manager.load_session(session_id)
        agent = self._build_agent()
        messages = self._build_messages(message, history)
        result = await agent.ainvoke({"messages": messages})

        final_messages = result.get("messages", [])
        for msg in reversed(final_messages):
            if hasattr(msg, "content") and msg.type == "ai" and msg.content:
                response = msg.content
                session_manager.save_message(session_id, "user", message)
                session_manager.save_message(session_id, "assistant", response)
                return response
        return "No response generated."


agent_manager = AgentManager()