mirror of
https://github.com/NCBM/plyngent.git
synced 2026-11-07 23:47:54 +08:00
33 lines
1.0 KiB
Python
33 lines
1.0 KiB
Python
from __future__ import annotations
|
|
|
|
from plyngent.agent.budget import estimate_message_chars, truncate_tool_result
|
|
from plyngent.lmproto.openai_compatible.model import ToolChatMessage, UserChatMessage
|
|
|
|
|
|
def test_truncate_tool_result_short() -> None:
|
|
assert truncate_tool_result("hello", 100) == "hello"
|
|
|
|
|
|
def test_truncate_tool_result_long() -> None:
|
|
text = "a" * 50
|
|
out = truncate_tool_result(text, 20)
|
|
assert out.startswith("a" * 20)
|
|
assert "Truncated" in out
|
|
assert "20" in out
|
|
assert "30" in out
|
|
assert "configuration file" in out
|
|
|
|
|
|
def test_truncate_tool_result_no_hint() -> None:
|
|
text = "a" * 50
|
|
out = truncate_tool_result(text, 20, hint=False)
|
|
assert out.startswith("a" * 20)
|
|
assert "truncated" in out
|
|
assert "configuration file" not in out
|
|
|
|
|
|
def test_estimate_message_chars() -> None:
|
|
assert estimate_message_chars(UserChatMessage(content="hello")) == 5
|
|
tool = ToolChatMessage(content="abc", tool_call_id="id1")
|
|
assert estimate_message_chars(tool) == 6
|