create the base chatbot agent

src/agents/__init__.py

@@ -1,3 +1,7 @@
+from src.agents.registry import BaseAgent
+from src.agents.chatbot.graph import ChatbotAgent
+from langchain_core.runnables import RunnableConfig
+from langchain_core.messages import AIMessageChunk, ToolMessage
 
 
 class AgentManager:
@@ -5,4 +9,45 @@ class AgentManager:
         self.agents = {}
 
     def add_agent(self, agent_id, agent):
-        self.agents[agent_id] = agent
+        self.agents[agent_id] = {
+            "agent": agent,
+            "configuration": agent.configuration
+        }
+
+def agent_cli(agent: BaseAgent, config: RunnableConfig = None):
+    config = config or {}
+    if "configurable" not in config:
+        config["configurable"] = {}
+
+    while True:
+        user_input = input("\nUser: ")
+        if user_input.lower() in ["quit", "exit", "q"]:
+            print("Goodbye!")
+            break
+
+        stream_flag = False
+        for msg, metadata in agent.stream_messages([{"role": "user", "content": user_input}], config):
+            if isinstance(msg, AIMessageChunk):
+                content = msg.content or msg.tool_calls
+
+                if not content:
+                    if stream_flag == True:
+                        print()
+                        stream_flag = False
+                    continue
+
+                if stream_flag == False and content:
+                    print(f"AI: {content}", end="", flush=True)
+                    stream_flag = True
+                    continue
+
+                elif content:
+                    print(f"{content}", end="", flush=True)
+
+            if isinstance(msg, ToolMessage):
+                print(f"Tool: {msg.content}")
+
+def get_agents():
+    agent_manager = AgentManager()
+    agent_manager.add_agent("chatbot", ChatbotAgent())
+    return agent_manager.agents

src/agents/chatbot/__init__.py

@@ -1,3 +1,3 @@
-from .graph import graph
+from .graph import ChatbotAgent
 
-__all__ = ["graph"]
+__all__ = ["ChatbotAgent"]

src/agents/chatbot/configuration.py

@@ -1,6 +1,7 @@
 from dataclasses import dataclass, field
 from typing import List
 
+from langchain_core.messages import SystemMessage
 from langchain_core.tools import Tool
 from langchain_openai import ChatOpenAI
 
@@ -8,9 +9,20 @@ from src import config
 from src.models import select_model
 from src.agents.registry import Configuration
 
+def get_default_llm():
+    return select_model(config).chat_open_ai
+
+def get_default_tools():
+    from langchain_community.tools.tavily_search import TavilySearchResults
+    return [TavilySearchResults(max_results=10)]
+
+def get_default_requirements():
+    return ["TAVILY_API_KEY"]
 
 @dataclass(kw_only=True)
 class ChatbotConfiguration(Configuration):
     name: str = "chatbot"
-    llm: ChatOpenAI = field(default_factory=select_model().chat_open_ai)
-    tools: List[Tool] = field(default_factory=list)
+    description: str = "A chatbot that can answer questions and help with tasks."
+    requirements: list[str] = field(default_factory=get_default_requirements)
+    tools: list[Tool] = field(default_factory=get_default_tools)
+    llm: ChatOpenAI = field(default_factory=get_default_llm)

src/agents/chatbot/graph.py

@@ -1,35 +1,83 @@
-"""Define a simple chatbot agent.
-
-This agent returns a predefined response without using an actual LLM.
-"""
-
-from typing import Any, Dict
+import asyncio
+import uuid
+from typing import Any
+from datetime import datetime
 
+from langchain_core.messages import AIMessageChunk, ToolMessage
 from langchain_core.runnables import RunnableConfig
-from langgraph.graph import StateGraph
+from langgraph.graph import StateGraph, START, END
+from langgraph.prebuilt import ToolNode, tools_condition
+from langgraph.checkpoint.memory import MemorySaver
 
 from src.agents.registry import State, BaseAgent
 from src.agents.chatbot.configuration import ChatbotConfiguration
-from src.agents.tools_factory import search_tool
 
 class ChatbotAgent(BaseAgent):
-    def __init__(self, configuration: ChatbotConfiguration):
+    _graph_cache = None
+
+    def __init__(self, configuration: ChatbotConfiguration = None):
         super().__init__(configuration)
+        self.configuration = configuration or ChatbotConfiguration()
         self.llm = configuration.llm
-        self.tools = [search_tool]
 
-    async def _get_tools(self):
-        return self.tools
+    def llm_call(self, state: State) -> dict[str, Any]:
+        tools = self.configuration.tools
+        system_prompt = {
+            "role": "system",
+            "content": (
+                f"Current time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n"
+            )
+        }
+        messages = [system_prompt] + state["messages"]
+        model = self.llm.bind_tools(tools)
 
-    async def llm_call(self, state: State) -> Dict[str, Any]:
-        return await self.llm.invoke(state.messages)
+        res = model.invoke(messages)
+        return {"messages": [res]}
 
