MindsDB 25.1.2.1__py3-none-any.whl → 25.1.5.0__py3-none-any.whl

mindsdb/integrations/utilities/rag/pipelines/rag.py

@@ -1,8 +1,9 @@
 from copy import copy
-from typing import Optional, Any
+from typing import Optional, Any, List
 
 from langchain_core.output_parsers import StrOutputParser
 from langchain.retrievers import ContextualCompressionRetriever
+from langchain_core.documents import Document
 
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_core.runnables import RunnableParallel, RunnablePassthrough, RunnableSerializable
@@ -28,6 +29,23 @@ from mindsdb.interfaces.agents.langchain_agent import create_chat_model
 class LangChainRAGPipeline:
     """
     Builds a RAG pipeline using langchain LCEL components
+
+    Args:
+        retriever_runnable: Base retriever component
+        prompt_template: Template for generating responses
+        llm: Language model for generating responses
+        reranker (bool): Whether to use reranking (default: False)
+        reranker_config (RerankerConfig): Configuration for the reranker, including:
+            - model: Model to use for reranking
+            - filtering_threshold: Minimum score to keep a document
+            - num_docs_to_keep: Maximum number of documents to keep
+            - max_concurrent_requests: Maximum concurrent API requests
+            - max_retries: Number of retry attempts for failed requests
+            - retry_delay: Delay between retries
+            - early_stop (bool): Whether to enable early stopping
+            - early_stop_threshold: Confidence threshold for early stopping
+        vector_store_config (VectorStoreConfig): Vector store configuration
+        summarization_config (SummarizationConfig): Summarization configuration
     """
 
     def __init__(
@@ -40,19 +58,15 @@ class LangChainRAGPipeline:
             vector_store_config: Optional[VectorStoreConfig] = None,
             summarization_config: Optional[SummarizationConfig] = None
     ):
-
         self.retriever_runnable = retriever_runnable
         self.prompt_template = prompt_template
         self.llm = llm
         if reranker:
             if reranker_config is None:
                 reranker_config = RerankerConfig()
-            self.reranker = LLMReranker(
-                model=reranker_config.model,
-                base_url=reranker_config.base_url,
-                filtering_threshold=reranker_config.filtering_threshold,
-                num_docs_to_keep=reranker_config.num_docs_to_keep
-            )
+            # Convert config to dict and initialize reranker
+            reranker_kwargs = reranker_config.model_dump(exclude_none=True)
+            self.reranker = LLMReranker(**reranker_kwargs)
         else:
             self.reranker = None
         self.summarizer = None
@@ -102,17 +116,45 @@ class LangChainRAGPipeline:
             raise ValueError("One of the required components (llm) is None")
 
         if self.reranker:
