streamlit_app.py

"""Streamlit demo: McKesson-branded RAG chat over uploaded documents.

Run with:
    streamlit run streamlit_app.py

Required environment (see .env):
    MCK_CLIENT_ID, MCK_CLIENT_SECRET   - Okta credentials for MuleSoft auth
    LANGSMITH_API_KEY                  - to send traces to LangSmith
    LANGSMITH_TRACING=true             - enable tracing
    LANGSMITH_PROJECT=document-rag-poc - LangSmith project name (optional)
"""

from __future__ import annotations

import base64
import os
import tempfile
import uuid
from pathlib import Path

from dotenv import load_dotenv

# Load env BEFORE importing anything that builds clients.
load_dotenv(Path(__file__).parent / ".env")

import streamlit as st

from agents.chat import analyze_with_agents, answer_question, index_document


def get_image_base64(image_path):
    """Convert image to base64 string for embedding in HTML"""
    try:
        with open(image_path, "rb") as image_file:
            return base64.b64encode(image_file.read()).decode()
    except Exception:
        return None


# ---------------------------------------------------------------------------
# Page config + McKesson-style chrome
# ---------------------------------------------------------------------------

st.set_page_config(
    page_title="McKesson Document Assistant",
    page_icon="🩺",
    layout="wide",
    initial_sidebar_state="expanded",
)


_MCK_BLUE = "#007DBA"
_MCK_NAVY = "#002F6C"
_MCK_BG_SOFT = "#F4F6F9"


_CUSTOM_CSS = f"""
<style>
    /* Hide Streamlit default chrome */
    #MainMenu {{visibility: hidden;}}
    footer {{visibility: hidden;}}

    /* Header band */
    .mck-header {{
        display: flex;
        align-items: center;
        gap: 1.5rem;
        background: {_MCK_BLUE};
        color: white;
        padding: 1.5rem 2rem;
        border-radius: 8px;
        margin-bottom: 2rem;
    }}
    .mck-logo {{
        width: 120px;
        height: auto;
        border-radius: 4px;
        background: white;
        padding: 0.5rem;
    }}
    .mck-title {{
        font-size: 1.8rem;
        font-weight: 600;
        margin: 0;
    }}
    .mck-sub {{
        font-size: 0.9rem;
        opacity: 0.85;
        margin: 0;
    }}

    /* Chat bubbles */
    [data-testid="stChatMessage"] {{
        border-radius: 12px;
        padding: 0.25rem 0.5rem;
    }}

    /* Buttons */
    .stButton > button {{
        background-color: {_MCK_BLUE};
        color: white;
        border: none;
        border-radius: 6px;
        font-weight: 600;
    }}
    .stButton > button:hover {{
        background-color: {_MCK_NAVY};
        color: white;
    }}

    /* Sidebar polish */
    section[data-testid="stSidebar"] {{
        background-color: {_MCK_BG_SOFT};
    }}
    section[data-testid="stSidebar"] h2 {{
        color: {_MCK_NAVY};
    }}

    /* Source pill */
    .src-pill {{
        display: inline-block;
        background: {_MCK_BG_SOFT};
        color: {_MCK_NAVY};
        border: 1px solid #d6dde6;
        border-radius: 999px;
        padding: 2px 10px;
        margin: 2px 4px 2px 0;
        font-size: 0.78rem;
        font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', monospace;
    }}
</style>
"""

st.markdown(_CUSTOM_CSS, unsafe_allow_html=True)


# Header
st.markdown(
    f"""
    <div class="mck-header">
        <div>
            <p class="mck-title">Document Assistant</p>
            <p class="mck-sub">Upload a document and ask questions — answers are grounded in your file.</p>
        </div>
    </div>
    """,
    unsafe_allow_html=True,
)


# ---------------------------------------------------------------------------
# Session state
# ---------------------------------------------------------------------------

if "messages" not in st.session_state:
    st.session_state.messages = []
if "indexed_files" not in st.session_state:
    st.session_state.indexed_files = []
