wehub-resource-sync
c3bf08ac8d
K8s Workspace Integration Tests / k8s-workspace-tests (push) Waiting to run
Pre-commit / run (ubuntu-latest) (push) Waiting to run
Python Unittest Coverage / test (macos-15, 3.11) (push) Waiting to run
Python Unittest Coverage / test (ubuntu-latest, 3.11) (push) Waiting to run
Python Unittest Coverage / test (windows-latest, 3.11) (push) Waiting to run
Web UI / check (push) Waiting to run
1084 lines
35 KiB
Python
1084 lines
35 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""A template test case."""
|
|
# pylint: disable=protected-access
|
|
import json
|
|
import os
|
|
import tempfile
|
|
from typing import Any
|
|
|
|
from unittest.async_case import IsolatedAsyncioTestCase
|
|
|
|
from utils import MockModel, AnyString
|
|
|
|
from agentscope.model import StructuredResponse
|
|
from agentscope.agent import Agent, ContextConfig
|
|
from agentscope.state import AgentState
|
|
from agentscope.message import (
|
|
UserMsg,
|
|
AssistantMsg,
|
|
TextBlock,
|
|
ToolCallBlock,
|
|
HintBlock,
|
|
Msg,
|
|
)
|
|
from agentscope.tool import Toolkit
|
|
|
|
|
|
class RecordingStructuredMockModel(MockModel):
|
|
"""A mock model that records structured-output compression calls."""
|
|
|
|
def __init__(
|
|
self,
|
|
*args: Any,
|
|
fail_structured_output_times: int = 0,
|
|
force_compression_overflow: bool = False,
|
|
**kwargs: Any,
|
|
) -> None:
|
|
"""Initialize the recording mock model."""
|
|
super().__init__(*args, **kwargs)
|
|
self.recorded_structured_messages: list[list[Msg]] = []
|
|
self._fail_structured_output_times = fail_structured_output_times
|
|
self._force_compression_overflow = force_compression_overflow
|
|
self._compression_count_calls = 0
|
|
|
|
async def count_tokens(
|
|
self,
|
|
messages: list[Msg],
|
|
tools: list[dict] | None,
|
|
) -> int:
|
|
"""Force the overflow branch when counting compression messages."""
|
|
is_compression_count = bool(
|
|
tools
|
|
and tools[0].get("function", {}).get("name")
|
|
== "generate_structured_output",
|
|
)
|
|
if self._force_compression_overflow and is_compression_count:
|
|
self._compression_count_calls += 1
|
|
if self._compression_count_calls == 1:
|
|
return self.context_size + 1
|
|
return 1
|
|
return await super().count_tokens(messages, tools)
|
|
|
|
async def _call_api_with_structured_output(
|
|
self,
|
|
model_name: str,
|
|
messages: list[Msg],
|
|
structured_model: Any,
|
|
**kwargs: Any,
|
|
) -> StructuredResponse:
|
|
"""Record the structured-output call and optionally fail first."""
|
|
self.recorded_structured_messages.append(list(messages))
|
|
if self._fail_structured_output_times > 0:
|
|
self._fail_structured_output_times -= 1
|
|
raise RuntimeError("simulated compression overflow")
|
|
return await super()._call_api_with_structured_output(
|
|
model_name,
|
|
messages,
|
|
structured_model,
|
|
**kwargs,
|
|
)
|
|
|
|
|
|
def _has_instruction_hint(
|
|
messages: list[Msg],
|
|
instructions: HintBlock,
|
|
) -> bool:
|
|
"""Return True if messages contain instructions as an assistant hint."""
|
|
for msg in messages:
|
|
if msg.role != "assistant":
|
|
continue
|
|
for hint_block in msg.get_content_blocks("hint"):
|
|
if hint_block.id == instructions.id:
|
|
return hint_block.hint == instructions.hint
|
|
return False
|
|
|
|
|
|
class ContextCompressionTest(IsolatedAsyncioTestCase):
|
|
"""The template test case."""
|
|
|
|
async def asyncSetUp(self) -> None:
|
|
"""The async setup method."""
|
|
|
|
async def test_split_function(self) -> None:
|
|
"""The template test."""
|
|
agent = Agent(
|
|
name="Friday",
|
|
system_prompt="".join(["0" for _ in range(60 * 4)]),
|
|
model=MockModel(),
|
|
context_config=ContextConfig(
|
|
trigger_ratio=0.8,
|
|
reserve_ratio=0.1,
|
|
),
|
|
state=AgentState(
|
|
session_id="123",
|
|
context=[
|
|
UserMsg(
|
|
"User",
|
|
"".join(["1" for _ in range(30 * 4)]),
|
|
id="1",
|
|
),
|
|
AssistantMsg(
|
|
"Friday",
|
|
"".join(["2" for _ in range(10 * 4)]),
|
|
id="2",
|
|
),
|
|
UserMsg(
|
|
"User",
|
|
"".join(["3" for _ in range(10 * 4)]),
|
|
id="3",
|
|
),
|
|
],
|
|
),
|
|
toolkit=Toolkit(),
|
|
)
|
|
|
|
# When the length of last two messages is exactly appropriate
|
|
to_compress, to_reserve = await agent._split_context_for_compression(
|
|
to_reserved_tokens=80,
|
|
tools=[],
|
|
)
|
|
|
|
self.assertListEqual(
|
|
[_.id for _ in to_compress],
|
|
["1"],
|
|
)
|
|
self.assertListEqual(
|
|
[_.id for _ in to_reserve],
|
|
["2", "3"],
|
|
)
|
|
|
|
# When one message is in the dividing line
|
|
agent.state.context = [
|
|
UserMsg("User", "".join(["2" for _ in range(30 * 4)]), id="1"),
|
|
AssistantMsg(
|
|
"Friday",
|
|
"".join(["3" for _ in range(15 * 4)]),
|
|
id="2",
|
|
),
|
|
UserMsg("User", "".join(["3" for _ in range(10 * 4)]), id="3"),
|
|
]
|
|
|
|
to_compress, to_reserve = await agent._split_context_for_compression(
|
|
to_reserved_tokens=80,
|
|
tools=[],
|
|
)
|
|
self.assertListEqual(
|
|
[_.id for _ in to_compress],
|
|
["1", "2"],
|
|
)
|
|
self.assertListEqual(
|
|
[_.id for _ in to_reserve],
|
|
["3"],
|
|
)
|
|
|
|
# When compress all messages
|
|
agent.state.context = [
|
|
UserMsg("User", "".join(["2" for _ in range(30 * 4)]), id="1"),
|
|
AssistantMsg(
|
|
"Friday",
|
|
"".join(["3" for _ in range(15 * 4)]),
|
|
id="2",
|
|
),
|
|
UserMsg("User", "".join(["3" for _ in range(30 * 4)]), id="3"),
|
|
]
|
|
to_compress, to_reserve = await agent._split_context_for_compression(
|
|
to_reserved_tokens=80,
|
|
tools=[],
|
|
)
|
|
self.assertListEqual(
|
|
[_.id for _ in to_compress],
|
|
["1", "2", "3"],
|
|
)
|
|
self.assertListEqual(
|
|
[_.id for _ in to_reserve],
|
|
[],
|
|
)
|
|
|
|
# When the boundary message has multiple blocks
|
|
agent.state.context = [
|
|
UserMsg("User", "".join(["a" for _ in range(30 * 4)]), id="1"),
|
|
AssistantMsg(
|
|
"Friday",
|
|
[
|
|
TextBlock(
|
|
text="".join(["b" for _ in range(10 * 4)]),
|
|
id="b",
|
|
),
|
|
TextBlock(
|
|
text="".join(["c" for _ in range(10 * 4)]),
|
|
id="c",
|
|
),
|
|
],
|
|
id="2",
|
|
),
|
|
UserMsg("User", "".join(["d" for _ in range(10 * 4)]), id="3"),
|
|
]
|
|
to_compress, to_reserve = await agent._split_context_for_compression(
|
|
to_reserved_tokens=80,
|
|
tools=[],
|
|
)
|
|
self.assertListEqual(
|
|
[_.model_dump() for _ in to_compress],
|
|
[
|
|
{
|
|
"id": "1",
|
|
"created_at": AnyString(),
|
|
"finished_at": AnyString(),
|
|
"name": "User",
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"id": AnyString(),
|
|
"type": "text",
|
|
"text": "a" * 120,
|
|
},
|
|
],
|
|
"metadata": {},
|
|
"usage": None,
|
|
},
|
|
{
|
|
"id": "2",
|
|
"created_at": AnyString(),
|
|
"finished_at": None,
|
|
"name": "Friday",
|
|
"role": "assistant",
|
|
"content": [
|
|
{
|
|
"id": "b",
|
|
"text": "b" * 40,
|
|
"type": "text",
|
|
},
|
|
],
|
|
"metadata": {},
|
|
"usage": None,
|
|
},
|
|
],
|
|
)
|
|
self.assertListEqual(
|
|
[_.model_dump() for _ in to_reserve],
|
|
[
|
|
{
|
|
"id": "2",
|
|
"created_at": AnyString(),
|
|
"finished_at": None,
|
|
"name": "Friday",
|
|
"role": "assistant",
|
|
"content": [
|
|
{
|
|
"id": AnyString(),
|
|
"text": "c" * 40,
|
|
"type": "text",
|
|
},
|
|
],
|
|
"metadata": {},
|
|
"usage": None,
|
|
},
|
|
{
|
|
"id": "3",
|
|
"created_at": AnyString(),
|
|
"finished_at": AnyString(),
|
|
"name": "User",
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"id": AnyString(),
|
|
"type": "text",
|
|
"text": "d" * 40,
|
|
},
|
|
],
|
|
"metadata": {},
|
|
"usage": None,
|
|
},
|
|
],
|
|
)
|
|
|
|
# When the boundary message has multiple blocks
|
|
# Cannot leave any blocks
|
|
agent.state.context = [
|
|
UserMsg("User", "".join(["a" for _ in range(30 * 4)]), id="1"),
|
|
AssistantMsg(
|
|
"Friday",
|
|
[
|
|
TextBlock(
|
|
text="".join(["b" for _ in range(10 * 4)]),
|
|
id="b",
|
|
),
|
|
TextBlock(
|
|
text="".join(["c" for _ in range(15 * 4)]),
|
|
id="c",
|
|
),
|
|
],
|
|
id="2",
|
|
),
|
|
UserMsg("User", "".join(["d" for _ in range(10 * 4)]), id="3"),
|
|
]
|
|
to_compress, to_reserve = await agent._split_context_for_compression(
|
|
to_reserved_tokens=80,
|
|
tools=[],
|
|
)
|
|
self.assertListEqual(
|
|
[_.model_dump() for _ in to_compress],
|
|
[
|
|
{
|
|
"id": "1",
|
|
"created_at": AnyString(),
|
|
"finished_at": AnyString(),
|
|
"name": "User",
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"id": AnyString(),
|
|
"type": "text",
|
|
"text": "a" * 120,
|
|
},
|
|
],
|
|
"metadata": {},
|
|
"usage": None,
|
|
},
|
|
{
|
|
"id": "2",
|
|
"created_at": AnyString(),
|
|
"finished_at": None,
|
|
"name": "Friday",
|
|
"role": "assistant",
|
|
"content": [
|
|
{
|
|
"id": "b",
|
|
"text": "b" * 40,
|
|
"type": "text",
|
|
},
|
|
{
|
|
"id": "c",
|
|
"text": "c" * 60,
|
|
"type": "text",
|
|
},
|
|
],
|
|
"metadata": {},
|
|
"usage": None,
|
|
},
|
|
],
|
|
)
|
|
self.assertListEqual(
|
|
[_.model_dump() for _ in to_reserve],
|
|
[
|
|
{
|
|
"id": "3",
|
|
"created_at": AnyString(),
|
|
"finished_at": AnyString(),
|
|
"name": "User",
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"id": AnyString(),
|
|
"type": "text",
|
|
"text": "d" * 40,
|
|
},
|
|
],
|
|
"metadata": {},
|
|
"usage": None,
|
|
},
|
|
],
|
|
)
|
|
|
|
# Leave the last block of the boundary message
|
|
agent.state.context = [
|
|
UserMsg("User", "".join(["a" for _ in range(30 * 4)]), id="1"),
|
|
AssistantMsg(
|
|
"Friday",
|
|
[
|
|
TextBlock(
|
|
text="".join(["b" for _ in range(10 * 4)]),
|
|
id="b",
|
|
),
|
|
TextBlock(
|
|
text="".join(["c" for _ in range(5 * 4)]),
|
|
id="c",
|
|
),
|
|
],
|
|
id="2",
|
|
),
|
|
UserMsg("User", "".join(["d" for _ in range(10 * 4)]), id="3"),
|
|
]
|
|
to_compress, to_reserve = await agent._split_context_for_compression(
|
|
to_reserved_tokens=80,
|
|
tools=[],
|
|
)
|
|
self.assertListEqual(
|
|
[_.model_dump() for _ in to_compress],
|
|
[
|
|
{
|
|
"id": "1",
|
|
"created_at": AnyString(),
|
|
"finished_at": AnyString(),
|
|
"name": "User",
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"id": AnyString(),
|
|
"type": "text",
|
|
"text": "a" * 120,
|
|
},
|
|
],
|
|
"metadata": {},
|
|
"usage": None,
|
|
},
|
|
{
|
|
"id": "2",
|
|
"created_at": AnyString(),
|
|
"finished_at": None,
|
|
"name": "Friday",
|
|
"role": "assistant",
|
|
"content": [
|
|
{
|
|
"id": "b",
|
|
"text": "b" * 40,
|
|
"type": "text",
|
|
},
|
|
],
|
|
"metadata": {},
|
|
"usage": None,
|
|
},
|
|
],
|
|
)
|
|
self.assertListEqual(
|
|
[_.model_dump() for _ in to_reserve],
|
|
[
|
|
{
|
|
"id": "2",
|
|
"created_at": AnyString(),
|
|
"finished_at": None,
|
|
"name": "Friday",
|
|
"role": "assistant",
|
|
"content": [
|
|
{
|
|
"id": "c",
|
|
"text": "c" * 20,
|
|
"type": "text",
|
|
},
|
|
],
|
|
"metadata": {},
|
|
"usage": None,
|
|
},
|
|
{
|
|
"id": "3",
|
|
"created_at": AnyString(),
|
|
"finished_at": AnyString(),
|
|
"name": "User",
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"id": AnyString(),
|
|
"type": "text",
|
|
"text": "d" * 40,
|
|
},
|
|
],
|
|
"metadata": {},
|
|
"usage": None,
|
|
},
|
|
],
|
|
)
|
|
|
|
# Leave all the blocks
|
|
agent.state.context = [
|
|
UserMsg("User", "".join(["a" for _ in range(30 * 4)]), id="1"),
|
|
AssistantMsg(
|
|
"Friday",
|
|
[
|
|
TextBlock(
|
|
text="".join(["b" for _ in range(5 * 4)]),
|
|
id="b",
|
|
),
|
|
TextBlock(
|
|
text="".join(["c" for _ in range(5 * 4)]),
|
|
id="c",
|
|
),
|
|
],
|
|
id="2",
|
|
),
|
|
UserMsg("User", "".join(["d" for _ in range(10 * 4)]), id="3"),
|
|
]
|
|
to_compress, to_reserve = await agent._split_context_for_compression(
|
|
to_reserved_tokens=80,
|
|
tools=[],
|
|
)
|
|
self.assertListEqual(
|
|
[_.model_dump() for _ in to_compress],
|
|
[
|
|
{
|
|
"id": "1",
|
|
"created_at": AnyString(),
|
|
"finished_at": AnyString(),
|
|
"name": "User",
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"id": AnyString(),
|
|
"type": "text",
|
|
"text": "a" * 120,
|
|
},
|
|
],
|
|
"metadata": {},
|
|
"usage": None,
|
|
},
|
|
],
|
|
)
|
|
self.assertListEqual(
|
|
[_.model_dump() for _ in to_reserve],
|
|
[
|
|
{
|
|
"id": "2",
|
|
"created_at": AnyString(),
|
|
"finished_at": None,
|
|
"name": "Friday",
|
|
"role": "assistant",
|
|
"content": [
|
|
{
|
|
"id": "b",
|
|
"text": "b" * 20,
|
|
"type": "text",
|
|
},
|
|
{
|
|
"id": "c",
|
|
"text": "c" * 20,
|
|
"type": "text",
|
|
},
|
|
],
|
|
"metadata": {},
|
|
"usage": None,
|
|
},
|
|
{
|
|
"id": "3",
|
|
"created_at": AnyString(),
|
|
"finished_at": AnyString(),
|
|
"name": "User",
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"id": AnyString(),
|
|
"type": "text",
|
|
"text": "d" * 40,
|
|
},
|
|
],
|
|
"metadata": {},
|
|
"usage": None,
|
|
},
|
|
],
|
|
)
|
|
|
|
# Leave all the messages
|
|
agent.state.context = [
|
|
AssistantMsg(
|
|
"Friday",
|
|
[
|
|
TextBlock(
|
|
text="".join(["b" for _ in range(5 * 4)]),
|
|
id="b",
|
|
),
|
|
TextBlock(
|
|
text="".join(["c" for _ in range(5 * 4)]),
|
|
id="c",
|
|
),
|
|
],
|
|
id="2",
|
|
),
|
|
UserMsg("User", "".join(["d" for _ in range(10 * 4)]), id="3"),
|
|
]
|
|
to_compress, to_reserve = await agent._split_context_for_compression(
|
|
to_reserved_tokens=80,
|
|
tools=[],
|
|
)
|
|
self.assertListEqual(
|
|
[_.model_dump() for _ in to_compress],
|
|
[],
|
|
)
|
|
self.assertListEqual(
|
|
[_.model_dump() for _ in to_reserve],
|
|
[
|
|
{
|
|
"id": "2",
|
|
"created_at": AnyString(),
|
|
"finished_at": None,
|
|
"name": "Friday",
|
|
"role": "assistant",
|
|
"content": [
|
|
{
|
|
"id": "b",
|
|
"text": "b" * 20,
|
|
"type": "text",
|
|
},
|
|
{
|
|
"id": "c",
|
|
"text": "c" * 20,
|
|
"type": "text",
|
|
},
|
|
],
|
|
"metadata": {},
|
|
"usage": None,
|
|
},
|
|
{
|
|
"id": "3",
|
|
"created_at": AnyString(),
|
|
"finished_at": AnyString(),
|
|
"name": "User",
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"id": AnyString(),
|
|
"type": "text",
|
|
"text": "d" * 40,
|
|
},
|
|
],
|
|
"metadata": {},
|
|
"usage": None,
|
|
},
|
|
],
|
|
)
|
|
|
|
async def test_context_compression(self) -> None:
|
|
"""Test the context compression logic."""
|
|
model = MockModel(context_size=100)
|
|
agent = Agent(
|
|
name="Friday",
|
|
system_prompt="".join(["0" for _ in range(20 * 4)]),
|
|
model=model,
|
|
context_config=ContextConfig(
|
|
trigger_ratio=0.7,
|
|
reserve_ratio=0.4,
|
|
),
|
|
state=AgentState(
|
|
session_id="123",
|
|
context=[
|
|
UserMsg(
|
|
"User",
|
|
"".join(["1" for _ in range(30 * 4)]),
|
|
id="1",
|
|
),
|
|
AssistantMsg(
|
|
"Friday",
|
|
"".join(["2" for _ in range(10 * 4)]),
|
|
id="2",
|
|
),
|
|
UserMsg(
|
|
"User",
|
|
"".join(["3" for _ in range(10 * 4)]),
|
|
id="3",
|
|
),
|
|
],
|
|
),
|
|
toolkit=Toolkit(),
|
|
)
|
|
|
|
model.set_structured_response(
|
|
StructuredResponse(
|
|
content={
|
|
"task_overview": "1",
|
|
"current_state": "2",
|
|
"important_discoveries": "3",
|
|
"next_steps": "4",
|
|
"context_to_preserve": "5",
|
|
},
|
|
),
|
|
)
|
|
|
|
await agent.compress_context()
|
|
|
|
self.assertEqual(
|
|
agent.state.summary,
|
|
"""<system-info>Here is a summary of your previous work
|
|
# Task Overview
|
|
1
|
|
|
|
# Current State
|
|
2
|
|
|
|
# Important Discoveries
|
|
3
|
|
|
|
# Next Steps
|
|
4
|
|
|
|
# Context to Preserve
|
|
5</system-info>""",
|
|
)
|
|
|
|
self.assertListEqual(
|
|
[_.model_dump() for _ in agent.state.context],
|
|
[
|
|
{
|
|
"id": "2",
|
|
"created_at": AnyString(),
|
|
"finished_at": None,
|
|
"name": "Friday",
|
|
"role": "assistant",
|
|
"content": [
|
|
{
|
|
"id": AnyString(),
|
|
"text": "2" * 40,
|
|
"type": "text",
|
|
},
|
|
],
|
|
"metadata": {},
|
|
"usage": None,
|
|
},
|
|
{
|
|
"id": "3",
|
|
"created_at": AnyString(),
|
|
"finished_at": AnyString(),
|
|
"name": "User",
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"id": AnyString(),
|
|
"text": "3" * 40,
|
|
"type": "text",
|
|
},
|
|
],
|
|
"metadata": {},
|
|
"usage": None,
|
|
},
|
|
],
|
|
)
|
|
|
|
async def test_context_compression_clears_evicted_read_cache(self) -> None:
|
|
"""Read cache is cleared when its Read block is compressed out."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
file_path = os.path.join(temp_dir, "test.txt")
|
|
with open(file_path, "w", encoding="utf-8") as f:
|
|
f.write("content\n")
|
|
|
|
model = MockModel(context_size=100)
|
|
agent = Agent(
|
|
name="Friday",
|
|
system_prompt="".join(["0" for _ in range(20 * 4)]),
|
|
model=model,
|
|
context_config=ContextConfig(
|
|
trigger_ratio=0.7,
|
|
reserve_ratio=0.4,
|
|
),
|
|
state=AgentState(
|
|
session_id="123",
|
|
context=[
|
|
AssistantMsg(
|
|
"Friday",
|
|
[
|
|
ToolCallBlock(
|
|
id="read-call-1",
|
|
name="Read",
|
|
input=json.dumps(
|
|
{"file_path": file_path},
|
|
),
|
|
),
|
|
],
|
|
id="1",
|
|
),
|
|
UserMsg(
|
|
"User",
|
|
"".join(["2" for _ in range(30 * 4)]),
|
|
id="2",
|
|
),
|
|
UserMsg(
|
|
"User",
|
|
"".join(["3" for _ in range(10 * 4)]),
|
|
id="3",
|
|
),
|
|
],
|
|
),
|
|
toolkit=Toolkit(),
|
|
)
|
|
await agent.state.tool_context.cache_file(
|
|
file_path=file_path,
|
|
lines=["content\n"],
|
|
)
|
|
self.assertIsNotNone(
|
|
await agent.state.tool_context.get_cache(file_path),
|
|
)
|
|
|
|
model.set_structured_response(
|
|
StructuredResponse(
|
|
content={
|
|
"task_overview": "1",
|
|
"current_state": "2",
|
|
"important_discoveries": "3",
|
|
"next_steps": "4",
|
|
"context_to_preserve": "5",
|
|
},
|
|
),
|
|
)
|
|
|
|
await agent.compress_context()
|
|
|
|
self.assertIsNone(
|
|
await agent.state.tool_context.get_cache(file_path),
|
|
)
|
|
|
|
async def test_context_compression_keeps_reserved_read_cache(
|
|
self,
|
|
) -> None:
|
|
"""Read cache is kept when the same file is still read in context."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
file_path = os.path.join(temp_dir, "test.txt")
|
|
with open(file_path, "w", encoding="utf-8") as f:
|
|
f.write("content\n")
|
|
|
|
model = MockModel(context_size=100)
|
|
agent = Agent(
|
|
name="Friday",
|
|
system_prompt="".join(["0" for _ in range(20 * 4)]),
|
|
model=model,
|
|
context_config=ContextConfig(
|
|
trigger_ratio=0.7,
|
|
reserve_ratio=0.6,
|
|
),
|
|
state=AgentState(
|
|
session_id="123",
|
|
context=[
|
|
AssistantMsg(
|
|
"Friday",
|
|
[
|
|
ToolCallBlock(
|
|
id="read-call-1",
|
|
name="Read",
|
|
input=json.dumps(
|
|
{"file_path": file_path},
|
|
),
|
|
),
|
|
],
|
|
id="1",
|
|
),
|
|
UserMsg(
|
|
"User",
|
|
"".join(["2" for _ in range(30 * 4)]),
|
|
id="2",
|
|
),
|
|
AssistantMsg(
|
|
"Friday",
|
|
[
|
|
ToolCallBlock(
|
|
id="read-call-2",
|
|
name="Read",
|
|
input=json.dumps(
|
|
{"file_path": file_path},
|
|
),
|
|
),
|
|
],
|
|
id="3",
|
|
),
|
|
],
|
|
),
|
|
toolkit=Toolkit(),
|
|
)
|
|
await agent.state.tool_context.cache_file(
|
|
file_path=file_path,
|
|
lines=["content\n"],
|
|
)
|
|
|
|
model.set_structured_response(
|
|
StructuredResponse(
|
|
content={
|
|
"task_overview": "1",
|
|
"current_state": "2",
|
|
"important_discoveries": "3",
|
|
"next_steps": "4",
|
|
"context_to_preserve": "5",
|
|
},
|
|
),
|
|
)
|
|
|
|
await agent.compress_context()
|
|
|
|
self.assertIsNotNone(
|
|
await agent.state.tool_context.get_cache(file_path),
|
|
)
|
|
|
|
async def test_context_compression_clears_unreferenced_read_cache(
|
|
self,
|
|
) -> None:
|
|
"""Read cache is cleared when no reserved Read references it."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
file_path = os.path.join(temp_dir, "test.txt")
|
|
with open(file_path, "w", encoding="utf-8") as f:
|
|
f.write("content\n")
|
|
|
|
model = MockModel(context_size=100)
|
|
agent = Agent(
|
|
name="Friday",
|
|
system_prompt="".join(["0" for _ in range(20 * 4)]),
|
|
model=model,
|
|
context_config=ContextConfig(
|
|
trigger_ratio=0.7,
|
|
reserve_ratio=0.4,
|
|
),
|
|
state=AgentState(
|
|
session_id="123",
|
|
context=[
|
|
UserMsg(
|
|
"User",
|
|
"".join(["2" for _ in range(60 * 4)]),
|
|
id="1",
|
|
),
|
|
UserMsg(
|
|
"User",
|
|
"".join(["3" for _ in range(30 * 4)]),
|
|
id="2",
|
|
),
|
|
],
|
|
),
|
|
toolkit=Toolkit(),
|
|
)
|
|
await agent.state.tool_context.cache_file(
|
|
file_path=file_path,
|
|
lines=["content\n"],
|
|
)
|
|
|
|
model.set_structured_response(
|
|
StructuredResponse(
|
|
content={
|
|
"task_overview": "1",
|
|
"current_state": "2",
|
|
"important_discoveries": "3",
|
|
"next_steps": "4",
|
|
"context_to_preserve": "5",
|
|
},
|
|
),
|
|
)
|
|
|
|
await agent.compress_context()
|
|
|
|
self.assertIsNone(
|
|
await agent.state.tool_context.get_cache(file_path),
|
|
)
|
|
|
|
async def test_context_compression_injects_instructions_as_hint(
|
|
self,
|
|
) -> None:
|
|
"""Instructions are injected as a HintBlock only for compression."""
|
|
model = RecordingStructuredMockModel(context_size=100)
|
|
agent = Agent(
|
|
name="Friday",
|
|
system_prompt="".join(["0" for _ in range(20 * 4)]),
|
|
model=model,
|
|
context_config=ContextConfig(
|
|
trigger_ratio=0.7,
|
|
reserve_ratio=0.4,
|
|
),
|
|
state=AgentState(
|
|
session_id="123",
|
|
context=[
|
|
UserMsg(
|
|
"User",
|
|
"".join(["1" for _ in range(30 * 4)]),
|
|
id="1",
|
|
),
|
|
AssistantMsg(
|
|
"Friday",
|
|
"".join(["2" for _ in range(10 * 4)]),
|
|
id="2",
|
|
),
|
|
UserMsg(
|
|
"User",
|
|
"".join(["3" for _ in range(10 * 4)]),
|
|
id="3",
|
|
),
|
|
],
|
|
),
|
|
toolkit=Toolkit(),
|
|
)
|
|
|
|
model.set_structured_response(
|
|
StructuredResponse(
|
|
content={
|
|
"task_overview": "1",
|
|
"current_state": "2",
|
|
"important_discoveries": "3",
|
|
"next_steps": "4",
|
|
"context_to_preserve": "5",
|
|
},
|
|
),
|
|
)
|
|
instructions = HintBlock(
|
|
hint="Keep user requirements and file paths.",
|
|
source="user",
|
|
)
|
|
|
|
await agent.compress_context(instructions=instructions)
|
|
|
|
self.assertEqual(len(model.recorded_structured_messages), 1)
|
|
self.assertTrue(
|
|
_has_instruction_hint(
|
|
model.recorded_structured_messages[0],
|
|
instructions,
|
|
),
|
|
)
|
|
self.assertFalse(
|
|
any(msg.get_content_blocks("hint") for msg in agent.state.context),
|
|
)
|
|
|
|
async def test_context_compression_overflow_retry_keeps_instructions(
|
|
self,
|
|
) -> None:
|
|
"""Overflow retry preserves instructions when rebuilding messages."""
|
|
model = RecordingStructuredMockModel(
|
|
context_size=100,
|
|
fail_structured_output_times=1,
|
|
force_compression_overflow=True,
|
|
)
|
|
agent = Agent(
|
|
name="Friday",
|
|
system_prompt="".join(["0" for _ in range(20 * 4)]),
|
|
model=model,
|
|
context_config=ContextConfig(
|
|
trigger_ratio=0.7,
|
|
reserve_ratio=0.4,
|
|
),
|
|
state=AgentState(
|
|
session_id="123",
|
|
context=[
|
|
UserMsg(
|
|
"User",
|
|
"".join(["1" for _ in range(30 * 4)]),
|
|
id="1",
|
|
),
|
|
AssistantMsg(
|
|
"Friday",
|
|
"".join(["2" for _ in range(10 * 4)]),
|
|
id="2",
|
|
),
|
|
UserMsg(
|
|
"User",
|
|
"".join(["3" for _ in range(10 * 4)]),
|
|
id="3",
|
|
),
|
|
],
|
|
),
|
|
toolkit=Toolkit(),
|
|
)
|
|
|
|
model.set_structured_response(
|
|
StructuredResponse(
|
|
content={
|
|
"task_overview": "1",
|
|
"current_state": "2",
|
|
"important_discoveries": "3",
|
|
"next_steps": "4",
|
|
"context_to_preserve": "5",
|
|
},
|
|
),
|
|
)
|
|
instructions = HintBlock(
|
|
hint="Keep the user's original success criteria.",
|
|
source="user",
|
|
)
|
|
|
|
await agent.compress_context(instructions=instructions)
|
|
|
|
self.assertEqual(len(model.recorded_structured_messages), 2)
|
|
self.assertTrue(
|
|
_has_instruction_hint(
|
|
model.recorded_structured_messages[-1],
|
|
instructions,
|
|
),
|
|
)
|
|
|
|
async def asyncTearDown(self) -> None:
|
|
"""The async teardown method."""
|