"""Sub-agent infrastructure for specialized tool routing. This module provides the SubAgent class that wraps specialized agents with their own tools and system prompts. """ import logging from typing import List, Optional, AsyncIterator from pathlib import Path from langchain_anthropic import ChatAnthropic from langchain_core.messages import HumanMessage, SystemMessage from langchain_core.runnables import RunnableConfig from langgraph.prebuilt import create_react_agent from langgraph.checkpoint.memory import MemorySaver logger = logging.getLogger(__name__) class SubAgent: """A specialized sub-agent with its own tools and system prompt. Sub-agents are lightweight, stateless agents that focus on specific domains. They use in-memory checkpointing since they don't need persistent state. """ def __init__( self, name: str, soul_file: str, tools: List, model_name: str = "claude-sonnet-4-20250514", temperature: float = 0.7, api_key: Optional[str] = None, base_dir: str = "." ): """Initialize a sub-agent. Args: name: Agent name (e.g., "chart", "data", "automation") soul_file: Filename in /soul directory (e.g., "chart_agent.md") tools: List of LangChain tools for this agent model_name: Anthropic model name temperature: Model temperature api_key: Anthropic API key base_dir: Base directory for resolving paths """ self.name = name self.soul_file = soul_file self.tools = tools self.model_name = model_name self.temperature = temperature self.api_key = api_key self.base_dir = base_dir # Load system prompt from soul file soul_path = Path(base_dir) / "soul" / soul_file if soul_path.exists(): with open(soul_path, "r") as f: self.system_prompt = f.read() logger.info(f"SubAgent '{name}': Loaded system prompt from {soul_path}") else: logger.warning(f"SubAgent '{name}': Soul file not found at {soul_path}, using default") self.system_prompt = f"You are a specialized {name} agent." # Initialize LLM self.llm = ChatAnthropic( model=model_name, temperature=temperature, api_key=api_key, streaming=True ) # Create agent with in-memory checkpointer (stateless) checkpointer = MemorySaver() self.agent = create_react_agent( self.llm, tools, checkpointer=checkpointer ) logger.info( f"SubAgent '{name}' initialized with {len(tools)} tools, " f"model={model_name}, temp={temperature}" ) async def execute( self, task: str, thread_id: Optional[str] = None ) -> str: """Execute a task with this sub-agent. Args: task: The task/prompt for this sub-agent thread_id: Optional thread ID for checkpointing (uses ephemeral ID if not provided) Returns: The agent's complete response as a string """ import uuid # Use ephemeral thread ID if not provided if thread_id is None: thread_id = f"subagent-{self.name}-{uuid.uuid4()}" logger.info(f"SubAgent '{self.name}': Executing task (thread_id={thread_id})") logger.debug(f"SubAgent '{self.name}': Task: {task[:200]}...") # Build messages with system prompt messages = [ HumanMessage(content=task) ] # Prepare config with system prompt injection config = RunnableConfig( configurable={ "thread_id": thread_id, "state_modifier": self.system_prompt }, metadata={ "subagent_name": self.name } ) # Execute and collect response full_response = "" event_count = 0 try: async for event in self.agent.astream_events( {"messages": messages}, config=config, version="v2" ): event_count += 1 # Log tool calls if event["event"] == "on_tool_start": tool_name = event.get("name", "unknown") logger.debug(f"SubAgent '{self.name}': Tool call started: {tool_name}") elif event["event"] == "on_tool_end": tool_name = event.get("name", "unknown") logger.debug(f"SubAgent '{self.name}': Tool call completed: {tool_name}") # Extract streaming tokens elif event["event"] == "on_chat_model_stream": chunk = event["data"]["chunk"] if hasattr(chunk, "content") and chunk.content: content = chunk.content # Handle both string and list content if isinstance(content, list): text_parts = [] for block in content: if isinstance(block, dict) and "text" in block: text_parts.append(block["text"]) elif hasattr(block, "text"): text_parts.append(block.text) content = "".join(text_parts) if content: full_response += content logger.info( f"SubAgent '{self.name}': Completed task " f"({event_count} events, {len(full_response)} chars)" ) except Exception as e: error_msg = f"SubAgent '{self.name}' execution error: {str(e)}" logger.error(error_msg, exc_info=True) return f"Error: {error_msg}" return full_response async def stream( self, task: str, thread_id: Optional[str] = None ) -> AsyncIterator[str]: """Execute a task with streaming response. Args: task: The task/prompt for this sub-agent thread_id: Optional thread ID for checkpointing Yields: Response chunks as they're generated """ import uuid # Use ephemeral thread ID if not provided if thread_id is None: thread_id = f"subagent-{self.name}-{uuid.uuid4()}" logger.info(f"SubAgent '{self.name}': Streaming task (thread_id={thread_id})") # Build messages with system prompt messages = [ HumanMessage(content=task) ] # Prepare config config = RunnableConfig( configurable={ "thread_id": thread_id, "state_modifier": self.system_prompt }, metadata={ "subagent_name": self.name } ) # Stream response try: async for event in self.agent.astream_events( {"messages": messages}, config=config, version="v2" ): # Log tool calls if event["event"] == "on_tool_start": tool_name = event.get("name", "unknown") logger.debug(f"SubAgent '{self.name}': Tool call started: {tool_name}") # Extract streaming tokens elif event["event"] == "on_chat_model_stream": chunk = event["data"]["chunk"] if hasattr(chunk, "content") and chunk.content: content = chunk.content # Handle both string and list content if isinstance(content, list): text_parts = [] for block in content: if isinstance(block, dict) and "text" in block: text_parts.append(block["text"]) elif hasattr(block, "text"): text_parts.append(block.text) content = "".join(text_parts) if content: yield content except Exception as e: error_msg = f"SubAgent '{self.name}' streaming error: {str(e)}" logger.error(error_msg, exc_info=True) yield f"Error: {error_msg}"