PraisonAI 3.0.0__py3-none-any.whl

praisonai/persistence/conversation/surrealdb.py

@@ -0,0 +1,287 @@
+"""
+SurrealDB implementation of ConversationStore.
+
+Requires: surrealdb
+Install: pip install surrealdb
+"""
+
+import json
+import logging
+from typing import List, Optional
+
+from .base import ConversationStore, ConversationSession, ConversationMessage
+
+logger = logging.getLogger(__name__)
+
+
+class SurrealDBConversationStore(ConversationStore):
+    """
+    SurrealDB-based conversation store.
+    
+    SurrealDB is a multi-model database supporting SQL-like queries.
+    
+    Example:
+        store = SurrealDBConversationStore(
+            url="ws://localhost:8000/rpc",
+            namespace="praisonai",
+            database="conversations"
+        )
+    """
+    
+    SCHEMA_VERSION = "1.0.0"
+    
+    def __init__(
+        self,
+        url: str = "ws://localhost:8000/rpc",
+        namespace: str = "praisonai",
+        database: str = "conversations",
+        username: Optional[str] = None,
+        password: Optional[str] = None,
+        table_prefix: str = "praison_",
+        auto_create_tables: bool = True,
+    ):
+        try:
+            from surrealdb import Surreal
+        except ImportError:
+            raise ImportError(
+                "surrealdb is required for SurrealDB support. "
+                "Install with: pip install surrealdb"
+            )
+        
+        self._Surreal = Surreal
+        self.url = url
+        self.namespace = namespace
+        self.database = database
+        self.username = username
+        self.password = password
+        self.table_prefix = table_prefix
+        self.sessions_table = f"{table_prefix}sessions"
+        self.messages_table = f"{table_prefix}messages"
+        self._client = None
+        
+        # Note: SurrealDB Python client is async, we'll use sync wrapper
+        self._init_sync()
+    
+    def _init_sync(self):
+        """Initialize synchronously using asyncio."""
+        import asyncio
+        
+        async def _init():
+            self._client = self._Surreal(self.url)
+            await self._client.connect()
+            if self.username and self.password:
+                await self._client.signin({"user": self.username, "pass": self.password})
+            await self._client.use(self.namespace, self.database)
+        
+        try:
+            loop = asyncio.get_event_loop()
+        except RuntimeError:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+        
+        loop.run_until_complete(_init())
+    
+    def _run_sync(self, coro):
+        """Run async coroutine synchronously."""
+        import asyncio
+        try:
+            loop = asyncio.get_event_loop()
+        except RuntimeError:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+        return loop.run_until_complete(coro)
+    
+    def create_session(self, session: ConversationSession) -> ConversationSession:
+        data = {
+            "session_id": session.session_id,
+            "user_id": session.user_id,
+            "agent_id": session.agent_id,
+            "name": session.name,
+            "state": session.state,
+            "metadata": session.metadata,
+            "created_at": session.created_at,
+            "updated_at": session.updated_at,
+        }
+        self._run_sync(self._client.create(self.sessions_table, data))
+        return session
+    
+    def get_session(self, session_id: str) -> Optional[ConversationSession]:
+        result = self._run_sync(
+            self._client.query(
+                f"SELECT * FROM {self.sessions_table} WHERE session_id = $sid",
+                {"sid": session_id}
+            )
+        )
+        if not result or not result[0].get("result"):
+            return None
+        row = result[0]["result"][0]
+        return ConversationSession(
+            session_id=row["session_id"],
+            user_id=row.get("user_id"),
+            agent_id=row.get("agent_id"),
+            name=row.get("name"),
+            state=row.get("state"),
+            metadata=row.get("metadata"),
+            created_at=row.get("created_at"),
+            updated_at=row.get("updated_at"),
+        )
+    
+    def update_session(self, session: ConversationSession) -> ConversationSession:
+        self._run_sync(
+            self._client.query(
+                f"""UPDATE {self.sessions_table} SET 
+                    user_id = $user_id, agent_id = $agent_id, name = $name,
+                    state = $state, metadata = $metadata, updated_at = $updated_at
+                    WHERE session_id = $session_id""",
+                {
+                    "session_id": session.session_id,
+                    "user_id": session.user_id,
+                    "agent_id": session.agent_id,
+                    "name": session.name,
+                    "state": session.state,
+                    "metadata": session.metadata,
+                    "updated_at": session.updated_at,
+                }
+            )
+        )
+        return session
+    
+    def delete_session(self, session_id: str) -> bool:
+        self._run_sync(
+            self._client.query(
+                f"DELETE FROM {self.messages_table} WHERE session_id = $sid",
+                {"sid": session_id}
+            )
+        )
+        result = self._run_sync(
+            self._client.query(
+                f"DELETE FROM {self.sessions_table} WHERE session_id = $sid",
+                {"sid": session_id}
+            )
+        )
+        return True
+    
+    def list_sessions(
+        self,
+        user_id: Optional[str] = None,
+        agent_id: Optional[str] = None,
+        limit: int = 100,
+        offset: int = 0
+    ) -> List[ConversationSession]:
+        conditions = []
+        params = {"limit": limit, "offset": offset}
+        
+        if user_id:
+            conditions.append("user_id = $user_id")
+            params["user_id"] = user_id
+        if agent_id:
+            conditions.append("agent_id = $agent_id")
+            params["agent_id"] = agent_id
+        
+        where = "WHERE " + " AND ".join(conditions) if conditions else ""
+        
+        result = self._run_sync(
+            self._client.query(
+                f"SELECT * FROM {self.sessions_table} {where} ORDER BY updated_at DESC LIMIT $limit START $offset",
+                params
+            )
+        )
+        
+        if not result or not result[0].get("result"):
+            return []
+        
+        return [
+            ConversationSession(
+                session_id=row["session_id"],
+                user_id=row.get("user_id"),
+                agent_id=row.get("agent_id"),
+                name=row.get("name"),
+                state=row.get("state"),
+                metadata=row.get("metadata"),
+                created_at=row.get("created_at"),
+                updated_at=row.get("updated_at"),
+            )
+            for row in result[0]["result"]
+        ]
+    
+    def add_message(self, session_id: str, message: ConversationMessage) -> ConversationMessage:
+        message.session_id = session_id
+        data = {
+            "id": message.id,
+            "session_id": session_id,
+            "role": message.role,
+            "content": message.content,
+            "tool_calls": message.tool_calls,
+            "tool_call_id": message.tool_call_id,
+            "metadata": message.metadata,
+            "created_at": message.created_at,
+        }
+        self._run_sync(self._client.create(self.messages_table, data))
+        return message
+    
+    def get_messages(
+        self,
+        session_id: str,
+        limit: Optional[int] = None,
+        before: Optional[float] = None,
+        after: Optional[float] = None
+    ) -> List[ConversationMessage]:
+        conditions = ["session_id = $session_id"]
+        params = {"session_id": session_id}
+        
+        if before:
+            conditions.append("created_at < $before")
+            params["before"] = before
+        if after:
+            conditions.append("created_at > $after")
+            params["after"] = after
+        
+        where = "WHERE " + " AND ".join(conditions)
+        limit_clause = f"LIMIT {limit}" if limit else ""
+        
+        result = self._run_sync(
+            self._client.query(
+                f"SELECT * FROM {self.messages_table} {where} ORDER BY created_at ASC {limit_clause}",
+                params
+            )
+        )
+        
+        if not result or not result[0].get("result"):
+            return []
+        
+        return [
+            ConversationMessage(
+                id=row["id"],
+                session_id=row["session_id"],
+                role=row["role"],
+                content=row.get("content", ""),
+                tool_calls=row.get("tool_calls"),
+                tool_call_id=row.get("tool_call_id"),
+                metadata=row.get("metadata"),
+                created_at=row.get("created_at"),
+            )
+            for row in result[0]["result"]
+        ]
+    
+    def delete_messages(self, session_id: str, message_ids: Optional[List[str]] = None) -> int:
+        if message_ids:
+            self._run_sync(
+                self._client.query(
+                    f"DELETE FROM {self.messages_table} WHERE session_id = $sid AND id IN $ids",
+                    {"sid": session_id, "ids": message_ids}
+                )
+            )
+            return len(message_ids)
+        else:
+            self._run_sync(
+                self._client.query(
+                    f"DELETE FROM {self.messages_table} WHERE session_id = $sid",
+                    {"sid": session_id}
+                )
+            )
+            return -1
+    
+    def close(self) -> None:
+        if self._client:
+            self._run_sync(self._client.close())
+            self._client = None

