Add custom subagent logic

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],
+)

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})

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

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