ForcePilot/src/agents/chatbot/graph.py

import os
import uuid
from typing import Any, cast, Annotated
from pathlib import Path
from datetime import datetime, timezone

from dataclasses import dataclass, field, fields
from langchain_core.messages import AIMessage, ToolMessage
from langgraph.graph import StateGraph, START, END
from langgraph.runtime import Runtime
from langgraph.prebuilt import ToolNode, tools_condition
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver, aiosqlite
from langgraph.checkpoint.memory import InMemorySaver

from src import config as sys_config
from src.utils import logger
from src.agents.common.utils import get_cur_time_with_utc
from src.agents.common.base import BaseAgent
from src.agents.common.models import load_chat_model

from .state import State
from .context import Context
from .tools import get_tools



class ChatbotAgent(BaseAgent):
    name = "智能体助手"
    description = "基础的对话机器人，可以回答问题，默认不使用任何工具，可在配置中启用需要的工具。"

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.graph = None
        self.context_schema = Context
        self.workdir = Path(sys_config.save_dir) / "agents" / self.module_name
        self.workdir.mkdir(parents=True, exist_ok=True)
        self.agent_tools = get_tools()

    def _get_tools(self, tools: list[str]):
        """根据配置获取工具。
        默认不使用任何工具。
        如果配置为列表，则使用列表中的工具。
        """
        self.agent_tools = get_tools()
        if tools is None or not isinstance(tools, list) or len(tools) == 0:
            # 默认不使用任何工具
            logger.info("未配置工具或配置为空，不使用任何工具")
            return []
        else:
            # 使用配置中指定的工具
            tools = [tool for tool in self.agent_tools if tool.name in tools]
            logger.info(f"使用工具: {[tool.name for tool in tools]}")
            return tools

    async def llm_call(self, state: State, runtime: Runtime[Context] = None) -> dict[str, Any]:
        """调用 llm 模型 - 异步版本以支持异步工具"""
        system_prompt = f"{runtime.context.system_prompt}. Current time is {get_cur_time_with_utc()}"
        model = load_chat_model(runtime.context.model)

        # 这里要根据配置动态获取工具
        if tools := self._get_tools(runtime.context.tools):
            model = model.bind_tools(tools)

        # 使用异步调用
        response = cast(
            AIMessage,
            await model.ainvoke(
                [{"role": "system", "content": system_prompt}, *state.messages]
            ),
        )
        return {"messages": [response]}


    async def dynamic_tools_node(
        self, state: State, runtime: Runtime[Context]
    ) -> dict[str, list[ToolMessage]]:
        """Execute tools dynamically based on configuration.

        This function gets the available tools based on the current configuration
        and executes the requested tool calls from the last message.
        """
        # Get available tools based on configuration
        available_tools = get_tools()

        # Create a ToolNode with the available tools
        tool_node = ToolNode(available_tools)

        # Execute the tool node
        result = await tool_node.ainvoke(state)

        return cast(dict[str, list[ToolMessage]], result)

    async def get_graph(self, **kwargs):
        """构建图"""
        if self.graph:
            return self.graph

        runnable_tools = get_tools()
        logger.debug(f"build graph `{self.id}` with {len(runnable_tools)} tools")

        builder = StateGraph(State, context_schema=self.context_schema)
        builder.add_node("chatbot", self.llm_call)
        builder.add_node("tools", self.dynamic_tools_node)
        builder.add_edge(START, "chatbot")
        builder.add_conditional_edges(
            "chatbot",
            tools_condition,
        )
        builder.add_edge("tools", "chatbot")
        builder.add_edge("chatbot", END)

        # 创建数据库连接并确保设置 checkpointer
        try:
            sqlite_checkpointer = AsyncSqliteSaver(await self.get_async_conn())
            graph = builder.compile(checkpointer=sqlite_checkpointer, name=self.name)
            self.graph = graph
            return graph
        except Exception as e:
            logger.error(f"构建 Graph 设置 checkpointer 时出错: {e}, 尝试使用内存存储")
            # 即使出错也返回一个可用的图实例，只是无法保存历史
            checkpointer = InMemorySaver()
            graph = builder.compile(checkpointer=checkpointer, name=self.name)
            self.graph = graph
            return graph

    async def get_async_conn(self) -> aiosqlite.Connection:
        """获取异步数据库连接"""
        return await aiosqlite.connect(os.path.join(self.workdir, "aio_history.db"))

    async def get_aio_memory(self) -> AsyncSqliteSaver:
        """获取异步存储实例"""
        return AsyncSqliteSaver(await self.get_async_conn())

def main():
    agent = ChatbotAgent(Context)

    thread_id = str(uuid.uuid4())
    config = {"configurable": {"thread_id": thread_id}}

    from src.agents.utils import agent_cli
    agent_cli(agent, config)


if __name__ == "__main__":
    main()
    # asyncio.run(main())