langchain-ai--langchain/libs/langchain/tests/unit_tests/agents/test_agent.py

"""Unit tests for agents."""

import asyncio
import json
import operator
from functools import reduce
from itertools import cycle
from typing import Any, cast

from langchain_core.agents import (
    AgentAction,
    AgentFinish,
    AgentStep,
)
from langchain_core.callbacks.manager import CallbackManagerForLLMRun
from langchain_core.language_models.llms import LLM
from langchain_core.messages import (
    AIMessage,
    AIMessageChunk,
    FunctionMessage,
    HumanMessage,
    ToolCall,
)
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables.utils import add
from langchain_core.tools import Tool, tool
from langchain_core.tracers import RunLog, RunLogPatch
from typing_extensions import override

from langchain_classic.agents import (
    AgentExecutor,
    AgentType,
    create_openai_functions_agent,
    create_openai_tools_agent,
    create_tool_calling_agent,
    initialize_agent,
)
from langchain_classic.agents.output_parsers.openai_tools import OpenAIToolAgentAction
from tests.unit_tests.callbacks.fake_callback_handler import FakeCallbackHandler
from tests.unit_tests.llms.fake_chat_model import GenericFakeChatModel
from tests.unit_tests.stubs import (
    _AnyIdAIMessageChunk,
)


class FakeListLLM(LLM):
    """Fake LLM for testing that outputs elements of a list."""

    responses: list[str]
    i: int = -1

    @override
    def _call(
        self,
        prompt: str,
        stop: list[str] | None = None,
        run_manager: CallbackManagerForLLMRun | None = None,
        **kwargs: Any,
    ) -> str:
        """Increment counter, and then return response in that index."""
        self.i += 1
        print(f"=== Mock Response #{self.i} ===")  # noqa: T201
        print(self.responses[self.i])  # noqa: T201
        return self.responses[self.i]

    def get_num_tokens(self, text: str) -> int:
        """Return number of tokens in text."""
        return len(text.split())

    async def _acall(self, *args: Any, **kwargs: Any) -> str:
        return self._call(*args, **kwargs)

    @property
    def _identifying_params(self) -> dict[str, Any]:
        return {}

    @property
    def _llm_type(self) -> str:
        """Return type of llm."""
        return "fake_list"


def _get_agent(**kwargs: Any) -> AgentExecutor:
    """Get agent for testing."""
    bad_action_name = "BadAction"
    responses = [
        f"I'm turning evil\nAction: {bad_action_name}\nAction Input: misalignment",
        "Oh well\nFinal Answer: curses foiled again",
    ]
    fake_llm = FakeListLLM(cache=False, responses=responses)

    tools = [
        Tool(
            name="Search",
            func=lambda x: x,
            description="Useful for searching",
        ),
        Tool(
            name="Lookup",
            func=lambda x: x,
            description="Useful for looking up things in a table",
        ),
    ]

    return initialize_agent(
        tools,
        fake_llm,
        agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
        verbose=True,
        **kwargs,
    )


def test_agent_bad_action() -> None:
    """Test react chain when bad action given."""
    agent = _get_agent()
    output = agent.run("when was langchain made")
    assert output == "curses foiled again"


def test_agent_stopped_early() -> None:
    """Test react chain when max iterations or max execution time is exceeded."""
    # iteration limit
    agent = _get_agent(max_iterations=0)
    output = agent.run("when was langchain made")
    assert output == "Agent stopped due to iteration limit or time limit."

    # execution time limit
    agent = _get_agent(max_execution_time=0.0)
    output = agent.run("when was langchain made")
    assert output == "Agent stopped due to iteration limit or time limit."


