PraisonAI 3.0.0__py3-none-any.whl

praisonai/persistence/knowledge/weaviate.py

@@ -0,0 +1,223 @@
+"""
+Weaviate implementation of KnowledgeStore.
+
+Requires: weaviate-client
+Install: pip install weaviate-client
+"""
+
+import logging
+import os
+from typing import Any, Dict, List, Optional
+
+from .base import KnowledgeStore, KnowledgeDocument
+
+logger = logging.getLogger(__name__)
+
+
+class WeaviateKnowledgeStore(KnowledgeStore):
+    """
+    Weaviate-based knowledge store for vector search.
+    
+    Example:
+        store = WeaviateKnowledgeStore(
+            url="http://localhost:8080"
+        )
+    """
+    
+    def __init__(
+        self,
+        url: Optional[str] = None,
+        api_key: Optional[str] = None,
+        additional_headers: Optional[Dict[str, str]] = None,
+    ):
+        try:
+            import weaviate
+        except ImportError:
+            raise ImportError(
+                "weaviate-client is required for Weaviate support. "
+                "Install with: pip install weaviate-client"
+            )
+        
+        url = url or os.getenv("WEAVIATE_URL", "http://localhost:8080")
+        api_key = api_key or os.getenv("WEAVIATE_API_KEY")
+        
+        auth = weaviate.auth.AuthApiKey(api_key) if api_key else None
+        self._client = weaviate.connect_to_custom(
+            http_host=url.replace("http://", "").replace("https://", "").split(":")[0],
+            http_port=int(url.split(":")[-1]) if ":" in url.split("/")[-1] else 8080,
+            http_secure=url.startswith("https"),
+            auth_credentials=auth,
+            additional_headers=additional_headers,
+        )
+        logger.info(f"Connected to Weaviate at {url}")
+    
+    def create_collection(
+        self,
+        name: str,
+        dimension: int,
+        distance: str = "cosine",
+        metadata: Optional[Dict[str, Any]] = None
+    ) -> None:
+        """Create a new collection (class in Weaviate)."""
+        from weaviate.classes.config import Configure, Property, DataType
+        
+        distance_map = {"cosine": "cosine", "euclidean": "l2-squared", "dot": "dot"}
+        
+        self._client.collections.create(
+            name=name,
+            vectorizer_config=Configure.Vectorizer.none(),
+            vector_index_config=Configure.VectorIndex.hnsw(
+                distance_metric=distance_map.get(distance, "cosine")
+            ),
+            properties=[
+                Property(name="content", data_type=DataType.TEXT),
+                Property(name="content_hash", data_type=DataType.TEXT),
+                Property(name="created_at", data_type=DataType.NUMBER),
+            ]
+        )
+        logger.info(f"Created Weaviate collection: {name}")
+    
+    def delete_collection(self, name: str) -> bool:
+        """Delete a collection."""
+        try:
+            self._client.collections.delete(name)
+            return True
+        except Exception as e:
+            logger.warning(f"Failed to delete collection {name}: {e}")
+            return False
+    
+    def collection_exists(self, name: str) -> bool:
+        """Check if a collection exists."""
+        return self._client.collections.exists(name)
+    
+    def list_collections(self) -> List[str]:
+        """List all collections."""
+        return [c.name for c in self._client.collections.list_all().values()]
+    
+    def insert(
+        self,
+        collection: str,
+        documents: List[KnowledgeDocument]
+    ) -> List[str]:
+        """Insert documents."""
+        col = self._client.collections.get(collection)
+        ids = []
+        
+        with col.batch.dynamic() as batch:
+            for doc in documents:
+                if doc.embedding is None:
+                    raise ValueError(f"Document {doc.id} has no embedding")
+                
+                properties = {
+                    "content": doc.content,
+                    "content_hash": doc.content_hash or "",
+                    "created_at": doc.created_at,
+                    **(doc.metadata or {})
+                }
+                
+                batch.add_object(
+                    properties=properties,
+                    vector=doc.embedding,
+                    uuid=doc.id,
+                )
+                ids.append(doc.id)
+        
+        return ids
+    
+    def upsert(
+        self,
+        collection: str,
+        documents: List[KnowledgeDocument]
+    ) -> List[str]:
+        """Insert or update documents."""
+        return self.insert(collection, documents)
+    
+    def search(
+        self,
+        collection: str,
+        query_embedding: List[float],
+        limit: int = 5,
+        filters: Optional[Dict[str, Any]] = None,
+        score_threshold: Optional[float] = None
+    ) -> List[KnowledgeDocument]:
+        """Search for similar documents."""
+        col = self._client.collections.get(collection)
+        
+        query = col.query.near_vector(
+            near_vector=query_embedding,
+            limit=limit,
+            return_metadata=["distance"],
+        )
+        
+        documents = []
+        for obj in query.objects:
+            if score_threshold:
+                score = 1 - (obj.metadata.distance or 0)
+                if score < score_threshold:
+                    continue
+            
+            props = obj.properties
+            doc = KnowledgeDocument(
+                id=str(obj.uuid),
+                content=props.get("content", ""),
+                embedding=None,
+                metadata={k: v for k, v in props.items() if k not in ("content", "content_hash", "created_at")},
+                content_hash=props.get("content_hash"),
+                created_at=props.get("created_at", 0),
+            )
+            documents.append(doc)
+        
+        return documents
+    
+    def get(
+        self,
+        collection: str,
+        ids: List[str]
+    ) -> List[KnowledgeDocument]:
+        """Get documents by IDs."""
+        col = self._client.collections.get(collection)
+        documents = []
+        
+        for doc_id in ids:
+            try:
+                obj = col.query.fetch_object_by_id(doc_id, include_vector=True)
+                if obj:
+                    props = obj.properties
+                    doc = KnowledgeDocument(
+                        id=str(obj.uuid),
+                        content=props.get("content", ""),
+                        embedding=list(obj.vector) if obj.vector else None,
+                        metadata={k: v for k, v in props.items() if k not in ("content", "content_hash", "created_at")},
+                        content_hash=props.get("content_hash"),
+                        created_at=props.get("created_at", 0),
+                    )
+                    documents.append(doc)
+            except Exception:
+                pass
+        
+        return documents
+    
+    def delete(
+        self,
+        collection: str,
+        ids: Optional[List[str]] = None,
+        filters: Optional[Dict[str, Any]] = None
+    ) -> int:
+        """Delete documents."""
+        col = self._client.collections.get(collection)
+        
+        if ids:
+            for doc_id in ids:
+                col.data.delete_by_id(doc_id)
+            return len(ids)
+        return 0
+    
+    def count(self, collection: str) -> int:
+        """Count documents."""
+        col = self._client.collections.get(collection)
+        return col.aggregate.over_all(total_count=True).total_count or 0
+    
+    def close(self) -> None:
+        """Close the store."""
+        if self._client:
+            self._client.close()