if "indexed_paths" not in st.session_state:
    st.session_state.indexed_paths = []  # full absolute paths for the agent
if "session_id" not in st.session_state:
    st.session_state.session_id = str(uuid.uuid4())[:8]
if "mode" not in st.session_state:
    st.session_state.mode = "Multi-agent (orchestrator + fact-checker)"
if "summary_type" not in st.session_state:
    st.session_state.summary_type = "executive"


# ---------------------------------------------------------------------------
# Sidebar — upload + status
# ---------------------------------------------------------------------------

with st.sidebar:
    # Add McKesson logo at the top of sidebar
    logo_base64 = get_image_base64("mc_logo.png")
    if logo_base64:
        st.markdown(
            f'<div style="text-align: center; margin-bottom: 20px;"><img src="data:image/png;base64,{logo_base64}" alt="McKesson Logo" style="width: 120px; height: auto; border-radius: 4px;"></div>',
            unsafe_allow_html=True,
        )
    
    st.markdown("## Document")
    uploaded = st.file_uploader(
        "Upload a PDF, TXT, or Markdown file",
        type=["pdf", "txt", "md"],
        accept_multiple_files=False,
    )

    if uploaded is not None:
        # Persist upload to a stable temp path so the indexer can read it.
        upload_dir = Path(tempfile.gettempdir()) / "mck_docs"
        upload_dir.mkdir(parents=True, exist_ok=True)
        dest = upload_dir / uploaded.name

        if dest.name not in st.session_state.indexed_files:
            dest.write_bytes(uploaded.getbuffer())
            with st.spinner(f"Indexing **{uploaded.name}**..."):
                try:
                    result = index_document(str(dest))
                    if result["indexed"]:
                        st.success(
                            f"Indexed **{uploaded.name}** — {result['chunks']} chunks"
                        )
                    else:
                        st.info(f"**{uploaded.name}** already indexed — ready.")
                    st.session_state.indexed_files.append(dest.name)
                    st.session_state.indexed_paths.append(str(dest.resolve()))
                except Exception as e:
                    st.error(f"Failed to index: {e}")

    st.markdown("---")
    st.markdown("### Session")
    st.code(st.session_state.session_id, language=None)

    st.markdown("### Indexed files")
    if st.session_state.indexed_files:
        for f in st.session_state.indexed_files:
            st.markdown(f"- `{f}`")
    else:
        st.caption("No files indexed yet.")

    st.markdown("---")
    st.markdown("### Analysis Mode")
    
    # Mode selection with clearer descriptions
    mode_options = [
        ("Document Analysis (AI + Fact-Checker)", "Multi-agent (orchestrator + fact-checker)"),
        ("Quick Q&A (Direct Answers)", "Quick Q&A (single-shot RAG)")
    ]
    
    selected_display = st.radio(
        "",
        options=[option[0] for option in mode_options],
        index=0 if st.session_state.mode.startswith("Multi") else 1,
        label_visibility="collapsed",
    )
    
    # Map display selection back to internal mode value
    for display, value in mode_options:
        if selected_display == display:
            st.session_state.mode = value
            break
    
    # Show description for selected mode
    if st.session_state.mode.startswith("Multi"):
        st.caption("📊 Generate comprehensive document summaries with fact verification")
    else:
        st.caption("❓ Get direct answers to specific questions with sources")
    
    # Show summary type options only for document analysis mode
    if st.session_state.mode.startswith("Multi"):
        st.markdown("**Summary Type:**")
        st.session_state.summary_type = st.selectbox(
            "",
            options=["executive", "detailed", "bullet"],
            index=["executive", "detailed", "bullet"].index(
                st.session_state.summary_type
            ),
            label_visibility="collapsed",
            format_func=lambda x: {
                "executive": "📝 Executive Summary (3-5 sentences)",
                "detailed": "📋 Detailed Analysis (structured sections)", 
                "bullet": "🔸 Bullet Points (key facts & findings)"
            }[x]
        )
        st.caption("💡 Your question guides what to focus on in the summary")
    else:
        st.caption("💡 Ask specific questions and get direct answers with sources")

    st.markdown("---")
    st.markdown("### Tracing")
    tracing_on = os.getenv("LANGSMITH_TRACING", "").lower() in {"true", "1", "yes"}
    project = os.getenv("LANGSMITH_PROJECT", "(default)")
    if tracing_on:
        st.success(f"LangSmith ON · project `{project}`")
    else:
        st.warning("LangSmith tracing disabled")

    if st.button("Clear chat history"):
        st.session_state.messages = []
        st.rerun()


