Integrating document search with RunnableWithMessageHistory in langchain

Other Questions langchain
1

I’m working with langchain js and trying to figure out how to add document context to my chat conversation. I have a retriever set up but can’t connect it properly to my chat history setup.

import { Ollama, OllamaEmbeddings } from '@langchain/ollama';
import { ChatPromptTemplate, MessagesPlaceholder } from '@langchain/core/prompts';
import { RunnableConfig, RunnableWithMessageHistory } from '@langchain/core/runnables';
import { ChatMessageHistory } from '@langchain/community/stores/message/in_memory';
import { PDFLoader } from '@langchain/community/document_loaders/fs/pdf';
import { RecursiveCharacterTextSplitter } from 'langchain/text_splitter';
import { MemoryVectorStore } from 'langchain/vectorstores/memory';

const docLoader = new PDFLoader('./document.pdf');
const llm = new Ollama({ model: 'mistral' });

const chatHistory = new ChatMessageHistory();

const template = ChatPromptTemplate.fromMessages([
    new MessagesPlaceholder('chat_history'),
    [ 'user', '{question}' ]
]);

const settings: RunnableConfig = { configurable: { sessionId: 'chat-1' } };

(async () => {
    const documents = await docLoader.load();
    
    const splitter = new RecursiveCharacterTextSplitter({
        chunkSize: 800,
        chunkOverlap: 150
    });
    const chunks = await splitter.splitDocuments(documents);
    const vectorDB = await MemoryVectorStore.fromDocuments(
        chunks,
        new OllamaEmbeddings()
    );
    const documentRetriever = vectorDB.asRetriever(); // How do I use this?
    
    const chain = template.pipe(llm);
    
    const withChatHistory = new RunnableWithMessageHistory({
        runnable: chain,
        getMessageHistory: (_id: string) => chatHistory,
        inputMessagesKey: 'question',
        historyMessagesKey: 'chat_history'
    });
    
    let response = await withChatHistory.invoke(
        { question: 'Hi, my name is Bob!' },
        settings
    );
    
    console.log(response);
    
    response = await withChatHistory.invoke(
        { question: 'Tell me about the document content' },
        settings
    );
    
    console.log(response);
})()

I want the AI to automatically search through my documents when needed without me having to manually pass context each time. How can I connect the retriever to work with the message history wrapper?

2

You need to add the retrieval step before your LLM call. Don’t pipe your template straight to the LLM - create a function that grabs the relevant docs first. I built a custom runnable sequence that takes the user’s question, runs it through the retriever, formats those results as context, then feeds everything to the chat template. Just add a context field to your prompt and make sure retrieval happens inside the chain that RunnableWithMessageHistory wraps. That way document search becomes part of the conversation instead of something you handle separately.

3

create a retrieval chain first, then wrap it with message history. use createRetrievalChain from langchain/chains/retrieval - it’ll automatically handle document fetching when users ask questions.

4

Modify your prompt template to add a context placeholder, then build a custom runnable that grabs documents before hitting your LLM. Chain your retriever with the chat setup using RunnablePassthrough and RunnableLambda. Here’s what worked for me: throw {context} into your system message, then write a retrieval function that takes the user’s question, searches your vector store, and formats what it finds. Chain this retrieval step before your existing template pipe - now documents get pulled automatically every time you invoke. Wrap the whole chain (including retrieval) with RunnableWithMessageHistory. Your chat history stays intact while documents get fetched seamlessly based on whatever the user just asked.

5

I hit this exact problem last month building a doc chat system for our internal knowledge base. The trick is using RunnableSequence to chain everything together properly.

Update your prompt template first:

const template = ChatPromptTemplate.fromMessages([
    ['system', 'Use the following context to answer questions: {context}'],
    new MessagesPlaceholder('chat_history'),
    ['user', '{question}']
]);

Then build a sequence that handles retrieval automatically:

import { RunnableSequence, RunnablePassthrough } from '@langchain/core/runnables';

const retrievalChain = RunnableSequence.from([
    {
        context: (input) => documentRetriever.getRelevantDocuments(input.question),
        question: new RunnablePassthrough(),
        chat_history: new RunnablePassthrough()
    },
    {
        context: (input) => input.context.map(doc => doc.pageContent).join('\n\n'),
        question: (input) => input.question,
        chat_history: (input) => input.chat_history
    },
    template,
    llm
]);

Now wrap this retrieval chain instead of your basic template chain:

const withChatHistory = new RunnableWithMessageHistory({
    runnable: retrievalChain,
    getMessageHistory: (_id: string) => chatHistory,
    inputMessagesKey: 'question',
    historyMessagesKey: 'chat_history'
});

Every question automatically triggers document search before hitting the LLM. Works like a charm and keeps your chat history intact.