app.py

import os

from dataset_loader import load_dataset

if not os.path.isdir("service_data"):
    load_dataset()

import uuid

import streamlit as st
from service_data.graph_llm_service import (
    GraphLLMService,  # Assuming your service is in a separate file
)

# Initialize the GraphLLM service
service = GraphLLMService()


def reset_chat():
    st.session_state.messages = []
    st.session_state.session_id = str(uuid.uuid4())


# Configure Streamlit page
st.set_page_config(page_title="JMB Chat", page_icon="💬")
st.title("JMB Chat")

# Initialize session state
if "messages" not in st.session_state:
    st.session_state.messages = []
if "session_id" not in st.session_state:
    st.session_state.session_id = str(uuid.uuid4())

# Sidebar with reset button
with st.sidebar:
    st.header("Chat Settings")
    if st.button("Reset Chat"):
        reset_chat()
    st.divider()
    st.write(f"Current Session ID: `{st.session_state.session_id}`")

# Display chat messages
for message in st.session_state.messages:
    role = message["role"]
    content = message["content"]
    documents = message.get("documents", [])

    with st.chat_message(role):
        st.markdown(content)
        if role == "assistant" and documents:
            with st.expander(f"📁 Sources ({len(documents)})"):
                for doc in documents:
                    st.write(f"- **{doc['file']}** (page {doc['page']})")

# Chat input and processing
if prompt := st.chat_input("Type your question..."):
    # Add user message to history
    st.session_state.messages.append({"role": "user", "content": prompt})

    # Display user message
    with st.chat_message("user"):
        st.markdown(prompt)

    # Get and display assistant response
    with st.chat_message("assistant"):
        try:
            with st.spinner("Thinking..."):
                response = service.get_answer(prompt, st.session_state.session_id)

            answer = response["answer"]
            documents = response["documents"]

            st.markdown(answer)
            if documents:
                with st.expander(f"📁 Sources ({len(documents)})"):
                    for doc in documents:
                        st.write(f"- **{doc['file']}** (page {doc['page']})")

            # Add assistant response to history
            st.session_state.messages.append(
                {"role": "assistant", "content": answer, "documents": documents}
            )

        except Exception as e:
            st.error(f"An error occurred: {str(e)}")
            st.session_state.messages.append(
                {"role": "assistant", "content": f"Error: {str(e)}"}
            )