def test_agent_with_callbacks() -> None:
    """Test react chain with callbacks by setting verbose globally."""
    handler1 = FakeCallbackHandler()
    handler2 = FakeCallbackHandler()

    tool = "Search"
    responses = [
        f"FooBarBaz\nAction: {tool}\nAction Input: misalignment",
        "Oh well\nFinal Answer: curses foiled again",
    ]
    # Only fake LLM gets callbacks for handler2
    fake_llm = FakeListLLM(responses=responses, callbacks=[handler2])
    tools = [
        Tool(
            name="Search",
            func=lambda x: x,
            description="Useful for searching",
        ),
    ]
    agent = initialize_agent(
        tools,
        fake_llm,
        agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
    )

    output = agent.run("when was langchain made", callbacks=[handler1])
    assert output == "curses foiled again"

    # 1 top level chain run runs, 2 LLMChain runs, 2 LLM runs, 1 tool run
    assert handler1.chain_starts == handler1.chain_ends == 3
    assert handler1.llm_starts == handler1.llm_ends == 2
    assert handler1.tool_starts == 1
    assert handler1.tool_ends == 1
    # 1 extra agent action
    assert handler1.starts == 7
    # 1 extra agent end
    assert handler1.ends == 7
    assert handler1.errors == 0
    # during LLMChain
    assert handler1.text == 2

    assert handler2.llm_starts == 2
    assert handler2.llm_ends == 2
    assert (
        handler2.chain_starts
        == handler2.tool_starts
        == handler2.tool_ends
        == handler2.chain_ends
        == 0
    )


def test_agent_stream() -> None:
    """Test react chain with callbacks by setting verbose globally."""
    tool = "Search"
    responses = [
        f"FooBarBaz\nAction: {tool}\nAction Input: misalignment",
        f"FooBarBaz\nAction: {tool}\nAction Input: something else",
        "Oh well\nFinal Answer: curses foiled again",
    ]
    # Only fake LLM gets callbacks for handler2
    fake_llm = FakeListLLM(responses=responses)
    tools = [
        Tool(
            name="Search",
            func=lambda x: f"Results for: {x}",
            description="Useful for searching",
        ),
    ]
    agent = initialize_agent(
        tools,
        fake_llm,
        agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
    )

    output = list(agent.stream("when was langchain made"))
    assert output == [
        {
            "actions": [
                AgentAction(
                    tool="Search",
                    tool_input="misalignment",
                    log="FooBarBaz\nAction: Search\nAction Input: misalignment",
                ),
            ],
            "messages": [
                AIMessage(
                    content="FooBarBaz\nAction: Search\nAction Input: misalignment",
                ),
            ],
        },
        {
            "steps": [
                AgentStep(
                    action=AgentAction(
                        tool="Search",
                        tool_input="misalignment",
                        log="FooBarBaz\nAction: Search\nAction Input: misalignment",
                    ),
                    observation="Results for: misalignment",
                ),
            ],
            "messages": [HumanMessage(content="Results for: misalignment")],
        },
        {
            "actions": [
                AgentAction(
                    tool="Search",
                    tool_input="something else",
                    log="FooBarBaz\nAction: Search\nAction Input: something else",
                ),
            ],
            "messages": [
                AIMessage(
                    content="FooBarBaz\nAction: Search\nAction Input: something else",
                ),
            ],
        },
        {
            "steps": [
                AgentStep(
                    action=AgentAction(
                        tool="Search",
                        tool_input="something else",
                        log="FooBarBaz\nAction: Search\nAction Input: something else",
                    ),
                    observation="Results for: something else",
                ),
            ],
            "messages": [HumanMessage(content="Results for: something else")],
        },
        {
            "output": "curses foiled again",
            "messages": [
                AIMessage(content="Oh well\nFinal Answer: curses foiled again"),
            ],
        },
    ]
    assert add(output) == {
        "actions": [
            AgentAction(
                tool="Search",
                tool_input="misalignment",
                log="FooBarBaz\nAction: Search\nAction Input: misalignment",
            ),
            AgentAction(
                tool="Search",
                tool_input="something else",
                log="FooBarBaz\nAction: Search\nAction Input: something else",
            ),
        ],
        "steps": [
            AgentStep(
                action=AgentAction(
                    tool="Search",
                    tool_input="misalignment",
                    log="FooBarBaz\nAction: Search\nAction Input: misalignment",
                ),
                observation="Results for: misalignment",
            ),
            AgentStep(
                action=AgentAction(
                    tool="Search",
                    tool_input="something else",
                    log="FooBarBaz\nAction: Search\nAction Input: something else",
                ),
                observation="Results for: something else",
            ),
        ],
        "messages": [
            AIMessage(content="FooBarBaz\nAction: Search\nAction Input: misalignment"),
            HumanMessage(content="Results for: misalignment"),
            AIMessage(
                content="FooBarBaz\nAction: Search\nAction Input: something else",
            ),
            HumanMessage(content="Results for: something else"),
            AIMessage(content="Oh well\nFinal Answer: curses foiled again"),
        ],
        "output": "curses foiled again",
    }


