diff --git a/src/agents/__init__.py b/src/agents/__init__.py index 3759c5b9..0ce6de15 100644 --- a/src/agents/__init__.py +++ b/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 diff --git a/src/agents/chatbot/__init__.py b/src/agents/chatbot/__init__.py index 92e9d0b0..d6f7031e 100644 --- a/src/agents/chatbot/__init__.py +++ b/src/agents/chatbot/__init__.py @@ -1,3 +1,3 @@ -from .graph import graph +from .graph import ChatbotAgent -__all__ = ["graph"] \ No newline at end of file +__all__ = ["ChatbotAgent"] \ No newline at end of file diff --git a/src/agents/chatbot/configuration.py b/src/agents/chatbot/configuration.py index bdd8b788..65d09178 100644 --- a/src/agents/chatbot/configuration.py +++ b/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) \ No newline at end of file + 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) diff --git a/src/agents/chatbot/graph.py b/src/agents/chatbot/graph.py index 99b00695..7ec47f45 100644 --- a/src/agents/chatbot/graph.py +++ b/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()) diff --git a/src/agents/tools_factory.py b/src/agents/tools_factory.py index b9f2780e..95d04426 100644 --- a/src/agents/tools_factory.py +++ b/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)