Initial commit
This commit is contained in:
LZH-YS1998
@@ -0,0 +1,569 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
from opc.core.models import AgentMemorySnapshotRecord, SessionMemorySnapshotRecord, Task, TaskStatus
|
||||
from opc.database.store import OPCStore
|
||||
from opc.layer1_perception.context_assembler import ContextAssembler
|
||||
from opc.layer5_memory.history_compactor import HistoryCompactor
|
||||
from opc.layer5_memory.memory_manager import MemoryManager
|
||||
|
||||
|
||||
class _StubLLM:
|
||||
class _Config:
|
||||
max_tokens = 1024
|
||||
|
||||
config = _Config()
|
||||
|
||||
def get_context_window(self, task_type: str | None = None, model: str | None = None) -> int:
|
||||
_ = (task_type, model)
|
||||
return 100
|
||||
|
||||
def count_input_tokens(
|
||||
self,
|
||||
messages: list[dict[str, str]],
|
||||
tools: list[dict[str, object]] | None = None,
|
||||
task_type: str | None = None,
|
||||
model: str | None = None,
|
||||
) -> int:
|
||||
_ = (tools, task_type, model)
|
||||
total = 0
|
||||
for message in messages:
|
||||
total += max(1, len(str(message.get("content", "")).strip()) // 20)
|
||||
return total
|
||||
|
||||
async def simple_chat(self, prompt: str, system: str | None = None, task_type: str | None = None) -> str:
|
||||
_ = (prompt, task_type)
|
||||
if system and "persisted session history" in system:
|
||||
return json.dumps(
|
||||
{
|
||||
"history_summary": "Session history summary before restart.",
|
||||
"memory_summary": (
|
||||
"## Primary Goal\n"
|
||||
"- Continue the project.\n\n"
|
||||
"## Active Rules\n"
|
||||
"- Keep memory concise.\n\n"
|
||||
"## Key Progress\n"
|
||||
"- Session summary stored.\n\n"
|
||||
"## Current State\n"
|
||||
"- Waiting for the latest tail.\n\n"
|
||||
"## Open Risks\n"
|
||||
"- Re-check raw history if needed."
|
||||
),
|
||||
},
|
||||
ensure_ascii=False,
|
||||
)
|
||||
if system and "employee-level process history" in system:
|
||||
return json.dumps(
|
||||
{
|
||||
"history_summary": "Agent history summary before restart.",
|
||||
"memory_summary": (
|
||||
"## Effective Patterns\n"
|
||||
"- Leave concise handoffs.\n\n"
|
||||
"## Watchouts\n"
|
||||
"- Avoid duplicate edits.\n\n"
|
||||
"## Current Progress\n"
|
||||
"- Task completed.\n\n"
|
||||
"## Current State\n"
|
||||
"- Waiting for final reflection."
|
||||
),
|
||||
},
|
||||
ensure_ascii=False,
|
||||
)
|
||||
if system and "finalizing employee memory" in system:
|
||||
return json.dumps(
|
||||
{
|
||||
"summary_text": "Final employee memory created.",
|
||||
"memory_text": (
|
||||
"## Effective Patterns\n"
|
||||
"- Leave concise handoffs.\n\n"
|
||||
"## Watchouts\n"
|
||||
"- Avoid duplicate edits.\n\n"
|
||||
"## Preferred Tools\n"
|
||||
"- Prefer targeted validation.\n\n"
|
||||
"## Reviewer Preferences\n"
|
||||
"- Surface exact artifacts.\n\n"
|
||||
"## Reusable Checklist\n"
|
||||
"- Include validation evidence."
|
||||
),
|
||||
"metadata": {
|
||||
"effective_patterns": ["Leave concise handoffs."],
|
||||
"watchouts": ["Avoid duplicate edits."],
|
||||
"preferred_tools": ["Prefer targeted validation."],
|
||||
"reviewer_preferences": ["Surface exact artifacts."],
|
||||
"reusable_checklist": ["Include validation evidence."],
|
||||
},
|
||||
},
|
||||
ensure_ascii=False,
|
||||
)
|
||||
raise AssertionError(f"Unexpected prompt type: {system}")
|
||||
|
||||
|
||||
class SessionContextCompressionTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_build_session_prompt_context_returns_full_visible_history(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
root = Path(tmpdir)
|
||||
store = OPCStore(root / "tasks.db")
|
||||
await store.initialize()
|
||||
memory = MemoryManager(root, "proj1", store=store)
|
||||
session_id = "session-1"
|
||||
|
||||
for idx in range(6):
|
||||
role = "user" if idx % 2 == 0 else "assistant"
|
||||
payload = (
|
||||
f"message {idx} " +
|
||||
"alpha beta gamma delta epsilon zeta eta theta iota kappa lambda mu"
|
||||
)
|
||||
await memory.append_session_message(session_id=session_id, role=role, text=payload)
|
||||
|
||||
context = await memory.build_session_prompt_context(session_id)
|
||||
history_messages = await memory.build_session_history_messages(session_id)
|
||||
|
||||
self.assertIn("Current Session History", context)
|
||||
self.assertIn("message 0", context)
|
||||
self.assertIn("message 5", context)
|
||||
self.assertEqual(len(history_messages), 6)
|
||||
self.assertEqual(history_messages[0]["role"], "user")
|
||||
self.assertIn("message 0", history_messages[0]["content"])
|
||||
self.assertEqual(history_messages[-1]["role"], "assistant")
|
||||
self.assertIn("message 5", history_messages[-1]["content"])
|
||||
|
||||
async def test_build_memory_context_includes_global_project_and_session_layers(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
root = Path(tmpdir)
|
||||
store = OPCStore(root / "tasks.db")
|
||||
await store.initialize()
|
||||
memory = MemoryManager(root, "proj1", store=store)
|
||||
session_id = "session-layers"
|
||||
|
||||
memory.save_memory("# Global rules\n- Use Chinese when appropriate.\n", project=False)
|
||||
memory.save_memory("# Project rules\n- Follow proj1 conventions.\n", project=True)
|
||||
await store.save_session_memory_snapshot(
|
||||
SessionMemorySnapshotRecord(
|
||||
project_id="proj1",
|
||||
session_id=session_id,
|
||||
summary_message_id="summary-1",
|
||||
source_boundary_message_id="msg-1",
|
||||
summary_text="session summary",
|
||||
memory_text="## Primary Goal\n- Finish the current ticket.",
|
||||
)
|
||||
)
|
||||
|
||||
context = await memory.build_memory_context(project_id="proj1", session_id=session_id)
|
||||
|
||||
self.assertIn("## Global Memory", context)
|
||||
self.assertIn("Use Chinese when appropriate", context)
|
||||
self.assertIn("## Project Memory (proj1)", context)
|
||||
self.assertIn("Follow proj1 conventions", context)
|
||||
self.assertIn("## Session Memory", context)
|
||||
self.assertIn("Finish the current ticket", context)
|
||||
|
||||
async def test_memory_manager_preserves_long_content_without_numeric_truncation(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
root = Path(tmpdir)
|
||||
store = OPCStore(root / "tasks.db")
|
||||
await store.initialize()
|
||||
memory = MemoryManager(root, "proj1", store=store)
|
||||
session_id = "session-long"
|
||||
|
||||
msg = await memory.append_session_message(
|
||||
session_id=session_id,
|
||||
role="assistant",
|
||||
text="",
|
||||
project_id="proj1",
|
||||
)
|
||||
assert msg is not None
|
||||
long_output = "very-long-output " * 300
|
||||
await memory.append_session_part(
|
||||
session_id,
|
||||
msg.message_id,
|
||||
"tool_output",
|
||||
{"tool_name": "test_tool", "output": long_output},
|
||||
)
|
||||
|
||||
context = await memory.build_session_prompt_context(session_id)
|
||||
history_messages = await memory.build_session_history_messages(session_id)
|
||||
compression_prompt = memory.get_compression_prompt(
|
||||
[{"role": "assistant", "content": long_output}],
|
||||
existing_memory="existing",
|
||||
)
|
||||
|
||||
self.assertIn(long_output.strip(), context)
|
||||
self.assertIn(long_output.strip(), history_messages[0]["content"])
|
||||
self.assertIn(long_output.strip(), compression_prompt)
|
||||
|
||||
async def test_restart_uses_summary_and_tail_after_compaction(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
root = Path(tmpdir)
|
||||
db_path = root / "tasks.db"
|
||||
store = OPCStore(db_path)
|
||||
await store.initialize()
|
||||
memory = MemoryManager(root, "proj1", store=store)
|
||||
memory.set_history_compactor(
|
||||
HistoryCompactor(
|
||||
llm=_StubLLM(),
|
||||
store=store,
|
||||
memory_manager=memory,
|
||||
compression_threshold=0.85,
|
||||
)
|
||||
)
|
||||
session_id = "session-restart"
|
||||
|
||||
for idx in range(6):
|
||||
role = "user" if idx % 2 == 0 else "assistant"
|
||||
payload = f"message {idx} " + ("alpha beta gamma delta " * 40)
|
||||
await memory.append_session_message(session_id=session_id, role=role, text=payload)
|
||||
|
||||
snapshot = await store.get_latest_session_memory_snapshot(session_id)
|
||||
self.assertIsNotNone(snapshot)
|
||||
assert snapshot is not None
|
||||
self.assertIn("## Primary Goal", snapshot.memory_text)
|
||||
|
||||
context = await memory.build_session_prompt_context(session_id)
|
||||
self.assertIn("## Session Memory", context)
|
||||
self.assertIn("Session history summary before restart", context)
|
||||
self.assertIn("message 5", context)
|
||||
|
||||
await store.close()
|
||||
|
||||
restarted_store = OPCStore(db_path)
|
||||
await restarted_store.initialize()
|
||||
restarted_memory = MemoryManager(root, "proj1", store=restarted_store)
|
||||
restarted_history = await restarted_memory.build_session_history_messages(session_id)
|
||||
restarted_context = await restarted_memory.build_session_prompt_context(session_id)
|
||||
|
||||
self.assertTrue(any("Session history summary before restart" in item["content"] for item in restarted_history))
|
||||
self.assertTrue(any("message 5" in item["content"] for item in restarted_history))
|
||||
self.assertFalse(any("message 0" in item["content"] for item in restarted_history))
|
||||
self.assertNotIn("message 0", restarted_context)
|
||||
|
||||
async def test_employee_history_tail_isolated_by_employee_id(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
root = Path(tmpdir)
|
||||
store = OPCStore(root / "tasks.db")
|
||||
await store.initialize()
|
||||
memory = MemoryManager(root, "proj1", store=store)
|
||||
session_id = "shared-session"
|
||||
|
||||
await memory.ensure_session(session_id, project_id="proj1", metadata={"role_id": "developer"})
|
||||
await memory.record_user_turn(
|
||||
session_id,
|
||||
"employee a asks a question",
|
||||
project_id="proj1",
|
||||
metadata={"employee_id": "emp-a", "role_id": "developer"},
|
||||
)
|
||||
await memory.record_assistant_turn(
|
||||
session_id,
|
||||
"employee a gets an answer",
|
||||
project_id="proj1",
|
||||
metadata={"employee_id": "emp-a", "role_id": "developer"},
|
||||
)
|
||||
await memory.record_user_turn(
|
||||
session_id,
|
||||
"employee b asks a different question",
|
||||
project_id="proj1",
|
||||
metadata={"employee_id": "emp-b", "role_id": "developer"},
|
||||
)
|
||||
|
||||
history_a = await memory.build_employee_history_tail_messages(
|
||||
project_id="proj1",
|
||||
session_id=session_id,
|
||||
employee_id="emp-a",
|
||||
)
|
||||
history_b = await memory.build_employee_history_tail_messages(
|
||||
project_id="proj1",
|
||||
session_id=session_id,
|
||||
employee_id="emp-b",
|
||||
)
|
||||
|
||||
self.assertEqual(len(history_a), 2)
|
||||
self.assertEqual(len(history_b), 1)
|
||||
self.assertTrue(all("employee a" in item["content"] for item in history_a))
|
||||
self.assertTrue(all("employee b" in item["content"] for item in history_b))
|
||||
|
||||
async def test_final_agent_memory_replaces_process_memory(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
root = Path(tmpdir)
|
||||
store = OPCStore(root / "tasks.db")
|
||||
await store.initialize()
|
||||
memory = MemoryManager(root, "proj1", store=store)
|
||||
memory.set_history_compactor(
|
||||
HistoryCompactor(
|
||||
llm=_StubLLM(),
|
||||
store=store,
|
||||
memory_manager=memory,
|
||||
compression_threshold=0.85,
|
||||
)
|
||||
)
|
||||
|
||||
await memory.ensure_session("session-a", project_id="proj1", metadata={"employee_id": "emp-a", "role_id": "developer"})
|
||||
await memory.ensure_session("session-b", project_id="proj1", metadata={"employee_id": "emp-b", "role_id": "developer"})
|
||||
|
||||
await store.save_agent_memory_snapshot(
|
||||
AgentMemorySnapshotRecord(
|
||||
project_id="proj1",
|
||||
session_id="session-a",
|
||||
employee_id="emp-a",
|
||||
role_id="developer",
|
||||
memory_kind="process",
|
||||
summary_message_id="summary-a",
|
||||
source_boundary_message_id="boundary-a",
|
||||
summary_text="process summary a",
|
||||
memory_text="## Effective Patterns\n- Temporary pattern A",
|
||||
)
|
||||
)
|
||||
await store.save_agent_memory_snapshot(
|
||||
AgentMemorySnapshotRecord(
|
||||
project_id="proj1",
|
||||
session_id="session-b",
|
||||
employee_id="emp-b",
|
||||
role_id="developer",
|
||||
memory_kind="process",
|
||||
summary_message_id="summary-b",
|
||||
source_boundary_message_id="boundary-b",
|
||||
summary_text="process summary b",
|
||||
memory_text="## Effective Patterns\n- Temporary pattern B",
|
||||
)
|
||||
)
|
||||
|
||||
task_a = Task(
|
||||
id="task-a",
|
||||
title="Implement feature A",
|
||||
project_id="proj1",
|
||||
session_id="session-a",
|
||||
assigned_to="developer",
|
||||
status=TaskStatus.DONE,
|
||||
result={"content": "done A"},
|
||||
metadata={
|
||||
"employee_assignment": {
|
||||
"employee_id": "emp-a",
|
||||
"role_id": "developer",
|
||||
"name": "Employee A",
|
||||
"domains": ["coding"],
|
||||
},
|
||||
"work_item_projection_id": "projection-a",
|
||||
"artifacts": ["file: a.py"],
|
||||
"decisions": ["Use API A"],
|
||||
},
|
||||
)
|
||||
task_b = Task(
|
||||
id="task-b",
|
||||
title="Implement feature B",
|
||||
project_id="proj1",
|
||||
session_id="session-b",
|
||||
assigned_to="developer",
|
||||
status=TaskStatus.DONE,
|
||||
result={"content": "done B"},
|
||||
metadata={
|
||||
"employee_assignment": {
|
||||
"employee_id": "emp-b",
|
||||
"role_id": "developer",
|
||||
"name": "Employee B",
|
||||
"domains": ["coding"],
|
||||
},
|
||||
"work_item_projection_id": "projection-b",
|
||||
"artifacts": ["file: b.py"],
|
||||
"decisions": ["Use API B"],
|
||||
},
|
||||
)
|
||||
delivery_task = Task(
|
||||
id="delivery-1",
|
||||
title="Project delivery",
|
||||
project_id="proj1",
|
||||
status=TaskStatus.DONE,
|
||||
metadata={"work_item_projection_id": "delivery"},
|
||||
)
|
||||
|
||||
await memory._record_project_reflections_and_finalize(
|
||||
delivery_task,
|
||||
[task_a, task_b],
|
||||
partial=False,
|
||||
)
|
||||
|
||||
final_a = await store.get_agent_memory_snapshot(
|
||||
project_id="proj1",
|
||||
employee_id="emp-a",
|
||||
memory_kind="final",
|
||||
memory_scope="project",
|
||||
)
|
||||
final_b = await store.get_agent_memory_snapshot(
|
||||
project_id="proj1",
|
||||
employee_id="emp-b",
|
||||
memory_kind="final",
|
||||
memory_scope="project",
|
||||
)
|
||||
process_a = await store.get_agent_memory_snapshot(
|
||||
project_id="proj1",
|
||||
session_id="session-a",
|
||||
employee_id="emp-a",
|
||||
memory_kind="process",
|
||||
memory_scope="session",
|
||||
)
|
||||
|
||||
self.assertIsNotNone(final_a)
|
||||
self.assertIsNotNone(final_b)
|
||||
self.assertIsNone(process_a)
|
||||
assert final_a is not None
|
||||
self.assertEqual(final_a.memory_scope, "project")
|
||||
self.assertEqual(final_a.session_id, "")
|
||||
self.assertIn("## Effective Patterns", final_a.memory_text)
|
||||
|
||||
context_a = await memory.build_employee_memory_context(
|
||||
project_id="proj1",
|
||||
session_id="session-a",
|
||||
employee_id="emp-a",
|
||||
role_id="developer",
|
||||
)
|
||||
context_b = await memory.build_employee_memory_context(
|
||||
project_id="proj1",
|
||||
session_id="session-b",
|
||||
employee_id="emp-b",
|
||||
role_id="developer",
|
||||
)
|
||||
self.assertIn("Employee Project Memory", context_a)
|
||||
self.assertIn("emp-a", context_a)
|
||||
self.assertNotIn("emp-b", context_a)
|
||||
self.assertIn("emp-b", context_b)
|
||||
|
||||
async def test_durable_compactor_persists_session_snapshot(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
root = Path(tmpdir)
|
||||
store = OPCStore(root / "tasks.db")
|
||||
await store.initialize()
|
||||
memory = MemoryManager(root, "proj1", store=store)
|
||||
compactor = HistoryCompactor(
|
||||
llm=_StubLLM(),
|
||||
store=store,
|
||||
memory_manager=memory,
|
||||
compression_threshold=0.85,
|
||||
)
|
||||
memory.set_history_compactor(compactor)
|
||||
|
||||
session_id = "session-loop"
|
||||
await memory.ensure_session(
|
||||
session_id,
|
||||
project_id="proj1",
|
||||
metadata={"employee_id": "emp-loop", "role_id": "developer"},
|
||||
)
|
||||
for idx in range(8):
|
||||
role = "user" if idx % 2 == 0 else "assistant"
|
||||
payload = (
|
||||
f"loop message {idx} "
|
||||
+ "alpha beta gamma delta epsilon zeta eta theta iota kappa lambda mu"
|
||||
)
|
||||
await memory.append_session_message(
|
||||
session_id=session_id,
|
||||
role=role,
|
||||
text=payload,
|
||||
project_id="proj1",
|
||||
metadata={"employee_id": "emp-loop", "role_id": "developer"},
|
||||
)
|
||||
|
||||
compacted = await compactor.maybe_compact_session(
|
||||
project_id="proj1",
|
||||
session_id=session_id,
|
||||
force=True,
|
||||
)
|
||||
|
||||
session_snapshot = await store.get_latest_session_memory_snapshot(session_id)
|
||||
rebuilt = await memory.build_session_history_tail_messages(session_id)
|
||||
self.assertTrue(compacted)
|
||||
self.assertIsNotNone(session_snapshot)
|
||||
assert session_snapshot is not None
|
||||
self.assertIn("Session history summary before restart", session_snapshot.summary_text)
|
||||
self.assertTrue(any("Session history summary before restart" in str(item.get("content", "")) for item in rebuilt))
|
||||
|
||||
async def test_legacy_session_final_memory_migrates_to_project_scope_on_read(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
root = Path(tmpdir)
|
||||
store = OPCStore(root / "tasks.db")
|
||||
await store.initialize()
|
||||
memory = MemoryManager(root, "proj1", store=store)
|
||||
|
||||
await store.save_agent_memory_snapshot(
|
||||
AgentMemorySnapshotRecord(
|
||||
project_id="proj1",
|
||||
session_id="legacy-session",
|
||||
employee_id="emp-legacy",
|
||||
role_id="developer",
|
||||
memory_scope="session",
|
||||
memory_kind="final",
|
||||
summary_message_id="legacy-summary",
|
||||
source_boundary_message_id="legacy-boundary",
|
||||
summary_text="legacy final summary",
|
||||
memory_text="## Effective Patterns\n- Legacy final memory",
|
||||
)
|
||||
)
|
||||
|
||||
context = await memory.build_employee_memory_context(
|
||||
project_id="proj1",
|
||||
session_id=None,
|
||||
employee_id="emp-legacy",
|
||||
role_id="developer",
|
||||
)
|
||||
migrated = await store.get_agent_memory_snapshot(
|
||||
project_id="proj1",
|
||||
employee_id="emp-legacy",
|
||||
memory_kind="final",
|
||||
memory_scope="project",
|
||||
)
|
||||
|
||||
self.assertIn("Employee Project Memory", context)
|
||||
self.assertIsNotNone(migrated)
|
||||
assert migrated is not None
|
||||
self.assertEqual(migrated.metadata.get("migrated_from_session_id"), "legacy-session")
|
||||
|
||||
async def test_focused_memory_context_prefers_relevant_sections(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
root = Path(tmpdir)
|
||||
store = OPCStore(root / "tasks.db")
|
||||
await store.initialize()
|
||||
memory = MemoryManager(root, "proj1", store=store)
|
||||
memory.save_memory(
|
||||
"# Global\n\n"
|
||||
"## Rust Build\n"
|
||||
"- Use cargo build for Rust projects.\n\n"
|
||||
"## Database Migrations\n"
|
||||
"- Always use migrations for schema changes.\n",
|
||||
project=False,
|
||||
)
|
||||
memory.save_memory(
|
||||
"# Project\n\n"
|
||||
"## Search Tooling\n"
|
||||
"- Prefer rg for code search.\n\n"
|
||||
"## Browser Checks\n"
|
||||
"- Use Playwright for browser verification.\n",
|
||||
project=True,
|
||||
)
|
||||
session_id = "session-focus"
|
||||
await memory.ensure_session(session_id, project_id="proj1")
|
||||
await memory.update_session_summary(session_id, "Remember the current runtime summary.")
|
||||
|
||||
focused = await memory.build_focused_memory_context(
|
||||
query="improve code search with rg",
|
||||
project_id="proj1",
|
||||
session_id=session_id,
|
||||
)
|
||||
self.assertIn("Focused Project Memory", focused)
|
||||
self.assertIn("Search Tooling", focused)
|
||||
self.assertNotIn("Browser Checks", focused)
|
||||
self.assertIn("## Session Memory", focused)
|
||||
|
||||
assembler = ContextAssembler(memory)
|
||||
core = await assembler.build_core_context(
|
||||
Task(
|
||||
title="Search fix",
|
||||
description="Use rg to inspect the repository",
|
||||
project_id="proj1",
|
||||
session_id=session_id,
|
||||
)
|
||||
)
|
||||
self.assertIn("Search Tooling", core)
|
||||
self.assertNotIn("Browser Checks", core)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user