frontend_fixed.py

"""
Frontend Streamlit application for the chatbot (fixed version).
"""
import streamlit as st
import os
# Import the entire backend module instead of specific functions
import backend
from utils.logging_config import setup_logging
import config

# Setup logging
logger = setup_logging()

# Constants for memory optimization
MAX_MESSAGES = 50

# Set page config
st.set_page_config(
    page_title="Document Chatbot",
    page_icon="📚",
    layout="centered",
    initial_sidebar_state="collapsed"
)

# Hide Streamlit's default elements and style sidebar
st.markdown("""
<style>
/* Hide default elements */
button[kind="deploy"], 
[data-testid="stToolbar"],
.stDeployButton,
#MainMenu,
footer {
    display: none !important;
}

/* Style the sidebar */
[data-testid="stSidebar"] {
    width: 120px !important;
    background-color: #0e1117 !important;
    border-right: 1px solid #1e1e1e !important;
}

/* Style the collapse button to prevent movement and maintain size */
button[kind="menuButton"] {
    left: 120px !important;
    margin-left: 0 !important;
    position: fixed !important;
    transform: translateX(0) !important;
    transition: none !important;
    background-color: #0e1117 !important;
    color: #fff !important;
    z-index: 999 !important;
}

/* Button styling */
[data-testid="stSidebar"] [data-testid="stButton"] button {
    background-color: #262730 !important;
    color: white !important;
    padding: 8px 10px !important;
    width: 100px !important;
    font-size: 0.9rem !important;
    margin: 1rem auto !important;
    display: block !important;
    border-radius: 4px !important;
    white-space: nowrap !important;
}

[data-testid="stSidebar"] [data-testid="stButton"] button p {
    text-align: center !important;
    white-space: nowrap !important;
    overflow: visible !important;
}

[data-testid="stSidebar"] [data-testid="stButton"] button:hover {
    background-color: #1E2130 !important;
}
</style>
""", unsafe_allow_html=True)

# Helper function to add messages and maintain history length
def add_message(role, content):
    """Add message to session state and trim if needed."""
    st.session_state.messages.append({"role": role, "content": content})
    # Trim message history if it gets too large
    if len(st.session_state.messages) > MAX_MESSAGES:
        # Keep the most recent messages
        st.session_state.messages = st.session_state.messages[-MAX_MESSAGES:]
    logger.info(f"Added {role} message. History length: {len(st.session_state.messages)}")

# Create a sidebar with just the button
with st.sidebar:
    def clear_chat():
        st.session_state.messages = []
        logger.info("Chat history cleared via sidebar button")
        st.rerun()
        
    st.button("Clear Chat", on_click=clear_chat)

# Initialize session state for chat history
if "messages" not in st.session_state:
    st.session_state.messages = []

# Initialize cached resources for all sessions
@st.cache_resource
def initialize_resources():
    """Initialize shared resources once for all sessions."""
    logger.info("Initializing shared resources...")
    
    # Get configuration from config module
    chatbot_config = config.get_chatbot_config()
    api_key = os.getenv("ANTHROPIC_API_KEY")
    
    # Load models using cache
    # Use the functions from the backend module
    llm = backend.load_llm_model(
        api_key,
        chatbot_config.get("model", "claude-3-7-sonnet-20250219"),
        chatbot_config.get("temperature", 0.1),
        chatbot_config.get("max_tokens", 2048)
    )
    
    embed_model = backend.load_embedding_model(
        chatbot_config.get("embedding_model", "sentence-transformers/all-MiniLM-L6-v2"),
        chatbot_config.get("device", "cpu"),
        chatbot_config.get("embed_batch_size", 8)
    )
    
    # Load or create index (shared across sessions)
    index = backend.load_or_create_index()
    
    return {
        "config": chatbot_config,
        "llm": llm,
        "embed_model": embed_model,
        "index": index
    }

# Get shared resources
resources = initialize_resources()

# Initialize chatbot only in session state if it doesn't exist
if "chatbot" not in st.session_state:
    with st.spinner("Initializing chatbot..."):
        logger.info("Initializing chatbot with shared resources...")
        st.session_state.chatbot = backend.Chatbot(
            resources["config"],
            resources["llm"],
            resources["embed_model"],
            resources["index"]
        )
        # Initialize query engine
        st.session_state.chatbot.initialize_query_engine()
        logger.info("Chatbot initialized successfully")

# Title and description
st.title("📚 Paul's Chatbot")
st.markdown("""
This chatbot can answer questions about your documents. 
Ask any question about the content in your documents!
""")

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

# Chat input
if prompt := st.chat_input("What would you like to know?"):
    # Add user message to chat history
    add_message("user", prompt)
    with st.chat_message("user"):
        st.markdown(prompt)

    # Get chatbot response
    with st.chat_message("assistant"):
        with st.spinner("Thinking..."):
            logger.info(f"User query: {prompt}")
            response = st.session_state.chatbot.query(prompt)
            st.markdown(response)
            add_message("assistant", response)
            logger.info("Response provided to user")