import copy
import inspect
import json
import logging
import textwrap
from contextlib import ExitStack
from inspect import isclass
from typing import Any, Dict, List, Optional, Self, Set, Tuple, Type, Union, cast

import openai
from pydantic import BaseModel, ValidationError
from pydantic.fields import ModelPrivateAttr
from rich import print
from rich.console import Console
from rich.markup import escape

from langroid.agent.base import Agent, AgentConfig, async_noop_fn, noop_fn
from langroid.agent.chat_document import ChatDocument
from langroid.agent.tool_message import (
    ToolMessage,
    format_schema_for_strict,
)
from langroid.agent.xml_tool_message import XMLToolMessage
from langroid.language_models.base import (
    LLMFunctionCall,
    LLMFunctionSpec,
    LLMMessage,
    LLMResponse,
    OpenAIJsonSchemaSpec,
    OpenAIToolSpec,
    Role,
    StreamingIfAllowed,
    ToolChoiceTypes,
)
from langroid.language_models.openai_gpt import OpenAIGPT
from langroid.mytypes import Entity, NonToolAction
from langroid.utils.configuration import settings
from langroid.utils.object_registry import ObjectRegistry
from langroid.utils.output import status
from langroid.utils.pydantic_utils import PydanticWrapper, get_pydantic_wrapper
from langroid.utils.types import is_callable

console = Console()

logger = logging.getLogger(__name__)


class ChatAgentConfig(AgentConfig):
    """
    Configuration for ChatAgent

    Attributes:
        system_message: system message to include in message sequence
             (typically defines role and task of agent).
             Used only if `task` is not specified in the constructor.
        user_message: user message to include in message sequence.
             Used only if `task` is not specified in the constructor.
        use_tools: whether to use our own ToolMessages mechanism
        handle_llm_no_tool (Any): desired agent_response when
            LLM generates non-tool msg.
        use_functions_api: whether to use functions/tools native to the LLM API
                (e.g. OpenAI's `function_call` or `tool_call` mechanism)
        use_tools_api: When `use_functions_api` is True, if this is also True,
            the OpenAI tool-call API is used, rather than the older/deprecated
            function-call API. However the tool-call API has some tricky aspects,
            hence we set this to False by default.
        strict_recovery: whether to enable strict schema recovery when there
            is a tool-generation error.
        enable_orchestration_tool_handling: whether to enable handling of orchestration
            tools, e.g. ForwardTool, DoneTool, PassTool, etc.
        output_format: When supported by the LLM (certain OpenAI LLMs
            and local LLMs served by providers such as vLLM), ensures
            that the output is a JSON matching the corresponding
            schema via grammar-based decoding
        handle_output_format: When `output_format` is a `ToolMessage` T,
            controls whether T is "enabled for handling".
        use_output_format: When `output_format` is a `ToolMessage` T,
            controls whether T is "enabled for use" (by LLM) and
            instructions on using T are added to the system message.
        instructions_output_format: Controls whether we generate instructions for
            `output_format` in the system message.
        use_tools_on_output_format: Controls whether to automatically switch
            to the Langroid-native tools mechanism when `output_format` is set.
            Note that LLMs may generate tool calls which do not belong to
            `output_format` even when strict JSON mode is enabled, so this should be
            enabled when such tool calls are not desired.
        output_format_include_defaults: Whether to include fields with default arguments
            in the output schema
        full_citations: Whether to show source reference citation + content for each
            citation, or just the main reference citation.
    """

    system_message: str = "You are a helpful assistant."
    user_message: Optional[str] = None
    handle_llm_no_tool: Any = None
    use_tools: bool = True
    use_functions_api: bool = False
    use_tools_api: bool = True
    strict_recovery: bool = True
    enable_orchestration_tool_handling: bool = True
    output_format: Optional[type] = None
    handle_output_format: bool = True
    use_output_format: bool = True
    instructions_output_format: bool = True
    output_format_include_defaults: bool = True
    use_tools_on_output_format: bool = True
    full_citations: bool = True  # show source + content for each citation?

    def _set_fn_or_tools(self) -> None:
        """
        Enable Langroid Tool or OpenAI-like fn-calling,
        depending on config settings.
        """
        if not self.use_functions_api or not self.use_tools:
            return
        if self.use_functions_api and self.use_tools:
            logger.debug(
                """
                You have enabled both `use_tools` and `use_functions_api`.
                Setting `use_functions_api` to False.
                """
            )
            self.use_tools = True
            self.use_functions_api = False


