How to display streaming response chunks from Azure OpenAI's LLM in Langchain on Gradio's new Chat-Interface
import os

import gradio as gr
import openai
from langchain.vectorstores import Chroma
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.chat_models import AzureChatOpenAI
from langchain.memory import ConversationBufferWindowMemory

from langchain.callbacks.manager import CallbackManager
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
from langchain.prompts import PromptTemplate
from langchain.chains import ConversationalRetrievalChain

from named_tuples import ConversationalRetrievalChainResponse


openai.api_base = os.getenv("OPENAI_API_BASE")
openai.api_key = os.getenv("OPENAI_API_KEY")
openai.api_type = "azure"
openai.api_version = os.getenv("OPENAI_API_VERSION")

# The LLM
llm = AzureChatOpenAI(
    streaming=True,
    callbacks=CallbackManager([StreamingStdOutCallbackHandler()]),
    temperature=0.0,
    deployment_name=os.getenv("OPENAI_DEPLOYMENT_NAME"),
)

# The Prompt
template = """

    ...

"""
QA_CHAIN_PROMPT = PromptTemplate.from_template(template)


# The Memory
buffer_window_memory = ConversationBufferWindowMemory(
    memory_key="chat_history", return_messages=True, k=4
)

# The Embeddings
openai_embeddings = OpenAIEmbeddings(
    deployment=os.getenv("OPENAI_EMBEDDINGS_MODEL"),
    openai_api_base=os.getenv("OPENAI_API_BASE"),
    openai_api_type="azure",
)

# The Knowledge Base
persist_directory = "vector_database/chromadb/"
vectordb = Chroma(
    persist_directory=persist_directory, embedding_function=openai_embeddings
)

# The Retriever
retriever = vectordb.as_retriever(search_type="mmr")

# The chain
chain = ConversationalRetrievalChain.from_llm(
    llm,
    retriever=retriever,
    memory=buffer_window_memory,
    combine_docs_chain_kwargs={"prompt": QA_CHAIN_PROMPT},
    verbose=True,
)


# The function that takes a message as an argument and returns a streaming LLM resposne
def ask_conversational_retrieval_chain(
    question,
) -> ConversationalRetrievalChainResponse:
    result = chain({"question": question})
    print(result)
    return ConversationalRetrievalChainResponse(result=result["answer"])


# The Gradio App
def respond(message, history):
    return str(ask_conversational_retrieval_chain(message).result)


with gr.Blocks(title="Chatbot") as demo:
    gr.ChatInterface(respond, autofocus=True)

if __name__ == "__main__":
    demo.queue(concurrency_count=5, max_size=20).launch()The function:ask_conversational_retrieval_chaingenerates a streaming response in the form of chunks and I would like to print those chunks as they’re being generated by the LLM directly to the chat interface.I’ve came across a few methods for streaming but those methods wait for the whole response to be generated before giving a “streaming like” response.P.S: I am new to all this, so any help would be appreciated!
2023-08-15T15:22:51Z
Hi@sqy-harshilthe basic idea would be to concatenate those chunks as they come in and simplyyieldthe concatenated chunks instead of returning in this function:def respond(message, history):
    return str(ask_conversational_retrieval_chain(message).result)Here’s a simplified but working example:import time
import gradio as gr

def slow_echo(message, history):
    for i in range(len(message)):
        time.sleep(0.05)
        yield "You typed: " + message[: i+1]

demo = gr.ChatInterface(slow_echo).queue()

if __name__ == "__main__":
    demo.launch()
2023-08-18T08:06:51Z
Hi@sqy-harshil, just wondering if you managed to solve it? My workflow is similar - using ConversationalRetrievalChain, but I’m unable to find a way to stream the response.
2023-08-21T10:59:52Z
I also tried streaming streaming the responses through the yield chunk method, but it actually first collects the entire response then starts sending in chunks through yield, which is actually worse. I tried many other things as well, now I am going to approach this with threads.@abidlabsDo you have any working example with openai llm and langchain in python? It would help a lot. Thanks.
2023-08-21T13:23:30Z
Hi@chitiiiyes, please take a look here:Creating A Chatbot Fast
2023-08-30T04:28:09Z
@DesmondChoyHi,I was able to resolve it using Langchain’s BaseCallbackHandler in streamlit,streaming.py:from langchain.callbacks.base import BaseCallbackHandler


class StreamHandler(BaseCallbackHandler):
    def __init__(self, container, initial_text=""):
        self.container = container
        self.text = initial_text

    def on_llm_new_token(self, token: str, **kwargs):
        self.text += token
        self.container.markdown(self.text)app.py:from streaming import StreamHandler

# The logic for generating the ConversationalRetrievalChain.from_llm...

callback = StreamHandler(st.empty())
response = qa_chain.run(user_input, callbacks=callback)

# Rest of the codeThis should give you a streaming response in chunks as the LLM is generating them.You could either implement this in Gradio as per your use case or refer toGradio’s official documentation for how to achieve a streaming response from LLMs