from langchain.tools import BaseTool
from langchain.tools.render import render_text_description
from langchain_core.language_models.base import LanguageModelLike
from langchain_core.messages import (
    AIMessage,
    FunctionMessage,
    HumanMessage,
    SystemMessage,
)
from langgraph.checkpoint import BaseCheckpointSaver
from langgraph.graph import END
from langgraph.graph.message import MessageGraph
from langgraph.prebuilt import ToolExecutor, ToolInvocation

from app.agent_types.prompts import xml_template
from app.message_types import LiberalFunctionMessage


def _collapse_messages(messages):
    log = ""
    if isinstance(messages[-1], AIMessage):
        scratchpad = messages[:-1]
        final = messages[-1]
    else:
        scratchpad = messages
        final = None
    if len(scratchpad) % 2 != 0:
        raise ValueError("Unexpected")
    for i in range(0, len(scratchpad), 2):
        action = messages[i]
        observation = messages[i + 1]
        log += f"{action.content}<observation>{observation.content}</observation>"
    if final is not None:
        log += final.content
    return AIMessage(content=log)


def construct_chat_history(messages):
    collapsed_messages = []
    temp_messages = []
    for message in messages:
        if isinstance(message, HumanMessage):
            if temp_messages:
                collapsed_messages.append(_collapse_messages(temp_messages))
                temp_messages = []
            collapsed_messages.append(message)
        elif isinstance(message, LiberalFunctionMessage):
            _dict = message.dict()
            _dict["content"] = str(_dict["content"])
            m_c = FunctionMessage(**_dict)
            temp_messages.append(m_c)
        else:
            temp_messages.append(message)

    # Don't forget to add the last non-human message if it exists
    if temp_messages:
        collapsed_messages.append(_collapse_messages(temp_messages))

    return collapsed_messages


def get_xml_agent_executor(
    tools: list[BaseTool],
    llm: LanguageModelLike,
    system_message: str,
    interrupt_before_action: bool,
    checkpoint: BaseCheckpointSaver,
):
    formatted_system_message = xml_template.format(
        system_message=system_message,
        tools=render_text_description(tools),
        tool_names=", ".join([t.name for t in tools]),
    )

    llm_with_stop = llm.bind(stop=["</tool_input>", "<observation>"])

    def _get_messages(messages):
        return [
            SystemMessage(content=formatted_system_message)
        ] + construct_chat_history(messages)

    agent = _get_messages | llm_with_stop
    tool_executor = ToolExecutor(tools)

    # Define the function that determines whether to continue or not
    def should_continue(messages):
        last_message = messages[-1]
        if "</tool>" in last_message.content:
            return "continue"
        else:
            return "end"

    # Define the function to execute tools
    async def call_tool(messages):
        # Based on the continue condition
        # we know the last message involves a function call
        last_message = messages[-1]
        # We construct an ToolInvocation from the function_call
        tool, tool_input = last_message.content.split("</tool>")
        _tool = tool.split("<tool>")[1]
        if "<tool_input>" not in tool_input:
            _tool_input = ""
        else:
            _tool_input = tool_input.split("<tool_input>")[1]
            if "</tool_input>" in _tool_input:
                _tool_input = _tool_input.split("</tool_input>")[0]
        action = ToolInvocation(
            tool=_tool,
            tool_input=_tool_input,
        )
        # We call the tool_executor and get back a response
        response = await tool_executor.ainvoke(action)
        # We use the response to create a FunctionMessage
        function_message = LiberalFunctionMessage(content=response, name=action.tool)
        # We return a list, because this will get added to the existing list
        return function_message

    workflow = MessageGraph()

    # Define the two nodes we will cycle between
    workflow.add_node("agent", agent)
    workflow.add_node("action", call_tool)

    # Set the entrypoint as `agent`
    # This means that this node is the first one called
    workflow.set_entry_point("agent")

    # We now add a conditional edge
    workflow.add_conditional_edges(
        # First, we define the start node. We use `agent`.
        # This means these are the edges taken after the `agent` node is called.
        "agent",
        # Next, we pass in the function that will determine which node is called next.
        should_continue,
        # Finally we pass in a mapping.
        # The keys are strings, and the values are other nodes.
        # END is a special node marking that the graph should finish.
        # What will happen is we will call `should_continue`, and then the output of that
        # will be matched against the keys in this mapping.
        # Based on which one it matches, that node will then be called.
        {
            # If `tools`, then we call the tool node.
            "continue": "action",
            # Otherwise we finish.
            "end": END,
        },
    )

    # We now add a normal edge from `tools` to `agent`.
    # This means that after `tools` is called, `agent` node is called next.
    workflow.add_edge("action", "agent")

    # Finally, we compile it!
    # This compiles it into a LangChain Runnable,
    # meaning you can use it as you would any other runnable
    app = workflow.compile(checkpointer=checkpoint)
    if interrupt_before_action:
        app.interrupt = ["action:inbox"]
    return app

langchain_core.messages.AIMessage langchain.tools.render.render_text_description langchain_core.messages.SystemMessage langchain_core.messages.FunctionMessage
[((1121, 1143), 'langchain_core.messages.AIMessage', 'AIMessage', ([], {'content': 'log'}), '(content=log)\n', (1130, 1143), False, 'from langchain_core.messages import AIMessage, FunctionMessage, HumanMessage, SystemMessage\n'), ((2644, 2663), 'langgraph.prebuilt.ToolExecutor', 'ToolExecutor', (['tools'], {}), '(tools)\n', (2656, 2663), False, 'from langgraph.prebuilt import ToolExecutor, ToolInvocation\n'), ((4058, 4072), 'langgraph.graph.message.MessageGraph', 'MessageGraph', ([], {}), '()\n', (4070, 4072), False, 'from langgraph.graph.message import MessageGraph\n'), ((3588, 3638), 'langgraph.prebuilt.ToolInvocation', 'ToolInvocation', ([], {'tool': '_tool', 'tool_input': '_tool_input'}), '(tool=_tool, tool_input=_tool_input)\n', (3602, 3638), False, 'from langgraph.prebuilt import ToolExecutor, ToolInvocation\n'), ((3874, 3932), 'app.message_types.LiberalFunctionMessage', 'LiberalFunctionMessage', ([], {'content': 'response', 'name': 'action.tool'}), '(content=response, name=action.tool)\n', (3896, 3932), False, 'from app.message_types import LiberalFunctionMessage\n'), ((2261, 2291), 'langchain.tools.render.render_text_description', 'render_text_description', (['tools'], {}), '(tools)\n', (2284, 2291), False, 'from langchain.tools.render import render_text_description\n'), ((1663, 1687), 'langchain_core.messages.FunctionMessage', 'FunctionMessage', ([], {}), '(**_dict)\n', (1678, 1687), False, 'from langchain_core.messages import AIMessage, FunctionMessage, HumanMessage, SystemMessage\n'), ((2488, 2535), 'langchain_core.messages.SystemMessage', 'SystemMessage', ([], {'content': 'formatted_system_message'}), '(content=formatted_system_message)\n', (2501, 2535), False, 'from langchain_core.messages import AIMessage, FunctionMessage, HumanMessage, SystemMessage\n')]