-            reranker = self.reranker
-            retriever = copy(self.retriever_runnable)
-            self.retriever_runnable = ContextualCompressionRetriever(
-                base_compressor=reranker, base_retriever=retriever
+            # Create a custom retriever that handles async operations properly
+            class AsyncRerankerRetriever(ContextualCompressionRetriever):
+                """Async-aware retriever that properly handles concurrent reranking operations."""
+
+                def __init__(self, base_retriever, reranker):
+                    super().__init__(
+                        base_compressor=reranker,
+                        base_retriever=base_retriever
+                    )
+
+                async def ainvoke(self, query: str) -> List[Document]:
+                    """Async retrieval with proper concurrency handling."""
+                    # Get initial documents
+                    if hasattr(self.base_retriever, 'ainvoke'):
+                        docs = await self.base_retriever.ainvoke(query)
+                    else:
+                        docs = await RunnablePassthrough(self.base_retriever.get_relevant_documents)(query)
+
+                    # Rerank documents
+                    if docs:
+                        docs = await self.base_compressor.acompress_documents(docs, query)
+                    return docs
+
+                def get_relevant_documents(self, query: str) -> List[Document]:
+                    """Sync wrapper for async retrieval."""
+                    import asyncio
+                    return asyncio.run(self.ainvoke(query))
+
+            # Use our custom async-aware retriever
+            self.retriever_runnable = AsyncRerankerRetriever(
+                base_retriever=copy(self.retriever_runnable),
+                reranker=self.reranker
             )
 
         rag_chain_from_docs = (
-                RunnablePassthrough.assign(context=(lambda x: format_docs(x["context"])))  # noqa: E126, E122
-                | prompt
-                | self.llm
-                | StrOutputParser()
+            RunnablePassthrough.assign(context=(lambda x: format_docs(x["context"])))
+            | prompt
+            | self.llm
+            | StrOutputParser()
         )
 
         retrieval_chain = RunnableParallel(
@@ -125,6 +167,16 @@ class LangChainRAGPipeline:
         rag_chain_with_source = retrieval_chain.assign(answer=rag_chain_from_docs)
         return rag_chain_with_source
 
+    async def ainvoke(self, input_dict: dict) -> dict:
+        """Async invocation of the RAG pipeline."""
+        chain = self.with_returned_sources()
+        return await chain.ainvoke(input_dict)
+
+    def invoke(self, input_dict: dict) -> dict:
+        """Sync invocation of the RAG pipeline."""
+        import asyncio
+        return asyncio.run(self.ainvoke(input_dict))
+
     @classmethod
     def _apply_search_kwargs(cls, retriever: Any, search_kwargs: Optional[SearchKwargs] = None, search_type: Optional[SearchType] = None) -> Any:
         """Apply search kwargs and search type to the retriever if they exist"""
@@ -235,6 +287,10 @@ class LangChainRAGPipeline:
         )
         vector_store_retriever = vector_store_operator.vector_store.as_retriever()
         vector_store_retriever = cls._apply_search_kwargs(vector_store_retriever, config.search_kwargs, config.search_type)
+        distance_function = DistanceFunction.SQUARED_EUCLIDEAN_DISTANCE
+        if config.vector_store_config.is_sparse and config.vector_store_config.vector_size is not None:
+            # Use negative dot product for sparse retrieval.
+            distance_function = DistanceFunction.NEGATIVE_DOT_PRODUCT
         retriever = SQLRetriever(
             fallback_retriever=vector_store_retriever,
             vector_store_handler=knowledge_base_table.get_vector_db(),
@@ -242,14 +298,11 @@ class LangChainRAGPipeline:
             examples=retriever_config.examples,
             embeddings_model=embeddings,
             rewrite_prompt_template=retriever_config.rewrite_prompt_template,
-            retry_prompt_template=retriever_config.query_retry_template,
+            metadata_filters_prompt_template=retriever_config.metadata_filters_prompt_template,
             num_retries=retriever_config.num_retries,
-            sql_prompt_template=retriever_config.sql_prompt_template,
-            query_checker_template=retriever_config.query_checker_template,
             embeddings_table=knowledge_base_table._kb.vector_database_table,
             source_table=retriever_config.source_table,
-            # Currently only similarity search is supported.
-            distance_function=DistanceFunction.SQUARED_EUCLIDEAN_DISTANCE,
+            distance_function=distance_function,
             search_kwargs=config.search_kwargs,
             llm=sql_llm
         )

mindsdb/integrations/utilities/rag/rerankers/reranker_compressor.py

@@ -4,16 +4,15 @@ import asyncio
 import logging
 import math
 import os
+import random
 from typing import Any, Dict, List, Optional, Sequence, Tuple
-from uuid import uuid4
+
 from langchain.retrievers.document_compressors.base import BaseDocumentCompressor
 from langchain_core.callbacks import Callbacks
-
-from mindsdb.integrations.utilities.rag.settings import DEFAULT_RERANKING_MODEL, DEFAULT_LLM_ENDPOINT
 from langchain_core.documents import Document
-from langchain_core.messages import HumanMessage, SystemMessage
-from langchain_openai import ChatOpenAI
+from openai import AsyncOpenAI
 
+from mindsdb.integrations.utilities.rag.settings import DEFAULT_RERANKING_MODEL, DEFAULT_LLM_ENDPOINT
 
 log = logging.getLogger(__name__)
 
@@ -23,128 +22,187 @@ class LLMReranker(BaseDocumentCompressor):
     model: str = DEFAULT_RERANKING_MODEL  # Model to use for reranking
     temperature: float = 0.0  # Temperature for the model
     openai_api_key: Optional[str] = None
-    remove_irrelevant: bool = True  # New flag to control removal of irrelevant documents,
+    remove_irrelevant: bool = True  # New flag to control removal of irrelevant documents
     base_url: str = DEFAULT_LLM_ENDPOINT
     num_docs_to_keep: Optional[int] = None  # How many of the top documents to keep after reranking & compressing.
-
     _api_key_var: str = "OPENAI_API_KEY"
-    client: Optional[Any] = None
+    client: Optional[AsyncOpenAI] = None
+    _semaphore: Optional[asyncio.Semaphore] = None
+    max_concurrent_requests: int = 20
+    max_retries: int = 3
+    retry_delay: float = 1.0
+    request_timeout: float = 20.0  # Timeout for API requests
+    early_stop: bool = True  # Whether to enable early stopping
+    early_stop_threshold: float = 0.8  # Confidence threshold for early stopping
 
     class Config:
         arbitrary_types_allowed = True
 
-    async def search_relevancy(self, query: str, document: str) -> Any:
-        openai_api_key = self.openai_api_key or os.getenv(self._api_key_var)
-
-        # Initialize the ChatOpenAI client
-        client = ChatOpenAI(openai_api_base=self.base_url, api_key=openai_api_key, model=self.model, temperature=0,
-                            logprobs=True)
-
-        # Create the message history for the conversation
-        message_history = [
-            SystemMessage(
-                content="""Your task is to classify whether the document is relevant to the search query provided below. Answer just "YES" or "NO"."""),
-            HumanMessage(content=f"""Document: ```{document}```; Search query: ```{query}```""")
-        ]
-
-        # Generate the response using LangChain's chat model
-        response = await client.agenerate(
-            messages=[message_history],
-            max_tokens=1
-        )
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self._semaphore = asyncio.Semaphore(self.max_concurrent_requests)
+
+    async def _init_client(self):
+        if self.client is None:
+            openai_api_key = self.openai_api_key or os.getenv(self._api_key_var)
+            if not openai_api_key:
+                raise ValueError(f"OpenAI API key not found in environment variable {self._api_key_var}")
+            self.client = AsyncOpenAI(
+                api_key=openai_api_key,
+                base_url=self.base_url,
+                timeout=self.request_timeout,
+                max_retries=2  # Client-level retries
+            )
 
-        # Return the response from the model
-        return response.generations[0]
+    async def search_relevancy(self, query: str, document: str) -> Any:
+        await self._init_client()
+
+        async with self._semaphore:
+            for attempt in range(self.max_retries):
+                try:
+                    response = await self.client.chat.completions.create(
+                        model=self.model,
+                        messages=[
+                            {"role": "system", "content": "Rate the relevance of the document to the query. Respond with 'yes' or 'no'."},
+                            {"role": "user", "content": f"Query: {query}\nDocument: {document}\nIs this document relevant?"}
+                        ],
+                        temperature=self.temperature,
+                        n=1,
+                        logprobs=True,
+                        max_tokens=1
+                    )
+
+                    # Extract response and logprobs
+                    answer = response.choices[0].message.content
+                    logprob = response.choices[0].logprobs.content[0].logprob
+
+                    return {"answer": answer, "logprob": logprob}
+
+                except Exception as e:
+                    if attempt == self.max_retries - 1:
+                        log.error(f"Failed after {self.max_retries} attempts: {str(e)}")
+                        raise
+                    # Exponential backoff with jitter
+                    retry_delay = self.retry_delay * (2 ** attempt) + random.uniform(0, 0.1)
+                    await asyncio.sleep(retry_delay)
 
     async def _rank(self, query_document_pairs: List[Tuple[str, str]]) -> List[Tuple[str, float]]:
-        # Gather results asynchronously for all query-document pairs
-        results = await asyncio.gather(
-            *[self.search_relevancy(query=query, document=document) for (query, document) in query_document_pairs]
-        )
-
         ranked_results = []
 
-        for idx, result in enumerate(results):
-            # Extract the log probability (assuming logprobs are provided in LangChain response)
-            msg = result[0].message
-            logprob = msg.response_metadata['logprobs']['content'][0]['logprob']
-            prob = math.exp(logprob)
-            answer = result[0].message.content  # The model's "YES" or "NO" response
-
-            # Calculate the score based on the model's response
-            if answer == "YES":
-                score = prob
-            elif answer.lower().strip().startswith("y"):
-                score = prob
-            elif answer == "NO":
-                score = 1 - prob
-            elif answer.lower().strip().startswith("n"):
-                score = 1 - prob
-            else:
-                score = 0.0  # Default if something unexpected happens
-
-            # Append the document and score to the result
-            ranked_results.append((query_document_pairs[idx][1], score))  # (document, score)
+        # Process in larger batches for better throughput
+        batch_size = min(self.max_concurrent_requests * 2, len(query_document_pairs))
+        for i in range(0, len(query_document_pairs), batch_size):
+            batch = query_document_pairs[i:i + batch_size]
+            try:
+                results = await asyncio.gather(
+                    *[self.search_relevancy(query=query, document=document) for (query, document) in batch],
+                    return_exceptions=True
+                )
+
+                for idx, result in enumerate(results):
+                    if isinstance(result, Exception):
+                        log.error(f"Error processing document {i+idx}: {str(result)}")
+                        ranked_results.append((batch[idx][1], 0.0))
+                        continue
+
+                    answer = result["answer"]
+                    logprob = result["logprob"]
+                    prob = math.exp(logprob)
+
+                    # Convert answer to score using the model's confidence
+                    if answer.lower().strip() == "yes":
+                        score = prob  # If yes, use the model's confidence
+                    elif answer.lower().strip() == "no":
+                        score = 1 - prob  # If no, invert the confidence
+                    else:
+                        score = 0.5 * prob  # For unclear answers, reduce confidence
+
+                    ranked_results.append((batch[idx][1], score))
+
+                    # Check if we should stop early
+                    high_scoring_docs = [r for r in ranked_results if r[1] >= self.filtering_threshold]
+                    can_stop_early = (
+                        self.early_stop  # Early stopping is enabled
+                        and self.num_docs_to_keep  # We have a target number of docs
+                        and len(high_scoring_docs) >= self.num_docs_to_keep  # Found enough good docs
+                        and score >= self.early_stop_threshold  # Current doc is good enough
+                    )
+
+                    if can_stop_early:
+                        log.info(f"Early stopping after finding {self.num_docs_to_keep} documents with high confidence")
+                        return ranked_results
+
+            except Exception as e:
+                log.error(f"Batch processing error: {str(e)}")
+                continue
 
         return ranked_results
 
-    def compress_documents(
-            self,
-            documents: Sequence[Document],
-            query: str,
-            callbacks: Optional[Callbacks] = None,
+    async def acompress_documents(
+        self,
+        documents: Sequence[Document],
+        query: str,
+        callbacks: Optional[Callbacks] = None,
     ) -> Sequence[Document]:
-        """Compress documents using OpenAI's rerank capability with individual document assessment."""
-        log.info(f"Compressing documents. Initial count: {len(documents)}")
-        if len(documents) == 0:
-            log.warning("No documents to compress. Returning empty list.")
+        """Async compress documents using reranking with proper error handling."""
+        if callbacks:
+            await callbacks.on_retriever_start({"query": query}, "Reranking documents")
+
+        log.info(f"Async compressing documents. Initial count: {len(documents)}")
+        if not documents:
+            if callbacks:
+                await callbacks.on_retriever_end({"documents": []})
             return []
 
-        doc_contents = [doc.page_content for doc in documents]
-        query_documents_pairs = [(query, doc) for doc in doc_contents]
-
-        # Create event loop and run async code
-        import asyncio
         try:
-            rankings = asyncio.get_event_loop().run_until_complete(self._rank(query_documents_pairs))
-        except RuntimeError:
-            # If no event loop is available, create a new one
-            loop = asyncio.new_event_loop()
-            asyncio.set_event_loop(loop)
-            rankings = loop.run_until_complete(self._rank(query_documents_pairs))
-
-        compressed = []
-        for ind, ranking in enumerate(rankings):
-            doc = documents[ind]
-            document_text, score = ranking
-            doc.metadata["relevance_score"] = score
-            doc.metadata["is_relevant"] = score > self.filtering_threshold
-            # Add the document to the compressed list if it is relevant or if we are not removing irrelevant documents
-            if not self.remove_irrelevant:
-                compressed.append(doc)
-            elif doc.metadata["is_relevant"]:
-                compressed.append(doc)
-
-        log.info(f"Compression complete. {len(compressed)} documents returned")
-        if not compressed:
-            log.warning("No documents found after compression")
-
-        if self.num_docs_to_keep is not None:
-            # Sort by relevance score with highest first.
-            compressed.sort(
-                key=lambda d: d.metadata.get('relevance_score', 0) if d.metadata else 0,
-                reverse=True
-            )
-            compressed = compressed[:self.num_docs_to_keep]
-
-        # Handle retrieval callbacks to account for reranked & compressed docs.
-        callbacks = callbacks if callbacks else []
-        run_id = uuid4().hex
-        if not isinstance(callbacks, list):
-            callbacks = callbacks.handlers
-        for callback in callbacks:
-            callback.on_retriever_end(compressed, run_id=run_id)
-        return compressed
+            # Prepare query-document pairs
+            query_document_pairs = [(query, doc.page_content) for doc in documents]
+
+            if callbacks:
+                await callbacks.on_text("Starting document reranking...")
+
+            # Get ranked results
+            ranked_results = await self._rank(query_document_pairs)
+
+            # Sort by score in descending order
+            ranked_results.sort(key=lambda x: x[1], reverse=True)
+
+            # Filter based on threshold and num_docs_to_keep
+            filtered_docs = []
+            for doc, score in ranked_results:
+                if score >= self.filtering_threshold:
+                    matching_doc = next(d for d in documents if d.page_content == doc)
+                    matching_doc.metadata = {**(matching_doc.metadata or {}), "relevance_score": score}
+                    filtered_docs.append(matching_doc)
+
+                    if callbacks:
+                        await callbacks.on_text(f"Document scored {score:.2f}")
+
+                    if self.num_docs_to_keep and len(filtered_docs) >= self.num_docs_to_keep:
+                        break
+
+            log.info(f"Async compression complete. Final count: {len(filtered_docs)}")
+
+            if callbacks:
+                await callbacks.on_retriever_end({"documents": filtered_docs})
+
+            return filtered_docs
+
+        except Exception as e:
+            error_msg = f"Error during async document compression: {str(e)}"
+            log.error(error_msg)
+            if callbacks:
+                await callbacks.on_retriever_error(error_msg)
+            return documents  # Return original documents on error
+
+    def compress_documents(
+        self,
+        documents: Sequence[Document],
+        query: str,
+        callbacks: Optional[Callbacks] = None,
+    ) -> Sequence[Document]:
+        """Sync wrapper for async compression."""
+        return asyncio.run(self.acompress_documents(documents, query, callbacks))
 
     @property
     def _identifying_params(self) -> Dict[str, Any]: