app.py

import streamlit as st
import pandas as pd
from sqlalchemy import create_engine, text
from dotenv import load_dotenv
import os

load_dotenv()

# ================== MySQL Configuration ==================
DB_USER = "root"
DB_PASSWORD = os.getenv("DB_PASSWORD")   # ← Load from .env only
DB_HOST = "localhost"
DB_NAME = "sensibleai_finance"

engine = create_engine(
    f"mysql+pymysql://{DB_USER}:{DB_PASSWORD}@{DB_HOST}/{DB_NAME}?charset=utf8mb4"
)

# ================== Groq SQL Agent ==================
from langchain_community.utilities import SQLDatabase
from langchain_groq import ChatGroq
from langchain_community.agent_toolkits import create_sql_agent
from langchain_core.prompts import ChatPromptTemplate

db = SQLDatabase(engine)

llm = ChatGroq(
    model="llama-3.3-70b-versatile",
    temperature=0,
    api_key=os.getenv("GROQ_API_KEY"),
)

SYSTEM_PROMPT = """You are SensibleAI, a senior financial analyst AI for OneStream platform.

CRITICAL: Use EXACT category names:
- 'Development Costs'
- 'Operational Costs'
- 'Marketing Costs'
- 'Travelling Cost'
- 'Training Cost'
- 'Maintenance Cost'

Rules:
- Use `transactions` table for actual spending.
- Use `budgets` table for variance calculations.
- Always use full exact category name in SQL.
- For totals: use SUM(amount)
- For variance: join on month_year and category.
- Never use dangerous SQL commands.
- Always give clear numbers and short business insight.

Be professional and concise."""

prompt = ChatPromptTemplate.from_messages([
    ("system", SYSTEM_PROMPT),
    ("human", "{input}"),
    ("placeholder", "{agent_scratchpad}"),
])

agent_executor = create_sql_agent(
    llm=llm,
    db=db,
    prompt=prompt,
    agent_type="tool-calling",
    verbose=True,
    handle_parsing_errors=True,
)

# ================== Streamlit App ==================
st.set_page_config(page_title="SensibleAI Finance Assistant", layout="wide")
st.title("🧠 SensibleAI Finance Assistant")
st.caption("MySQL + Groq | Mini OneStream SensibleAI Demo")

# Sidebar
st.sidebar.title("Demo for OneStream Senior AI Engineer")
st.sidebar.markdown("**Natural Language → SQL + Visualizations**")

st.sidebar.subheader("Quick Questions")
examples = [
    "What is the total Development Costs?",
    "Calculate the variance for Marketing Costs in 2025-01",
    "Which category overspent the most?",
    "Give me a high-level summary of Q1 2025 expenses",
    "Show Marketing Costs by month"
]

for ex in examples:
    if st.sidebar.button(ex, use_container_width=True):
        if "messages" not in st.session_state:
            st.session_state.messages = []
        st.session_state.messages.append({"role": "user", "content": ex})

# Database status
if st.button("🔍 Check MySQL Connection & Data", type="primary"):
    try:
        with engine.connect() as conn:
            trans = conn.execute(text("SELECT COUNT(*) FROM transactions")).scalar()
            budg = conn.execute(text("SELECT COUNT(*) FROM budgets")).scalar()
        st.success(f"✅ Connected! Transactions: {trans} rows | Budgets: {budg} rows")
    except Exception as e:
        st.error(f"Connection error: {e}")

# Chat Interface
if "messages" not in st.session_state:
    st.session_state.messages = []

for msg in st.session_state.messages:
    with st.chat_message(msg["role"]):
        st.markdown(msg["content"])

if user_input := st.chat_input("Ask anything about the financial data..."):
    st.session_state.messages.append({"role": "user", "content": user_input})
    with st.chat_message("user"):
        st.markdown(user_input)

    with st.chat_message("assistant"):
        with st.spinner("Thinking..."):
            response = agent_executor.invoke({"input": user_input})
            answer = response["output"]
            st.markdown(answer)
    
    st.session_state.messages.append({"role": "assistant", "content": answer})

# ================== VISUALIZATIONS SECTION ==================
st.divider()
st.subheader("📊 Visual Analytics")

col1, col2 = st.columns(2)

with col1:
    if st.button("Show Expense Breakdown by Category", use_container_width=True):
        df_cat = pd.read_sql("""
            SELECT category, SUM(amount) as total_spend 
            FROM transactions 
            GROUP BY category 
            ORDER BY total_spend DESC
        """, engine)
        st.bar_chart(df_cat.set_index("category"), use_container_width=True)

with col2:
    if st.button("Show Monthly Spending Trend", use_container_width=True):
        df_month = pd.read_sql("""
            SELECT month_year, SUM(amount) as total 
            FROM transactions 
            GROUP BY month_year 
            ORDER BY month_year
        """, engine)
        st.line_chart(df_month.set_index("month_year"), use_container_width=True)

# Budget vs Actual Comparison
if st.button("Show Budget vs Actual Comparison", use_container_width=True):
    df_var = pd.read_sql("""
        SELECT 
            t.category,
            SUM(t.amount) as actual,
            b.budget_amount as budget,
            (SUM(t.amount) - b.budget_amount) as variance
        FROM transactions t
        JOIN budgets b ON t.month_year = b.month_year AND t.category = b.category
        GROUP BY t.category, b.budget_amount
        ORDER BY variance DESC
    """, engine)
    
    st.dataframe(df_var, use_container_width=True)
    
    st.bar_chart(df_var.set_index("category")[["actual", "budget"]], 
                 use_container_width=True)

st.divider()
st.info("This demo showcases full-stack AI development, natural language to SQL, and data visualization — key skills for the Senior AI Engineer role at OneStream.")