Issue with RetrievalQAWithSourcesChain Deployment in Databricks using MLflow

[UnnamalaiNarayanan]

Checked other resources

• I added a very descriptive title to this question.
• I searched the LangChain documentation with the integrated search.
• I used the GitHub search to find a similar question and didn't find it.

Commit to Help

• I commit to help with one of those options 👆

Example Code

from typing import Any, List, Dict, Union
from langchain.callbacks.manager import CallbackManagerForChainRun
from langchain.chains import RetrievalQA
from langchain.docstore.document import Document
from langchain.schema import Document
from langchain.vectorstores.base import VectorStoreRetriever
import inspect

class RetrievalQAFilter(RetrievalQAWithSourcesChain):
    retriever: BaseRetriever = Field(exclude=False)

    #initializingVectorRetriever
    # Function to create and return the configured retriever
    def get_retriever(persist_dir: str = None):
        
        # Initialize Vector Search Client using Databricks credentials
        vs_client = VectorSearchClient(
            workspace_url=host,
            service_principal_client_id=os.environ["DATABRICKS_CLIENT_ID"],
            service_principal_client_secret=os.environ["DATABRICKS_CLIENT_SECRET"],
            disable_notice=True
        )

        # Get the vector search index
        vs_index = vs_client.get_index(
            endpoint_name=VECTOR_SEARCH_ENDPOINT_NAME,
            index_name=index_name
        )

        # Initialize vector search as retriever
        vector_search_as_retriever = DatabricksVectorSearch(
            vs_index,
            text_column=retriever_text_column,
            columns=retriever_vectorstore_columns
        )

        # Configure the retriever with filtering and hybrid query capabilities
        retriever = VectorStoreRetrieverFilter(
            vectorstore=vector_search_as_retriever, 
            search_type=retriever_search_type,
            query_type="hybrid"
        )

        return retriever

    def _get_docs(
        self, 
        question: str, 
        filter: dict[str, str] = None, 
        *, 
        run_manager: CallbackManagerForChainRun
    ) -> list[Document]:
        self.retriever = self.get_retriever()
        if filter is not None:
            return self.retriever._get_relevant_documents(question, filter)
        return self.retriever._get_relevant_documents(question)

class VectorStoreRetrieverFilter(VectorStoreRetriever):
    query_type: str
    def _get_relevant_documents(self, query: str, filter: Dict[str, Any]=None)-> List[Document]:
        """
        Retrieve relevant documents using similarity search, optionally applying a filter.
        """
        question = query.get("question", "")
        filter = query.get("filter", None)
        query = question

        print('Query: ',query)
        print('Filter: ',filter)

        # Number of documents to retrieve
        k = model_config.get("retriever_config").get("k")

        # If a filter is provided, apply it to the search
        if filter is not None:
            new_filter = {k: v for k, v in filter.items() if v}  # Remove empty filter values
            if self.search_type == "similarity":
                if self.query_type == "hybrid":
                    docs = self.vectorstore.similarity_search(
                        query, query_type="hybrid", filter=new_filter, k=k, **self.search_kwargs
                    )
                elif self.query_type == "ann":
                    docs = self.vectorstore.similarity_search(
                        query, filter=new_filter, k=k, **self.search_kwargs
                    )
        else:
            # If no filter, perform standard similarity search
            if self.search_type == "similarity":
                if self.query_type == "hybrid":
                    docs = self.vectorstore.similarity_search(
                        query, k=k, query_type="hybrid", **self.search_kwargs
                    )
                elif self.query_type == "ann":
                    docs = self.vectorstore.similarity_search(query, k=k, **self.search_kwargs)

        return docs

chain = RetrievalQAFilter.from_chain_type(
    llm=qa_model,  
    retriever=vector_search_as_retriever, 
    chain_type="stuff",  
    return_source_documents=True,  
    chain_type_kwargs={
        "prompt": PromptTemplate(
            template=model_config.get("llm_config").get("llm_prompt_template"),
            input_variables=["context", "chat_history", "question"]
        ),
        "document_variable_name": "context" 
    },
)

with mlflow.start_run(run_name = "llama_model_run") as run:
    signature = infer_signature(input,response["answer"])
    model_info = mlflow.langchain.log_model(
        chain,
        loader_fn=RetrievalQAFilter.get_retriever,
        artifact_path = "chain",
        registered_model_name=model_name,
        pip_requirements = [
            "mlflow==" + mlflow.__version__,
            "langchain==0.3.1",
             "databricks-vectorsearch",
             "langchain_core==0.3.6" ,
             "langchain_community==0.3.1", 
             "databricks-sdk"
        ],
        input_example=input,
        signature=signature
    )

Description

We are building an application using RetrievalQAWithSourcesChain to extract information from PDFs and return the relevant source documents used for generating responses. To support filtering, we developed a custom class (RetrievalQAFilter) that overrides the functionality of RetrievalQAWithSourcesChain, based on the guidance from this GitHub issue.

However, we are encountering an error when deploying the model as a serving endpoint in Databricks using MLflow:
Error: retriever must be present.

We would appreciate any insights or suggestions to resolve this. Are we missing any critical steps in configuring or serving the model? Please let us know if further information is needed.
Thank you for your support!

System Info

langchain==0.3.1 langchain_core==0.3.6 langchain_community==0.3.1 databricks-vectorsearch