import os
import time

import streamlit as st
from langchain_core.runnables import RunnableLambda


def stream_reply(question: str):
    delay = float(os.getenv("STREAM_DELAY_SECONDS", "0.03"))
    for word in f"LangChain streamed this response for: {question}".split():
        if delay:
            time.sleep(delay)
        yield word + " "


chain = RunnableLambda(stream_reply)

st.set_page_config(page_title="LangChain Streaming Chatbot")
st.title("LangChain Streaming Chatbot")

if "messages" not in st.session_state:
    st.session_state.messages = [
        {
            "role": "assistant",
            "content": "Ask a question to see a streamed LangChain response.",
        }
    ]

for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

if prompt := st.chat_input("Ask about LangChain streaming"):
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    with st.chat_message("assistant"):
        response = st.write_stream(chain.stream(prompt))

    st.session_state.messages.append({"role": "assistant", "content": response})