naas-abi-core 1.4.1__py3-none-any.whl

naas_abi_core/services/agent/Agent.py

@@ -0,0 +1,1619 @@
+from __future__ import annotations
+
+# Standard library imports for type hints
+import os
+import re
+import uuid
+
+# Dataclass imports for configuration
+from dataclasses import dataclass, field
+from enum import Enum
+from queue import Empty, Queue
+from typing import (
+    TYPE_CHECKING,
+    Annotated,
+    Any,
+    Callable,
+    Dict,
+    Generator,
+    Literal,
+    Optional,
+    Sequence,
+    Union,
+    cast,
+)
+
+import pydash as pd
+from langchain_core.language_models import BaseChatModel
+from langchain_core.tools import BaseTool, StructuredTool, Tool, tool
+from langgraph.prebuilt import InjectedState
+from naas_abi_core.models.Model import ChatModel
+from naas_abi_core.utils.Expose import Expose
+from naas_abi_core.utils.Logger import logger
+
+# Pydantic imports for schema validation (keep - it's already loaded by other modules)
+from pydantic import BaseModel, Field
+
+###
+
+
+# Only import heavy modules for type checking
+if TYPE_CHECKING:
+    from fastapi import APIRouter
+    from langchain_core.runnables import Runnable
+    from langgraph.checkpoint.base import BaseCheckpointSaver
+    from langgraph.graph.state import CompiledStateGraph
+
+from langchain_core.messages import (
+    AIMessage,
+    AnyMessage,
+    BaseMessage,
+    HumanMessage,
+    SystemMessage,
+    ToolCall,
+    ToolMessage,
+)
+from langgraph.checkpoint.memory import MemorySaver
+from langgraph.graph import START, StateGraph
+from langgraph.graph.message import MessagesState
+from langgraph.types import Command
+from sse_starlette.sse import EventSourceResponse
+
+
+def create_checkpointer() -> BaseCheckpointSaver:
+    """Create a checkpointer based on environment configuration.
+
+    Returns a PostgreSQL-backed checkpointer if POSTGRES_URL is set,
+    otherwise returns an in-memory checkpointer.
+    """
+    postgres_url = os.getenv("POSTGRES_URL")
+
+    if postgres_url:
+        try:
+            import time
+
+            from langgraph.checkpoint.postgres import PostgresSaver
+            from psycopg import Connection
+            from psycopg.rows import dict_row
+
+            logger.debug(
+                f"Using PostgreSQL checkpointer for persistent memory: {postgres_url}"
+            )
+
+            # Try connection with retries (PostgreSQL might still be starting)
+            max_retries = 3
+            for attempt in range(max_retries):
+                try:
+                    # Create connection with proper configuration (matching from_conn_string)
+                    conn = Connection.connect(
+                        postgres_url,
+                        autocommit=True,
+                        prepare_threshold=0,
+                        row_factory=dict_row,
+                    )
+                    checkpointer = PostgresSaver(conn)
+
+                    # Setup tables if they don't exist
+                    checkpointer.setup()
+                    logger.debug("PostgreSQL checkpointer tables initialized")
+
+                    return checkpointer
+
+                except Exception as conn_error:
+                    if attempt < max_retries - 1:
+                        logger.warning(
+                            f"PostgreSQL connection attempt {attempt + 1} failed, retrying in 2 seconds..."
+                        )
+                        time.sleep(2)
+                    else:
+                        raise conn_error
+
+        except ImportError:
+            logger.error(
+                "PostgreSQL checkpointer requested but langgraph.checkpoint.postgres not available. Falling back to in-memory."
+            )
+        except Exception as e:
+            # Provide more helpful error messages
+            error_msg = str(e)
+            if "nodename nor servname provided" in error_msg:
+                logger.error(
+                    f"PostgreSQL connection failed - cannot resolve hostname. Check if PostgreSQL is running and hostname is correct in POSTGRES_URL: {postgres_url}"
+                )
+                logger.error(
+                    "Hint: If running outside Docker, use 'localhost' instead of 'postgres' in POSTGRES_URL"
+                )
+            else:
+                logger.error(f"Failed to initialize PostgreSQL checkpointer: {e}")
+            logger.error("Falling back to in-memory checkpointer")
+
+        # Fallback to in-memory checkpointer
+        return MemorySaver()
+    else:
+        logger.debug(
+            "Using in-memory checkpointer (set POSTGRES_URL for persistent memory)"
+        )
+        return MemorySaver()
+
+
+class AgentSharedState:
+    _thread_id: str
+    _current_active_agent: Optional[str]
+    _supervisor_agent: Optional[str]
+    _requesting_help: bool
+
+    def __init__(
+        self,
+        thread_id: str = "1",
+        current_active_agent: Optional[str] = None,
+        supervisor_agent: Optional[str] = None,
+    ):
+        assert isinstance(thread_id, str)
+
+        self._thread_id = thread_id
+        self._current_active_agent = current_active_agent
+        self._supervisor_agent = supervisor_agent
+        self._requesting_help = False
+
+    @property
+    def thread_id(self) -> str:
+        return self._thread_id
+
+    def set_thread_id(self, thread_id: str):
+        self._thread_id = thread_id
+
+    @property
+    def current_active_agent(self) -> Optional[str]:
+        return self._current_active_agent
+
+    def set_current_active_agent(self, agent_name: Optional[str]):
+        self._current_active_agent = agent_name
+
+    @property
+    def supervisor_agent(self) -> Optional[str]:
+        return self._supervisor_agent
+
+    def set_supervisor_agent(self, agent_name: Optional[str]):
+        self._supervisor_agent = agent_name
+
+    @property
+    def requesting_help(self) -> bool:
+        return self._requesting_help
+
+    def set_requesting_help(self, requesting_help: bool):
+        self._requesting_help = requesting_help
+
+
+@dataclass
+class Event:
+    payload: Any
+
+
+@dataclass
+class ToolUsageEvent(Event):
+    pass
+
+
+@dataclass
+class ToolResponseEvent(Event):
+    pass
+
+
+@dataclass
+class AIMessageEvent(Event):
+    agent_name: str
+
+
+@dataclass
+class FinalStateEvent(Event):
+    pass
+
+
+@dataclass
+class AgentConfiguration:
+    on_tool_usage: Callable[[AnyMessage], None] = field(
+        default_factory=lambda: lambda _: None
+    )
+    on_tool_response: Callable[[AnyMessage], None] = field(
+        default_factory=lambda: lambda _: None
+    )
+    on_ai_message: Callable[[AnyMessage, str], None] = field(
+        default_factory=lambda: lambda _, __: None
+    )
+    system_prompt: str = field(
+        default="You are a helpful assistant. If a tool you used did not return the result you wanted, look for another tool that might be able to help you. If you don't find a suitable tool. Just output 'I DONT KNOW'"
+    )
+
+
+class CompletionQuery(BaseModel):
+    prompt: str = Field(..., description="The prompt to send to the agent")
+    thread_id: str | int = Field(
+        ..., description="The thread ID to use for the conversation"
+    )
+
+
+class Agent(Expose):
+    """An Agent class that orchestrates interactions between a language model and tools.
+
+    Performance Features:
+        • Lazy Initialization: Efficient resource utilization through lazy loading
+        • Connection Pooling: Optimized database connections for memory backends
+        • Parallel Execution: Concurrent tool execution where dependencies allow
+        • Caching: Intelligent caching of tool results and model responses
+        • Resource Management: Proper cleanup and resource management
+
+    Attributes:
+        _name (str): Unique identifier for the agent
+        _description (str): Human-readable description of the agent's purpose
+        _system_prompt (str): System prompt that defines the agent's behavior
+        _chat_model (BaseChatModel): The underlying language model with tool binding
+        _chat_model_with_tools (Runnable): Language model configured with available tools
+        _tools (list[Union[Tool, Agent]]): Original list of provided tools and agents
+        _structured_tools (list[Union[Tool, BaseTool]]): Processed and validated tools
+        _tools_by_name (dict[str, Union[Tool, BaseTool]]): Tool lookup dictionary
+        _native_tools (list[dict]): Native tools compatible with the language model
+        _agents (list[Agent]): List of sub-agents for delegation
+        _checkpointer (BaseCheckpointSaver): Memory backend for conversation persistence
+        _state (AgentSharedState): Shared state management for conversation threads
+        graph (CompiledStateGraph): Compiled workflow graph for conversation execution
+        _configuration (AgentConfiguration): Configuration settings for agent behavior
+        _event_queue (Queue): Event queue for real-time event streaming
+        _chat_model_output_version (str|None): Version identifier for model output format
+    """
+
+    _name: str
+    _description: str
+    _system_prompt: str
+
+    _chat_model: BaseChatModel
+    _chat_model_with_tools: Runnable[
+        Any
+        | str
+        | Sequence[BaseMessage | list[str] | tuple[str, str] | str | dict[str, Any]],
+        BaseMessage,
+    ]
+    _tools: list[Union[Tool, BaseTool, "Agent"]]
+    _original_tools: list[Union[Tool, BaseTool, "Agent"]]
+    _tools_by_name: dict[str, Union[Tool, BaseTool]]
+    _native_tools: list[dict]
+
+    # An agent can have other agents.
+    # He will be responsible to load them as tools.
+    _agents: list["Agent"] = []
+
+    _chekpointer: BaseCheckpointSaver
+    _state: AgentSharedState
+
+    graph: CompiledStateGraph
+    _workflow: StateGraph
+    _configuration: AgentConfiguration
+
+    _on_tool_usage: Callable[[AnyMessage], None]
+    _on_tool_response: Callable[[AnyMessage], None]
+    _on_ai_message: Callable[[AnyMessage, str], None]
+
+    # Avent queue used to stream tool usage and responses.
+    _event_queue: Queue
+
+    _chat_model_output_version: Union[str, None] = None
+
+    @classmethod
+    def New(
+        cls,
+        agent_shared_state: Optional[AgentSharedState] = None,
+        agent_configuration: Optional[AgentConfiguration] = None,
+    ) -> "Agent":
+        """Create a new instance of the agent.
+
+        Args:
+            agent_shared_state: Optional[AgentSharedState]: The shared state of the agent.
+            agent_configuration: Optional[AgentConfiguration]: The configuration of the agent.
+
+        Returns:
+        raise NotImplementedError("This method is not implemented")
+        """
+        raise NotImplementedError("This method is not implemented")
+
+    def __init__(
+        self,
+        name: str,
+        description: str,
+        chat_model: BaseChatModel | ChatModel,
+        tools: list[Union[Tool, BaseTool, "Agent"]] = [],
+        agents: list["Agent"] = [],
+        memory: BaseCheckpointSaver | None = None,
+        state: AgentSharedState = AgentSharedState(),
+        configuration: AgentConfiguration = AgentConfiguration(),
+        event_queue: Queue | None = None,
+        native_tools: list[dict] = [],
+    ):
+        """Initialize a new Agent instance.
+
+        Args:
+            chat_model (BaseChatModel): The language model to use for chat interactions.
+                Should support tool binding.
+            tools (list[Tool]): List of tools to make available to the agent.
+            memory (BaseCheckpointSaver, optional): Component to save conversation state.
+                If None, will use PostgreSQL if POSTGRES_URL env var is set, otherwise in-memory.
+        """
+        logger.debug(f"Initializing agent: {name}")
+        self._name = name
+        self._description = description
+        self._system_prompt = configuration.system_prompt
+        self._state = state
+        self._original_tools = tools
+        self._original_agents = agents
+
+        # We set the supervisor agent and current active agent before the default tools are injected.
+        if self._state.supervisor_agent is not None:
+            self._state.set_supervisor_agent(self._state.supervisor_agent)
+        logger.debug(f"Supervisor agent: {self._state.supervisor_agent}")
+
+        agent_names = [a.name for a in self._original_agents] + [name]
+        if (
+            self._state.current_active_agent is not None
+            and self._state.current_active_agent in agent_names
+        ):
+            self._state.set_current_active_agent(self._state.current_active_agent)
+        logger.debug(f"Current active agent: {self._state.current_active_agent}")
+
+        # We inject default tools
+        tools += self.default_tools()
+
+        # We store the original list of provided tools. This will be usefull for duplication.
+        self._tools = tools
+        self._native_tools = native_tools
+
+        # Assertions
+        assert isinstance(name, str)
+        assert isinstance(description, str)
+        assert isinstance(chat_model, BaseChatModel | ChatModel)
+
+        # We assert agents
+        for agent in agents:
+            assert isinstance(agent, Agent)
+
+        # We store the provided tools in __structured_tools because we will need to know which ones are provided by the user and which one are agents.
+        # This is needed when we duplicate the agent.
+        _structured_tools, _agents = self.prepare_tools(
+            cast(list[Union[Tool, BaseTool, "Agent"]], tools), agents
+        )
+        self._structured_tools = _structured_tools
+        self._agents = _agents
+
+        # We assert that the tool that are provided are valid.
+        for t in self._structured_tools:
+            assert isinstance(t, StructuredTool)
+            assert hasattr(t, "name")
+            assert hasattr(t, "description")
+            assert hasattr(t, "func")
+            assert hasattr(t, "args_schema")
+
+        self._tools_by_name: dict[str, Union[Tool, BaseTool]] = {
+            tool.name: tool for tool in self._structured_tools
+        }
+
+        base_chat_model: BaseChatModel = (
+            chat_model if isinstance(chat_model, BaseChatModel) else chat_model.model
+        )
+        assert isinstance(base_chat_model, BaseChatModel)
+
+        self._chat_model = base_chat_model
+        if hasattr(base_chat_model, "output_version"):
+            self._chat_model_output_version = base_chat_model.output_version
+
+        self._chat_model_with_tools = base_chat_model
+        if self._tools or self._native_tools:
+            tools_to_bind: list[Union[Tool, BaseTool, Dict]] = []
+            tools_to_bind.extend(self._structured_tools)
+            tools_to_bind.extend(self._native_tools)
+
+            # Test if the chat model can bind tools by trying with a default tool first
+            if self._can_bind_tools(base_chat_model):
+                self._chat_model_with_tools = base_chat_model.bind_tools(tools_to_bind)
+            else:
+                logger.warning(
+                    f"Chat model {type(base_chat_model).__name__} does not support tool calling. Tools will not be available for agent '{self._name}'."
+                )
+                # Keep the original model without tools
+                self._chat_model_with_tools = base_chat_model
+
+        # Use provided memory or create based on environment
+        if memory is None:
+            self._checkpointer = create_checkpointer()
+        else:
+            self._checkpointer = memory
+
+        # Randomize the thread_id to prevent the same thread_id to be used by multiple agents.
+        if os.getenv("ENV") == "dev":
+            self._state.set_thread_id(str(uuid.uuid4()))
+
+        # We set the configuration.
+        self._configuration = configuration
+        self._on_tool_usage = configuration.on_tool_usage
+        self._on_tool_response = configuration.on_tool_response
+        self._on_ai_message = configuration.on_ai_message
+
+        # Initialize the event queue.
+        if event_queue is None:
+            self._event_queue = Queue()
+        else:
+            self._event_queue = event_queue
+
+        # We build the graph.
+        self.build_graph()
+
+    def _can_bind_tools(self, chat_model: BaseChatModel) -> bool:
+        """Test if the chat model can bind tools by attempting to bind the get_time_date default tool.
+
+        Args:
+            chat_model (BaseChatModel): The chat model to test
+
+        Returns:
+            bool: True if the model can bind tools, False otherwise
+        """
+        try:
+            # Create the get_time_date tool that's used in default_tools()
+            @tool(return_direct=True)
+            def get_time_date(timezone: str = "Europe/Paris") -> str:
+                """Get the current time and date."""
+                from datetime import datetime
+                from zoneinfo import ZoneInfo
+
+                return datetime.now(ZoneInfo(timezone)).strftime("%H:%M:%S %Y-%m-%d")
+
+            # Try to bind this single tool to test if the model supports tool binding
+            chat_model.bind_tools([get_time_date])
+
+            # If we get here without an exception, the model supports tool binding
+            # logger.debug(f"Chat model {type(chat_model).__name__} supports tool calling.")
+            return True
+
+        except Exception as e:
+            # If binding tools raises an exception, the model doesn't support tools
+            logger.debug(
+                f"Chat model {type(chat_model).__name__} does not support tool calling: {e}"
+            )
+            return False
+
+    def default_tools(self) -> list[Tool | BaseTool]:
+        @tool(return_direct=True)
+        def request_help(reason: str):
+            """
+            Request help from the supervisor agent when you (the LLM) are uncertain about the next step or do not have the required capability to fulfill the user's request.
+
+            Use this tool if:
+            - You are unsure how to proceed.
+            - You lack the necessary knowledge or ability to complete the task.
+            - The user's request is outside your capabilities or unclear.
+
+            Args:
+                reason (str): A brief explanation of why you are requesting help (e.g., "I am uncertain about the next step", "I do not have the required capability", "The user's request is unclear").
+
+            The supervisor agent will review your reason and provide assistance or take over the conversation.
+            """
+            logger.debug(f"'{self.name}' is requesting help from the supervisor agent")
+            return "Requesting help from the supervisor agent."
+
+        @tool(return_direct=False)
+        def get_time_date(timezone: str = "Europe/Paris") -> str:
+            """Returns the current date and time for a given timezone."""
+            from datetime import datetime
+            from zoneinfo import ZoneInfo
+
+            return datetime.now(ZoneInfo(timezone)).strftime("%H:%M:%S %Y-%m-%d")
+
+        @tool(return_direct=True)
+        def get_current_active_agent() -> str:
+            """Returns the current active agent."""
+            return "The current active agent is: " + (
+                self._state.current_active_agent or self.name
+            )
+
+        @tool(return_direct=True)
+        def get_supervisor_agent() -> str:
+            """Returns the supervisor agent."""
+            if self._state.supervisor_agent is None:
+                return "I don't have a supervisor agent."
+            return "The supervisor agent is: " + self._state.supervisor_agent
+
+        @tool(return_direct=True)
+        def list_tools_available() -> str:
+            """Displays a formatted list of all available tools."""
+            if (
+                not hasattr(self, "_structured_tools")
+                or len(self._structured_tools) == 0
+            ):
+                return "I don't have any tools available to help you at the moment."
+
+            tools_text = "Here are the tools I can use to help you:\n\n"
+            for t in self._structured_tools:
+                if not t.name.startswith("transfer_to"):
+                    tools_text += f"- `{t.name}`: {t.description.splitlines()[0]}\n"
+            return tools_text.rstrip()
+
+        @tool(return_direct=True)
+        def list_subagents_available() -> str:
+            """Displays a formatted list of all available sub-agents."""
+            if not hasattr(self, "_agents") or len(self._agents) == 0:
+                return "I don't have any sub-agents that can assist me at the moment."
+
+            agents_text = "I can collaborate with these sub-agents:\n"
+            for agent in self._agents:
+                agents_text += f"- `{agent.name}`: {agent.description}\n"
+            return agents_text.rstrip()
+
+        @tool(return_direct=True)
+        def list_intents_available() -> str:
+            """Displays a formatted list of all available intents."""
+            if not hasattr(self, "_intents") or len(self._intents) == 0:
+                return "I haven't been configured with any specific intents yet."
+
+            from naas_abi_core.services.agent.IntentAgent import (
+                Intent,
+                IntentScope,
+                IntentType,
+            )
+
+            # Group intents by scope and type
+            intents_by_scope: dict[
+                Optional[IntentScope], dict[IntentType, list[Intent]]
+            ] = {}
+            for intent in self._intents:
+                if intent.intent_scope not in intents_by_scope:
+                    intents_by_scope[intent.intent_scope] = {}
+                if intent.intent_type not in intents_by_scope[intent.intent_scope]:
+                    intents_by_scope[intent.intent_scope][intent.intent_type] = []
+                intents_by_scope[intent.intent_scope][intent.intent_type].append(intent)
+
+            intents_text = "Here are all the intents I'm configured with:\n\n"
+            for scope, types_dict in intents_by_scope.items():
+                intents_text += f"### Intents for {str(scope)}\n\n"
+                for intent_type, intents in types_dict.items():
+                    intents_text += f"#### {str(intent_type)}\n\n"
+                    intents_text += "| Intent | Target |\n"
+                    intents_text += "|--------|--------|\n"
+                    for intent in intents:
+                        if intent.intent_scope == IntentType.RAW:
+                            intents_text += (
+                                f"| {intent.intent_value} | {intent.intent_target} |\n"
+                            )
+                        else:
+                            intents_text += f"| {intent.intent_value} | `{intent.intent_target}` |\n"
+                    intents_text += "\n"
+            return intents_text.rstrip()
+
+        @tool(return_direct=False)
+        def read_makefile() -> str:
+            """Read the Makefile and return the content."""
+            try:
+                with open("Makefile", "r") as f:
+                    makefile_content = f.read()
+
+                return "Here are the make commands available:\n\n" + makefile_content
+
+            except FileNotFoundError:
+                return "Could not find Makefile in the root directory."
+            except Exception as e:
+                return f"Error reading Makefile: {str(e)}"
+
+        tools: list[Tool | BaseTool] = [
+            get_time_date,
+            list_tools_available,
+            list_subagents_available,
+            list_intents_available,
+        ]
+        if self.state.supervisor_agent and self.state.supervisor_agent != self.name:
+            tools.append(request_help)
+
+        if self.state.supervisor_agent is not None or len(self._agents) > 0:
+            tools.append(get_current_active_agent)
+            tools.append(get_supervisor_agent)
+
+        if self.state.supervisor_agent == self.name and os.getenv("ENV") == "dev":
+            tools.append(read_makefile)
+        return tools
+
+    @property
+    def system_prompt(self) -> str:
+        return self._system_prompt
+
+    def set_system_prompt(self, system_prompt: str):
+        self._system_prompt = system_prompt
+
+    @property
+    def structured_tools(self) -> list[Tool | BaseTool]:
+        return self._structured_tools
+
+    @property
+    def state(self) -> AgentSharedState:
+        return self._state
+
+    def validate_tool_name(self, tool: BaseTool) -> BaseTool:
+        if not re.match(r"^[a-zA-Z0-9_-]+$", tool.name):
+            # Replace invalid characters with '_'
+            valid_name = re.sub(r"[^a-zA-Z0-9_-]", "_", tool.name)
+            logger.warning(
+                f"Tool name '{tool.name}' does not comply with '^[a-zA-Z0-9_-]+$'. Renaming to '{valid_name}'."
+            )
+            tool.name = valid_name
+        return tool
+
+    def prepare_tools(
+        self, tools: list[Union[Tool, BaseTool, "Agent"]], agents: list
+    ) -> tuple[list[Tool | BaseTool], list["Agent"]]:
+        """
+        If we have Agents in tools, we are properly loading them as handoff tools.
+        It will effectively make the 'self' agent a supervisor agent.
+
+        Ensures no duplicate tools or agents are added by tracking unique names/instances.
+        """
+        _tools: list[Tool | BaseTool] = []
+        _agents: list["Agent"] = []
+        _tool_names: set[str] = set()
+        _agent_names: set[str] = set()
+
+        # We process tools knowing that they can either be StructutedTools or Agent.
+        for t in tools:
+            if isinstance(t, Agent):
+                # TODO: We might want to duplicate the agent first.
+                # logger.debug(f"Agent passed as tool: {t}")
+                if t.name not in _agent_names:
+                    _agents.append(t)
+                    _agent_names.add(t.name)
+                    for tool in t.as_tools():
+                        tool = self.validate_tool_name(tool)
+                        if tool.name not in _tool_names:
+                            _tools.append(tool)
+                            _tool_names.add(tool.name)
+            else:
+                # Accept both Tool and BaseTool
+                if hasattr(t, "name"):
+                    if t.name not in _tool_names:
+                        _tools.append(t)
+                        _tool_names.add(t.name)
+
+        # We process agents that are not provided in tools.
+        for agent in agents:
+            if agent.name not in _agent_names:
+                _agents.append(agent)
+                _agent_names.add(agent.name)
+                for tool in agent.as_tools():
+                    tool = self.validate_tool_name(tool)
+                    if tool.name not in _tool_names:
+                        _tools.append(tool)
+                        _tool_names.add(tool.name)
+
+        return _tools, _agents
+
+    def as_tools(self, parent_graph: bool = False) -> list[BaseTool]:
+        return [make_handoff_tool(agent=self, parent_graph=parent_graph)]
+
+    def build_graph(self, patcher: Optional[Callable] = None):
+        graph = StateGraph(MessagesState)
+
+        graph.add_node(self.current_active_agent)
+        graph.add_edge(START, "current_active_agent")
+
+        graph.add_node(self.continue_conversation)
+
+        graph.add_node(self.call_model)
+
+        graph.add_node(self.call_tools)
+
+        for agent in self._agents:
+            graph.add_node(agent.name, agent.graph)
+
+        # Patcher is callable that can be passed and that will impact the graph before we compile it.
+        # This is used to be able to give more flexibility about how the graph is being built.
+        if patcher:
+            graph = patcher(graph)
+
+        self.graph = graph.compile(checkpointer=self._checkpointer)
+
+    def get_last_human_message(self, state: MessagesState) -> Any | None:
+        """Get the appropriate human message based on AI message context.
+
+        Args:
+            state (MessagesState): Current conversation state
+
+        Returns:
+            Any | None: The relevant human message
+        """
+        last_ai_message: Any | None = pd.find(
+            state["messages"][::-1], lambda m: isinstance(m, AIMessage)
+        )
+        if (
+            last_ai_message is not None
+            and last_ai_message.additional_kwargs.get("owner") == self.name
+        ):
+            return pd.find(
+                state["messages"][::-1], lambda m: isinstance(m, HumanMessage)
+            )
+        elif (
+            last_ai_message is not None
+            and hasattr(last_ai_message, "additional_kwargs")
+            and last_ai_message.additional_kwargs is not None
+            and "owner" in last_ai_message.additional_kwargs
+        ):
+            return pd.filter_(
+                state["messages"][::-1], lambda m: isinstance(m, HumanMessage)
+            )[1]
+        else:
+            return pd.find(
+                state["messages"][::-1], lambda m: isinstance(m, HumanMessage)
+            )
+
+    def current_active_agent(self, state: MessagesState) -> Command:
+        """Goto the current active agent.
+
+        Args:
+            state (MessagesState): Current conversation state
+
+        Returns:
+            Command: Command to goto the current active agent
+        """
+        # Log the current active agent
+        logger.debug(f"😏 Supervisor agent: '{self.state.supervisor_agent}'")
+        logger.debug(f"🟢 Active agent: '{self.state.current_active_agent}'")
+        logger.debug(f"🤖 Current Agent: '{self.name}'")
+
+        # Get the last human message
+        last_human_message = self.get_last_human_message(state)
+
+        # Handle agent routing via @mention
+        if (
+            last_human_message is not None
+            and isinstance(last_human_message.content, str)
+            and last_human_message.content.startswith("@")
+            and last_human_message.content.split(" ")[0].split("@")[1] != self.name
+        ):
+            at_mention = last_human_message.content.split(" ")[0].split("@")[1]
+            logger.debug(f"🔀 Handle agent routing via @mention to '{at_mention}'")
+
+            # Check if we have an agent with this name.
+            agent = pd.find(
+                self._agents, lambda a: a.name.lower() == at_mention.lower()
+            )
+
+            # Remove mention from the last human message with re.sub
+            import re
+
+            last_human_message.content = re.sub(
+                r"^@[^ ]* ", "", last_human_message.content
+            )
+
+            if agent is not None:
+                self._state.set_current_active_agent(agent.name)
+                return Command(goto=agent.name, update={"messages": state["messages"]})
+            else:
+                logger.debug(f"❌ Agent '{at_mention}' not found")
+
+        if (
+            self._state.current_active_agent is not None
+            and self._state.current_active_agent != self.name
+        ):
+            logger.debug(
+                f"⏩ Continuing conversation with agent '{self._state.current_active_agent}'"
+            )
+            # Check if current active agent is in list of agents.
+            if self._state.current_active_agent in [a.name for a in self._agents]:
+                self._state.set_current_active_agent(self._state.current_active_agent)
+                return Command(goto=self._state.current_active_agent)
+            else:
+                logger.debug(f"❌ Agent '{self._state.current_active_agent}' not found")
+
+        # self._state.set_current_active_agent(self.name)
+        logger.debug(f"💬 Starting chatting with agent '{self.name}'")
+        if (
+            self.state.supervisor_agent != self.name
+            and "SUPERVISOR SYSTEM PROMPT" not in self._system_prompt
+        ):
+            # This agent is a subagent with a supervisor
+            subagent_prompt = f"""
+SUPERVISOR SYSTEM PROMPT:
+
+Remember, you are a specialized agent working under the supervision of {self.state.supervisor_agent}.
+
+1. Stay focused on your specialized role and core capabilities.
+2. Follow your system prompt instructions precisely.
+3. For EVERY user message, first evaluate if you can handle it within your core capabilities.
+4. If the user message is not clear, ask for clarification and evaluate again if you can handle it within your core capabilities.
+5. If you encounter ANY of these situations:
+   - You are uncertain about how to proceed
+   - The task seems outside your core capabilities 
+   - You need clarification about requirements
+   - You want to confirm a critical action
+   - You are not 100% confident in your ability to handle the task
+   Then you MUST use the `request_help` tool to ask your supervisor for help.
+   IMPORTANT: Do NOT explain to the user that you cannot perform the task or say it is outside your scope. Instead, IMMEDIATELY and directly use `request_help`.
+6. Do not attempt tasks beyond your defined role.
+7. Always maintain consistency with your system prompt rules.
+8. When in doubt, ALWAYS request help rather than risk mistakes. Do not type or explain your inability—just use the tool.
+
+Your supervisor will help ensure you operate effectively within your role while providing guidance for complex scenarios.
+
+--------------------------------
+
+SUBAGENT SYSTEM PROMPT:
+
+{self._system_prompt}
+"""
+            self.set_system_prompt(subagent_prompt)
+
+        if (
+            self.state.supervisor_agent == self.name
+            and os.getenv("ENV") == "dev"
+            and "DEVELOPPER SYSTEM PROMPT" not in self._system_prompt
+        ):
+            dev_prompt = f"""
+DEVELOPPER SYSTEM PROMPT:
+
+For any questions/commands related to the project, use tool: `read_makefile` to get the information you need.
+
+--------------------------------
+
+AGENT SYSTEM PROMPT:
+
+{self._system_prompt}
+"""
+
+            self.set_system_prompt(dev_prompt)
+
+        if "CURRENT_DATE" not in self._system_prompt:
+            from datetime import datetime
+
+            current_date_str = f"CURRENT_DATE: The current date is {datetime.now().strftime('%Y-%m-%d')}\n"
+            self._system_prompt = self._system_prompt + "\n" + current_date_str
+            self.set_system_prompt(self._system_prompt)
+            return Command(goto="current_active_agent")
+
+        # logger.debug(f"💬 System prompt: {self._system_prompt}")
+        return Command(goto="continue_conversation")
+
+    def continue_conversation(self, state: MessagesState) -> Command:
+        return Command(goto="call_model")
+
+    def handle_openai_response_v1(self, response: BaseMessage) -> Command:
+        content_str: str = ""
+        tool_call: list[ToolCall] = []
+        logger.debug(f"Chat model output version is responses/v1: {response}")
+
+        if isinstance(response.content, list):
+            # Parse response content
+            for item in response.content:
+                # Ensure item is a dict before accessing attributes
+                if isinstance(item, dict):
+                    # Get text content
+                    if item.get("type") == "text":
+                        text_content = item.get("text", "")
+                        if isinstance(text_content, str):
+                            content_str += text_content
+
+                        # Add sources from annotations if any
+                        annotations = item.get("annotations", [])
+                        if isinstance(annotations, list) and len(annotations) > 0:
+                            content_str += "\n\n\n\n*Annotations:*\n"
+                            for annotation in annotations:
+                                if (
+                                    isinstance(annotation, dict)
+                                    and annotation.get("type") == "url_citation"
+                                ):
+                                    title = annotation.get("title", "")
+                                    url = annotation.get("url", "")
+                                    content_str += f"- [{title}]({url})\n"
+
+                    if "action" in item:
+                        tool_call.append(
+                            ToolCall(
+                                name=item["type"],
+                                args={"query": item["action"].get("query", "")},
+                                id=item.get("id"),
+                                type="tool_call",
+                            )
+                        )
+
+        # Create AIMessage with the content
+        usage_metadata = None
+        if hasattr(response, "usage_metadata"):
+            usage_metadata = response.usage_metadata
+        ai_message = AIMessage(content=content_str, usage_metadata=usage_metadata)
+
+        # If action was detected, notify tool usage
+        if len(tool_call) > 0:
+            # Use the ai_message which is already the correct type
+            ai_message.tool_calls = tool_call
+            self._notify_tool_usage(ai_message)
+            tool_message = ToolMessage(
+                content=content_str, tool_call_id=tool_call[0].get("id")
+            )
+            self._notify_tool_response(tool_message)
+            return Command(
+                goto="__end__", update={"messages": [tool_message, ai_message]}
+            )
+
+        return Command(goto="__end__", update={"messages": [ai_message]})
+
+    def call_model(
+        self,
+        state: MessagesState,
+    ) -> Command[Literal["call_tools", "__end__"]]:
+        self._state.set_current_active_agent(self.name)
+        logger.debug(f"🧠 Calling model on current active agent: {self.name}")
+
+        # Inserting system prompt before messages.
+        messages = state["messages"]
+        if self._system_prompt:
+            messages = [
+                SystemMessage(content=self._system_prompt),
+            ] + messages
+        logger.debug(f"Messages before calling model: {messages}")
+
+        # Calling model
+        response: BaseMessage = self._chat_model_with_tools.invoke(messages)
+        logger.debug(
+            f"Model response content: {response.content if hasattr(response, 'content') else response}"
+        )
+
+        # Handle tool calls if present
+        if (
+            isinstance(response, AIMessage)
+            and hasattr(response, "tool_calls")
+            and len(response.tool_calls) > 0
+        ):
+            logger.debug("⏩ Calling tools")
+            # TODO: Rethink this.
+            # This is done to prevent an LLM to call multiple tools at once.
+            # It's important because, as some tools are subgraphs, and that we are passing the full state, the subgraph will be able to mess with the state.
+            # Therefore, if the LLM calls a tool here like the "add" method, and at the same time request the multiplication agent, the agent will mess with the state, and the result of the "add" tool will be lost.
+            #### ----->  A solution would be to rebuild the state to make sure that the following message of a tool call it the response of that call. If we do that we should theroetically be able to call multiple tools at once, which would be more effective.
+            # response.tool_calls = [response.tool_calls[0]]
+
+            return Command(goto="call_tools", update={"messages": [response]})
+
+        elif self._chat_model_output_version == "responses/v1":
+            return self.handle_openai_response_v1(response)
+
+        return Command(goto="__end__", update={"messages": [response]})
+
+    def call_tools(self, state: MessagesState) -> list[Command]:
+        # Check if messages are present in the state.
+        if (
+            "messages" not in state
+            or not isinstance(state["messages"], list)
+            or len(state["messages"]) == 0
+        ):
+            logger.warning("No messages in state, cannot call tools")
+            return [Command(goto="__end__")]
+
+        # Check if the last message is an AIMessage and has tool calls.
+        last_message: AnyMessage = state["messages"][-1]
+        if (
+            not isinstance(last_message, AIMessage)
+            or not hasattr(last_message, "tool_calls")
+            or len(last_message.tool_calls) == 0
+        ):
+            logger.warning(
+                f"No tool calls found in last message but call_tools was called: {last_message}"
+            )
+            return [Command(goto="__end__", update={"messages": [last_message]})]
+
+        # Get the tool calls from the last message.
+        tool_calls: list[ToolCall] = last_message.tool_calls
+        assert len(tool_calls) > 0, state["messages"][-1]
+
+        # Initialize the results list.
+        results: list[Command] = []
+
+        # Initialize the called tools list.
+        called_tools: list[BaseTool] = []
+        for tool_call in tool_calls:
+            tool_name: str = tool_call["name"]
+            logger.debug(f"🛠️  Calling tool: {tool_name}")
+            tool_: BaseTool = self._tools_by_name[tool_name]
+
+            tool_input_fields = tool_.get_input_schema().model_json_schema()[
+                "properties"
+            ]
+
+            # For tools with InjectedToolCallId, we must pass the full ToolCall object
+            # according to LangChain's requirements
+            args: dict[str, Any] | ToolCall = tool_call
+
+            # Check if tool needs state injection
+            if "state" in tool_input_fields:
+                args = {**tool_call, "state": state}  # inject state
+
+            # Check if tool is a handoff tool
+            is_handoff = tool_call["name"].startswith("transfer_to_")
+            if is_handoff is True:
+                args = {"state": state, "tool_call": {**tool_call, "role": "tool_call"}}
+
+            # Try to invoke the tool.
+            try:
+                logger.debug(f"🔧 Tool arguments: {args.get('args')}")
+                tool_response = tool_.invoke(args)
+                logger.debug(
+                    f"📦 Tool response: {tool_response.content if hasattr(tool_response, 'content') else tool_response}"
+                )
+                if (
+                    tool_response is not None
+                    and hasattr(tool_response, "name")
+                    and tool_response.name == "request_help"
+                    and self._state.supervisor_agent != self.name
+                ):
+                    self._state.set_current_active_agent(self._state.supervisor_agent)
+                    self._state.set_requesting_help(True)
+                    results.append(
+                        Command(goto="current_active_agent", graph=Command.PARENT)
+                    )
+                    return results
+
+                called_tools.append(tool_)
+
+                # Handle tool response.
+                if isinstance(tool_response, ToolMessage):
+                    results.append(Command(update={"messages": [tool_response]}))
+                elif isinstance(tool_response, Command):
+                    results.append(tool_response)
+                else:
+                    logger.warning(
+                        f"Tool call {tool_name} returned an unexpected type: {type(tool_response)}"
+                    )
+                    results.append(
+                        Command(
+                            goto="__end__",
+                            update={
+                                "messages": [
+                                    ToolMessage(
+                                        content=str(tool_response),
+                                        tool_call_id=tool_call["id"],
+                                    )
+                                ]
+                            },
+                        )
+                    )
+            except Exception as e:
+                logger.error(f"🚨 Tool call {tool_name} failed: {e}")
+                results.append(
+                    Command(
+                        goto="__end__",
+                        update={
+                            "messages": [
+                                ToolMessage(
+                                    content=f"Tool call {tool_name} failed: {str(e)}",
+                                    tool_call_id=tool_call["id"],
+                                )
+                            ]
+                        },
+                    )
+                )
+
+        assert len(results) > 0, state
+
+        # Checking if every called tools has return_direct set to True.
+        # This is used to know if we can send the ToolMessage to the call_model node.
+        return_direct: bool = True
+        for t in called_tools:
+            if hasattr(t, "return_direct") and t.return_direct is False:
+                return_direct = False
+                break
+
+        # If the last response is a ToolMessage, we want the model to interpret it.
+        last_tool_reponse: ToolMessage | Command | None = pd.get(
+            results[-1], "update.messages[-1]", None
+        )
+        logger.debug(f"last_tool_reponse: {last_tool_reponse}")
+        if (
+            isinstance(last_tool_reponse, ToolMessage)
+            and hasattr(last_tool_reponse, "name")
+            and last_tool_reponse.name is not None
+            and not last_tool_reponse.name.startswith("transfer_to_")
+        ):
+            if return_direct is False:
+                logger.debug("⏩ Calling model to interpret the tool response.")
+                results.append(Command(goto="call_model"))
+            else:
+                logger.debug(
+                    "📧 Injecting ToolMessage into AIMessage for the user to see."
+                )
+                results.append(
+                    Command(
+                        update={
+                            "messages": [AIMessage(content=last_tool_reponse.content)]
+                        }
+                    )
+                )
+
+        logger.debug(f"✅ Tool results: {results}")
+        return results
+
+    @property
+    def workflow(self) -> StateGraph:
+        return self._workflow
+
+    def _notify_tool_usage(self, message: AnyMessage):
+        self._event_queue.put(ToolUsageEvent(payload=message))
+        self._on_tool_usage(message)
+
+    def _notify_tool_response(self, message: AnyMessage):
+        self._event_queue.put(ToolResponseEvent(payload=message))
+        self._on_tool_response(message)
+
+    def _notify_ai_message(self, message: AnyMessage, agent_name: str):
+        self._event_queue.put(AIMessageEvent(payload=message, agent_name=agent_name))
+        self._on_ai_message(message, agent_name)
+
+    def on_tool_usage(self, callback: Callable[[AnyMessage], None]):
+        """Register a callback to be called when a tool is used.
+
+        The callback will be invoked whenever the model makes a tool call,
+        before the tool is actually executed.
+
+        Args:
+            callback (Callable[[AnyMessage], None]): Function to call with the message
+                containing the tool call
+        """
+        self._on_tool_usage = callback
+        # # Also set the callback on all sub-agents to ensure they notify properly
+        # for agent in self._agents:
+        #     agent.on_tool_usage(callback)
+
+    def on_tool_response(self, callback: Callable[[AnyMessage], None]):
+        """Register a callback to be called when a tool response is received.
+
+        The callback will be invoked whenever a tool response message is processed,
+        before passing the messages to the model.
+
+        Args:
+            callback (Callable[[AnyMessage], None]): Function to call with the message
+                containing the tool response
+        """
+        self._on_tool_response = callback
+        # # Also set the callback on all sub-agents to ensure they notify properly
+        # for agent in self._agents:
+        #     agent.on_tool_response(callback)
+
+    def on_ai_message(self, callback: Callable[[AnyMessage, str], None]):
+        """Register a callback to be called when an AI message is received."""
+        self._on_ai_message = callback
+        # Also set the callback on all sub-agents to ensure they notify properly
+        for agent in self._agents:
+            agent.on_ai_message(callback)
+
+    @property
+    def app(self):
+        """Get the underlying Langchain app.
+        This property exposes the underlying Langchain app for advanced usage scenarios.
+        Users can call app.invoke() directly with custom message sequences and configurations
+        if they need more control than the standard invoke() method provides.
+
+        Returns:
+            RunnableSequence: The Langchain runnable sequence that processes messages
+        """
+        return self._app
+
+    def stream(self, prompt: str) -> Generator[dict[str, Any] | Any, None, None]:
+        """Process a user prompt through the agent and yield responses as they come.
+
+        Args:
+            prompt (str): The user's text prompt to process
+
+        Yields:
+            str: The model's response text
+        """
+        notified = {}
+
+        for chunk in self.graph.stream(
+            {"messages": [HumanMessage(content=prompt)]},
+            config={"configurable": {"thread_id": self._state.thread_id}},
+            subgraphs=True,
+        ):
+            source, payload = chunk
+            agent_name = self._name if len(source) == 0 else source[0].split(":")[0]
+            if isinstance(payload, dict):
+                last_messages = []
+                v = list(payload.values())[0]
+
+                if v is None:
+                    continue
+
+                if isinstance(v, dict):
+                    if (
+                        "messages" in v
+                        and isinstance(v["messages"], list)
+                        and len(v["messages"]) > 0
+                    ):
+                        last_messages = [v["messages"][-1]]
+                    else:
+                        continue
+                elif isinstance(v, list):
+                    last_messages = []
+                    for e in v:
+                        if (
+                            isinstance(e, dict)
+                            and "messages" in e
+                            and isinstance(e["messages"], list)
+                            and len(e["messages"]) > 0
+                        ):
+                            last_messages.append(e["messages"][-1])
+                else:
+                    continue
+
+                for last_message in last_messages:
+                    if isinstance(last_message, AIMessage):
+                        if pd.get(last_message, "additional_kwargs.tool_calls"):
+                            # This is a tool call.
+                            self._notify_tool_usage(last_message)
+                        else:
+                            # This if is here to filter each source of AIMessage. Which means that it will notify ai message only if the methods:
+                            # - call_model
+                            # - call_tools
+                            # are called.
+                            # If you need another method to be able to return an AIMessage or a Command(..., update={"messages": [AIMessage(...)]}) we either need to add it to the list or have this specific method calling self._notify_ai_message directly.
+
+                            allowed_sources_of_ai_message = ["call_model", "call_tools"]
+                            if any(
+                                source in payload
+                                for source in allowed_sources_of_ai_message
+                            ):
+                                self._notify_ai_message(last_message, agent_name)
+
+                    elif isinstance(last_message, ToolMessage):
+                        if last_message.id not in notified:
+                            self._notify_tool_response(last_message)
+                            notified[last_message.id] = True
+                    else:
+                        if "tool_call_id" in last_message:
+                            if last_message["tool_call_id"] not in notified:
+                                self._notify_tool_response(last_message)
+                                notified[last_message["tool_call_id"]] = True
+            yield chunk
+
+    def invoke(self, prompt: str) -> str:
+        """Process a user prompt through the agent and return the response.
+
+        This method takes a text prompt from the user, processes it through the underlying
+        Langchain app, and returns the model's response. The prompt is wrapped in a
+        HumanMessage and processed in a new message sequence.
+
+        Args:
+            prompt (str): The user's text prompt to process
+
+
+            str: The model's response text
+        """
+
+        chunks: list[dict[str, Any]] = []
+        for chunk in self.stream(prompt):
+            if isinstance(chunk, tuple):
+                chunk = chunk[1]
+
+            assert isinstance(chunk, dict)
+
+            chunks.append(chunk)
+
+        if len(chunks) == 0:
+            return ""
+
+        last_chunk_values = list(chunks[-1].values())
+        if len(last_chunk_values) == 0:
+            return ""
+        value = last_chunk_values[0]
+        messages = []
+        if (
+            isinstance(value, dict)
+            and "messages" in value
+            and isinstance(value["messages"], list)
+        ):
+            messages = value["messages"]
+        elif isinstance(value, list) and len(value) > 0:
+            last_item = value[-1]
+            if (
+                isinstance(last_item, dict)
+                and "messages" in last_item
+                and isinstance(last_item["messages"], list)
+            ):
+                messages = last_item["messages"]
+
+        if len(messages) == 0:
+            return ""
+
+        last_message = messages[-1]
+        if hasattr(last_message, "content"):
+            content = last_message.content
+        else:
+            content = str(last_message) if last_message is not None else ""
+        # content = list(chunks[-1].values())[0]["messages"][-1].content
+
+        return content
+
+    def reset(self):
+        """Reset the agent's conversation state.
+
+        This method increments the internal thread ID counter, effectively starting a new
+        conversation thread. Any subsequent invocations will be processed as part of a
+        new conversation context.
+        """
+        try:
+            current_thread_id = int(self._state.thread_id)
+            self._state.set_thread_id(str(current_thread_id + 1))
+        except (ValueError, TypeError):
+            # If thread_id is not a valid integer, generate a new UUID
+            self._state.set_thread_id(str(uuid.uuid4()))
+
+    def duplicate(
+        self,
+        queue: Queue | None = None,
+        agent_shared_state: AgentSharedState | None = None,
+    ) -> "Agent":
+        """Create a new instance of the agent with the same configuration.
+
+        This method creates a deep copy of the agent with the same configuration
+        but with its own independent state. This is useful when you need to run
+        multiple instances of the same agent concurrently.
+
+        Returns:
+            Agent: A new Agent instance with the same configuration
+        """
+        shared_state = agent_shared_state or AgentSharedState()
+
+        if queue is None:
+            queue = Queue()
+
+        # We duplicated each agent and add them as tools.
+        # This will be recursively done for each sub agents.
+        agents: list[Agent] = [
+            agent.duplicate(queue, shared_state) for agent in self._original_agents
+        ]
+
+        new_agent = self.__class__(
+            name=self._name,
+            description=self._description,
+            chat_model=self._chat_model,
+            tools=self._original_tools,
+            agents=agents,
+            memory=self._checkpointer,
+            state=shared_state,  # Create new state instance
+            configuration=self._configuration,
+            event_queue=queue,
+        )
+
+        return new_agent
+
+    def as_api(
+        self,
+        router: APIRouter,
+        route_name: str = "",
+        name: str = "",
+        description: str = "",
+        description_stream: str = "",
+        tags: list[str | Enum] | None = None,
+    ) -> None:
+        """Adds API endpoints for this agent to the given router.
+
+        Args:
+            router (APIRouter): The router to add endpoints to
+            route_name (str): Optional prefix for route names. Defaults to ""
+            name (str): Optional name to add to the endpoints. Defaults to ""
+            description (str): Optional description to add to the endpoints. Defaults to ""
+            description_stream (str): Optional description to add to the stream endpoints. Defaults to ""
+            tags (list[str]): Optional list of tags to add to the endpoints. Defaults to None
+        """
+
+        route_name = route_name or self._name
+        name = name or self._name.capitalize().replace("_", " ")
+        description = description or self._description
+        description_stream = description_stream or self._description
+
+        @router.post(
+            f"/{route_name}/completion" if route_name else "/completion",
+            name=f"{name} completion",
+            description=description,
+            tags=tags,
+        )
+        def completion(query: CompletionQuery):
+            if isinstance(query.thread_id, int):
+                query.thread_id = str(query.thread_id)
+            logger.debug(
+                f"completion - current active agent: {self._state.current_active_agent}"
+            )
+            logger.debug(
+                f"completion - supervisor agent: {self._state.supervisor_agent}"
+            )
+
+            new_agent = self.duplicate(
+                queue=self._event_queue, agent_shared_state=self._state
+            )
+            new_agent.state.set_thread_id(query.thread_id)
+            return new_agent.invoke(query.prompt)
+
+        @router.post(
+            f"/{route_name}/stream-completion" if route_name else "/stream-completion",
+            name=f"{name} stream completion",
+            description=description_stream,
+            tags=tags,
+        )
+        async def stream_completion(query: CompletionQuery):
+            if isinstance(query.thread_id, int):
+                query.thread_id = str(query.thread_id)
+            logger.debug(
+                f"stream_completion - current active agent: {self._state.current_active_agent}"
+            )
+            logger.debug(
+                f"stream_completion - supervisor agent: {self._state.supervisor_agent}"
+            )
+
+            new_agent = self.duplicate(
+                queue=self._event_queue, agent_shared_state=self._state
+            )
+            new_agent.state.set_thread_id(query.thread_id)
+            return EventSourceResponse(
+                new_agent.stream_invoke(query.prompt),
+                media_type="text/event-stream; charset=utf-8",
+            )
+
+    def stream_invoke(self, prompt: str):
+        """Process a user prompt through the agent and yield responses as they come.
+
+        Args:
+            prompt (str): The user's text prompt to process
+
+        Yields:
+            dict: Event data formatted for SSE
+        """
+
+        # Start a thread to run the invoke and put results in queue
+        def run_invoke():
+            final_state = self.invoke(prompt)
+            self._event_queue.put(FinalStateEvent(payload=final_state))
+
+        from threading import Thread
+
+        thread = Thread(target=run_invoke)
+        thread.start()
+
+        final_state = None
+        while True:
+            try:
+                message = self._event_queue.get(timeout=0.05)
+                if isinstance(message, ToolUsageEvent):
+                    yield {
+                        "event": "tool_usage",
+                        "data": str(message.payload.tool_calls[0]["name"]),
+                    }
+                elif isinstance(message, ToolResponseEvent):
+                    yield {
+                        "event": "tool_response",
+                        "data": str(pd.get(message, "payload.content", "NULL")),
+                    }
+                elif isinstance(message, AIMessageEvent):
+                    yield {
+                        "event": "ai_message",
+                        "data": str(message.payload.content),
+                    }
+                elif isinstance(message, FinalStateEvent):
+                    final_state = message.payload
+                    break
+
+                if (
+                    not thread.is_alive()
+                    and self._event_queue.empty()
+                    and final_state is None
+                ):
+                    # We have a problem.
+                    raise Exception(
+                        "Agent thread has died and no final state event was received."
+                    )
+            except Empty:
+                pass
+
+        response = final_state
+        logger.debug(f"Response: {response}")
+
+        # Use a buffer to handle text chunks
+        buffer = ""
+        for char in response:
+            buffer += char
+            if char in ["\n", "\r"]:
+                # if buffer.strip():  # Only send non-empty lines
+                yield {"event": "message", "data": buffer.rstrip()}
+                buffer = ""
+
+        # Don't forget remaining text
+        if buffer.strip():
+            yield {"event": "message", "data": buffer}
+
+        yield {"event": "done", "data": "[DONE]"}
+
+    @property
+    def tools(self) -> list[Union[Tool, BaseTool]]:
+        """Get the list of tools available to the agent.
+
+        Returns:
+            list[Tool]: List of tools configured for this agent
+        """
+        return self._structured_tools
+
+    @property
+    def name(self) -> str:
+        """Get the name of the agent.
+
+        Returns:
+            str: The agent's name
+        """
+        return self._name
+
+    @property
+    def description(self) -> str:
+        """Get the description of the agent.
+
+        Returns:
+            str: The agent's description
+        """
+        return self._description
+
+    @property
+    def agents(self) -> list["Agent"]:
+        return self._agents
+
+    @property
+    def chat_model(self) -> BaseChatModel:
+        """Get the chat model used by the agent.
+
+        Returns:
+            BaseChatModel: The agent's chat model
+        """
+        if isinstance(self._chat_model, ChatModel):
+            return self._chat_model.model
+        return self._chat_model
+
+    @property
+    def configuration(self) -> AgentConfiguration:
+        """Get the configuration used by the agent.
+
+        Returns:
+            AgentConfiguration: The agent's configuration
+        """
+        return self._configuration
+
+    def hello(self) -> str:
+        return "Hello"
+
+
+def make_handoff_tool(*, agent: Agent, parent_graph: bool = False) -> BaseTool:
+    """Create a tool that can return handoff via a Command"""
+    tool_name = f"transfer_to_{agent.name}"
+
+    @tool(tool_name)
+    def handoff_to_agent(
+        # # optionally pass current graph state to the tool (will be ignored by the LLM)
+        state: Annotated[dict, InjectedState],
+        # optionally pass the current tool call ID (will be ignored by the LLM)
+        tool_call: Annotated[ToolCall, ToolCall],
+    ):
+        """Ask another agent for help."""
+        agent_label = " ".join(
+            word.capitalize() for word in agent.name.replace("_", " ").split()
+        )
+
+        tool_message = ToolMessage(
+            content=f"Conversation transferred to {agent_label}",
+            name=tool_name,
+            tool_call_id=tool_call["id"],
+        )
+
+        agent.state.set_current_active_agent(agent.name)
+
+        return Command(
+            # navigate to another agent node in the PARENT graph
+            goto=agent.name,
+            graph=Command.PARENT if parent_graph else None,
+            # This is the state update that the agent `agent_name` will see when it is invoked.
+            # We're passing agent's FULL internal message history AND adding a tool message to make sure
+            # the resulting chat history is valid. See the paragraph above for more information.
+            update={"messages": state["messages"] + [tool_message]},
+        )
+
+    assert isinstance(handoff_to_agent, BaseTool)
+
+    return handoff_to_agent