def test_agent_tool_return_direct() -> None:
    """Test agent using tools that return directly."""
    tool = "Search"
    responses = [
        f"FooBarBaz\nAction: {tool}\nAction Input: misalignment",
        "Oh well\nFinal Answer: curses foiled again",
    ]
    fake_llm = FakeListLLM(responses=responses)
    tools = [
        Tool(
            name="Search",
            func=lambda x: x,
            description="Useful for searching",
            return_direct=True,
        ),
    ]
    agent = initialize_agent(
        tools,
        fake_llm,
        agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
    )

    output = agent.run("when was langchain made")
    assert output == "misalignment"


def test_agent_tool_return_direct_in_intermediate_steps() -> None:
    """Test agent using tools that return directly."""
    tool = "Search"
    responses = [
        f"FooBarBaz\nAction: {tool}\nAction Input: misalignment",
        "Oh well\nFinal Answer: curses foiled again",
    ]
    fake_llm = FakeListLLM(responses=responses)
    tools = [
        Tool(
            name="Search",
            func=lambda x: x,
            description="Useful for searching",
            return_direct=True,
        ),
    ]
    agent = initialize_agent(
        tools,
        fake_llm,
        agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
        return_intermediate_steps=True,
    )

    resp = agent("when was langchain made")
    assert isinstance(resp, dict)
    assert resp["output"] == "misalignment"
    assert len(resp["intermediate_steps"]) == 1
    action, _action_intput = resp["intermediate_steps"][0]
    assert action.tool == "Search"


def test_agent_with_new_prefix_suffix() -> None:
    """Test agent initialization kwargs with new prefix and suffix."""
    fake_llm = FakeListLLM(
        responses=["FooBarBaz\nAction: Search\nAction Input: misalignment"],
    )
    tools = [
        Tool(
            name="Search",
            func=lambda x: x,
            description="Useful for searching",
            return_direct=True,
        ),
    ]
    prefix = "FooBarBaz"

    suffix = "Begin now!\nInput: {input}\nThought: {agent_scratchpad}"

    agent = initialize_agent(
        tools=tools,
        llm=fake_llm,
        agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
        agent_kwargs={"prefix": prefix, "suffix": suffix},
    )

    # avoids "BasePromptTemplate" has no attribute "template" error
    assert hasattr(agent.agent.llm_chain.prompt, "template")  # type: ignore[union-attr]
    prompt_str = agent.agent.llm_chain.prompt.template  # type: ignore[union-attr]
    assert prompt_str.startswith(prefix), "Prompt does not start with prefix"
    assert prompt_str.endswith(suffix), "Prompt does not end with suffix"


def test_agent_lookup_tool() -> None:
    """Test agent lookup tool."""
    fake_llm = FakeListLLM(
        responses=["FooBarBaz\nAction: Search\nAction Input: misalignment"],
    )
    tools = [
        Tool(
            name="Search",
            func=lambda x: x,
            description="Useful for searching",
            return_direct=True,
        ),
    ]
    agent = initialize_agent(
        tools=tools,
        llm=fake_llm,
        agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
    )

    assert agent.lookup_tool("Search") == tools[0]


def test_agent_invalid_tool() -> None:
    """Test agent invalid tool and correct suggestions."""
    fake_llm = FakeListLLM(responses=["FooBarBaz\nAction: Foo\nAction Input: Bar"])
    tools = [
        Tool(
            name="Search",
            func=lambda x: x,
            description="Useful for searching",
            return_direct=True,
        ),
    ]
    agent = initialize_agent(
        tools=tools,
        llm=fake_llm,
        agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
        return_intermediate_steps=True,
        max_iterations=1,
    )

    resp = agent("when was langchain made")
    assert (
        resp["intermediate_steps"][0][1]
        == "Foo is not a valid tool, try one of [Search]."
    )