-    async def get_graph(self):
-        workflow = StateGraph(State, config_schema=ChatbotConfiguration)
-        workflow.add_node("chat", self.chat)
-        workflow.add_edge("__start__", "chat")
-        return workflow.compile()
+    def get_graph(self):
+        """构建图"""
+        if ChatbotAgent._graph_cache is None:
+            workflow = StateGraph(State)
+            workflow.add_node("chatbot", self.llm_call)
+            workflow.add_node("tools", ToolNode(tools=self.configuration.tools))
+            workflow.add_edge(START, "chatbot")
+            workflow.add_conditional_edges(
+                "chatbot",
+                tools_condition,
+            )
+            workflow.add_edge("tools", "chatbot")
+            workflow.add_edge("chatbot", END)
 
-    async def invoke(self, messages: List[str], config: RunnableConfig) -> Dict[str, Any]:
-        return await self.get_graph().invoke({"messages": messages}, config)
+            graph = workflow.compile(checkpointer=MemorySaver())
+            ChatbotAgent._graph_cache = graph
 
+        return ChatbotAgent._graph_cache
+
+    def stream_values(self, messages: list[str], config: RunnableConfig = None):
+        graph = self.get_graph()
+        for event in graph.stream({"messages": messages}, stream_mode="values", config=config):
+            yield event["messages"]
+
+    def stream_messages(self, messages: list[str], config: RunnableConfig = None):
+        graph = self.get_graph()
+        for msg, metadata in graph.stream({"messages": messages}, stream_mode="messages", config=config):
+            msg_type = msg.type
+
+            return_keys = config.get("configurable", {}).get("return_keys", [])
+            if not return_keys or msg_type in return_keys:
+                yield msg, metadata
+
+
+def main():
+    agent = ChatbotAgent(ChatbotConfiguration())
+
+    thread_id = str(uuid.uuid4())
+    config = {"configurable": {"thread_id": thread_id}}
+
+    from src.agents import agent_cli
+    agent_cli(agent, config)
+
+
+if __name__ == "__main__":
+    main()
+    # asyncio.run(main())

src/agents/tools_factory.py

@@ -1,16 +1,97 @@
+import json
+import re
+import os
+from typing import Any, Callable, Optional, Type, Union
+
+
 from pydantic import BaseModel, Field
-from langchain_core.tools import Tool, StructuredTool
+from langchain.agents import tool
+from langchain_core.tools import Tool, BaseTool
 
-from src.utils.web_search import WebSearcher
 
-class SearchArgsSchema(BaseModel):
-    query: str = Field(..., description="The query to search the web for")
+_TOOLS_REGISTRY = {}
 
-search_tool = Tool(
-    name="search",
-    description="Search the web for information",
-    func=WebSearcher().search,
-    args_schema=SearchArgsSchema,
-    return_direct=True,
-)
+# refs https://github.com/chatchat-space/LangGraph-Chatchat chatchat-server/chatchat/server/agent/tools_factory/tools_registry.py
+def regist_tool(
+    *args: Any,
+    title: str = "",
+    description: str = "",
+    return_direct: bool = False,
+    args_schema: Optional[Type[BaseModel]] = None,
+    infer_schema: bool = True,
+) -> Union[Callable, BaseTool]:
+    """
+    wrapper of langchain tool decorator
+    add tool to registry automatically
+    """
+
+    def _parse_tool(t: BaseTool):
+        nonlocal description, title
+
+        _TOOLS_REGISTRY[t.name] = t
+
+        # change default description
+        if not description:
+            if t.func is not None:
+                description = t.func.__doc__
+            elif t.coroutine is not None:
+                description = t.coroutine.__doc__
+        t.description = " ".join(re.split(r"\n+\s*", description))
+        # set a default title for human
+        if not title:
+            title = "".join([x.capitalize() for x in t.name.split("_")])
+        setattr(t, "_title", title)
+
+    def wrapper(def_func: Callable) -> BaseTool:
+        partial_ = tool(
+            *args,
+            return_direct=return_direct,
+            args_schema=args_schema,
+            infer_schema=infer_schema,
+        )
+        t = partial_(def_func)
+        _parse_tool(t)
+        return t
+
+    if len(args) == 0:
+        return wrapper
+    else:
+        t = tool(
+            *args,
+            return_direct=return_direct,
+            args_schema=args_schema,
+            infer_schema=infer_schema,
+        )
+        _parse_tool(t)
+        return t
+
+
+class BaseToolOutput:
+    """
+    LLM 要求 Tool 的输出为 str，但 Tool 用在别处时希望它正常返回结构化数据。
+    只需要将 Tool 返回值用该类封装，能同时满足两者的需要。
+    基类简单的将返回值字符串化，或指定 format="json" 将其转为 json。
+    用户也可以继承该类定义自己的转换方法。
+    """
+
+    def __init__(
+        self,
+        data: Any,
+        format: str | Callable = None,
+        data_alias: str = "",
+        **extras: Any,
+    ) -> None:
+        self.data = data
+        self.format = format
+        self.extras = extras
+        if data_alias:
+            setattr(self, data_alias, property(lambda obj: obj.data))
+
+    def __str__(self) -> str:
+        if self.format == "json":
+            return json.dumps(self.data, ensure_ascii=False, indent=2)
+        elif callable(self.format):
+            return self.format(self)
+        else:
+            return str(self.data)