chat-proxy/streamlit/app.py

import streamlit as st
import streamlit.components.v1 as components
import redis
import json
import uuid
import hashlib
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

# ---------------------
# Redis connection
# ---------------------
r = redis.Redis(host="localhost", port=6379, decode_responses=True)

# ---------------------
# Session ID helpers
# ---------------------
def get_or_set_session_id():
    """Get or set session_id from browser localStorage and trigger rerun when first set."""
    if "session_id" in st.session_state:
        return st.session_state.session_id

    components.html(f"""
        <script>
        const key = 'egalware_session_id';
        let sid = window.localStorage.getItem(key);
        if (!sid) {{
            sid = '{uuid.uuid4()}';
            window.localStorage.setItem(key, sid);
        }}
        const streamlitDoc = window.parent.document;
        let hidden = streamlitDoc.querySelector('#session_id_input');
        if (!hidden) {{
            hidden = document.createElement('input');
            hidden.type = 'hidden';
            hidden.id = 'session_id_input';
            streamlitDoc.body.appendChild(hidden);
        }}
        hidden.value = sid;
        hidden.dispatchEvent(new Event('input', {{ bubbles: true }}));
        </script>
    """, height=0)

    sid = st.text_input("session_id_input", label_visibility="collapsed")
    if sid and sid != st.session_state.get("session_id"):
        st.session_state.session_id = sid
        st.rerun()  # immediately rerun with the new ID
    return st.session_state.get("session_id")

# ---------------------
# Chat history persistence
# ---------------------
def load_history(session_id):
    raw = r.get(f"chatbot:history:{session_id}")
    if raw:
        messages = json.loads(raw)
        return [
            AIMessage(content=m["content"]) if m["type"] == "ai"
            else HumanMessage(content=m["content"])
            for m in messages
        ]
    return [AIMessage(content="Hello, I am EgalWare's ChatBot. How can I help you? (puoi fare domande in italiano, ma in inglese funziona meglio...)")]

def save_history(session_id, history):
    messages = [
        {"type": "ai" if isinstance(m, AIMessage) else "human", "content": m.content}
        for m in history
    ]
    r.set(f"chatbot:history:{session_id}", json.dumps(messages))
    r.expire(f"chatbot:history:{session_id}", 60*60*24*7)

def delete_history(session_id):
    r.delete(f"chatbot:history:{session_id}")
    for key in r.scan_iter(f"chatbot:cache:{session_id}:*"):
        r.delete(key)

# ---------------------
# Caching
# ---------------------
def get_cache_key(session_id, prompt):
    digest = hashlib.sha256(prompt.encode()).hexdigest()
    return f"chatbot:cache:{session_id}:{digest}"

def get_cached_response(session_id, prompt):
    return r.get(get_cache_key(session_id, prompt))

def set_cached_response(session_id, prompt, response):
    r.setex(get_cache_key(session_id, prompt), 300, response)

# ---------------------
# LLM
# ---------------------
def get_response(session_id, user_query, chat_history):
    cached = get_cached_response(session_id, user_query)
    if cached:
        yield cached
        return

    prompt = ChatPromptTemplate.from_template(
        "You are a helpful assistant. Answer the following considering the history:\n\n"
        "Chat history: {chat_history}\n\nUser question: {user_question}"
    )
    llm = ChatOpenAI(
        base_url="http://10.74.83.100:1234/v1",
        api_key="lm-studio",
        model="qwen/qwen3-4b-2507"
    )
    chain = prompt | llm | StrOutputParser()

    full_resp = ""
    for chunk in chain.stream({
        "chat_history": chat_history,
        "user_question": user_query
    }):
        full_resp += chunk
        yield chunk

    set_cached_response(session_id, user_query, full_resp)

# ---------------------
# UI
# ---------------------
st.set_page_config(page_title="Egalware Chatbot", page_icon="🤖")
session_id = get_or_set_session_id()

# Sticky header CSS
st.markdown("""
    <style>
    .sticky-header {
        position: sticky;
        top: 0;
        background-color: white;
        padding-top: 0.5rem;
        padding-bottom: 0.5rem;
        z-index: 999;
        border-bottom: 1px solid #ddd;
    }
    </style>
""", unsafe_allow_html=True)

# Header
st.markdown('<div class="sticky-header">', unsafe_allow_html=True)
col_title, col_btn = st.columns([0.9, 0.1])
with col_title:
    st.title("Egalware's Chatbot")
with col_btn:
    if st.button("🗑️", help="Clear conversation", type="secondary"):
        delete_history(session_id)
        st.session_state.chat_history = load_history(session_id)
        st.rerun()
st.markdown('</div>', unsafe_allow_html=True)

# If still no ID, display placeholder history so UI doesn't look empty
if not session_id:
    if "chat_history" not in st.session_state:
        st.session_state.chat_history = [AIMessage(content="Initializing session… please wait")]
else:
    if "chat_history" not in st.session_state:
        st.session_state.chat_history = load_history(session_id)

# Conversation display
for message in st.session_state.chat_history:
    role = "AI" if isinstance(message, AIMessage) else "Human"
    with st.chat_message(role):
        st.write(message.content)

# Input
user_query = st.chat_input("Type your message here…")
if session_id and user_query:
    st.session_state.chat_history.append(HumanMessage(content=user_query))
    with st.chat_message("Human"):
        st.markdown(user_query)
    with st.chat_message("AI"):
        response_text = st.write_stream(
            get_response(session_id, user_query, st.session_state.chat_history)
        )
    st.session_state.chat_history.append(AIMessage(content=response_text))
    save_history(session_id, st.session_state.chat_history)