update agnt
This commit is contained in:
Abdelrahman Elsayed Mahmoud Abdallah
@@ -0,0 +1,563 @@
|
||||
"""Session-memory-based compaction — uses an on-disk session memory summary
|
||||
instead of an API call to compact the conversation.
|
||||
|
||||
Mirrors the npm ``src/services/compact/sessionMemoryCompact.ts`` module.
|
||||
|
||||
When session memory is available and up-to-date, this avoids the cost
|
||||
of an API round-trip by reusing the background-maintained summary as
|
||||
the compaction text. Falls back to ``None`` so the caller can use
|
||||
the legacy API-based compact instead.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from .agent_context_usage import estimate_tokens
|
||||
from .agent_session import AgentMessage
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .compact import CompactionResult
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Configuration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@dataclass
|
||||
class SessionMemoryCompactConfig:
|
||||
"""Configuration for session-memory compaction thresholds."""
|
||||
|
||||
min_tokens: int = 10_000
|
||||
"""Minimum tokens to preserve after compaction."""
|
||||
|
||||
min_text_block_messages: int = 5
|
||||
"""Minimum number of messages with text content to keep."""
|
||||
|
||||
max_tokens: int = 40_000
|
||||
"""Hard cap — never preserve more than this many tokens."""
|
||||
|
||||
max_section_tokens: int = 2_000
|
||||
"""Maximum tokens per section in the session memory."""
|
||||
|
||||
max_total_tokens: int = 12_000
|
||||
"""Maximum total tokens for the session memory summary."""
|
||||
|
||||
|
||||
DEFAULT_CONFIG = SessionMemoryCompactConfig()
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Session memory file management
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
SESSION_MEMORY_TEMPLATE_SECTIONS = (
|
||||
'## User Profile',
|
||||
'## Project Context',
|
||||
'## Key Decisions & Rationale',
|
||||
'## Current Task Context',
|
||||
'## Important Patterns & Preferences',
|
||||
'## Learned Corrections',
|
||||
'## Tool Usage Patterns',
|
||||
'## Conversation Flow',
|
||||
'## Open Questions & Uncertainties',
|
||||
)
|
||||
"""Section headers used in the session memory template."""
|
||||
|
||||
|
||||
def get_session_memory_dir() -> Path:
|
||||
"""Return the directory where session memory files are stored."""
|
||||
home = Path.home()
|
||||
return home / '.claude' / 'session-memory'
|
||||
|
||||
|
||||
def get_session_memory_path() -> Path:
|
||||
"""Return the path to the session memory file."""
|
||||
return get_session_memory_dir() / 'session.md'
|
||||
|
||||
|
||||
def load_session_memory() -> str | None:
|
||||
"""Load session memory from disk, returning None if absent or empty."""
|
||||
path = get_session_memory_path()
|
||||
if not path.exists():
|
||||
return None
|
||||
try:
|
||||
content = path.read_text(encoding='utf-8').strip()
|
||||
if not content:
|
||||
return None
|
||||
return content
|
||||
except (OSError, UnicodeDecodeError):
|
||||
return None
|
||||
|
||||
|
||||
def save_session_memory(content: str) -> None:
|
||||
"""Save session memory to disk."""
|
||||
path = get_session_memory_path()
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(content, encoding='utf-8')
|
||||
|
||||
|
||||
def is_template_only(content: str) -> bool:
|
||||
"""Check if the session memory is just the empty template.
|
||||
|
||||
Returns True if all sections are present but contain no real content
|
||||
beyond the template headers and italic descriptions.
|
||||
"""
|
||||
lines = content.strip().splitlines()
|
||||
for line in lines:
|
||||
stripped = line.strip()
|
||||
if not stripped:
|
||||
continue
|
||||
# Skip section headers
|
||||
if stripped.startswith('##'):
|
||||
continue
|
||||
# Skip italic template descriptions
|
||||
if stripped.startswith('*') and stripped.endswith('*'):
|
||||
continue
|
||||
if stripped.startswith('_') and stripped.endswith('_'):
|
||||
continue
|
||||
# Found non-template content
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Message analysis helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _has_text_content(msg: AgentMessage) -> bool:
|
||||
"""Check if a message has meaningful text content."""
|
||||
content = msg.content.strip()
|
||||
if not content:
|
||||
return False
|
||||
if content == '[Old tool result content cleared]':
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _get_tool_result_ids(msg: AgentMessage) -> list[str]:
|
||||
"""Extract tool_call_ids from a tool-result message."""
|
||||
if msg.role == 'tool' and msg.tool_call_id:
|
||||
return [msg.tool_call_id]
|
||||
return []
|
||||
|
||||
|
||||
def _has_tool_use_with_ids(msg: AgentMessage, ids: set[str]) -> bool:
|
||||
"""Check if an assistant message contains tool_use blocks matching any of the given IDs."""
|
||||
if msg.role != 'assistant':
|
||||
return False
|
||||
tool_calls = msg.metadata.get('tool_calls') or msg.tool_calls
|
||||
if not tool_calls:
|
||||
return False
|
||||
for tc in tool_calls:
|
||||
tc_id = tc.get('id', '') if isinstance(tc, dict) else ''
|
||||
if tc_id in ids:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Index calculation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def adjust_index_to_preserve_api_invariants(
|
||||
messages: list[AgentMessage],
|
||||
keep_from: int,
|
||||
) -> int:
|
||||
"""Walk backwards to ensure kept messages maintain API-valid structure.
|
||||
|
||||
Specifically:
|
||||
1. All tool_result messages must have corresponding tool_use messages.
|
||||
2. Assistant messages sharing the same message_id (thinking blocks)
|
||||
must be kept together.
|
||||
"""
|
||||
if keep_from <= 0:
|
||||
return 0
|
||||
|
||||
# Step 1: Ensure tool_use/tool_result pairs
|
||||
kept_tool_result_ids: set[str] = set()
|
||||
for msg in messages[keep_from:]:
|
||||
for tid in _get_tool_result_ids(msg):
|
||||
kept_tool_result_ids.add(tid)
|
||||
|
||||
if kept_tool_result_ids:
|
||||
idx = keep_from - 1
|
||||
while idx >= 0:
|
||||
msg = messages[idx]
|
||||
if _has_tool_use_with_ids(msg, kept_tool_result_ids):
|
||||
keep_from = idx
|
||||
# Include any new tool_results this brings in
|
||||
for m in messages[idx:keep_from]:
|
||||
for tid in _get_tool_result_ids(m):
|
||||
kept_tool_result_ids.add(tid)
|
||||
idx -= 1
|
||||
|
||||
# Step 2: Ensure thinking block continuity (same message_id)
|
||||
kept_msg_ids: set[str] = set()
|
||||
for msg in messages[keep_from:]:
|
||||
if msg.role == 'assistant' and msg.message_id:
|
||||
kept_msg_ids.add(msg.message_id)
|
||||
|
||||
if kept_msg_ids:
|
||||
idx = keep_from - 1
|
||||
while idx >= 0:
|
||||
msg = messages[idx]
|
||||
if msg.role == 'assistant' and msg.message_id in kept_msg_ids:
|
||||
keep_from = idx
|
||||
idx -= 1
|
||||
|
||||
return keep_from
|
||||
|
||||
|
||||
def calculate_messages_to_keep_index(
|
||||
messages: list[AgentMessage],
|
||||
last_summarized_index: int,
|
||||
model: str = '',
|
||||
config: SessionMemoryCompactConfig | None = None,
|
||||
) -> int:
|
||||
"""Calculate the index from which to preserve messages.
|
||||
|
||||
Starts from ``last_summarized_index + 1`` and expands backwards
|
||||
to meet the configured minimums (token count, text block count).
|
||||
Stops if the hard max_tokens cap is reached.
|
||||
"""
|
||||
if config is None:
|
||||
config = DEFAULT_CONFIG
|
||||
|
||||
start_index = last_summarized_index + 1
|
||||
if start_index >= len(messages):
|
||||
return len(messages)
|
||||
|
||||
keep_from = start_index
|
||||
token_count = 0
|
||||
text_block_count = 0
|
||||
|
||||
# Count forward from keep_from to end
|
||||
for msg in messages[keep_from:]:
|
||||
token_count += estimate_tokens(msg.content, model)
|
||||
if _has_text_content(msg):
|
||||
text_block_count += 1
|
||||
|
||||
# Expand backwards if minimums not met
|
||||
idx = keep_from - 1
|
||||
while idx >= 0:
|
||||
# Check if we've reached a compact boundary — don't go past it
|
||||
if messages[idx].metadata.get('kind') == 'compact_boundary':
|
||||
break
|
||||
|
||||
msg_tokens = estimate_tokens(messages[idx].content, model)
|
||||
|
||||
# Hard cap: stop if adding this would exceed max_tokens
|
||||
if token_count + msg_tokens > config.max_tokens:
|
||||
break
|
||||
|
||||
# Expand backwards
|
||||
keep_from = idx
|
||||
token_count += msg_tokens
|
||||
if _has_text_content(messages[idx]):
|
||||
text_block_count += 1
|
||||
|
||||
# Check if minimums are met
|
||||
if (token_count >= config.min_tokens
|
||||
and text_block_count >= config.min_text_block_messages):
|
||||
break
|
||||
|
||||
idx -= 1
|
||||
|
||||
# Ensure API invariants
|
||||
keep_from = adjust_index_to_preserve_api_invariants(messages, keep_from)
|
||||
|
||||
return keep_from
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Session memory truncation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def truncate_session_memory(
|
||||
content: str,
|
||||
config: SessionMemoryCompactConfig | None = None,
|
||||
) -> tuple[str, bool]:
|
||||
"""Truncate session memory sections to fit within token limits.
|
||||
|
||||
Returns ``(truncated_content, was_truncated)``.
|
||||
"""
|
||||
if config is None:
|
||||
config = DEFAULT_CONFIG
|
||||
|
||||
total_tokens = estimate_tokens(content, '')
|
||||
if total_tokens <= config.max_total_tokens:
|
||||
return content, False
|
||||
|
||||
# Split by section headers and truncate each
|
||||
lines = content.splitlines(keepends=True)
|
||||
sections: list[list[str]] = []
|
||||
current_section: list[str] = []
|
||||
|
||||
for line in lines:
|
||||
if line.strip().startswith('## ') and current_section:
|
||||
sections.append(current_section)
|
||||
current_section = [line]
|
||||
else:
|
||||
current_section.append(line)
|
||||
if current_section:
|
||||
sections.append(current_section)
|
||||
|
||||
truncated_sections: list[str] = []
|
||||
was_truncated = False
|
||||
|
||||
for section in sections:
|
||||
section_text = ''.join(section)
|
||||
section_tokens = estimate_tokens(section_text, '')
|
||||
if section_tokens > config.max_section_tokens:
|
||||
# Truncate to fit within section limit
|
||||
truncated_lines: list[str] = []
|
||||
running_tokens = 0
|
||||
for line in section:
|
||||
line_tokens = estimate_tokens(line, '')
|
||||
if running_tokens + line_tokens > config.max_section_tokens:
|
||||
truncated_lines.append('...(truncated)\n')
|
||||
was_truncated = True
|
||||
break
|
||||
truncated_lines.append(line)
|
||||
running_tokens += line_tokens
|
||||
truncated_sections.append(''.join(truncated_lines))
|
||||
else:
|
||||
truncated_sections.append(section_text)
|
||||
|
||||
result = ''.join(truncated_sections)
|
||||
|
||||
# Check total again
|
||||
if estimate_tokens(result, '') > config.max_total_tokens:
|
||||
was_truncated = True
|
||||
|
||||
return result, was_truncated
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Core session memory compaction
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def try_session_memory_compaction(
|
||||
messages: list[AgentMessage],
|
||||
model: str = '',
|
||||
last_summarized_message_id: str | None = None,
|
||||
auto_compact_threshold: int | None = None,
|
||||
config: SessionMemoryCompactConfig | None = None,
|
||||
) -> 'CompactionResult | None':
|
||||
"""Attempt session-memory-based compaction.
|
||||
|
||||
Returns a :class:`CompactionResult` if session memory is available
|
||||
and the compaction succeeds, or ``None`` to signal the caller should
|
||||
fall back to API-based compaction.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
messages:
|
||||
The current session messages.
|
||||
model:
|
||||
Model name for token estimation.
|
||||
last_summarized_message_id:
|
||||
The message_id of the last message that was included in the
|
||||
session memory summary. Messages after this are preserved.
|
||||
auto_compact_threshold:
|
||||
If provided, return None if post-compact tokens exceed this.
|
||||
config:
|
||||
Compaction configuration thresholds.
|
||||
"""
|
||||
from .compact import CompactionResult
|
||||
|
||||
if config is None:
|
||||
config = DEFAULT_CONFIG
|
||||
|
||||
# Check environment gates
|
||||
if os.environ.get('DISABLE_CLAUDE_CODE_SM_COMPACT'):
|
||||
return None
|
||||
|
||||
# Load session memory
|
||||
session_memory = load_session_memory()
|
||||
if session_memory is None:
|
||||
return None
|
||||
|
||||
if is_template_only(session_memory):
|
||||
return None
|
||||
|
||||
# Find the boundary message
|
||||
last_summarized_index: int | None = None
|
||||
|
||||
if last_summarized_message_id:
|
||||
for i, msg in enumerate(messages):
|
||||
if msg.message_id == last_summarized_message_id:
|
||||
last_summarized_index = i
|
||||
break
|
||||
|
||||
if last_summarized_index is None:
|
||||
# No boundary found — can't determine what's already summarized
|
||||
# Fall back to legacy compact
|
||||
return None
|
||||
|
||||
# Calculate which messages to preserve
|
||||
keep_from = calculate_messages_to_keep_index(
|
||||
messages, last_summarized_index, model=model, config=config,
|
||||
)
|
||||
|
||||
messages_to_keep = list(messages[keep_from:])
|
||||
|
||||
# Filter out old compact boundaries from kept messages
|
||||
messages_to_keep = [
|
||||
m for m in messages_to_keep
|
||||
if m.metadata.get('kind') != 'compact_boundary'
|
||||
]
|
||||
|
||||
# Truncate session memory if needed
|
||||
truncated_memory, was_truncated = truncate_session_memory(
|
||||
session_memory, config=config,
|
||||
)
|
||||
|
||||
# Build the compaction result
|
||||
pre_tokens = sum(estimate_tokens(m.content, model) for m in messages)
|
||||
|
||||
boundary = AgentMessage(
|
||||
role='user',
|
||||
content=(
|
||||
'<system-reminder>\n'
|
||||
f'Earlier conversation was compacted using session memory. '
|
||||
f'{len(messages) - len(messages_to_keep)} messages summarized.\n'
|
||||
'</system-reminder>'
|
||||
),
|
||||
message_id='compact_boundary',
|
||||
metadata={
|
||||
'kind': 'compact_boundary',
|
||||
'source': 'session_memory',
|
||||
'pre_compact_token_count': pre_tokens,
|
||||
},
|
||||
)
|
||||
|
||||
summary_content = (
|
||||
'Here is a summary of our conversation so far:\n\n'
|
||||
f'{truncated_memory}'
|
||||
)
|
||||
if was_truncated:
|
||||
memory_path = get_session_memory_path()
|
||||
summary_content += (
|
||||
f'\n\n(Session memory was truncated. '
|
||||
f'Full version at: {memory_path})'
|
||||
)
|
||||
|
||||
summary_msg = AgentMessage(
|
||||
role='user',
|
||||
content=summary_content,
|
||||
message_id='compact_summary',
|
||||
metadata={
|
||||
'kind': 'compact_summary',
|
||||
'is_compact_summary': True,
|
||||
'source': 'session_memory',
|
||||
},
|
||||
)
|
||||
|
||||
post_messages = [boundary, summary_msg] + messages_to_keep
|
||||
post_tokens = sum(estimate_tokens(m.content, model) for m in post_messages)
|
||||
|
||||
# Check threshold
|
||||
if auto_compact_threshold is not None and post_tokens > auto_compact_threshold:
|
||||
return None
|
||||
|
||||
return CompactionResult(
|
||||
boundary_message=boundary,
|
||||
summary_messages=[summary_msg],
|
||||
messages_to_keep=messages_to_keep,
|
||||
pre_compact_token_count=pre_tokens,
|
||||
post_compact_token_count=post_tokens,
|
||||
true_post_compact_token_count=post_tokens,
|
||||
summary_text=truncated_memory,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Session memory extraction (lightweight background summary updater)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
SESSION_MEMORY_EXTRACTION_PROMPT = """Analyze the conversation so far and update the session memory.
|
||||
Extract key information into these sections:
|
||||
|
||||
## User Profile
|
||||
Who the user is, their role, expertise level, and preferences.
|
||||
|
||||
## Project Context
|
||||
What project/codebase is being worked on, its structure, and tech stack.
|
||||
|
||||
## Key Decisions & Rationale
|
||||
Important decisions made during this session and why.
|
||||
|
||||
## Current Task Context
|
||||
What the user is currently working on and the state of that work.
|
||||
|
||||
## Important Patterns & Preferences
|
||||
Coding style, conventions, or preferences observed.
|
||||
|
||||
## Learned Corrections
|
||||
Mistakes made and corrections applied — things to avoid repeating.
|
||||
|
||||
## Tool Usage Patterns
|
||||
Which tools work well, preferred approaches for common tasks.
|
||||
|
||||
## Conversation Flow
|
||||
Major topic transitions and how the conversation has progressed.
|
||||
|
||||
## Open Questions & Uncertainties
|
||||
Unresolved questions or areas of ambiguity.
|
||||
|
||||
Write concisely. Focus on information that would be useful for continuing
|
||||
this conversation after a context reset. Omit sections with no content."""
|
||||
|
||||
|
||||
def extract_session_memory_from_messages(
|
||||
messages: list[AgentMessage],
|
||||
model: str = '',
|
||||
) -> str:
|
||||
"""Build a session memory summary from conversation messages.
|
||||
|
||||
This is a lightweight local extraction — it walks the messages and
|
||||
builds a structured summary without an API call. For richer
|
||||
summaries, the full LLM-based extraction should be used.
|
||||
"""
|
||||
user_messages: list[str] = []
|
||||
tool_names_used: set[str] = set()
|
||||
file_paths: set[str] = set()
|
||||
corrections: list[str] = []
|
||||
|
||||
for msg in messages:
|
||||
if msg.role == 'user' and _has_text_content(msg):
|
||||
user_messages.append(msg.content[:200])
|
||||
elif msg.role == 'tool' and msg.name:
|
||||
tool_names_used.add(msg.name)
|
||||
path = msg.metadata.get('path')
|
||||
if isinstance(path, str):
|
||||
file_paths.add(path)
|
||||
|
||||
sections: list[str] = []
|
||||
|
||||
if user_messages:
|
||||
sections.append('## Current Task Context')
|
||||
# Use recent user messages as task context
|
||||
recent = user_messages[-5:]
|
||||
for um in recent:
|
||||
sections.append(f'- {um[:100]}')
|
||||
|
||||
if tool_names_used:
|
||||
sections.append('\n## Tool Usage Patterns')
|
||||
sections.append(f'Tools used: {", ".join(sorted(tool_names_used))}')
|
||||
|
||||
if file_paths:
|
||||
sections.append('\n## Project Context')
|
||||
sections.append(f'Files accessed: {", ".join(sorted(list(file_paths)[:20]))}')
|
||||
|
||||
sections.append('\n## Conversation Flow')
|
||||
sections.append(f'Total messages: {len(messages)}')
|
||||
sections.append(
|
||||
f'User messages: {sum(1 for m in messages if m.role == "user")}'
|
||||
)
|
||||
|
||||
return '\n'.join(sections)
|
||||
Reference in New Issue
Block a user