async def test_runnable_agent() -> None:
    """Simple test to verify that an agent built via composition works."""
    # Will alternate between responding with hello and goodbye
    infinite_cycle = cycle([AIMessage(content="hello world!")])
    # When streaming GenericFakeChatModel breaks AIMessage into chunks based on spaces
    model = GenericFakeChatModel(messages=infinite_cycle)

    template = ChatPromptTemplate.from_messages(
        [
            ("system", "You are Cat Agent 007"),
            ("human", "{question}"),
        ],
    )

    def fake_parse(_: AIMessage) -> AgentFinish | AgentAction:
        """A parser."""
        return AgentFinish(return_values={"foo": "meow"}, log="hard-coded-message")

    agent = template | model | fake_parse
    executor = AgentExecutor(agent=agent, tools=[])

    # Invoke
    result: Any = await asyncio.to_thread(executor.invoke, {"question": "hello"})
    assert result == {"foo": "meow", "question": "hello"}

    # ainvoke
    result = await executor.ainvoke({"question": "hello"})
    assert result == {"foo": "meow", "question": "hello"}

    # Batch
    result = await asyncio.to_thread(
        executor.batch,
        [{"question": "hello"}, {"question": "hello"}],
    )
    assert result == [
        {"foo": "meow", "question": "hello"},
        {"foo": "meow", "question": "hello"},
    ]

    # abatch
    result = await executor.abatch([{"question": "hello"}, {"question": "hello"}])
    assert result == [
        {"foo": "meow", "question": "hello"},
        {"foo": "meow", "question": "hello"},
    ]

    # Stream
    results = await asyncio.to_thread(list, executor.stream({"question": "hello"}))
    assert results == [
        {"foo": "meow", "messages": [AIMessage(content="hard-coded-message")]},
    ]

    # astream
    results = [r async for r in executor.astream({"question": "hello"})]
    assert results == [
        {
            "foo": "meow",
            "messages": [
                AIMessage(content="hard-coded-message"),
            ],
        },
    ]

    # stream log
    log_results: list[RunLogPatch] = [
        r async for r in executor.astream_log({"question": "hello"})
    ]
    # # Let's stream just the llm tokens.
    messages = []
    for log_record in log_results:
        for op in log_record.ops:
            if op["op"] == "add" and isinstance(op["value"], AIMessageChunk):
                messages.append(op["value"])  # noqa: PERF401

    assert messages != []

    # Aggregate state
    run_log = reduce(operator.add, log_results)

    assert isinstance(run_log, RunLog)

    assert run_log.state["final_output"] == {
        "foo": "meow",
        "messages": [AIMessage(content="hard-coded-message")],
    }


async def test_runnable_agent_with_function_calls() -> None:
    """Test agent with intermediate agent actions."""
    # Will alternate between responding with hello and goodbye
    infinite_cycle = cycle(
        [
            AIMessage(content="looking for pet..."),
            AIMessage(content="Found Pet"),
        ],
    )
    model = GenericFakeChatModel(messages=infinite_cycle)

    template = ChatPromptTemplate.from_messages(
        [
            ("system", "You are Cat Agent 007"),
            ("human", "{question}"),
        ],
    )

    parser_responses = cycle(
        [
            AgentAction(
                tool="find_pet",
                tool_input={
                    "pet": "cat",
                },
                log="find_pet()",
            ),
            AgentFinish(
                return_values={"foo": "meow"},
                log="hard-coded-message",
            ),
        ],
    )

    def fake_parse(_: AIMessage) -> AgentFinish | AgentAction:
        """A parser."""
        return cast("AgentFinish | AgentAction", next(parser_responses))

    @tool
    def find_pet(pet: str) -> str:
        """Find the given pet."""
        if pet != "cat":
            msg = "Only cats allowed"
            raise ValueError(msg)
        return "Spying from under the bed."

    agent = template | model | fake_parse
    executor = AgentExecutor(agent=agent, tools=[find_pet])

    # Invoke
    result = await asyncio.to_thread(executor.invoke, {"question": "hello"})
    assert result == {"foo": "meow", "question": "hello"}

    # ainvoke
    result = await executor.ainvoke({"question": "hello"})
    assert result == {"foo": "meow", "question": "hello"}

    # astream
    results = [r async for r in executor.astream({"question": "hello"})]
    assert results == [
        {
            "actions": [
                AgentAction(
                    tool="find_pet",
                    tool_input={"pet": "cat"},
                    log="find_pet()",
                ),
            ],
            "messages": [AIMessage(content="find_pet()")],
        },
        {
            "messages": [HumanMessage(content="Spying from under the bed.")],
            "steps": [
                AgentStep(
                    action=AgentAction(
                        tool="find_pet",
                        tool_input={"pet": "cat"},
                        log="find_pet()",
                    ),
                    observation="Spying from under the bed.",
                ),
            ],
        },
        {"foo": "meow", "messages": [AIMessage(content="hard-coded-message")]},
    ]

    # astream log

    messages = []
    async for patch in executor.astream_log({"question": "hello"}):
        messages.extend(
            [
                op["value"].content
                for op in patch.ops
                if op["op"] == "add"
                and isinstance(op["value"], AIMessageChunk)
                and op["value"].content != ""
            ]
        )

    assert messages == ["looking", " ", "for", " ", "pet...", "Found", " ", "Pet"]


