Streaming for MultiRetrievalQAChain on langchain JS

[alexgitENI]

Hello everyone, I'm developing a chatbot that uses the MultiRetrievalQAChain function to provide the most appropriate response. I've managed to get it to work in "normal" mode`

I now want to switch to stream mode to improve response time, the problem is that all intermediate actions are streamed, I only want to stream the last response and not all intermediate responses.

here is my working synchronous code:

const model = new ChatOpenAI({ modelName: "gpt-3.5-turbo",verbose:true,});


const retrieverNames = ["retrieverVols", "retriever"];
const retrieverDescriptions = [
  "Good for answering questions on the theme of Latvia",
  "Good for answering questions about the shopping list",
];


const retrievers = [
  vectorStoreVols.asRetriever(3),
  vectorStore.asRetriever(3)

];

const multiRetrievalQAChain = MultiRetrievalQAChain.fromLLMAndRetrievers(model, {
  retrieverNames,
  retrieverDescriptions,
  retrievers,
  retrievalQAChainOpts: {
    returnSourceDocuments: true,
  },
});

multiRetrievalQAChain.defaultChain.prompt= prompt

const testPromise1 = await multiRetrievalQAChain.call({
  input:
    "How tall is the Eiffeil Tower?",
});

  console.log(testPromise1);

here my streamcode that not work

const model = new ChatOpenAI({ modelName: "gpt-3.5-turbo",verbose:true,
streaming: true,
maxTokens: 200,
callbackManager: CallbackManager.fromHandlers({
  async handleLLMNewToken(token) {
    console.log({ token });
  },
  async handleLLMEnd(output) {
    console.log("End of stream.", output);
  },
}),
});


const retrieverNames = ["retrieverVols", "retriever"];
const retrieverDescriptions = [
  "Good for answering questions on the theme of Latvia",
  "Good for answering questions about the shopping list",
];


const retrievers = [
  vectorStoreVols.asRetriever(3),
  vectorStore.asRetriever(3)

];

const multiRetrievalQAChain = MultiRetrievalQAChain.fromLLMAndRetrievers(model, {
  retrieverNames,
  retrieverDescriptions,
  retrievers,


  retrievalQAChainOpts: {
    returnSourceDocuments: true,
  },
});

multiRetrievalQAChain.defaultChain.prompt= prompt

const testPromise1 = await multiRetrievalQAChain.call({
  input:
    "How tall is the Eiffeil Tower?",
});

  console.log(testPromise1);

here is the console output for the stream

[llm/start] [1:llm:ChatOpenAI] Entering LLM run with input: {
  "messages": [
    [
      {
        "lc": 1,
        "type": "constructor",
        "id": [
          "langchain",
          "schema",
          "HumanMessage"
        ],
        "kwargs": {
          "content": "Given a query to a question answering system, select the system best suited for the input. You will be given the names of the available systems and a description of what questions the system is best suited for. You may also revise the original input if you think that revising it will ultimately lead to a better response.\n\n<< FORMATTING >>\nReturn a markdown code snippet with a JSON object formatted to look like:\n```json\n{\n  \"destination\": string // name of the question answering system to use or \"DEFAULT\" (optional)\n  \"next_inputs\": {\n    \"query\": string // a potentially modified version of the original input\n  } // input to be fed to the next model\n}\n```\n\nREMEMBER: \"destination\" MUST be one of the candidate prompt names specified below OR it can be \"DEFAULT\" if the input is not well suited for any of the candidate prompts.\nREMEMBER: \"next_inputs.query\" can just be the original input if you don't think any modifications are needed.\n\n<< CANDIDATE PROMPTS >>\nretrieverVols: Good for answering questions on the theme of Latvia\nretriever: Good for answering questions about the shopping list\n\n<< INPUT >>\nHow tall is the Eiffeil Tower?\n\n<< OUTPUT >>\n",
          "additional_kwargs": {}
        }
      }
    ]
  ]
}
{ token: '' }
{ token: '{\n' }
{ token: ' ' }
{ token: ' "' }
{ token: 'destination' }
{ token: '":' }
{ token: ' "' }
{ token: 'DEFAULT' }
{ token: '",\n' }
{ token: ' ' }
{ token: ' "' }
{ token: 'next' }
{ token: '_inputs' }
{ token: '":' }
{ token: ' {\n' }
{ token: '   ' }
{ token: ' "' }
{ token: 'query' }
{ token: '":' }
{ token: ' "' }
{ token: 'How' }
{ token: ' tall' }
{ token: ' is' }
{ token: ' the' }
{ token: ' E' }
{ token: 'iff' }
{ token: 'el' }
{ token: ' Tower' }
{ token: '?"\n' }
{ token: ' ' }
{ token: ' }\n' }
{ token: '}' }
{ token: '' }
End of stream. {
  generations: [ [ [ChatGenerationChunk] ] ],
  llmOutput: {
    estimatedTokenUsage: { promptTokens: 269, completionTokens: 31, totalTokens: 300 }
  }
}
[llm/end] [1:llm:ChatOpenAI] [4.92s] Exiting LLM run with output: {
  "generations": [
    [
      {
        "text": "{\n  \"destination\": \"DEFAULT\",\n  \"next_inputs\": {\n    \"query\": \"How tall is the Eiffel Tower?\"\n  }\n}",
        "generationInfo": {
          "prompt": 0,
          "completion": 0
        },
        "message": {
          "lc": 1,
          "type": "constructor",
          "id": [
            "langchain",
            "schema",
            "AIMessageChunk"
          ],
          "kwargs": {
            "content": "{\n  \"destination\": \"DEFAULT\",\n  \"next_inputs\": {\n    \"query\": \"How tall is the Eiffel Tower?\"\n  }\n}",
            "additional_kwargs": {}
          }
        }
      }
    ]
  ],
  "llmOutput": {
    "estimatedTokenUsage": {
      "promptTokens": 269,
      "completionTokens": 31,
      "totalTokens": 300
    }
  }
}
[llm/start] [1:llm:ChatOpenAI] Entering LLM run with input: {
  "messages": [
    [
      {
        "lc": 1,
        "type": "constructor",
        "id": [
          "langchain",
          "schema",
          "HumanMessage"
        ],
        "kwargs": {
          "content": "Reply: I'm an IA specializing in travel, so I can't answer your question!",
          "additional_kwargs": {}
        }
      }
    ]
  ]
}
{ token: '' }
{ token: 'That' }
{ token: "'s" }
{ token: ' alright' }
{ token: '!' }
{ token: ' If' }
{ token: ' you' }
{ token: ' have' }
{ token: ' any' }
{ token: ' other' }
{ token: ' questions' }
{ token: ' related' }
{ token: ' to' }
{ token: ' travel' }
{ token: ',' }
{ token: ' feel' }
{ token: ' free' }
{ token: ' to' }
{ token: ' ask' }
{ token: ',' }
{ token: ' and' }
{ token: ' I' }
{ token: "'ll" }
{ token: ' be' }
{ token: ' happy' }
{ token: ' to' }
{ token: ' help' }
{ token: ' you' }
{ token: ' out' }
{ token: '.' }
{ token: '' }
End of stream. {
  generations: [ [ [ChatGenerationChunk] ] ],
  llmOutput: {
    estimatedTokenUsage: { promptTokens: 25, completionTokens: 29, totalTokens: 54 }
  }
}
[llm/end] [1:llm:ChatOpenAI] [4.55s] Exiting LLM run with output: {
  "generations": [
    [
      {
        "text": "That's alright! If you have any other questions related to travel, feel free to ask, and I'll be happy to help you out.",
        "generationInfo": {
          "prompt": 0,
          "completion": 0
        },
        "message": {
          "lc": 1,
          "type": "constructor",
          "id": [
            "langchain",
            "schema",
            "AIMessageChunk"
          ],
          "kwargs": {
            "content": "That's alright! If you have any other questions related to travel, feel free to ask, and I'll be happy to help you out.",
            "additional_kwargs": {}
          }
        }
      }
    ]
  ],
  "llmOutput": {
    "estimatedTokenUsage": {
      "promptTokens": 25,
      "completionTokens": 29,
      "totalTokens": 54
    }
  }
}
{
  result: "That's alright! If you have any other questions related to travel, feel free to ask, and I'll be happy to help you out."
}

