diff --git a/examples/research/research_agent.py b/examples/research/research_agent.py index f98457d7e6..06518734c2 100644 --- a/examples/research/research_agent.py +++ b/examples/research/research_agent.py @@ -8,6 +8,8 @@ from langchain_core.messages import ToolMessage, tool from claude_everything.graph import create_deep_agent +from claude_everything.sub_agent import create_task_tool +from claude_everything.tools import write_todos from langgraph.prebuilt.chat_agent_executor import AgentState @@ -17,15 +19,20 @@ class ResearchAgentState(AgentState): # Search tool to use to do research -def internet_search(query, max_results: int = 5, topic: Literal["general", "news", "finance"] = "general", include_raw_content: bool = False): +def internet_search( + query, + max_results: int = 5, + topic: Literal["general", "news", "finance"] = "general", + include_raw_content: bool = False, +): """Run a web search""" tavily_async_client = TavilyClient(api_key=os.environ["TAVILY_API_KEY"]) search_docs = tavily_async_client.search( - query, - max_results=max_results, - include_raw_content=include_raw_content, - topic=topic - ) + query, + max_results=max_results, + include_raw_content=include_raw_content, + topic=topic, + ) return search_docs @@ -34,12 +41,12 @@ def write_report(report: str, tool_call_id: Annotated[str, InjectedToolCallId]): """Use this to write your final report to a file. the `report` argument should be the whole report. Make sure it is comprehensive.""" - return Command(update={ - "report": report, - "messages": [ - ToolMessage(f"Wrote final report", tool_call_id=tool_call_id) - ] - }) + return Command( + update={ + "report": report, + "messages": [ToolMessage(f"Wrote final report", tool_call_id=tool_call_id)], + } + ) # Prompt prefix to steer the agent to be an expert researcher @@ -52,5 +59,14 @@ def write_report(report: str, tool_call_id: Annotated[str, InjectedToolCallId]): Use this to run an internet search for a given query. You can specify the number of results, the topic, and whether raw content should be included. """ -# Create the agent -agent = create_deep_agent([internet_search], research_prompt_prefix, state_schema=ResearchAgentState, main_agent_tools=[write_report]) +# use-case specific tools +research_tools = [internet_search] + +task_tool = create_task_tool(research_tools) + +agent = create_deep_agent( + [task_tool], + research_prompt_prefix, + state_schema=ResearchAgentState, + main_agent_tools=[write_report], +) diff --git a/src/claude_everything/graph.py b/src/claude_everything/graph.py index fb267161ce..7bf5479a6b 100644 --- a/src/claude_everything/graph.py +++ b/src/claude_everything/graph.py @@ -1,4 +1,3 @@ -from claude_everything.sub_agent import create_task_tool from claude_everything.model import model from claude_everything.tools import write_todos @@ -17,13 +16,12 @@ - When doing web search, prefer to use the `task` tool in order to reduce context usage.""" - def create_deep_agent(tools, prompt_prefix, state_schema=None, main_agent_tools=None): prompt = prompt_prefix + base_prompt main_extra_tools = main_agent_tools or [] return create_react_agent( model, prompt=prompt, - tools=[create_task_tool(tools, prompt_prefix), write_todos] + tools + main_extra_tools, - state_schema=state_schema + tools=[write_todos] + tools + main_extra_tools, + state_schema=state_schema, ).with_config({"recursion_limit": 1000}) diff --git a/src/claude_everything/sub_agent.py b/src/claude_everything/sub_agent.py index 3ce5ed168a..3fd1391431 100644 --- a/src/claude_everything/sub_agent.py +++ b/src/claude_everything/sub_agent.py @@ -2,22 +2,30 @@ from claude_everything.prompts import TASK_DESCRIPTION from langgraph.prebuilt import create_react_agent from langchain_core.tools import tool +from .sub_agent_registry import SUBAGENT_REGISTRY, register_subagent -def create_task_tool(tools, prompt_prefix): - sub_agent = create_react_agent( +@register_subagent("general-purpose") +def make_general_purpose_agent(tools): + return create_react_agent( model, - prompt=prompt_prefix, - tools=tools + prompt=TASK_DESCRIPTION, + tools=tools or [], ) + +def create_task_tool(tools): @tool(description=TASK_DESCRIPTION) - def task(description: str, subagent_type: str): - if subagent_type != "general-purpose": - return f"Error: invoked agent of type {subagent_type}, the only allowed type is `general-purpose`" - result = sub_agent.invoke({ - "messages": [{"role": "user", "content": description}] - }) - return result['messages'][-1].content + def task(description: str, subagent_type: str = None): + if not subagent_type: + subagent_type = "general-purpose" + agent_factory = SUBAGENT_REGISTRY.get(subagent_type) + if not agent_factory: + return f"Error: unknown agent type {subagent_type}" + sub_agent = agent_factory(tools) + result = sub_agent.invoke( + {"messages": [{"role": "user", "content": description}]} + ) + return result["messages"][-1].content return task diff --git a/src/claude_everything/sub_agent_registry.py b/src/claude_everything/sub_agent_registry.py new file mode 100644 index 0000000000..ac574ac2cf --- /dev/null +++ b/src/claude_everything/sub_agent_registry.py @@ -0,0 +1,10 @@ +# in-memory registry of sub-agents +SUBAGENT_REGISTRY = {} + + +def register_subagent(agent_type): + def decorator(factory_fn): + SUBAGENT_REGISTRY[agent_type] = factory_fn + return factory_fn + + return decorator