104 lines
3.1 KiB
Python
104 lines
3.1 KiB
Python
from typing import Any
|
|
|
|
from src.llm.conversation import (
|
|
_is_tool_result_message, # pyright: ignore[reportPrivateUsage]
|
|
_is_tool_use_message, # pyright: ignore[reportPrivateUsage]
|
|
truncate_messages_to_fit,
|
|
)
|
|
|
|
|
|
def test_truncate_messages_to_fit_keeps_last_unit_when_over_limit() -> None:
|
|
messages = [
|
|
{"role": "user", "content": "x " * 2000},
|
|
]
|
|
|
|
truncated = truncate_messages_to_fit(messages, max_tokens=1)
|
|
|
|
assert truncated == messages
|
|
|
|
|
|
def test_truncate_messages_to_fit_preserves_tool_result_pair() -> None:
|
|
messages = [
|
|
{"role": "user", "content": "old context " * 1000},
|
|
{
|
|
"role": "assistant",
|
|
"content": None,
|
|
"tool_calls": [
|
|
{
|
|
"id": "call_1",
|
|
"type": "function",
|
|
"function": {"name": "lookup", "arguments": "{}"},
|
|
}
|
|
],
|
|
},
|
|
{"role": "tool", "tool_call_id": "call_1", "content": "result"},
|
|
]
|
|
|
|
truncated = truncate_messages_to_fit(messages, max_tokens=5)
|
|
|
|
assert truncated == messages[1:]
|
|
|
|
|
|
def test_is_tool_use_message_detects_gemini_function_call_in_parts() -> None:
|
|
msg: dict[str, Any] = {
|
|
"role": "model",
|
|
"parts": [
|
|
{"function_call": {"name": "search", "args": {"q": "honcho"}}},
|
|
],
|
|
}
|
|
assert _is_tool_use_message(msg) is True
|
|
|
|
|
|
def test_is_tool_result_message_detects_gemini_function_response_in_parts() -> None:
|
|
msg: dict[str, Any] = {
|
|
"role": "user",
|
|
"parts": [
|
|
{"function_response": {"name": "search", "response": {"result": "ok"}}},
|
|
],
|
|
}
|
|
assert _is_tool_result_message(msg) is True
|
|
|
|
|
|
def test_is_tool_use_message_detects_anthropic_tool_use_block() -> None:
|
|
msg: dict[str, Any] = {
|
|
"role": "assistant",
|
|
"content": [
|
|
{"type": "text", "text": "calling lookup"},
|
|
{"type": "tool_use", "id": "t_1", "name": "lookup", "input": {}},
|
|
],
|
|
}
|
|
assert _is_tool_use_message(msg) is True
|
|
|
|
|
|
def test_truncate_messages_to_fit_preserves_gemini_tool_pair() -> None:
|
|
"""A Gemini-shaped function_call / function_response pair must stay
|
|
grouped when older units get dropped. Regression: before adding the
|
|
parts-based detection, neither message would be recognized as a tool
|
|
unit, and truncation could split or drop them individually."""
|
|
messages: list[dict[str, Any]] = [
|
|
{"role": "user", "parts": [{"text": "old context " * 1000}]},
|
|
{
|
|
"role": "model",
|
|
"parts": [
|
|
{"function_call": {"name": "lookup", "args": {}}},
|
|
],
|
|
},
|
|
{
|
|
"role": "user",
|
|
"parts": [
|
|
{
|
|
"function_response": {
|
|
"name": "lookup",
|
|
"response": {"result": "found"},
|
|
}
|
|
}
|
|
],
|
|
},
|
|
]
|
|
|
|
truncated = truncate_messages_to_fit(messages, max_tokens=20)
|
|
|
|
# The oldest (bulk-text) message should be dropped; the function_call +
|
|
# function_response pair stays intact together.
|
|
assert truncated == messages[1:]
|