Custom Retriever with vectorstore

[chamarthi-ks]

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

class CustomRetriever(BaseRetriever):

    vectorstore: VectorStoreRetriever
    search_type: str = "similarity"
    search_kwargs: dict = Field(default_factory=dict)


    def _get_relevant_documents(self, query: str, *, run_manager: Optional[CallbackManagerForRetrieverRun] = None):

        # Perform similarity search (retrieve top 3 most relevant documents)
        similar_docs = self.vectorstore.similarity_search_with_relevance_scores(query, k=10)

        unique_file = []
        for header in similar_docs:
            if header[0].metadata['Header 1']=='Title letter':
                unique_file.append(header[0].metadata.get('source'))
                print(header[0].metadata.get('source'))



        custom_chunks = []
        headers_to_split_on = [("#", "Header 1")]
        markdown_splitter = MarkdownHeaderTextSplitter(headers_to_split_on=headers_to_split_on, strip_headers=False)

        for location in unique_file:
            # loading the markdown
            loader = TextLoader(location,encoding="utf-8")
            doc = loader.load()
            print(location)
            md_header_splits = markdown_splitter.split_text(doc[0].page_content)
            # search for the required sections of the document
            print(md_header_splits[0].metadata)
            for split in md_header_splits:
                if split.metadata.get('Header 1') in ['Title letter','Supporting Document']:

                    custom_chunks.append([split.page_content,split.metadata.get('Header 1'),location])


        documents = [
            Document(page_content=custom_chunk[0], metadata={"header": custom_chunk[1], "source": custom_chunk[2]})
            for custom_chunk in custom_chunks]

        
        return documents

memory = ConversationBufferMemory(memory_key='chat_history', return_messages=True, output_key='answer')
qachat = ConversationalRetrievalChain.from_llm(
    llm=model,
    memory=memory,
    retriever=CustomRetriever,
    return_source_documents=True,
    # condense_question_prompt=condense_question_prompt,
    chain_type="stuff",
    combine_docs_chain_kwargs={"prompt": prompt_template_answer}
    )

Description

The objective is to give the custom context to the llm to answer the question using ConversationalRetrievalChain.from_llm wrapper

I have created a vector store, for a given question to get similar documents, using similarity search to find the similar chunks. Now here is the challenge, i need to search for source of the chunks and read the files to specific content as context. Because i will get high similarity search result on the Title letter, i want the context to be Supporting Document as a context for that title letter. It is not possible to get the chunk of the Supporting Document using similarity search.

System Info

langchain==0.3.18
langchain-community==0.3.5
langchain-core==0.3.34
langchain-openai==0.2.6
langchain-text-splitters==0.3.6
langchain-unstructured==0.1.5
langchain-upstage==0.3.0