chuk-ai-session-manager 0.7.1__py3-none-any.whl → 0.8.1__py3-none-any.whl

chuk_ai_session_manager/procedural_memory/manager.py

@@ -0,0 +1,523 @@
+# chuk_ai_session_manager/procedural_memory/manager.py
+"""
+Tool Memory Manager - orchestrates procedural memory operations.
+
+Handles:
+- Recording tool invocations
+- Detecting fix relationships (failure -> success)
+- Updating aggregated patterns
+- Searching and retrieving history
+- Session persistence integration
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import logging
+from datetime import datetime, timezone
+from typing import Any, Callable, Optional, TYPE_CHECKING
+
+from pydantic import BaseModel, Field
+
+from chuk_ai_session_manager.procedural_memory.models import (
+    ProceduralMemory,
+    ToolFixRelation,
+    ToolLogEntry,
+    ToolOutcome,
+    ToolPattern,
+)
+
+if TYPE_CHECKING:
+    from chuk_ai_session_manager.models.session import Session
+
+log = logging.getLogger(__name__)
+
+
+class ToolMemoryManager(BaseModel):
+    """
+    Manages procedural memory - tool traces, patterns, and fixes.
+
+    This is the main interface for:
+    - Recording tool calls and their outcomes
+    - Detecting when a success fixes a prior failure
+    - Building aggregated patterns from traces
+    - Retrieving relevant history for context injection
+    - Persisting to/from Session state
+    """
+
+    memory: ProceduralMemory
+
+    # Configuration
+    max_log_entries: int = Field(default=1000)
+    max_patterns_per_tool: int = Field(default=10)
+    fix_detection_window: int = Field(default=10)  # Look back N calls for fixes
+
+    # Optional callbacks
+    on_fix_detected: Optional[Callable[[ToolFixRelation], None]] = Field(
+        default=None, exclude=True
+    )
+
+    model_config = {"arbitrary_types_allowed": True}
+
+    @classmethod
+    def create(cls, session_id: str, **kwargs: Any) -> "ToolMemoryManager":
+        """Create a new manager for a session."""
+        memory = ProceduralMemory(session_id=session_id)
+        return cls(memory=memory, **kwargs)
+
+    @classmethod
+    async def from_session(
+        cls, session: "Session", **kwargs: Any
+    ) -> "ToolMemoryManager":
+        """
+        Load procedural memory from a session's state.
+
+        Args:
+            session: The Session to load from
+            **kwargs: Additional configuration
+
+        Returns:
+            ToolMemoryManager with restored state
+        """
+        state_key = "procedural_memory"
+        stored = await session.get_state(state_key)
+
+        if stored:
+            try:
+                memory = ProceduralMemory.model_validate(stored)
+                log.info(f"Restored procedural memory for session {session.id}")
+                return cls(memory=memory, **kwargs)
+            except Exception as e:
+                log.warning(f"Failed to restore procedural memory: {e}")
+
+        # Create fresh
+        return cls.create(session_id=session.id, **kwargs)
+
+    async def save_to_session(self, session: "Session") -> None:
+        """
+        Persist procedural memory to a session's state.
+
+        Args:
+            session: The Session to save to
+        """
+        state_key = "procedural_memory"
+        await session.set_state(state_key, self.memory.model_dump(mode="json"))
+        log.debug(f"Saved procedural memory to session {session.id}")
+
+    @property
+    def session_id(self) -> str:
+        """Get the session ID."""
+        return self.memory.session_id
+
+    # --- Recording ---
+
+    async def record_call(
+        self,
+        tool_name: str,
+        arguments: dict[str, Any],
+        result: Any,
+        outcome: ToolOutcome,
+        context_goal: Optional[str] = None,
+        error_type: Optional[str] = None,
+        error_message: Optional[str] = None,
+        execution_time_ms: Optional[int] = None,
+        preceding_call_id: Optional[str] = None,
+    ) -> ToolLogEntry:
+        """
+        Record a tool invocation.
+
+        Args:
+            tool_name: Name of the tool called
+            arguments: Arguments passed to the tool
+            result: The result (will be summarized)
+            outcome: Success/failure status
+            context_goal: What the user was trying to accomplish
+            error_type: Type of error if failure
+            error_message: Error details if failure
+            execution_time_ms: How long the call took
+            preceding_call_id: ID of the call that came before this in a chain
+
+        Returns:
+            The created ToolLogEntry
+        """
+        call_id = self.memory.allocate_call_id()
+
+        entry = ToolLogEntry(
+            id=call_id,
+            timestamp=datetime.now(timezone.utc),
+            tool_name=tool_name,
+            arguments=arguments,
+            arguments_hash=self._hash_arguments(arguments),
+            context_goal=context_goal,
+            preceding_call_id=preceding_call_id,
+            outcome=outcome,
+            result_summary=self._summarize_result(result),
+            result_type=self._classify_result(result),
+            execution_time_ms=execution_time_ms,
+            error_type=error_type,
+            error_message=error_message,
+        )
+
+        # Add to log
+        self.memory.tool_log.append(entry)
+
+        # Enforce max size
+        if len(self.memory.tool_log) > self.max_log_entries:
+            self.memory.tool_log.pop(0)
+
+        # Update patterns
+        pattern = self.memory.get_pattern(tool_name)
+        pattern.record_call(entry)
+
+        # Handle success - check if it fixes a prior failure
+        if outcome == ToolOutcome.SUCCESS:
+            await self._check_if_fixes_prior(entry)
+
+            # Record success pattern if it was a fix
+            if entry.is_fix():
+                pattern.add_success_pattern(
+                    goal_match=context_goal,
+                    delta_that_fixed=entry.delta_args,
+                    example_call_id=call_id,
+                )
+
+        # Handle failure - record error pattern
+        elif outcome in (ToolOutcome.FAILURE, ToolOutcome.TIMEOUT):
+            pattern.add_error_pattern(
+                error_type=error_type or "unknown",
+                context=context_goal,
+                example_args=arguments,
+            )
+
+        self.memory.updated_at = datetime.now(timezone.utc)
+
+        log.debug(f"Recorded tool call: {entry.format_compact()}")
+        return entry
+
+    async def _check_if_fixes_prior(self, success_entry: ToolLogEntry) -> None:
+        """
+        Check if this successful call fixes a recent failure.
+
+        Looks back through recent calls for the same tool that failed,
+        and links them if found.
+        """
+        # Look at recent calls (excluding this one)
+        window = self.memory.tool_log[-(self.fix_detection_window + 1) : -1]
+
+        for prior in reversed(window):
+            # Must be same tool
+            if prior.tool_name != success_entry.tool_name:
+                continue
+
+            # Must be a failure
+            if not prior.is_failure():
+                continue
+
+            # Must not already be fixed
+            if prior.was_fixed():
+                continue
+
+            # Compute what changed
+            delta = self._compute_arg_delta(prior.arguments, success_entry.arguments)
+
+            if delta:
+                # Link the fix
+                prior.fixed_by = success_entry.id
+                success_entry.fix_for = prior.id
+                success_entry.delta_args = delta
+
+                # Record the fix relation
+                relation = ToolFixRelation(
+                    failed_call_id=prior.id,
+                    success_call_id=success_entry.id,
+                    delta_args=delta,
+                )
+                self.memory.fix_relations.append(relation)
+
+                # Update error pattern with fix info
+                if prior.error_type:
+                    pattern = self.memory.get_pattern(success_entry.tool_name)
+                    pattern.record_fix(prior.error_type, delta)
+
+                log.info(
+                    f"Detected fix: {prior.id} -> {success_entry.id}, delta: {delta}"
+                )
+
+                # Callback
+                if self.on_fix_detected:
+                    self.on_fix_detected(relation)
+
+                # Only link to most recent failure
+                break
+
+    def _compute_arg_delta(
+        self, failed_args: dict[str, Any], success_args: dict[str, Any]
+    ) -> Optional[dict[str, Any]]:
+        """Compute what changed between failed and successful call."""
+        delta: dict[str, Any] = {}
+
+        failed_keys = set(failed_args.keys())
+        success_keys = set(success_args.keys())
+
+        # Added keys
+        added = success_keys - failed_keys
+        if added:
+            delta["added"] = {k: success_args[k] for k in added}
+
+        # Removed keys
+        removed = failed_keys - success_keys
+        if removed:
+            delta["removed"] = list(removed)
+
+        # Changed values
+        changed = {}
+        for k in failed_keys & success_keys:
+            if not self._values_equal(failed_args[k], success_args[k]):
+                changed[k] = {"from": failed_args[k], "to": success_args[k]}
+        if changed:
+            delta["changed"] = changed
+
+        return delta if delta else None
+
+    def _values_equal(self, a: Any, b: Any) -> bool:
+        """Check if two values are equal (handling nested structures)."""
+        try:
+            return json.dumps(a, sort_keys=True) == json.dumps(b, sort_keys=True)
+        except (TypeError, ValueError):
+            return a == b
+
+    def _hash_arguments(self, arguments: dict[str, Any]) -> str:
+        """Create hash of arguments for quick comparison."""
+        try:
+            args_str = json.dumps(arguments, sort_keys=True, default=str)
+            return hashlib.sha256(args_str.encode()).hexdigest()[:12]
+        except (TypeError, ValueError):
+            return ""
+
+    def _summarize_result(self, result: Any, max_len: int = 100) -> str:
+        """Create compact summary of a result."""
+        if result is None:
+            return "null"
+
+        if isinstance(result, str):
+            if len(result) <= max_len:
+                return result
+            return result[: max_len - 3] + "..."
+
+        if isinstance(result, bool):
+            return "true" if result else "false"
+
+        if isinstance(result, (int, float)):
+            return str(result)
+
+        if isinstance(result, list):
+            return f"list[{len(result)}]"
+
+        if isinstance(result, dict):
+            # Check for common patterns
+            if "status" in result:
+                return f"status={result['status']}"
+            if "error" in result:
+                return f"error: {result['error']}"
+            if "result" in result:
+                return self._summarize_result(result["result"], max_len)
+            return f"object[{len(result)} keys]"
+
+        return str(type(result).__name__)
+
+    def _classify_result(self, result: Any) -> str:
+        """Classify the type of result."""
+        if result is None:
+            return "null"
+        if isinstance(result, bool):
+            return "boolean"
+        if isinstance(result, (int, float)):
+            return "number"
+        if isinstance(result, str):
+            return "string"
+        if isinstance(result, list):
+            return "list"
+        if isinstance(result, dict):
+            if "status" in result:
+                return result.get("status", "object")
+            return "object"
+        return "unknown"
+
+    # --- Retrieval ---
+
+    def get_recent_calls(
+        self,
+        tool_name: Optional[str] = None,
+        limit: int = 5,
+        outcome: Optional[ToolOutcome] = None,
+    ) -> list[ToolLogEntry]:
+        """
+        Get recent tool calls.
+
+        Args:
+            tool_name: Filter by tool name (None = all tools)
+            limit: Maximum entries to return
+            outcome: Filter by outcome (None = all outcomes)
+
+        Returns:
+            List of entries, most recent first
+        """
+        entries = self.memory.tool_log
+
+        if tool_name:
+            entries = [e for e in entries if e.tool_name == tool_name]
+
+        if outcome:
+            entries = [e for e in entries if e.outcome == outcome]
+
+        return list(reversed(entries[-limit:]))
+
+    def get_pattern(self, tool_name: str) -> Optional[ToolPattern]:
+        """Get aggregated pattern for a tool."""
+        return self.memory.tool_patterns.get(tool_name)
+
+    def get_all_patterns(self) -> dict[str, ToolPattern]:
+        """Get all tool patterns."""
+        return self.memory.tool_patterns.copy()
+
+    def search_calls(
+        self,
+        tool_name: Optional[str] = None,
+        goal_contains: Optional[str] = None,
+        outcome: Optional[ToolOutcome] = None,
+        error_type: Optional[str] = None,
+        only_fixes: bool = False,
+        only_fixed: bool = False,
+        limit: int = 10,
+    ) -> list[ToolLogEntry]:
+        """
+        Search tool log with filters.
+
+        Args:
+            tool_name: Filter by tool name
+            goal_contains: Filter by goal containing text
+            outcome: Filter by outcome
+            error_type: Filter by error type
+            only_fixes: Only return calls that fixed something
+            only_fixed: Only return failures that were fixed
+            limit: Maximum results
+
+        Returns:
+            Matching entries, most recent first
+        """
+        results = []
+
+        for entry in reversed(self.memory.tool_log):
+            # Apply filters
+            if tool_name and entry.tool_name != tool_name:
+                continue
+
+            if goal_contains:
+                if not entry.context_goal:
+                    continue
+                if goal_contains.lower() not in entry.context_goal.lower():
+                    continue
+
+            if outcome and entry.outcome != outcome:
+                continue
+
+            if error_type and entry.error_type != error_type:
+                continue
+
+            if only_fixes and not entry.is_fix():
+                continue
+
+            if only_fixed and not entry.was_fixed():
+                continue
+
+            results.append(entry)
+
+            if len(results) >= limit:
+                break
+
+        return results
+
+    def get_fix_for_error(
+        self, tool_name: str, error_type: str
+    ) -> Optional[dict[str, Any]]:
+        """
+        Get the typical fix for an error type.
+
+        Returns the delta_args that typically fixes this error,
+        based on observed fix relationships.
+        """
+        pattern = self.get_pattern(tool_name)
+        if not pattern:
+            return None
+
+        for ep in pattern.error_patterns:
+            if ep.error_type == error_type and ep.fix_delta:
+                return ep.fix_delta
+
+        return None
+
+    def get_successful_args_for_goal(
+        self, tool_name: str, goal: str
+    ) -> Optional[dict[str, Any]]:
+        """
+        Get argument hints for a goal based on past successes.
+
+        Searches success patterns for matching goals and returns
+        argument hints.
+        """
+        pattern = self.get_pattern(tool_name)
+        if not pattern:
+            return None
+
+        goal_lower = goal.lower()
+        for sp in pattern.success_patterns:
+            if sp.goal_match and goal_lower in sp.goal_match.lower():
+                return sp.arg_hints
+
+        return None
+
+    # --- Statistics ---
+
+    def get_stats(self) -> dict[str, Any]:
+        """Get memory statistics."""
+        total_calls = len(self.memory.tool_log)
+        total_fixes = len(self.memory.fix_relations)
+        tools_tracked = len(self.memory.tool_patterns)
+
+        success_count = sum(
+            1 for e in self.memory.tool_log if e.outcome == ToolOutcome.SUCCESS
+        )
+        failure_count = sum(1 for e in self.memory.tool_log if e.is_failure())
+
+        return {
+            "session_id": self.session_id,
+            "total_calls": total_calls,
+            "success_count": success_count,
+            "failure_count": failure_count,
+            "success_rate": success_count / total_calls if total_calls > 0 else 0,
+            "total_fixes_detected": total_fixes,
+            "tools_tracked": tools_tracked,
+            "created_at": self.memory.created_at.isoformat(),
+            "updated_at": self.memory.updated_at.isoformat(),
+        }
+
+    # --- Persistence hooks ---
+
+    def to_dict(self) -> dict[str, Any]:
+        """Serialize to dictionary for persistence."""
+        return self.memory.model_dump(mode="json")
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any], **kwargs: Any) -> "ToolMemoryManager":
+        """Restore from dictionary."""
+        memory = ProceduralMemory.model_validate(data)
+        return cls(memory=memory, **kwargs)
+
+    def reset(self) -> None:
+        """Clear all memory."""
+        self.memory.tool_log.clear()
+        self.memory.tool_patterns.clear()
+        self.memory.fix_relations.clear()
+        self.memory.next_call_id = 1
+        self.memory.updated_at = datetime.now(timezone.utc)