praisonai/persistence/migrations/__init__.py

@@ -0,0 +1,10 @@
+"""
+Schema versioning and migrations for PraisonAI persistence.
+
+This module provides tools for managing database schema versions
+and applying migrations when upgrading.
+"""
+
+from .manager import MigrationManager, SchemaVersion
+
+__all__ = ["MigrationManager", "SchemaVersion"]

praisonai/persistence/migrations/manager.py

@@ -0,0 +1,251 @@
+"""
+Migration manager for PraisonAI persistence schema versioning.
+"""
+
+import logging
+from dataclasses import dataclass, field
+from typing import Any, Callable, Dict, List, Optional, Tuple
+from packaging import version as pkg_version
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class SchemaVersion:
+    """Represents a schema version."""
+    version: str
+    description: str = ""
+    applied_at: Optional[float] = None
+    
+    def __lt__(self, other: "SchemaVersion") -> bool:
+        return pkg_version.parse(self.version) < pkg_version.parse(other.version)
+    
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, SchemaVersion):
+            return False
+        return self.version == other.version
+
+
+@dataclass
+class Migration:
+    """Represents a single migration."""
+    version: str
+    description: str
+    up: Callable[[Any], None]
+    down: Optional[Callable[[Any], None]] = None
+
+
+class MigrationManager:
+    """
+    Manages database schema versions and migrations.
+    
+    Example:
+        manager = MigrationManager(db_adapter)
+        manager.register_migration(
+            version="1.0.0",
+            description="Initial schema",
+            up=lambda db: db.execute("CREATE TABLE ..."),
+            down=lambda db: db.execute("DROP TABLE ...")
+        )
+        manager.migrate_up()  # Apply all pending migrations
+    """
+    
+    # Current schema version
+    CURRENT_VERSION = "1.0.0"
+    
+    def __init__(self, db_adapter: Any = None):
+        """
+        Initialize migration manager.
+        
+        Args:
+            db_adapter: Database adapter (PraisonDB or similar)
+        """
+        self._db = db_adapter
+        self._migrations: List[Migration] = []
+        self._versions_table = "praison_schema_versions"
+        
+        # Register built-in migrations
+        self._register_builtin_migrations()
+    
+    def _register_builtin_migrations(self):
+        """Register built-in migrations."""
+        self.register_migration(
+            version="1.0.0",
+            description="Initial schema with sessions, messages, runs",
+            up=self._migrate_1_0_0_up,
+            down=self._migrate_1_0_0_down
+        )
+    
+    def register_migration(
+        self,
+        version: str,
+        description: str,
+        up: Callable[[Any], None],
+        down: Optional[Callable[[Any], None]] = None
+    ):
+        """
+        Register a new migration.
+        
+        Args:
+            version: Version string (e.g., "1.0.0")
+            description: Human-readable description
+            up: Function to apply migration
+            down: Optional function to rollback migration
+        """
+        migration = Migration(
+            version=version,
+            description=description,
+            up=up,
+            down=down
+        )
+        self._migrations.append(migration)
+        self._migrations.sort(key=lambda m: pkg_version.parse(m.version))
+    
+    def get_current_version(self) -> Optional[str]:
+        """Get the current schema version from database."""
+        if not self._db:
+            return None
+        
+        try:
+            # Try to get version from state store
+            if hasattr(self._db, '_state_store') and self._db._state_store:
+                self._db._init_stores()
+                version_data = self._db._state_store.get(f"{self._versions_table}:current")
+                if version_data:
+                    return version_data.get("version")
+        except Exception as e:
+            logger.warning(f"Failed to get schema version: {e}")
+        
+        return None
+    
+    def set_current_version(self, version: str):
+        """Set the current schema version in database."""
+        if not self._db:
+            return
+        
+        try:
+            if hasattr(self._db, '_state_store') and self._db._state_store:
+                self._db._init_stores()
+                import time
+                self._db._state_store.set(f"{self._versions_table}:current", {
+                    "version": version,
+                    "applied_at": time.time()
+                })
+        except Exception as e:
+            logger.warning(f"Failed to set schema version: {e}")
+    
+    def get_pending_migrations(self) -> List[Migration]:
+        """Get list of migrations that haven't been applied."""
+        current = self.get_current_version()
+        if not current:
+            return self._migrations.copy()
+        
+        current_ver = pkg_version.parse(current)
+        return [
+            m for m in self._migrations
+            if pkg_version.parse(m.version) > current_ver
+        ]
+    
+    def migrate_up(self, target_version: Optional[str] = None) -> List[str]:
+        """
+        Apply pending migrations up to target version.
+        
+        Args:
+            target_version: Optional target version (default: latest)
+            
+        Returns:
+            List of applied migration versions
+        """
+        pending = self.get_pending_migrations()
+        if not pending:
+            logger.info("No pending migrations")
+            return []
+        
+        target = pkg_version.parse(target_version) if target_version else None
+        applied = []
+        
+        for migration in pending:
+            if target and pkg_version.parse(migration.version) > target:
+                break
+            
+            logger.info(f"Applying migration {migration.version}: {migration.description}")
+            try:
+                migration.up(self._db)
+                self.set_current_version(migration.version)
+                applied.append(migration.version)
+                logger.info(f"Successfully applied migration {migration.version}")
+            except Exception as e:
+                logger.error(f"Failed to apply migration {migration.version}: {e}")
+                raise
+        
+        return applied
+    
+    def migrate_down(self, target_version: str) -> List[str]:
+        """
+        Rollback migrations down to target version.
+        
+        Args:
+            target_version: Target version to rollback to
+            
+        Returns:
+            List of rolled back migration versions
+        """
+        current = self.get_current_version()
+        if not current:
+            logger.info("No migrations to rollback")
+            return []
+        
+        current_ver = pkg_version.parse(current)
+        target_ver = pkg_version.parse(target_version)
+        
+        if target_ver >= current_ver:
+            logger.info("Target version is not lower than current")
+            return []
+        
+        # Get migrations to rollback (in reverse order)
+        to_rollback = [
+            m for m in reversed(self._migrations)
+            if pkg_version.parse(m.version) > target_ver and pkg_version.parse(m.version) <= current_ver
+        ]
+        
+        rolled_back = []
+        for migration in to_rollback:
+            if not migration.down:
+                logger.warning(f"Migration {migration.version} has no down function, skipping")
+                continue
+            
+            logger.info(f"Rolling back migration {migration.version}")
+            try:
+                migration.down(self._db)
+                rolled_back.append(migration.version)
+                logger.info(f"Successfully rolled back migration {migration.version}")
+            except Exception as e:
+                logger.error(f"Failed to rollback migration {migration.version}: {e}")
+                raise
+        
+        self.set_current_version(target_version)
+        return rolled_back
+    
+    def get_migration_status(self) -> Dict[str, Any]:
+        """Get current migration status."""
+        current = self.get_current_version()
+        pending = self.get_pending_migrations()
+        
+        return {
+            "current_version": current,
+            "latest_version": self._migrations[-1].version if self._migrations else None,
+            "pending_count": len(pending),
+            "pending_versions": [m.version for m in pending],
+            "all_versions": [m.version for m in self._migrations]
+        }
+    
+    # --- Built-in Migrations ---
+    
+    def _migrate_1_0_0_up(self, db: Any):
+        """Initial schema migration."""
+        # This is a no-op since tables are created on first use
+        logger.info("Initial schema (1.0.0) - tables created on first use")
+    
+    def _migrate_1_0_0_down(self, db: Any):
+        """Rollback initial schema."""
+        logger.info("Rolling back initial schema (1.0.0)")