154 lines
4.7 KiB
Python
154 lines
4.7 KiB
Python
from __future__ import annotations
|
|
|
|
from copy import deepcopy
|
|
from typing import Any, cast
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
|
|
from src.config import ModelConfig
|
|
from src.llm import tool_loop
|
|
from src.llm.runtime import AttemptPlan
|
|
from src.llm.tool_loop import execute_tool_loop
|
|
from src.llm.types import HonchoLLMCallResponse, ProviderClient
|
|
|
|
|
|
def _openai_plan() -> AttemptPlan:
|
|
return AttemptPlan(
|
|
provider="openai",
|
|
model="mistral-small",
|
|
client=cast(ProviderClient, object()),
|
|
thinking_budget_tokens=None,
|
|
reasoning_effort=None,
|
|
selected_config=ModelConfig(
|
|
model="mistral-small",
|
|
transport="openai",
|
|
),
|
|
attempt=1,
|
|
retry_attempts=1,
|
|
is_fallback=False,
|
|
)
|
|
|
|
|
|
def _response(
|
|
content: str,
|
|
*,
|
|
tool_calls: list[dict[str, Any]] | None = None,
|
|
) -> HonchoLLMCallResponse[Any]:
|
|
return HonchoLLMCallResponse(
|
|
content=content,
|
|
output_tokens=1,
|
|
finish_reasons=["tool_calls" if tool_calls else "stop"],
|
|
tool_calls_made=tool_calls or [],
|
|
)
|
|
|
|
|
|
def _tool_call() -> dict[str, Any]:
|
|
return {
|
|
"id": "call_1",
|
|
"name": "search",
|
|
"input": {"query": "honcho"},
|
|
}
|
|
|
|
|
|
def _reject_tool_to_user(messages: list[dict[str, Any]]) -> None:
|
|
for previous, current in zip(messages, messages[1:], strict=False):
|
|
if previous.get("role") == "tool" and current.get("role") == "user":
|
|
raise RuntimeError("Unexpected role 'user' after role 'tool'")
|
|
|
|
|
|
async def _execute_search(_name: str, _input: dict[str, Any]) -> str:
|
|
return "result"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_max_iteration_synthesis_continues_after_tool_result() -> None:
|
|
calls: list[list[dict[str, Any]]] = []
|
|
responses = iter([_response("", tool_calls=[_tool_call()]), _response("done")])
|
|
|
|
async def strict_call(*_args: Any, **kwargs: Any) -> HonchoLLMCallResponse[Any]:
|
|
messages = deepcopy(kwargs["messages"])
|
|
calls.append(messages)
|
|
_reject_tool_to_user(messages)
|
|
return next(responses)
|
|
|
|
with patch.object(tool_loop, "honcho_llm_call_inner", new=strict_call):
|
|
result = await execute_tool_loop(
|
|
prompt="hi",
|
|
max_tokens=64,
|
|
messages=[{"role": "user", "content": "hi"}],
|
|
tools=[
|
|
{
|
|
"name": "search",
|
|
"description": "Search",
|
|
"input_schema": {"type": "object"},
|
|
}
|
|
],
|
|
tool_choice="auto",
|
|
tool_executor=_execute_search,
|
|
max_tool_iterations=1,
|
|
response_model=None,
|
|
json_mode=False,
|
|
temperature=None,
|
|
stop_seqs=None,
|
|
verbosity=None,
|
|
enable_retry=False,
|
|
retry_attempts=1,
|
|
max_input_tokens=None,
|
|
get_attempt_plan=_openai_plan,
|
|
before_retry_callback=lambda _retry_state: None,
|
|
)
|
|
|
|
assert isinstance(result, HonchoLLMCallResponse)
|
|
assert [message["role"] for message in calls[-1][-2:]] == ["tool", "assistant"]
|
|
assert "maximum number of tool calls" in calls[-1][-1]["content"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_empty_response_retry_continues_after_tool_result() -> None:
|
|
calls: list[list[dict[str, Any]]] = []
|
|
responses = iter(
|
|
[
|
|
_response("", tool_calls=[_tool_call()]),
|
|
_response(""),
|
|
_response("done"),
|
|
]
|
|
)
|
|
|
|
async def strict_call(*_args: Any, **kwargs: Any) -> HonchoLLMCallResponse[Any]:
|
|
messages = deepcopy(kwargs["messages"])
|
|
calls.append(messages)
|
|
_reject_tool_to_user(messages)
|
|
return next(responses)
|
|
|
|
with patch.object(tool_loop, "honcho_llm_call_inner", new=strict_call):
|
|
result = await execute_tool_loop(
|
|
prompt="hi",
|
|
max_tokens=64,
|
|
messages=[{"role": "user", "content": "hi"}],
|
|
tools=[
|
|
{
|
|
"name": "search",
|
|
"description": "Search",
|
|
"input_schema": {"type": "object"},
|
|
}
|
|
],
|
|
tool_choice="auto",
|
|
tool_executor=_execute_search,
|
|
max_tool_iterations=3,
|
|
response_model=None,
|
|
json_mode=False,
|
|
temperature=None,
|
|
stop_seqs=None,
|
|
verbosity=None,
|
|
enable_retry=False,
|
|
retry_attempts=1,
|
|
max_input_tokens=None,
|
|
get_attempt_plan=_openai_plan,
|
|
before_retry_callback=lambda _retry_state: None,
|
|
)
|
|
|
|
assert isinstance(result, HonchoLLMCallResponse)
|
|
assert [message["role"] for message in calls[-1][-2:]] == ["tool", "assistant"]
|
|
assert "last response was empty" in calls[-1][-1]["content"]
|