PraisonAI 3.0.0__py3-none-any.whl

praisonai/persistence/knowledge/redis_vector.py

@@ -0,0 +1,291 @@
+"""
+Redis Vector implementation of KnowledgeStore.
+
+Requires: redis
+Install: pip install redis
+"""
+
+import json
+import logging
+from typing import Any, Dict, List, Optional
+
+from .base import KnowledgeStore, KnowledgeDocument
+
+logger = logging.getLogger(__name__)
+
+
+class RedisVectorKnowledgeStore(KnowledgeStore):
+    """
+    Redis-based knowledge store using RediSearch for vector search.
+    
+    Requires Redis Stack or Redis with RediSearch module.
+    
+    Example:
+        store = RedisVectorKnowledgeStore(
+            url="redis://localhost:6379"
+        )
+    """
+    
+    def __init__(
+        self,
+        url: Optional[str] = None,
+        host: str = "localhost",
+        port: int = 6379,
+        password: Optional[str] = None,
+        prefix: str = "praison:vec:",
+    ):
+        try:
+            import redis
+            from redis.commands.search.field import VectorField, TextField, NumericField
+            from redis.commands.search.indexDefinition import IndexDefinition, IndexType
+            from redis.commands.search.query import Query
+        except ImportError:
+            raise ImportError(
+                "redis is required for Redis Vector support. "
+                "Install with: pip install redis"
+            )
+        
+        self._redis = redis
+        self._VectorField = VectorField
+        self._TextField = TextField
+        self._NumericField = NumericField
+        self._IndexDefinition = IndexDefinition
+        self._IndexType = IndexType
+        self._Query = Query
+        
+        self.prefix = prefix
+        
+        if url:
+            self._client = redis.from_url(url, decode_responses=False)
+        else:
+            self._client = redis.Redis(
+                host=host, port=port, password=password, decode_responses=False
+            )
+        
+        self._client.ping()
+        logger.info(f"Connected to Redis at {url or f'{host}:{port}'}")
+    
+    def _index_name(self, collection: str) -> str:
+        return f"{self.prefix}{collection}:idx"
+    
+    def _doc_key(self, collection: str, doc_id: str) -> str:
+        return f"{self.prefix}{collection}:{doc_id}"
+    
+    def create_collection(
+        self,
+        name: str,
+        dimension: int,
+        distance: str = "cosine",
+        metadata: Optional[Dict[str, Any]] = None
+    ) -> None:
+        """Create a new index."""
+        index_name = self._index_name(name)
+        
+        distance_map = {"cosine": "COSINE", "euclidean": "L2", "dot": "IP"}
+        
+        schema = (
+            self._TextField("content"),
+            self._TextField("content_hash"),
+            self._NumericField("created_at"),
+            self._VectorField(
+                "embedding",
+                "HNSW",
+                {
+                    "TYPE": "FLOAT32",
+                    "DIM": dimension,
+                    "DISTANCE_METRIC": distance_map.get(distance, "COSINE"),
+                }
+            ),
+        )
+        
+        definition = self._IndexDefinition(
+            prefix=[f"{self.prefix}{name}:"],
+            index_type=self._IndexType.HASH,
+        )
+        
+        try:
+            self._client.ft(index_name).create_index(schema, definition=definition)
+            logger.info(f"Created Redis index: {index_name}")
+        except Exception as e:
+            if "Index already exists" not in str(e):
+                raise
+    
+    def delete_collection(self, name: str) -> bool:
+        """Delete an index and all documents."""
+        index_name = self._index_name(name)
+        try:
+            self._client.ft(index_name).dropindex(delete_documents=True)
+            return True
+        except Exception as e:
+            logger.warning(f"Failed to delete index {index_name}: {e}")
+            return False
+    
+    def collection_exists(self, name: str) -> bool:
+        """Check if an index exists."""
+        index_name = self._index_name(name)
+        try:
+            self._client.ft(index_name).info()
+            return True
+        except Exception:
+            return False
+    
+    def list_collections(self) -> List[str]:
+        """List all collections."""
+        try:
+            indexes = self._client.execute_command("FT._LIST")
+            prefix = f"{self.prefix}"
+            return [
+                idx.decode().replace(prefix, "").replace(":idx", "")
+                for idx in indexes
+                if idx.decode().startswith(prefix)
+            ]
+        except Exception:
+            return []
+    
+    def insert(
+        self,
+        collection: str,
+        documents: List[KnowledgeDocument]
+    ) -> List[str]:
+        """Insert documents."""
+        import numpy as np
+        
+        ids = []
+        pipe = self._client.pipeline()
+        
+        for doc in documents:
+            if doc.embedding is None:
+                raise ValueError(f"Document {doc.id} has no embedding")
+            
+            key = self._doc_key(collection, doc.id)
+            embedding_bytes = np.array(doc.embedding, dtype=np.float32).tobytes()
+            
+            pipe.hset(key, mapping={
+                "content": doc.content,
+                "content_hash": doc.content_hash or "",
+                "created_at": doc.created_at,
+                "embedding": embedding_bytes,
+            })
+            ids.append(doc.id)
+        
+        pipe.execute()
+        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."""
+        import numpy as np
+        
+        index_name = self._index_name(collection)
+        query_bytes = np.array(query_embedding, dtype=np.float32).tobytes()
+        
+        query_str = f"*=>[KNN {limit} @embedding $vec AS score]"
+        
+        query = (
+            self._Query(query_str)
+            .return_fields("content", "content_hash", "created_at", "score")
+            .sort_by("score")
+            .dialect(2)
+        )
+        
+        results = self._client.ft(index_name).search(
+            query, query_params={"vec": query_bytes}
+        )
+        
+        documents = []
+        for doc in results.docs:
+            score = float(doc.score) if hasattr(doc, "score") else 0
+            # Redis returns distance, convert to similarity for cosine
+            similarity = 1 - score
+            
+            if score_threshold and similarity < score_threshold:
+                continue
+            
+            doc_id = doc.id.decode() if isinstance(doc.id, bytes) else doc.id
+            doc_id = doc_id.split(":")[-1]  # Remove prefix
+            
+            knowledge_doc = KnowledgeDocument(
+                id=doc_id,
+                content=doc.content.decode() if isinstance(doc.content, bytes) else doc.content,
+                embedding=None,
+                metadata={},
+                content_hash=doc.content_hash.decode() if isinstance(doc.content_hash, bytes) else doc.content_hash,
+                created_at=float(doc.created_at) if doc.created_at else 0,
+            )
+            documents.append(knowledge_doc)
+        
+        return documents
+    
+    def get(
+        self,
+        collection: str,
+        ids: List[str]
+    ) -> List[KnowledgeDocument]:
+        """Get documents by IDs."""
+        import numpy as np
+        
+        documents = []
+        for doc_id in ids:
+            key = self._doc_key(collection, doc_id)
+            data = self._client.hgetall(key)
+            
+            if data:
+                embedding = None
+                if b"embedding" in data:
+                    embedding = np.frombuffer(data[b"embedding"], dtype=np.float32).tolist()
+                
+                doc = KnowledgeDocument(
+                    id=doc_id,
+                    content=data.get(b"content", b"").decode(),
+                    embedding=embedding,
+                    metadata={},
+                    content_hash=data.get(b"content_hash", b"").decode(),
+                    created_at=float(data.get(b"created_at", 0)),
+                )
+                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:
+            pipe = self._client.pipeline()
+            for doc_id in ids:
+                key = self._doc_key(collection, doc_id)
+                pipe.delete(key)
+            pipe.execute()
+            return len(ids)
+        return 0
+    
+    def count(self, collection: str) -> int:
+        """Count documents."""
+        index_name = self._index_name(collection)
+        try:
+            info = self._client.ft(index_name).info()
+            return int(info.get("num_docs", 0))
+        except Exception:
+            return 0
+    
+    def close(self) -> None:
+        """Close the store."""
+        if self._client:
+            self._client.close()

praisonai/persistence/knowledge/singlestore_vector.py

@@ -0,0 +1,299 @@
+"""
+SingleStore Vector implementation of KnowledgeStore.
+
+Requires: singlestoredb
+Install: pip install singlestoredb
+"""
+
+import logging
+from typing import Any, Dict, List, Optional
+
+from .base import KnowledgeStore, KnowledgeDocument
+
+logger = logging.getLogger(__name__)
+
+
+class SingleStoreVectorKnowledgeStore(KnowledgeStore):
+    """
+    SingleStore vector store for knowledge/RAG.
+    
+    Uses SingleStore's native vector capabilities.
+    
+    Example:
+        store = SingleStoreVectorKnowledgeStore(
+            url="singlestoredb://user:pass@host:3306/db"
+        )
+    """
+    
+    def __init__(
+        self,
+        url: Optional[str] = None,
+        host: str = "localhost",
+        port: int = 3306,
+        database: str = "praisonai",
+        user: str = "root",
+        password: str = "",
+        table_prefix: str = "praisonai_",
+        embedding_dim: int = 1536,
+    ):
+        """
+        Initialize SingleStore Vector store.
+        
+        Args:
+            url: SingleStore connection URL
+            host: Database host
+            port: Database port
+            database: Database name
+            user: Database user
+            password: Database password
+            table_prefix: Prefix for table names
+            embedding_dim: Embedding dimension
+        """
+        self.url = url
+        self.host = host
+        self.port = port
+        self.database = database
+        self.user = user
+        self.password = password
+        self.table_prefix = table_prefix
+        self.embedding_dim = embedding_dim
+        
+        self._conn = None
+        self._initialized = False
+    
+    def _init_client(self):
+        """Initialize SingleStore client lazily."""
+        if self._initialized:
+            return
+        
+        try:
+            import singlestoredb as s2
+        except ImportError:
+            raise ImportError(
+                "singlestoredb is required for SingleStore Vector support. "
+                "Install with: pip install singlestoredb"
+            )
+        
+        if self.url:
+            self._conn = s2.connect(self.url)
+        else:
+            self._conn = s2.connect(
+                host=self.host,
+                port=self.port,
+                database=self.database,
+                user=self.user,
+                password=self.password
+            )
+        
+        self._initialized = True
+    
+    def create_collection(
+        self,
+        name: str,
+        dimension: int = 1536,
+        distance: str = "cosine",
+        metadata: Optional[Dict[str, Any]] = None
+    ) -> None:
+        """Create a vector table."""
+        self._init_client()
+        
+        table_name = f"{self.table_prefix}{name}"
+        
+        with self._conn.cursor() as cur:
+            cur.execute(f"""
+                CREATE TABLE IF NOT EXISTS {table_name} (
+                    id VARCHAR(255) PRIMARY KEY,
+                    content TEXT,
+                    embedding VECTOR({dimension}),
+                    metadata JSON,
+                    content_hash VARCHAR(64),
+                    created_at DOUBLE
+                )
+            """)
+    
+    def delete_collection(self, name: str) -> bool:
+        """Delete a vector table."""
+        self._init_client()
+        
+        table_name = f"{self.table_prefix}{name}"
+        try:
+            with self._conn.cursor() as cur:
+                cur.execute(f"DROP TABLE IF EXISTS {table_name}")
+            return True
+        except Exception as e:
+            logger.error(f"Failed to delete collection {name}: {e}")
+            return False
+    
+    def collection_exists(self, name: str) -> bool:
+        """Check if table exists."""
+        self._init_client()
+        
+        table_name = f"{self.table_prefix}{name}"
+        with self._conn.cursor() as cur:
+            cur.execute(f"""
+                SELECT COUNT(*) FROM information_schema.tables 
+                WHERE table_schema = %s AND table_name = %s
+            """, (self.database, table_name))
+            return cur.fetchone()[0] > 0
+    
+    def list_collections(self) -> List[str]:
+        """List all vector tables."""
+        self._init_client()
+        
+        with self._conn.cursor() as cur:
+            cur.execute(f"""
+                SELECT table_name FROM information_schema.tables 
+                WHERE table_schema = %s AND table_name LIKE %s
+            """, (self.database, f"{self.table_prefix}%"))
+            return [row[0][len(self.table_prefix):] for row in cur.fetchall()]
+    
+    def insert(
+        self,
+        collection: str,
+        documents: List[KnowledgeDocument]
+    ) -> List[str]:
+        """Insert documents."""
+        self._init_client()
+        
+        table_name = f"{self.table_prefix}{collection}"
+        ids = []
+        
+        with self._conn.cursor() as cur:
+            for doc in documents:
+                embedding_str = str(doc.embedding) if doc.embedding else None
+                cur.execute(f"""
+                    INSERT INTO {table_name} (id, content, embedding, metadata, content_hash, created_at)
+                    VALUES (%s, %s, %s, %s, %s, %s)
+                """, (doc.id, doc.content, embedding_str,
+                      str(doc.metadata) if doc.metadata else None,
+                      doc.content_hash, doc.created_at))
+                ids.append(doc.id)
+        
+        return ids
+    
+    def upsert(
+        self,
+        collection: str,
+        documents: List[KnowledgeDocument]
+    ) -> List[str]:
+        """Upsert documents."""
+        self._init_client()
+        
+        table_name = f"{self.table_prefix}{collection}"
+        ids = []
+        
+        with self._conn.cursor() as cur:
+            for doc in documents:
+                embedding_str = str(doc.embedding) if doc.embedding else None
+                cur.execute(f"""
+                    REPLACE INTO {table_name} (id, content, embedding, metadata, content_hash, created_at)
+                    VALUES (%s, %s, %s, %s, %s, %s)
+                """, (doc.id, doc.content, embedding_str,
+                      str(doc.metadata) if doc.metadata else None,
+                      doc.content_hash, doc.created_at))
+                ids.append(doc.id)
+        
+        return ids
+    
+    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 vector search."""
+        self._init_client()
+        
+        table_name = f"{self.table_prefix}{collection}"
+        embedding_str = str(query_embedding)
+        
+        with self._conn.cursor() as cur:
+            cur.execute(f"""
+                SELECT id, content, metadata, content_hash, created_at,
+                       DOT_PRODUCT(embedding, %s) as score
+                FROM {table_name}
+                ORDER BY score DESC
+                LIMIT %s
+            """, (embedding_str, limit))
+            
+            documents = []
+            for row in cur.fetchall():
+                score = row[5] if len(row) > 5 else 0
+                if score_threshold and score < score_threshold:
+                    continue
+                
+                documents.append(KnowledgeDocument(
+                    id=row[0],
+                    content=row[1],
+                    metadata=eval(row[2]) if row[2] else None,
+                    content_hash=row[3],
+                    created_at=row[4]
+                ))
+            
+            return documents
+    
+    def get(
+        self,
+        collection: str,
+        ids: List[str]
+    ) -> List[KnowledgeDocument]:
+        """Get documents by IDs."""
+        self._init_client()
+        
+        table_name = f"{self.table_prefix}{collection}"
+        placeholders = ','.join(['%s'] * len(ids))
+        
+        with self._conn.cursor() as cur:
+            cur.execute(f"""
+                SELECT id, content, metadata, content_hash, created_at
+                FROM {table_name}
+                WHERE id IN ({placeholders})
+            """, ids)
+            
+            return [
+                KnowledgeDocument(
+                    id=row[0],
+                    content=row[1],
+                    metadata=eval(row[2]) if row[2] else None,
+                    content_hash=row[3],
+                    created_at=row[4]
+                )
+                for row in cur.fetchall()
+            ]
+    
+    def delete(
+        self,
+        collection: str,
+        ids: Optional[List[str]] = None,
+        filters: Optional[Dict[str, Any]] = None
+    ) -> int:
+        """Delete documents."""
+        self._init_client()
+        
+        table_name = f"{self.table_prefix}{collection}"
+        
+        with self._conn.cursor() as cur:
+            if ids:
+                placeholders = ','.join(['%s'] * len(ids))
+                cur.execute(f"DELETE FROM {table_name} WHERE id IN ({placeholders})", ids)
+            else:
+                cur.execute(f"DELETE FROM {table_name}")
+            return cur.rowcount
+    
+    def count(self, collection: str) -> int:
+        """Count documents."""
+        self._init_client()
+        
+        table_name = f"{self.table_prefix}{collection}"
+        with self._conn.cursor() as cur:
+            cur.execute(f"SELECT COUNT(*) FROM {table_name}")
+            return cur.fetchone()[0]
+    
+    def close(self) -> None:
+        """Close the connection."""
+        if self._conn:
+            self._conn.close()
+            self._conn = None
+            self._initialized = False