async def test_runnable_with_multi_action_per_step() -> None:
    """Test an agent that can make multiple function calls at once."""
    # Will alternate between responding with hello and goodbye
    infinite_cycle = cycle(
        [
            AIMessage(content="looking for pet..."),
            AIMessage(content="Found Pet"),
        ],
    )
    model = GenericFakeChatModel(messages=infinite_cycle)

    template = ChatPromptTemplate.from_messages(
        [
            ("system", "You are Cat Agent 007"),
            ("human", "{question}"),
        ],
    )

    parser_responses = cycle(
        [
            [
                AgentAction(
                    tool="find_pet",
                    tool_input={
                        "pet": "cat",
                    },
                    log="find_pet()",
                ),
                AgentAction(
                    tool="pet_pet",  # A function that allows you to pet the given pet.
                    tool_input={
                        "pet": "cat",
                    },
                    log="pet_pet()",
                ),
            ],
            AgentFinish(
                return_values={"foo": "meow"},
                log="hard-coded-message",
            ),
        ],
    )

    def fake_parse(_: AIMessage) -> AgentFinish | AgentAction:
        """A parser."""
        return cast("AgentFinish | AgentAction", next(parser_responses))

    @tool
    def find_pet(pet: str) -> str:
        """Find the given pet."""
        if pet != "cat":
            msg = "Only cats allowed"
            raise ValueError(msg)
        return "Spying from under the bed."

    @tool
    def pet_pet(pet: str) -> str:
        """Pet the given pet."""
        if pet != "cat":
            msg = "Only cats should be petted."
            raise ValueError(msg)
        return "purrrr"

    agent = template | model | fake_parse
    executor = AgentExecutor(agent=agent, tools=[find_pet])

    # Invoke
    result = await asyncio.to_thread(executor.invoke, {"question": "hello"})
    assert result == {"foo": "meow", "question": "hello"}

    # ainvoke
    result = await executor.ainvoke({"question": "hello"})
    assert result == {"foo": "meow", "question": "hello"}

    # astream
    results = [r async for r in executor.astream({"question": "hello"})]
    assert results == [
        {
            "actions": [
                AgentAction(
                    tool="find_pet",
                    tool_input={"pet": "cat"},
                    log="find_pet()",
                ),
            ],
            "messages": [AIMessage(content="find_pet()")],
        },
        {
            "actions": [
                AgentAction(tool="pet_pet", tool_input={"pet": "cat"}, log="pet_pet()"),
            ],
            "messages": [AIMessage(content="pet_pet()")],
        },
        {
            # By-default observation gets converted into human message.
            "messages": [HumanMessage(content="Spying from under the bed.")],
            "steps": [
                AgentStep(
                    action=AgentAction(
                        tool="find_pet",
                        tool_input={"pet": "cat"},
                        log="find_pet()",
                    ),
                    observation="Spying from under the bed.",
                ),
            ],
        },
        {
            "messages": [
                HumanMessage(
                    content="pet_pet is not a valid tool, try one of [find_pet].",
                ),
            ],
            "steps": [
                AgentStep(
                    action=AgentAction(
                        tool="pet_pet",
                        tool_input={"pet": "cat"},
                        log="pet_pet()",
                    ),
                    observation="pet_pet is not a valid tool, try one of [find_pet].",
                ),
            ],
        },
        {"foo": "meow", "messages": [AIMessage(content="hard-coded-message")]},
    ]

    # astream log

    messages = []
    async for patch in executor.astream_log({"question": "hello"}):
        for op in patch.ops:
            if op["op"] != "add":
                continue

            value = op["value"]

            if not isinstance(value, AIMessageChunk):
                continue

            if value.content == "":  # Then it's a function invocation message
                continue

            messages.append(value.content)

    assert messages == ["looking", " ", "for", " ", "pet...", "Found", " ", "Pet"]


def _make_func_invocation(name: str, **kwargs: Any) -> AIMessage:
    """Create an AIMessage that represents a function invocation.

    Args:
        name: Name of the function to invoke.
        kwargs: Keyword arguments to pass to the function.

    Returns:
        AIMessage that represents a request to invoke a function.
    """
    return AIMessage(
        content="",
        additional_kwargs={
            "function_call": {
                "name": name,
                "arguments": json.dumps(kwargs),
            },
        },
    )


def _recursive_dump(obj: Any) -> Any:
    """Recursively dump the object if encountering any pydantic models."""
    if isinstance(obj, dict):
        return {
            k: _recursive_dump(v)
            for k, v in obj.items()
            if k != "id"  # Remove the id field for testing purposes
        }
    if isinstance(obj, list):
        return [_recursive_dump(v) for v in obj]
    if hasattr(obj, "dict"):
        # if the object contains an ID field, we'll remove it for testing purposes
        if hasattr(obj, "id"):
            d = obj.model_dump()
            d.pop("id")
            return _recursive_dump(d)
        return _recursive_dump(obj.model_dump())
    return obj


async def test_openai_agent_with_streaming() -> None:
    """Test openai agent with streaming."""
    infinite_cycle = cycle(
        [
            _make_func_invocation("find_pet", pet="cat"),
            AIMessage(content="The cat is spying from under the bed."),
        ],
    )

    model = GenericFakeChatModel(messages=infinite_cycle)

    @tool
    def find_pet(pet: str) -> str:
        """Find the given pet."""
        if pet != "cat":
            msg = "Only cats allowed"
            raise ValueError(msg)
        return "Spying from under the bed."

    template = ChatPromptTemplate.from_messages(
        [
            ("system", "You are a helpful AI bot. Your name is kitty power meow."),
            ("human", "{question}"),
            MessagesPlaceholder(
                variable_name="agent_scratchpad",
            ),
        ],
    )

    # type error due to base tool type below -- would need to be adjusted on tool
    # decorator.
    agent = create_openai_functions_agent(
        model,
        [find_pet],
        template,
    )
    executor = AgentExecutor(agent=agent, tools=[find_pet])

    # Invoke
    result = await asyncio.to_thread(executor.invoke, {"question": "hello"})
    assert result == {
        "output": "The cat is spying from under the bed.",
        "question": "hello",
    }

    # astream
    chunks = [chunk async for chunk in executor.astream({"question": "hello"})]
    assert _recursive_dump(chunks) == [
        {
            "actions": [
                {
                    "log": "\nInvoking: `find_pet` with `{'pet': 'cat'}`\n\n\n",
                    "message_log": [
                        {
                            "additional_kwargs": {
                                "function_call": {
                                    "arguments": '{"pet": "cat"}',
                                    "name": "find_pet",
                                },
                            },
                            "content": "",
                            "name": None,
                            "response_metadata": {},
                            "type": "AIMessageChunk",
                        },
                    ],
                    "tool": "find_pet",
                    "tool_input": {"pet": "cat"},
                    "type": "AgentActionMessageLog",
                },
            ],
            "messages": [
                {
                    "additional_kwargs": {
                        "function_call": {
                            "arguments": '{"pet": "cat"}',
                            "name": "find_pet",
                        },
                    },
                    "chunk_position": "last",
                    "content": "",
                    "invalid_tool_calls": [],
                    "name": None,
                    "response_metadata": {},
                    "tool_call_chunks": [],
                    "tool_calls": [],
                    "type": "AIMessageChunk",
                    "usage_metadata": None,
                },
            ],
        },
        {
            "messages": [
                {
                    "additional_kwargs": {},
                    "content": "Spying from under the bed.",
                    "name": "find_pet",
                    "response_metadata": {},
                    "type": "function",
                },
            ],
            "steps": [
                {
                    "action": {
                        "log": "\nInvoking: `find_pet` with `{'pet': 'cat'}`\n\n\n",
                        "tool": "find_pet",
                        "tool_input": {"pet": "cat"},
                        "type": "AgentActionMessageLog",
                    },
                    "observation": "Spying from under the bed.",
                },
            ],
        },
        {
            "messages": [
                {
                    "additional_kwargs": {},
                    "content": "The cat is spying from under the bed.",
                    "invalid_tool_calls": [],
                    "name": None,
                    "response_metadata": {},
                    "tool_calls": [],
                    "type": "ai",
                    "usage_metadata": None,
                },
            ],
            "output": "The cat is spying from under the bed.",
        },
    ]

    #
    # # astream_log
    log_patches = [
        log_patch async for log_patch in executor.astream_log({"question": "hello"})
    ]

    messages = []

    for log_patch in log_patches:
        for op in log_patch.ops:
            if op["op"] == "add" and isinstance(op["value"], AIMessageChunk):
                value = op["value"]
                if value.content:  # Filter out function call messages
                    messages.append(value.content)

    assert messages == [
        "The",
        " ",
        "cat",
        " ",
        "is",
        " ",
        "spying",
        " ",
        "from",
        " ",
        "under",
        " ",
        "the",
        " ",
        "bed.",
    ]