class ChatAgent(Agent):
    """
    Chat Agent interacting with external env
    (could be human, or external tools).
    The agent (the LLM actually) is provided with an optional "Task Spec",
    which is a sequence of `LLMMessage`s. These are used to initialize
    the `task_messages` of the agent.
    In most applications we will use a `ChatAgent` rather than a bare `Agent`.
    The `Agent` class mainly exists to hold various common methods and attributes.
    One difference between `ChatAgent` and `Agent` is that `ChatAgent`'s
    `llm_response` method uses "chat mode" API (i.e. one that takes a
    message sequence rather than a single message),
    whereas the same method in the `Agent` class uses "completion mode" API (i.e. one
    that takes a single message).
    """

    def __init__(
        self,
        config: ChatAgentConfig = ChatAgentConfig(),
        task: Optional[List[LLMMessage]] = None,
    ):
        """
        Chat-mode agent initialized with task spec as the initial message sequence
        Args:
            config: settings for the agent

        """
        super().__init__(config)
        self.config: ChatAgentConfig = config
        self.config._set_fn_or_tools()
        self.message_history: List[LLMMessage] = []
        self.init_state()
        # An agent's "task" is defined by a system msg and an optional user msg;
        # These are "priming" messages that kick off the agent's conversation.
        self.system_message: str = self.config.system_message
        self.user_message: str | None = self.config.user_message

        if task is not None:
            # if task contains a system msg, we override the config system msg
            if len(task) > 0 and task[0].role == Role.SYSTEM:
                self.system_message = task[0].content
            # if task contains a user msg, we override the config user msg
            if len(task) > 1 and task[1].role == Role.USER:
                self.user_message = task[1].content

        # system-level instructions for using tools/functions:
        # We maintain these as tools/functions are enabled/disabled,
        # and whenever an LLM response is sought, these are used to
        # recreate the system message (via `_create_system_and_tools_message`)
        # each time, so it reflects the current set of enabled tools/functions.
        # (a) these are general instructions on using certain tools/functions,
        #   if they are specified in a ToolMessage class as a classmethod `instructions`
        self.system_tool_instructions: str = ""
        # (b) these are only for the builtin in Langroid TOOLS mechanism:
        self.system_tool_format_instructions: str = ""

        self.llm_functions_map: Dict[str, LLMFunctionSpec] = {}
        self.llm_functions_handled: Set[str] = set()
        self.llm_functions_usable: Set[str] = set()
        self.llm_function_force: Optional[Dict[str, str]] = None

        self.output_format: Optional[type[ToolMessage | BaseModel]] = None

        self.saved_requests_and_tool_setings = self._requests_and_tool_settings()
        # This variable is not None and equals a `ToolMessage` T, if and only if:
        # (a) T has been set as the output_format of this agent, AND
        # (b) T has been "enabled for use" ONLY for enforcing this output format, AND
        # (c) T has NOT been explicitly "enabled for use" by this Agent.
        self.enabled_use_output_format: Optional[type[ToolMessage]] = None
        # As above but deals with "enabled for handling" instead of "enabled for use".
        self.enabled_handling_output_format: Optional[type[ToolMessage]] = None
        if config.output_format is not None:
            self.set_output_format(config.output_format)
        # instructions specifically related to enforcing `output_format`
        self.output_format_instructions = ""

        # controls whether to disable strict schemas for this agent if
        # strict mode causes exception
        self.disable_strict = False
        # Tracks whether any strict tool is enabled; used to determine whether to set
        # `self.disable_strict` on an exception
        self.any_strict = False
        # Tracks the set of tools on which we force-disable strict decoding
        self.disable_strict_tools_set: set[str] = set()

        if self.config.enable_orchestration_tool_handling:
            # Only enable HANDLING by `agent_response`, NOT LLM generation of these.
            # This is useful where tool-handlers or agent_response generate these
            # tools, and need to be handled.
            # We don't want enable orch tool GENERATION by default, since that
            # might clutter-up the LLM system message unnecessarily.
            from langroid.agent.tools.orchestration import (
                AgentDoneTool,
                AgentSendTool,
                DonePassTool,
                DoneTool,
                ForwardTool,
                PassTool,
                ResultTool,
                SendTool,
            )

            self.enable_message(ForwardTool, use=False, handle=True)
            self.enable_message(DoneTool, use=False, handle=True)
            self.enable_message(AgentDoneTool, use=False, handle=True)
            self.enable_message(PassTool, use=False, handle=True)
            self.enable_message(DonePassTool, use=False, handle=True)
            self.enable_message(SendTool, use=False, handle=True)
            self.enable_message(AgentSendTool, use=False, handle=True)
            self.enable_message(ResultTool, use=False, handle=True)

    def init_state(self) -> None:
        """
        Initialize the state of the agent. Just conversation state here,
        but subclasses can override this to initialize other state.
        """
        super().init_state()
        self.clear_history(0)
        self.clear_dialog()

    @staticmethod
    def from_id(id: str) -> "ChatAgent":
        """
        Get an agent from its ID
        Args:
            agent_id (str): ID of the agent
        Returns:
            ChatAgent: The agent with the given ID
        """
        return cast(ChatAgent, Agent.from_id(id))

    def clone(self, i: int = 0) -> "ChatAgent":
        """Create i'th clone of this agent, ensuring tool use/handling is cloned.
        Important: We assume all member variables are in the __init__ method here
        and in the Agent class.
        TODO: We are attempting to clone an agent after its state has been
        changed in possibly many ways. Below is an imperfect solution. Caution advised.
        Revisit later.
        """
        agent_cls = type(self)
        # Use model_copy to preserve Pydantic subclass types (like MockLMConfig)
        # instead of deepcopy which loses subclass information
        config_copy = self.config.model_copy(deep=True)
        config_copy.name = f"{config_copy.name}-{i}"
        new_agent = agent_cls(config_copy)
        new_agent.system_tool_instructions = self.system_tool_instructions
        new_agent.system_tool_format_instructions = self.system_tool_format_instructions
        new_agent.llm_tools_map = self.llm_tools_map
        new_agent.llm_tools_known = self.llm_tools_known
        new_agent.llm_tools_handled = self.llm_tools_handled
        new_agent.llm_tools_usable = self.llm_tools_usable
        new_agent.llm_functions_map = self.llm_functions_map
        new_agent.llm_functions_handled = self.llm_functions_handled
        new_agent.llm_functions_usable = self.llm_functions_usable
        new_agent.llm_function_force = self.llm_function_force
        # Caution - we are copying the vector-db, maybe we don't always want this?
        new_agent.vecdb = self.vecdb
        new_agent.id = ObjectRegistry.new_id()
        if self.config.add_to_registry:
            ObjectRegistry.register_object(new_agent)
        return new_agent

    def _strict_mode_for_tool(self, tool: str | type[ToolMessage]) -> bool:
        """Should we enable strict mode for a given tool?"""
        if isinstance(tool, str):
            tool_class = self.llm_tools_map[tool]
        else:
            tool_class = tool
        name = tool_class.default_value("request")
        if name in self.disable_strict_tools_set or self.disable_strict:
            return False
        strict: Optional[bool] = tool_class.default_value("strict")
        if strict is None:
            strict = self._strict_tools_available()

        return strict

    def _fn_call_available(self) -> bool:
        """Does this agent's LLM support function calling?"""
        return self.llm is not None and self.llm.supports_functions_or_tools()

    def _strict_tools_available(self) -> bool:
        """Does this agent's LLM support strict tools?"""
        return (
            not self.disable_strict
            and self.llm is not None
            and isinstance(self.llm, OpenAIGPT)
            and self.llm.config.parallel_tool_calls is False
            and self.llm.supports_strict_tools
        )

    def _json_schema_available(self) -> bool:
        """Does this agent's LLM support strict JSON schema output format?"""
        return (
            not self.disable_strict
            and self.llm is not None
            and isinstance(self.llm, OpenAIGPT)
            and self.llm.supports_json_schema
        )

    def set_system_message(self, msg: str) -> None:
        self.system_message = msg
        if len(self.message_history) > 0:
            # if there is message history, update the system message in it
            self.message_history[0].content = msg

    def set_user_message(self, msg: str) -> None:
        self.user_message = msg

    @property
    def task_messages(self) -> List[LLMMessage]:
        """
        The task messages are the initial messages that define the task
        of the agent. There will be at least a system message plus possibly a user msg.
        Returns:
            List[LLMMessage]: the task messages
        """
        msgs = [self._create_system_and_tools_message()]
        if self.user_message:
            msgs.append(LLMMessage(role=Role.USER, content=self.user_message))
        return msgs

    def _drop_msg_update_tool_calls(self, msg: LLMMessage) -> None:
        id2idx = {t.id: i for i, t in enumerate(self.oai_tool_calls)}
        if msg.role == Role.TOOL:
            # dropping tool result, so ADD the corresponding tool-call back
            # to the list of pending calls!
            id = msg.tool_call_id
            if id in self.oai_tool_id2call:
                self.oai_tool_calls.append(self.oai_tool_id2call[id])
        elif msg.tool_calls is not None:
            # dropping a msg with tool-calls, so DROP these from pending list
            # as well as from id -> call map
            for tool_call in msg.tool_calls:
                if tool_call.id in id2idx:
                    self.oai_tool_calls.pop(id2idx[tool_call.id])
                if tool_call.id in self.oai_tool_id2call:
                    del self.oai_tool_id2call[tool_call.id]

    def clear_history(self, start: int = -2, end: int = -1) -> None:
        """
        Clear the message history, deleting  messages from index `start`,
        up to index `end`.

        Args:
            start (int): index of first message to delete; default = -2
                    (i.e. delete last 2 messages, typically these
                    are the last user and assistant messages)
            end (int): index of last message to delete; Default = -1
                    (i.e. delete all messages up to the last one)
        """
        n = len(self.message_history)
        if start < 0:
            start = max(0, n + start)
        end_ = n if end == -1 else end + 1
        dropped = self.message_history[start:end_]
        # consider the dropped msgs in REVERSE order, so we are
        # carefully updating self.oai_tool_calls
        for msg in reversed(dropped):
            self._drop_msg_update_tool_calls(msg)
            # clear out the chat document from the ObjectRegistry
            ChatDocument.delete_id(msg.chat_document_id)
        del self.message_history[start:end_]

    def update_history(self, message: str, response: str) -> None:
        """
        Update the message history with the latest user message and LLM response.
        Args:
            message (str): user message
            response: (str): LLM response
        """
        self.message_history.extend(
            [
                LLMMessage(role=Role.USER, content=message),
                LLMMessage(role=Role.ASSISTANT, content=response),
            ]
        )

    def tool_format_rules(self) -> str:
        """
        Specification of tool formatting rules
        (typically JSON-based but can be non-JSON, e.g. XMLToolMessage),
        based on the currently enabled usable `ToolMessage`s

        Returns:
            str: formatting rules
        """
        # ONLY Usable tools (i.e. LLM-generation allowed),
        usable_tool_classes: List[Type[ToolMessage]] = [
            t
            for t in list(self.llm_tools_map.values())
            if t.default_value("request") in self.llm_tools_usable
        ]

        if len(usable_tool_classes) == 0:
            return ""
        format_instructions = "\n\n".join(
            [
                msg_cls.format_instructions(tool=self.config.use_tools)
                for msg_cls in usable_tool_classes
            ]
        )
        # if any of the enabled classes has json_group_instructions, then use that,
        # else fall back to ToolMessage.json_group_instructions
        for msg_cls in usable_tool_classes:
            if hasattr(msg_cls, "json_group_instructions") and callable(
                getattr(msg_cls, "json_group_instructions")
            ):
                return msg_cls.group_format_instructions().format(
                    format_instructions=format_instructions
                )
        return ToolMessage.group_format_instructions().format(
            format_instructions=format_instructions
        )

    def tool_instructions(self) -> str:
        """
        Instructions for tools or function-calls, for enabled and usable Tools.
        These are inserted into system prompt regardless of whether we are using
        our own ToolMessage mechanism or the LLM's function-call mechanism.

        Returns:
            str: concatenation of instructions for all usable tools
        """
        enabled_classes: List[Type[ToolMessage]] = list(self.llm_tools_map.values())
        if len(enabled_classes) == 0:
            return ""
        instructions = []
        for msg_cls in enabled_classes:
            if msg_cls.default_value("request") in self.llm_tools_usable:
                class_instructions = ""
                if hasattr(msg_cls, "instructions") and inspect.ismethod(
                    msg_cls.instructions
                ):
                    class_instructions = msg_cls.instructions()
                if (
                    self.config.use_tools
                    and hasattr(msg_cls, "langroid_tools_instructions")
                    and inspect.ismethod(msg_cls.langroid_tools_instructions)
                ):
                    class_instructions += msg_cls.langroid_tools_instructions()
                # example will be shown in tool_format_rules() when using TOOLs,
                # so we don't need to show it here.
                example = "" if self.config.use_tools else (msg_cls.usage_examples())
                if example != "":
                    example = "EXAMPLES:\n" + example
                guidance = (
                    ""
                    if class_instructions == ""
                    else ("GUIDANCE: " + class_instructions)
                )
                if guidance == "" and example == "":
                    continue
                instructions.append(
                    textwrap.dedent(
                        f"""
                        TOOL: {msg_cls.default_value("request")}:
                        {guidance}
                        {example}
                        """.lstrip()
                    )
                )
        if len(instructions) == 0:
            return ""
        instructions_str = "\n\n".join(instructions)
        return textwrap.dedent(
            f"""
            === GUIDELINES ON SOME TOOLS/FUNCTIONS USAGE ===
            {instructions_str}
            """.lstrip()
        )

    def augment_system_message(self, message: str) -> None:
        """
        Augment the system message with the given message.
        Args:
            message (str): system message
        """
        self.system_message += "\n\n" + message

    def last_message_with_role(self, role: Role) -> LLMMessage | None:
        """from `message_history`, return the last message with role `role`"""
        n_role_msgs = len([m for m in self.message_history if m.role == role])
        if n_role_msgs == 0:
            return None
        idx = self.nth_message_idx_with_role(role, n_role_msgs)
        return self.message_history[idx]

    def last_message_idx_with_role(self, role: Role) -> int:
        """Index of last message in message_history, with specified role.
        Return -1 if not found. Index = 0 is the first message in the history.
        """
        indices_with_role = [
            i for i, m in enumerate(self.message_history) if m.role == role
        ]
        if len(indices_with_role) == 0:
            return -1
        return indices_with_role[-1]

    def nth_message_idx_with_role(self, role: Role, n: int) -> int:
        """Index of `n`th message in message_history, with specified role.
        (n is assumed to be 1-based, i.e. 1 is the first message with that role).
        Return -1 if not found. Index = 0 is the first message in the history.
        """
        indices_with_role = [
            i for i, m in enumerate(self.message_history) if m.role == role
        ]

        if len(indices_with_role) < n:
            return -1
        return indices_with_role[n - 1]

    def update_last_message(self, message: str, role: str = Role.USER) -> None:
        """
        Update the last message that has role `role` in the message history.
        Useful when we want to replace a long user prompt, that may contain context
        documents plus a question, with just the question.
        Args:
            message (str): new message to replace with
            role (str): role of message to replace
        """
        if len(self.message_history) == 0:
            return
        # find last message in self.message_history with role `role`
        for i in range(len(self.message_history) - 1, -1, -1):
            if self.message_history[i].role == role:
                self.message_history[i].content = message
                break

    def delete_last_message(self, role: str = Role.USER) -> None:
        """
        Delete the last message that has role `role` from the message history.
        Args:
            role (str): role of message to delete
        """
        if len(self.message_history) == 0:
            return
        # find last message in self.message_history with role `role`
        for i in range(len(self.message_history) - 1, -1, -1):
            if self.message_history[i].role == role:
                self.message_history.pop(i)
                break

    def _create_system_and_tools_message(self) -> LLMMessage:
        """
        (Re-)Create the system message for the LLM of the agent,
        taking into account any tool instructions that have been added
        after the agent was initialized.

        The system message will consist of:
        (a) the system message from the `task` arg in constructor, if any,
            otherwise the default system message from the config
        (b) the system tool instructions, if any
        (c) the system json tool instructions, if any

        Returns:
            LLMMessage object
        """
        content = self.system_message
        if self.system_tool_instructions != "":
            content += "\n\n" + self.system_tool_instructions
        if self.system_tool_format_instructions != "":
            content += "\n\n" + self.system_tool_format_instructions
        if self.output_format_instructions != "":
            content += "\n\n" + self.output_format_instructions

        # remove leading and trailing newlines and other whitespace
        return LLMMessage(role=Role.SYSTEM, content=content.strip())

    def handle_message_fallback(self, msg: str | ChatDocument) -> Any:
        """
        Fallback method for the "no-tools" scenario, i.e., the current `msg`
        (presumably emitted by the LLM) does not have any tool that the agent
        can handle.
        NOTE: The `msg` may contain tools but either (a) the agent is not
        enabled to handle them, or (b) there's an explicit `recipient` field
        in the tool that doesn't match the agent's name.

        Uses the self.config.non_tool_routing to determine the action to take.

        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
        """
        if (
            isinstance(msg, str)
            or msg.metadata.sender != Entity.LLM
            or self.config.handle_llm_no_tool is None
            or self.has_only_unhandled_tools(msg)
        ):
            return None
        # we ONLY use the `handle_llm_no_tool` config option when
        # the msg is from LLM and does not contain ANY tools at all.
        from langroid.agent.tools.orchestration import AgentDoneTool, ForwardTool

        no_tool_option = self.config.handle_llm_no_tool
        if no_tool_option in list(NonToolAction):
            # in case the `no_tool_option` is one of the special NonToolAction vals
            match self.config.handle_llm_no_tool:
                case NonToolAction.FORWARD_USER:
                    return ForwardTool(agent="User")
                case NonToolAction.DONE:
                    return AgentDoneTool(content=msg.content, tools=msg.tool_messages)
        elif is_callable(no_tool_option):
            return no_tool_option(msg)
        # Otherwise just return `no_tool_option` as is:
        # This can be any string, such as a specific nudge/reminder to the LLM,
        # or even something like ResultTool etc.
        return no_tool_option

    def unhandled_tools(self) -> set[str]:
        """The set of tools that are known but not handled.
        Useful in task flow: an agent can refuse to accept an incoming msg
        when it only has unhandled tools.
        """
        return self.llm_tools_known - self.llm_tools_handled

    def enable_message(
        self,
        message_class: Optional[Type[ToolMessage] | List[Type[ToolMessage]]],
        use: bool = True,
        handle: bool = True,
        force: bool = False,
        require_recipient: bool = False,
        include_defaults: bool = True,
    ) -> None:
        """
        Add the tool (message class) to the agent, and enable either
        - tool USE (i.e. the LLM can generate JSON to use this tool),
        - tool HANDLING (i.e. the agent can handle JSON from this tool),

        Args:
            message_class: The ToolMessage class OR List of such classes to enable,
                for USE, or HANDLING, or both.
                If this is a list of ToolMessage classes, then the remain args are
                applied to all classes.
                Optional; if None, then apply the enabling to all tools in the
                agent's toolset that have been enabled so far.
            use: IF True, allow the agent (LLM) to use this tool (or all tools),
                else disallow
            handle: if True, allow the agent (LLM) to handle (i.e. respond to) this
                tool (or all tools)
            force: whether to FORCE the agent (LLM) to USE the specific
                 tool represented by `message_class`.
                 `force` is ignored if `message_class` is None.
            require_recipient: whether to require that recipient be specified
                when using the tool message (only applies if `use` is True).
            include_defaults: whether to include fields that have default values,
                in the "properties" section of the JSON format instructions.
                (Normally the OpenAI completion API ignores these fields,
                but the Assistant fn-calling seems to pay attn to these,
                and if we don't want this, we should set this to False.)
        """
        if message_class is not None and isinstance(message_class, list):
            for mc in message_class:
                self.enable_message(
                    mc,
                    use=use,
                    handle=handle,
                    force=force,
                    require_recipient=require_recipient,
                    include_defaults=include_defaults,
                )
            return None
        if require_recipient and message_class is not None:
            message_class = message_class.require_recipient()
        if isinstance(message_class, XMLToolMessage):
            # XMLToolMessage is not compatible with OpenAI's Tools/functions API,
            # so we disable use of functions API, enable langroid-native Tools,
            # which are prompt-based.
            self.config.use_functions_api = False
            self.config.use_tools = True
        super().enable_message_handling(message_class)  # enables handling only
        tools = self._get_tool_list(message_class)
        if message_class is not None:
            request = message_class.default_value("request")
            if request == "":
                raise ValueError(
                    f"""
                    ToolMessage class {message_class} must have a non-empty
                    'request' field if it is to be enabled as a tool.
                    """
                )
            llm_function = message_class.llm_function_schema(defaults=include_defaults)
            self.llm_functions_map[request] = llm_function
            if force:
                self.llm_function_force = dict(name=request)
            else:
                self.llm_function_force = None

        for t in tools:
            self.llm_tools_known.add(t)

            if handle:
                self.llm_tools_handled.add(t)
                self.llm_functions_handled.add(t)

                if (
                    self.enabled_handling_output_format is not None
                    and self.enabled_handling_output_format.name() == t
                ):
                    # `t` was designated as "enabled for handling" ONLY for
                    # output_format enforcement, but we are explicitly ]
                    # enabling it for handling here, so we set the variable to None.
                    self.enabled_handling_output_format = None
            else:
                self.llm_tools_handled.discard(t)
                self.llm_functions_handled.discard(t)

            if use:
                tool_class = self.llm_tools_map[t]
                allow_llm_use = tool_class._allow_llm_use
                if isinstance(allow_llm_use, ModelPrivateAttr):
                    allow_llm_use = allow_llm_use.default
                if allow_llm_use:
                    self.llm_tools_usable.add(t)
                    self.llm_functions_usable.add(t)
                else:
                    logger.warning(
                        f"""
                        ToolMessage class {tool_class} does not allow LLM use,
                        because `_allow_llm_use=False` either in the Tool or a
                        parent class of this tool;
                        so not enabling LLM use for this tool!
                        If you intended an LLM to use this tool,
                        set `_allow_llm_use=True` when you define the tool.
                        """
                    )
                if (
                    self.enabled_use_output_format is not None
                    and self.enabled_use_output_format.default_value("request") == t
                ):
                    # `t` was designated as "enabled for use" ONLY for output_format
                    # enforcement, but we are explicitly enabling it for use here,
                    # so we set the variable to None.
                    self.enabled_use_output_format = None
            else:
                self.llm_tools_usable.discard(t)
                self.llm_functions_usable.discard(t)

        self._update_tool_instructions()

    def _update_tool_instructions(self) -> None:
        # Set tool instructions and JSON format instructions,
        # in case Tools have been enabled/disabled.
        if self.config.use_tools:
            self.system_tool_format_instructions = self.tool_format_rules()
        self.system_tool_instructions = self.tool_instructions()

    def _requests_and_tool_settings(self) -> tuple[Optional[set[str]], bool, bool]:
        """
        Returns the current set of enabled requests for inference and tools configs.
        Used for restoring setings overriden by `set_output_format`.
        """
        return (
            self.enabled_requests_for_inference,
            self.config.use_functions_api,
            self.config.use_tools,
        )

    @property
    def all_llm_tools_known(self) -> set[str]:
        """All known tools; we include `output_format` if it is a `ToolMessage`."""
        known = self.llm_tools_known

        if self.output_format is not None and issubclass(
            self.output_format, ToolMessage
        ):
            return known.union({self.output_format.default_value("request")})

        return known

    def set_output_format(
        self,
        output_type: Optional[type],
        force_tools: Optional[bool] = None,
        use: Optional[bool] = None,
        handle: Optional[bool] = None,
        instructions: Optional[bool] = None,
        is_copy: bool = False,
    ) -> None:
        """
        Sets `output_format` to `output_type` and, if `force_tools` is enabled,
        switches to the native Langroid tools mechanism to ensure that no tool
        calls not of `output_type` are generated. By default, `force_tools`
        follows the `use_tools_on_output_format` parameter in the config.

        If `output_type` is None, restores to the state prior to setting
        `output_format`.

        If `use`, we enable use of `output_type` when it is a subclass
        of `ToolMesage`. Note that this primarily controls instruction
        generation: the model will always generate `output_type` regardless
        of whether `use` is set. Defaults to the `use_output_format`
        parameter in the config. Similarly, handling of `output_type` is
        controlled by `handle`, which defaults to the
        `handle_output_format` parameter in the config.

        `instructions` controls whether we generate instructions specifying
        the output format schema. Defaults to the `instructions_output_format`
        parameter in the config.

        `is_copy` is set when called via `__getitem__`. In that case, we must
        copy certain fields to ensure that we do not overwrite the main agent's
        setings.
        """
        # Disable usage of an output format which was not specifically enabled
        # by `enable_message`
        if self.enabled_use_output_format is not None:
            self.disable_message_use(self.enabled_use_output_format)
            self.enabled_use_output_format = None

        # Disable handling of an output format which did not specifically have
        # handling enabled via `enable_message`
        if self.enabled_handling_output_format is not None:
            self.disable_message_handling(self.enabled_handling_output_format)
            self.enabled_handling_output_format = None

        # Reset any previous instructions
        self.output_format_instructions = ""

        if output_type is None:
            self.output_format = None
            (
                requests_for_inference,
                use_functions_api,
                use_tools,
            ) = self.saved_requests_and_tool_setings
            self.config = self.config.model_copy()
            self.enabled_requests_for_inference = requests_for_inference
            self.config.use_functions_api = use_functions_api
            self.config.use_tools = use_tools
        else:
            if force_tools is None:
                force_tools = self.config.use_tools_on_output_format

            if not any(
                (isclass(output_type) and issubclass(output_type, t))
                for t in [ToolMessage, BaseModel]
            ):
                output_type = get_pydantic_wrapper(output_type)

            if self.output_format is None and force_tools:
                self.saved_requests_and_tool_setings = (
                    self._requests_and_tool_settings()
                )

            self.output_format = output_type
            if issubclass(output_type, ToolMessage):
                name = output_type.default_value("request")
                if use is None:
                    use = self.config.use_output_format

                if handle is None:
                    handle = self.config.handle_output_format

                if use or handle:
                    is_usable = name in self.llm_tools_usable.union(
                        self.llm_functions_usable
                    )
                    is_handled = name in self.llm_tools_handled.union(
                        self.llm_functions_handled
                    )

                    if is_copy:
                        if use:
                            # We must copy `llm_tools_usable` so the base agent
                            # is unmodified
                            self.llm_tools_usable = self.llm_tools_usable.copy()
                            self.llm_functions_usable = self.llm_functions_usable.copy()
                        if handle:
                            # If handling the tool, do the same for `llm_tools_handled`
                            self.llm_tools_handled = self.llm_tools_handled.copy()
                            self.llm_functions_handled = (
                                self.llm_functions_handled.copy()
                            )
                    # Enable `output_type`
                    self.enable_message(
                        output_type,
                        # Do not override existing settings
                        use=use or is_usable,
                        handle=handle or is_handled,
                    )

                    # If the `output_type` ToilMessage was not already enabled for
                    # use, this means we are ONLY enabling it for use specifically
                    # for enforcing this output format, so we set the
                    # `enabled_use_output_forma  to this output_type, to
                    # record that it should be disabled when `output_format` is changed
                    if not is_usable:
                        self.enabled_use_output_format = output_type

                    # (same reasoning as for use-enabling)
                    if not is_handled:
                        self.enabled_handling_output_format = output_type

                generated_tool_instructions = name in self.llm_tools_usable.union(
                    self.llm_functions_usable
                )
            else:
                generated_tool_instructions = False

            if instructions is None:
                instructions = self.config.instructions_output_format
            if issubclass(output_type, BaseModel) and instructions:
                if generated_tool_instructions:
                    # Already generated tool instructions as part of "enabling for use",
                    # so only need to generate a reminder to use this tool.
                    name = cast(ToolMessage, output_type).default_value("request")
                    self.output_format_instructions = textwrap.dedent(
                        f"""
                        === OUTPUT FORMAT INSTRUCTIONS ===

                        Please provide output using the `{name}` tool/function.
                        """
                    )
                else:
                    if issubclass(output_type, ToolMessage):
                        output_format_schema = output_type.llm_function_schema(
                            request=True,
                            defaults=self.config.output_format_include_defaults,
                        ).parameters
                    else:
                        output_format_schema = output_type.model_json_schema()

                    format_schema_for_strict(output_format_schema)

                    self.output_format_instructions = textwrap.dedent(
                        f"""
                        === OUTPUT FORMAT INSTRUCTIONS ===
                        Please provide output as JSON with the following schema:

                        {output_format_schema}
                        """
                    )

            if force_tools:
                if issubclass(output_type, ToolMessage):
                    self.enabled_requests_for_inference = {
                        output_type.default_value("request")
                    }
                if self.config.use_functions_api:
                    self.config = self.config.model_copy()
                    self.config.use_functions_api = False
                    self.config.use_tools = True

    def __getitem__(self, output_type: type) -> Self:
        """
        Returns a (shallow) copy of `self` with a forced output type.
        """
        clone = copy.copy(self)
        clone.set_output_format(output_type, is_copy=True)
        return clone

    def disable_message_handling(
        self,
        message_class: Optional[Type[ToolMessage]] = None,
    ) -> None:
        """
        Disable this agent from RESPONDING to a `message_class` (Tool). If
            `message_class` is None, then disable this agent from responding to ALL.
        Args:
            message_class: The ToolMessage class to disable; Optional.
        """
        super().disable_message_handling(message_class)
        for t in self._get_tool_list(message_class):
            self.llm_tools_handled.discard(t)
            self.llm_functions_handled.discard(t)

    def disable_message_use(
        self,
        message_class: Optional[Type[ToolMessage]],
    ) -> None:
        """
        Disable this agent from USING a message class (Tool).
        If `message_class` is None, then disable this agent from USING ALL tools.
        Args:
            message_class: The ToolMessage class to disable.
                If None, disable all.
        """
        for t in self._get_tool_list(message_class):
            self.llm_tools_usable.discard(t)
            self.llm_functions_usable.discard(t)

        self._update_tool_instructions()

    def disable_message_use_except(self, message_class: Type[ToolMessage]) -> None:
        """
        Disable this agent from USING ALL messages EXCEPT a message class (Tool)
        Args:
            message_class: The only ToolMessage class to allow
        """
        request = message_class.model_fields["request"].default
        to_remove = [r for r in self.llm_tools_usable if r != request]
        for r in to_remove:
            self.llm_tools_usable.discard(r)
            self.llm_functions_usable.discard(r)
        self._update_tool_instructions()

    def _load_output_format(self, message: ChatDocument) -> None:
        """
        If set, attempts to parse a value of type `self.output_format` from the message
        contents or any tool/function call and assigns it to `content_any`.
        """
        if self.output_format is not None:
            any_succeeded = False
            attempts: list[str | LLMFunctionCall] = [
                message.content,
            ]

            if message.function_call is not None:
                attempts.append(message.function_call)

            if message.oai_tool_calls is not None:
                attempts.extend(
                    [
                        c.function
                        for c in message.oai_tool_calls
                        if c.function is not None
                    ]
                )

            for attempt in attempts:
                try:
                    if isinstance(attempt, str):
                        content = json.loads(attempt)
                    else:
                        if not (
                            issubclass(self.output_format, ToolMessage)
                            and attempt.name
                            == self.output_format.default_value("request")
                        ):
                            continue

                        content = attempt.arguments

                    content_any = self.output_format.model_validate(content)

                    if issubclass(self.output_format, PydanticWrapper):
                        message.content_any = content_any.value  # type: ignore
                    else:
                        message.content_any = content_any
                    any_succeeded = True
                    break
                except (ValidationError, json.JSONDecodeError):
                    continue

            if not any_succeeded:
                self.disable_strict = True
                logging.warning(
                    """
                    Validation error occured with strict output format enabled.
                    Disabling strict mode.
                    """
                )

    def get_tool_messages(
        self,
        msg: str | ChatDocument | None,
        all_tools: bool = False,
    ) -> List[ToolMessage]:
        """
        Extracts messages and tracks whether any errors occurred. If strict mode
        was enabled, disables it for the tool, else triggers strict recovery.
        """
        self.tool_error = False
        most_recent_sent_by_llm = (
            len(self.message_history) > 0
            and self.message_history[-1].role == Role.ASSISTANT
        )
        was_llm = most_recent_sent_by_llm or (
            isinstance(msg, ChatDocument) and msg.metadata.sender == Entity.LLM
        )
        try:
            tools = super().get_tool_messages(msg, all_tools)
        except ValidationError as ve:
            # Check if tool class was attached to the exception
            if hasattr(ve, "tool_class") and ve.tool_class:
                tool_class = ve.tool_class  # type: ignore
                if issubclass(tool_class, ToolMessage):
                    was_strict = (
                        self.config.use_functions_api
                        and self.config.use_tools_api
                        and self._strict_mode_for_tool(tool_class)
                    )
                    # If the result of strict output for a tool using the
                    # OpenAI tools API fails to parse, we infer that the
                    # schema edits necessary for compatibility prevented
                    # adherence to the underlying `ToolMessage` schema and
                    # disable strict output for the tool
                    if was_strict:
                        name = tool_class.default_value("request")
                        self.disable_strict_tools_set.add(name)
                        logging.warning(
                            f"""
                            Validation error occured with strict tool format.
                            Disabling strict mode for the {name} tool.
                            """
                        )
                    else:
                        # We will trigger the strict recovery mechanism to force
                        # the LLM to correct its output, allowing us to parse
                        if isinstance(msg, ChatDocument):
                            self.tool_error = msg.metadata.sender == Entity.LLM
                        else:
                            self.tool_error = most_recent_sent_by_llm

            if was_llm:
                raise ve
            else:
                self.tool_error = False
                return []

        if not was_llm:
            self.tool_error = False

        return tools

    def _get_any_tool_message(self, optional: bool = True) -> type[ToolMessage] | None:
        """
        Returns a `ToolMessage` which wraps all enabled tools, excluding those
        where strict recovery is disabled. Used in strict recovery.
        """
        possible_tools = tuple(
            self.llm_tools_map[t]
            for t in self.llm_tools_usable
            if t not in self.disable_strict_tools_set
        )
        if len(possible_tools) == 0:
            return None
        any_tool_type = Union.__getitem__(possible_tools)  # type ignore

        maybe_optional_type = Optional[any_tool_type] if optional else any_tool_type

        class AnyTool(ToolMessage):
            purpose: str = "To call a tool/function."
            request: str = "tool_or_function"
            tool: maybe_optional_type  # type: ignore

            def response(self, agent: ChatAgent) -> None | str | ChatDocument:
                # One-time use
                agent.set_output_format(None)

                if self.tool is None:
                    return None

                # As the ToolMessage schema accepts invalid
                # `tool.request` values, reparse with the
                # corresponding tool
                request = self.tool.request
                if request not in agent.llm_tools_map:
                    return None
                tool = agent.llm_tools_map[request].model_validate_json(
                    self.tool.to_json()
                )

                return agent.handle_tool_message(tool)

            async def response_async(
                self, agent: ChatAgent
            ) -> None | str | ChatDocument:
                # One-time use
                agent.set_output_format(None)

                if self.tool is None:
                    return None

                # As the ToolMessage schema accepts invalid
                # `tool.request` values, reparse with the
                # corresponding tool
                request = self.tool.request
                if request not in agent.llm_tools_map:
                    return None
                tool = agent.llm_tools_map[request].model_validate_json(
                    self.tool.to_json()
                )

                return await agent.handle_tool_message_async(tool)

        return AnyTool

    def _strict_recovery_instructions(
        self,
        tool_type: Optional[type[ToolMessage]] = None,
        optional: bool = True,
    ) -> str:
        """Returns instructions for strict recovery."""
        optional_instructions = (
            (
                "\n"
                + """
        If you did NOT intend to do so, `tool` should be null.
        """
            )
            if optional
            else ""
        )
        response_prefix = "If you intended to make such a call, r" if optional else "R"
        instruction_prefix = "If you do so, b" if optional else "B"

        schema_instructions = (
            f"""
        The schema for `tool_or_function` is as follows:
        {tool_type.llm_function_schema(defaults=True, request=True).parameters}
        """
            if tool_type
            else ""
        )

        return textwrap.dedent(
            f"""
        Your previous attempt to make a tool/function call appears to have failed.
        {response_prefix}espond with your desired tool/function. Do so with the
        `tool_or_function` tool/function where `tool` is set to your intended call.
        {schema_instructions}

        {instruction_prefix}e sure that your corrected call matches your intention
        in your previous request. For any field with a default value which
        you did not intend to override in your previous attempt, be sure
        to set that field to its default value. {optional_instructions}
        """
        )

    def truncate_message(
        self,
        idx: int,
        tokens: int = 5,
        warning: str = "...[Contents truncated!]",
        inplace: bool = True,
    ) -> LLMMessage:
        """
        Truncate message at idx in msg history to `tokens` tokens.

        If inplace is True, the message is truncated in place, else
        it LEAVES the original message INTACT and returns a new message
        """
        if inplace:
            llm_msg = self.message_history[idx]
        else:
            llm_msg = copy.deepcopy(self.message_history[idx])
        orig_content = llm_msg.content
        new_content = (
            self.parser.truncate_tokens(orig_content, tokens)
            if self.parser is not None
            else orig_content[: tokens * 4]  # approx truncation
        )
        llm_msg.content = new_content + "\n" + warning
        return llm_msg

    def _reduce_raw_tool_results(self, message: ChatDocument) -> None:
        """
        If message is the result of a ToolMessage that had
        a `_max_retained_tokens` set to a non-None value, then we replace contents
        with a placeholder message.
        """
        parent_message: ChatDocument | None = message.parent
        tools = [] if parent_message is None else parent_message.tool_messages
        truncate_tools = []
        for t in tools:
            max_retained_tokens = t._max_retained_tokens
            if isinstance(max_retained_tokens, ModelPrivateAttr):
                max_retained_tokens = max_retained_tokens.default
            if max_retained_tokens is not None:
                truncate_tools.append(t)
        limiting_tool = truncate_tools[0] if len(truncate_tools) > 0 else None
        if limiting_tool is not None:
            max_retained_tokens = limiting_tool._max_retained_tokens
            if isinstance(max_retained_tokens, ModelPrivateAttr):
                max_retained_tokens = max_retained_tokens.default
            if max_retained_tokens is not None:
                tool_name = limiting_tool.default_value("request")
                max_tokens: int = max_retained_tokens
                truncation_warning = f"""
                    The result of the {tool_name} tool were too large,
                    and has been truncated to {max_tokens} tokens.
                    To obtain the full result, the tool needs to be re-used.
                """
                self.truncate_message(
                    message.metadata.msg_idx, max_tokens, truncation_warning
                )

    def llm_response(
        self, message: Optional[str | ChatDocument] = None
    ) -> Optional[ChatDocument]:
        """
        Respond to a single user message, appended to the message history,
        in "chat" mode
        Args:
            message (str|ChatDocument): message or ChatDocument object to respond to.
                If None, use the self.task_messages
        Returns:
            LLM response as a ChatDocument object
        """
        if self.llm is None:
            return None

        # If enabled and a tool error occurred, we recover by generating the tool in
        # strict json mode
        if (
            self.tool_error
            and self.output_format is None
            and self._json_schema_available()
            and self.config.strict_recovery
        ):
            self.tool_error = False
            AnyTool = self._get_any_tool_message()
            if AnyTool is None:
                return None
            self.set_output_format(
                AnyTool,
                force_tools=True,
                use=True,
                handle=True,
                instructions=True,
            )
            recovery_message = self._strict_recovery_instructions(AnyTool)
            augmented_message = message
            if augmented_message is None:
                augmented_message = recovery_message
            elif isinstance(augmented_message, str):
                augmented_message = augmented_message + recovery_message
            else:
                augmented_message.content = augmented_message.content + recovery_message

            # only use the augmented message for this one response...
            result = self.llm_response(augmented_message)
            # ... restore the original user message so that the AnyTool recover
            # instructions don't persist in the message history
            # (this can cause the LLM to use the AnyTool directly as a tool)
            if message is None:
                self.delete_last_message(role=Role.USER)
            else:
                msg = message if isinstance(message, str) else message.content
                self.update_last_message(msg, role=Role.USER)
            return result

        hist, output_len = self._prep_llm_messages(message)
        if len(hist) == 0:
            return None
        tool_choice = (
            "auto"
            if isinstance(message, str)
            else (message.oai_tool_choice if message is not None else "auto")
        )
        with StreamingIfAllowed(self.llm, self.llm.get_stream()):
            try:
                response = self.llm_response_messages(hist, output_len, tool_choice)
            except openai.BadRequestError as e:
                if self.any_strict:
                    self.disable_strict = True
                    self.set_output_format(None)
                    logging.warning(
                        f"""
                        OpenAI BadRequestError raised with strict mode enabled.
                        Message: {e.message}
                        Disabling strict mode and retrying.
                        """
                    )
                    return self.llm_response(message)
                else:
                    raise e
        self.message_history.extend(ChatDocument.to_LLMMessage(response))
        response.metadata.msg_idx = len(self.message_history) - 1
        response.metadata.agent_id = self.id
        if isinstance(message, ChatDocument):
            self._reduce_raw_tool_results(message)
        # Preserve trail of tool_ids for OpenAI Assistant fn-calls
        response.metadata.tool_ids = (
            []
            if isinstance(message, str)
            else message.metadata.tool_ids if message is not None else []
        )

        return response

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

        # If enabled and a tool error occurred, we recover by generating the tool in
        # strict json mode
        if (
            self.tool_error
            and self.output_format is None
            and self._json_schema_available()
            and self.config.strict_recovery
        ):
            self.tool_error = False
            AnyTool = self._get_any_tool_message()
            self.set_output_format(
                AnyTool,
                force_tools=True,
                use=True,
                handle=True,
                instructions=True,
            )
            recovery_message = self._strict_recovery_instructions(AnyTool)
            augmented_message = message
            if augmented_message is None:
                augmented_message = recovery_message
            elif isinstance(augmented_message, str):
                augmented_message = augmented_message + recovery_message
            else:
                augmented_message.content = augmented_message.content + recovery_message

            # only use the augmented message for this one response...
            result = self.llm_response(augmented_message)
            # ... restore the original user message so that the AnyTool recover
            # instructions don't persist in the message history
            # (this can cause the LLM to use the AnyTool directly as a tool)
            if message is None:
                self.delete_last_message(role=Role.USER)
            else:
                msg = message if isinstance(message, str) else message.content
                self.update_last_message(msg, role=Role.USER)
            return result

        hist, output_len = self._prep_llm_messages(message)
        if len(hist) == 0:
            return None
        tool_choice = (
            "auto"
            if isinstance(message, str)
            else (message.oai_tool_choice if message is not None else "auto")
        )
        with StreamingIfAllowed(self.llm, self.llm.get_stream()):
            try:
                response = await self.llm_response_messages_async(
                    hist, output_len, tool_choice
                )
            except openai.BadRequestError as e:
                if self.any_strict:
                    self.disable_strict = True
                    self.set_output_format(None)
                    logging.warning(
                        f"""
                        OpenAI BadRequestError raised with strict mode enabled.
                        Message: {e.message}
                        Disabling strict mode and retrying.
                        """
                    )
                    return await self.llm_response_async(message)
                else:
                    raise e
        self.message_history.extend(ChatDocument.to_LLMMessage(response))
        response.metadata.msg_idx = len(self.message_history) - 1
        response.metadata.agent_id = self.id
        if isinstance(message, ChatDocument):
            self._reduce_raw_tool_results(message)
        # Preserve trail of tool_ids for OpenAI Assistant fn-calls
        response.metadata.tool_ids = (
            []
            if isinstance(message, str)
            else message.metadata.tool_ids if message is not None else []
        )

        return response

    def init_message_history(self) -> None:
        """
        Initialize the message history with the system message and user message
        """
        self.message_history = [self._create_system_and_tools_message()]
        if self.user_message:
            self.message_history.append(
                LLMMessage(role=Role.USER, content=self.user_message)
            )

    def _prep_llm_messages(
        self,
        message: Optional[str | ChatDocument] = None,
        truncate: bool = True,
    ) -> Tuple[List[LLMMessage], int]:
        """
        Prepare messages to be sent to self.llm_response_messages,
            which is the main method that calls the LLM API to get a response.
            If desired output tokens + message history exceeds the model context length,
            then first the max output tokens is reduced to fit, and if that is not
            possible, older messages may be truncated to accommodate at least
            self.config.llm.min_output_tokens of output.

        Returns:
            Tuple[List[LLMMessage], int]: (messages, output_len)
                messages = Full list of messages to send
                output_len = max expected number of tokens in response
        """

        if (
            not self.llm_can_respond(message)
            or self.config.llm is None
            or self.llm is None
        ):
            return [], 0

        if message is None and len(self.message_history) > 0:
            # this means agent has been used to get LLM response already,
            # and so the last message is an "assistant" response.
            # We delete this last assistant response and re-generate it.
            self.clear_history(-1)
            logger.warning(
                "Re-generating the last assistant response since message is None"
            )

        if len(self.message_history) == 0:
            # initial messages have not yet been loaded, so load them
            self.init_message_history()

            # for debugging, show the initial message history
            if settings.debug:
                print(
                    f"""
                [grey37]LLM Initial Msg History:
                {escape(self.message_history_str())}
                [/grey37]
                """
                )
        else:
            assert self.message_history[0].role == Role.SYSTEM
            # update the system message with the latest tool instructions
            self.message_history[0] = self._create_system_and_tools_message()

        if message is not None:
            if (
                isinstance(message, str)
                or message.id() != self.message_history[-1].chat_document_id
            ):
                # either the message is a str, or it is a fresh ChatDocument
                # different from the last message in the history
                llm_msgs = ChatDocument.to_LLMMessage(message, self.oai_tool_calls)
                # LLM only responds to the content, so only those msgs with
                # non-empty content should be kept
                llm_msgs = [m for m in llm_msgs if m.content.strip() != ""]
                if len(llm_msgs) == 0:
                    return [], 0
                # process tools if any
                done_tools = [m.tool_call_id for m in llm_msgs if m.role == Role.TOOL]
                self.oai_tool_calls = [
                    t for t in self.oai_tool_calls if t.id not in done_tools
                ]
                self.message_history.extend(llm_msgs)

        hist = self.message_history
        output_len = self.config.llm.model_max_output_tokens
        if (
            truncate
            and output_len > self.llm.chat_context_length() - self.chat_num_tokens(hist)
        ):
            CHAT_HISTORY_BUFFER = 300
            # chat + output > max context length,
            # so first try to shorten requested output len to fit;
            # use an extra margin of CHAT_HISTORY_BUFFER tokens
            # in case our calcs are off (and to allow for some extra tokens)
            output_len = (
                self.llm.chat_context_length()
                - self.chat_num_tokens(hist)
                - CHAT_HISTORY_BUFFER
            )
            if output_len > self.config.llm.min_output_tokens:
                logger.debug(
                    f"""
                    Chat Model context length is {self.llm.chat_context_length()},
                    but the current message history is {self.chat_num_tokens(hist)}
                    tokens long, which does not allow
                    {self.config.llm.model_max_output_tokens} output tokens.
                    Therefore we reduced `max_output_tokens` to {output_len} tokens,
                    so they can fit within the model's context length
                    """
                )
            else:
                # unacceptably small output len, so compress early parts of conv
                # history if output_len is still too long.
                # TODO we should really be doing summarization or other types of
                #   prompt-size reduction
                msg_idx_to_compress = 1  # don't touch system msg
                # we will try compressing msg indices up to but not including
                # last user msg
                last_msg_idx_to_compress = (
                    self.last_message_idx_with_role(
                        role=Role.USER,
                    )
                    - 1
                )
                n_truncated = 0
                while (
                    self.chat_num_tokens(hist)
                    > self.llm.chat_context_length() - self.config.llm.min_output_tokens
                ):
                    # try dropping early parts of conv history
                    # TODO we should really be doing summarization or other types of
                    #   prompt-size reduction
                    if msg_idx_to_compress > last_msg_idx_to_compress:
                        # We want to preserve the first message (typically system msg)
                        # and last message (user msg).
                        raise ValueError(
                            """
                        The (message history + max_output_tokens) is longer than the
                        max chat context length of this model, and we have tried
                        reducing the requested max output tokens, as well as truncating
                        early parts of the message history, to accommodate the model
                        context length, but we have run out of msgs to drop.

                        HINT: In the `llm` field of your `ChatAgentConfig` object,
                        which is of type `LLMConfig/OpenAIGPTConfig`, try
                        - increasing `chat_context_length`
                            (if accurate for the model), or
                        - decreasing `max_output_tokens`
                        """
                        )
                    n_truncated += 1
                    # compress the msg at idx `msg_idx_to_compress`
                    hist[msg_idx_to_compress] = self.truncate_message(
                        msg_idx_to_compress,
                        tokens=30,
                        warning="... [Contents truncated!]",
                    )

                    msg_idx_to_compress += 1

                output_len = min(
                    self.config.llm.model_max_output_tokens,
                    self.llm.chat_context_length()
                    - self.chat_num_tokens(hist)
                    - CHAT_HISTORY_BUFFER,
                )
                if output_len < self.config.llm.min_output_tokens:
                    raise ValueError(
                        f"""
                        Tried to shorten prompt history for chat mode
                        but even after truncating all messages except system msg and
                        last (user) msg,
                        the history token len {self.chat_num_tokens(hist)} is
                        too long to accommodate the desired minimum output tokens
                        {self.config.llm.min_output_tokens} within the
                        model's context length {self.llm.chat_context_length()}.
                        Please try shortening the system msg or user prompts,
                        or adjust `config.llm.min_output_tokens` to be smaller.
                        """
                    )
                else:
                    # we MUST have truncated at least one msg
                    msg_tokens = self.chat_num_tokens()
                    logger.warning(
                        f"""
                    Chat Model context length is {self.llm.chat_context_length()}
                    tokens, but the current message history is {msg_tokens} tokens long,
                    which does not allow {self.config.llm.model_max_output_tokens}
                    output tokens.
                    Therefore we truncated the first {n_truncated} messages
                    in the conversation history so that history token
                    length is reduced to {self.chat_num_tokens(hist)}, and
                    we use `max_output_tokens = {output_len}`,
                    so they can fit within the model's context length
                    of {self.llm.chat_context_length()} tokens.
                    """
                    )

        if isinstance(message, ChatDocument):
            # record the position of the corresponding LLMMessage in
            # the message_history
            message.metadata.msg_idx = len(hist) - 1
            message.metadata.agent_id = self.id
        return hist, output_len

    def _function_args(
        self,
    ) -> Tuple[
        Optional[List[LLMFunctionSpec]],
        str | Dict[str, str],
        Optional[List[OpenAIToolSpec]],
        Optional[Dict[str, Dict[str, str] | str]],
        Optional[OpenAIJsonSchemaSpec],
    ]:
        """
        Get function/tool spec/output format arguments for
        OpenAI-compatible LLM API call
        """
        functions: Optional[List[LLMFunctionSpec]] = None
        fun_call: str | Dict[str, str] = "none"
        tools: Optional[List[OpenAIToolSpec]] = None
        force_tool: Optional[Dict[str, Dict[str, str] | str]] = None
        self.any_strict = False
        if self.config.use_functions_api and len(self.llm_functions_usable) > 0:
            if not self.config.use_tools_api:
                functions = [
                    self.llm_functions_map[f] for f in self.llm_functions_usable
                ]
                fun_call = (
                    "auto"
                    if self.llm_function_force is None
                    else self.llm_function_force
                )
            else:

                def to_maybe_strict_spec(function: str) -> OpenAIToolSpec:
                    spec = self.llm_functions_map[function]
                    strict = self._strict_mode_for_tool(function)
                    if strict:
                        self.any_strict = True
                        strict_spec = copy.deepcopy(spec)
                        format_schema_for_strict(strict_spec.parameters)
                    else:
                        strict_spec = spec

                    return OpenAIToolSpec(
                        type="function",
                        strict=strict,
                        function=strict_spec,
                    )

                tools = [to_maybe_strict_spec(f) for f in self.llm_functions_usable]
                force_tool = (
                    None
                    if self.llm_function_force is None
                    else {
                        "type": "function",
                        "function": {"name": self.llm_function_force["name"]},
                    }
                )
        output_format = None
        if self.output_format is not None and self._json_schema_available():
            self.any_strict = True
            if issubclass(self.output_format, ToolMessage) and not issubclass(
                self.output_format, XMLToolMessage
            ):
                spec = self.output_format.llm_function_schema(
                    request=True,
                    defaults=self.config.output_format_include_defaults,
                )
                format_schema_for_strict(spec.parameters)

                output_format = OpenAIJsonSchemaSpec(
                    # We always require that outputs strictly match the schema
                    strict=True,
                    function=spec,
                )
            elif issubclass(self.output_format, BaseModel):
                param_spec = self.output_format.model_json_schema()
                format_schema_for_strict(param_spec)

                output_format = OpenAIJsonSchemaSpec(
                    # We always require that outputs strictly match the schema
                    strict=True,
                    function=LLMFunctionSpec(
                        name="json_output",
                        description="Strict Json output format.",
                        parameters=param_spec,
                    ),
                )

        return functions, fun_call, tools, force_tool, output_format

    def llm_response_messages(
        self,
        messages: List[LLMMessage],
        output_len: Optional[int] = None,
        tool_choice: ToolChoiceTypes | Dict[str, str | Dict[str, str]] = "auto",
    ) -> ChatDocument:
        """
        Respond to a series of messages, e.g. with OpenAI ChatCompletion
        Args:
            messages: seq of messages (with role, content fields) sent to LLM
            output_len: max number of tokens expected in response.
                    If None, use the LLM's default model_max_output_tokens.
        Returns:
            Document (i.e. with fields "content", "metadata")
        """
        assert self.config.llm is not None and self.llm is not None
        output_len = output_len or self.config.llm.model_max_output_tokens
        streamer = noop_fn
        if self.llm.get_stream():
            streamer = self.callbacks.start_llm_stream()
        self.llm.config.streamer = streamer
        with ExitStack() as stack:  # for conditionally using rich spinner
            if not self.llm.get_stream() and not settings.quiet:
                # show rich spinner only if not streaming!
                # (Why? b/c the intent of showing a spinner is to "show progress",
                # and we don't need to do that when streaming, since
                # streaming output already shows progress.)
                cm = status(
                    "LLM responding to messages...",
                    log_if_quiet=False,
                )
                stack.enter_context(cm)
            if self.llm.get_stream() and not settings.quiet:
                console.print(f"[green]{self.indent}", end="")
            functions, fun_call, tools, force_tool, output_format = (
                self._function_args()
            )
            assert self.llm is not None
            response = self.llm.chat(
                messages,
                output_len,
                tools=tools,
                tool_choice=force_tool or tool_choice,
                functions=functions,
                function_call=fun_call,
                response_format=output_format,
            )
        if self.llm.get_stream():
            self.callbacks.finish_llm_stream(
                content=str(response),
                is_tool=self.has_tool_message_attempt(
                    ChatDocument.from_LLMResponse(response, displayed=True),
                ),
            )
        self.llm.config.streamer = noop_fn
        if response.cached:
            self.callbacks.cancel_llm_stream()
        self._render_llm_response(response)
        self.update_token_usage(
            response,  # .usage attrib is updated!
            messages,
            self.llm.get_stream(),
            chat=True,
            print_response_stats=self.config.show_stats and not settings.quiet,
        )
        chat_doc = ChatDocument.from_LLMResponse(response, displayed=True)
        self.oai_tool_calls = response.oai_tool_calls or []
        self.oai_tool_id2call.update(
            {t.id: t for t in self.oai_tool_calls if t.id is not None}
        )

        # If using strict output format, parse the output JSON
        self._load_output_format(chat_doc)

        return chat_doc

    async def llm_response_messages_async(
        self,
        messages: List[LLMMessage],
        output_len: Optional[int] = None,
        tool_choice: ToolChoiceTypes | Dict[str, str | Dict[str, str]] = "auto",
    ) -> ChatDocument:
        """
        Async version of `llm_response_messages`. See there for details.
        """
        assert self.config.llm is not None and self.llm is not None
        output_len = output_len or self.config.llm.model_max_output_tokens
        functions, fun_call, tools, force_tool, output_format = self._function_args()
        assert self.llm is not None

        streamer_async = async_noop_fn
        if self.llm.get_stream():
            streamer_async = await self.callbacks.start_llm_stream_async()
        self.llm.config.streamer_async = streamer_async

        response = await self.llm.achat(
            messages,
            output_len,
            tools=tools,
            tool_choice=force_tool or tool_choice,
            functions=functions,
            function_call=fun_call,
            response_format=output_format,
        )
        if self.llm.get_stream():
            self.callbacks.finish_llm_stream(
                content=str(response),
                is_tool=self.has_tool_message_attempt(
                    ChatDocument.from_LLMResponse(response, displayed=True),
                ),
            )
        self.llm.config.streamer_async = async_noop_fn
        if response.cached:
            self.callbacks.cancel_llm_stream()
        self._render_llm_response(response)
        self.update_token_usage(
            response,  # .usage attrib is updated!
            messages,
            self.llm.get_stream(),
            chat=True,
            print_response_stats=self.config.show_stats and not settings.quiet,
        )
        chat_doc = ChatDocument.from_LLMResponse(response, displayed=True)
        self.oai_tool_calls = response.oai_tool_calls or []
        self.oai_tool_id2call.update(
            {t.id: t for t in self.oai_tool_calls if t.id is not None}
        )

        # If using strict output format, parse the output JSON
        self._load_output_format(chat_doc)

        return chat_doc

    def _render_llm_response(
        self, response: ChatDocument | LLMResponse, citation_only: bool = False
    ) -> None:
        is_cached = (
            response.cached
            if isinstance(response, LLMResponse)
            else response.metadata.cached
        )
        if self.llm is None:
            return
        if not citation_only and (not self.llm.get_stream() or is_cached):
            # 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 is_cached else ""
            chat_doc = (
                response
                if isinstance(response, ChatDocument)
                else ChatDocument.from_LLMResponse(response, displayed=True)
            )
            # TODO: prepend TOOL: or OAI-TOOL: if it's a tool-call
            if not settings.quiet:
                print(cached + "[green]" + escape(str(response)))
            self.callbacks.show_llm_response(
                content=str(response),
                is_tool=self.has_tool_message_attempt(chat_doc),
                cached=is_cached,
            )
        if isinstance(response, LLMResponse):
            # we are in the context immediately after an LLM responded,
            # we won't have citations yet, so we're done
            return
        if response.metadata.has_citation:
            citation = (
                response.metadata.source_content
                if self.config.full_citations
                else response.metadata.source
            )
            if not settings.quiet:
                print("[grey37]SOURCES:\n" + escape(citation) + "[/grey37]")
            self.callbacks.show_llm_response(
                content=str(citation),
                is_tool=False,
                cached=False,
                language="text",
            )

    def _llm_response_temp_context(self, message: str, prompt: str) -> ChatDocument:
        """
        Get LLM response to `prompt` (which presumably includes the `message`
        somewhere, along with possible large "context" passages),
        but only include `message` as the USER message, and not the
        full `prompt`, in the message history.
        Args:
            message: the original, relatively short, user request or query
            prompt: the full prompt potentially containing `message` plus context

        Returns:
            Document object containing the response.
        """
        # we explicitly call THIS class's respond method,
        # not a derived class's (or else there would be infinite recursion!)
        with StreamingIfAllowed(self.llm, self.llm.get_stream()):  # type: ignore
            answer_doc = cast(ChatDocument, ChatAgent.llm_response(self, prompt))
        self.update_last_message(message, role=Role.USER)
        return answer_doc

    async def _llm_response_temp_context_async(
        self, message: str, prompt: str
    ) -> ChatDocument:
        """
        Async version of `_llm_response_temp_context`. See there for details.
        """
        # we explicitly call THIS class's respond method,
        # not a derived class's (or else there would be infinite recursion!)
        with StreamingIfAllowed(self.llm, self.llm.get_stream()):  # type: ignore
            answer_doc = cast(
                ChatDocument,
                await ChatAgent.llm_response_async(self, prompt),
            )
        self.update_last_message(message, role=Role.USER)
        return answer_doc

    def llm_response_forget(
        self, message: Optional[str | ChatDocument] = None
    ) -> ChatDocument:
        """
        LLM Response to single message, and restore message_history.
        In effect a "one-off" message & response that leaves agent
        message history state intact.

        Args:
            message (str|ChatDocument): message to respond to.

        Returns:
            A Document object with the response.

        """
        # explicitly call THIS class's respond method,
        # not a derived class's (or else there would be infinite recursion!)
        n_msgs = len(self.message_history)
        with StreamingIfAllowed(self.llm, self.llm.get_stream()):  # type: ignore
            response = cast(ChatDocument, ChatAgent.llm_response(self, message))
        # If there is a response, then we will have two additional
        # messages in the message history, i.e. the user message and the
        # assistant response. We want to (carefully) remove these two messages.
        if len(self.message_history) > n_msgs:
            msg = self.message_history.pop()
            self._drop_msg_update_tool_calls(msg)

        if len(self.message_history) > n_msgs:
            msg = self.message_history.pop()
            self._drop_msg_update_tool_calls(msg)

        # If using strict output format, parse the output JSON
        self._load_output_format(response)

        return response

    async def llm_response_forget_async(
        self, message: Optional[str | ChatDocument] = None
    ) -> ChatDocument:
        """
        Async version of `llm_response_forget`. See there for details.
        """
        # explicitly call THIS class's respond method,
        # not a derived class's (or else there would be infinite recursion!)
        n_msgs = len(self.message_history)
        with StreamingIfAllowed(self.llm, self.llm.get_stream()):  # type: ignore
            response = cast(
                ChatDocument, await ChatAgent.llm_response_async(self, message)
            )
        # If there is a response, then we will have two additional
        # messages in the message history, i.e. the user message and the
        # assistant response. We want to (carefully) remove these two messages.
        if len(self.message_history) > n_msgs:
            msg = self.message_history.pop()
            self._drop_msg_update_tool_calls(msg)

        if len(self.message_history) > n_msgs:
            msg = self.message_history.pop()
            self._drop_msg_update_tool_calls(msg)
        return response

    def chat_num_tokens(self, messages: Optional[List[LLMMessage]] = None) -> int:
        """
        Total number of tokens in the message history so far.

        Args:
            messages: if provided, compute the number of tokens in this list of
                messages, rather than the current message history.
        Returns:
            int: number of tokens in message history
        """
        if self.parser is None:
            raise ValueError(
                "ChatAgent.parser is None. "
                "You must set ChatAgent.parser "
                "before calling chat_num_tokens()."
            )
        hist = messages if messages is not None else self.message_history
        return sum([self.parser.num_tokens(m.content) for m in hist])

    def message_history_str(self, i: Optional[int] = None) -> str:
        """
        Return a string representation of the message history
        Args:
            i: if provided, return only the i-th message when i is postive,
                or last k messages when i = -k.
        Returns:
        """
        if i is None:
            return "\n".join([str(m) for m in self.message_history])
        elif i > 0:
            return str(self.message_history[i])
        else:
            return "\n".join([str(m) for m in self.message_history[i:]])

    def __del__(self) -> None:
        """
        Cleanup method called when the ChatAgent is garbage collected.
        Note: We don't close LLM clients here because they may be shared
        across multiple agents when client caching is enabled.
        The clients are managed centrally and cleaned up via atexit hooks.
        """
        # Previously we closed clients here, but this caused issues when
        # multiple agents shared the same cached client instance.
        # Clients are now managed centrally in langroid.language_models.client_cache
        pass
