PraisonAI 3.0.0__py3-none-any.whl

praisonai/persistence/knowledge/langchain_adapter.py

@@ -0,0 +1,291 @@
+"""
+LangChain VectorStore adapter implementation of KnowledgeStore.
+
+Wraps any LangChain VectorStore for use with PraisonAI.
+
+Requires: langchain-core
+Install: pip install langchain-core
+"""
+
+import logging
+from typing import Any, Dict, List, Optional
+
+from .base import KnowledgeStore, KnowledgeDocument
+
+logger = logging.getLogger(__name__)
+
+
+class LangChainKnowledgeStore(KnowledgeStore):
+    """
+    LangChain VectorStore adapter for knowledge/RAG.
+    
+    Wraps any LangChain VectorStore for use with PraisonAI's interface.
+    
+    Example:
+        from langchain_chroma import Chroma
+        
+        lc_store = Chroma(persist_directory="./chroma_db")
+        store = LangChainKnowledgeStore(vectorstore=lc_store)
+    """
+    
+    def __init__(
+        self,
+        vectorstore: Any = None,
+        vectorstore_cls: Optional[str] = None,
+        vectorstore_kwargs: Optional[Dict[str, Any]] = None,
+    ):
+        """
+        Initialize LangChain adapter.
+        
+        Args:
+            vectorstore: Pre-initialized LangChain VectorStore instance
+            vectorstore_cls: Class name to import (e.g., "langchain_chroma.Chroma")
+            vectorstore_kwargs: Kwargs for vectorstore initialization
+        """
+        self._vectorstore = vectorstore
+        self._vectorstore_cls = vectorstore_cls
+        self._vectorstore_kwargs = vectorstore_kwargs or {}
+        self._initialized = vectorstore is not None
+    
+    def _init_client(self):
+        """Initialize vectorstore lazily."""
+        if self._initialized:
+            return
+        
+        if not self._vectorstore_cls:
+            raise ValueError(
+                "Either vectorstore or vectorstore_cls must be provided"
+            )
+        
+        # Dynamic import of vectorstore class
+        parts = self._vectorstore_cls.rsplit(".", 1)
+        if len(parts) == 2:
+            module_name, class_name = parts
+            import importlib
+            module = importlib.import_module(module_name)
+            cls = getattr(module, class_name)
+            self._vectorstore = cls(**self._vectorstore_kwargs)
+        else:
+            raise ValueError(f"Invalid vectorstore_cls: {self._vectorstore_cls}")
+        
+        self._initialized = True
+    
+    def create_collection(
+        self,
+        name: str,
+        dimension: int = 1536,
+        distance: str = "cosine",
+        metadata: Optional[Dict[str, Any]] = None
+    ) -> None:
+        """Create collection (handled by underlying vectorstore)."""
+        self._init_client()
+        logger.info(f"Collection creation handled by LangChain vectorstore")
+    
+    def delete_collection(self, name: str) -> bool:
+        """Delete collection."""
+        self._init_client()
+        try:
+            if hasattr(self._vectorstore, 'delete_collection'):
+                self._vectorstore.delete_collection()
+                return True
+            logger.warning("Underlying vectorstore doesn't support delete_collection")
+            return False
+        except Exception as e:
+            logger.error(f"Failed to delete collection: {e}")
+            return False
+    
+    def collection_exists(self, name: str) -> bool:
+        """Check if collection exists."""
+        self._init_client()
+        return self._vectorstore is not None
+    
+    def list_collections(self) -> List[str]:
+        """List collections."""
+        self._init_client()
+        if hasattr(self._vectorstore, 'list_collections'):
+            return self._vectorstore.list_collections()
+        return ["default"]
+    
+    def insert(
+        self,
+        collection: str,
+        documents: List[KnowledgeDocument]
+    ) -> List[str]:
+        """Insert documents."""
+        self._init_client()
+        
+        try:
+            from langchain_core.documents import Document as LCDocument
+        except ImportError:
+            raise ImportError(
+                "langchain-core is required. Install with: pip install langchain-core"
+            )
+        
+        lc_docs = []
+        ids = []
+        for doc in documents:
+            lc_doc = LCDocument(
+                page_content=doc.content,
+                metadata={
+                    **(doc.metadata or {}),
+                    "id": doc.id,
+                    "content_hash": doc.content_hash,
+                    "created_at": doc.created_at,
+                }
+            )
+            lc_docs.append(lc_doc)
+            ids.append(doc.id)
+        
+        if hasattr(self._vectorstore, 'add_documents'):
+            self._vectorstore.add_documents(lc_docs, ids=ids)
+        else:
+            logger.warning("Vectorstore doesn't support add_documents")
+        
+        return ids
+    
+    def upsert(
+        self,
+        collection: str,
+        documents: List[KnowledgeDocument]
+    ) -> List[str]:
+        """Upsert documents."""
+        # Most LangChain vectorstores handle upsert via add_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."""
+        self._init_client()
+        
+        try:
+            if hasattr(self._vectorstore, 'similarity_search_by_vector'):
+                results = self._vectorstore.similarity_search_by_vector(
+                    query_embedding,
+                    k=limit,
+                    filter=filters
+                )
+            else:
+                logger.warning("Vectorstore doesn't support vector search")
+                return []
+            
+            documents = []
+            for lc_doc in results:
+                metadata = lc_doc.metadata or {}
+                documents.append(KnowledgeDocument(
+                    id=metadata.get("id", ""),
+                    content=lc_doc.page_content,
+                    metadata=metadata,
+                    content_hash=metadata.get("content_hash"),
+                    created_at=metadata.get("created_at", 0)
+                ))
+            
+            return documents
+        except Exception as e:
+            logger.error(f"Search failed: {e}")
+            return []
+    
+    def similarity_search(
+        self,
+        query_text: str,
+        limit: int = 5,
+        filters: Optional[Dict[str, Any]] = None
+    ) -> List[KnowledgeDocument]:
+        """
+        Text-based similarity search (native LangChain method).
+        
+        Args:
+            query_text: Query string
+            limit: Maximum results
+            filters: Metadata filters
+        
+        Returns:
+            List of matching documents
+        """
+        self._init_client()
+        
+        try:
+            results = self._vectorstore.similarity_search(
+                query_text,
+                k=limit,
+                filter=filters
+            )
+            
+            documents = []
+            for lc_doc in results:
+                metadata = lc_doc.metadata or {}
+                documents.append(KnowledgeDocument(
+                    id=metadata.get("id", ""),
+                    content=lc_doc.page_content,
+                    metadata=metadata,
+                    content_hash=metadata.get("content_hash"),
+                    created_at=metadata.get("created_at", 0)
+                ))
+            
+            return documents
+        except Exception as e:
+            logger.error(f"Similarity search failed: {e}")
+            return []
+    
+    def get(
+        self,
+        collection: str,
+        ids: List[str]
+    ) -> List[KnowledgeDocument]:
+        """Get documents by IDs."""
+        self._init_client()
+        
+        if hasattr(self._vectorstore, 'get'):
+            try:
+                results = self._vectorstore.get(ids)
+                documents = []
+                for lc_doc in results:
+                    metadata = lc_doc.metadata or {}
+                    documents.append(KnowledgeDocument(
+                        id=metadata.get("id", ""),
+                        content=lc_doc.page_content,
+                        metadata=metadata
+                    ))
+                return documents
+            except Exception as e:
+                logger.warning(f"Get by ID failed: {e}")
+        
+        return []
+    
+    def delete(
+        self,
+        collection: str,
+        ids: Optional[List[str]] = None,
+        filters: Optional[Dict[str, Any]] = None
+    ) -> int:
+        """Delete documents."""
+        self._init_client()
+        
+        if ids and hasattr(self._vectorstore, 'delete'):
+            try:
+                self._vectorstore.delete(ids)
+                return len(ids)
+            except Exception as e:
+                logger.error(f"Delete failed: {e}")
+        
+        return 0
+    
+    def count(self, collection: str) -> int:
+        """Count documents."""
+        self._init_client()
+        
+        if hasattr(self._vectorstore, '__len__'):
+            return len(self._vectorstore)
+        
+        logger.warning("Vectorstore doesn't support counting")
+        return -1
+    
+    def close(self) -> None:
+        """Close the client."""
+        self._vectorstore = None
+        self._initialized = False

