LZH-YS1998
608 lines
24 KiB
Python
608 lines
24 KiB
Python
"""Persistent history compaction for session and employee memory."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from datetime import datetime
|
|
from typing import Any
|
|
|
|
from loguru import logger
|
|
|
|
from opc.core.models import (
|
|
AgentCompactionRecord,
|
|
AgentMemorySnapshotRecord,
|
|
SessionCompactionRecord,
|
|
SessionMemorySnapshotRecord,
|
|
SessionMessageRecord,
|
|
)
|
|
|
|
|
|
class HistoryCompactor:
|
|
"""Compacts persisted history into summary + memory snapshots."""
|
|
|
|
_COMPACTION_MESSAGE_CHAR_BUDGET = 4_000
|
|
_COMPACTION_TRUNCATION_MARKER = "[history compaction input truncated]"
|
|
_RETRY_TRUNCATION_MARKER = "[history compaction retry truncated]"
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
llm: Any | None,
|
|
store: Any | None,
|
|
memory_manager: Any,
|
|
task_type: str = "quick_tasks",
|
|
compression_threshold: float = 0.85,
|
|
) -> None:
|
|
self.llm = llm
|
|
self.store = store
|
|
self.memory_manager = memory_manager
|
|
self.task_type = task_type
|
|
self.compression_threshold = compression_threshold
|
|
|
|
async def maybe_compact_after_message(self, message: SessionMessageRecord) -> None:
|
|
_ = message
|
|
return
|
|
|
|
async def maybe_compact_session(
|
|
self,
|
|
*,
|
|
project_id: str,
|
|
session_id: str,
|
|
force: bool = False,
|
|
) -> bool:
|
|
if not self.llm or not self.store:
|
|
return False
|
|
visible_items = await self.memory_manager._get_visible_session_transcript(session_id)
|
|
if not visible_items:
|
|
return False
|
|
raw_items = [item for item in visible_items if not getattr(item["message"], "summary_flag", False)]
|
|
if not raw_items:
|
|
return False
|
|
visible_messages = self._items_to_messages(visible_items)
|
|
if not self._should_compact(visible_messages, force=force):
|
|
return False
|
|
compact_items, boundary_message_id = self._select_compaction_items(raw_items, force=force)
|
|
if not compact_items or not boundary_message_id:
|
|
return False
|
|
messages = self._items_to_messages(
|
|
compact_items,
|
|
per_message_budget=self._COMPACTION_MESSAGE_CHAR_BUDGET,
|
|
)
|
|
|
|
existing = await self.store.get_latest_session_memory_snapshot(session_id)
|
|
result = await self._summarize_session(
|
|
project_id=project_id,
|
|
session_id=session_id,
|
|
messages=messages,
|
|
existing_memory=(existing.memory_text if existing else ""),
|
|
existing_summary=(existing.summary_text if existing else ""),
|
|
)
|
|
summary_message = await self.memory_manager.append_session_message(
|
|
session_id=session_id,
|
|
role="assistant",
|
|
text=result["history_summary"],
|
|
project_id=project_id,
|
|
summary_flag=True,
|
|
parent_message_id=boundary_message_id,
|
|
metadata={
|
|
"kind": "session_history_summary",
|
|
"summary_scope": "session",
|
|
"skip_compaction": True,
|
|
},
|
|
)
|
|
if not summary_message:
|
|
return False
|
|
|
|
await self.store.save_session_compaction(
|
|
SessionCompactionRecord(
|
|
session_id=session_id,
|
|
compaction_message_id=summary_message.message_id,
|
|
source_boundary_message_id=boundary_message_id,
|
|
metadata={
|
|
"project_id": project_id,
|
|
"raw_message_count": len(compact_items),
|
|
"summary_scope": "session",
|
|
},
|
|
)
|
|
)
|
|
await self.store.save_session_memory_snapshot(
|
|
SessionMemorySnapshotRecord(
|
|
project_id=project_id,
|
|
session_id=session_id,
|
|
summary_message_id=summary_message.message_id,
|
|
source_boundary_message_id=boundary_message_id,
|
|
summary_text=result["history_summary"],
|
|
memory_text=result["memory_summary"],
|
|
metadata={
|
|
"summary_scope": "session",
|
|
"raw_message_count": len(compact_items),
|
|
},
|
|
)
|
|
)
|
|
await self.memory_manager.update_session_summary(session_id, result["history_summary"])
|
|
return True
|
|
|
|
async def maybe_compact_agent(
|
|
self,
|
|
*,
|
|
project_id: str,
|
|
session_id: str,
|
|
employee_id: str,
|
|
role_id: str = "",
|
|
force: bool = False,
|
|
) -> bool:
|
|
if not self.llm or not self.store or not employee_id:
|
|
return False
|
|
visible_items = await self.memory_manager._get_visible_agent_transcript(
|
|
project_id=project_id,
|
|
session_id=session_id,
|
|
employee_id=employee_id,
|
|
)
|
|
if not visible_items:
|
|
return False
|
|
raw_items = [item for item in visible_items if not getattr(item["message"], "summary_flag", False)]
|
|
if not raw_items:
|
|
return False
|
|
visible_messages = self._items_to_messages(visible_items)
|
|
if not self._should_compact(visible_messages, force=force):
|
|
return False
|
|
compact_items, boundary_message_id = self._select_compaction_items(raw_items, force=force)
|
|
if not compact_items or not boundary_message_id:
|
|
return False
|
|
messages = self._items_to_messages(
|
|
compact_items,
|
|
per_message_budget=self._COMPACTION_MESSAGE_CHAR_BUDGET,
|
|
)
|
|
|
|
existing = await self.store.get_agent_memory_snapshot(
|
|
project_id=project_id,
|
|
session_id=session_id,
|
|
employee_id=employee_id,
|
|
memory_kind="process",
|
|
memory_scope="session",
|
|
)
|
|
result = await self._summarize_agent_process(
|
|
project_id=project_id,
|
|
session_id=session_id,
|
|
employee_id=employee_id,
|
|
role_id=role_id,
|
|
messages=messages,
|
|
existing_memory=(existing.memory_text if existing else ""),
|
|
existing_summary=(existing.summary_text if existing else ""),
|
|
)
|
|
summary_message = await self.memory_manager.append_session_message(
|
|
session_id=session_id,
|
|
role="assistant",
|
|
text=result["history_summary"],
|
|
project_id=project_id,
|
|
summary_flag=True,
|
|
parent_message_id=boundary_message_id,
|
|
metadata={
|
|
"kind": "agent_history_summary",
|
|
"summary_scope": "agent",
|
|
"employee_id": employee_id,
|
|
"role_id": role_id,
|
|
"skip_compaction": True,
|
|
},
|
|
)
|
|
if not summary_message:
|
|
return False
|
|
|
|
await self.store.save_agent_compaction(
|
|
AgentCompactionRecord(
|
|
project_id=project_id,
|
|
session_id=session_id,
|
|
employee_id=employee_id,
|
|
role_id=role_id,
|
|
compaction_message_id=summary_message.message_id,
|
|
source_boundary_message_id=boundary_message_id,
|
|
metadata={
|
|
"summary_scope": "agent",
|
|
"raw_message_count": len(compact_items),
|
|
},
|
|
)
|
|
)
|
|
await self.store.save_agent_memory_snapshot(
|
|
AgentMemorySnapshotRecord(
|
|
project_id=project_id,
|
|
session_id=session_id,
|
|
employee_id=employee_id,
|
|
role_id=role_id,
|
|
memory_scope="session",
|
|
memory_kind="process",
|
|
summary_message_id=summary_message.message_id,
|
|
source_boundary_message_id=boundary_message_id,
|
|
summary_text=result["history_summary"],
|
|
memory_text=result["memory_summary"],
|
|
metadata={
|
|
"summary_scope": "agent",
|
|
"raw_message_count": len(compact_items),
|
|
},
|
|
)
|
|
)
|
|
return True
|
|
|
|
async def finalize_agent_memory(
|
|
self,
|
|
*,
|
|
project_id: str,
|
|
session_id: str,
|
|
employee_id: str,
|
|
role_id: str,
|
|
process_memory: str,
|
|
reflection_payload: dict[str, Any],
|
|
) -> dict[str, Any]:
|
|
if not self.llm:
|
|
return self._fallback_final_agent_memory(process_memory, reflection_payload)
|
|
prompt_payload = {
|
|
"project_id": project_id,
|
|
"session_id": session_id,
|
|
"employee_id": employee_id,
|
|
"role_id": role_id,
|
|
"process_memory": process_memory,
|
|
"reflection": reflection_payload,
|
|
}
|
|
raw = await self.llm.simple_chat(
|
|
prompt=json.dumps(prompt_payload, ensure_ascii=False),
|
|
system=(
|
|
"You are finalizing employee memory for a multi-agent runtime.\n"
|
|
"Return strict JSON with keys `summary_text`, `memory_text`, and `metadata`.\n"
|
|
"`memory_text` must be concise markdown with sections:\n"
|
|
"## Effective Patterns\n## Watchouts\n## Preferred Tools\n## Reviewer Preferences\n## Reusable Checklist\n"
|
|
"`metadata` must contain arrays with keys `effective_patterns`, `watchouts`, "
|
|
"`preferred_tools`, `reviewer_preferences`, `reusable_checklist`.\n"
|
|
"Merge the process memory with the reflection, remove duplication, and keep only durable guidance."
|
|
),
|
|
task_type=self.task_type,
|
|
)
|
|
parsed = self._parse_json_response(raw)
|
|
if not parsed:
|
|
return self._fallback_final_agent_memory(process_memory, reflection_payload)
|
|
metadata = parsed.get("metadata", {})
|
|
return {
|
|
"summary_text": str(parsed.get("summary_text", "")).strip() or str(parsed.get("memory_text", "")).strip(),
|
|
"memory_text": str(parsed.get("memory_text", "")).strip() or process_memory.strip(),
|
|
"metadata": metadata if isinstance(metadata, dict) else {},
|
|
}
|
|
|
|
def _get_token_threshold(self, *, reserve_tokens: int = 0) -> int | None:
|
|
if not self.llm:
|
|
return None
|
|
context_limit = self.llm.get_context_window(task_type=self.task_type)
|
|
if context_limit is None:
|
|
return None
|
|
threshold = int(context_limit * self.compression_threshold)
|
|
if reserve_tokens:
|
|
threshold = max(0, threshold - reserve_tokens)
|
|
return threshold
|
|
|
|
def should_compact_prompt(
|
|
self,
|
|
messages: list[dict[str, Any]],
|
|
*,
|
|
tools: list[dict[str, Any]] | None = None,
|
|
force: bool = False,
|
|
reserve_tokens: int = 0,
|
|
) -> bool:
|
|
_ = messages
|
|
_ = tools
|
|
_ = force
|
|
_ = reserve_tokens
|
|
return False
|
|
|
|
def _is_context_overflow_error(self, error: Exception) -> bool:
|
|
detector = getattr(self.llm, "is_context_overflow_error", None)
|
|
if callable(detector):
|
|
try:
|
|
return bool(detector(error))
|
|
except Exception:
|
|
return False
|
|
return False
|
|
|
|
@staticmethod
|
|
def _truncate_messages_for_retry(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
|
if len(messages) <= 1:
|
|
return messages
|
|
drop_count = max(1, len(messages) // 5)
|
|
if drop_count >= len(messages):
|
|
drop_count = len(messages) - 1
|
|
return messages[drop_count:]
|
|
|
|
@classmethod
|
|
def _truncate_message_content(cls, content: str, *, budget: int, marker: str) -> str:
|
|
if len(content) <= budget:
|
|
return content
|
|
clipped = max(120, budget - len(marker) - 1)
|
|
return content[:clipped].rstrip() + "\n" + marker
|
|
|
|
@classmethod
|
|
def _compact_messages_for_retry(
|
|
cls,
|
|
messages: list[dict[str, Any]],
|
|
*,
|
|
budget: int,
|
|
) -> list[dict[str, Any]]:
|
|
compacted: list[dict[str, Any]] = []
|
|
for message in messages:
|
|
compacted.append({
|
|
**message,
|
|
"content": cls._truncate_message_content(
|
|
str(message.get("content", "") or ""),
|
|
budget=budget,
|
|
marker=cls._RETRY_TRUNCATION_MARKER,
|
|
),
|
|
})
|
|
return compacted
|
|
|
|
async def _simple_chat_with_retry(
|
|
self,
|
|
*,
|
|
payload: dict[str, Any],
|
|
system: str,
|
|
) -> str:
|
|
retries = 0
|
|
working_payload = dict(payload)
|
|
while True:
|
|
try:
|
|
return await self.llm.simple_chat(
|
|
prompt=json.dumps(working_payload, ensure_ascii=False),
|
|
system=system,
|
|
task_type=self.task_type,
|
|
)
|
|
except Exception as exc:
|
|
messages = list(working_payload.get("messages", []) or [])
|
|
if retries >= 3 or not self._is_context_overflow_error(exc) or not messages:
|
|
raise
|
|
retries += 1
|
|
retry_budget = max(1_200, 4_000 // (2 ** (retries - 1)))
|
|
retry_messages = self._compact_messages_for_retry(messages, budget=retry_budget)
|
|
if len(retry_messages) > 1:
|
|
retry_messages = self._truncate_messages_for_retry(retry_messages)
|
|
working_payload["messages"] = retry_messages
|
|
logger.debug(
|
|
"History compactor retrying after context overflow with "
|
|
f"{len(working_payload['messages'])} messages preserved and budget={retry_budget} chars."
|
|
)
|
|
|
|
def _should_compact(self, messages: list[dict[str, Any]], *, force: bool = False) -> bool:
|
|
if force:
|
|
return True
|
|
if not messages or not self.llm:
|
|
return False
|
|
counted_tokens = self.llm.count_input_tokens(messages, task_type=self.task_type)
|
|
threshold_tokens = self._get_token_threshold()
|
|
if counted_tokens is None or threshold_tokens is None:
|
|
return False
|
|
return counted_tokens >= threshold_tokens
|
|
|
|
def _select_compaction_items(
|
|
self,
|
|
raw_items: list[dict[str, Any]],
|
|
*,
|
|
force: bool = False,
|
|
) -> tuple[list[dict[str, Any]], str]:
|
|
if not raw_items:
|
|
return [], ""
|
|
if force or not self.llm:
|
|
return raw_items, raw_items[-1]["message"].message_id
|
|
|
|
threshold_tokens = self._get_token_threshold()
|
|
if threshold_tokens is None:
|
|
return [], ""
|
|
|
|
for keep_start in range(len(raw_items)):
|
|
tail_messages = self._items_to_messages(raw_items[keep_start:])
|
|
tail_tokens = self.llm.count_input_tokens(tail_messages, task_type=self.task_type)
|
|
if tail_tokens is None:
|
|
continue
|
|
if tail_tokens < threshold_tokens:
|
|
compact_items = raw_items if keep_start == 0 else raw_items[:keep_start]
|
|
return compact_items, compact_items[-1]["message"].message_id
|
|
|
|
return raw_items, raw_items[-1]["message"].message_id
|
|
|
|
def _items_to_messages(
|
|
self,
|
|
items: list[dict[str, Any]],
|
|
*,
|
|
per_message_budget: int | None = None,
|
|
) -> list[dict[str, Any]]:
|
|
messages: list[dict[str, Any]] = []
|
|
for item in items:
|
|
message = item["message"]
|
|
content = self.memory_manager._render_session_parts(item["parts"]).strip()
|
|
if not content:
|
|
continue
|
|
if per_message_budget:
|
|
content = self._truncate_message_content(
|
|
content,
|
|
budget=per_message_budget,
|
|
marker=self._COMPACTION_TRUNCATION_MARKER,
|
|
)
|
|
role = "user" if message.role == "user" else "assistant"
|
|
messages.append({"role": role, "content": content})
|
|
return messages
|
|
|
|
async def _summarize_session(
|
|
self,
|
|
*,
|
|
project_id: str,
|
|
session_id: str,
|
|
messages: list[dict[str, Any]],
|
|
existing_memory: str,
|
|
existing_summary: str,
|
|
) -> dict[str, str]:
|
|
if not self.llm:
|
|
return self._fallback_session_summary(messages, existing_memory)
|
|
payload = {
|
|
"project_id": project_id,
|
|
"session_id": session_id,
|
|
"existing_memory": existing_memory,
|
|
"existing_summary": existing_summary,
|
|
"messages": messages,
|
|
}
|
|
raw = await self._simple_chat_with_retry(
|
|
payload=payload,
|
|
system=(
|
|
"You are compacting persisted session history.\n"
|
|
"Return strict JSON with keys `history_summary` and `memory_summary`.\n"
|
|
"`history_summary` should help another agent continue the session after restart.\n"
|
|
"`memory_summary` should be concise markdown with sections:\n"
|
|
"## Primary Goal\n## Active Rules\n## Key Progress\n## Current State\n## Open Risks\n"
|
|
"Merge with existing memory and remove duplication."
|
|
),
|
|
)
|
|
parsed = self._parse_json_response(raw)
|
|
if not parsed:
|
|
return self._fallback_session_summary(messages, existing_memory)
|
|
return {
|
|
"history_summary": str(parsed.get("history_summary", "")).strip() or self._fallback_session_summary(messages, existing_memory)["history_summary"],
|
|
"memory_summary": str(parsed.get("memory_summary", "")).strip() or self._fallback_session_summary(messages, existing_memory)["memory_summary"],
|
|
}
|
|
|
|
async def _summarize_agent_process(
|
|
self,
|
|
*,
|
|
project_id: str,
|
|
session_id: str,
|
|
employee_id: str,
|
|
role_id: str,
|
|
messages: list[dict[str, Any]],
|
|
existing_memory: str,
|
|
existing_summary: str,
|
|
) -> dict[str, str]:
|
|
if not self.llm:
|
|
return self._fallback_agent_summary(messages, existing_memory)
|
|
payload = {
|
|
"project_id": project_id,
|
|
"session_id": session_id,
|
|
"employee_id": employee_id,
|
|
"role_id": role_id,
|
|
"existing_memory": existing_memory,
|
|
"existing_summary": existing_summary,
|
|
"messages": messages,
|
|
}
|
|
raw = await self._simple_chat_with_retry(
|
|
payload=payload,
|
|
system=(
|
|
"You are compacting employee-level process history.\n"
|
|
"Return strict JSON with keys `history_summary` and `memory_summary`.\n"
|
|
"`history_summary` should summarize what this employee already did in this session.\n"
|
|
"`memory_summary` should be concise markdown with sections:\n"
|
|
"## Effective Patterns\n## Watchouts\n## Current Progress\n## Current State\n"
|
|
"Keep it durable, specific, and deduplicated."
|
|
),
|
|
)
|
|
parsed = self._parse_json_response(raw)
|
|
if not parsed:
|
|
return self._fallback_agent_summary(messages, existing_memory)
|
|
return {
|
|
"history_summary": str(parsed.get("history_summary", "")).strip() or self._fallback_agent_summary(messages, existing_memory)["history_summary"],
|
|
"memory_summary": str(parsed.get("memory_summary", "")).strip() or self._fallback_agent_summary(messages, existing_memory)["memory_summary"],
|
|
}
|
|
|
|
def _fallback_session_summary(self, messages: list[dict[str, Any]], existing_memory: str) -> dict[str, str]:
|
|
snippets = [str(item.get("content", "")).strip() for item in messages if str(item.get("content", "")).strip()]
|
|
summary = "\n".join(f"- {snippet}" for snippet in snippets[-6:]) or "- No prior details captured."
|
|
memory_parts = [
|
|
"## Primary Goal",
|
|
f"- {snippets[0]}" if snippets else "- (unknown)",
|
|
"",
|
|
"## Active Rules",
|
|
"- Reuse durable constraints from earlier turns.",
|
|
"",
|
|
"## Key Progress",
|
|
*([f"- {snippet}" for snippet in snippets[-4:]] or ["- (none)"]),
|
|
"",
|
|
"## Current State",
|
|
f"- Existing memory length: {len(existing_memory.strip())} characters",
|
|
"",
|
|
"## Open Risks",
|
|
"- Re-check older transcript if the summary omits key details.",
|
|
]
|
|
return {
|
|
"history_summary": summary.strip(),
|
|
"memory_summary": "\n".join(memory_parts).strip(),
|
|
}
|
|
|
|
def _fallback_agent_summary(self, messages: list[dict[str, Any]], existing_memory: str) -> dict[str, str]:
|
|
snippets = [str(item.get("content", "")).strip() for item in messages if str(item.get("content", "")).strip()]
|
|
memory_parts = [
|
|
"## Effective Patterns",
|
|
*([f"- {snippet}" for snippet in snippets[-3:]] or ["- (none yet)"]),
|
|
"",
|
|
"## Watchouts",
|
|
"- Avoid repeating failed or already-compacted paths without new evidence.",
|
|
"",
|
|
"## Current Progress",
|
|
*([f"- {snippet}" for snippet in snippets[-2:]] or ["- (none)"]),
|
|
"",
|
|
"## Current State",
|
|
f"- Existing process memory length: {len(existing_memory.strip())} characters",
|
|
]
|
|
return {
|
|
"history_summary": "\n".join(f"- {snippet}" for snippet in snippets[-5:]) or "- No agent history captured.",
|
|
"memory_summary": "\n".join(memory_parts).strip(),
|
|
}
|
|
|
|
def _fallback_final_agent_memory(
|
|
self,
|
|
process_memory: str,
|
|
reflection_payload: dict[str, Any],
|
|
) -> dict[str, Any]:
|
|
what_worked = [str(item).strip() for item in reflection_payload.get("what_worked", []) if str(item).strip()]
|
|
watchouts = [str(item).strip() for item in reflection_payload.get("mistakes_to_avoid", []) if str(item).strip()]
|
|
preferred_tools = [str(item).strip() for item in reflection_payload.get("tool_preferences", []) if str(item).strip()]
|
|
reviewer_preferences = [str(item).strip() for item in reflection_payload.get("reviewer_preferences", []) if str(item).strip()]
|
|
checklist = [str(item).strip() for item in reflection_payload.get("reusable_checklist", []) if str(item).strip()]
|
|
parts = [
|
|
"## Effective Patterns",
|
|
*([f"- {item}" for item in what_worked[:6]] or ["- (none)"]),
|
|
"",
|
|
"## Watchouts",
|
|
*([f"- {item}" for item in watchouts[:6]] or ["- (none)"]),
|
|
"",
|
|
"## Preferred Tools",
|
|
*([f"- {item}" for item in preferred_tools[:6]] or ["- (none)"]),
|
|
"",
|
|
"## Reviewer Preferences",
|
|
*([f"- {item}" for item in reviewer_preferences[:6]] or ["- (none)"]),
|
|
"",
|
|
"## Reusable Checklist",
|
|
*([f"- {item}" for item in checklist[:6]] or ["- (none)"]),
|
|
]
|
|
summary_text = str(reflection_payload.get("project_summary", "")).strip() or process_memory.strip()
|
|
return {
|
|
"summary_text": summary_text,
|
|
"memory_text": "\n".join(parts).strip(),
|
|
"metadata": {
|
|
"effective_patterns": what_worked[:6],
|
|
"watchouts": watchouts[:6],
|
|
"preferred_tools": preferred_tools[:6],
|
|
"reviewer_preferences": reviewer_preferences[:6],
|
|
"reusable_checklist": checklist[:6],
|
|
},
|
|
}
|
|
|
|
def _parse_json_response(self, raw: str) -> dict[str, Any]:
|
|
text = raw.strip()
|
|
if text.startswith("```"):
|
|
parts = text.split("\n", 1)
|
|
text = parts[1] if len(parts) == 2 else text[3:]
|
|
if text.endswith("```"):
|
|
text = text[:-3]
|
|
text = text.strip()
|
|
try:
|
|
data = json.loads(text)
|
|
return data if isinstance(data, dict) else {}
|
|
except Exception:
|
|
start = text.find("{")
|
|
end = text.rfind("}")
|
|
if start >= 0 and end > start:
|
|
try:
|
|
data = json.loads(text[start : end + 1])
|
|
return data if isinstance(data, dict) else {}
|
|
except Exception as e:
|
|
logger.warning(f"Failed to parse compaction JSON: {e}")
|
|
return {}
|