app.py

import streamlit as st
import asyncio
import sys
from pathlib import Path
import base64
import pandas as pd
from typing import Literal, Tuple, Optional
from wiki import render_wiki_tab
from search_handlers import run_global_search, run_local_search, run_drift_search
import auth


import graphrag.api as api
from graphrag.config import GraphRagConfig, load_config, resolve_paths
from graphrag.index.create_pipeline_config import create_pipeline_config
from graphrag.logging import PrintProgressReporter
from graphrag.utils.storage import _create_storage, _load_table_from_storage


st.set_page_config(page_title="GraphRAG Chat Interface", page_icon="🔍", layout="wide")

# Define default avatars at the module level
DEFAULT_USER_AVATAR = "👤"
DEFAULT_BOT_AVATAR = "🤖"

# Initialize session state for avatars
if "user_avatar" not in st.session_state:
    st.session_state.user_avatar = DEFAULT_USER_AVATAR
if "bot_avatar" not in st.session_state:
    st.session_state.bot_avatar = DEFAULT_BOT_AVATAR

# Define avatar images
USER_AVATAR = "👤"  # Default user emoji
BOT_AVATAR = "🤖"  # Default bot emoji


class StreamlitProgressReporter(PrintProgressReporter):
    def __init__(self, placeholder):
        super().__init__("")
        self.placeholder = placeholder

    def success(self, message: str):
        self.placeholder.success(message)


def render_chat_tab():
    """Render the Chat tab content."""    
    format_message_history()

    # Chat input
    if prompt := st.chat_input("Enter your query..."):
        # Add user message to history with timestamp
        st.session_state.messages.append(
            {
                "role": "user",
                "content": prompt,
                "timestamp": pd.Timestamp.now().strftime("%H:%M"),
            }
        )

        # Process query
        with st.spinner("Processing your query..."):
            response_placeholder = st.empty()
            try:
                if st.session_state.search_type == "global":
                    response, context = run_global_search(
                        config_filepath=st.session_state.config_filepath,
                        data_dir=st.session_state.data_dir,
                        root_dir=st.session_state.root_dir,
                        community_level=st.session_state.community_level,
                        response_type=st.session_state.response_type,
                        streaming=st.session_state.streaming,
                        query=prompt,
                        progress_placeholder=response_placeholder,
                    )
                elif st.session_state.search_type == "drift":
                    response, context = run_drift_search(
                        config_filepath=st.session_state.config_filepath,
                        data_dir=st.session_state.data_dir,
                        root_dir=st.session_state.root_dir,
                        community_level=st.session_state.community_level,
                        response_type=st.session_state.response_type,
                        streaming=st.session_state.streaming,
                        query=prompt,
                        progress_placeholder=response_placeholder,
                    )
                else:
                    response, context = run_local_search(
                        config_filepath=st.session_state.config_filepath,
                        data_dir=st.session_state.data_dir,
                        root_dir=st.session_state.root_dir,
                        community_level=st.session_state.community_level,
                        response_type=st.session_state.response_type,
                        streaming=st.session_state.streaming,
                        query=prompt,
                        progress_placeholder=response_placeholder,
                    )

                # Clear the placeholder before adding the final response
                response_placeholder.empty()

                # Add assistant response to history with timestamp
                st.session_state.messages.append(
                    {
                        "role": "assistant",
                        "content": response,
                        "timestamp": pd.Timestamp.now().strftime("%H:%M"),
                    }
                )

                # Show context in expander
                with st.expander("View Search Context"):
                    st.json(context)

            except Exception as e:
                error_message = f"Error processing query: {str(e)}"
                st.session_state.messages.append(
                    {
                        "role": "assistant",
                        "content": error_message,
                        "timestamp": pd.Timestamp.now().strftime("%H:%M"),
                    }
                )

        
        st.rerun()


def display_message(msg: str, is_user: bool = False, timestamp: str = "") -> None:
    """Display a chat message with avatar and consistent formatting."""
    role = "user" if is_user else "assistant"
    message_class = "user-message" if is_user else "assistant-message"
    avatar = st.session_state.user_avatar if is_user else st.session_state.bot_avatar

    message_container = f"""
        <div class="chat-message {message_class}">
            <div class="avatar">
                <div style="font-size: 25px; text-align: center;">{avatar}</div>
            </div>
            <div class="message-content-wrapper">
                <div class="message-bubble">
                    <div class="message-content">
                        {msg}
                    </div>
                </div>
                <div class="timestamp">{timestamp}</div>
            </div>
        </div>
    """
    st.markdown(message_container, unsafe_allow_html=True)


