PraisonAI 3.0.0__py3-none-any.whl

praisonai/persistence/hooks/agent_hooks.py

@@ -0,0 +1,297 @@
+"""
+Agent hooks for automatic persistence integration.
+
+Provides wrapper functions to add persistence capabilities to PraisonAI agents
+without modifying the core SDK.
+"""
+
+import time
+import uuid
+import logging
+from typing import Any, Dict, List, Optional, Callable
+from functools import wraps
+
+from ..orchestrator import PersistenceOrchestrator
+from ..conversation.base import ConversationMessage, ConversationSession
+
+logger = logging.getLogger(__name__)
+
+
+def wrap_agent_with_persistence(
+    agent,
+    orchestrator: PersistenceOrchestrator,
+    session_id: Optional[str] = None,
+    user_id: Optional[str] = None,
+    auto_resume: bool = True,
+):
+    """
+    Wrap an agent to add automatic persistence hooks.
+    
+    This function wraps the agent's chat method to automatically:
+    1. Load/create session on first message
+    2. Persist each user message
+    3. Persist each assistant response
+    4. Update session metadata after each interaction
+    
+    Args:
+        agent: PraisonAI Agent instance
+        orchestrator: PersistenceOrchestrator instance
+        session_id: Session ID (auto-generated if None)
+        user_id: User ID for the session
+        auto_resume: Whether to auto-resume existing session
+    
+    Returns:
+        The wrapped agent with persistence capabilities
+    
+    Example:
+        from praisonai.persistence import PersistenceOrchestrator
+        from praisonai.persistence.hooks import wrap_agent_with_persistence
+        from praisonaiagents import Agent
+        
+        orchestrator = PersistenceOrchestrator(...)
+        agent = Agent(name="Assistant", role="Helper")
+        
+        # Wrap agent with persistence
+        agent = wrap_agent_with_persistence(agent, orchestrator, session_id="my-session")
+        
+        # Now all chats are automatically persisted
+        response = agent.chat("Hello!")
+    """
+    session_id = session_id or f"session-{uuid.uuid4().hex[:8]}"
+    user_id = user_id or getattr(agent, 'user_id', 'default')
+    
+    # Store original chat method
+    original_chat = agent.chat
+    _session_initialized = [False]  # Use list to allow mutation in closure
+    
+    @wraps(original_chat)
+    def chat_with_persistence(prompt, *args, **kwargs):
+        nonlocal _session_initialized
+        
+        # Initialize session on first call
+        if not _session_initialized[0]:
+            history = orchestrator.on_agent_start(
+                agent, 
+                session_id=session_id, 
+                user_id=user_id,
+                resume=auto_resume
+            )
+            
+            # Inject history into agent's chat_history if resuming
+            if auto_resume and history:
+                agent.chat_history = [
+                    {"role": msg.role, "content": msg.content}
+                    for msg in history
+                ]
+                logger.info(f"Resumed session {session_id} with {len(history)} messages")
+            
+            _session_initialized[0] = True
+        
+        # Persist user message
+        orchestrator.on_message(session_id, "user", prompt if isinstance(prompt, str) else str(prompt))
+        
+        # Call original chat
+        response = original_chat(prompt, *args, **kwargs)
+        
+        # Persist assistant response
+        if response:
+            orchestrator.on_message(session_id, "assistant", response)
+        
+        return response
+    
+    # Replace chat method
+    agent.chat = chat_with_persistence
+    
+    # Add persistence-related methods to agent
+    agent._persistence_orchestrator = orchestrator
+    agent._persistence_session_id = session_id
+    
+    def get_session():
+        """Get the current session."""
+        return orchestrator.get_session(session_id)
+    
+    def get_messages(limit: Optional[int] = None):
+        """Get messages from the current session."""
+        return orchestrator.get_messages(session_id, limit=limit)
+    
+    def end_session():
+        """End the current session."""
+        orchestrator.on_agent_end(agent, session_id)
+    
+    agent.get_session = get_session
+    agent.get_messages = get_messages
+    agent.end_session = end_session
+    
+    return agent
+
+
+class PersistentAgent:
+    """
+    A wrapper class that provides persistence capabilities for any agent.
+    
+    This is an alternative to wrap_agent_with_persistence that uses composition
+    instead of monkey-patching.
+    
+    Example:
+        from praisonai.persistence import PersistenceOrchestrator
+        from praisonai.persistence.hooks import PersistentAgent
+        from praisonaiagents import Agent
+        
+        orchestrator = PersistenceOrchestrator(...)
+        base_agent = Agent(name="Assistant", role="Helper")
+        
+        agent = PersistentAgent(base_agent, orchestrator, session_id="my-session")
+        response = agent.chat("Hello!")
+    """
+    
+    def __init__(
+        self,
+        agent,
+        orchestrator: PersistenceOrchestrator,
+        session_id: Optional[str] = None,
+        user_id: Optional[str] = None,
+        auto_resume: bool = True,
+    ):
+        self._agent = agent
+        self._orchestrator = orchestrator
+        self._session_id = session_id or f"session-{uuid.uuid4().hex[:8]}"
+        self._user_id = user_id or getattr(agent, 'user_id', 'default')
+        self._auto_resume = auto_resume
+        self._session_initialized = False
+    
+    def __getattr__(self, name):
+        """Delegate attribute access to wrapped agent."""
+        return getattr(self._agent, name)
+    
+    def _ensure_session(self):
+        """Ensure session is initialized."""
+        if not self._session_initialized:
+            history = self._orchestrator.on_agent_start(
+                self._agent,
+                session_id=self._session_id,
+                user_id=self._user_id,
+                resume=self._auto_resume
+            )
+            
+            if self._auto_resume and history:
+                self._agent.chat_history = [
+                    {"role": msg.role, "content": msg.content}
+                    for msg in history
+                ]
+                logger.info(f"Resumed session {self._session_id} with {len(history)} messages")
+            
+            self._session_initialized = True
+    
+    def chat(self, prompt, *args, **kwargs):
+        """Chat with automatic persistence."""
+        self._ensure_session()
+        
+        # Persist user message
+        self._orchestrator.on_message(
+            self._session_id, "user", 
+            prompt if isinstance(prompt, str) else str(prompt)
+        )
+        
+        # Call agent chat
+        response = self._agent.chat(prompt, *args, **kwargs)
+        
+        # Persist assistant response
+        if response:
+            self._orchestrator.on_message(self._session_id, "assistant", response)
+        
+        return response
+    
+    def get_session(self) -> Optional[ConversationSession]:
+        """Get the current session."""
+        return self._orchestrator.get_session(self._session_id)
+    
+    def get_messages(self, limit: Optional[int] = None) -> List[ConversationMessage]:
+        """Get messages from the current session."""
+        return self._orchestrator.get_messages(self._session_id, limit=limit)
+    
+    def end_session(self):
+        """End the current session."""
+        self._orchestrator.on_agent_end(self._agent, self._session_id)
+    
+    @property
+    def session_id(self) -> str:
+        """Get the session ID."""
+        return self._session_id
+
+
+def create_persistent_session(
+    orchestrator: PersistenceOrchestrator,
+    session_id: Optional[str] = None,
+    user_id: Optional[str] = None,
+):
+    """
+    Create a context manager for persistent sessions.
+    
+    Example:
+        from praisonai.persistence import PersistenceOrchestrator
+        from praisonai.persistence.hooks import create_persistent_session
+        from praisonaiagents import Agent
+        
+        orchestrator = PersistenceOrchestrator(...)
+        
+        with create_persistent_session(orchestrator, session_id="my-session") as session:
+            agent = Agent(name="Assistant", role="Helper")
+            
+            # Messages are persisted within the session context
+            session.persist_message("user", "Hello!")
+            response = agent.chat("Hello!")
+            session.persist_message("assistant", response)
+    """
+    
+    class PersistentSession:
+        def __init__(self, orchestrator, session_id, user_id):
+            self.orchestrator = orchestrator
+            self.session_id = session_id or f"session-{uuid.uuid4().hex[:8]}"
+            self.user_id = user_id or "default"
+            self._started = False
+        
+        def __enter__(self):
+            self.start()
+            return self
+        
+        def __exit__(self, exc_type, exc_val, exc_tb):
+            self.end()
+            return False
+        
+        def start(self, agent=None):
+            """Start the session."""
+            if not self._started:
+                class MockAgent:
+                    name = "session-agent"
+                
+                self.orchestrator.on_agent_start(
+                    agent or MockAgent(),
+                    session_id=self.session_id,
+                    user_id=self.user_id
+                )
+                self._started = True
+        
+        def end(self, agent=None):
+            """End the session."""
+            if self._started:
+                class MockAgent:
+                    name = "session-agent"
+                
+                self.orchestrator.on_agent_end(
+                    agent or MockAgent(),
+                    self.session_id
+                )
+        
+        def persist_message(self, role: str, content: str):
+            """Persist a message to the session."""
+            self.orchestrator.on_message(self.session_id, role, content)
+        
+        def get_messages(self, limit: Optional[int] = None):
+            """Get messages from the session."""
+            return self.orchestrator.get_messages(self.session_id, limit=limit)
+        
+        def get_session(self):
+            """Get the session object."""
+            return self.orchestrator.get_session(self.session_id)
+    
+    return PersistentSession(orchestrator, session_id, user_id)

