import asyncio
import copy
import inspect
import json
import logging
import re
from abc import ABC
from collections import OrderedDict
from contextlib import ExitStack
from types import SimpleNamespace
from typing import (
    Any,
    Callable,
    Coroutine,
    Dict,
    List,
    Optional,
    Set,
    Tuple,
    Type,
    TypeVar,
    cast,
    get_args,
    get_origin,
    no_type_check,
)

from pydantic import Field, ValidationError, field_validator
from pydantic_settings import BaseSettings
from rich import print
from rich.console import Console
from rich.markup import escape
from rich.prompt import Prompt

from langroid.agent.chat_document import ChatDocMetaData, ChatDocument
from langroid.agent.tool_message import ToolMessage
from langroid.agent.xml_tool_message import XMLToolMessage
from langroid.exceptions import XMLException
from langroid.language_models.base import (
    LanguageModel,
    LLMConfig,
    LLMFunctionCall,
    LLMMessage,
    LLMResponse,
    LLMTokenUsage,
    OpenAIToolCall,
    StreamingIfAllowed,
    ToolChoiceTypes,
)
from langroid.language_models.openai_gpt import OpenAIGPT, OpenAIGPTConfig
from langroid.mytypes import Entity
from langroid.parsing.file_attachment import FileAttachment
from langroid.parsing.parse_json import extract_top_level_json
from langroid.parsing.parser import Parser, ParsingConfig
from langroid.prompts.prompts_config import PromptsConfig
from langroid.utils.configuration import settings
from langroid.utils.constants import (
    DONE,
    NO_ANSWER,
    PASS,
    PASS_TO,
    SEND_TO,
)
from langroid.utils.object_registry import ObjectRegistry
from langroid.utils.output import status
from langroid.utils.types import from_string, to_string
from langroid.vector_store.base import VectorStore, VectorStoreConfig

ORCHESTRATION_STRINGS = [DONE, PASS, PASS_TO, SEND_TO]
console = Console(quiet=settings.quiet)

logger = logging.getLogger(__name__)

T = TypeVar("T")


class AgentConfig(BaseSettings):
    """
    General config settings for an LLM agent. This is nested, combining configs of
    various components.
    """

    name: str = "LLM-Agent"
    debug: bool = False
    vecdb: Optional[VectorStoreConfig] = None
    llm: Optional[LLMConfig] = OpenAIGPTConfig()
    parsing: Optional[ParsingConfig] = ParsingConfig()
    prompts: Optional[PromptsConfig] = PromptsConfig()
    show_stats: bool = True  # show token usage/cost stats?
    hide_agent_response: bool = False  # hide agent response?
    add_to_registry: bool = True  # register agent in ObjectRegistry?
    respond_tools_only: bool = False  # respond only to tool messages (not plain text)?
    # allow multiple tool messages in a single response?
    allow_multiple_tools: bool = True
    human_prompt: str = (
        "Human (respond or q, x to exit current level, " "or hit enter to continue)"
    )

    @field_validator("name")
    @classmethod
    def check_name_alphanum(cls, v: str) -> str:
        if not re.match(r"^[a-zA-Z0-9_-]+$", v):
            raise ValueError(
                "The name must only contain alphanumeric characters, "
                "underscores, or hyphens, with no spaces"
            )
        return v


def noop_fn(*args: List[Any], **kwargs: Dict[str, Any]) -> None:
    pass


async def async_noop_fn(*args: List[Any], **kwargs: Dict[str, Any]) -> None:
    pass


async def async_lambda_noop_fn() -> Callable[..., Coroutine[Any, Any, None]]:
    return async_noop_fn


