PraisonAI 3.0.0__py3-none-any.whl

praisonai/adapters/rerankers.py

@@ -0,0 +1,315 @@
+"""
+Reranker Implementations for PraisonAI.
+
+Provides concrete implementations of RerankerProtocol:
+- LLMReranker: Uses LLM to rerank documents
+- CrossEncoderReranker: Uses sentence-transformers cross-encoder
+- CohereReranker: Uses Cohere rerank API
+"""
+
+import logging
+from typing import Any, List, Optional
+
+logger = logging.getLogger(__name__)
+
+# Lazy import flags
+_SENTENCE_TRANSFORMERS_AVAILABLE = None
+_COHERE_AVAILABLE = None
+
+
+def _check_sentence_transformers():
+    """Check if sentence-transformers is available."""
+    global _SENTENCE_TRANSFORMERS_AVAILABLE
+    if _SENTENCE_TRANSFORMERS_AVAILABLE is None:
+        import importlib.util
+        _SENTENCE_TRANSFORMERS_AVAILABLE = importlib.util.find_spec("sentence_transformers") is not None
+    return _SENTENCE_TRANSFORMERS_AVAILABLE
+
+
+def _check_cohere():
+    """Check if cohere is available."""
+    global _COHERE_AVAILABLE
+    if _COHERE_AVAILABLE is None:
+        import importlib.util
+        _COHERE_AVAILABLE = importlib.util.find_spec("cohere") is not None
+    return _COHERE_AVAILABLE
+
+
+class LLMReranker:
+    """
+    LLM-based reranker.
+    
+    Uses an LLM to score document relevance to a query.
+    """
+    
+    name: str = "llm"
+    
+    def __init__(
+        self,
+        llm: Optional[Any] = None,
+        model: Optional[str] = None,
+        batch_size: int = 5,
+        **kwargs
+    ):
+        self.llm = llm
+        self.model = model or "gpt-4o-mini"
+        self.batch_size = batch_size
+        self._client = None
+    
+    def _get_client(self):
+        """Get LLM client - returns self.llm if provided, otherwise None (use litellm)."""
+        if self.llm:
+            return self.llm
+        return None  # Will use litellm directly
+    
+    def rerank(
+        self,
+        query: str,
+        documents: List[str],
+        top_k: Optional[int] = None,
+        **kwargs
+    ) -> List[Any]:
+        """Rerank documents using LLM scoring."""
+        from praisonaiagents.knowledge.rerankers import RerankResult
+        
+        results = []
+        
+        # Process in batches
+        for i in range(0, len(documents), self.batch_size):
+            batch = documents[i:i + self.batch_size]
+            batch_results = self._score_batch(query, batch, start_idx=i)
+            results.extend(batch_results)
+        
+        # Sort by score
+        results.sort(key=lambda x: x.score, reverse=True)
+        
+        if top_k is not None:
+            results = results[:top_k]
+        
+        return results
+    
+    async def arerank(
+        self,
+        query: str,
+        documents: List[str],
+        top_k: Optional[int] = None,
+        **kwargs
+    ) -> List[Any]:
+        """Async version (just calls sync)."""
+        return self.rerank(query, documents, top_k, **kwargs)
+    
+    def _score_batch(self, query: str, documents: List[str], start_idx: int = 0) -> List[Any]:
+        """Score a batch of documents."""
+        from praisonaiagents.knowledge.rerankers import RerankResult
+        
+        results = []
+        
+        for i, doc in enumerate(documents):
+            score = self._score_document(query, doc)
+            results.append(RerankResult(
+                text=doc,
+                score=score,
+                original_index=start_idx + i
+            ))
+        
+        return results
+    
+    def _score_document(self, query: str, document: str) -> float:
+        """Score a single document's relevance to query using litellm."""
+        try:
+            client = self._get_client()
+            
+            prompt = f"""Rate the relevance of this document to the query on a scale of 0 to 10.
+Return ONLY a number between 0 and 10.
+
+Query: {query}
+
+Document: {document[:1000]}
+
+Relevance score (0-10):"""
+            
+            # Check if client has a callable 'chat' method (PraisonAI Agent)
+            if client is not None and hasattr(client, 'chat') and callable(getattr(client, 'chat', None)):
+                # PraisonAI Agent or similar with callable chat method
+                response = client.chat(prompt)
+            else:
+                # Use litellm for unified multi-provider support
+                import litellm
+                
+                # litellm handles model-specific quirks automatically
+                completion = litellm.completion(
+                    model=self.model,
+                    messages=[{"role": "user", "content": prompt}],
+                    max_tokens=10
+                )
+                response = completion.choices[0].message.content
+            
+            # Parse score
+            if response:
+                try:
+                    score = float(response.strip().split()[0])
+                    return min(max(score / 10.0, 0.0), 1.0)  # Normalize to 0-1
+                except (ValueError, IndexError):
+                    pass
+        except Exception as e:
+            logger.warning(f"LLM scoring failed: {e}")
+        
+        return 0.5  # Default score
+
+
+class CrossEncoderReranker:
+    """
+    Cross-encoder reranker using sentence-transformers.
+    
+    Uses a cross-encoder model for accurate relevance scoring.
+    """
+    
+    name: str = "cross_encoder"
+    
+    def __init__(
+        self,
+        model_name: str = "cross-encoder/ms-marco-MiniLM-L-6-v2",
+        **kwargs
+    ):
+        if not _check_sentence_transformers():
+            raise ImportError(
+                "sentence-transformers is required for CrossEncoderReranker. "
+                "Install with: pip install sentence-transformers"
+            )
+        
+        from sentence_transformers import CrossEncoder
+        self.model = CrossEncoder(model_name)
+    
+    def rerank(
+        self,
+        query: str,
+        documents: List[str],
+        top_k: Optional[int] = None,
+        **kwargs
+    ) -> List[Any]:
+        """Rerank documents using cross-encoder."""
+        from praisonaiagents.knowledge.rerankers import RerankResult
+        
+        # Create query-document pairs
+        pairs = [[query, doc] for doc in documents]
+        
+        # Get scores
+        scores = self.model.predict(pairs)
+        
+        # Create results
+        results = []
+        for i, (doc, score) in enumerate(zip(documents, scores)):
+            results.append(RerankResult(
+                text=doc,
+                score=float(score),
+                original_index=i
+            ))
+        
+        # Sort by score
+        results.sort(key=lambda x: x.score, reverse=True)
+        
+        if top_k is not None:
+            results = results[:top_k]
+        
+        return results
+    
+    async def arerank(
+        self,
+        query: str,
+        documents: List[str],
+        top_k: Optional[int] = None,
+        **kwargs
+    ) -> List[Any]:
+        """Async version (just calls sync)."""
+        return self.rerank(query, documents, top_k, **kwargs)
+
+
+class CohereReranker:
+    """
+    Cohere reranker using Cohere's rerank API.
+    """
+    
+    name: str = "cohere"
+    
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        model: str = "rerank-english-v3.0",
+        **kwargs
+    ):
+        if not _check_cohere():
+            raise ImportError(
+                "cohere is required for CohereReranker. "
+                "Install with: pip install cohere"
+            )
+        
+        import os
+        import cohere
+        
+        self.api_key = api_key or os.environ.get("COHERE_API_KEY")
+        if not self.api_key:
+            raise ValueError("COHERE_API_KEY environment variable or api_key parameter required")
+        
+        self.model = model
+        self._client = cohere.Client(self.api_key)
+    
+    def rerank(
+        self,
+        query: str,
+        documents: List[str],
+        top_k: Optional[int] = None,
+        **kwargs
+    ) -> List[Any]:
+        """Rerank documents using Cohere API."""
+        from praisonaiagents.knowledge.rerankers import RerankResult
+        
+        response = self._client.rerank(
+            query=query,
+            documents=documents,
+            top_n=top_k or len(documents),
+            model=self.model
+        )
+        
+        results = []
+        for r in response.results:
+            results.append(RerankResult(
+                text=documents[r.index],
+                score=r.relevance_score,
+                original_index=r.index
+            ))
+        
+        return results
+    
+    async def arerank(
+        self,
+        query: str,
+        documents: List[str],
+        top_k: Optional[int] = None,
+        **kwargs
+    ) -> List[Any]:
+        """Async version (just calls sync)."""
+        return self.rerank(query, documents, top_k, **kwargs)
+
+
+def register_default_rerankers():
+    """Register all default rerankers with the registry."""
+    from praisonaiagents.knowledge.rerankers import get_reranker_registry
+    
+    registry = get_reranker_registry()
+    
+    # LLM reranker is always available
+    registry.register("llm", LLMReranker)
+    
+    # Cross-encoder if available
+    if _check_sentence_transformers():
+        registry.register("cross_encoder", CrossEncoderReranker)
+    
+    # Cohere if available
+    if _check_cohere():
+        registry.register("cohere", CohereReranker)
+    
+    logger.debug("Registered default rerankers")
+
+
+# Auto-register on import
+register_default_rerankers()