praisonai/persistence/knowledge/__init__.py

@@ -0,0 +1,26 @@
+"""
+KnowledgeStore implementations for vector embeddings and semantic search.
+
+Supported backends:
+- Qdrant (qdrant)
+- Pinecone (pinecone)
+- ChromaDB (chroma)
+- Weaviate (weaviate)
+- LanceDB (lancedb)
+- Milvus (milvus)
+- PGVector (pgvector)
+- Redis Vector (redis)
+- Cassandra Vector (cassandra)
+- ClickHouse Vector (clickhouse)
+"""
+
+__all__ = [
+    "KnowledgeStore",
+    "KnowledgeDocument",
+]
+
+def __getattr__(name: str):
+    if name in ("KnowledgeStore", "KnowledgeDocument"):
+        from .base import KnowledgeStore, KnowledgeDocument
+        return locals()[name]
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

praisonai/persistence/knowledge/base.py

@@ -0,0 +1,144 @@
+"""
+Base interfaces for KnowledgeStore.
+
+KnowledgeStore handles vector embeddings and semantic search for RAG.
+"""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional
+import time
+import uuid
+import hashlib
+
+
+@dataclass
+class KnowledgeDocument:
+    """A document with vector embedding for semantic search."""
+    id: str = field(default_factory=lambda: str(uuid.uuid4()))
+    content: str = ""
+    embedding: Optional[List[float]] = None
+    metadata: Optional[Dict[str, Any]] = None
+    content_hash: Optional[str] = None
+    created_at: float = field(default_factory=time.time)
+    
+    def __post_init__(self):
+        if self.content and not self.content_hash:
+            self.content_hash = hashlib.sha256(self.content.encode()).hexdigest()[:16]
+    
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "id": self.id,
+            "content": self.content,
+            "embedding": self.embedding,
+            "metadata": self.metadata,
+            "content_hash": self.content_hash,
+            "created_at": self.created_at,
+        }
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "KnowledgeDocument":
+        return cls(**{k: v for k, v in data.items() if k in cls.__dataclass_fields__})
+
+
+class KnowledgeStore(ABC):
+    """
+    Abstract base class for vector/knowledge persistence.
+    
+    Implementations handle document storage and semantic search for different backends:
+    - Qdrant, Pinecone, ChromaDB, Weaviate (dedicated vector DBs)
+    - LanceDB, Milvus (embedded/distributed vector DBs)
+    - PGVector, Redis, Cassandra, ClickHouse (vector extensions)
+    """
+    
+    @abstractmethod
+    def create_collection(
+        self, 
+        name: str, 
+        dimension: int,
+        distance: str = "cosine",
+        metadata: Optional[Dict[str, Any]] = None
+    ) -> None:
+        """Create a new collection/index for vectors."""
+        raise NotImplementedError
+    
+    @abstractmethod
+    def delete_collection(self, name: str) -> bool:
+        """Delete a collection and all its documents."""
+        raise NotImplementedError
+    
+    @abstractmethod
+    def collection_exists(self, name: str) -> bool:
+        """Check if a collection exists."""
+        raise NotImplementedError
+    
+    @abstractmethod
+    def list_collections(self) -> List[str]:
+        """List all collections."""
+        raise NotImplementedError
+    
+    @abstractmethod
+    def insert(
+        self, 
+        collection: str, 
+        documents: List[KnowledgeDocument]
+    ) -> List[str]:
+        """Insert documents into a collection. Returns list of IDs."""
+        raise NotImplementedError
+    
+    @abstractmethod
+    def upsert(
+        self, 
+        collection: str, 
+        documents: List[KnowledgeDocument]
+    ) -> List[str]:
+        """Insert or update documents. Returns list of IDs."""
+        raise NotImplementedError
+    
+    @abstractmethod
+    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 by embedding."""
+        raise NotImplementedError
+    
+    @abstractmethod
+    def get(
+        self, 
+        collection: str, 
+        ids: List[str]
+    ) -> List[KnowledgeDocument]:
+        """Get documents by IDs."""
+        raise NotImplementedError
+    
+    @abstractmethod
+    def delete(
+        self, 
+        collection: str, 
+        ids: Optional[List[str]] = None,
+        filters: Optional[Dict[str, Any]] = None
+    ) -> int:
+        """Delete documents by IDs or filters. Returns count deleted."""
+        raise NotImplementedError
+    
+    @abstractmethod
+    def count(self, collection: str) -> int:
+        """Count documents in a collection."""
+        raise NotImplementedError
+    
+    @abstractmethod
+    def close(self) -> None:
+        """Close the store and release resources."""
+        raise NotImplementedError
+    
+    def __enter__(self):
+        return self
+    
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+        return False

praisonai/persistence/knowledge/cassandra.py

@@ -0,0 +1,232 @@
+"""
+Cassandra Vector implementation of KnowledgeStore.
+
+Requires: cassandra-driver
+Install: pip install cassandra-driver
+"""
+
+import logging
+from typing import Any, Dict, List, Optional
+
+from .base import KnowledgeStore, KnowledgeDocument
+
+logger = logging.getLogger(__name__)
+
+
+class CassandraKnowledgeStore(KnowledgeStore):
+    """
+    Cassandra-based knowledge store using vector search (SAI).
+    
+    Requires Cassandra 5.0+ or DataStax Astra with vector search.
+    
+    Example:
+        store = CassandraKnowledgeStore(
+            hosts=["localhost"],
+            keyspace="praisonai"
+        )
+    """
+    
+    def __init__(
+        self,
+        hosts: Optional[List[str]] = None,
+        keyspace: str = "praisonai",
+        username: Optional[str] = None,
+        password: Optional[str] = None,
+        port: int = 9042,
+        secure_connect_bundle: Optional[str] = None,
+    ):
+        try:
+            from cassandra.cluster import Cluster
+            from cassandra.auth import PlainTextAuthProvider
+        except ImportError:
+            raise ImportError(
+                "cassandra-driver is required for Cassandra support. "
+                "Install with: pip install cassandra-driver"
+            )
+        
+        self._Cluster = Cluster
+        self.keyspace = keyspace
+        
+        auth = None
+        if username and password:
+            auth = PlainTextAuthProvider(username=username, password=password)
+        
+        if secure_connect_bundle:
+            # Astra DB connection
+            from cassandra.cluster import Cluster
+            cloud_config = {"secure_connect_bundle": secure_connect_bundle}
+            self._cluster = Cluster(cloud=cloud_config, auth_provider=auth)
+        else:
+            hosts = hosts or ["localhost"]
+            self._cluster = Cluster(hosts, port=port, auth_provider=auth)
+        
+        self._session = self._cluster.connect()
+        
+        # Create keyspace if not exists
+        self._session.execute(f"""
+            CREATE KEYSPACE IF NOT EXISTS {keyspace}
+            WITH replication = {{'class': 'SimpleStrategy', 'replication_factor': 1}}
+        """)
+        self._session.set_keyspace(keyspace)
+        logger.info(f"Connected to Cassandra keyspace: {keyspace}")
+    
+    def create_collection(
+        self,
+        name: str,
+        dimension: int,
+        distance: str = "cosine",
+        metadata: Optional[Dict[str, Any]] = None
+    ) -> None:
+        """Create a new table with vector column."""
+        similarity_map = {"cosine": "COSINE", "euclidean": "EUCLIDEAN", "dot": "DOT_PRODUCT"}
+        
+        self._session.execute(f"""
+            CREATE TABLE IF NOT EXISTS {name} (
+                id text PRIMARY KEY,
+                content text,
+                content_hash text,
+                created_at double,
+                embedding vector<float, {dimension}>
+            )
+        """)
+        
+        # Create SAI index for vector search
+        self._session.execute(f"""
+            CREATE CUSTOM INDEX IF NOT EXISTS {name}_embedding_idx ON {name} (embedding)
+            USING 'StorageAttachedIndex'
+            WITH OPTIONS = {{'similarity_function': '{similarity_map.get(distance, "COSINE")}'}}
+        """)
+        logger.info(f"Created Cassandra table: {name}")
+    
+    def delete_collection(self, name: str) -> bool:
+        """Delete a table."""
+        try:
+            self._session.execute(f"DROP TABLE IF EXISTS {name}")
+            return True
+        except Exception as e:
+            logger.warning(f"Failed to delete table {name}: {e}")
+            return False
+    
+    def collection_exists(self, name: str) -> bool:
+        """Check if a table exists."""
+        result = self._session.execute("""
+            SELECT table_name FROM system_schema.tables 
+            WHERE keyspace_name = %s AND table_name = %s
+        """, (self.keyspace, name))
+        return len(list(result)) > 0
+    
+    def list_collections(self) -> List[str]:
+        """List all tables."""
+        result = self._session.execute("""
+            SELECT table_name FROM system_schema.tables WHERE keyspace_name = %s
+        """, (self.keyspace,))
+        return [row.table_name for row in result]
+    
+    def insert(
+        self,
+        collection: str,
+        documents: List[KnowledgeDocument]
+    ) -> List[str]:
+        """Insert documents."""
+        ids = []
+        for doc in documents:
+            if doc.embedding is None:
+                raise ValueError(f"Document {doc.id} has no embedding")
+            
+            self._session.execute(f"""
+                INSERT INTO {collection} (id, content, content_hash, created_at, embedding)
+                VALUES (%s, %s, %s, %s, %s)
+            """, (doc.id, doc.content, doc.content_hash or "", doc.created_at, doc.embedding))
+            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)  # Cassandra INSERT is upsert
+    
+    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 using ANN."""
+        result = self._session.execute(f"""
+            SELECT id, content, content_hash, created_at, similarity_cosine(embedding, %s) as score
+            FROM {collection}
+            ORDER BY embedding ANN OF %s
+            LIMIT %s
+        """, (query_embedding, query_embedding, limit))
+        
+        documents = []
+        for row in result:
+            if score_threshold and row.score < score_threshold:
+                continue
+            
+            doc = KnowledgeDocument(
+                id=row.id,
+                content=row.content,
+                embedding=None,
+                metadata={},
+                content_hash=row.content_hash,
+                created_at=row.created_at,
+            )
+            documents.append(doc)
+        
+        return documents
+    
+    def get(
+        self,
+        collection: str,
+        ids: List[str]
+    ) -> List[KnowledgeDocument]:
+        """Get documents by IDs."""
+        documents = []
+        for doc_id in ids:
+            result = self._session.execute(f"""
+                SELECT id, content, content_hash, created_at, embedding
+                FROM {collection} WHERE id = %s
+            """, (doc_id,))
+            
+            for row in result:
+                doc = KnowledgeDocument(
+                    id=row.id,
+                    content=row.content,
+                    embedding=list(row.embedding) if row.embedding else None,
+                    metadata={},
+                    content_hash=row.content_hash,
+                    created_at=row.created_at,
+                )
+                documents.append(doc)
+        
+        return documents
+    
+    def delete(
+        self,
+        collection: str,
+        ids: Optional[List[str]] = None,
+        filters: Optional[Dict[str, Any]] = None
+    ) -> int:
+        """Delete documents."""
+        if ids:
+            for doc_id in ids:
+                self._session.execute(f"DELETE FROM {collection} WHERE id = %s", (doc_id,))
+            return len(ids)
+        return 0
+    
+    def count(self, collection: str) -> int:
+        """Count documents."""
+        result = self._session.execute(f"SELECT COUNT(*) FROM {collection}")
+        return result.one()[0]
+    
+    def close(self) -> None:
+        """Close the store."""
+        if self._cluster:
+            self._cluster.shutdown()