app.py

import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent))
sys.path.insert(0, str(Path(__file__).parent / "core"))
sys.path.insert(0, str(Path(__file__).parent / "core" / "tools"))

from dotenv import load_dotenv
load_dotenv()

import streamlit as st
from langchain_core.messages import HumanMessage

# ── Page config ────────────────────────────────────────────────────────────────
st.set_page_config(
    page_title="FinancAgent",
    page_icon="💼",
    layout="wide",
    initial_sidebar_state="expanded",
)

# ── Lazy-load the compiled graph (cached across reruns) ────────────────────────
@st.cache_resource(show_spinner="Initializing financial intelligence system…")
def load_graph():
    from graph import app
    return app

# ── Constants ──────────────────────────────────────────────────────────────────
ROLES = ["CFO", "CEO", "Finance Analyst", "Department Head"]
DEPARTMENTS = ["Engineering", "Finance", "Marketing", "Operations", "Product", "Sales"]

AGENT_LABELS = {
    "analyst":  "📊 Analyst",
    "budget":   "📋 Budget",
    "anomaly":  "🔍 Anomaly",
    "search":   "🔎 Search",
    "forecast": "📈 Forecast",
}

EXAMPLE_QUESTIONS = [
    "How much did we spend in 2024 broken down by department?",
    "Were we over budget last year? Highlight anything over 10%.",
    "Are there any anomalies or suspicious transactions in 2024?",
    "What is our projected total spend for the next 6 months?",
    "Show me revenue vs cash flow for 2024.",
    "Find all Airbnb transactions in Operations.",
]

# ── Session state defaults ─────────────────────────────────────────────────────
if "messages" not in st.session_state:
    st.session_state.messages = []          # list of {role, content, agents_used}
if "user_role" not in st.session_state:
    st.session_state.user_role = "CFO"
if "user_department" not in st.session_state:
    st.session_state.user_department = None

# ── Sidebar ────────────────────────────────────────────────────────────────────
with st.sidebar:
    st.title("💼 FinancAgent")
    st.caption("AI-powered financial intelligence")
    st.divider()

    st.subheader("Session Setup")

    selected_role = st.selectbox(
        "Your Role",
        ROLES,
        index=ROLES.index(st.session_state.user_role),
    )

    selected_department = None
    if selected_role == "Department Head":
        selected_department = st.selectbox("Your Department", DEPARTMENTS)

    # Clear chat when role changes
    if selected_role != st.session_state.user_role:
        st.session_state.messages = []
        st.session_state.user_role = selected_role
    if selected_department != st.session_state.user_department:
        st.session_state.messages = []
        st.session_state.user_department = selected_department

    st.divider()

    st.subheader("Try an example")
    for example in EXAMPLE_QUESTIONS:
        if st.button(example, use_container_width=True, key=f"ex_{example[:20]}"):
            st.session_state.pending_question = example

    st.divider()

    if st.button("🗑️ Clear Chat", use_container_width=True):
        st.session_state.messages = []
        st.rerun()

    st.caption(f"Role: **{selected_role}**" + (f" · {selected_department}" if selected_department else ""))

# ── Main chat area ─────────────────────────────────────────────────────────────
st.header("Financial Intelligence Assistant")

# Display conversation history
for msg in st.session_state.messages:
    with st.chat_message(msg["role"]):
        st.markdown(msg["content"])
        if msg["role"] == "assistant" and msg.get("agents_used"):
            labels = [AGENT_LABELS.get(a, a) for a in msg["agents_used"]]
            st.caption("Agents: " + "  ·  ".join(labels))

# Handle example button clicks
pending = st.session_state.pop("pending_question", None)

# Chat input
user_input = st.chat_input("Ask a financial question…") or pending

if user_input:
    # Show user message immediately
    st.session_state.messages.append({"role": "user", "content": user_input})
    with st.chat_message("user"):
        st.markdown(user_input)

    # Run the graph
    with st.chat_message("assistant"):
        with st.spinner("Thinking…"):
            try:
                graph = load_graph()
                from core.state import FinanceSystemState

                state = FinanceSystemState(
                    messages=[HumanMessage(user_input)],
                    user_role=selected_role,
                    user_department=selected_department,
                )
                result = graph.invoke(state)

                response = result.get("final_response") or "No response generated."
                agents_used = result.get("routing_decision") or []

            except Exception as e:
                response = f"Something went wrong: {e}"
                agents_used = []

        st.markdown(response)
        if agents_used:
            labels = [AGENT_LABELS.get(a, a) for a in agents_used]
            st.caption("Agents: " + "  ·  ".join(labels))

    st.session_state.messages.append({
        "role": "assistant",
        "content": response,
        "agents_used": agents_used,
    })