mirror of
https://github.com/langchain-ai/streamlit-agent.git
synced 2026-07-01 09:25:05 -04:00
Joshua Carroll
39 lines
1.3 KiB
Python
39 lines
1.3 KiB
Python
from langchain.callbacks.base import BaseCallbackHandler
|
|
from langchain.schema import ChatMessage
|
|
from langchain_openai import ChatOpenAI
|
|
import streamlit as st
|
|
|
|
|
|
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) -> None:
|
|
self.text += token
|
|
self.container.markdown(self.text)
|
|
|
|
|
|
with st.sidebar:
|
|
openai_api_key = st.text_input("OpenAI API Key", type="password")
|
|
|
|
if "messages" not in st.session_state:
|
|
st.session_state["messages"] = [ChatMessage(role="assistant", content="How can I help you?")]
|
|
|
|
for msg in st.session_state.messages:
|
|
st.chat_message(msg.role).write(msg.content)
|
|
|
|
if prompt := st.chat_input():
|
|
st.session_state.messages.append(ChatMessage(role="user", content=prompt))
|
|
st.chat_message("user").write(prompt)
|
|
|
|
if not openai_api_key:
|
|
st.info("Please add your OpenAI API key to continue.")
|
|
st.stop()
|
|
|
|
with st.chat_message("assistant"):
|
|
stream_handler = StreamHandler(st.empty())
|
|
llm = ChatOpenAI(openai_api_key=openai_api_key, streaming=True, callbacks=[stream_handler])
|
|
response = llm.invoke(st.session_state.messages)
|
|
st.session_state.messages.append(ChatMessage(role="assistant", content=response.content))
|