MindsDB 1.3.1__py3-none-any.whl → 25.1.2.0__py3-none-any.whl

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

@@ -0,0 +1,253 @@
+from copy import copy
+from typing import Optional, Any
+
+from langchain_core.output_parsers import StrOutputParser
+from langchain.retrievers import ContextualCompressionRetriever
+
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.runnables import RunnableParallel, RunnablePassthrough, RunnableSerializable
+
+from mindsdb.integrations.handlers.langchain_embedding_handler.langchain_embedding_handler import construct_model_from_args
+from mindsdb.integrations.libs.vectordatabase_handler import DistanceFunction
+from mindsdb.integrations.utilities.rag.chains.map_reduce_summarizer_chain import MapReduceSummarizerChain
+from mindsdb.integrations.utilities.rag.retrievers.auto_retriever import AutoRetriever
+from mindsdb.integrations.utilities.rag.retrievers.multi_vector_retriever import MultiVectorRetriever
+from mindsdb.integrations.utilities.rag.retrievers.sql_retriever import SQLRetriever
+from mindsdb.integrations.utilities.rag.rerankers.reranker_compressor import LLMReranker
+from mindsdb.integrations.utilities.rag.settings import (RAGPipelineModel,
+                                                         DEFAULT_AUTO_META_PROMPT_TEMPLATE,
+                                                         SearchKwargs, SearchType,
+                                                         RerankerConfig,
+                                                         SummarizationConfig, VectorStoreConfig)
+from mindsdb.integrations.utilities.rag.settings import DEFAULT_RERANKER_FLAG
+
+from mindsdb.integrations.utilities.rag.vector_store import VectorStoreOperator
+from mindsdb.interfaces.agents.langchain_agent import create_chat_model
+
+
+class LangChainRAGPipeline:
+    """
+    Builds a RAG pipeline using langchain LCEL components
+    """
+
+    def __init__(
+            self,
+            retriever_runnable,
+            prompt_template,
+            llm,
+            reranker: bool = DEFAULT_RERANKER_FLAG,
+            reranker_config: Optional[RerankerConfig] = None,
+            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
+            )
+        else:
+            self.reranker = None
+        self.summarizer = None
+        self.vector_store_config = vector_store_config
+        knowledge_base_table = self.vector_store_config.kb_table if self.vector_store_config is not None else None
+        if summarization_config is not None and knowledge_base_table is not None:
+            self.summarizer = MapReduceSummarizerChain(
+                vector_store_handler=knowledge_base_table.get_vector_db(),
+                table_name=knowledge_base_table.get_vector_db_table_name(),
+                summarization_config=summarization_config
+            )
+
+    def with_returned_sources(self) -> RunnableSerializable:
+        """
+        Builds a RAG pipeline with returned sources
+        :return:
+        """
+
+        def format_docs(docs):
+            if isinstance(docs, str):
+                # this is to handle the case where the retriever returns a string
+                # instead of a list of documents e.g. SQLRetriever
+                return docs
+            if not docs:
+                return ''
+            # Sort by original document so we can group source summaries together.
+            docs.sort(key=lambda d: d.metadata.get('original_row_id') if d.metadata else 0)
+            original_document_id = None
+            summary_prepended_text = 'Summary of the original document that the below context was taken from:\n'
+            document_content = ''
+            for d in docs:
+                metadata = d.metadata or {}
+                if metadata.get('original_row_id') != original_document_id and metadata.get('summary'):
+                    # We have a summary of a new document to prepend.
+                    original_document_id = metadata.get('original_row_id')
+                    summary = f"{summary_prepended_text}{metadata.get('summary')}\n"
+                    document_content += summary
+                document_content += f'{d.page_content}\n\n'
+            return document_content
+
+        prompt = ChatPromptTemplate.from_template(self.prompt_template)
+
+        # Ensure all the required components are not None
+        if prompt is None:
+            raise ValueError("One of the required components (prompt) is None")
+        if self.llm is None:
+            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
+            )
+
+        rag_chain_from_docs = (
+                RunnablePassthrough.assign(context=(lambda x: format_docs(x["context"])))  # noqa: E126, E122
+                | prompt
+                | self.llm
+                | StrOutputParser()
+        )
+
+        retrieval_chain = RunnableParallel(
+            context=self.retriever_runnable,
+            question=RunnablePassthrough()
+        )
+        if self.summarizer is not None:
+            retrieval_chain = retrieval_chain | self.summarizer
+
+        rag_chain_with_source = retrieval_chain.assign(answer=rag_chain_from_docs)
+        return rag_chain_with_source
+
+    @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"""
+        if hasattr(retriever, 'search_kwargs') and search_kwargs:
+            # Convert search kwargs to dict, excluding None values
+            kwargs_dict = search_kwargs.model_dump(exclude_none=True)
+
+            # Only include relevant parameters based on search type
+            if search_type == SearchType.SIMILARITY:
+                # Remove MMR and similarity threshold specific params
+                kwargs_dict.pop('fetch_k', None)
+                kwargs_dict.pop('lambda_mult', None)
+                kwargs_dict.pop('score_threshold', None)
+            elif search_type == SearchType.MMR:
+                # Remove similarity threshold specific params
+                kwargs_dict.pop('score_threshold', None)
+            elif search_type == SearchType.SIMILARITY_SCORE_THRESHOLD:
+                # Remove MMR specific params
+                kwargs_dict.pop('fetch_k', None)
+                kwargs_dict.pop('lambda_mult', None)
+
+            retriever.search_kwargs.update(kwargs_dict)
+
+            # Set search type if supported by the retriever
+            if hasattr(retriever, 'search_type') and search_type:
+                retriever.search_type = search_type.value
+
+        return retriever
+
+    @classmethod
+    def from_retriever(cls, config: RAGPipelineModel):
+        """
+        Builds a RAG pipeline with returned sources using a simple vector store retriever
+        :param config: RAGPipelineModel
+        :return:
+        """
+        vector_store_operator = VectorStoreOperator(
+            vector_store=config.vector_store,
+            documents=config.documents,
+            embedding_model=config.embedding_model,
+            vector_store_config=config.vector_store_config
+        )
+        retriever = vector_store_operator.vector_store.as_retriever()
+        retriever = cls._apply_search_kwargs(retriever, config.search_kwargs, config.search_type)
+
+        return cls(
+            retriever,
+            config.rag_prompt_template,
+            config.llm,
+            vector_store_config=config.vector_store_config,
+            reranker=config.reranker,
+            reranker_config=config.reranker_config,
+            summarization_config=config.summarization_config
+        )
+
+    @classmethod
+    def from_auto_retriever(cls, config: RAGPipelineModel):
+        if not config.retriever_prompt_template:
+            config.retriever_prompt_template = DEFAULT_AUTO_META_PROMPT_TEMPLATE
+
+        retriever = AutoRetriever(config=config).as_runnable()
+        retriever = cls._apply_search_kwargs(retriever, config.search_kwargs, config.search_type)
+        return cls(
+            retriever,
+            config.rag_prompt_template,
+            config.llm,
+            reranker_config=config.reranker_config,
+            reranker=config.reranker,
+            vector_store_config=config.vector_store_config,
+            summarization_config=config.summarization_config
+        )
+
+    @classmethod
+    def from_multi_vector_retriever(cls, config: RAGPipelineModel):
+        retriever = MultiVectorRetriever(config=config).as_runnable()
+        retriever = cls._apply_search_kwargs(retriever, config.search_kwargs, config.search_type)
+        return cls(
+            retriever,
+            config.rag_prompt_template,
+            config.llm,
+            reranker_config=config.reranker_config,
+            reranker=config.reranker,
+            vector_store_config=config.vector_store_config,
+            summarization_config=config.summarization_config
+        )
+
+    @classmethod
+    def from_sql_retriever(cls, config: RAGPipelineModel):
+        retriever_config = config.sql_retriever_config
+        if retriever_config is None:
+            raise ValueError('Must provide "sql_retriever_config" for RAG pipeline config')
+        vector_store_config = config.vector_store_config
+        knowledge_base_table = vector_store_config.kb_table if vector_store_config is not None else None
+        if knowledge_base_table is None:
+            raise ValueError('Must provide valid "vector_store_config" for RAG pipeline config')
+        embedding_args = knowledge_base_table._kb.embedding_model.learn_args.get('using', {})
+        embeddings = construct_model_from_args(embedding_args)
+        sql_llm = create_chat_model({
+            'model_name': retriever_config.llm_config.model_name,
+            'provider': retriever_config.llm_config.provider,
+            **retriever_config.llm_config.params
+        })
+        retriever = SQLRetriever(
+            vector_store_handler=knowledge_base_table.get_vector_db(),
+            metadata_schemas=retriever_config.metadata_schemas,
+            examples=retriever_config.examples,
+            embeddings_model=embeddings,
+            rewrite_prompt_template=retriever_config.rewrite_prompt_template,
+            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,
+            search_kwargs=config.search_kwargs,
+            llm=sql_llm
+        )
+        return cls(
+            retriever,
+            config.rag_prompt_template,
+            config.llm,
+            reranker_config=config.reranker_config,
+            reranker=config.reranker,
+            vector_store_config=config.vector_store_config,
+            summarization_config=config.summarization_config
+        )

mindsdb/integrations/utilities/rag/rag_pipeline_builder.py

@@ -0,0 +1,83 @@
+import pandas as pd
+from langchain.storage import InMemoryByteStore
+from langchain_core.runnables import RunnableSerializable
+from mindsdb.integrations.utilities.rag.pipelines.rag import LangChainRAGPipeline
+from mindsdb.integrations.utilities.rag.settings import (
+    RetrieverType,
+    RAGPipelineModel
+)
+from mindsdb.integrations.utilities.rag.utils import documents_to_df
+from mindsdb.utilities.log import getLogger
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+
+logger = getLogger(__name__)
+
+_retriever_strategies = {
+    RetrieverType.VECTOR_STORE: lambda config: _create_pipeline_from_vector_store(config),
+    RetrieverType.AUTO: lambda config: _create_pipeline_from_auto_retriever(config),
+    RetrieverType.MULTI: lambda config: _create_pipeline_from_multi_retriever(config),
+    RetrieverType.SQL: lambda config: _create_pipeline_from_sql_retriever(config)
+}
+
+
+def _create_pipeline_from_vector_store(config: RAGPipelineModel) -> LangChainRAGPipeline:
+
+    return LangChainRAGPipeline.from_retriever(
+        config=config
+    )
+
+
+def _create_pipeline_from_auto_retriever(config: RAGPipelineModel) -> LangChainRAGPipeline:
+    return LangChainRAGPipeline.from_auto_retriever(
+        config=config
+    )
+
+
+def _create_pipeline_from_multi_retriever(config: RAGPipelineModel) -> LangChainRAGPipeline:
+
+    if config.text_splitter is None:
+        config.text_splitter = RecursiveCharacterTextSplitter(
+            chunk_size=config.chunk_size, chunk_overlap=config.chunk_overlap
+        )
+    if config.parent_store is None:
+        config.parent_store = InMemoryByteStore()
+
+    return LangChainRAGPipeline.from_multi_vector_retriever(
+        config=config
+    )
+
+
+def _create_pipeline_from_sql_retriever(config: RAGPipelineModel) -> LangChainRAGPipeline:
+    return LangChainRAGPipeline.from_sql_retriever(
+        config=config
+    )
+
+
+def _process_documents_to_df(config: RAGPipelineModel) -> pd.DataFrame:
+    return documents_to_df(config.content_column_name,
+                           config.documents,
+                           embedding_model=config.embedding_model,
+                           with_embeddings=True)
+
+
+def get_pipeline_from_retriever(config: RAGPipelineModel) -> RunnableSerializable:
+    retriever_strategy = _retriever_strategies.get(config.retriever_type)
+    if retriever_strategy:
+        return retriever_strategy(config).with_returned_sources()
+    else:
+        raise ValueError(
+            f'Invalid retriever type, must be one of: {list(_retriever_strategies.keys())}. Got {config.retriever_type}')
+
+
+class RAG:
+    def __init__(self, config: RAGPipelineModel):
+        self.pipeline = get_pipeline_from_retriever(config)
+
+    def __call__(self, question: str) -> dict:
+        logger.info(f"Processing question using rag pipeline: {question}")
+        result = self.pipeline.invoke(question)
+
+        returned_sources = [docs.page_content for docs in result['context']]
+        logger.info(f"retrieved context used to answer question: {returned_sources}")
+
+        return result

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

@@ -0,0 +1,156 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+import math
+import os
+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
+
+
+log = logging.getLogger(__name__)
+
+
+class LLMReranker(BaseDocumentCompressor):
+    filtering_threshold: float = 0.5  # Default threshold for filtering
+    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,
+    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
+
+    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
+        )
+
+        # Return the response from the model
+        return response.generations[0]
+
+    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)
+
+        return ranked_results
+
+    def compress_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.")
+            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
+
+    @property
+    def _identifying_params(self) -> Dict[str, Any]:
+        """Get the identifying parameters."""
+        return {
+            "model": self.model,
+            "temperature": self.temperature,
+            "remove_irrelevant": self.remove_irrelevant,
+        }

mindsdb/integrations/utilities/rag/retrievers/auto_retriever.py

@@ -0,0 +1,117 @@
+from typing import List
+import json
+
+
+from langchain.retrievers.self_query.base import SelfQueryRetriever
+
+import pandas as pd
+
+from mindsdb.integrations.utilities.rag.retrievers.base import BaseRetriever
+
+from mindsdb.integrations.utilities.rag.utils import documents_to_df
+from mindsdb.integrations.utilities.rag.vector_store import VectorStoreOperator
+
+from mindsdb.integrations.utilities.rag.settings import RAGPipelineModel
+
+
+class AutoRetriever(BaseRetriever):
+    """
+    AutoRetrieval is a class that uses langchain to extract metadata from a List of Document and query it using self retrievers.
+
+    """
+
+    def __init__(
+            self,
+            config: RAGPipelineModel
+    ):
+        """
+
+        :param config: RAGPipelineModel
+
+
+        """
+
+        self.documents = config.documents
+        self.content_column_name = config.content_column_name
+        self.vectorstore = config.vector_store
+        self.filter_columns = config.auto_retriever_filter_columns
+        self.document_description = config.dataset_description
+        self.llm = config.llm
+        self.embedding_model = config.embedding_model
+        self.prompt_template = config.retriever_prompt_template
+        self.cardinality_threshold = config.cardinality_threshold
+
+    def _get_low_cardinality_columns(self, data: pd.DataFrame):
+        """
+        Given a dataframe, return a list of columns with low cardinality if datatype is not bool.
+        :return:
+        """
+        low_cardinality_columns = []
+        columns = data.columns if self.filter_columns is None else self.filter_columns
+        for column in columns:
+            if data[column].dtype != "bool":
+                if data[column].nunique() < self.cardinality_threshold:
+                    low_cardinality_columns.append(column)
+        return low_cardinality_columns
+
+    def get_metadata_field_info(self):
+        """
+        Given a list of Document, use llm to extract metadata from it.
+        :return:
+        """
+
+        def _alter_description(data: pd.DataFrame,
+                               low_cardinality_columns: list,
+                               result: List[dict]):
+            """
+            For low cardinality columns, alter the description to include the sorted valid values.
+            :param data: pd.DataFrame
+            :param low_cardinality_columns: list
+            :param result: List[dict]
+            """
+            for column_name in low_cardinality_columns:
+                valid_values = sorted(data[column_name].unique())
+                for entry in result:
+                    if entry["name"] == column_name:
+                        entry["description"] += f". Valid values: {valid_values}"
+
+        data = documents_to_df(
+            self.content_column_name,
+            self.documents
+        )
+
+        prompt = self.prompt_template.format(dataframe=data.head().to_json(),
+                                             description=self.document_description)
+        result: List[dict] = json.loads(self.llm.invoke(input=prompt).content)
+
+        _alter_description(
+            data,
+            self._get_low_cardinality_columns(data),
+            result
+        )
+
+        return result
+
+    def get_vectorstore(self):
+        """
+
+        :return:
+        """
+        return VectorStoreOperator(vector_store=self.vectorstore,
+                                   documents=self.documents,
+                                   embedding_model=self.embedding_model).vector_store
+
+    def as_runnable(self) -> BaseRetriever:
+        """
+        return the self-query retriever
+        :return:
+        """
+        vectorstore = self.get_vectorstore()
+
+        return SelfQueryRetriever.from_llm(
+            llm=self.llm,
+            vectorstore=vectorstore,
+            document_contents=self.document_description,
+            metadata_field_info=self.get_metadata_field_info(),
+            verbose=True
+        )

mindsdb/integrations/utilities/rag/retrievers/base.py

@@ -0,0 +1,10 @@
+from abc import ABC, abstractmethod
+from langchain_core.runnables import RunnableSerializable
+
+
+class BaseRetriever(ABC):
+    """Represents a base retriever for a RAG pipeline"""
+
+    @abstractmethod
+    def as_runnable(self) -> RunnableSerializable:
+        pass