replypier/backend/tests/utils/test_agno_patches.py

"""
Copyright 2024-2026 ReplyPier

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

from types import SimpleNamespace

from agno.models.base import Model
from agno.tools.function import Function, FunctionCall

from app.utils.agno_patches import _strip_tools, apply_agno_patches


def _model_stub():
    # create_tool_call_limit_error_result only reads tool_message_role from self,
    # and the patch only sets its own attributes, so a stub is enough.
    return SimpleNamespace(tool_message_role="tool")


def _next_round(model, tools=None):
    """One model round-trip, the way agno's loop does it."""
    return _strip_tools(
        model, {"tools": tools if tools is not None else [{"name": "search"}], "tool_choice": "auto"}
    )


def _limit_error_message(model_stub, call_id="call_1"):
    apply_agno_patches()
    function_call = FunctionCall(
        function=Function(name="search_knowledge_base"),
        arguments={"query": "hello"},
        call_id=call_id,
    )
    return Model.create_tool_call_limit_error_result(model_stub, function_call)


def test_limit_does_not_end_the_run_before_the_model_answers():
    """The whole point of the limit handling: bound the run WITHOUT costing the
    answer.

    Stopping at the tool result ends the loop before the model writes anything,
    so the turn comes back empty and the visitor gets the generic
    "I didn't quite catch that" fallback. That was the production regression this
    replaces — most turns on any agent with a few knowledge sources.
    """
    message = _limit_error_message(_model_stub())
    assert not message.stop_after_tool_call


def test_tools_are_withdrawn_once_the_budget_is_spent():
    """With tools cleared, the next completion has to be plain text — which both
    produces an answer and makes the runaway loop of #269 impossible."""
    model = _model_stub()
    _next_round(model)
    _limit_error_message(model)

    kwargs = _next_round(model)
    assert kwargs["tools"] is None
    assert kwargs["tool_choice"] is None


def test_the_structured_output_tool_is_never_withdrawn():
    """On Groq the reply IS a tool call named `json`, not text.

    Withdrawing every tool would leave that model no way to answer at all, so a
    Groq agent reaching the limit would return an empty turn every single time —
    the exact failure this patch removes for everyone else.
    """
    model = _model_stub()
    _next_round(model)
    _limit_error_message(model)

    kwargs = _strip_tools(model, {
        "tools": [
            {"type": "function", "function": {"name": "search_knowledge_base"}},
            {"type": "function", "function": {"name": "json"}},
        ],
        "tool_choice": "auto",
    })
    assert [t["function"]["name"] for t in kwargs["tools"]] == ["json"]
    # tool_choice must survive — it is how the capture tool gets called.
    assert kwargs["tool_choice"] == "auto"


def test_tools_survive_until_the_limit_is_hit():
    """Before the limit, the model call must be completely untouched."""
    model = _model_stub()
    original = {"tools": [{"name": "search"}], "tool_choice": "auto"}
    assert _strip_tools(model, dict(original)) == original


def test_stripping_leaves_the_other_arguments_alone():
    model = _model_stub()
    _next_round(model)
    _limit_error_message(model)

    kwargs = _strip_tools(
        model, {"tools": [{"name": "search"}], "tool_choice": "auto", "messages": ["m"]}
    )
    assert kwargs["messages"] == ["m"]


def test_parallel_tool_calls_in_one_batch_do_not_end_the_run():
    """The regression that made the first version of this patch useless.

    A model can ask for several tools in ONE assistant message; agno walks the
    whole batch before calling the model again, so every over-budget call in that
    batch lands here. Treating the second as a repeat terminated the run and the
    visitor got the empty-turn fallback anyway. Caught in local testing, not by
    the original tests — they called this twice and unwittingly simulated two
    rounds.
    """
    model = _model_stub()
    _next_round(model)

    first = _limit_error_message(model, call_id="call_1")
    second = _limit_error_message(model, call_id="call_2")
    third = _limit_error_message(model, call_id="call_3")

    assert not first.stop_after_tool_call
    assert not second.stop_after_tool_call
    assert not third.stop_after_tool_call


def test_the_run_is_never_hard_stopped():
    """No branch may set stop_after_tool_call — that is what cost the answer.

    A previous version kept a hard stop for the case where tool withdrawal
    somehow failed. It fired occasionally in real traffic and every time it did,
    the visitor got the empty-turn fallback instead of a reply. The hang it
    guarded is already prevented by the AGENT_RUN_TIMEOUT wrapper around every
    agno run, which does not depend on this module patching the right methods.

    Covers a later round as well as the same batch, since the old backstop keyed
    on exactly that distinction.
    """
    model = _model_stub()
    _next_round(model)
    assert not _limit_error_message(model).stop_after_tool_call

    _next_round(model)
    assert not _limit_error_message(model, call_id="call_2").stop_after_tool_call

    _next_round(model)
    assert not _limit_error_message(model, call_id="call_3").stop_after_tool_call


def test_exhaustion_does_not_leak_between_runs():
    """The state lives on the Model instance, and create_model() builds a fresh
    one per request. A limit hit in one conversation must not withdraw tools from
    the next."""
    spent = _model_stub()
    _next_round(spent)
    _limit_error_message(spent)

    fresh = _model_stub()
    original = {"tools": [{"name": "search"}], "tool_choice": "auto"}
    assert _strip_tools(fresh, dict(original)) == original


def test_limit_error_result_keeps_original_shape():
    """The patch only adds behaviour; the message agno built is unchanged."""
    message = _limit_error_message(_model_stub())
    assert message.role == "tool"
    assert message.tool_call_error is True
    assert message.tool_call_id == "call_1"
    assert message.tool_name == "search_knowledge_base"
    assert "Answer now" in message.content


def test_apply_is_idempotent():
    """Repeated application must not stack wrappers."""
    apply_agno_patches()
    patched_once = Model.create_tool_call_limit_error_result
    apply_agno_patches()
    assert Model.create_tool_call_limit_error_result is patched_once