File size: 1,378 Bytes
19b1388 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 |
from langchain.callbacks.base import BaseCallbackHandler
from langchain.chat_models import ChatOpenAI
from langchain.schema import ChatMessage
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(st.session_state.messages)
st.session_state.messages.append(ChatMessage(role="assistant", content=response.content))
|