class Agent(ABC):
    """
    An Agent is an abstraction that typically (but not necessarily)
    encapsulates an LLM.
    """

    id: str = Field(default_factory=lambda: ObjectRegistry.new_id())
    # OpenAI tool-calls awaiting response; update when a tool result with Role.TOOL
    # is added to self.message_history
    oai_tool_calls: List[OpenAIToolCall] = []
    # Index of ALL tool calls generated by the agent
    oai_tool_id2call: Dict[str, OpenAIToolCall] = {}

    def __init__(self, config: AgentConfig = AgentConfig()):
        self.config = config
        self.id = ObjectRegistry.new_id()  # Initialize agent ID
        self.lock = asyncio.Lock()  # for async access to update self.llm.usage_cost
        self.dialog: List[Tuple[str, str]] = []  # seq of LLM (prompt, response) tuples
        self.llm_tools_map: Dict[str, Type[ToolMessage]] = {}
        self.llm_tools_handled: Set[str] = set()
        self.llm_tools_usable: Set[str] = set()
        self.llm_tools_known: Set[str] = set()  # all known tools, handled/used or not
        # Indicates which tool-names are allowed to be inferred when
        # the LLM "forgets" to include the request field in its tool-call.
        self.enabled_requests_for_inference: Optional[Set[str]] = (
            None  # If None, we allow all
        )
        self.interactive: bool = True  # may be modified by Task wrapper
        self.token_stats_str = ""
        self.default_human_response: Optional[str] = None
        self._indent = ""
        self.llm = LanguageModel.create(config.llm)
        self.vecdb = VectorStore.create(config.vecdb) if config.vecdb else None
        self.tool_error = False
        if config.parsing is not None and self.config.llm is not None:
            # token_encoding_model is used to obtain the tokenizer,
            # so in case it's an OpenAI model, we ensure that the tokenizer
            # corresponding to the model is used.
            if isinstance(self.llm, OpenAIGPT) and self.llm.is_openai_chat_model():
                config.parsing.token_encoding_model = self.llm.config.chat_model
        self.parser: Optional[Parser] = (
            Parser(config.parsing) if config.parsing else None
        )
        if config.add_to_registry:
            ObjectRegistry.register_object(self)

        self.callbacks = SimpleNamespace(
            start_llm_stream=lambda: noop_fn,
            start_llm_stream_async=async_lambda_noop_fn,
            cancel_llm_stream=noop_fn,
            finish_llm_stream=noop_fn,
            show_llm_response=noop_fn,
            show_agent_response=noop_fn,
            get_user_response=None,
            get_user_response_async=None,
            get_last_step=noop_fn,
            set_parent_agent=noop_fn,
            show_error_message=noop_fn,
            show_start_response=noop_fn,
        )
        Agent.init_state(self)

    def init_state(self) -> None:
        """Initialize all state vars. Called by Task.run() if restart is True"""
        self.total_llm_token_cost = 0.0
        self.total_llm_token_usage = 0

    @staticmethod
    def from_id(id: str) -> "Agent":
        return cast(Agent, ObjectRegistry.get(id))

    @staticmethod
    def delete_id(id: str) -> None:
        ObjectRegistry.remove(id)

    def entity_responders(
        self,
    ) -> List[
        Tuple[Entity, Callable[[None | str | ChatDocument], None | ChatDocument]]
    ]:
        """
        Sequence of (entity, response_method) pairs. This sequence is used
            in a `Task` to respond to the current pending message.
            See `Task.step()` for details.
        Returns:
            Sequence of (entity, response_method) pairs.
        """
        return [
            (Entity.AGENT, self.agent_response),
            (Entity.LLM, self.llm_response),
            (Entity.USER, self.user_response),
        ]

    def entity_responders_async(
        self,
    ) -> List[
        Tuple[
            Entity,
            Callable[
                [None | str | ChatDocument], Coroutine[Any, Any, None | ChatDocument]
            ],
        ]
    ]:
        """
        Async version of `entity_responders`. See there for details.
        """
        return [
            (Entity.AGENT, self.agent_response_async),
            (Entity.LLM, self.llm_response_async),
            (Entity.USER, self.user_response_async),
        ]

    @property
    def indent(self) -> str:
        """Indentation to print before any responses from the agent's entities."""
        return self._indent

    @indent.setter
    def indent(self, value: str) -> None:
        self._indent = value

    def update_dialog(self, prompt: str, output: str) -> None:
        self.dialog.append((prompt, output))

    def get_dialog(self) -> List[Tuple[str, str]]:
        return self.dialog

    def clear_dialog(self) -> None:
        self.dialog = []

    def _analyze_handler_params(
        self, handler_method: Any
    ) -> Tuple[bool, Optional[str], Optional[str]]:
        """
        Analyze parameters of a handler method to determine their types.

        Returns:
            Tuple of (has_annotations, agent_param_name, chat_doc_param_name)
            - has_annotations: True if useful type annotations were found
            - agent_param_name: Name of the agent parameter if found
            - chat_doc_param_name: Name of the chat_doc parameter if found
        """
        sig = inspect.signature(handler_method)
        params = list(sig.parameters.values())
        # Remove the first 'self' parameter
        params = params[1:]
        # Don't use name
        # [p for p in params if p.name != "self"]

        agent_param = None
        chat_doc_param = None
        has_annotations = False

        for param in params:
            # First try type annotations
            if param.annotation != inspect.Parameter.empty:
                ann_str = str(param.annotation)
                # Check for Agent-like types
                if (
                    inspect.isclass(param.annotation)
                    and issubclass(param.annotation, Agent)
                ) or (
                    not inspect.isclass(param.annotation)
                    and (
                        "Agent" in ann_str
                        or (
                            hasattr(param.annotation, "__name__")
                            and "Agent" in param.annotation.__name__
                        )
                    )
                ):
                    agent_param = param.name
                    has_annotations = True
                # Check for ChatDocument-like types
                elif (
                    param.annotation is ChatDocument
                    or "ChatDocument" in ann_str
                    or "ChatDoc" in ann_str
                ):
                    chat_doc_param = param.name
                    has_annotations = True

            # Fallback to parameter names
            elif param.name == "agent":
                agent_param = param.name
            elif param.name == "chat_doc":
                chat_doc_param = param.name

        return has_annotations, agent_param, chat_doc_param

    @no_type_check
    def _create_handler_wrapper(
        self,
        handler_method: Any,
        is_async: bool = False,
    ) -> Any:
        """
        Create a wrapper function for a handler method based on its signature.

        Args:
            message_class: The ToolMessage class
            handler_method: The handle/handle_async method
            is_async: Whether this is for an async handler

        Returns:
            Appropriate wrapper function
        """
        sig = inspect.signature(handler_method)
        params = list(sig.parameters.values())
        params = params[1:]
        # params = [p for p in params if p.name != "self"]

        has_annotations, agent_param, chat_doc_param = self._analyze_handler_params(
            handler_method,
        )

        # Build wrapper based on found parameters
        if len(params) == 0:
            if is_async:

                async def wrapper(obj: Any) -> Any:
                    return await obj.handle_async()

            else:

                def wrapper(obj: Any) -> Any:
                    return obj.handle()

        elif agent_param and chat_doc_param:
            # Both parameters present - build wrapper respecting their order
            param_names = [p.name for p in params]
            if param_names.index(agent_param) < param_names.index(chat_doc_param):
                # agent is first parameter
                if is_async:

                    async def wrapper(obj: Any, chat_doc: Any) -> Any:
                        return await obj.handle_async(self, chat_doc)

                else:

                    def wrapper(obj: Any, chat_doc: Any) -> Any:
                        return obj.handle(self, chat_doc)

            else:
                # chat_doc is first parameter
                if is_async:

                    async def wrapper(obj: Any, chat_doc: Any) -> Any:
                        return await obj.handle_async(chat_doc, self)

                else:

                    def wrapper(obj: Any, chat_doc: Any) -> Any:
                        return obj.handle(chat_doc, self)

        elif agent_param and not chat_doc_param:
            # Only agent parameter
            if is_async:

                async def wrapper(obj: Any) -> Any:
                    return await obj.handle_async(self)

            else:

                def wrapper(obj: Any) -> Any:
                    return obj.handle(self)

        elif chat_doc_param and not agent_param:
            # Only chat_doc parameter
            if is_async:

                async def wrapper(obj: Any, chat_doc: Any) -> Any:
                    return await obj.handle_async(chat_doc)

            else:

                def wrapper(obj: Any, chat_doc: Any) -> Any:
                    return obj.handle(chat_doc)

        else:
            # No recognized parameters - backward compatibility
            # Assume single parameter is chat_doc (legacy behavior)
            if len(params) == 1:
                if is_async:

                    async def wrapper(obj: Any, chat_doc: Any) -> Any:
                        return await obj.handle_async(chat_doc)

                else:

                    def wrapper(obj: Any, chat_doc: Any) -> Any:
                        return obj.handle(chat_doc)

            else:
                # Multiple unrecognized parameters - best guess
                if is_async:

                    async def wrapper(obj: Any, chat_doc: Any) -> Any:
                        return await obj.handle_async(chat_doc)

                else:

                    def wrapper(obj: Any, chat_doc: Any) -> Any:
                        return obj.handle(chat_doc)

        return wrapper

    def _get_tool_list(
        self, message_class: Optional[Type[ToolMessage]] = None
    ) -> List[str]:
        """
        If `message_class` is None, return a list of all known tool names.
        Otherwise, first add the tool name corresponding to the message class
        (which is the value of the `request` field of the message class),
        to the `self.llm_tools_map` dict, and then return a list
        containing this tool name.

        Args:
            message_class (Optional[Type[ToolMessage]]): The message class whose tool
                name is to be returned; Optional, default is None.
                if None, return a list of all known tool names.

        Returns:
            List[str]: List of tool names: either just the tool name corresponding
                to the message class, or all known tool names
                (when `message_class` is None).

        """
        if message_class is None:
            return list(self.llm_tools_map.keys())

        if not issubclass(message_class, ToolMessage):
            raise ValueError("message_class must be a subclass of ToolMessage")
        tool = message_class.default_value("request")

        """
        if tool has handler method explicitly defined - use it,
        otherwise use the tool name as the handler
        """
        if hasattr(message_class, "_handler"):
            handler = getattr(message_class, "_handler", tool)
        else:
            handler = tool

        self.llm_tools_map[tool] = message_class
        if (
            hasattr(message_class, "handle")
            and inspect.isfunction(message_class.handle)
            and not hasattr(self, handler)
        ):
            """
            If the message class has a `handle` method,
            and agent does NOT have a tool handler method,
            then we create a method for the agent whose name
            is the value of `handler`, and whose body is the `handle` method.
            This removes a separate step of having to define this method
            for the agent, and also keeps the tool definition AND handling
            in one place, i.e. in the message class.
            See `tests/main/test_stateless_tool_messages.py` for an example.
            """
            wrapper = self._create_handler_wrapper(
                message_class.handle,
                is_async=False,
            )
            setattr(self, handler, wrapper)
        elif (
            hasattr(message_class, "response")
            and inspect.isfunction(message_class.response)
            and not hasattr(self, handler)
        ):
            has_chat_doc_arg = (
                len(inspect.signature(message_class.response).parameters) > 2
            )
            if has_chat_doc_arg:

                def response_wrapper_with_chat_doc(obj: Any, chat_doc: Any) -> Any:
                    return obj.response(self, chat_doc)

                setattr(self, handler, response_wrapper_with_chat_doc)
            else:

                def response_wrapper_no_chat_doc(obj: Any) -> Any:
                    return obj.response(self)

                setattr(self, handler, response_wrapper_no_chat_doc)

        if hasattr(message_class, "handle_message_fallback") and (
            inspect.isfunction(message_class.handle_message_fallback)
        ):
            # When a ToolMessage has a `handle_message_fallback` method,
            # we inject it into the agent as a method, overriding the default
            # `handle_message_fallback` method (which does nothing).
            # It's possible multiple tool messages have a `handle_message_fallback`,
            # in which case, the last one inserted will be used.
            def fallback_wrapper(msg: Any) -> Any:
                return message_class.handle_message_fallback(self, msg)

            setattr(
                self,
                "handle_message_fallback",
                fallback_wrapper,
            )

        async_handler_name = f"{handler}_async"
        if (
            hasattr(message_class, "handle_async")
            and inspect.isfunction(message_class.handle_async)
            and not hasattr(self, async_handler_name)
        ):
            wrapper = self._create_handler_wrapper(
                message_class.handle_async,
                is_async=True,
            )
            setattr(self, async_handler_name, wrapper)
        elif (
            hasattr(message_class, "response_async")
            and inspect.isfunction(message_class.response_async)
            and not hasattr(self, async_handler_name)
        ):
            has_chat_doc_arg = (
                len(inspect.signature(message_class.response_async).parameters) > 2
            )

            if has_chat_doc_arg:

                @no_type_check
                async def handler(obj, chat_doc):
                    return await obj.response_async(self, chat_doc)

            else:

                @no_type_check
                async def handler(obj):
                    return await obj.response_async(self)

            setattr(self, async_handler_name, handler)

        return [tool]

    def enable_message_handling(
        self, message_class: Optional[Type[ToolMessage]] = None
    ) -> None:
        """
        Enable an agent to RESPOND (i.e. handle) a "tool" message of a specific type
            from LLM. Also "registers" (i.e. adds) the `message_class` to the
            `self.llm_tools_map` dict.

        Args:
            message_class (Optional[Type[ToolMessage]]): The message class to enable;
                Optional; if None, all known message classes are enabled for handling.

        """
        for t in self._get_tool_list(message_class):
            self.llm_tools_handled.add(t)

    def disable_message_handling(
        self,
        message_class: Optional[Type[ToolMessage]] = None,
    ) -> None:
        """
        Disable a message class from being handled by this Agent.

        Args:
            message_class (Optional[Type[ToolMessage]]): The message class to disable.
                If None, all message classes are disabled.
        """
        for t in self._get_tool_list(message_class):
            self.llm_tools_handled.discard(t)

    def sample_multi_round_dialog(self) -> str:
        """
        Generate a sample multi-round dialog based on enabled message classes.
        Returns:
            str: The sample dialog string.
        """
        enabled_classes: List[Type[ToolMessage]] = list(self.llm_tools_map.values())
        # use at most 2 sample conversations, no need to be exhaustive;
        sample_convo = [
            msg_cls().usage_examples(random=True)  # type: ignore
            for i, msg_cls in enumerate(enabled_classes)
            if i < 2
        ]
        return "\n\n".join(sample_convo)

    def create_agent_response(
        self,
        content: str | None = None,
        files: List[FileAttachment] = [],
        content_any: Any = None,
        tool_messages: List[ToolMessage] = [],
        oai_tool_calls: Optional[List[OpenAIToolCall]] = None,
        oai_tool_choice: ToolChoiceTypes | Dict[str, Dict[str, str] | str] = "auto",
        oai_tool_id2result: OrderedDict[str, str] | None = None,
        function_call: LLMFunctionCall | None = None,
        recipient: str = "",
    ) -> ChatDocument:
        """Template for agent_response."""
        return self.response_template(
            Entity.AGENT,
            content=content,
            files=files,
            content_any=content_any,
            tool_messages=tool_messages,
            oai_tool_calls=oai_tool_calls,
            oai_tool_choice=oai_tool_choice,
            oai_tool_id2result=oai_tool_id2result,
            function_call=function_call,
            recipient=recipient,
        )

    def render_agent_response(
        self,
        results: Optional[str | OrderedDict[str, str] | ChatDocument],
    ) -> None:
        """
        Render the response from the agent, typically from tool-handling.
        Args:
            results: results from tool-handling, which may be a string,
                a dict of tool results, or a ChatDocument.
        """
        if self.config.hide_agent_response or results is None:
            return
        if isinstance(results, str):
            results_str = results
        elif isinstance(results, ChatDocument):
            results_str = results.content
        elif isinstance(results, dict):
            results_str = json.dumps(results, indent=2)
        if not settings.quiet:
            console.print(f"[red]{self.indent}", end="")
            print(f"[red]Agent: {escape(results_str)}")

    def _agent_response_final(
        self,
        msg: Optional[str | ChatDocument],
        results: Optional[str | OrderedDict[str, str] | ChatDocument],
    ) -> Optional[ChatDocument]:
        """
        Convert results to final response.
        """
        if results is None:
            return None
        if isinstance(results, str):
            results_str = results
        elif isinstance(results, ChatDocument):
            results_str = results.content
        elif isinstance(results, dict):
            results_str = json.dumps(results, indent=2)
        if not settings.quiet:
            self.render_agent_response(results)
        maybe_json = len(extract_top_level_json(results_str)) > 0
        self.callbacks.show_agent_response(
            content=results_str,
            language="json" if maybe_json else "text",
            is_tool=(
                isinstance(results, ChatDocument)
                and self.has_tool_message_attempt(results)
            ),
        )
        if isinstance(results, ChatDocument):
            # Preserve trail of tool_ids for OpenAI Assistant fn-calls
            results.metadata.tool_ids = (
                [] if msg is None or isinstance(msg, str) else msg.metadata.tool_ids
            )
            results.metadata.agent_id = self.id
            return results
        sender_name = self.config.name
        if isinstance(msg, ChatDocument) and msg.function_call is not None:
            # if result was from handling an LLM `function_call`,
            # set sender_name to name of the function_call
            sender_name = msg.function_call.name

        results_str, id2result, oai_tool_id = self.process_tool_results(
            results if isinstance(results, str) else "",
            id2result=None if isinstance(results, str) else results,
            tool_calls=(msg.oai_tool_calls if isinstance(msg, ChatDocument) else None),
        )
        return ChatDocument(
            content=results_str,
            oai_tool_id2result=id2result,
            metadata=ChatDocMetaData(
                source=Entity.AGENT,
                sender=Entity.AGENT,
                agent_id=self.id,
                sender_name=sender_name,
                oai_tool_id=oai_tool_id,
                # preserve trail of tool_ids for OpenAI Assistant fn-calls
                tool_ids=(
                    [] if msg is None or isinstance(msg, str) else msg.metadata.tool_ids
                ),
            ),
        )

    async def agent_response_async(
        self,
        msg: Optional[str | ChatDocument] = None,
    ) -> Optional[ChatDocument]:
        """
        Asynch version of `agent_response`. See there for details.
        """
        if msg is None:
            return None

        results = await self.handle_message_async(msg)

        return self._agent_response_final(msg, results)

    def agent_response(
        self,
        msg: Optional[str | ChatDocument] = None,
    ) -> Optional[ChatDocument]:
        """
        Response from the "agent itself", typically (but not only)
        used to handle LLM's "tool message" or `function_call`
        (e.g. OpenAI `function_call`).
        Args:
            msg (str|ChatDocument): the input to respond to: if msg is a string,
                and it contains a valid JSON-structured "tool message", or
                if msg is a ChatDocument, and it contains a `function_call`.
        Returns:
            Optional[ChatDocument]: the response, packaged as a ChatDocument

        """
        if msg is None:
            return None

        results = self.handle_message(msg)

        return self._agent_response_final(msg, results)

    def process_tool_results(
        self,
        results: str,
        id2result: OrderedDict[str, str] | None,
        tool_calls: List[OpenAIToolCall] | None = None,
    ) -> Tuple[str, Dict[str, str] | None, str | None]:
        """
        Process results from a response, based on whether
        they are results of OpenAI tool-calls from THIS agent, so that
        we can construct an appropriate LLMMessage that contains tool results.

        Args:
            results (str): A possible string result from handling tool(s)
            id2result (OrderedDict[str,str]|None): A dict of OpenAI tool id -> result,
                if there are multiple tool results.
            tool_calls (List[OpenAIToolCall]|None): List of OpenAI tool-calls that the
                results are a response to.

        Return:
            - str: The response string
            - Dict[str,str]|None: A dict of OpenAI tool id -> result, if there are
                multiple tool results.
            - str|None: tool_id if there was a single tool result

        """
        id2result_ = copy.deepcopy(id2result) if id2result is not None else None
        results_str = ""
        oai_tool_id = None

        if results != "":
            # in this case ignore id2result
            assert (
                id2result is None
            ), "id2result should be None when results string is non-empty!"
            results_str = results
            if len(self.oai_tool_calls) > 0:
                # We only have one result, so in case there is a
                # "pending" OpenAI tool-call, we expect no more than 1 such.
                assert (
                    len(self.oai_tool_calls) == 1
                ), "There are multiple pending tool-calls, but only one result!"
                # We record the tool_id of the tool-call that
                # the result is a response to, so that ChatDocument.to_LLMMessage
                # can properly set the `tool_call_id` field of the LLMMessage.
                oai_tool_id = self.oai_tool_calls[0].id
        elif id2result is not None and id2result_ is not None:  # appease mypy
            if len(id2result_) == len(self.oai_tool_calls):
                # if the number of pending tool calls equals the number of results,
                # then ignore the ids in id2result, and use the results in order,
                # which is preserved since id2result is an OrderedDict.
                assert len(id2result_) > 1, "Expected to see > 1 result in id2result!"
                results_str = ""
                id2result_ = OrderedDict(
                    zip(
                        [tc.id or "" for tc in self.oai_tool_calls], id2result_.values()
                    )
                )
            else:
                assert (
                    tool_calls is not None
                ), "tool_calls cannot be None when id2result is not None!"
                # This must be an OpenAI tool id -> result map;
                # However some ids may not correspond to the tool-calls in the list of
                # pending tool-calls (self.oai_tool_calls).
                # Such results are concatenated into a simple string, to store in the
                # ChatDocument.content, and the rest
                # (i.e. those that DO correspond to tools in self.oai_tool_calls)
                # are stored as a dict in ChatDocument.oai_tool_id2result.

                # OAI tools from THIS agent, awaiting response
                pending_tool_ids = [tc.id for tc in self.oai_tool_calls]
                # tool_calls that the results are a response to
                # (but these may have been sent from another agent, hence may not be in
                # self.oai_tool_calls)
                parent_tool_id2name = {
                    tc.id: tc.function.name
                    for tc in tool_calls or []
                    if tc.function is not None
                }

                # (id, result) for result NOT corresponding to self.oai_tool_calls,
                # i.e. these are results of EXTERNAL tool-calls from another agent.
                external_tool_id_results = []

                for tc_id, result in id2result.items():
                    if tc_id not in pending_tool_ids:
                        external_tool_id_results.append((tc_id, result))
                        id2result_.pop(tc_id)
                if len(external_tool_id_results) == 0:
                    results_str = ""
                elif len(external_tool_id_results) == 1:
                    results_str = external_tool_id_results[0][1]
                else:
                    results_str = "\n\n".join(
                        [
                            f"Result from tool/function "
                            f"{parent_tool_id2name[id]}: {result}"
                            for id, result in external_tool_id_results
                        ]
                    )

                if len(id2result_) == 0:
                    id2result_ = None
                elif len(id2result_) == 1 and len(external_tool_id_results) == 0:
                    results_str = list(id2result_.values())[0]
                    oai_tool_id = list(id2result_.keys())[0]
                    id2result_ = None

        return results_str, id2result_, oai_tool_id

    def response_template(
        self,
        e: Entity,
        content: str | None = None,
        files: List[FileAttachment] = [],
        content_any: Any = None,
        tool_messages: List[ToolMessage] = [],
        oai_tool_calls: Optional[List[OpenAIToolCall]] = None,
        oai_tool_choice: ToolChoiceTypes | Dict[str, Dict[str, str] | str] = "auto",
        oai_tool_id2result: OrderedDict[str, str] | None = None,
        function_call: LLMFunctionCall | None = None,
        recipient: str = "",
    ) -> ChatDocument:
        """Template for response from entity `e`."""
        return ChatDocument(
            content=content or "",
            files=files,
            content_any=content_any,
            tool_messages=tool_messages,
            oai_tool_calls=oai_tool_calls,
            oai_tool_id2result=oai_tool_id2result,
            function_call=function_call,
            oai_tool_choice=oai_tool_choice,
            metadata=ChatDocMetaData(
                source=e, sender=e, sender_name=self.config.name, recipient=recipient
            ),
        )

    def create_user_response(
        self,
        content: str | None = None,
        files: List[FileAttachment] = [],
        content_any: Any = None,
        tool_messages: List[ToolMessage] = [],
        oai_tool_calls: List[OpenAIToolCall] | None = None,
        oai_tool_choice: ToolChoiceTypes | Dict[str, Dict[str, str] | str] = "auto",
        oai_tool_id2result: OrderedDict[str, str] | None = None,
        function_call: LLMFunctionCall | None = None,
        recipient: str = "",
    ) -> ChatDocument:
        """Template for user_response."""
        return self.response_template(
            e=Entity.USER,
            content=content,
            files=files,
            content_any=content_any,
            tool_messages=tool_messages,
            oai_tool_calls=oai_tool_calls,
            oai_tool_choice=oai_tool_choice,
            oai_tool_id2result=oai_tool_id2result,
            function_call=function_call,
            recipient=recipient,
        )

    def user_can_respond(self, msg: Optional[str | ChatDocument] = None) -> bool:
        """
        Whether the user can respond to a message.

        Args:
            msg (str|ChatDocument): the string to respond to.

        Returns:

        """
        # When msg explicitly addressed to user, this means an actual human response
        # is being sought.
        need_human_response = (
            isinstance(msg, ChatDocument) and msg.metadata.recipient == Entity.USER
        )

        if not self.interactive and not need_human_response:
            return False

        return True

    def _user_response_final(
        self, msg: Optional[str | ChatDocument], user_msg: str
    ) -> Optional[ChatDocument]:
        """
        Convert user_msg to final response.
        """
        if not user_msg:
            need_human_response = (
                isinstance(msg, ChatDocument) and msg.metadata.recipient == Entity.USER
            )
            user_msg = (
                (self.default_human_response or "null") if need_human_response else ""
            )
        user_msg = user_msg.strip()

        tool_ids = []
        if msg is not None and isinstance(msg, ChatDocument):
            tool_ids = msg.metadata.tool_ids

        # only return non-None result if user_msg not empty
        if not user_msg:
            return None
        else:
            if user_msg.startswith("SYSTEM"):
                user_msg = user_msg.replace("SYSTEM", "").strip()
                source = Entity.SYSTEM
                sender = Entity.SYSTEM
            else:
                source = Entity.USER
                sender = Entity.USER
            return ChatDocument(
                content=user_msg,
                metadata=ChatDocMetaData(
                    agent_id=self.id,
                    source=source,
                    sender=sender,
                    # preserve trail of tool_ids for OpenAI Assistant fn-calls
                    tool_ids=tool_ids,
                ),
            )

    async def user_response_async(
        self,
        msg: Optional[str | ChatDocument] = None,
    ) -> Optional[ChatDocument]:
        """
        Asynch version of `user_response`. See there for details.
        """
        if not self.user_can_respond(msg):
            return None

        if self.default_human_response is not None:
            user_msg = self.default_human_response
        else:
            if (
                self.callbacks.get_user_response_async is not None
                and self.callbacks.get_user_response_async is not async_noop_fn
            ):
                user_msg = await self.callbacks.get_user_response_async(prompt="")
            elif self.callbacks.get_user_response is not None:
                user_msg = self.callbacks.get_user_response(prompt="")
            else:
                user_msg = Prompt.ask(
                    f"[blue]{self.indent}"
                    + self.config.human_prompt
                    + f"\n{self.indent}"
                )

        return self._user_response_final(msg, user_msg)

    def user_response(
        self,
        msg: Optional[str | ChatDocument] = None,
    ) -> Optional[ChatDocument]:
        """
        Get user response to current message. Could allow (human) user to intervene
        with an actual answer, or quit using "q" or "x"

        Args:
            msg (str|ChatDocument): the string to respond to.

        Returns:
            (str) User response, packaged as a ChatDocument

        """

        if not self.user_can_respond(msg):
            return None

        if self.default_human_response is not None:
            user_msg = self.default_human_response
        else:
            if self.callbacks.get_user_response is not None:
                # ask user with empty prompt: no need for prompt
                # since user has seen the conversation so far.
                # But non-empty prompt can be useful when Agent
                # uses a tool that requires user input, or in other scenarios.
                user_msg = self.callbacks.get_user_response(prompt="")
            else:
                user_msg = Prompt.ask(
                    f"[blue]{self.indent}"
                    + self.config.human_prompt
                    + f"\n{self.indent}"
                )

        return self._user_response_final(msg, user_msg)

    @no_type_check
    def llm_can_respond(self, message: Optional[str | ChatDocument] = None) -> bool:
        """
        Whether the LLM can respond to a message.
        Args:
            message (str|ChatDocument): message or ChatDocument object to respond to.

        Returns:

        """
        if self.llm is None:
            return False

        if message is not None and len(self.try_get_tool_messages(message)) > 0:
            # if there is a valid "tool" message (either JSON or via `function_call`)
            # then LLM cannot respond to it
            return False

        return True

    def can_respond(self, message: Optional[str | ChatDocument] = None) -> bool:
        """
        Whether the agent can respond to a message.
        Used in Task.py to skip a sub-task when we know it would not respond.
        Args:
            message (str|ChatDocument): message or ChatDocument object to respond to.
        """
        tools = self.try_get_tool_messages(message)
        if len(tools) == 0 and self.config.respond_tools_only:
            return False
        if message is not None and self.has_only_unhandled_tools(message):
            # The message has tools that are NOT enabled to be handled by this agent,
            # which means the agent cannot respond to it.
            return False
        return True

    def create_llm_response(
        self,
        content: str | None = None,
        content_any: Any = None,
        tool_messages: List[ToolMessage] = [],
        oai_tool_calls: None | List[OpenAIToolCall] = None,
        oai_tool_choice: ToolChoiceTypes | Dict[str, Dict[str, str] | str] = "auto",
        oai_tool_id2result: OrderedDict[str, str] | None = None,
        function_call: LLMFunctionCall | None = None,
        recipient: str = "",
    ) -> ChatDocument:
        """Template for llm_response."""
        return self.response_template(
            Entity.LLM,
            content=content,
            content_any=content_any,
            tool_messages=tool_messages,
            oai_tool_calls=oai_tool_calls,
            oai_tool_choice=oai_tool_choice,
            oai_tool_id2result=oai_tool_id2result,
            function_call=function_call,
            recipient=recipient,
        )

    @no_type_check
    async def llm_response_async(
        self,
        message: Optional[str | ChatDocument] = None,
    ) -> Optional[ChatDocument]:
        """
        Asynch version of `llm_response`. See there for details.
        """
        if message is None or not self.llm_can_respond(message):
            return None

        if isinstance(message, ChatDocument):
            prompt = message.content
        else:
            prompt = message

        output_len = self.config.llm.model_max_output_tokens
        if self.num_tokens(prompt) + output_len > self.llm.completion_context_length():
            output_len = self.llm.completion_context_length() - self.num_tokens(prompt)
            if output_len < self.config.llm.min_output_tokens:
                raise ValueError(
                    """
                Token-length of Prompt + Output is longer than the
                completion context length of the LLM!
                """
                )
            else:
                logger.warning(
                    f"""
                Requested output length has been shortened to {output_len}
                so that the total length of Prompt + Output is less than
                the completion context length of the LLM.
                """
                )

        with StreamingIfAllowed(self.llm, self.llm.get_stream()):
            response = await self.llm.agenerate(prompt, output_len)

        if not self.llm.get_stream() or response.cached and not settings.quiet:
            # We would have already displayed the msg "live" ONLY if
            # streaming was enabled, AND we did not find a cached response.
            # If we are here, it means the response has not yet been displayed.
            cached = f"[red]{self.indent}(cached)[/red]" if response.cached else ""
            print(cached + "[green]" + escape(response.message))
        async with self.lock:
            self.update_token_usage(
                response,
                prompt,
                self.llm.get_stream(),
                chat=False,  # i.e. it's a completion model not chat model
                print_response_stats=self.config.show_stats and not settings.quiet,
            )
        cdoc = ChatDocument.from_LLMResponse(response, displayed=True)
        # Preserve trail of tool_ids for OpenAI Assistant fn-calls
        cdoc.metadata.tool_ids = (
            [] if isinstance(message, str) else message.metadata.tool_ids
        )
        return cdoc

    @no_type_check
    def llm_response(
        self,
        message: Optional[str | ChatDocument] = None,
    ) -> Optional[ChatDocument]:
        """
        LLM response to a prompt.
        Args:
            message (str|ChatDocument): prompt string, or ChatDocument object

        Returns:
            Response from LLM, packaged as a ChatDocument
        """
        if message is None or not self.llm_can_respond(message):
            return None

        if isinstance(message, ChatDocument):
            prompt = message.content
        else:
            prompt = message

        with ExitStack() as stack:  # for conditionally using rich spinner
            if not self.llm.get_stream():
                # show rich spinner only if not streaming!
                cm = status("LLM responding to message...")
                stack.enter_context(cm)
            output_len = self.config.llm.model_max_output_tokens
            if (
                self.num_tokens(prompt) + output_len
                > self.llm.completion_context_length()
            ):
                output_len = self.llm.completion_context_length() - self.num_tokens(
                    prompt
                )
                if output_len < self.config.llm.min_output_tokens:
                    raise ValueError(
                        """
                    Token-length of Prompt + Output is longer than the
                    completion context length of the LLM!
                    """
                    )
                else:
                    logger.warning(
                        f"""
                    Requested output length has been shortened to {output_len}
                    so that the total length of Prompt + Output is less than
                    the completion context length of the LLM.
                    """
                    )
            if self.llm.get_stream() and not settings.quiet:
                console.print(f"[green]{self.indent}", end="")
            response = self.llm.generate(prompt, output_len)

        if not self.llm.get_stream() or response.cached and not settings.quiet:
            # we would have already displayed the msg "live" ONLY if
            # streaming was enabled, AND we did not find a cached response
            # If we are here, it means the response has not yet been displayed.
            cached = "[red](cached)[/red]" if response.cached else ""
            console.print(f"[green]{self.indent}", end="")
            print(cached + "[green]" + escape(response.message))
        self.update_token_usage(
            response,
            prompt,
            self.llm.get_stream(),
            chat=False,  # i.e. it's a completion model not chat model
            print_response_stats=self.config.show_stats and not settings.quiet,
        )
        cdoc = ChatDocument.from_LLMResponse(response, displayed=True)
        # Preserve trail of tool_ids for OpenAI Assistant fn-calls
        cdoc.metadata.tool_ids = (
            [] if isinstance(message, str) else message.metadata.tool_ids
        )
        return cdoc

    def has_tool_message_attempt(self, msg: str | ChatDocument | None) -> bool:
        """
        Check whether msg contains a Tool/fn-call attempt (by the LLM).

        CAUTION: This uses self.get_tool_messages(msg) which as a side-effect
        may update msg.tool_messages when msg is a ChatDocument, if there are
        any tools in msg.
        """
        if msg is None:
            return False
        if isinstance(msg, ChatDocument):
            if len(msg.tool_messages) > 0:
                return True
            if msg.metadata.sender != Entity.LLM:
                return False
        try:
            tools = self.get_tool_messages(msg)
            return len(tools) > 0
        except (ValidationError, XMLException):
            # there is a tool/fn-call attempt but had a validation error,
            # so we still consider this a tool message "attempt"
            return True
        return False

    def _tool_recipient_match(self, tool: ToolMessage) -> bool:
        """Is tool enabled for handling by this agent and intended for this
        agent to handle (i.e. if there's any explicit `recipient` field exists in
        tool, then it matches this agent's name)?
        """
        if tool.default_value("request") not in self.llm_tools_handled:
            return False
        if hasattr(tool, "recipient") and isinstance(tool.recipient, str):
            return tool.recipient == "" or tool.recipient == self.config.name
        return True

    def has_only_unhandled_tools(self, msg: str | ChatDocument) -> bool:
        """
        Does the msg have at least one tool, and none of the tools in the msg are
        handleable by this agent?
        """
        if msg is None:
            return False
        tools = self.try_get_tool_messages(msg, all_tools=True)
        if len(tools) == 0:
            return False
        return all(not self._tool_recipient_match(t) for t in tools)

    def try_get_tool_messages(
        self,
        msg: str | ChatDocument | None,
        all_tools: bool = False,
    ) -> List[ToolMessage]:
        try:
            return self.get_tool_messages(msg, all_tools)
        except (ValidationError, XMLException):
            return []

    def get_tool_messages(
        self,
        msg: str | ChatDocument | None,
        all_tools: bool = False,
    ) -> List[ToolMessage]:
        """
        Get ToolMessages recognized in msg, handle-able by this agent.
        NOTE: as a side-effect, this will update msg.tool_messages
        when msg is a ChatDocument and msg contains tool messages.
        The intent here is that update=True should be set ONLY within agent_response()
        or agent_response_async() methods. In other words, we want to persist the
        msg.tool_messages only AFTER the agent has had a chance to handle the tools.

        Args:
            msg (str|ChatDocument): the message to extract tools from.
            all_tools (bool):
                - if True, return all tools,
                    i.e. any recognized tool in self.llm_tools_known,
                    whether it is handled by this agent or not;
                - otherwise, return only the tools handled by this agent.

        Returns:
            List[ToolMessage]: list of ToolMessage objects
        """

        if msg is None:
            return []

        if isinstance(msg, str):
            json_tools = self.get_formatted_tool_messages(msg)
            if all_tools:
                return json_tools
            else:
                return [
                    t
                    for t in json_tools
                    if self._tool_recipient_match(t) and t.default_value("request")
                ]

        if all_tools and len(msg.all_tool_messages) > 0:
            # We've already identified all_tool_messages in the msg;
            # return the corresponding ToolMessage objects
            return msg.all_tool_messages
        if len(msg.tool_messages) > 0:
            # We've already found tool_messages,
            # (either via OpenAI Fn-call or Langroid-native ToolMessage);
            # or they were added by an agent_response.
            # note these could be from a forwarded msg from another agent,
            # so return ONLY the messages THIS agent to enabled to handle.
            if all_tools:
                return msg.tool_messages
            return [t for t in msg.tool_messages if self._tool_recipient_match(t)]
        assert isinstance(msg, ChatDocument)
        if (
            msg.content != ""
            and msg.oai_tool_calls is None
            and msg.function_call is None
        ):

            tools = self.get_formatted_tool_messages(
                msg.content, from_llm=msg.metadata.sender == Entity.LLM
            )
            msg.all_tool_messages = tools
            # filter for actually handle-able tools, and recipient is this agent
            my_tools = [t for t in tools if self._tool_recipient_match(t)]
            msg.tool_messages = my_tools

            if all_tools:
                return tools
            else:
                return my_tools

        # otherwise, we look for `tool_calls` (possibly multiple)
        tools = self.get_oai_tool_calls_classes(msg)
        msg.all_tool_messages = tools
        my_tools = [t for t in tools if self._tool_recipient_match(t)]
        msg.tool_messages = my_tools

        if len(tools) == 0:
            # otherwise, we look for a `function_call`
            fun_call_cls = self.get_function_call_class(msg)
            tools = [fun_call_cls] if fun_call_cls is not None else []
            msg.all_tool_messages = tools
            my_tools = [t for t in tools if self._tool_recipient_match(t)]
            msg.tool_messages = my_tools
        if all_tools:
            return tools
        else:
            return my_tools

    def get_formatted_tool_messages(
        self, input_str: str, from_llm: bool = True
    ) -> List[ToolMessage]:
        """
        Returns ToolMessage objects (tools) corresponding to
        tool-formatted substrings, if any.
        ASSUMPTION - These tools are either ALL JSON-based, or ALL XML-based
        (i.e. not a mix of both).
        Terminology: a "formatted tool msg" is one which the LLM generates as
            part of its raw string output, rather than within a JSON object
            in the API response (i.e. this method does not extract tools/fns returned
            by OpenAI's tools/fns API or similar APIs).

        Args:
            input_str (str): input string, typically a message sent by an LLM
            from_llm (bool): whether the input was generated by the LLM. If so,
                we track malformed tool calls.

        Returns:
            List[ToolMessage]: list of ToolMessage objects
        """
        self.tool_error = False
        substrings = XMLToolMessage.find_candidates(input_str)
        is_json = False
        if len(substrings) == 0:
            substrings = extract_top_level_json(input_str)
            is_json = len(substrings) > 0
            if not is_json:
                return []

        results = [self._get_one_tool_message(j, is_json, from_llm) for j in substrings]
        valid_results = [r for r in results if r is not None]
        # If any tool is correctly formed we do not set the flag
        if len(valid_results) > 0:
            self.tool_error = False
        return valid_results

    def get_function_call_class(self, msg: ChatDocument) -> Optional[ToolMessage]:
        """
        From ChatDocument (constructed from an LLM Response), get the `ToolMessage`
        corresponding to the `function_call` if it exists.
        """
        if msg.function_call is None:
            return None
        tool_name = msg.function_call.name
        tool_msg = msg.function_call.arguments or {}
        self.tool_error = False
        if tool_name not in self.llm_tools_handled:
            logger.warning(
                f"""
                The function_call '{tool_name}' is not handled
                by the agent named '{self.config.name}'!
                If you intended this agent to handle this function_call,
                either the fn-call name is incorrectly generated by the LLM,
                (in which case you may need to adjust your LLM instructions),
                or you need to enable this agent to handle this fn-call.
                """
            )
            if (
                tool_name not in self.all_llm_tools_known
                and msg.metadata.sender == Entity.LLM
            ):
                self.tool_error = True
            return None
        tool_class = self.llm_tools_map[tool_name]
        tool_msg.update(dict(request=tool_name))
        try:
            tool = tool_class.model_validate(tool_msg)
        except ValidationError as ve:
            # Store tool class as an attribute on the exception
            ve.tool_class = tool_class  # type: ignore
            raise ve
        return tool

    def get_oai_tool_calls_classes(self, msg: ChatDocument) -> List[ToolMessage]:
        """
        From ChatDocument (constructed from an LLM Response), get
         a list of ToolMessages corresponding to the `tool_calls`, if any.
        """

        if msg.oai_tool_calls is None:
            return []
        tools = []
        all_errors = True
        for tc in msg.oai_tool_calls:
            if tc.function is None:
                continue
            tool_name = tc.function.name
            tool_msg = tc.function.arguments or {}
            if tool_name not in self.llm_tools_handled:
                logger.warning(
                    f"""
                    The tool_call '{tool_name}' is not handled
                    by the agent named '{self.config.name}'!
                    If you intended this agent to handle this function_call,
                    either the fn-call name is incorrectly generated by the LLM,
                    (in which case you may need to adjust your LLM instructions),
                    or you need to enable this agent to handle this fn-call.
                    """
                )
                continue
            all_errors = False
            tool_class = self.llm_tools_map[tool_name]
            tool_msg.update(dict(request=tool_name))
            try:
                tool = tool_class.model_validate(tool_msg)
            except ValidationError as ve:
                # Store tool class as an attribute on the exception
                ve.tool_class = tool_class  # type: ignore
                raise ve
            tool.id = tc.id or ""
            tools.append(tool)
        # When no tool is valid and the message was produced
        # by the LLM, set the recovery flag
        self.tool_error = all_errors and msg.metadata.sender == Entity.LLM
        return tools

    def tool_validation_error(
        self, ve: ValidationError, tool_class: Optional[Type[ToolMessage]] = None
    ) -> str:
        """
        Handle a validation error raised when parsing a tool message,
            when there is a legit tool name used, but it has missing/bad fields.
        Args:
            ve (ValidationError): The exception raised
            tool_class (Optional[Type[ToolMessage]]): The tool class that
                failed validation

        Returns:
            str: The error message to send back to the LLM
        """
        # First try to get tool class from the exception itself
        if hasattr(ve, "tool_class") and ve.tool_class:
            tool_name = ve.tool_class.default_value("request")  # type: ignore
        elif tool_class is not None:
            tool_name = tool_class.default_value("request")
        else:
            # Fallback: try to extract from error context if available
            tool_name = "Unknown Tool"
        bad_field_errors = "\n".join(
            [f"{e['loc']}: {e['msg']}" for e in ve.errors() if "loc" in e]
        )
        return f"""
        There were one or more errors in your attempt to use the
        TOOL or function_call named '{tool_name}':
        {bad_field_errors}
        Please write your message again, correcting the errors.
        """

    def _get_multiple_orch_tool_errs(
        self, tools: List[ToolMessage]
    ) -> List[str | ChatDocument | None]:
        """
        Return error document if the message contains multiple orchestration tools
        """
        # check whether there are multiple orchestration-tools (e.g. DoneTool etc),
        # in which case set result to error-string since we don't yet support
        # multi-tools with one or more orch tools.
        from langroid.agent.tools.orchestration import (
            AgentDoneTool,
            AgentSendTool,
            DonePassTool,
            DoneTool,
            ForwardTool,
            PassTool,
            SendTool,
        )
        from langroid.agent.tools.recipient_tool import RecipientTool

        ORCHESTRATION_TOOLS = (
            AgentDoneTool,
            DoneTool,
            PassTool,
            DonePassTool,
            ForwardTool,
            RecipientTool,
            SendTool,
            AgentSendTool,
        )

        has_orch = any(isinstance(t, ORCHESTRATION_TOOLS) for t in tools)
        if has_orch and len(tools) > 1:
            return ["ERROR: Use ONE tool at a time!"] * len(tools)

        return []

    def _handle_message_final(
        self, tools: List[ToolMessage], results: List[str | ChatDocument | None]
    ) -> None | str | OrderedDict[str, str] | ChatDocument:
        """
        Convert results to final response
        """
        # extract content from ChatDocument results so we have all str|None
        results = [r.content if isinstance(r, ChatDocument) else r for r in results]

        tool_names = [t.default_value("request") for t in tools]

        has_ids = all([t.id != "" for t in tools])
        if has_ids:
            id2result = OrderedDict(
                (t.id, r)
                for t, r in zip(tools, results)
                if r is not None and isinstance(r, str)
            )
            result_values = list(id2result.values())
            if len(id2result) > 1 and any(
                orch_str in r
                for r in result_values
                for orch_str in ORCHESTRATION_STRINGS
            ):
                # Cannot support multi-tool results containing orchestration strings!
                # Replace results with err string to force LLM to retry
                err_str = "ERROR: Please use ONE tool at a time!"
                id2result = OrderedDict((id, err_str) for id in id2result.keys())

        name_results_list = [
            (name, r) for name, r in zip(tool_names, results) if r is not None
        ]
        if len(name_results_list) == 0:
            return None

        # there was a non-None result

        if has_ids and len(id2result) > 1:
            # if there are multiple OpenAI Tool results, return them as a dict
            return id2result

        # multi-results: prepend the tool name to each result
        str_results = [f"Result from {name}: {r}" for name, r in name_results_list]
        final = "\n\n".join(str_results)
        return final

    async def handle_message_async(
        self, msg: str | ChatDocument
    ) -> None | str | OrderedDict[str, str] | ChatDocument:
        """
        Asynch version of `handle_message`. See there for details.
        """
        try:
            tools = self.get_tool_messages(msg)
            tools = [t for t in tools if self._tool_recipient_match(t)]
        except ValidationError as ve:
            # correct tool name but bad fields
            return self.tool_validation_error(ve)
        except XMLException as xe:  # from XMLToolMessage parsing
            return str(xe)
        except ValueError:
            # invalid tool name
            # We return None since returning "invalid tool name" would
            # be considered a valid result in task loop, and would be treated
            # as a response to the tool message even though the tool was not intended
            # for this agent.
            return None
        if len(tools) > 1 and not self.config.allow_multiple_tools:
            return self.to_ChatDocument("ERROR: Use ONE tool at a time!")
        if len(tools) == 0:
            fallback_result = self.handle_message_fallback(msg)
            if fallback_result is None:
                return None
            return self.to_ChatDocument(
                fallback_result,
                chat_doc=msg if isinstance(msg, ChatDocument) else None,
            )
        chat_doc = msg if isinstance(msg, ChatDocument) else None

        results = self._get_multiple_orch_tool_errs(tools)
        if not results:
            results = [
                await self.handle_tool_message_async(t, chat_doc=chat_doc)
                for t in tools
            ]
            # if there's a solitary ChatDocument|str result, return it as is
            if len(results) == 1 and isinstance(results[0], (str, ChatDocument)):
                return results[0]

        return self._handle_message_final(tools, results)

    def handle_message(
        self, msg: str | ChatDocument
    ) -> None | str | OrderedDict[str, str] | ChatDocument:
        """
        Handle a "tool" message either a string containing one or more
        valid "tool" JSON substrings,  or a
        ChatDocument containing a `function_call` attribute.
        Handle with the corresponding handler method, and return
        the results as a combined string.

        Args:
            msg (str | ChatDocument): The string or ChatDocument to handle

        Returns:
            The result of the handler method can be:
             - None if no tools successfully handled, or no tools present
             - str if langroid-native JSON tools were handled, and results concatenated,
                 OR there's a SINGLE OpenAI tool-call.
                (We do this so the common scenario of a single tool/fn-call
                 has a simple behavior).
             - Dict[str, str] if multiple OpenAI tool-calls were handled
                 (dict is an id->result map)
             - ChatDocument if a handler returned a ChatDocument, intended to be the
                 final response of the `agent_response` method.
        """
        try:
            tools = self.get_tool_messages(msg)
            tools = [t for t in tools if self._tool_recipient_match(t)]
        except ValidationError as ve:
            # correct tool name but bad fields
            return self.tool_validation_error(ve)
        except XMLException as xe:  # from XMLToolMessage parsing
            return str(xe)
        except ValueError:
            # invalid tool name
            # We return None since returning "invalid tool name" would
            # be considered a valid result in task loop, and would be treated
            # as a response to the tool message even though the tool was not intended
            # for this agent.
            return None
        if len(tools) == 0:
            fallback_result = self.handle_message_fallback(msg)
            if fallback_result is None:
                return None
            return self.to_ChatDocument(
                fallback_result,
                chat_doc=msg if isinstance(msg, ChatDocument) else None,
            )

        results: List[str | ChatDocument | None] = []
        if len(tools) > 1 and not self.config.allow_multiple_tools:
            results = ["ERROR: Use ONE tool at a time!"] * len(tools)
        if not results:
            results = self._get_multiple_orch_tool_errs(tools)
        if not results:
            chat_doc = msg if isinstance(msg, ChatDocument) else None
            results = [self.handle_tool_message(t, chat_doc=chat_doc) for t in tools]
            # if there's a solitary ChatDocument|str result, return it as is
            if len(results) == 1 and isinstance(results[0], (str, ChatDocument)):
                return results[0]

        return self._handle_message_final(tools, results)

    @property
    def all_llm_tools_known(self) -> set[str]:
        """All known tools; this may extend self.llm_tools_known."""
        return self.llm_tools_known

    def handle_message_fallback(self, msg: str | ChatDocument) -> Any:
        """
        Fallback method for the case where the msg has no tools that
        can be handled by this agent.
        This method can be overridden by subclasses, e.g.,
        to create a "reminder" message when a tool is expected but the LLM "forgot"
        to generate one.

        Args:
            msg (str | ChatDocument): The input msg to handle
        Returns:
            Any: The result of the handler method
        """
        return None

    def _get_one_tool_message(
        self, tool_candidate_str: str, is_json: bool = True, from_llm: bool = True
    ) -> Optional[ToolMessage]:
        """
        Parse the tool_candidate_str into ANY ToolMessage KNOWN to agent --
        This includes non-used/handled tools, i.e. any tool in self.all_llm_tools_known.
        The exception to this is below where we try our best to infer the tool
        when the LLM has "forgotten" to include the "request" field in the tool str ---
        in this case we ONLY look at the possible set of HANDLED tools, i.e.
        self.llm_tools_handled.
        """
        if is_json:
            maybe_tool_dict = json.loads(tool_candidate_str)
        else:
            try:
                maybe_tool_dict = XMLToolMessage.extract_field_values(
                    tool_candidate_str
                )
            except Exception as e:
                from langroid.exceptions import XMLException

                raise XMLException(f"Error extracting XML fields:\n {str(e)}")
        # check if the maybe_tool_dict contains a "properties" field
        # which further contains the actual tool-call
        # (some weak LLMs do this). E.g. gpt-4o sometimes generates this:
        # TOOL: {
        #     "type": "object",
        #     "properties": {
        #         "request": "square",
        #         "number": 9
        #     },
        #     "required": [
        #         "number",
        #         "request"
        #     ]
        # }

        if not isinstance(maybe_tool_dict, dict):
            self.tool_error = from_llm
            return None

        properties = maybe_tool_dict.get("properties")
        if isinstance(properties, dict):
            maybe_tool_dict = properties
        request = maybe_tool_dict.get("request")
        if request is None:
            if self.enabled_requests_for_inference is None:
                possible = [self.llm_tools_map[r] for r in self.llm_tools_handled]
            else:
                allowable = self.enabled_requests_for_inference.intersection(
                    self.llm_tools_handled
                )
                possible = [self.llm_tools_map[r] for r in allowable]

            default_keys = set(ToolMessage.model_fields.keys())
            request_keys = set(maybe_tool_dict.keys())

            def maybe_parse(tool: type[ToolMessage]) -> Optional[ToolMessage]:
                all_keys = set(tool.model_fields.keys())
                non_inherited_keys = all_keys.difference(default_keys)
                # If the request has any keys not valid for the tool and
                # does not specify some key specific to the type
                # (e.g. not just `purpose`), the LLM must explicitly specify `request`
                if not (
                    request_keys.issubset(all_keys)
                    and len(request_keys.intersection(non_inherited_keys)) > 0
                ):
                    return None

                try:
                    return tool.model_validate(maybe_tool_dict)
                except ValidationError:
                    return None

            candidate_tools = list(
                filter(
                    lambda t: t is not None,
                    map(maybe_parse, possible),
                )
            )

            # If only one valid candidate exists, we infer
            # "request" to be the only possible value
            if len(candidate_tools) == 1:
                return candidate_tools[0]
            else:
                self.tool_error = from_llm
                return None

        if not isinstance(request, str) or request not in self.all_llm_tools_known:
            self.tool_error = from_llm
            return None

        message_class = self.llm_tools_map.get(request)
        if message_class is None:
            logger.warning(f"No message class found for request '{request}'")
            self.tool_error = from_llm
            return None

        try:
            message = message_class.model_validate(maybe_tool_dict)
        except ValidationError as ve:
            self.tool_error = from_llm
            # Store tool class as an attribute on the exception
            ve.tool_class = message_class  # type: ignore
            raise ve
        return message

    def to_ChatDocument(
        self,
        msg: Any,
        orig_tool_name: str | None = None,
        chat_doc: Optional[ChatDocument] = None,
        author_entity: Entity = Entity.AGENT,
    ) -> Optional[ChatDocument]:
        """
        Convert result of a responder (agent_response or llm_response, or task.run()),
        or tool handler, or handle_message_fallback,
        to a ChatDocument, to enable handling by other
        responders/tasks in a task loop possibly involving multiple agents.

        Args:
            msg (Any): The result of a responder or tool handler or task.run()
            orig_tool_name (str): The original tool name that generated the response,
                if any.
            chat_doc (ChatDocument): The original ChatDocument object that `msg`
                is a response to.
            author_entity (Entity): The intended author of the result ChatDocument
        """
        if msg is None or isinstance(msg, ChatDocument):
            return msg

        is_agent_author = author_entity == Entity.AGENT

        if isinstance(msg, str):
            return self.response_template(author_entity, content=msg, content_any=msg)
        elif isinstance(msg, ToolMessage):
            # result is a ToolMessage, so...
            result_tool_name = msg.default_value("request")
            if (
                is_agent_author
                and result_tool_name in self.llm_tools_handled
                and (orig_tool_name is None or orig_tool_name != result_tool_name)
            ):
                # TODO: do we need to remove the tool message from the chat_doc?
                # if (chat_doc is not None and
                #     msg in chat_doc.tool_messages):
                #    chat_doc.tool_messages.remove(msg)
                # if we can handle it, do so
                result = self.handle_tool_message(msg, chat_doc=chat_doc)
                if result is not None and isinstance(result, ChatDocument):
                    return result
            else:
                # else wrap it in an agent response and return it so
                # orchestrator can find a respondent
                return self.response_template(author_entity, tool_messages=[msg])
        else:
            result = to_string(msg)

        return (
            None
            if result is None
            else self.response_template(author_entity, content=result, content_any=msg)
        )

    def from_ChatDocument(self, msg: ChatDocument, output_type: Type[T]) -> Optional[T]:
        """
        Extract a desired output_type from a ChatDocument object.
        We use this fallback order:
        - if `msg.content_any` exists and matches the output_type, return it
        - if `msg.content` exists and output_type is str return it
        - if output_type is a ToolMessage, return the first tool in `msg.tool_messages`
        - if output_type is a list of ToolMessage,
            return all tools in `msg.tool_messages`
        - search for a tool in `msg.tool_messages` that has a field of output_type,
             and if found, return that field value
        - return None if all the above fail
        """
        content = msg.content
        if output_type is str and content != "":
            return cast(T, content)
        content_any = msg.content_any
        if content_any is not None and isinstance(content_any, output_type):
            return cast(T, content_any)

        tools = self.try_get_tool_messages(msg, all_tools=True)

        if get_origin(output_type) is list:
            list_element_type = get_args(output_type)[0]
            if issubclass(list_element_type, ToolMessage):
                # list_element_type is a subclass of ToolMessage:
                # We output a list of objects derived from list_element_type
                return cast(
                    T,
                    [t for t in tools if isinstance(t, list_element_type)],
                )
        elif get_origin(output_type) is None and issubclass(output_type, ToolMessage):
            # output_type is a subclass of ToolMessage:
            # return the first tool that has this specific output_type
            for tool in tools:
                if isinstance(tool, output_type):
                    return cast(T, tool)
            return None
        elif get_origin(output_type) is None and output_type in (str, int, float, bool):
            # attempt to get the output_type from the content,
            # if it's a primitive type
            primitive_value = from_string(content, output_type)  # type: ignore
            if primitive_value is not None:
                return cast(T, primitive_value)

        # then search for output_type as a field in a tool
        for tool in tools:
            value = tool.get_value_of_type(output_type)
            if value is not None:
                return cast(T, value)
        return None

    def _maybe_truncate_result(
        self,
        result: str | ChatDocument | None,
        max_tokens: int | None,
    ) -> str | ChatDocument | None:
        """
        Truncate the result string to `max_tokens` tokens.
        """

        if result is None or max_tokens is None:
            return result
        result_str = result.content if isinstance(result, ChatDocument) else result
        num_tokens = (
            self.parser.num_tokens(result_str)
            if self.parser is not None
            else len(result_str) / 4.0
        )
        if num_tokens <= max_tokens:
            return result
        truncate_warning = f"""
        The TOOL result was large, so it was truncated to {max_tokens} tokens.
        To get the full result, the TOOL must be called again.
        """
        if isinstance(result, str):
            return (
                self.parser.truncate_tokens(result, max_tokens)
                if self.parser is not None
                else result[: max_tokens * 4]  # approx truncate
            ) + truncate_warning
        elif isinstance(result, ChatDocument):
            result.content = (
                self.parser.truncate_tokens(result.content, max_tokens)
                if self.parser is not None
                else result.content[: max_tokens * 4]  # approx truncate
            ) + truncate_warning
            return result

    async def handle_tool_message_async(
        self,
        tool: ToolMessage,
        chat_doc: Optional[ChatDocument] = None,
    ) -> None | str | ChatDocument:
        """
        Asynch version of `handle_tool_message`. See there for details.
        """
        tool_name = tool.default_value("request")
        if hasattr(tool, "_handler"):
            handler_name = getattr(tool, "_handler", tool_name)
        else:
            handler_name = tool_name
        handler_method = getattr(self, handler_name + "_async", None)
        if handler_method is None:
            return self.handle_tool_message(tool, chat_doc=chat_doc)
        has_chat_doc_arg = (
            chat_doc is not None
            and "chat_doc" in inspect.signature(handler_method).parameters
        )
        try:
            if has_chat_doc_arg:
                maybe_result = await handler_method(tool, chat_doc=chat_doc)
            else:
                maybe_result = await handler_method(tool)
            result = self.to_ChatDocument(maybe_result, tool_name, chat_doc)
        except Exception as e:
            # raise the error here since we are sure it's
            # not a pydantic validation error,
            # which we check in `handle_message`
            raise e
        return self._maybe_truncate_result(
            result, tool._max_result_tokens
        )  # type: ignore

    def handle_tool_message(
        self,
        tool: ToolMessage,
        chat_doc: Optional[ChatDocument] = None,
    ) -> None | str | ChatDocument:
        """
        Respond to a tool request from the LLM, in the form of an ToolMessage object.
        Args:
            tool: ToolMessage object representing the tool request.
            chat_doc: Optional ChatDocument object containing the tool request.
                This is passed to the tool-handler method only if it has a `chat_doc`
                argument.

        Returns:

        """
        tool_name = tool.default_value("request")
        if hasattr(tool, "_handler"):
            handler_name = getattr(tool, "_handler", tool_name)
        else:
            handler_name = tool_name
        handler_method = getattr(self, handler_name, None)
        if handler_method is None:
            return None
        has_chat_doc_arg = (
            chat_doc is not None
            and "chat_doc" in inspect.signature(handler_method).parameters
        )
        try:
            if has_chat_doc_arg:
                maybe_result = handler_method(tool, chat_doc=chat_doc)
            else:
                maybe_result = handler_method(tool)
            result = self.to_ChatDocument(maybe_result, tool_name, chat_doc)
        except Exception as e:
            # raise the error here since we are sure it's
            # not a pydantic validation error,
            # which we check in `handle_message`
            raise e
        return self._maybe_truncate_result(
            result, tool._max_result_tokens
        )  # type: ignore

    def num_tokens(self, prompt: str | List[LLMMessage]) -> int:
        if self.parser is None:
            raise ValueError("Parser must be set, to count tokens")
        if isinstance(prompt, str):
            return self.parser.num_tokens(prompt)
        else:
            return sum(
                [
                    self.parser.num_tokens(m.content)
                    + self.parser.num_tokens(str(m.function_call or ""))
                    for m in prompt
                ]
            )

    def _get_response_stats(
        self, chat_length: int, tot_cost: float, response: LLMResponse
    ) -> str:
        """
        Get LLM response stats as a string

        Args:
            chat_length (int): number of messages in the chat
            tot_cost (float): total cost of the chat so far
            response (LLMResponse): LLMResponse object
        """

        if self.config.llm is None:
            logger.warning("LLM config is None, cannot get response stats")
            return ""
        if response.usage:
            in_tokens = response.usage.prompt_tokens
            out_tokens = response.usage.completion_tokens
            llm_response_cost = format(response.usage.cost, ".4f")
            cumul_cost = format(tot_cost, ".4f")
            assert isinstance(self.llm, LanguageModel)
            context_length = self.llm.chat_context_length()
            max_out = self.config.llm.model_max_output_tokens

            llm_model = (
                "no-LLM" if self.config.llm is None else self.llm.config.chat_model
            )
            # tot cost across all LLMs, agents
            all_cost = format(self.llm.tot_tokens_cost()[1], ".4f")
            return (
                f"[bold]Stats:[/bold] [magenta]N_MSG={chat_length}, "
                f"TOKENS: in={in_tokens}, out={out_tokens}, "
                f"max={max_out}, ctx={context_length}, "
                f"COST: now=${llm_response_cost}, cumul=${cumul_cost}, "
                f"tot=${all_cost} "
                f"[bold]({llm_model})[/bold][/magenta]"
            )
        return ""

    def update_token_usage(
        self,
        response: LLMResponse,
        prompt: str | List[LLMMessage],
        stream: bool,
        chat: bool = True,
        print_response_stats: bool = True,
    ) -> None:
        """
        Updates `response.usage` obj (token usage and cost fields) if needed.
        An update is needed only if:
        - stream is True (i.e. streaming was enabled), and
        - the response was NOT obtained from cached, and
        - the API did NOT provide the usage/cost fields during streaming
          (As of Sep 2024, the OpenAI API started providing these; for other APIs
            this may not necessarily be the case).

        Args:
            response (LLMResponse): LLMResponse object
            prompt (str | List[LLMMessage]): prompt or list of LLMMessage objects
            stream (bool): whether to update the usage in the response object
                if the response is not cached.
            chat (bool): whether this is a chat model or a completion model
            print_response_stats (bool): whether to print the response stats
        """
        if response is None or self.llm is None:
            return

        no_usage_info = response.usage is None or response.usage.prompt_tokens == 0
        # Note: If response was not streamed, then
        # `response.usage` would already have been set by the API,
        # so we only need to update in the stream case.
        if stream and no_usage_info:
            # usage, cost = 0 when response is from cache
            prompt_tokens = 0
            completion_tokens = 0
            cost = 0.0
            if not response.cached:
                prompt_tokens = self.num_tokens(prompt)
                completion_tokens = self.num_tokens(response.message)
                if response.function_call is not None:
                    completion_tokens += self.num_tokens(str(response.function_call))
                cost = self.compute_token_cost(prompt_tokens, 0, completion_tokens)
            response.usage = LLMTokenUsage(
                prompt_tokens=prompt_tokens,
                completion_tokens=completion_tokens,
                cost=cost,
            )

        # update total counters
        if response.usage is not None:
            self.total_llm_token_cost += response.usage.cost
            self.total_llm_token_usage += response.usage.total_tokens
            self.llm.update_usage_cost(
                chat,
                response.usage.prompt_tokens,
                response.usage.completion_tokens,
                response.usage.cost,
            )
            chat_length = 1 if isinstance(prompt, str) else len(prompt)
            self.token_stats_str = self._get_response_stats(
                chat_length, self.total_llm_token_cost, response
            )
            if print_response_stats:
                print(self.indent + self.token_stats_str)

    def compute_token_cost(self, prompt: int, cached: int, completion: int) -> float:
        price = cast(LanguageModel, self.llm).chat_cost()
        return (
            price[0] * (prompt - cached) + price[1] * cached + price[2] * completion
        ) / 1000

    def ask_agent(
        self,
        agent: "Agent",
        request: str,
        no_answer: str = NO_ANSWER,
        user_confirm: bool = True,
    ) -> Optional[str]:
        """
        Send a request to another agent, possibly after confirming with the user.
        This is not currently used, since we rely on the task loop and
        `RecipientTool` to address requests to other agents. It is generally best to
        avoid using this method.

        Args:
            agent (Agent): agent to ask
            request (str): request to send
            no_answer (str): expected response when agent does not know the answer
            user_confirm (bool): whether to gate the request with a human confirmation

        Returns:
            str: response from agent
        """
        agent_type = type(agent).__name__
        if user_confirm:
            user_response = Prompt.ask(
                f"""[magenta]Here is the request or message:
                {request}
                Should I forward this to {agent_type}?""",
                default="y",
                choices=["y", "n"],
            )
            if user_response not in ["y", "yes"]:
                return None
        answer = agent.llm_response(request)
        if answer != no_answer:
            return (f"{agent_type} says: " + str(answer)).strip()
        return None
