####
#### Streamlit Streaming using LM Studio as OpenAI Standin
#### run with `streamlit run app.py`

# !pip install pypdf langchain langchain-core langchain-openai

import streamlit as st
from langchain_core.messages import AIMessage, HumanMessage
from langchain_openai import ChatOpenAI
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate

# app config
st.set_page_config(page_title="Egalware Chatbot", page_icon="🤖")
st.title("Egalware's Live Chatbot")

def get_response(user_query, chat_history):
    template = """
    You are a helpful assistant. Answer the following questions considering the history of the conversation:

    Chat history: {chat_history}

    User question: {user_question}
    """
    prompt = ChatPromptTemplate.from_template(template)

    # Using LM Studio Local Inference Server
    llm = ChatOpenAI(
        base_url="http://10.74.83.100:1234/v1",
        api_key="lm-studio",
        model="qwen/qwen3-4b-2507"
    )

    chain = prompt | llm | StrOutputParser()

    # Return a generator for streaming
    return chain.stream({
        "chat_history": chat_history,
        "user_question": user_query,
    })

# session state
if "chat_history" not in st.session_state:
    st.session_state.chat_history = [
        AIMessage(content="Hello, I am EgalWare's Live & Stateless ChatBot. "
                          "How can I help you? (puoi fare domande in italiano, "
                          "ma in inglese funziona meglio...)"),
    ]

# conversation history display
for message in st.session_state.chat_history:
    if isinstance(message, AIMessage):
        with st.chat_message("AI"):
            st.write(message.content)
    elif isinstance(message, HumanMessage):
        with st.chat_message("Human"):
            st.write(message.content)

# user input
user_query = st.chat_input("Type your message here...")
if user_query:
    # store human message
    st.session_state.chat_history.append(HumanMessage(content=user_query))

    with st.chat_message("Human"):
        st.markdown(user_query)

    # stream AI response and capture it
    with st.chat_message("AI"):
        chunks = []
        for chunk in get_response(user_query, st.session_state.chat_history):
            st.write(chunk)
            chunks.append(chunk)
        full_response = "".join(chunks)

    # store AI message with actual text
    st.session_state.chat_history.append(AIMessage(content=full_response))