Streamlit, Langchain and VertexAI integration

[AnhNgDo]

Issues with getting Vertex AI models to work with Streamlit callbacks

Hi, I've got a Streamlit app that can switch between OpenAI and VertexAI models.

One little issue was the Streamlit integrations (StreamlitCallbackHandler) - VertexAI models does not output into the Streamlit container as expected. OpenAI models seems to work fine.

My work around is to manually call st.write if using VertexAI - just wondering if there is a way to make the Streamlit callback work?

Many thanks!

Here is extract for the code:

...
prompt_default = ChatPromptTemplate.from_messages([
    SystemMessagePromptTemplate.from_template(system_tempate_default),
    HumanMessagePromptTemplate.from_template(human_template_default),
])    

if st.session_state["llm_provider"] == "OpenAI":
    llm_default = ChatOpenAI(model="gpt-3.5-turbo", temperature=0.6, streaming=True)
elif st.session_state["llm_provider"] == "VertexAI":
    llm_default = VertexAI(temperature=0.4, max_output_tokens=512)

chain_default = LLMChain(llm=llm_default, prompt=prompt_default, memory=memory, verbose=True)
...

if query := st.chat_input("Ask me any thing..."):
    # Display user message in chat message container
    with st.chat_message("user"):
        st.write(query)
    # Display assistant response in chat message container
    with st.chat_message("assistant"):
        st_callback = StreamlitCallbackHandler(st.container()) # streamlit container for callback's output
        message_placeholder = st.empty()
        stream_handler = StreamHandler(message_placeholder) # streaming llm output to streamlit
         ........
         ......
        convo_state = chain_analyzer.run({"learner_name":user, "chat_extract": chat_extract_str, "input": query}, 
                                         callbacks=[langfuse_handler, st_callback])
        
        # Write convo_state to container - Vertex AI only (does not work with st_callback)
        if (st.session_state["llm_provider"] == "VertexAI"):
            with st.container():
                st.write(convo_state)
        ........
        .....
            response = chain_default.run({"learner_name": st.session_state["user"], "input": query}, 
                                         callbacks=[stream_handler, langfuse_handler])
            # Vertex AI only (does not work with st stream_handler)
            if (st.session_state["llm_provider"] == "VertexAI"):
                st.write(response)

[dosubot[bot]]

🤖

Hi there,

Thanks for reaching out with your question about integrating VertexAI models with Streamlit callbacks in the LangChain framework.

From your code, it seems like you're trying to use the StreamlitCallbackHandler with the VertexAI model. However, the StreamlitCallbackHandler might not be working as expected because the _stream method in the VertexAI class does not currently include the callback manager.

To make the Streamlit callback work with VertexAI models, you need to modify the _stream method in the VertexAI class to include the callback manager. Here's an example of how you might modify the _stream method:

def _stream(
        self,
        prompt: str,
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> Iterator[GenerationChunk]:
        params = self._prepare_params(stop=stop, **kwargs)
        for stream_resp in stream_completion_with_retry(
            self, prompt, run_manager=run_manager, **params
        ):
            chunk = _response_to_generation(stream_resp)
            yield chunk
            if run_manager:
                run_manager.on_llm_new_token(
                    chunk.text,
                    chunk=chunk,
                    verbose=self.verbose,
                )

In this method, the run_manager parameter is an instance of CallbackManagerForLLMRun, which is used to manage callbacks during the run of a large language model. However, it is not currently being used to manage callbacks for the VertexAI model.

To make the Streamlit callback work with VertexAI models, you need to modify the _stream method to call the appropriate methods on the run_manager object at the appropriate times. For example, you might need to call run_manager.on_llm_new_token(chunk.text, chunk=chunk, verbose=self.verbose) after each new token is generated by the model, and run_manager.on_llm_end() after the model has finished generating tokens.

Please note that the exact modifications you need to make will depend on the specifics of how the Streamlit callback and the VertexAI model are supposed to interact.

I hope this helps! Let me know if you have any other questions.

Sources

• libs/langchain/langchain/chat_models/vertexai.py
• libs/langchain/langchain/llms/vertexai.py
• libs/langchain/langchain/callbacks/streamlit/streamlit_callback_handler.py

---

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.

[AnhNgDo]

Hmm, I'm confused. The above suggestion is already implemented in Langchain - not sure what I'm supposed to change?

This is what I can see from Langchain's code (vertexai.py):

def _stream(
        self,
        prompt: str,
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> Iterator[GenerationChunk]:
        params = self._prepare_params(stop=stop, **kwargs)
        for stream_resp in stream_completion_with_retry(
            self, prompt, run_manager=run_manager, **params
        ):
            chunk = _response_to_generation(stream_resp)
            yield chunk
            if run_manager:
                run_manager.on_llm_new_token(
                    chunk.text,
                    chunk=chunk,
                    verbose=self.verbose,
                )