stabilize 0.9.2__py3-none-any.whl

stabilize/rag/cache.py

@@ -0,0 +1,294 @@
+"""Embedding cache implementations for RAG."""
+
+from __future__ import annotations
+
+import json
+import os
+import sqlite3
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    import psycopg
+
+
+@dataclass
+class CachedEmbedding:
+    """A cached embedding with its metadata."""
+
+    doc_id: str
+    content: str
+    embedding: list[float]
+    embedding_model: str
+    chunk_index: int
+
+
+class EmbeddingCache(ABC):
+    """Abstract interface for embedding cache."""
+
+    @abstractmethod
+    def store(self, embeddings: list[CachedEmbedding]) -> None:
+        """Store embeddings in cache."""
+        ...
+
+    @abstractmethod
+    def load(self, embedding_model: str) -> list[CachedEmbedding]:
+        """Load embeddings from cache for a specific model."""
+        ...
+
+    @abstractmethod
+    def is_initialized(self, embedding_model: str) -> bool:
+        """Check if cache has embeddings for the given model."""
+        ...
+
+    @abstractmethod
+    def clear(self) -> None:
+        """Clear all cached embeddings."""
+        ...
+
+
+class SqliteEmbeddingCache(EmbeddingCache):
+    """Store embeddings in SQLite database.
+
+    Default location: ~/.stabilize/embeddings.db
+    """
+
+    def __init__(self, db_path: str | None = None):
+        if db_path is None:
+            cache_dir = Path.home() / ".stabilize"
+            cache_dir.mkdir(parents=True, exist_ok=True)
+            db_path = str(cache_dir / "embeddings.db")
+
+        self.db_path = db_path
+        self._init_schema()
+
+    def _get_connection(self) -> sqlite3.Connection:
+        """Get a database connection."""
+        conn = sqlite3.connect(self.db_path)
+        conn.row_factory = sqlite3.Row
+        return conn
+
+    def _init_schema(self) -> None:
+        """Initialize the database schema."""
+        with self._get_connection() as conn:
+            conn.execute(
+                """
+                CREATE TABLE IF NOT EXISTS rag_embeddings (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    doc_id TEXT NOT NULL,
+                    content TEXT NOT NULL,
+                    embedding TEXT NOT NULL,
+                    embedding_model TEXT NOT NULL,
+                    chunk_index INTEGER DEFAULT 0,
+                    created_at TEXT DEFAULT (datetime('now')),
+                    UNIQUE(doc_id, chunk_index, embedding_model)
+                )
+            """
+            )
+            conn.execute(
+                """
+                CREATE INDEX IF NOT EXISTS idx_rag_embeddings_model
+                ON rag_embeddings(embedding_model)
+            """
+            )
+            conn.commit()
+
+    def store(self, embeddings: list[CachedEmbedding]) -> None:
+        """Store embeddings in SQLite."""
+        if not embeddings:
+            return
+
+        with self._get_connection() as conn:
+            # Clear existing embeddings for this model
+            model = embeddings[0].embedding_model
+            conn.execute(
+                "DELETE FROM rag_embeddings WHERE embedding_model = ?",
+                (model,),
+            )
+
+            # Insert new embeddings
+            for emb in embeddings:
+                conn.execute(
+                    """
+                    INSERT INTO rag_embeddings
+                    (doc_id, content, embedding, embedding_model, chunk_index)
+                    VALUES (?, ?, ?, ?, ?)
+                    """,
+                    (
+                        emb.doc_id,
+                        emb.content,
+                        json.dumps(emb.embedding),
+                        emb.embedding_model,
+                        emb.chunk_index,
+                    ),
+                )
+            conn.commit()
+
+    def load(self, embedding_model: str) -> list[CachedEmbedding]:
+        """Load embeddings from SQLite."""
+        with self._get_connection() as conn:
+            cursor = conn.execute(
+                """
+                SELECT doc_id, content, embedding, embedding_model, chunk_index
+                FROM rag_embeddings
+                WHERE embedding_model = ?
+                ORDER BY doc_id, chunk_index
+                """,
+                (embedding_model,),
+            )
+            results = []
+            for row in cursor:
+                results.append(
+                    CachedEmbedding(
+                        doc_id=row["doc_id"],
+                        content=row["content"],
+                        embedding=json.loads(row["embedding"]),
+                        embedding_model=row["embedding_model"],
+                        chunk_index=row["chunk_index"],
+                    )
+                )
+            return results
+
+    def is_initialized(self, embedding_model: str) -> bool:
+        """Check if cache has embeddings for the given model."""
+        with self._get_connection() as conn:
+            cursor = conn.execute(
+                "SELECT COUNT(*) FROM rag_embeddings WHERE embedding_model = ?",
+                (embedding_model,),
+            )
+            row = cursor.fetchone()
+            count: int = row[0] if row else 0
+            return count > 0
+
+    def clear(self) -> None:
+        """Clear all cached embeddings."""
+        with self._get_connection() as conn:
+            conn.execute("DELETE FROM rag_embeddings")
+            conn.commit()
+
+
+class PostgresEmbeddingCache(EmbeddingCache):
+    """Store embeddings in PostgreSQL database.
+
+    Uses the rag_embeddings table created by migration.
+    """
+
+    def __init__(self, connection_string: str):
+        self.connection_string = connection_string
+        self._conn: psycopg.Connection[tuple[Any, ...]] | None = None
+
+    def _get_connection(self) -> psycopg.Connection[tuple[Any, ...]]:
+        """Get a database connection."""
+        try:
+            import psycopg as psycopg_module
+        except ImportError as e:
+            raise ImportError("PostgreSQL support requires: pip install stabilize[postgres]") from e
+
+        if self._conn is None or self._conn.closed:
+            self._conn = psycopg_module.connect(self.connection_string)
+        return self._conn
+
+    def store(self, embeddings: list[CachedEmbedding]) -> None:
+        """Store embeddings in PostgreSQL."""
+        if not embeddings:
+            return
+
+        conn = self._get_connection()
+        with conn.cursor() as cur:
+            # Clear existing embeddings for this model
+            model = embeddings[0].embedding_model
+            cur.execute(
+                "DELETE FROM rag_embeddings WHERE embedding_model = %s",
+                (model,),
+            )
+
+            # Insert new embeddings
+            for emb in embeddings:
+                cur.execute(
+                    """
+                    INSERT INTO rag_embeddings
+                    (doc_id, content, embedding, embedding_model, chunk_index)
+                    VALUES (%s, %s, %s, %s, %s)
+                    """,
+                    (
+                        emb.doc_id,
+                        emb.content,
+                        json.dumps(emb.embedding),
+                        emb.embedding_model,
+                        emb.chunk_index,
+                    ),
+                )
+            conn.commit()
+
+    def load(self, embedding_model: str) -> list[CachedEmbedding]:
+        """Load embeddings from PostgreSQL."""
+        conn = self._get_connection()
+        with conn.cursor() as cur:
+            cur.execute(
+                """
+                SELECT doc_id, content, embedding, embedding_model, chunk_index
+                FROM rag_embeddings
+                WHERE embedding_model = %s
+                ORDER BY doc_id, chunk_index
+                """,
+                (embedding_model,),
+            )
+            results = []
+            for row in cur.fetchall():
+                embedding_data = row[2]
+                if isinstance(embedding_data, str):
+                    embedding_data = json.loads(embedding_data)
+                results.append(
+                    CachedEmbedding(
+                        doc_id=row[0],
+                        content=row[1],
+                        embedding=embedding_data,
+                        embedding_model=row[3],
+                        chunk_index=row[4],
+                    )
+                )
+            return results
+
+    def is_initialized(self, embedding_model: str) -> bool:
+        """Check if cache has embeddings for the given model."""
+        conn = self._get_connection()
+        with conn.cursor() as cur:
+            cur.execute(
+                "SELECT COUNT(*) FROM rag_embeddings WHERE embedding_model = %s",
+                (embedding_model,),
+            )
+            row = cur.fetchone()
+            count: int = row[0] if row else 0
+            return count > 0
+
+    def clear(self) -> None:
+        """Clear all cached embeddings."""
+        conn = self._get_connection()
+        with conn.cursor() as cur:
+            cur.execute("DELETE FROM rag_embeddings")
+            conn.commit()
+
+
+def get_cache(db_url: str | None = None) -> EmbeddingCache:
+    """Get appropriate cache based on configuration.
+
+    Priority:
+    1. Explicit db_url parameter
+    2. MG_DATABASE_URL environment variable (if postgres)
+    3. Default SQLite in ~/.stabilize/embeddings.db
+    """
+    if db_url:
+        if db_url.startswith("postgres"):
+            return PostgresEmbeddingCache(db_url)
+        else:
+            return SqliteEmbeddingCache(db_url)
+
+    # Check environment
+    env_url = os.environ.get("MG_DATABASE_URL")
+    if env_url and env_url.startswith("postgres"):
+        return PostgresEmbeddingCache(env_url)
+
+    # Default: SQLite
+    return SqliteEmbeddingCache()