praisonai/adapters/retrievers.py

@@ -0,0 +1,394 @@
+"""
+Retriever Implementations for PraisonAI.
+
+Provides concrete implementations of RetrieverProtocol:
+- BasicRetriever: Simple vector similarity retrieval
+- FusionRetriever: Multi-query with RRF fusion
+- RecursiveRetriever: Depth-limited sub-retrieval
+- AutoMergeRetriever: Merges adjacent chunks
+"""
+
+import logging
+from typing import Any, Callable, Dict, List, Optional
+
+logger = logging.getLogger(__name__)
+
+
+class BasicRetriever:
+    """
+    Basic vector similarity retriever.
+    
+    Simply queries the vector store and returns results.
+    """
+    
+    name: str = "basic"
+    
+    def __init__(
+        self,
+        vector_store: Any,
+        embedding_fn: Optional[Callable[[str], List[float]]] = None,
+        top_k: int = 10,
+        **kwargs
+    ):
+        self.vector_store = vector_store
+        self.embedding_fn = embedding_fn
+        self.top_k = top_k
+    
+    def retrieve(
+        self,
+        query: str,
+        top_k: Optional[int] = None,
+        filter: Optional[Dict[str, Any]] = None,
+        **kwargs
+    ) -> List[Any]:
+        """Retrieve documents by vector similarity."""
+        from praisonaiagents.knowledge.retrieval import RetrievalResult
+        
+        k = top_k or self.top_k
+        
+        # Get query embedding
+        if self.embedding_fn:
+            query_embedding = self.embedding_fn(query)
+        else:
+            raise ValueError("embedding_fn is required for BasicRetriever")
+        
+        # Query vector store
+        results = self.vector_store.query(
+            embedding=query_embedding,
+            top_k=k,
+            filter=filter
+        )
+        
+        # Convert to RetrievalResult
+        retrieval_results = []
+        for r in results:
+            retrieval_results.append(RetrievalResult(
+                text=r.text if hasattr(r, 'text') else str(r),
+                score=r.score if hasattr(r, 'score') else 1.0,
+                metadata=r.metadata if hasattr(r, 'metadata') else {},
+                doc_id=r.id if hasattr(r, 'id') else None,
+                chunk_index=r.metadata.get('chunk_index') if hasattr(r, 'metadata') else None
+            ))
+        
+        return retrieval_results
+    
+    async def aretrieve(
+        self,
+        query: str,
+        top_k: Optional[int] = None,
+        filter: Optional[Dict[str, Any]] = None,
+        **kwargs
+    ) -> List[Any]:
+        """Async version (just calls sync)."""
+        return self.retrieve(query, top_k, filter, **kwargs)
+
+
+class FusionRetriever:
+    """
+    Fusion retriever using Reciprocal Rank Fusion.
+    
+    Generates multiple query variations and fuses results.
+    """
+    
+    name: str = "fusion"
+    
+    def __init__(
+        self,
+        vector_store: Any,
+        embedding_fn: Optional[Callable[[str], List[float]]] = None,
+        llm: Optional[Any] = None,
+        num_queries: int = 3,
+        top_k: int = 10,
+        rrf_k: int = 60,
+        **kwargs
+    ):
+        self.vector_store = vector_store
+        self.embedding_fn = embedding_fn
+        self.llm = llm
+        self.num_queries = num_queries
+        self.top_k = top_k
+        self.rrf_k = rrf_k
+    
+    def retrieve(
+        self,
+        query: str,
+        top_k: Optional[int] = None,
+        filter: Optional[Dict[str, Any]] = None,
+        **kwargs
+    ) -> List[Any]:
+        """Retrieve with query fusion."""
+        from praisonaiagents.knowledge.retrieval import RetrievalResult, reciprocal_rank_fusion
+        
+        k = top_k or self.top_k
+        
+        # Generate query variations
+        queries = self._generate_query_variations(query)
+        
+        # Retrieve for each query
+        all_results = []
+        for q in queries:
+            if self.embedding_fn:
+                q_embedding = self.embedding_fn(q)
+                results = self.vector_store.query(
+                    embedding=q_embedding,
+                    top_k=k,
+                    filter=filter
+                )
+                
+                # Convert to RetrievalResult
+                query_results = []
+                for r in results:
+                    query_results.append(RetrievalResult(
+                        text=r.text if hasattr(r, 'text') else str(r),
+                        score=r.score if hasattr(r, 'score') else 1.0,
+                        metadata=r.metadata if hasattr(r, 'metadata') else {},
+                        doc_id=r.id if hasattr(r, 'id') else None
+                    ))
+                all_results.append(query_results)
+        
+        # Fuse results using RRF
+        fused = reciprocal_rank_fusion(all_results, k=self.rrf_k)
+        
+        return fused[:k]
+    
+    async def aretrieve(
+        self,
+        query: str,
+        top_k: Optional[int] = None,
+        filter: Optional[Dict[str, Any]] = None,
+        **kwargs
+    ) -> List[Any]:
+        """Async version (just calls sync)."""
+        return self.retrieve(query, top_k, filter, **kwargs)
+    
+    def _generate_query_variations(self, query: str) -> List[str]:
+        """Generate query variations."""
+        queries = [query]
+        
+        # Try LLM-based variation if available
+        if self.llm and len(queries) < self.num_queries:
+            try:
+                prompt = f"""Generate {self.num_queries - 1} alternative phrasings of this search query.
+Return only the queries, one per line.
+
+Original query: {query}"""
+                response = self.llm.chat(prompt)
+                if response:
+                    for line in response.strip().split("\n"):
+                        line = line.strip()
+                        if line and line != query:
+                            queries.append(line)
+                            if len(queries) >= self.num_queries:
+                                break
+            except Exception as e:
+                logger.warning(f"LLM query variation failed: {e}")
+        
+        # Fallback: simple variations
+        if len(queries) < self.num_queries:
+            # Add question form
+            if not query.endswith("?"):
+                queries.append(f"What is {query}?")
+            # Add keyword form
+            keywords = " ".join(query.split()[:5])
+            if keywords != query:
+                queries.append(keywords)
+        
+        return queries[:self.num_queries]
+
+
+class RecursiveRetriever:
+    """
+    Recursive retriever with depth-limited sub-retrieval.
+    
+    Retrieves, expands queries based on results, and retrieves again.
+    """
+    
+    name: str = "recursive"
+    
+    def __init__(
+        self,
+        vector_store: Any,
+        embedding_fn: Optional[Callable[[str], List[float]]] = None,
+        llm: Optional[Any] = None,
+        max_depth: int = 2,
+        top_k: int = 10,
+        **kwargs
+    ):
+        self.vector_store = vector_store
+        self.embedding_fn = embedding_fn
+        self.llm = llm
+        self.max_depth = max_depth
+        self.top_k = top_k
+    
+    def retrieve(
+        self,
+        query: str,
+        top_k: Optional[int] = None,
+        filter: Optional[Dict[str, Any]] = None,
+        **kwargs
+    ) -> List[Any]:
+        """Retrieve with recursive expansion."""
+        from praisonaiagents.knowledge.retrieval import RetrievalResult
+        
+        k = top_k or self.top_k
+        seen_ids = set()
+        all_results = []
+        
+        current_queries = [query]
+        
+        for depth in range(self.max_depth):
+            next_queries = []
+            
+            for q in current_queries:
+                if self.embedding_fn:
+                    q_embedding = self.embedding_fn(q)
+                    results = self.vector_store.query(
+                        embedding=q_embedding,
+                        top_k=k,
+                        filter=filter
+                    )
+                    
+                    for r in results:
+                        r_id = r.id if hasattr(r, 'id') else r.text[:50]
+                        if r_id not in seen_ids:
+                            seen_ids.add(r_id)
+                            all_results.append(RetrievalResult(
+                                text=r.text if hasattr(r, 'text') else str(r),
+                                score=r.score if hasattr(r, 'score') else 1.0,
+                                metadata=r.metadata if hasattr(r, 'metadata') else {},
+                                doc_id=r.id if hasattr(r, 'id') else None
+                            ))
+                            
+                            # Generate follow-up query from result
+                            if depth < self.max_depth - 1 and self.llm:
+                                follow_up = self._generate_follow_up(query, r.text if hasattr(r, 'text') else str(r))
+                                if follow_up:
+                                    next_queries.append(follow_up)
+            
+            current_queries = next_queries[:3]  # Limit expansion
+            if not current_queries:
+                break
+        
+        # Sort by score
+        all_results.sort(key=lambda x: x.score, reverse=True)
+        return all_results[:k]
+    
+    async def aretrieve(
+        self,
+        query: str,
+        top_k: Optional[int] = None,
+        filter: Optional[Dict[str, Any]] = None,
+        **kwargs
+    ) -> List[Any]:
+        """Async version (just calls sync)."""
+        return self.retrieve(query, top_k, filter, **kwargs)
+    
+    def _generate_follow_up(self, original_query: str, context: str) -> Optional[str]:
+        """Generate a follow-up query based on retrieved context."""
+        if not self.llm:
+            return None
+        
+        try:
+            prompt = f"""Based on this search result, generate a follow-up query to find more relevant information.
+Return only the query, nothing else.
+
+Original query: {original_query}
+Search result: {context[:500]}"""
+            response = self.llm.chat(prompt)
+            if response:
+                return response.strip()
+        except Exception as e:
+            logger.warning(f"Follow-up query generation failed: {e}")
+        
+        return None
+
+
+class AutoMergeRetriever:
+    """
+    Auto-merge retriever that combines adjacent chunks.
+    
+    Retrieves chunks and merges those from the same document.
+    """
+    
+    name: str = "auto_merge"
+    
+    def __init__(
+        self,
+        vector_store: Any,
+        embedding_fn: Optional[Callable[[str], List[float]]] = None,
+        top_k: int = 10,
+        max_gap: int = 1,
+        **kwargs
+    ):
+        self.vector_store = vector_store
+        self.embedding_fn = embedding_fn
+        self.top_k = top_k
+        self.max_gap = max_gap
+    
+    def retrieve(
+        self,
+        query: str,
+        top_k: Optional[int] = None,
+        filter: Optional[Dict[str, Any]] = None,
+        **kwargs
+    ) -> List[Any]:
+        """Retrieve and merge adjacent chunks."""
+        from praisonaiagents.knowledge.retrieval import RetrievalResult, merge_adjacent_chunks
+        
+        k = top_k or self.top_k
+        
+        # Get more results than needed for merging
+        fetch_k = k * 3
+        
+        if self.embedding_fn:
+            q_embedding = self.embedding_fn(query)
+            results = self.vector_store.query(
+                embedding=q_embedding,
+                top_k=fetch_k,
+                filter=filter
+            )
+            
+            # Convert to RetrievalResult
+            retrieval_results = []
+            for r in results:
+                retrieval_results.append(RetrievalResult(
+                    text=r.text if hasattr(r, 'text') else str(r),
+                    score=r.score if hasattr(r, 'score') else 1.0,
+                    metadata=r.metadata if hasattr(r, 'metadata') else {},
+                    doc_id=r.metadata.get('doc_id') if hasattr(r, 'metadata') else None,
+                    chunk_index=r.metadata.get('chunk_index') if hasattr(r, 'metadata') else None
+                ))
+            
+            # Merge adjacent chunks
+            merged = merge_adjacent_chunks(retrieval_results, max_gap=self.max_gap)
+            return merged[:k]
+        
+        return []
+    
+    async def aretrieve(
+        self,
+        query: str,
+        top_k: Optional[int] = None,
+        filter: Optional[Dict[str, Any]] = None,
+        **kwargs
+    ) -> List[Any]:
+        """Async version (just calls sync)."""
+        return self.retrieve(query, top_k, filter, **kwargs)
+
+
+def register_default_retrievers():
+    """Register all default retrievers with the registry."""
+    from praisonaiagents.knowledge.retrieval import get_retriever_registry
+    
+    registry = get_retriever_registry()
+    
+    registry.register("basic", BasicRetriever)
+    registry.register("fusion", FusionRetriever)
+    registry.register("recursive", RecursiveRetriever)
+    registry.register("auto_merge", AutoMergeRetriever)
+    
+    logger.debug("Registered default retrievers")
+
+
+# Auto-register on import
+register_default_retrievers()