# ---------------------------------------------------------------------------
# Chat area
# ---------------------------------------------------------------------------

for msg in st.session_state.messages:
    with st.chat_message(msg["role"], avatar="🩺" if msg["role"] == "assistant" else None):
        st.markdown(msg["content"])
        if msg.get("sources"):
            pills = "".join(
                f'<span class="src-pill">{s["source"]} · p.{s["page"]}</span>'
                for s in msg["sources"]
            )
            st.markdown(pills, unsafe_allow_html=True)


# Adaptive interface based on mode
if st.session_state.mode.startswith("Multi"):
    # Document Analysis Mode - Show button instead of chat input
    st.markdown("---")
    st.markdown("### 📊 Document Analysis")
    
    if not st.session_state.indexed_paths:
        st.warning("📄 Please upload a document first to begin analysis.")
    else:
        if st.button("🔍 Analyze Document", type="primary", use_container_width=True):
            # Create a user message for the analysis
            analysis_prompt = f"Generate a {st.session_state.summary_type} analysis of the document"
            st.session_state.messages.append({"role": "user", "content": analysis_prompt})
            with st.chat_message("user"):
                st.markdown(f"📊 Analyze document ({st.session_state.summary_type} summary)")
            
            with st.chat_message("assistant", avatar="🩺"):
                try:
                    # Use the most recently uploaded document for analysis.
                    doc_path = st.session_state.indexed_paths[-1]
                    with st.status(
                        "Running document analysis...", expanded=True
                    ) as status:
                        st.write("**1.** Orchestrator analyzing document")
                        st.write(
                            f"**2.** Generating {st.session_state.summary_type} summary"
                        )
                        st.write("**3.** Fact-checking the analysis")
                        result = analyze_with_agents(
                            document_path=doc_path,
                            query="comprehensive document analysis",
                            summary_type=st.session_state.summary_type,
                            thread_id=st.session_state.session_id,
                        )
                        status.update(
                            label="Document analysis complete", state="complete"
                    )
                    st.markdown(result["answer"])
                    st.session_state.messages.append(
                        {
                            "role": "assistant",
                            "content": result["answer"],
                            "sources": [],
                        }
                    )
                except Exception as e:
                    err = f"Sorry — something went wrong: `{type(e).__name__}: {e}`"
                    st.error(err)
                    st.session_state.messages.append(
                        {"role": "assistant", "content": err, "sources": []}
                    )

else:
    # Quick Q&A Mode - Traditional chat interface
    prompt = st.chat_input("Ask a question about the uploaded document...")
    if prompt:
        st.session_state.messages.append({"role": "user", "content": prompt})
        with st.chat_message("user"):
            st.markdown(prompt)

        with st.chat_message("assistant", avatar="🩺"):
            try:
                with st.spinner("Thinking..."):
                    result = answer_question(prompt)
                st.markdown(result["answer"])
                if result["sources"]:
                    pills = "".join(
                        f'<span class="src-pill">{s["source"]} · p.{s["page"]}</span>'
                        for s in result["sources"]
                    )
                    st.markdown(pills, unsafe_allow_html=True)
                st.session_state.messages.append(
                    {
                        "role": "assistant",
                        "content": result["answer"],
                        "sources": result["sources"],
                    }
                )
            except Exception as e:
                err = f"Sorry — something went wrong: `{type(e).__name__}: {e}`"
                st.error(err)
                st.session_state.messages.append(
                    {"role": "assistant", "content": err, "sources": []}
                )