stabilize/stages/__init__.py

@@ -0,0 +1,11 @@
+"""Stage definition builders for pipeline execution."""
+
+from stabilize.stages.builder import (
+    StageDefinitionBuilder,
+    StageDefinitionBuilderFactory,
+)
+
+__all__ = [
+    "StageDefinitionBuilder",
+    "StageDefinitionBuilderFactory",
+]

stabilize/stages/builder.py

@@ -0,0 +1,253 @@
+"""
+Stage definition builders.
+
+Stage definition builders are responsible for:
+1. Building tasks for a stage
+2. Building before stages (setup, validation)
+3. Building after stages (cleanup, notification)
+4. Building on-failure stages (rollback, alerts)
+"""
+
+from __future__ import annotations
+
+from abc import ABC
+from typing import TYPE_CHECKING
+
+from stabilize.dag.graph import StageGraphBuilder
+from stabilize.models.task import TaskExecution
+
+if TYPE_CHECKING:
+    from stabilize.models.stage import StageExecution
+
+
+class StageDefinitionBuilder(ABC):
+    """
+    Abstract base class for stage definition builders.
+
+    Each stage type has a corresponding builder that defines:
+    - What tasks the stage should execute
+    - What synthetic stages run before/after
+    - What happens on failure
+
+    Example:
+        class DeployStageBuilder(StageDefinitionBuilder):
+            @property
+            def type(self) -> str:
+                return "deploy"
+
+            def build_tasks(self, stage: StageExecution) -> List[TaskExecution]:
+                return [
+                    TaskExecution.create(
+                        name="Determine Target Server Group",
+                        implementing_class="DetermineTargetServerGroupTask",
+                        stage_start=True,
+                    ),
+                    TaskExecution.create(
+                        name="Deploy Server Group",
+                        implementing_class="DeployServerGroupTask",
+                        stage_end=True,
+                    ),
+                ]
+
+            def before_stages(
+                self,
+                stage: StageExecution,
+                graph: StageGraphBuilder,
+            ) -> None:
+                # Add validation stage before deploy
+                validation = StageExecution.create_synthetic(
+                    type="validation",
+                    name="Validate Deploy Configuration",
+                    parent=stage,
+                    owner=SyntheticStageOwner.STAGE_BEFORE,
+                )
+                graph.add(validation)
+    """
+
+    @property
+    def type(self) -> str:
+        """
+        Get the stage type this builder handles.
+
+        Defaults to lowercase class name without "StageBuilder" suffix.
+        """
+        name = self.__class__.__name__
+        if name.endswith("StageBuilder"):
+            name = name[:-12]
+        return name.lower()
+
+    @property
+    def aliases(self) -> list[str]:
+        """Get alternative names for this stage type."""
+        return []
+
+    def build_tasks(self, stage: StageExecution) -> list[TaskExecution]:
+        """
+        Build the tasks for this stage.
+
+        Override to define what tasks the stage should execute.
+
+        Args:
+            stage: The stage being built
+
+        Returns:
+            List of tasks to execute
+        """
+        return []
+
+    def before_stages(
+        self,
+        stage: StageExecution,
+        graph: StageGraphBuilder,
+    ) -> None:
+        """
+        Build synthetic stages that run before this stage's tasks.
+
+        Override to add setup, validation, or other pre-requisite stages.
+
+        Args:
+            stage: The parent stage
+            graph: Builder for adding synthetic stages
+        """
+        pass
+
+    def after_stages(
+        self,
+        stage: StageExecution,
+        graph: StageGraphBuilder,
+    ) -> None:
+        """
+        Build synthetic stages that run after this stage completes.
+
+        Override to add cleanup, notification, or other post-processing stages.
+
+        Args:
+            stage: The parent stage
+            graph: Builder for adding synthetic stages
+        """
+        pass
+
+    def on_failure_stages(
+        self,
+        stage: StageExecution,
+        graph: StageGraphBuilder,
+    ) -> None:
+        """
+        Build synthetic stages that run when this stage fails.
+
+        Override to add rollback, alerting, or other failure-handling stages.
+
+        Args:
+            stage: The failed stage
+            graph: Builder for adding synthetic stages
+        """
+        pass
+
+    def add_context_flags(self, stage: StageExecution) -> None:
+        """
+        Add any required context flags to the stage.
+
+        Called before task execution to set up stage context.
+
+        Args:
+            stage: The stage to modify
+        """
+        pass
+
+
+class NoOpStageBuilder(StageDefinitionBuilder):
+    """A stage builder that does nothing."""
+
+    @property
+    def type(self) -> str:
+        return "noop"
+
+
+class WaitStageBuilder(StageDefinitionBuilder):
+    """Builder for wait stages."""
+
+    @property
+    def type(self) -> str:
+        return "wait"
+
+    def build_tasks(self, stage: StageExecution) -> list[TaskExecution]:
+        return [
+            TaskExecution.create(
+                name="Wait",
+                implementing_class="WaitTask",
+                stage_start=True,
+                stage_end=True,
+            ),
+        ]
+
+
+class StageDefinitionBuilderFactory:
+    """Factory for resolving stage definition builders."""
+
+    def __init__(self) -> None:
+        self._builders: dict[str, StageDefinitionBuilder] = {}
+        self._default_builder = NoOpStageBuilder()
+
+        # Register built-in builders
+        self.register(NoOpStageBuilder())
+        self.register(WaitStageBuilder())
+
+    def register(self, builder: StageDefinitionBuilder) -> None:
+        """
+        Register a stage definition builder.
+
+        Args:
+            builder: The builder to register
+        """
+        self._builders[builder.type] = builder
+        for alias in builder.aliases:
+            self._builders[alias] = builder
+
+    def register_class(
+        self,
+        builder_class: type[StageDefinitionBuilder],
+    ) -> None:
+        """Register a builder by class."""
+        self.register(builder_class())
+
+    def get(self, stage_type: str) -> StageDefinitionBuilder:
+        """
+        Get the builder for a stage type.
+
+        Args:
+            stage_type: The stage type
+
+        Returns:
+            The builder (or default if not found)
+        """
+        return self._builders.get(stage_type, self._default_builder)
+
+    def has(self, stage_type: str) -> bool:
+        """Check if a builder is registered for a stage type."""
+        return stage_type in self._builders
+
+    def list_types(self) -> list[str]:
+        """Get all registered stage types."""
+        return list(self._builders.keys())
+
+
+# Global factory instance
+_default_factory: StageDefinitionBuilderFactory | None = None
+
+
+def get_default_factory() -> StageDefinitionBuilderFactory:
+    """Get the default global stage definition builder factory."""
+    global _default_factory
+    if _default_factory is None:
+        _default_factory = StageDefinitionBuilderFactory()
+    return _default_factory
+
+
+def register_builder(builder: StageDefinitionBuilder) -> None:
+    """Register a builder in the default factory."""
+    get_default_factory().register(builder)
+
+
+def get_builder(stage_type: str) -> StageDefinitionBuilder:
+    """Get a builder from the default factory."""
+    return get_default_factory().get(stage_type)

stabilize/tasks/__init__.py

@@ -0,0 +1,19 @@
+"""Task system for pipeline execution."""
+
+from stabilize.tasks.interface import (
+    CallableTask,
+    RetryableTask,
+    SkippableTask,
+    Task,
+)
+from stabilize.tasks.registry import TaskRegistry
+from stabilize.tasks.result import TaskResult
+
+__all__ = [
+    "TaskResult",
+    "Task",
+    "RetryableTask",
+    "CallableTask",
+    "SkippableTask",
+    "TaskRegistry",
+]