praisonai/persistence/knowledge/lightrag_adapter.py

@@ -0,0 +1,212 @@
+"""
+LightRAG adapter implementation of KnowledgeStore.
+
+Requires: lightrag
+Install: pip install lightrag
+"""
+
+import logging
+from typing import Any, Dict, List, Optional
+
+from .base import KnowledgeStore, KnowledgeDocument
+
+logger = logging.getLogger(__name__)
+
+
+class LightRAGKnowledgeStore(KnowledgeStore):
+    """
+    LightRAG adapter for knowledge/RAG.
+    
+    Wraps LightRAG for use with PraisonAI's KnowledgeStore interface.
+    
+    Example:
+        store = LightRAGKnowledgeStore(
+            working_dir="./lightrag_data"
+        )
+    """
+    
+    def __init__(
+        self,
+        working_dir: str = "./lightrag_data",
+        llm_model: str = "gpt-4o-mini",
+        embedding_model: str = "text-embedding-3-small",
+        **kwargs
+    ):
+        """
+        Initialize LightRAG adapter.
+        
+        Args:
+            working_dir: Directory for LightRAG data
+            llm_model: LLM model for graph extraction
+            embedding_model: Embedding model name
+            **kwargs: Additional LightRAG options
+        """
+        self.working_dir = working_dir
+        self.llm_model = llm_model
+        self.embedding_model = embedding_model
+        self.kwargs = kwargs
+        
+        self._rag = None
+        self._initialized = False
+    
+    def _init_client(self):
+        """Initialize LightRAG client lazily."""
+        if self._initialized:
+            return
+        
+        try:
+            from lightrag import LightRAG, QueryParam
+            from lightrag.llm import gpt_4o_mini_complete, openai_embedding
+        except ImportError:
+            raise ImportError(
+                "lightrag is required for LightRAG support. "
+                "Install with: pip install lightrag"
+            )
+        
+        self._QueryParam = QueryParam
+        self._rag = LightRAG(
+            working_dir=self.working_dir,
+            llm_model_func=gpt_4o_mini_complete,
+            embedding_func=openai_embedding,
+            **self.kwargs
+        )
+        self._initialized = True
+    
+    def create_collection(
+        self,
+        name: str,
+        dimension: int = 1536,
+        distance: str = "cosine",
+        metadata: Optional[Dict[str, Any]] = None
+    ) -> None:
+        """Create collection (LightRAG uses working directory)."""
+        self._init_client()
+        logger.info(f"LightRAG uses working_dir: {self.working_dir}")
+    
+    def delete_collection(self, name: str) -> bool:
+        """Delete collection data."""
+        import shutil
+        import os
+        
+        try:
+            if os.path.exists(self.working_dir):
+                shutil.rmtree(self.working_dir)
+            return True
+        except Exception as e:
+            logger.error(f"Failed to delete collection: {e}")
+            return False
+    
+    def collection_exists(self, name: str) -> bool:
+        """Check if working directory exists."""
+        import os
+        return os.path.exists(self.working_dir)
+    
+    def list_collections(self) -> List[str]:
+        """List collections (single collection per working_dir)."""
+        import os
+        if os.path.exists(self.working_dir):
+            return ["default"]
+        return []
+    
+    def insert(
+        self,
+        collection: str,
+        documents: List[KnowledgeDocument]
+    ) -> List[str]:
+        """Insert documents into LightRAG."""
+        self._init_client()
+        
+        ids = []
+        for doc in documents:
+            try:
+                self._rag.insert(doc.content)
+                ids.append(doc.id)
+            except Exception as e:
+                logger.warning(f"Failed to insert document {doc.id}: {e}")
+        
+        return ids
+    
+    def upsert(
+        self,
+        collection: str,
+        documents: List[KnowledgeDocument]
+    ) -> List[str]:
+        """Upsert documents (same as insert for LightRAG)."""
+        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 using LightRAG's query method."""
+        self._init_client()
+        
+        # LightRAG uses text query, not embeddings directly
+        # This is a limitation - we need the original query text
+        logger.warning(
+            "LightRAG search requires text query, not embeddings. "
+            "Use query() method directly for best results."
+        )
+        return []
+    
+    def query(
+        self,
+        query_text: str,
+        mode: str = "hybrid",
+        limit: int = 5
+    ) -> str:
+        """
+        Query LightRAG with text (native method).
+        
+        Args:
+            query_text: The query string
+            mode: Query mode (naive, local, global, hybrid)
+            limit: Maximum results
+        
+        Returns:
+            Query response as string
+        """
+        self._init_client()
+        
+        try:
+            result = self._rag.query(
+                query_text,
+                param=self._QueryParam(mode=mode, top_k=limit)
+            )
+            return result
+        except Exception as e:
+            logger.error(f"Query failed: {e}")
+            return ""
+    
+    def get(
+        self,
+        collection: str,
+        ids: List[str]
+    ) -> List[KnowledgeDocument]:
+        """Get documents by IDs (not supported by LightRAG)."""
+        logger.warning("LightRAG does not support get by ID")
+        return []
+    
+    def delete(
+        self,
+        collection: str,
+        ids: Optional[List[str]] = None,
+        filters: Optional[Dict[str, Any]] = None
+    ) -> int:
+        """Delete documents (not directly supported)."""
+        logger.warning("LightRAG does not support individual document deletion")
+        return 0
+    
+    def count(self, collection: str) -> int:
+        """Count documents (not directly supported)."""
+        logger.warning("LightRAG does not support document counting")
+        return -1
+    
+    def close(self) -> None:
+        """Close the client."""
+        self._rag = None
+        self._initialized = False