Does anyone have any clues?
Many thanks to anyone who will take the time to answer my questions.

[dosubot[bot]]

🤖

Based on the context provided, it seems like you want to stream only the final response from the MultiRetrievalQAChain function, instead of all the intermediate responses. This can be achieved by using the 'end' event of the stream, which is emitted when there is no more data to be consumed from the stream.

Here's how you can modify your code:

// Create a retrieval QA chain that uses a Llama 2-powered QA stuff chain with a custom prompt.
const chain = new RetrievalQAChain({
  combineDocumentsChain: loadQAStuffChain(model, { prompt: QA_CHAIN_PROMPT }),
  retriever,
  returnSourceDocuments: true,
  inputKey: "question",
});

const stream = chain.call({
  question: "What are the approaches to Task Decomposition?",
});

let finalResponse;
stream.on('data', (data) => {
  finalResponse = data;
});

stream.on('end', () => {
  console.log({ finalResponse });
});

In this code, we're creating a stream from the call to the chain. We then listen for 'data' events on the stream, and each time we receive data, we update our finalResponse variable with the new data. Finally, when the 'end' event is emitted, we log the final response. This way, we're only logging the final response, not all the intermediate responses.

This solution was suggested in a similar issue titled Stream only the final answer when using an agent in the langchainjs repository.