def format_message_history() -> None:
    """Display all messages in the chat history with consistent formatting."""
    st.markdown('<div class="chat-container">', unsafe_allow_html=True)
    for message in st.session_state.messages:
        timestamp = message.get("timestamp", "")
        display_message(
            msg=message["content"],
            is_user=(message["role"] == "user"),
            timestamp=timestamp,
        )
    st.markdown("</div>", unsafe_allow_html=True)


@st.cache_resource
def load_css():
    with open("styles.css", "r") as f:
        return f.read()


def initialize_session_state():
    """Initialize session state variables if they don't exist."""
    if "messages" not in st.session_state:
        st.session_state.messages = []
    if "response_placeholder" not in st.session_state:
        st.session_state.response_placeholder = None
    if "config_filepath" not in st.session_state:
        st.session_state.config_filepath = None
    if "data_dir" not in st.session_state:
        st.session_state.data_dir = None
    if "root_dir" not in st.session_state:
        st.session_state.root_dir = "."
    if "community_level" not in st.session_state:
        st.session_state.community_level = 2
    if "response_type" not in st.session_state:
        st.session_state.response_type = "concise"
    if "search_type" not in st.session_state:
        st.session_state.search_type = "global"
    if "streaming" not in st.session_state:
        st.session_state.streaming = True
    if "authenticated" not in st.session_state:
        st.session_state.authenticated = False


def main():
    initialize_session_state()

    # Authentication check
    if not st.session_state.authenticated:
        if auth.check_credentials():
            st.session_state.authenticated = True
            st.rerun()  # Rerun to reflect the authentication state
        else:
            st.stop()  # Stop further execution if authentication fails

    # If authenticated, proceed with the main app
    if st.session_state.authenticated:
        # Main application content
        st.title("PWC Home Assignment #1, Graphrag")

        css = load_css()
        st.markdown(f"<style>{css}</style>", unsafe_allow_html=True)

        # Sidebar configuration
        with st.sidebar:
            # Display logos side by side at the top of the sidebar
            col1, col2 = st.columns(2)            
            with col1:
                st.markdown(
                    '<div class="logo-container"><img class="logo-image" src="https://nexttech.pwc.co.il/wp-content/uploads/2023/12/image-2.png"></div>',
                    unsafe_allow_html=True,
                )
            with col2:
                st.markdown(
                    '<div class="logo-container"><img class="logo-image" src="https://nexttech.pwc.co.il/wp-content/uploads/2023/12/Frame.png"></div>',
                    unsafe_allow_html=True,
                )

            st.header("Configuration")
            st.session_state.community_level = st.number_input(
                "Community Level",
                min_value=0,
                max_value=10,
                value=st.session_state.community_level,
                help="Controls the granularity of the search...",
            )

            # Only show response type for global and local search
            if st.session_state.search_type != "drift":
                st.session_state.response_type = st.selectbox(
                    "Response Type",
                    options=["concise", "detailed"],
                    index=0 if st.session_state.response_type == "concise" else 1,
                    help="Style of response generation",
                )

            st.session_state.search_type = st.selectbox(
                "Search Type",
                options=["global", "local", "drift"],
                index=(
                    0
                    if st.session_state.search_type == "global"
                    else 1 if st.session_state.search_type == "local" else 2
                ),
                help="""Search Types:
                - Local Search: "Focuses on finding specific information by searching through direct connections in the knowledge graph. Best for precise, fact-based queries."
                - Global Search: "Analyzes the entire document collection at a high level using community summaries. Best for understanding broad themes and general policies."
                - DRIFT Search: "Combines local and global search capabilities, dynamically exploring connections while gathering detailed information. Best for complex queries requiring both specific details and broader context."
                """,
            )

            # Show streaming option only for supported search types
            if st.session_state.search_type != "drift":
                st.session_state.streaming = st.checkbox(
                    "Enable Streaming",
                    value=st.session_state.streaming,
                    help="Stream response tokens as they're generated",
                )
            else:
                st.session_state.streaming = False
                st.info("Streaming is not available for DRIFT search")

            # logout button
            if st.button("Logout"):
                st.session_state.clear()  # Clear all session state data
                initialize_session_state()  # Reinitialize the session state
                st.query_params = {"restart": "true"}  # Refresh the UI
                st.rerun()

        # Create tabs
        tab1, tab2 = st.tabs(["Assignment Documentation", "Chat"])

        # readme tab content
        with tab1:            
            render_wiki_tab()

        # Chat tab content
        with tab2:
            render_chat_tab()

        st.sidebar.markdown(
            """
            <div style="position: absolute; bottom: 0; width: 100%; text-align: center; font-size: 14px; margin-bottom: -200px;">
                Liran Baba | 
                <a href="https://linkedin.com/in/liranba" target="_blank">LinkedIn</a> | 
                <a href="https://huggingface.co/CordwainerSmith" target="_blank">HuggingFace</a>
            </div>
            """,
            unsafe_allow_html=True,
        )


if __name__ == "__main__":
    main()