steerdev 0.4.27__py3-none-any.whl

steerdev_agent/workflow/memory.py

@@ -0,0 +1,185 @@
+"""Phase memory storage for workflow execution."""
+
+import json
+from pathlib import Path
+from typing import Any
+
+from loguru import logger
+from pydantic import BaseModel, Field
+
+
+class PhaseMemory(BaseModel):
+    """Memory stored for a completed phase."""
+
+    phase_id: str
+    phase_name: str
+    phase_type: str
+    input_context: dict[str, Any] = Field(default_factory=dict)
+    output_context: dict[str, Any] = Field(default_factory=dict)
+    result_summary: str | None = None
+
+
+class WorkflowMemoryManager:
+    """Manages local storage for phase context between executions.
+
+    Stores phase memory in `.steerdev/phases/{workflow_id}/{run_id}/{phase_id}.json`
+    to enable context propagation between phases. The two-level directory structure
+    ensures that multiple runs of the same workflow don't overwrite each other,
+    and allows debugging of failed runs by preserving their memory.
+    """
+
+    def __init__(self, working_directory: str | Path) -> None:
+        """Initialize the memory manager.
+
+        Args:
+            working_directory: Base directory for the project.
+        """
+        self.working_directory = Path(working_directory)
+        self.base_path = self.working_directory / ".steerdev" / "phases"
+
+    def _get_workflow_path(self, workflow_id: str) -> Path:
+        """Get the path for a workflow's memory directory."""
+        return self.base_path / workflow_id
+
+    def _get_run_path(self, workflow_id: str, run_id: str) -> Path:
+        """Get the path for a workflow run's memory files."""
+        return self._get_workflow_path(workflow_id) / run_id
+
+    def _get_phase_path(self, workflow_id: str, run_id: str, phase_id: str) -> Path:
+        """Get the path for a specific phase's memory file."""
+        return self._get_run_path(workflow_id, run_id) / f"{phase_id}.json"
+
+    def save_phase_memory(
+        self,
+        workflow_id: str,
+        run_id: str,
+        memory: PhaseMemory,
+    ) -> Path:
+        """Save phase memory after completion.
+
+        Args:
+            workflow_id: The workflow definition ID.
+            run_id: The workflow run ID.
+            memory: Phase memory to save.
+
+        Returns:
+            Path to the saved memory file.
+        """
+        run_path = self._get_run_path(workflow_id, run_id)
+        run_path.mkdir(parents=True, exist_ok=True)
+
+        phase_path = self._get_phase_path(workflow_id, run_id, memory.phase_id)
+        phase_path.write_text(memory.model_dump_json(indent=2))
+
+        logger.debug(f"Saved phase memory: {phase_path}")
+        return phase_path
+
+    def load_phase_memory(
+        self,
+        workflow_id: str,
+        run_id: str,
+        phase_id: str,
+    ) -> PhaseMemory | None:
+        """Load memory for a specific phase.
+
+        Args:
+            workflow_id: The workflow definition ID.
+            run_id: The workflow run ID.
+            phase_id: The phase ID to load.
+
+        Returns:
+            Phase memory or None if not found.
+        """
+        phase_path = self._get_phase_path(workflow_id, run_id, phase_id)
+
+        if not phase_path.exists():
+            return None
+
+        try:
+            data = json.loads(phase_path.read_text())
+            return PhaseMemory(**data)
+        except (json.JSONDecodeError, ValueError) as e:
+            logger.error(f"Failed to load phase memory: {e}")
+            return None
+
+    def load_all_phase_memories(
+        self,
+        workflow_id: str,
+        run_id: str,
+    ) -> list[PhaseMemory]:
+        """Load all phase memories for a workflow run.
+
+        Args:
+            workflow_id: The workflow definition ID.
+            run_id: The workflow run ID.
+
+        Returns:
+            List of phase memories, ordered by file modification time.
+        """
+        run_path = self._get_run_path(workflow_id, run_id)
+
+        if not run_path.exists():
+            return []
+
+        memories = []
+        for phase_file in sorted(run_path.glob("*.json"), key=lambda p: p.stat().st_mtime):
+            try:
+                data = json.loads(phase_file.read_text())
+                memories.append(PhaseMemory(**data))
+            except (json.JSONDecodeError, ValueError) as e:
+                logger.warning(f"Skipping invalid phase memory {phase_file}: {e}")
+
+        return memories
+
+    def build_accumulated_context(
+        self,
+        workflow_id: str,
+        run_id: str,
+    ) -> dict[str, Any]:
+        """Build accumulated context from all completed phases.
+
+        Merges output contexts from all phases, with later phases
+        taking precedence for duplicate keys.
+
+        Args:
+            workflow_id: The workflow definition ID.
+            run_id: The workflow run ID.
+
+        Returns:
+            Merged context dictionary.
+        """
+        memories = self.load_all_phase_memories(workflow_id, run_id)
+
+        accumulated: dict[str, Any] = {}
+        for memory in memories:
+            accumulated.update(memory.output_context)
+
+        return accumulated
+
+    def cleanup_run(self, workflow_id: str, run_id: str) -> bool:
+        """Remove memory files for a completed workflow run.
+
+        Note: This should only be called for successful workflow runs.
+        Failed runs should preserve their memory for debugging.
+
+        Args:
+            workflow_id: The workflow definition ID.
+            run_id: The workflow run ID to clean up.
+
+        Returns:
+            True if cleanup was successful.
+        """
+        run_path = self._get_run_path(workflow_id, run_id)
+
+        if not run_path.exists():
+            return True
+
+        try:
+            import shutil
+
+            shutil.rmtree(run_path)
+            logger.debug(f"Cleaned up workflow memory: {run_path}")
+            return True
+        except OSError as e:
+            logger.error(f"Failed to cleanup workflow memory: {e}")
+            return False