def _make_tools_invocation(name_to_arguments: dict[str, dict[str, Any]]) -> AIMessage:
    """Create an AIMessage that represents a tools invocation.

    Args:
        name_to_arguments: A dictionary mapping tool names to an invocation.

    Returns:
        AIMessage that represents a request to invoke a tool.
    """
    raw_tool_calls = [
        {"function": {"name": name, "arguments": json.dumps(arguments)}, "id": str(idx)}
        for idx, (name, arguments) in enumerate(name_to_arguments.items())
    ]
    tool_calls = [
        ToolCall(name=name, args=args, id=str(idx), type="tool_call")
        for idx, (name, args) in enumerate(name_to_arguments.items())
    ]
    return AIMessage(
        content="",
        additional_kwargs={
            "tool_calls": raw_tool_calls,
        },
        tool_calls=tool_calls,
    )


async def test_openai_agent_tools_agent() -> None:
    """Test OpenAI tools agent."""
    infinite_cycle = cycle(
        [
            _make_tools_invocation(
                {
                    "find_pet": {"pet": "cat"},
                    "check_time": {},
                },
            ),
            AIMessage(content="The cat is spying from under the bed."),
        ],
    )

    GenericFakeChatModel.bind_tools = lambda self, _: self  # type: ignore[assignment,misc]
    model = GenericFakeChatModel(messages=infinite_cycle)

    @tool
    def find_pet(pet: str) -> str:
        """Find the given pet."""
        if pet != "cat":
            msg = "Only cats allowed"
            raise ValueError(msg)
        return "Spying from under the bed."

    @tool
    def check_time() -> str:
        """Find the given pet."""
        return "It's time to pet the cat."

    template = ChatPromptTemplate.from_messages(
        [
            ("system", "You are a helpful AI bot. Your name is kitty power meow."),
            ("human", "{question}"),
            MessagesPlaceholder(
                variable_name="agent_scratchpad",
            ),
        ],
    )

    # type error due to base tool type below -- would need to be adjusted on tool
    # decorator.
    openai_agent = create_openai_tools_agent(
        model,
        [find_pet],
        template,
    )
    tool_calling_agent = create_tool_calling_agent(
        model,
        [find_pet],
        template,
    )
    for agent in [openai_agent, tool_calling_agent]:
        executor = AgentExecutor(agent=agent, tools=[find_pet])

        # Invoke
        result = await asyncio.to_thread(executor.invoke, {"question": "hello"})
        assert result == {
            "output": "The cat is spying from under the bed.",
            "question": "hello",
        }

        # astream
        chunks = [chunk async for chunk in executor.astream({"question": "hello"})]
        assert chunks == [
            {
                "actions": [
                    OpenAIToolAgentAction(
                        tool="find_pet",
                        tool_input={"pet": "cat"},
                        log="\nInvoking: `find_pet` with `{'pet': 'cat'}`\n\n\n",
                        message_log=[
                            _AnyIdAIMessageChunk(
                                content="",
                                additional_kwargs={
                                    "tool_calls": [
                                        {
                                            "function": {
                                                "name": "find_pet",
                                                "arguments": '{"pet": "cat"}',
                                            },
                                            "id": "0",
                                        },
                                        {
                                            "function": {
                                                "name": "check_time",
                                                "arguments": "{}",
                                            },
                                            "id": "1",
                                        },
                                    ],
                                },
                                chunk_position="last",
                            ),
                        ],
                        tool_call_id="0",
                    ),
                ],
                "messages": [
                    _AnyIdAIMessageChunk(
                        content="",
                        additional_kwargs={
                            "tool_calls": [
                                {
                                    "function": {
                                        "name": "find_pet",
                                        "arguments": '{"pet": "cat"}',
                                    },
                                    "id": "0",
                                },
                                {
                                    "function": {
                                        "name": "check_time",
                                        "arguments": "{}",
                                    },
                                    "id": "1",
                                },
                            ],
                        },
                        chunk_position="last",
                    ),
                ],
            },
            {
                "actions": [
                    OpenAIToolAgentAction(
                        tool="check_time",
                        tool_input={},
                        log="\nInvoking: `check_time` with `{}`\n\n\n",
                        message_log=[
                            _AnyIdAIMessageChunk(
                                content="",
                                additional_kwargs={
                                    "tool_calls": [
                                        {
                                            "function": {
                                                "name": "find_pet",
                                                "arguments": '{"pet": "cat"}',
                                            },
                                            "id": "0",
                                        },
                                        {
                                            "function": {
                                                "name": "check_time",
                                                "arguments": "{}",
                                            },
                                            "id": "1",
                                        },
                                    ],
                                },
                                chunk_position="last",
                            ),
                        ],
                        tool_call_id="1",
                    ),
                ],
                "messages": [
                    _AnyIdAIMessageChunk(
                        content="",
                        additional_kwargs={
                            "tool_calls": [
                                {
                                    "function": {
                                        "name": "find_pet",
                                        "arguments": '{"pet": "cat"}',
                                    },
                                    "id": "0",
                                },
                                {
                                    "function": {
                                        "name": "check_time",
                                        "arguments": "{}",
                                    },
                                    "id": "1",
                                },
                            ],
                        },
                        chunk_position="last",
                    ),
                ],
            },
            {
                "messages": [
                    FunctionMessage(
                        content="Spying from under the bed.",
                        name="find_pet",
                    ),
                ],
                "steps": [
                    AgentStep(
                        action=OpenAIToolAgentAction(
                            tool="find_pet",
                            tool_input={"pet": "cat"},
                            log="\nInvoking: `find_pet` with `{'pet': 'cat'}`\n\n\n",
                            message_log=[
                                _AnyIdAIMessageChunk(
                                    content="",
                                    additional_kwargs={
                                        "tool_calls": [
                                            {
                                                "function": {
                                                    "name": "find_pet",
                                                    "arguments": '{"pet": "cat"}',
                                                },
                                                "id": "0",
                                            },
                                            {
                                                "function": {
                                                    "name": "check_time",
                                                    "arguments": "{}",
                                                },
                                                "id": "1",
                                            },
                                        ],
                                    },
                                    chunk_position="last",
                                ),
                            ],
                            tool_call_id="0",
                        ),
                        observation="Spying from under the bed.",
                    ),
                ],
            },
            {
                "messages": [
                    FunctionMessage(
                        content="check_time is not a valid tool, "
                        "try one of [find_pet].",
                        name="check_time",
                    ),
                ],
                "steps": [
                    AgentStep(
                        action=OpenAIToolAgentAction(
                            tool="check_time",
                            tool_input={},
                            log="\nInvoking: `check_time` with `{}`\n\n\n",
                            message_log=[
                                _AnyIdAIMessageChunk(
                                    content="",
                                    additional_kwargs={
                                        "tool_calls": [
                                            {
                                                "function": {
                                                    "name": "find_pet",
                                                    "arguments": '{"pet": "cat"}',
                                                },
                                                "id": "0",
                                            },
                                            {
                                                "function": {
                                                    "name": "check_time",
                                                    "arguments": "{}",
                                                },
                                                "id": "1",
                                            },
                                        ],
                                    },
                                    chunk_position="last",
                                ),
                            ],
                            tool_call_id="1",
                        ),
                        observation="check_time is not a valid tool, "
                        "try one of [find_pet].",
                    ),
                ],
            },
            {
                "messages": [
                    AIMessage(content="The cat is spying from under the bed."),
                ],
                "output": "The cat is spying from under the bed.",
            },
        ]

        # astream_log
        log_patches = [
            log_patch async for log_patch in executor.astream_log({"question": "hello"})
        ]

        # Get the tokens from the astream log response.
        messages = []

        for log_patch in log_patches:
            for op in log_patch.ops:
                if op["op"] == "add" and isinstance(op["value"], AIMessageChunk):
                    value = op["value"]
                    if value.content:  # Filter out function call messages
                        messages.append(value.content)

        assert messages == [
            "The",
            " ",
            "cat",
            " ",
            "is",
            " ",
            "spying",
            " ",
            "from",
            " ",
            "under",
            " ",
            "the",
            " ",
            "bed.",
        ]