Add the ability to run the map_reduce chains process results step as async

[MalhotraVijay]

This will add the ability to add an AsyncCallbackManager (handler) for the reducer chain, which would be able to stream the tokens via the async def on_llm_new_token callback method

Fixes # (issue) 5532

@hwchase17 @agola11
The following code snippet explains how this change would be used to enable reduce_llm with streaming support in a map_reduce chain

I have tested this change and it works for the streaming use-case of reducer responses. I am happy to share more information if this makes solution sense.

AsyncHandler
..........................
class StreamingLLMCallbackHandler(AsyncCallbackHandler):
    """Callback handler for streaming LLM responses."""

    def __init__(self, websocket):
        self.websocket = websocket
    
    # This callback method is to be executed in async
    async def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
        resp = ChatResponse(sender="bot", message=token, type="stream")
        await self.websocket.send_json(resp.dict())


Chain
..........
stream_handler = StreamingLLMCallbackHandler(websocket)
stream_manager = AsyncCallbackManager([stream_handler])

streaming_llm = ChatOpenAI(
        streaming=True,
        callback_manager=stream_manager,
        verbose=False,
        temperature=0,
    )
    main_llm = OpenAI(
        temperature=0,
        verbose=False,
    )

    doc_chain = load_qa_chain(
        llm=main_llm,
        reduce_llm=streaming_llm,
        chain_type="map_reduce", 
        callback_manager=manager
    )
    qa_chain = ConversationalRetrievalChain(
        retriever=vectorstore.as_retriever(),
        combine_docs_chain=doc_chain,
        question_generator=question_generator,
        callback_manager=manager,
    )
    
    # Here `acall` will trigger `acombine_docs` on `map_reduce` which should then call `_aprocess_result` which in turn will call `self.combine_document_chain.arun` hence async callback will be awaited
    result = await qa_chain.acall(
         {"question": question, "chat_history": chat_history}
      )

[agola11]

Thanks!

[MalhotraVijay]

Hi @agola11 - Thank you for reviewing earlier. I have added a few lint fixes. Kindly review that when you get a chance. Thanks.

[MalhotraVijay]

Hi @agola11 - Friendly reminder, if you could check and approve the PR again, as I've added some linting fixes which was causing the automation to fail. Thanks.