src/streamlit_app.py

# personal_finance_chatbot.py
import streamlit as st
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import json
from datetime import datetime

# Configuration
MODEL_NAME = "ibm-granite/granite-7b-base"  # Correct Granite HF name
USER_TYPES = ["student", "professional"]

# Initialize NLP pipeline
@st.cache_resource
def load_model():
    """Load and cache Granite model for text generation"""
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
    return pipeline("text-generation", model=model, tokenizer=tokenizer)

# ------------------ USER PROFILE ------------------
class UserProfile:
    def __init__(self, user_type, financial_goals=None, income=0, expenses=None):
        self.user_type = user_type
        self.financial_goals = financial_goals or []
        self.income = income
        self.expenses = expenses or {}
        self.transaction_history = []

    def add_transaction(self, amount, category, description=""):
        """Record a financial transaction"""
        self.transaction_history.append({
            "date": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            "amount": amount,
            "category": category,
            "description": description
        })

    def get_budget_summary(self):
        """Generate a budget summary"""
        total_expenses = sum(t["amount"] for t in self.transaction_history if t["amount"] < 0)
        total_income = sum(t["amount"] for t in self.transaction_history if t["amount"] > 0)

        return {
            "total_income": total_income,
            "total_expenses": abs(total_expenses),
            "net_savings": total_income + total_expenses,  # total_expenses is negative already
            "category_breakdown": self._get_category_breakdown()
        }

    def _get_category_breakdown(self):
        breakdown = {}
        for t in self.transaction_history:
            if t["amount"] < 0:
                cat = t["category"]
                breakdown[cat] = breakdown.get(cat, 0) + abs(t["amount"])
        return breakdown

# ------------------ CHATBOT CORE ------------------
class FinanceChatbot:
    def __init__(self):
        self.nlp = load_model()
        self.user_profiles = {}
        self.current_user = None

    def set_user(self, user_id, user_type):
        if user_id not in self.user_profiles:
            self.user_profiles[user_id] = UserProfile(user_type=user_type)
        self.current_user = user_id

    def generate_response(self, query):
        if not self.current_user:
            return "⚠️ Please set up your profile first (student or professional)."

        profile = self.user_profiles[self.current_user]
        context = self._build_context(profile)

        tone_instruction = (
            "Use simple, encouraging language for a student."
            if profile.user_type == "student" else
            "Use concise, professional language for a working professional."
        )

        prompt = f"""
        You are an AI-powered financial assistant.
        User profile: {context}
        Instruction: {tone_instruction}

        User asked: "{query}"

        Respond with:
        1. Direct and clear answer
        2. 1-2 actionable suggestions
        3. Keep it under 3 sentences unless more detail is needed.
        """

        try:
            result = self.nlp(prompt, max_new_tokens=200, do_sample=True, temperature=0.7)
            response = result[0]['generated_text'].replace(prompt, "").strip()
            return response
        except Exception as e:
            return f"❌ Error: {str(e)}"

    def _build_context(self, profile):
        budget = profile.get_budget_summary()
        return json.dumps({
            "user_type": profile.user_type,
            "income": profile.income,
            "net_savings": budget["net_savings"],
            "top_expenses": sorted(budget["category_breakdown"].items(),
                                   key=lambda x: x[1], reverse=True)[:3],
            "recent_transactions": profile.transaction_history[-3:]
        })

    def analyze_spending(self):
        if not self.current_user:
            return "⚠️ No user profile selected."

        profile = self.user_profiles[self.current_user]
        budget = profile.get_budget_summary()
        if not budget["category_breakdown"]:
            return "ℹ️ No spending data yet."

        prompt = f"""
        Analyze the spending breakdown: {json.dumps(budget['category_breakdown'])}.
        User type: {profile.user_type}, Income: {profile.income}.

        Provide:
        1. One key spending insight
        2. One actionable saving tip
        3. Tone adapted for {profile.user_type}
        """

        try:
            result = self.nlp(prompt, max_new_tokens=150, do_sample=True, temperature=0.7)
            return result[0]['generated_text'].replace(prompt, "").strip()
        except Exception as e:
            return f"❌ Error: {str(e)}"

# ------------------ STREAMLIT UI ------------------
def main():
    st.set_page_config(page_title="Personal Finance Chatbot", layout="wide")

    if 'chatbot' not in st.session_state:
        st.session_state.chatbot = FinanceChatbot()
    if 'user_id' not in st.session_state:
        st.session_state.user_id = None
    if 'messages' not in st.session_state:
        st.session_state.messages = []

    # Sidebar
    with st.sidebar:
        st.title("👤 User Profile")
        user_id = st.text_input("Your ID", value=st.session_state.get('user_id', ''))
        user_type = st.selectbox("I am a:", USER_TYPES)

        if st.button("Save Profile"):
            st.session_state.user_id = user_id
            st.session_state.chatbot.set_user(user_id, user_type)
            st.success(f"Profile saved as {user_type}")

        st.divider()
        st.subheader("📊 Quick Actions")

        if st.session_state.user_id:
            if st.button("View Budget Summary"):
                profile = st.session_state.chatbot.user_profiles[st.session_state.user_id]
                summary = profile.get_budget_summary()
                st.session_state.messages.append(
                    {"role": "assistant", "content":
                        f"### Budget Summary\n- Income: ${summary['total_income']:.2f}\n"
                        f"- Expenses: ${summary['total_expenses']:.2f}\n"
                        f"- Net Savings: ${summary['net_savings']:.2f}\n"
                        f"- Top Expenses: {summary['category_breakdown']}"}
                )

            if st.button("Get Spending Insights"):
                insights = st.session_state.chatbot.analyze_spending()
                st.session_state.messages.append(
                    {"role": "assistant", "content": f"### Spending Insights\n{insights}"}
                )

    # Main chat
    st.title("💰 Personal Finance Chatbot")
    st.write("Ask me about **savings, taxes, investments, or budgeting!**")

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

    if prompt := st.chat_input("Type your financial question here..."):
        if not st.session_state.user_id:
            st.error("⚠️ Please set up your profile first!")
        else:
            st.session_state.messages.append({"role": "user", "content": prompt})
            with st.chat_message("user"):
                st.markdown(prompt)

            with st.spinner("💡 Thinking..."):
                response = st.session_state.chatbot.generate_response(prompt)
                with st.chat_message("assistant"):
                    st.markdown(response)
                st.session_state.messages.append({"role": "assistant", "content": response})

if __name__ == "__main__":
    main()