PraisonAI 3.0.0__py3-none-any.whl

praisonai/persistence/knowledge/pinecone.py

@@ -0,0 +1,253 @@
+"""
+Pinecone implementation of KnowledgeStore.
+
+Requires: pinecone-client
+Install: pip install pinecone-client
+"""
+
+import logging
+import os
+from typing import Any, Dict, List, Optional
+
+from .base import KnowledgeStore, KnowledgeDocument
+
+logger = logging.getLogger(__name__)
+
+
+class PineconeKnowledgeStore(KnowledgeStore):
+    """
+    Pinecone-based knowledge store for vector search.
+    
+    Cloud-only managed vector database.
+    
+    Example:
+        store = PineconeKnowledgeStore(
+            api_key="your-api-key",
+            index_name="praisonai"
+        )
+    """
+    
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        index_name: Optional[str] = None,
+        environment: Optional[str] = None,
+        namespace: str = "",
+    ):
+        try:
+            from pinecone import Pinecone
+        except ImportError:
+            raise ImportError(
+                "pinecone-client is required for Pinecone support. "
+                "Install with: pip install pinecone-client"
+            )
+        
+        api_key = api_key or os.getenv("PINECONE_API_KEY")
+        if not api_key:
+            raise ValueError("Pinecone API key is required")
+        
+        self._pc = Pinecone(api_key=api_key)
+        self._index_name = index_name or os.getenv("PINECONE_INDEX", "praisonai")
+        self._namespace = namespace
+        self._index = None
+        
+        # Try to connect to existing index
+        try:
+            self._index = self._pc.Index(self._index_name)
+            logger.info(f"Connected to Pinecone index: {self._index_name}")
+        except Exception as e:
+            logger.warning(f"Index {self._index_name} not found: {e}")
+    
+    def create_collection(
+        self,
+        name: str,
+        dimension: int,
+        distance: str = "cosine",
+        metadata: Optional[Dict[str, Any]] = None
+    ) -> None:
+        """Create a new index (Pinecone calls them indexes, not collections)."""
+        from pinecone import ServerlessSpec
+        
+        metric_map = {"cosine": "cosine", "euclidean": "euclidean", "dot": "dotproduct"}
+        
+        self._pc.create_index(
+            name=name,
+            dimension=dimension,
+            metric=metric_map.get(distance, "cosine"),
+            spec=ServerlessSpec(cloud="aws", region="us-east-1")
+        )
+        self._index_name = name
+        self._index = self._pc.Index(name)
+        logger.info(f"Created Pinecone index: {name}")
+    
+    def delete_collection(self, name: str) -> bool:
+        """Delete an index."""
+        try:
+            self._pc.delete_index(name)
+            if name == self._index_name:
+                self._index = None
+            return True
+        except Exception as e:
+            logger.warning(f"Failed to delete index {name}: {e}")
+            return False
+    
+    def collection_exists(self, name: str) -> bool:
+        """Check if an index exists."""
+        try:
+            indexes = self._pc.list_indexes()
+            return any(idx.name == name for idx in indexes)
+        except Exception:
+            return False
+    
+    def list_collections(self) -> List[str]:
+        """List all indexes."""
+        indexes = self._pc.list_indexes()
+        return [idx.name for idx in indexes]
+    
+    def insert(
+        self,
+        collection: str,
+        documents: List[KnowledgeDocument]
+    ) -> List[str]:
+        """Insert documents."""
+        if collection != self._index_name:
+            self._index = self._pc.Index(collection)
+            self._index_name = collection
+        
+        vectors = []
+        ids = []
+        
+        for doc in documents:
+            if doc.embedding is None:
+                raise ValueError(f"Document {doc.id} has no embedding")
+            
+            vectors.append({
+                "id": doc.id,
+                "values": doc.embedding,
+                "metadata": {
+                    "content": doc.content,
+                    "content_hash": doc.content_hash or "",
+                    "created_at": doc.created_at,
+                    **(doc.metadata or {})
+                }
+            })
+            ids.append(doc.id)
+        
+        self._index.upsert(vectors=vectors, namespace=self._namespace)
+        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."""
+        if collection != self._index_name:
+            self._index = self._pc.Index(collection)
+            self._index_name = collection
+        
+        results = self._index.query(
+            vector=query_embedding,
+            top_k=limit,
+            filter=filters,
+            include_metadata=True,
+            namespace=self._namespace,
+        )
+        
+        documents = []
+        for match in results.matches:
+            if score_threshold and match.score < score_threshold:
+                continue
+            
+            metadata = match.metadata or {}
+            content = metadata.pop("content", "")
+            content_hash = metadata.pop("content_hash", None)
+            created_at = metadata.pop("created_at", 0)
+            
+            doc = KnowledgeDocument(
+                id=match.id,
+                content=content,
+                embedding=None,
+                metadata=metadata,
+                content_hash=content_hash,
+                created_at=created_at,
+            )
+            documents.append(doc)
+        
+        return documents
+    
+    def get(
+        self,
+        collection: str,
+        ids: List[str]
+    ) -> List[KnowledgeDocument]:
+        """Get documents by IDs."""
+        if collection != self._index_name:
+            self._index = self._pc.Index(collection)
+            self._index_name = collection
+        
+        results = self._index.fetch(ids=ids, namespace=self._namespace)
+        
+        documents = []
+        for doc_id, vector in results.vectors.items():
+            metadata = vector.metadata or {}
+            content = metadata.pop("content", "")
+            content_hash = metadata.pop("content_hash", None)
+            created_at = metadata.pop("created_at", 0)
+            
+            doc = KnowledgeDocument(
+                id=doc_id,
+                content=content,
+                embedding=list(vector.values),
+                metadata=metadata,
+                content_hash=content_hash,
+                created_at=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 collection != self._index_name:
+            self._index = self._pc.Index(collection)
+            self._index_name = collection
+        
+        if ids:
+            self._index.delete(ids=ids, namespace=self._namespace)
+            return len(ids)
+        elif filters:
+            self._index.delete(filter=filters, namespace=self._namespace)
+            return -1
+        return 0
+    
+    def count(self, collection: str) -> int:
+        """Count documents."""
+        if collection != self._index_name:
+            self._index = self._pc.Index(collection)
+            self._index_name = collection
+        
+        stats = self._index.describe_index_stats()
+        if self._namespace:
+            return stats.namespaces.get(self._namespace, {}).get("vector_count", 0)
+        return stats.total_vector_count
+    
+    def close(self) -> None:
+        """Close the store."""
+        self._index = None

praisonai/persistence/knowledge/qdrant.py

@@ -0,0 +1,301 @@
+"""
+Qdrant implementation of KnowledgeStore.
+
+Requires: qdrant-client
+Install: pip install qdrant-client
+"""
+
+import logging
+from typing import Any, Dict, List, Optional
+
+from .base import KnowledgeStore, KnowledgeDocument
+
+logger = logging.getLogger(__name__)
+
+
+class QdrantKnowledgeStore(KnowledgeStore):
+    """
+    Qdrant-based knowledge store for vector search.
+    
+    Example:
+        store = QdrantKnowledgeStore(
+            url="http://localhost:6333"
+        )
+    """
+    
+    def __init__(
+        self,
+        url: Optional[str] = None,
+        host: str = "localhost",
+        port: int = 6333,
+        grpc_port: int = 6334,
+        api_key: Optional[str] = None,
+        prefer_grpc: bool = False,
+        timeout: int = 30,
+    ):
+        """
+        Initialize Qdrant knowledge store.
+        
+        Args:
+            url: Full URL (overrides host/port)
+            host: Qdrant host
+            port: Qdrant HTTP port
+            grpc_port: Qdrant gRPC port
+            api_key: API key for Qdrant Cloud
+            prefer_grpc: Use gRPC instead of HTTP
+            timeout: Request timeout in seconds
+        """
+        try:
+            from qdrant_client import QdrantClient
+            from qdrant_client.models import (
+                Distance, VectorParams, PointStruct,
+                Filter, FieldCondition, MatchValue
+            )
+        except ImportError:
+            raise ImportError(
+                "qdrant-client is required for Qdrant support. "
+                "Install with: pip install qdrant-client"
+            )
+        
+        self._QdrantClient = QdrantClient
+        self._Distance = Distance
+        self._VectorParams = VectorParams
+        self._PointStruct = PointStruct
+        self._Filter = Filter
+        self._FieldCondition = FieldCondition
+        self._MatchValue = MatchValue
+        
+        if url:
+            self._client = QdrantClient(url=url, api_key=api_key, timeout=timeout)
+        else:
+            self._client = QdrantClient(
+                host=host,
+                port=port,
+                grpc_port=grpc_port,
+                api_key=api_key,
+                prefer_grpc=prefer_grpc,
+                timeout=timeout,
+            )
+        
+        logger.info(f"Connected to Qdrant at {url or f'{host}:{port}'}")
+    
+    def _get_distance(self, distance: str) -> Any:
+        """Convert distance string to Qdrant Distance enum."""
+        distance_map = {
+            "cosine": self._Distance.COSINE,
+            "euclidean": self._Distance.EUCLID,
+            "dot": self._Distance.DOT,
+            "manhattan": self._Distance.MANHATTAN,
+        }
+        return distance_map.get(distance.lower(), self._Distance.COSINE)
+    
+    def create_collection(
+        self,
+        name: str,
+        dimension: int,
+        distance: str = "cosine",
+        metadata: Optional[Dict[str, Any]] = None
+    ) -> None:
+        """Create a new collection."""
+        self._client.create_collection(
+            collection_name=name,
+            vectors_config=self._VectorParams(
+                size=dimension,
+                distance=self._get_distance(distance)
+            )
+        )
+        logger.info(f"Created Qdrant collection: {name} (dim={dimension}, distance={distance})")
+    
+    def delete_collection(self, name: str) -> bool:
+        """Delete a collection."""
+        try:
+            self._client.delete_collection(collection_name=name)
+            logger.info(f"Deleted Qdrant collection: {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."""
+        try:
+            collections = self._client.get_collections()
+            return any(c.name == name for c in collections.collections)
+        except Exception:
+            return False
+    
+    def list_collections(self) -> List[str]:
+        """List all collections."""
+        collections = self._client.get_collections()
+        return [c.name for c in collections.collections]
+    
+    def insert(
+        self,
+        collection: str,
+        documents: List[KnowledgeDocument]
+    ) -> List[str]:
+        """Insert documents into a collection."""
+        points = []
+        ids = []
+        
+        for i, doc in enumerate(documents):
+            if doc.embedding is None:
+                raise ValueError(f"Document {doc.id} has no embedding")
+            
+            point = self._PointStruct(
+                id=i,  # Qdrant prefers numeric IDs for performance
+                vector=doc.embedding,
+                payload={
+                    "doc_id": doc.id,
+                    "content": doc.content,
+                    "content_hash": doc.content_hash,
+                    "created_at": doc.created_at,
+                    **(doc.metadata or {})
+                }
+            )
+            points.append(point)
+            ids.append(doc.id)
+        
+        self._client.upsert(collection_name=collection, points=points)
+        logger.debug(f"Inserted {len(documents)} documents into {collection}")
+        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."""
+        search_filter = None
+        if filters:
+            conditions = [
+                self._FieldCondition(
+                    key=k,
+                    match=self._MatchValue(value=v)
+                )
+                for k, v in filters.items()
+            ]
+            search_filter = self._Filter(must=conditions)
+        
+        results = self._client.query_points(
+            collection_name=collection,
+            query=query_embedding,
+            limit=limit,
+            query_filter=search_filter,
+            score_threshold=score_threshold,
+        )
+        
+        documents = []
+        for point in results.points:
+            payload = point.payload or {}
+            doc = KnowledgeDocument(
+                id=payload.get("doc_id", str(point.id)),
+                content=payload.get("content", ""),
+                embedding=None,  # Don't return embeddings by default
+                metadata={
+                    k: v for k, v in payload.items()
+                    if k not in ("doc_id", "content", "content_hash", "created_at")
+                },
+                content_hash=payload.get("content_hash"),
+                created_at=payload.get("created_at", 0),
+            )
+            documents.append(doc)
+        
+        return documents
+    
+    def get(
+        self,
+        collection: str,
+        ids: List[str]
+    ) -> List[KnowledgeDocument]:
+        """Get documents by IDs."""
+        # Qdrant uses numeric IDs internally, so we search by doc_id in payload
+        documents = []
+        for doc_id in ids:
+            results = self._client.scroll(
+                collection_name=collection,
+                scroll_filter=self._Filter(
+                    must=[
+                        self._FieldCondition(
+                            key="doc_id",
+                            match=self._MatchValue(value=doc_id)
+                        )
+                    ]
+                ),
+                limit=1,
+            )
+            for point in results[0]:
+                payload = point.payload or {}
+                doc = KnowledgeDocument(
+                    id=payload.get("doc_id", str(point.id)),
+                    content=payload.get("content", ""),
+                    embedding=list(point.vector) if point.vector else None,
+                    metadata={
+                        k: v for k, v in payload.items()
+                        if k not in ("doc_id", "content", "content_hash", "created_at")
+                    },
+                    content_hash=payload.get("content_hash"),
+                    created_at=payload.get("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 by IDs or filters."""
+        if ids:
+            # Delete by doc_id in payload
+            for doc_id in ids:
+                self._client.delete(
+                    collection_name=collection,
+                    points_selector=self._Filter(
+                        must=[
+                            self._FieldCondition(
+                                key="doc_id",
+                                match=self._MatchValue(value=doc_id)
+                            )
+                        ]
+                    )
+                )
+            return len(ids)
+        elif filters:
+            conditions = [
+                self._FieldCondition(
+                    key=k,
+                    match=self._MatchValue(value=v)
+                )
+                for k, v in filters.items()
+            ]
+            self._client.delete(
+                collection_name=collection,
+                points_selector=self._Filter(must=conditions)
+            )
+            return -1  # Unknown count
+        return 0
+    
+    def count(self, collection: str) -> int:
+        """Count documents in a collection."""
+        info = self._client.get_collection(collection_name=collection)
+        return info.points_count or 0
+    
+    def close(self) -> None:
+        """Close the store."""
+        if self._client:
+            self._client.close()
+            self._client = None