LZH-YS1998
100 lines
3.3 KiB
Python
100 lines
3.3 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import unittest
|
|
|
|
from opc.layer5_memory.history_compactor import HistoryCompactor
|
|
|
|
|
|
class _OverflowLLM:
|
|
def __init__(self) -> None:
|
|
self.calls = 0
|
|
|
|
def is_context_overflow_error(self, error: Exception) -> bool:
|
|
return "prompt too long" in str(error).lower()
|
|
|
|
async def simple_chat(self, prompt: str, system: str | None = None, task_type: str | None = None) -> str:
|
|
_ = (system, task_type)
|
|
self.calls += 1
|
|
payload = json.loads(prompt)
|
|
messages = list(payload.get("messages", []) or [])
|
|
if self.calls == 1:
|
|
raise RuntimeError("prompt too long")
|
|
return json.dumps(
|
|
{
|
|
"history_summary": f"summarized {len(messages)} messages",
|
|
"memory_summary": "## Primary Goal\n- Continue safely.",
|
|
},
|
|
ensure_ascii=False,
|
|
)
|
|
|
|
|
|
class _SingleMessageOverflowLLM:
|
|
def __init__(self) -> None:
|
|
self.calls = 0
|
|
|
|
def is_context_overflow_error(self, error: Exception) -> bool:
|
|
return "prompt too long" in str(error).lower()
|
|
|
|
async def simple_chat(self, prompt: str, system: str | None = None, task_type: str | None = None) -> str:
|
|
_ = (system, task_type)
|
|
self.calls += 1
|
|
if len(prompt) > 4500:
|
|
raise RuntimeError("prompt too long")
|
|
payload = json.loads(prompt)
|
|
messages = list(payload.get("messages", []) or [])
|
|
return json.dumps(
|
|
{
|
|
"history_summary": f"single-message summarized {len(messages)} messages",
|
|
"memory_summary": "## Primary Goal\n- Continue safely.",
|
|
},
|
|
ensure_ascii=False,
|
|
)
|
|
|
|
|
|
class HistoryCompactorRetryTests(unittest.IsolatedAsyncioTestCase):
|
|
async def test_simple_chat_retry_truncates_messages_after_overflow(self) -> None:
|
|
llm = _OverflowLLM()
|
|
compactor = HistoryCompactor(
|
|
llm=llm,
|
|
store=object(),
|
|
memory_manager=object(),
|
|
)
|
|
result = await compactor._summarize_session(
|
|
project_id="proj1",
|
|
session_id="sess1",
|
|
messages=[
|
|
{"role": "user", "content": "a"},
|
|
{"role": "assistant", "content": "b"},
|
|
{"role": "user", "content": "c"},
|
|
{"role": "assistant", "content": "d"},
|
|
{"role": "user", "content": "e"},
|
|
],
|
|
existing_memory="",
|
|
existing_summary="",
|
|
)
|
|
|
|
self.assertEqual(llm.calls, 2)
|
|
self.assertIn("summarized", result["history_summary"])
|
|
self.assertIn("Primary Goal", result["memory_summary"])
|
|
|
|
async def test_simple_chat_retry_truncates_single_huge_message_content(self) -> None:
|
|
llm = _SingleMessageOverflowLLM()
|
|
compactor = HistoryCompactor(
|
|
llm=llm,
|
|
store=object(),
|
|
memory_manager=object(),
|
|
)
|
|
result = await compactor._summarize_session(
|
|
project_id="proj1",
|
|
session_id="sess1",
|
|
messages=[
|
|
{"role": "assistant", "content": "x" * 5000},
|
|
],
|
|
existing_memory="",
|
|
existing_summary="",
|
|
)
|
|
|
|
self.assertGreaterEqual(llm.calls, 2)
|
|
self.assertIn("single-message summarized", result["history_summary"])
|