I hope this helps! If you have any other questions, feel free to ask.

Sources

• Stream only the final answer when using an agent
• examples/src/chains/retrieval_qa_legacy.ts
• examples/src/use_cases/local_retrieval_qa/qa_chain.ts
• langchain/src/chains/router/multi_retrieval_qa.ts

---

This response is meant to be useful and save you time. It is not meant to be a precise solution, but rather a starting point for your own research.

Help me be more useful! Please leave a 👍 if this is helpful and 👎 if it is irrelevant.

If you want to continue the conversation, start your reply with @dosu-bot.

[alexgitENI]

@dosu-bot. Your example contains a RetrievalQAChain, I want to use the MultiRetrievalQAChain which is a set of several queries performed.

I saw that the python library has a callback named FinalStreamingStdOutCallbackHandler. Do you know of anything similar for langchainJS?

[dosubot[bot]]

🤖

Based on the information you've provided and the context from the langchainjs repository, there is indeed a workaround to only stream the final response when using the MultiRetrievalQAChain function in stream mode. You can modify the callback function that handles the stream to only log or process the final chunk of the stream. Here's how you can do it:

import { ChatOpenAI } from "langchain/chat_models/openai";
import { HumanMessage } from "langchain/schema";

const chat = new ChatOpenAI({
  maxTokens: 25,
  streaming: true,
});

let finalToken = "";
const response = await chat.call([new HumanMessage("Tell me a joke.")], {
  callbacks: [
    {
      handleLLMNewToken(token: string) {
        finalToken = token;
      },
    },
  ],
});

console.log(finalToken);

In this modified code, instead of logging each token, we replace finalToken with the latest token in each callback. This way, finalToken will always hold the last token when the callback finishes, which is the final response.

This solution was suggested in a similar issue titled Stream only the final answer when using an agent and you can find more details about streaming in the LangChainJS documentation.

As for the MultiRetrievalQAChain function, it is a class in the LangChainJS framework that represents a multi-retrieval question answering chain. It extends the MultiRouteChain class and provides additional functionality specific to multi-retrieval QA chains. You can find more details about this class in the LangChainJS source code and its test file.

I hope this helps! If you have any other questions or need further clarification, feel free to ask.

Sources

• Stream only the final answer when using an agent
• examples/src/chains/conversational_qa_streaming.ts
• examples/src/models/chat/chat_streaming.ts
• langchain/src/chains/router/multi_retrieval_qa.ts
• langchain/src/chains/router/tests/multi_retrieval_qa.int.test.ts

---

This response is meant to be useful and save you time. It is not meant to be a precise solution, but rather a starting point for your own research.

Help me be more useful! Please leave a 👍 if this is helpful and 👎 if it is irrelevant.

If you want to continue the conversation, start your reply with @dosu-bot.