How do you change the source field in RetrievalQAWithSourcesChain

[jordanparker6]

I have a VectorDB with sources saved as doucment_id. How do I use the document_id as the source instead of using a source key in the RetrievalQAWithSourcesChain?

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
Cell In[12], line 1
----> 1 chain({ "question": "What is hello" }, return_only_outputs=True)

File ~/Programs/search-service/.venv/lib/python3.10/site-packages/langchain/chains/base.py:116, in Chain.__call__(self, inputs, return_only_outputs)
    114 except (KeyboardInterrupt, Exception) as e:
    115     self.callback_manager.on_chain_error(e, verbose=self.verbose)
--> 116     raise e
    117 self.callback_manager.on_chain_end(outputs, verbose=self.verbose)
    118 return self.prep_outputs(inputs, outputs, return_only_outputs)

File ~/Programs/search-service/.venv/lib/python3.10/site-packages/langchain/chains/base.py:113, in Chain.__call__(self, inputs, return_only_outputs)
    107 self.callback_manager.on_chain_start(
    108     {"name": self.__class__.__name__},
    109     inputs,
    110     verbose=self.verbose,
    111 )
    112 try:
--> 113     outputs = self._call(inputs)
    114 except (KeyboardInterrupt, Exception) as e:
    115     self.callback_manager.on_chain_error(e, verbose=self.verbose)

File ~/Programs/search-service/.venv/lib/python3.10/site-packages/langchain/chains/qa_with_sources/base.py:118, in BaseQAWithSourcesChain._call(self, inputs)
    116 def _call(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
...
     66     }
     67     doc_dicts.append(document_info)
     68 # Format each document according to the prompt

KeyError: 'source'

[mswoff]

I ran into this issue when trying to switch over some code I had written to use langchain instead of my custom scripts. Specifically, RetrievalQAWithSourcesChain.from_chain_type(llm=chat_model, chain_type='stuff', retriever=vectorstore.as_retriever())(<some question here>) would error because I had setup my Pinecone vector store to use metadata['path'] rather than metadata['source'].

My workaround for this was to create a custom RetrievalQAWithSourcesChain which modifies the _get_docs() to do this step.

    class CustomRetrievalQAWithSourcesChain(RetrievalQAWithSourcesChain):
        def _get_docs(self, inputs: Dict[str, Any]) -> List[Document]:
            # Call the parent class's method to get the documents
            docs = super()._get_docs(inputs)
            # Modify the document metadata
            for doc in docs:
                doc.metadata['source'] = doc.metadata.pop('path')
            return docs

    qa_chain = CustomRetrievalQAWithSourcesChain.from_chain_type(llm=chat_model, chain_type='stuff', retriever=vectorstore.as_retriever())

I would love to know if there is a better more langchainy way to add this step into the process (rather than a custom class). Or does this functionality need to be added?