app.py

import os
import pandas as pd
from pydantic import BaseModel
import needle
import streamlit as st

# Define schema for Needle 2 on-device extraction
class ExpenseRecord(BaseModel):
    merchant: str
    amount: float
    category: str
    payment_method: str

CSV_FILE = "expenses.csv"

def initialize_csv():
    """Ensure expenses.csv exists with correct columns."""
    if not os.path.exists(CSV_FILE):
        df = pd.DataFrame({
            "merchant": pd.Series(dtype="str"),
            "amount": pd.Series(dtype="float"),
            "category": pd.Series(dtype="str"),
            "payment_method": pd.Series(dtype="str")
        })
        df.to_csv(CSV_FILE, index=False)

def extract_expense(note_text: str) -> ExpenseRecord:
    """Extract structured financial data using Needle 2."""
    try:
        extracted = needle.extract(note_text, ExpenseRecord)
        return extracted
    except Exception:
        return ExpenseRecord(
            merchant="Uncategorized",
            amount=0.0,
            category="Miscellaneous",
            payment_method="Unknown"
        )

def save_expense(record: ExpenseRecord):
    """Save an expense record to local CSV."""
    initialize_csv()
    df = pd.read_csv(CSV_FILE)
    new_row = pd.DataFrame([{
        "merchant": str(record.merchant),
        "amount": round(float(record.amount), 2),
        "category": str(record.category),
        "payment_method": str(record.payment_method)
    }])
    if df.empty:
        df = new_row
    else:
        df = pd.concat([df, new_row], ignore_index=True)
    df.to_csv(CSV_FILE, index=False)

def main():
    """Streamlit dashboard interface."""
    st.set_page_config(page_title="Needle 2 Finance Parser", page_icon="💳", layout="wide")
    st.title("💳 Needle 2 Offline Finance Parser")
    st.caption("Zero-cloud, high-speed expense extraction powered by 45M Parameter Needle 2 Engine")

    col1, col2 = st.columns([1, 1])

    with col1:
        st.subheader("📥 Input Financial Note")
        note_input = st.text_area("Type or paste raw expense note:", placeholder="e.g. Spent Rs 450 at Blinkit for groceries using UPI")

        if st.button("⚡ Extract & Save Expense", type="primary"):
            if note_input.strip():
                with st.spinner("Extracting with Needle 2..."):
                    record = extract_expense(note_input)
                    save_expense(record)
                    st.success(f"Extracted: {record.merchant} - ₹{record.amount:,.2f} ({record.category})")
            else:
                st.warning("Please enter an expense note first.")

    with col2:
        st.subheader("📊 Spending Summary")
        initialize_csv()
        df = pd.read_csv(CSV_FILE)

        if not df.empty:
            st.metric("Total Expenses", f"₹{df['amount'].sum():,.2f}", delta=f"{len(df)} transactions")
            category_totals = df.groupby("category")["amount"].sum()
            st.bar_chart(category_totals)
        else:
            st.info("No expenses recorded yet. Input a note to start tracking.")

    st.divider()
    st.subheader("🧾 Recent Transactions")
    if os.path.exists(CSV_FILE):
        df_display = pd.read_csv(CSV_FILE)
        st.dataframe(df_display, use_container_width=True)

if __name__ == "__main__":
    main()