praisonai/persistence/factory.py

@@ -0,0 +1,301 @@
+"""
+Factory functions for creating store instances.
+
+Provides lazy loading of backend implementations to avoid importing
+unused dependencies.
+"""
+
+import logging
+from typing import Any, Dict, Optional
+
+from .config import PersistenceConfig
+
+logger = logging.getLogger(__name__)
+
+
+def create_conversation_store(
+    backend: str,
+    url: Optional[str] = None,
+    **options: Any
+):
+    """
+    Create a ConversationStore instance.
+    
+    Args:
+        backend: Backend type (postgres, mysql, sqlite, singlestore, supabase, surrealdb)
+        url: Connection URL
+        **options: Backend-specific options
+    
+    Returns:
+        ConversationStore instance
+    
+    Example:
+        store = create_conversation_store(
+            "postgres",
+            url="postgresql://localhost:5432/praisonai"
+        )
+    """
+    backend = backend.lower()
+    
+    if backend == "postgres":
+        from .conversation.postgres import PostgresConversationStore
+        return PostgresConversationStore(url=url, **options)
+    
+    elif backend == "mysql":
+        from .conversation.mysql import MySQLConversationStore
+        return MySQLConversationStore(url=url, **options)
+    
+    elif backend == "sqlite":
+        from .conversation.sqlite import SQLiteConversationStore
+        path = url or options.pop("path", None)
+        return SQLiteConversationStore(path=path, **options)
+    
+    elif backend == "singlestore":
+        from .conversation.singlestore import SingleStoreConversationStore
+        return SingleStoreConversationStore(url=url, **options)
+    
+    elif backend == "supabase":
+        from .conversation.supabase import SupabaseConversationStore
+        return SupabaseConversationStore(url=url, **options)
+    
+    elif backend == "surrealdb":
+        from .conversation.surrealdb import SurrealDBConversationStore
+        return SurrealDBConversationStore(url=url, **options)
+    
+    elif backend == "json":
+        from .conversation.json_store import JSONConversationStore
+        path = url or options.pop("path", None) or "./praisonai_conversations"
+        return JSONConversationStore(path=path, **options)
+    
+    elif backend in ("async_postgres", "asyncpg", "postgres_async"):
+        from .conversation.async_postgres import AsyncPostgresConversationStore
+        return AsyncPostgresConversationStore(url=url, **options)
+    
+    elif backend in ("async_sqlite", "aiosqlite", "sqlite_async"):
+        from .conversation.async_sqlite import AsyncSQLiteConversationStore
+        path = url or options.pop("path", None)
+        return AsyncSQLiteConversationStore(path=path, **options)
+    
+    elif backend in ("async_mysql", "aiomysql", "mysql_async"):
+        from .conversation.async_mysql import AsyncMySQLConversationStore
+        return AsyncMySQLConversationStore(url=url, **options)
+    
+    else:
+        raise ValueError(
+            f"Unknown conversation store backend: {backend}. "
+            f"Supported: postgres, mysql, sqlite, json, singlestore, supabase, surrealdb, "
+            f"async_postgres, async_sqlite, async_mysql"
+        )
+
+
+def create_knowledge_store(
+    backend: str,
+    url: Optional[str] = None,
+    **options: Any
+):
+    """
+    Create a KnowledgeStore instance.
+    
+    Args:
+        backend: Backend type (qdrant, pinecone, chroma, weaviate, lancedb, milvus, pgvector, redis, cassandra, clickhouse)
+        url: Connection URL
+        **options: Backend-specific options
+    
+    Returns:
+        KnowledgeStore instance
+    
+    Example:
+        store = create_knowledge_store(
+            "qdrant",
+            url="http://localhost:6333"
+        )
+    """
+    backend = backend.lower()
+    
+    if backend == "qdrant":
+        from .knowledge.qdrant import QdrantKnowledgeStore
+        return QdrantKnowledgeStore(url=url, **options)
+    
+    elif backend == "pinecone":
+        from .knowledge.pinecone import PineconeKnowledgeStore
+        return PineconeKnowledgeStore(**options)
+    
+    elif backend == "chroma":
+        from .knowledge.chroma import ChromaKnowledgeStore
+        path = url or options.pop("path", None)
+        return ChromaKnowledgeStore(path=path, **options)
+    
+    elif backend == "weaviate":
+        from .knowledge.weaviate import WeaviateKnowledgeStore
+        return WeaviateKnowledgeStore(url=url, **options)
+    
+    elif backend == "lancedb":
+        from .knowledge.lancedb import LanceDBKnowledgeStore
+        path = url or options.pop("path", None)
+        return LanceDBKnowledgeStore(path=path, **options)
+    
+    elif backend == "milvus":
+        from .knowledge.milvus import MilvusKnowledgeStore
+        return MilvusKnowledgeStore(url=url, **options)
+    
+    elif backend == "pgvector":
+        from .knowledge.pgvector import PGVectorKnowledgeStore
+        return PGVectorKnowledgeStore(url=url, **options)
+    
+    elif backend == "redis":
+        from .knowledge.redis_vector import RedisVectorKnowledgeStore
+        return RedisVectorKnowledgeStore(url=url, **options)
+    
+    elif backend == "cassandra":
+        from .knowledge.cassandra import CassandraKnowledgeStore
+        return CassandraKnowledgeStore(**options)
+    
+    elif backend == "clickhouse":
+        from .knowledge.clickhouse import ClickHouseKnowledgeStore
+        return ClickHouseKnowledgeStore(**options)
+    
+    elif backend == "couchbase":
+        from .knowledge.couchbase import CouchbaseKnowledgeStore
+        return CouchbaseKnowledgeStore(**options)
+    
+    elif backend in ("mongodb_vector", "mongodb_atlas", "mongo_vector"):
+        from .knowledge.mongodb_vector import MongoDBVectorKnowledgeStore
+        return MongoDBVectorKnowledgeStore(url=url, **options)
+    
+    elif backend in ("singlestore_vector", "singlestore_v"):
+        from .knowledge.singlestore_vector import SingleStoreVectorKnowledgeStore
+        return SingleStoreVectorKnowledgeStore(url=url, **options)
+    
+    elif backend in ("surrealdb_vector", "surrealdb_v"):
+        from .knowledge.surrealdb_vector import SurrealDBVectorKnowledgeStore
+        return SurrealDBVectorKnowledgeStore(url=url, **options)
+    
+    elif backend in ("upstash_vector", "upstash_v"):
+        from .knowledge.upstash_vector import UpstashVectorKnowledgeStore
+        return UpstashVectorKnowledgeStore(url=url, **options)
+    
+    elif backend == "lightrag":
+        from .knowledge.lightrag_adapter import LightRAGKnowledgeStore
+        return LightRAGKnowledgeStore(**options)
+    
+    elif backend in ("langchain", "langchain_adapter"):
+        from .knowledge.langchain_adapter import LangChainKnowledgeStore
+        return LangChainKnowledgeStore(**options)
+    
+    elif backend in ("llamaindex", "llama_index", "llamaindex_adapter"):
+        from .knowledge.llamaindex_adapter import LlamaIndexKnowledgeStore
+        return LlamaIndexKnowledgeStore(**options)
+    
+    elif backend in ("cosmosdb", "cosmos", "azure_cosmos", "cosmosdb_vector"):
+        from .knowledge.cosmosdb_vector import CosmosDBVectorKnowledgeStore
+        return CosmosDBVectorKnowledgeStore(**options)
+    
+    else:
+        raise ValueError(
+            f"Unknown knowledge store backend: {backend}. "
+            f"Supported: qdrant, pinecone, chroma, weaviate, lancedb, milvus, pgvector, redis, cassandra, clickhouse, "
+            f"couchbase, mongodb_vector, singlestore_vector, surrealdb_vector, upstash_vector, lightrag, langchain, llamaindex, cosmosdb"
+        )
+
+
+def create_state_store(
+    backend: str,
+    url: Optional[str] = None,
+    **options: Any
+):
+    """
+    Create a StateStore instance.
+    
+    Args:
+        backend: Backend type (redis, dynamodb, firestore, mongodb, upstash, memory)
+        url: Connection URL
+        **options: Backend-specific options
+    
+    Returns:
+        StateStore instance
+    
+    Example:
+        store = create_state_store(
+            "redis",
+            url="redis://localhost:6379"
+        )
+    """
+    backend = backend.lower()
+    
+    if backend == "redis":
+        from .state.redis import RedisStateStore
+        return RedisStateStore(url=url, **options)
+    
+    elif backend == "dynamodb":
+        from .state.dynamodb import DynamoDBStateStore
+        return DynamoDBStateStore(**options)
+    
+    elif backend == "firestore":
+        from .state.firestore import FirestoreStateStore
+        return FirestoreStateStore(**options)
+    
+    elif backend == "mongodb":
+        from .state.mongodb import MongoDBStateStore
+        return MongoDBStateStore(url=url, **options)
+    
+    elif backend == "upstash":
+        from .state.upstash import UpstashStateStore
+        return UpstashStateStore(url=url, **options)
+    
+    elif backend == "memory":
+        from .state.memory import MemoryStateStore
+        return MemoryStateStore(**options)
+    
+    elif backend == "gcs":
+        from .state.gcs import GCSStateStore
+        bucket = options.pop("bucket_name", None) or options.pop("bucket", None)
+        if not bucket:
+            raise ValueError("GCS state store requires 'bucket_name' option")
+        return GCSStateStore(bucket_name=bucket, **options)
+    
+    elif backend in ("async_mongodb", "motor", "mongodb_async"):
+        from .state.async_mongodb import AsyncMongoDBStateStore
+        return AsyncMongoDBStateStore(url=url, **options)
+    
+    else:
+        raise ValueError(
+            f"Unknown state store backend: {backend}. "
+            f"Supported: redis, dynamodb, firestore, mongodb, upstash, memory, gcs, async_mongodb"
+        )
+
+
+def create_stores_from_config(config: PersistenceConfig) -> Dict[str, Any]:
+    """
+    Create all configured stores from a PersistenceConfig.
+    
+    Returns:
+        Dict with keys: conversation, knowledge, state (values may be None)
+    """
+    stores = {
+        "conversation": None,
+        "knowledge": None,
+        "state": None,
+    }
+    
+    if config.conversation_store:
+        stores["conversation"] = create_conversation_store(
+            config.conversation_store,
+            url=config.conversation_url,
+            **config.conversation_options
+        )
+    
+    if config.knowledge_store:
+        stores["knowledge"] = create_knowledge_store(
+            config.knowledge_store,
+            url=config.knowledge_url,
+            **config.knowledge_options
+        )
+    
+    if config.state_store:
+        stores["state"] = create_state_store(
+            config.state_store,
+            url=config.state_url,
+            **config.state_options
+        )
+    
+    return stores

praisonai/persistence/hooks/__init__.py

@@ -0,0 +1,18 @@
+"""
+Agent hooks for automatic persistence integration.
+
+Provides wrapper functions to add persistence capabilities to PraisonAI agents
+without modifying the core SDK.
+"""
+
+from .agent_hooks import (
+    wrap_agent_with_persistence,
+    PersistentAgent,
+    create_persistent_session,
+)
+
+__all__ = [
+    "wrap_agent_with_persistence",
+    "PersistentAgent",
+    "create_persistent_session",
+]