Coverage for src / harnessutils / manager.py: 96%
296 statements
« prev ^ index » next coverage.py v7.13.2, created at 2026-04-07 09:04 -0600
« prev ^ index » next coverage.py v7.13.2, created at 2026-04-07 09:04 -0600
1"""Main ConversationManager API for harness-utils."""
3import time
4from collections.abc import Callable
5from pathlib import Path
6from typing import Any
8from harnessutils.compaction.pruning import PruningDecision, prune_tool_outputs
9from harnessutils.compaction.summarization import is_overflow, summarize_conversation
10from harnessutils.compaction.truncation import truncate_output
11from harnessutils.config import HarnessConfig, StorageConfig
12from harnessutils.conversion.to_model import to_model_messages
13from harnessutils.hooks import MessageHooks
14from harnessutils.inspection import ContextInspector
15from harnessutils.memory import SemanticMemoryBackend
16from harnessutils.models.conversation import Conversation
17from harnessutils.models.message import Message
18from harnessutils.models.usage import Usage
19from harnessutils.snapshots import Snapshot, SnapshotDiff, SnapshotManager
20from harnessutils.storage.filesystem import FilesystemStorage
21from harnessutils.storage.memory import MemoryStorage
22from harnessutils.types import LLMClient, StorageBackend
23from harnessutils.utils.ids import generate_id
24from harnessutils.workspace import resolve_workspace
27class ConversationManager:
28 """Main interface for managing conversations with context window management.
30 Provides high-level API for:
31 - Creating and managing conversations
32 - Adding messages
33 - Automatic context compaction (truncation, pruning, summarization)
34 - Message storage and retrieval
35 """
37 def __init__(
38 self,
39 storage: StorageBackend | None = None,
40 config: HarnessConfig | None = None,
41 *,
42 workspace_root: Path | str | None = None,
43 harness_name: str = ".harness",
44 semantic_memory: SemanticMemoryBackend | None = None,
45 message_hooks: MessageHooks | None = None,
46 ):
47 """Initialize conversation manager.
49 Args:
50 storage: Storage backend (uses in-memory if None)
51 config: Configuration (uses defaults if None)
52 workspace_root: Optional workspace root. Creates harness dir with stable
53 project UUID. Defaults storage to FilesystemStorage under <harness>/sessions/.
54 harness_name: Name of the harness directory (default: ".harness")
55 semantic_memory: Optional semantic memory backend. Stored as attribute;
56 harness-utils does not call it internally.
57 message_hooks: Optional hooks for the add_message() lifecycle.
58 on_before_add_message fires before storage and can modify or reject.
59 on_after_add_message fires after storage for side effects.
60 """
61 self.config = config or HarnessConfig()
62 self.project_id: str | None = None
63 self.semantic_memory = semantic_memory
65 if workspace_root is not None:
66 harness_dir, self.project_id = resolve_workspace(Path(workspace_root), harness_name)
67 if storage is None:
68 sessions_path = harness_dir / "sessions"
69 storage_config = StorageConfig(base_path=sessions_path)
70 storage = FilesystemStorage(storage_config)
72 self.storage = storage or MemoryStorage()
73 self.message_hooks = message_hooks
74 self._message_cache: dict[str, list[Message]] = {}
75 # Pass storage (not self.storage) so SnapshotManager receives None when no
76 # explicit backend was given — it falls back to its internal in-memory dict.
77 self.snapshot_manager = SnapshotManager(storage)
79 def create_conversation(
80 self,
81 conversation_id: str | None = None,
82 project_id: str | None = None,
83 ) -> Conversation:
84 """Create a new conversation.
86 Args:
87 conversation_id: Optional conversation ID (generated if None)
88 project_id: Optional project ID for grouping
90 Returns:
91 New conversation object
92 """
93 if conversation_id is None:
94 conversation_id = generate_id("conv")
96 now = int(time.time() * 1000)
97 conversation = Conversation(
98 id=conversation_id,
99 project_id=project_id,
100 created=now,
101 updated=now,
102 )
104 self.storage.save_conversation(conversation_id, conversation.to_dict())
105 self._message_cache[conversation_id] = []
107 return conversation
109 def add_message(self, conversation_id: str, message: Message) -> None:
110 """Add a message to a conversation.
112 If message_hooks.on_before_add_message is set, it is called first.
113 Its return value replaces the message that will be stored. Raising
114 from it rejects the message entirely (nothing is stored).
116 If message_hooks.on_after_add_message is set, it is called after
117 successful storage. Its return value is ignored.
119 Args:
120 conversation_id: Conversation to add message to
121 message: Message to add
122 """
123 # Pre-hook — may modify message or raise to reject
124 if self.message_hooks and self.message_hooks.on_before_add_message:
125 message = self.message_hooks.on_before_add_message(conversation_id, message)
127 # Inject timestamp if not set — required for cleanup_stale_data to work
128 if "timestamp" not in message.metadata:
129 message.metadata["timestamp"] = int(time.time() * 1000)
131 self.storage.save_message(conversation_id, message.id, message.to_dict())
133 # Persist parts separately so they survive process restart
134 for idx, part in enumerate(message.parts):
135 self.storage.save_part(message.id, f"{idx:04d}", part.to_dict())
137 if conversation_id not in self._message_cache:
138 self._message_cache[conversation_id] = []
139 self._message_cache[conversation_id].append(message)
141 # Load conversation and update
142 conv_data = self.storage.load_conversation(conversation_id)
143 conv = Conversation.from_dict(conv_data)
144 conv.updated = int(time.time() * 1000)
146 # Track velocity if message has token count
147 if message.tokens:
148 tokens_added = message.tokens.total
149 conv.update_velocity(tokens_added)
151 self.storage.save_conversation(conversation_id, conv.to_dict())
153 # Post-hook — side effects only, message already persisted
154 if self.message_hooks and self.message_hooks.on_after_add_message:
155 self.message_hooks.on_after_add_message(conversation_id, message)
157 def get_messages(self, conversation_id: str) -> list[Message]:
158 """Get all messages for a conversation.
160 Args:
161 conversation_id: Conversation ID
163 Returns:
164 List of messages in chronological order
165 """
166 if conversation_id in self._message_cache:
167 return self._message_cache[conversation_id]
169 message_ids = self.storage.list_messages(conversation_id)
170 messages = []
171 for msg_id in message_ids:
172 msg = Message.from_dict(self.storage.load_message(conversation_id, msg_id))
173 # Reload persisted parts (saved by add_message)
174 part_ids = self.storage.list_parts(msg_id)
175 if part_ids:
176 from harnessutils.models.parts import part_from_dict
177 for part_id in part_ids:
178 try:
179 part_data = self.storage.load_part(msg_id, part_id)
180 msg.parts.append(part_from_dict(part_data))
181 except Exception:
182 pass
183 messages.append(msg)
185 self._message_cache[conversation_id] = messages
186 return messages
188 def query_messages(
189 self,
190 conversation_id: str,
191 limit: int | None = None,
192 offset: int = 0,
193 order: str = "asc",
194 after: int | None = None,
195 before: int | None = None,
196 filter: dict[str, Any] | None = None,
197 ) -> list[Message]:
198 """Query messages with filtering and pagination.
200 Args:
201 conversation_id: Conversation to query
202 limit: Maximum messages to return
203 offset: Skip first N messages
204 order: "asc" (oldest first) or "desc" (newest first)
205 after: Unix ms timestamp - messages after this time
206 before: Unix ms timestamp - messages before this time
207 filter: Filter criteria dict:
208 - has_errors (bool): Only messages with errors
209 - has_warnings (bool): Only messages with warnings
210 - min_importance (float): Minimum importance score
211 - tools (list[str]): Specific tool types
212 - roles (list[str]): Message roles (user/assistant)
213 - has_tool_outputs (bool): Only messages with tool outputs
214 - is_summary (bool): Only summary messages
216 Returns:
217 Filtered and paginated messages
219 Examples:
220 # Get last 10 messages
221 messages = manager.query_messages(conv_id, limit=10, order="desc")
223 # Get messages with errors
224 messages = manager.query_messages(
225 conv_id,
226 filter={"has_errors": True}
227 )
229 # Get recent high-importance messages
230 messages = manager.query_messages(
231 conv_id,
232 filter={"min_importance": 50.0},
233 after=timestamp_24h_ago,
234 limit=20
235 )
236 """
237 from harnessutils.query import MessageFilter, QueryOptions, query_messages
239 # Build filter
240 msg_filter = None
241 if filter is not None:
242 msg_filter = MessageFilter(
243 has_errors=filter.get("has_errors"),
244 has_warnings=filter.get("has_warnings"),
245 min_importance=filter.get("min_importance"),
246 tools=filter.get("tools"),
247 roles=filter.get("roles"),
248 has_tool_outputs=filter.get("has_tool_outputs"),
249 is_summary=filter.get("is_summary"),
250 )
252 # Build query options
253 options = QueryOptions(
254 limit=limit,
255 offset=offset,
256 order=order, # type: ignore
257 after=after,
258 before=before,
259 filter=msg_filter,
260 )
262 # Get all messages and query
263 messages = self.get_messages(conversation_id)
264 return query_messages(messages, options)
266 def get_context_summary(
267 self, conversation_id: str, recent_limit: int = 5
268 ) -> dict[str, Any]:
269 """Get lightweight context summary without loading full content.
271 Provides overview of conversation state including:
272 - Message count and token usage
273 - Summary messages
274 - Recent activity (last N messages)
275 - Key messages (high importance)
276 - Error messages
278 Args:
279 conversation_id: Conversation to summarize
280 recent_limit: Number of recent messages to include (default: 5)
282 Returns:
283 Dictionary with context summary:
284 - conversation_id (str)
285 - message_count (int)
286 - total_tokens (int)
287 - summaries (list): Summary messages
288 - recent_activity (list): Last N message summaries
289 - key_messages (list): High-importance messages
290 - errors (list): Error messages
292 Examples:
293 # Get quick overview
294 summary = manager.get_context_summary(conv_id)
295 print(f"Messages: {summary['message_count']}")
296 print(f"Tokens: {summary['total_tokens']}")
297 print(f"Errors: {len(summary['errors'])}")
298 """
299 from harnessutils.query import build_context_summary
301 messages = self.get_messages(conversation_id)
302 summary = build_context_summary(conversation_id, messages, recent_limit)
303 return summary.to_dict()
305 def inspect_context(self, conversation_id: str) -> ContextInspector:
306 """Create inspector for querying context state.
308 Provides observability into:
309 - What's currently in context
310 - Token counts and breakdowns
311 - Pruning predictions
312 - Audit trail of decisions
314 Args:
315 conversation_id: Conversation to inspect
317 Returns:
318 ContextInspector instance with full query capabilities
320 Example:
321 >>> inspector = manager.inspect_context(conv_id)
322 >>> summary = inspector.summary()
323 >>> print(f"Total tokens: {summary['total_tokens']}")
324 >>> impact = inspector.predict_impact(5000)
325 >>> if impact['would_trigger_pruning']:
326 ... print(f"Would prune ~{impact['estimated_pruned_count']} outputs")
327 """
328 messages = self.get_messages(conversation_id)
330 # Load conversation metadata
331 conv_data = self.storage.load_conversation(conversation_id)
332 conversation = Conversation.from_dict(conv_data)
334 # Load persisted pruning decisions for audit trail
335 decisions_data = conversation.metadata.get("latest_pruning_decisions", [])
336 decisions = [PruningDecision.from_dict(d) for d in decisions_data]
338 return ContextInspector(messages, self.config, conversation, decisions)
340 def prune_before_turn(
341 self,
342 conversation_id: str,
343 auto_mode: bool = False,
344 query: str | None = None,
345 summarizer: Callable[[str, str], str] | None = None,
346 ) -> dict[str, Any]:
347 """Proactively prune old tool outputs before processing a turn.
349 This is Tier 2 compaction - removes old tool outputs while
350 preserving conversation structure.
352 Args:
353 conversation_id: Conversation to prune
354 auto_mode: Whether this was auto-triggered
355 query: Current user query for relevance-based importance boosting
356 summarizer: Optional callable(tool_name, output) -> summary_text.
357 When provided, called before clearing each output so a
358 one-sentence summary is preserved in the model context.
360 Returns:
361 Detailed pruning result with token tracking and breakdown:
362 - pruned: Total outputs removed
363 - tokens_saved: Total tokens saved
364 - tokens_before: Token count before pruning
365 - tokens_after: Token count after pruning
366 - duplicates_pruned: Outputs removed due to duplication
367 - importance_pruned: Outputs removed due to low importance
368 - duplicate_tokens_saved: Tokens saved from deduplication
369 - importance_tokens_saved: Tokens saved from importance pruning
370 - reduction_percent: Percentage reduction in token usage
371 """
372 if not self.config.compaction.prune and auto_mode:
373 return {
374 "pruned": 0,
375 "tokens_saved": 0,
376 "tokens_before": 0,
377 "tokens_after": 0,
378 "duplicates_pruned": 0,
379 "importance_pruned": 0,
380 "duplicate_tokens_saved": 0,
381 "importance_tokens_saved": 0,
382 "reduction_percent": 0,
383 }
385 messages = self.get_messages(conversation_id)
386 result = prune_tool_outputs(
387 messages,
388 self.config.pruning,
389 query=query,
390 summarizer=summarizer,
391 )
393 for msg in messages:
394 self.storage.save_message(conversation_id, msg.id, msg.to_dict())
396 # Update cache with pruned messages — parts are preserved since
397 # prune_tool_outputs only clears tool output strings in-place
398 self._message_cache[conversation_id] = messages
400 # Persist decisions so audit trail and quality metrics work
401 conv_data = self.storage.load_conversation(conversation_id)
402 conv = Conversation.from_dict(conv_data)
403 conv.metadata["latest_pruning_decisions"] = [d.to_dict() for d in result.decisions]
404 self.storage.save_conversation(conversation_id, conv.to_dict())
406 return result.to_dict()
408 def predict_overflow(
409 self,
410 conversation_id: str,
411 current_usage: Usage,
412 ) -> bool:
413 """Predict if conversation will overflow in next N turns.
415 Args:
416 conversation_id: Conversation to check
417 current_usage: Current token usage
419 Returns:
420 True if overflow predicted within lookahead window
421 """
422 if not self.config.compaction.use_predictive:
423 return False
425 # Load conversation and get velocity
426 conv_data = self.storage.load_conversation(conversation_id)
427 conv = Conversation.from_dict(conv_data)
428 velocity = conv.get_velocity()
430 if velocity is None or not velocity.turn_deltas:
431 return False # No velocity data yet
433 # Project tokens ahead
434 lookahead = self.config.compaction.predictive_lookahead
435 predicted_growth = velocity.predict_tokens_ahead(lookahead)
437 # Calculate current total and projected total
438 current_total = current_usage.input + current_usage.cache.read
439 projected_total = current_total + predicted_growth
441 # Check against safety margin
442 usable_space = (
443 self.config.model_limits.default_context_limit
444 - self.config.model_limits.default_output_limit
445 )
446 safety_threshold = usable_space * self.config.compaction.predictive_safety_margin
448 return projected_total > safety_threshold
450 def needs_compaction(
451 self,
452 conversation_id: str,
453 usage: Usage,
454 ) -> bool:
455 """Check if conversation needs summarization (Tier 3).
457 Uses both reactive (overflow) and predictive checks.
459 Args:
460 conversation_id: Conversation to check
461 usage: Token usage from last turn
463 Returns:
464 True if summarization needed
465 """
466 # Reactive check: already overflowed
467 if is_overflow(
468 usage,
469 self.config.model_limits.default_context_limit,
470 self.config.model_limits.default_output_limit,
471 ):
472 return True
474 # Predictive check: will overflow soon
475 return self.predict_overflow(conversation_id, usage)
477 def compact(
478 self,
479 conversation_id: str,
480 llm_client: LLMClient,
481 parent_message_id: str,
482 model: str | None = None,
483 auto_mode: bool = False,
484 ) -> dict[str, Any]:
485 """Compact conversation using LLM summarization (Tier 3).
487 Args:
488 conversation_id: Conversation to compact
489 llm_client: LLM client for summarization
490 parent_message_id: Message that triggered compaction
491 model: Optional model to use for summarization
492 auto_mode: Whether this was auto-triggered
494 Returns:
495 Compaction result with summary message and metrics
496 """
497 if not self.config.compaction.auto and auto_mode:
498 return {"summarized": False}
500 messages = self.get_messages(conversation_id)
501 summary_id = generate_id("msg")
503 result = summarize_conversation(
504 messages=messages,
505 llm_client=llm_client,
506 parent_message_id=parent_message_id,
507 message_id=summary_id,
508 model=model,
509 auto_mode=auto_mode,
510 config=self.config.summarization,
511 )
513 self.add_message(conversation_id, result.summary_message)
515 return {
516 "summarized": True,
517 "summary_message_id": summary_id,
518 "tokens_used": result.tokens_used.total,
519 "cost": result.cost,
520 }
522 def to_model_format(self, conversation_id: str) -> list[dict[str, Any]]:
523 """Convert conversation messages to model format for LLM requests.
525 Args:
526 conversation_id: Conversation to convert
528 Returns:
529 List of messages in model format
530 """
531 messages = self.get_messages(conversation_id)
532 return to_model_messages(messages)
534 def calculate_context_usage(
535 self,
536 conversation_id: str,
537 model: str | None = None,
538 ) -> int:
539 """Calculate exact token count for conversation using tiktoken.
541 This counts ALL tokens in the conversation (user messages, assistant
542 responses, tool outputs, etc.) that will be sent to the model.
544 Uses cl100k_base tokenizer (GPT-4/Claude) which works for most modern LLMs.
546 Args:
547 conversation_id: Conversation to calculate usage for
548 model: Optional model name (currently unused, defaults to cl100k_base)
550 Returns:
551 Exact token count that will be used in context window
552 """
553 from harnessutils.tokens.exact import count_tokens_exact
555 messages = self.to_model_format(conversation_id)
556 return count_tokens_exact(messages, model)
558 def get_tool_output_tokens(self, conversation_id: str) -> dict[str, Any]:
559 """Get detailed breakdown of token usage for tool outputs.
561 Args:
562 conversation_id: Conversation to analyze
564 Returns:
565 Dictionary with token breakdown:
566 - total: Total tokens in tool outputs
567 - by_tool: Token count per tool type
568 - prunable: Tokens that could be pruned
569 - protected: Tokens in protected outputs
570 """
571 from harnessutils.compaction.pruning import calculate_context_tokens
573 messages = self.get_messages(conversation_id)
575 total = calculate_context_tokens(messages)
576 by_tool: dict[str, int] = {}
577 prunable = 0
578 protected = 0
579 turns_skipped = 0
581 for msg in reversed(messages):
582 if msg.role == "user":
583 turns_skipped += 1
585 for part in msg.parts:
586 from harnessutils.models.parts import ToolPart
587 from harnessutils.tokens.exact import count_tokens_fast
589 if not isinstance(part, ToolPart):
590 continue
592 if part.state.status != "completed":
593 continue
595 if not part.state.output:
596 continue
598 tokens = count_tokens_fast(part.state.output)
600 # Track by tool type
601 by_tool[part.tool] = by_tool.get(part.tool, 0) + tokens
603 # Determine if prunable
604 is_protected = (
605 turns_skipped < self.config.pruning.protect_turns
606 or part.tool in self.config.pruning.protected_tools
607 or (part.state.time and part.state.time.compacted)
608 )
610 if is_protected:
611 protected += tokens
612 else:
613 prunable += tokens
615 return {
616 "total": total,
617 "by_tool": by_tool,
618 "prunable": prunable,
619 "protected": protected,
620 "prunability_percent": round((prunable / total * 100) if total > 0 else 0, 1),
621 }
623 def get_context_quality(self, conversation_id: str) -> dict[str, Any]:
624 """Get current quality metrics for conversation.
626 Calculates all quality metrics, updates history, and returns assessment
627 with health status and actionable recommendations.
629 Args:
630 conversation_id: Conversation to analyze
632 Returns:
633 Dictionary with all metrics, health status, and recommendations.
634 Keys: information_density, redundancy_ratio, staleness_score,
635 error_preservation_rate, protected_ratio, health, recommendations
636 """
637 from harnessutils.quality import assess_quality
639 conv_data = self.storage.load_conversation(conversation_id)
640 conv = Conversation.from_dict(conv_data)
641 messages = self.get_messages(conversation_id)
643 # Get pruning decisions from metadata if available
644 decisions = conv.metadata.get("latest_pruning_decisions")
646 snapshot = assess_quality(
647 messages=messages,
648 config=self.config.pruning,
649 decisions=decisions,
650 )
652 # Update history
653 conv.update_quality_history(snapshot)
654 self.storage.save_conversation(conversation_id, conv.to_dict())
656 return snapshot.to_dict()
658 def track_quality_metric(
659 self,
660 conversation_id: str,
661 metric_name: str,
662 value: float,
663 timestamp: int | None = None,
664 ) -> None:
665 """Track a single quality metric value.
667 Note: Prefer get_context_quality() which calculates all metrics.
668 This is for custom/external metrics.
670 Args:
671 conversation_id: Conversation to track
672 metric_name: Name of metric (e.g., "information_density")
673 value: Metric value
674 timestamp: Unix ms timestamp (defaults to now)
675 """
676 from harnessutils.quality import QualitySnapshot
678 if timestamp is None:
679 timestamp = int(time.time() * 1000)
681 conv_data = self.storage.load_conversation(conversation_id)
682 conv = Conversation.from_dict(conv_data)
684 # Create minimal snapshot with just this metric
685 # (Other metrics set to 0.0, empty recommendations)
686 snapshot = QualitySnapshot(
687 timestamp=timestamp,
688 information_density=value if metric_name == "information_density" else 0.0,
689 redundancy_ratio=value if metric_name == "redundancy_ratio" else 0.0,
690 staleness_score=value if metric_name == "staleness_score" else 0.0,
691 error_preservation_rate=value
692 if metric_name == "error_preservation_rate"
693 else 0.0,
694 protected_ratio=value if metric_name == "protected_ratio" else 0.0,
695 health="unknown",
696 recommendations=[],
697 )
699 conv.update_quality_history(snapshot)
700 self.storage.save_conversation(conversation_id, conv.to_dict())
702 def get_quality_trend(
703 self,
704 conversation_id: str,
705 metric: str,
706 window: int = 20,
707 ) -> list[tuple[int, float]]:
708 """Get trend data for a specific metric.
710 Args:
711 conversation_id: Conversation to query
712 metric: Metric name (e.g., "information_density")
713 window: Number of most recent snapshots to return
715 Returns:
716 List of (timestamp, value) tuples, most recent first
717 """
718 conv_data = self.storage.load_conversation(conversation_id)
719 conv = Conversation.from_dict(conv_data)
721 history = conv.get_quality_history()
722 if history is None:
723 return []
725 return history.get_trend(metric, window)
727 def truncate_tool_output(
728 self,
729 output: str,
730 tool_name: str,
731 ) -> str:
732 """Truncate tool output if it exceeds limits (Tier 1).
734 Args:
735 output: Tool output to truncate
736 tool_name: Name of the tool
738 Returns:
739 Potentially truncated output
740 """
741 output_id = generate_id(f"output_{tool_name}")
743 result = truncate_output(
744 output=output,
745 config=self.config.truncation,
746 output_id=output_id,
747 )
749 if result.truncated and result.output_path:
750 self.storage.save_truncated_output(result.output_path, output)
752 return result.content
754 def create_snapshot(
755 self,
756 conversation_id: str,
757 snapshot_id: str | None = None,
758 metadata: dict[str, Any] | None = None,
759 ) -> Snapshot:
760 """Create snapshot of conversation state for reproducibility.
762 Captures full conversation state including:
763 - All messages with content
764 - Conversation metadata (velocity, etc.)
765 - Current configuration
767 Useful for:
768 - Debugging (save state before/after changes)
769 - A/B testing (compare different strategies)
770 - Reproducibility (restore exact state)
772 Args:
773 conversation_id: Conversation to snapshot
774 snapshot_id: Optional snapshot ID (auto-generated if None)
775 metadata: Optional metadata (e.g., {"reason": "before_refactor"})
777 Returns:
778 Created snapshot
780 Example:
781 >>> snap = manager.create_snapshot(conv_id, metadata={"test": "baseline"})
782 >>> # Make changes...
783 >>> snap2 = manager.create_snapshot(conv_id, metadata={"test": "optimized"})
784 >>> diff = manager.compare_snapshots(snap.snapshot_id, snap2.snapshot_id)
785 """
786 messages = self.get_messages(conversation_id)
788 conv_data = self.storage.load_conversation(conversation_id)
789 conversation = Conversation.from_dict(conv_data)
791 # Serialize config to dict recursively (convert dataclasses, Paths, etc.)
792 from dataclasses import asdict, is_dataclass
793 from pathlib import Path
795 def _serialize_config(obj: Any) -> Any:
796 """Recursively convert config to JSON-safe dict."""
797 if isinstance(obj, Path):
798 return str(obj)
799 elif is_dataclass(obj) and not isinstance(obj, type):
800 # Convert dataclass instance to dict, then recursively serialize values
801 obj_dict = asdict(obj)
802 return {k: _serialize_config(v) for k, v in obj_dict.items()}
803 elif isinstance(obj, dict):
804 return {k: _serialize_config(v) for k, v in obj.items()}
805 elif isinstance(obj, (list, tuple)):
806 return [_serialize_config(item) for item in obj]
807 else:
808 return obj
810 config_dict = _serialize_config(self.config)
812 return self.snapshot_manager.create_snapshot(
813 conversation_id=conversation_id,
814 messages=messages,
815 conversation=conversation,
816 config=config_dict,
817 snapshot_id=snapshot_id,
818 metadata=metadata,
819 )
821 def restore_snapshot(self, snapshot_id: str) -> str:
822 """Restore conversation from snapshot.
824 WARNING: This replaces current conversation state with snapshot state.
825 Consider creating a snapshot of current state first.
827 Args:
828 snapshot_id: Snapshot to restore
830 Returns:
831 Conversation ID of restored conversation
833 Raises:
834 FileNotFoundError: If snapshot not found
835 """
836 snapshot = self.snapshot_manager.get_snapshot(snapshot_id)
837 if not snapshot:
838 raise FileNotFoundError(f"Snapshot {snapshot_id} not found")
840 messages, conversation = self.snapshot_manager.restore_snapshot(snapshot)
842 # Clear existing messages (replace entire conversation state)
843 # For memory storage, clear the messages dict for this conversation
844 if hasattr(self.storage, "messages"):
845 if snapshot.conversation_id in self.storage.messages:
846 self.storage.messages[snapshot.conversation_id] = {}
848 # Save conversation
849 self.storage.save_conversation(snapshot.conversation_id, conversation.to_dict())
851 # Save messages
852 for msg in messages:
853 self.storage.save_message(snapshot.conversation_id, msg.id, msg.to_dict())
855 # Clear cache
856 if snapshot.conversation_id in self._message_cache:
857 del self._message_cache[snapshot.conversation_id]
859 return snapshot.conversation_id
861 def compare_snapshots(
862 self, snapshot1_id: str, snapshot2_id: str
863 ) -> SnapshotDiff | None:
864 """Compare two snapshots to see what changed.
866 Args:
867 snapshot1_id: First snapshot ID (earlier)
868 snapshot2_id: Second snapshot ID (later)
870 Returns:
871 SnapshotDiff describing changes:
872 - messages_added: Number of messages added
873 - messages_removed: Number of messages removed
874 - tokens_delta: Change in token count
875 - message_changes: List of specific changes
876 - config_changes: Configuration modifications
877 - metadata_changes: Additional metadata
879 None if snapshots not found
880 """
881 return self.snapshot_manager.compare_snapshots(snapshot1_id, snapshot2_id)
883 def export_snapshot(self, snapshot_id: str, file_path: str) -> None:
884 """Export snapshot to JSON file for version control.
886 Args:
887 snapshot_id: Snapshot to export
888 file_path: Path to write JSON file
890 Raises:
891 FileNotFoundError: If snapshot not found
893 Example:
894 >>> snap = manager.create_snapshot(conv_id)
895 >>> manager.export_snapshot(snap.snapshot_id, "snapshots/baseline.json")
896 >>> # Commit to git for reproducibility
897 """
898 self.snapshot_manager.export_snapshot(snapshot_id, file_path)
900 def import_snapshot(self, file_path: str) -> Snapshot:
901 """Import snapshot from JSON file.
903 Args:
904 file_path: Path to JSON file
906 Returns:
907 Imported snapshot
908 """
909 return self.snapshot_manager.import_snapshot(file_path)
911 def cleanup_stale_data(
912 self,
913 conversation_id: str,
914 max_age_hours: float = 24,
915 keep_errors: bool = True,
916 execute: bool = False,
917 ) -> dict[str, Any]:
918 """Clean up stale conversation data.
920 By default, identifies old messages and tool outputs for cleanup based on age
921 without making any changes. Set execute=True to actually clear stale outputs.
923 Args:
924 conversation_id: Conversation to clean up
925 max_age_hours: Maximum age in hours before considering stale (default: 24)
926 keep_errors: Whether to preserve error messages (default: True)
927 execute: When True, actually clears stale outputs (default: False)
929 Returns:
930 Dictionary with cleanup statistics:
931 - messages_archived (int): Messages eligible for archival
932 - tokens_freed (int): Tokens that would be freed
933 - duplicates_removed (int): Duplicate outputs found
934 - stale_outputs_pruned (int): Stale outputs identified/cleared
935 - operations (list[str]): Description of operations
937 Example:
938 >>> result = manager.cleanup_stale_data(conv_id, max_age_hours=48)
939 >>> print(f"Can free {result['tokens_freed']} tokens")
940 >>> result = manager.cleanup_stale_data(conv_id, max_age_hours=48, execute=True)
941 >>> print(f"Freed {result['tokens_freed']} tokens")
942 """
943 from harnessutils.maintenance import cleanup_stale_data
945 messages = self.get_messages(conversation_id)
946 result = cleanup_stale_data(
947 messages=messages,
948 config=self.config.pruning,
949 max_age_hours=max_age_hours,
950 keep_errors=keep_errors,
951 execute=execute,
952 )
954 if execute:
955 for msg in messages:
956 self.storage.save_message(conversation_id, msg.id, msg.to_dict())
957 # Parts are stored separately — persist modified parts back to storage
958 for idx, part in enumerate(msg.parts):
959 self.storage.save_part(msg.id, f"{idx:04d}", part.to_dict())
960 # Invalidate cache — outputs were cleared in-place
961 if conversation_id in self._message_cache:
962 del self._message_cache[conversation_id]
964 return result.to_dict()
966 def scan_and_deduplicate(self, conversation_id: str) -> dict[str, Any]:
967 """Scan for duplicate outputs and return statistics.
969 Identifies duplicate tool outputs without removing them.
970 Use prune_before_turn() to actually remove duplicates.
972 Args:
973 conversation_id: Conversation to scan
975 Returns:
976 Dictionary with deduplication statistics:
977 - duplicates_removed (int): Number of duplicates found
978 - tokens_freed (int): Tokens that would be freed
979 - operations (list[str]): Description of findings
981 Example:
982 >>> result = manager.scan_and_deduplicate(conv_id)
983 >>> if result['duplicates_removed'] > 0:
984 ... print(f"Found {result['duplicates_removed']} duplicates")
985 ... # Run compaction to actually remove them
986 ... manager.prune_before_turn(conv_id)
987 """
988 from harnessutils.maintenance import scan_and_deduplicate
990 messages = self.get_messages(conversation_id)
991 result = scan_and_deduplicate(
992 messages=messages,
993 config=self.config.pruning,
994 )
996 return result.to_dict()
998 def get_memory(self, project_id: str, key: str, default: Any = None) -> Any:
999 """Get a value from project-scoped memory.
1001 Args:
1002 project_id: Project to retrieve memory from
1003 key: Memory key
1004 default: Default value if key not found
1006 Returns:
1007 Stored value or default
1008 """
1009 data = self._load_project_memory(project_id)
1010 return data.get(key, default)
1012 def set_memory(self, project_id: str, key: str, value: Any) -> None:
1013 """Set a value in project-scoped memory.
1015 Args:
1016 project_id: Project to store memory in
1017 key: Memory key
1018 value: Value to store
1019 """
1020 data = self._load_project_memory(project_id)
1021 data[key] = value
1022 try:
1023 self.storage.save_project_memory(project_id, data)
1024 except AttributeError:
1025 pass # Backend doesn't support project memory
1027 def delete_memory(self, project_id: str, key: str) -> None:
1028 """Delete a key from project-scoped memory.
1030 Args:
1031 project_id: Project to delete memory from
1032 key: Memory key to delete
1033 """
1034 data = self._load_project_memory(project_id)
1035 data.pop(key, None)
1036 try:
1037 self.storage.save_project_memory(project_id, data)
1038 except AttributeError:
1039 pass # Backend doesn't support project memory
1041 def list_memory(self, project_id: str) -> dict[str, Any]:
1042 """List all memory keys and values for a project.
1044 Args:
1045 project_id: Project to list memory for
1047 Returns:
1048 Dictionary of all memory key-value pairs
1049 """
1050 return self._load_project_memory(project_id)
1052 def _load_project_memory(self, project_id: str) -> dict[str, Any]:
1053 """Load project memory, returning empty dict on failure.
1055 Args:
1056 project_id: Project to load memory for
1058 Returns:
1059 Memory dictionary (empty if not found or backend unsupported)
1060 """
1061 try:
1062 return self.storage.load_project_memory(project_id)
1063 except (FileNotFoundError, AttributeError):
1064 return {}
1066 def detect_context_issues(self, conversation_id: str) -> list[dict[str, Any]]:
1067 """Detect quality and drift issues in conversation context.
1069 Analyzes conversation for common problems:
1070 - High redundancy (duplicate content)
1071 - Staleness accumulation (old messages)
1072 - Low information density
1073 - Error preservation status
1074 - Excessive protection (limiting pruning)
1076 Args:
1077 conversation_id: Conversation to analyze
1079 Returns:
1080 List of issue dictionaries, each containing:
1081 - issue_type (str): Type of issue
1082 - severity (str): "info", "warning", or "error"
1083 - description (str): Human-readable description
1084 - affected_count (int): Number of items affected
1085 - suggested_fix (str): Recommended action
1086 - metadata (dict): Additional details
1088 Example:
1089 >>> issues = manager.detect_context_issues(conv_id)
1090 >>> for issue in issues:
1091 ... if issue['severity'] == 'warning':
1092 ... print(f"⚠️ {issue['description']}")
1093 ... print(f" Fix: {issue['suggested_fix']}")
1094 """
1095 from harnessutils.maintenance import detect_context_issues
1097 conv_data = self.storage.load_conversation(conversation_id)
1098 conv = Conversation.from_dict(conv_data)
1099 messages = self.get_messages(conversation_id)
1101 issues = detect_context_issues(
1102 messages=messages,
1103 conversation=conv,
1104 config=self.config.pruning,
1105 )
1107 return [issue.to_dict() for issue in issues]