Update app.py

app.py

@@ -1,23 +1,12 @@
 # main.py
 """
 Financial Health Score Service (FastAPI)
-
-Features:
-- Fetches user budgets and last-30-day transactions from MongoDB
-- Looks up transaction currency (from `currencies` collection) and embeds currency code
-- Builds a careful prompt for OpenAI (gpt-4o-mini) that instructs usage of currency code
-- Calls OpenAI and reliably extracts JSON output
-- Returns {"userId", "score", "explanation"}
-
-IMPORTANT:
-- Ensure MONGO_URI and OPENAI_API_KEY are set in your environment.
-- The `currencies` collection name is assumed to be "currencies".
-- This file uses the OpenAI Python client (OpenAI(api_key=...)) per your earlier setup.
+Now using GPT-4.1 with strict JSON mode.
+This guarantees the model ALWAYS returns valid JSON.
 """
 
 import json
 import os
-import re
 from datetime import datetime, timedelta
 
 from bson import ObjectId
@@ -43,12 +32,9 @@ if not OPENAI_API_KEY:
 mongo_client = MongoClient(MONGO_URI)
 default_db = mongo_client.get_default_database()
 
-if default_db is None:
-    raise RuntimeError("Unable to determine default database from MONGO_URI")
-
 budget_collection = default_db["budget"]
 transaction_collection = default_db["transactions"]
-currencies_collection = default_db["currencies"]  # <-- currencies collection
+currencies_collection = default_db["currencies"]
 
 # OpenAI client
 openai = OpenAI(api_key=OPENAI_API_KEY)
@@ -80,15 +66,14 @@ def normalize_budgets(budgets):
         head_categories = []
         heads = budget.get("headCategories") or []
 
-        if isinstance(heads, list):
-            for head in heads:
-                head_categories.append({
-                    "spendLimitType": head.get("spendLimitType"),
-                    "spendAmount": safe_number(head.get("spendAmount")),
-                    "maxAmount": safe_number(head.get("maxAmount")),
-                    "remainingAmount": safe_number(head.get("remainingAmount")),
-                    "notifications": head.get("notifications") or [],
-                })
+        for head in heads:
+            head_categories.append({
+                "spendLimitType": head.get("spendLimitType"),
+                "spendAmount": safe_number(head.get("spendAmount")),
+                "maxAmount": safe_number(head.get("maxAmount")),
+                "remainingAmount": safe_number(head.get("remainingAmount")),
+                "notifications": head.get("notifications") or [],
+            })
 
         normalized.append({
             "name": budget.get("name"),
@@ -106,42 +91,26 @@ def normalize_budgets(budgets):
 
 
 def normalize_transactions(transactions):
-    """
-    Trim transactions and attach currency code (e.g., 'EUR', 'INR', 'USD') when possible.
-    """
     trimmed = []
     for txn in transactions:
         date_value = txn.get("date")
+        date_str = date_value.date().isoformat() if isinstance(date_value, datetime) else None
 
-        if isinstance(date_value, datetime):
-            date_str = date_value.date().isoformat()
-        else:
-            date_str = None
-
-        # ---- Currency lookup ----
         currency_code = None
         currency_id = txn.get("currency")
 
         try:
-            # currency may be stored as an ObjectId already; handle strings too
+            # Handle ObjectId or string forms
             if isinstance(currency_id, ObjectId):
                 currency_doc = currencies_collection.find_one({"_id": currency_id})
-            elif isinstance(currency_id, dict) and "$oid" in currency_id:
-                # sometimes the raw export contains {"$oid": "..."}
-                try:
-                    currency_doc = currencies_collection.find_one({"_id": ObjectId(currency_id["$oid"])})
-                except Exception:
-                    currency_doc = None
             elif isinstance(currency_id, str):
-                try:
-                    currency_doc = currencies_collection.find_one({"_id": ObjectId(currency_id)})
-                except Exception:
-                    currency_doc = currencies_collection.find_one({"code": currency_id}) or currencies_collection.find_one({"currency": currency_id})
+                currency_doc = currencies_collection.find_one({"_id": ObjectId(currency_id)})
+            elif isinstance(currency_id, dict) and "$oid" in currency_id:
+                currency_doc = currencies_collection.find_one({"_id": ObjectId(currency_id["$oid"])})
             else:
                 currency_doc = None
 
             if currency_doc:
-                # Option A: use currency code only (e.g., "EUR")
                 currency_code = currency_doc.get("code") or currency_doc.get("currency")
 
         except Exception:
@@ -150,23 +119,23 @@ def normalize_transactions(transactions):
         trimmed.append({
             "type": txn.get("type"),
             "amount": safe_number(txn.get("amount")),
-            "currency": currency_code,  # <-- added
+            "currency": currency_code,
             "date": date_str,
         })
     return trimmed
 
 
 def score_prompt(budgets, transactions):
-    # We instruct the model to use currency code when mentioning amounts (Option A)
-    return f"""
-You are a succinct financial wellness expert. Using the budgets and last 30 days of transactions below,
-rate the user's financial health on a scale from 0 to 100 (higher is better).
+    return {
+        "role": "user",
+        "content": f"""
+You are a financial wellness expert. Using the budgets and last 30 days of transactions below, 
+rate the user’s financial health from 0 to 100 (higher = better).
 
-IMPORTANT:
-- When referring to monetary amounts, ALWAYS prefix with the currency code if available.
-  Example: "EUR 10,000", "INR 5,000", "USD 200".
-- If a transaction has no currency code, you may use the number only (e.g., 1000).
-- Keep the explanation short (one or two sentences) and directly related to budgets and transactions.
+Rules:
+- Always prefix amounts with currency code when available (e.g., INR 5000).
+- Keep the explanation short (1–2 sentences).
+- Consider income, expenses, spending consistency, and remaining budgets.
 
 Budgets:
 {json.dumps(normalize_budgets(budgets), indent=2)}
@@ -174,64 +143,10 @@ Budgets:
 Transactions (last 30 days):
 {json.dumps(normalize_transactions(transactions), indent=2)}
 
-Respond only with valid JSON, nothing else, using this exact shape:
-{{ "score": number, "explanation": "short explanation" }}
+Respond ONLY using this exact JSON structure:
+{{ "score": number, "explanation": "string" }}
 """
-    
-
-def extract_json_payload(text):
-    """Extract JSON from plain text or fenced code blocks."""
-    trimmed = (text or "").strip()
-
-    # try fenced json block first
-    fenced = re.search(r"```(?:json)?\s*([\s\S]*?)```", trimmed)
-    if fenced:
-        return json.loads(fenced.group(1).strip())
-
-    # try to find first { ... } substring
-    first_obj = re.search(r"(\{[\s\S]*\})", trimmed)
-    if first_obj:
-        return json.loads(first_obj.group(1))
-
-    # last resort: direct JSON load
-    return json.loads(trimmed)
-
-
-# ===========================
-# BULLETPROOF OPENAI EXTRACTOR
-# ===========================
-def extract_openai_text(response):
-    """
-    Robust extractor for OpenAI SDK responses.
-    Handles several possible message wrappers and returns the assistant text.
-    """
-    try:
-        # Best-effort to access nested choice message content
-        choices = getattr(response, "choices", None) or response.get("choices") if isinstance(response, dict) else None
-        if not choices:
-            # fallback: maybe response is a dict-like structure
-            return str(response)
-
-        msg = choices[0].get("message") if isinstance(choices[0], dict) else getattr(choices[0], "message", None)
-        if not msg:
-            return str(choices[0])
-
-        # If message exposes 'content'
-        if isinstance(msg, dict) and "content" in msg:
-            return msg["content"].strip()
-        if hasattr(msg, "content"):
-            return msg.content.strip()
-
-        # If message is a repr like ChatCompletionMessage(content='...'), extract via regex
-        msg_str = str(msg)
-        match = re.search(r"content='([\s\S]*?)'", msg_str)
-        if match:
-            return match.group(1).strip()
-
-        # fallback
-        return msg_str.strip()
-    except Exception:
-        return str(response)
+    }
 
 
 # ===========================
@@ -250,7 +165,7 @@ def financial_score(payload: ScoreRequest):
     except Exception:
         raise HTTPException(status_code=400, detail="Invalid userId")
 
-    # Fetch budgets (all budgets created by this user)
+    # Fetch budgets
     budgets = list(budget_collection.find({"createdBy": user_id}))
 
     # Fetch last 30 days transactions
@@ -259,10 +174,9 @@ def financial_score(payload: ScoreRequest):
         transaction_collection.find({
             "user": user_id,
             "date": {"$gte": thirty_days_ago}
-        }).sort("date", -1).limit(100)
+        }).sort("date", -1).limit(120)
     )
 
-    # If neither budgets nor recent transactions exist -> score 0
     if not budgets and not transactions:
         return {
             "userId": payload.userId,
@@ -270,45 +184,29 @@ def financial_score(payload: ScoreRequest):
             "explanation": "No budgets or recent transactions found."
         }
 
-    prompt = score_prompt(budgets, transactions)
+    # Build messages
+    messages = [score_prompt(budgets, transactions)]
 
-    # Call OpenAI
+    # -- Strict JSON Mode using GPT-4.1 --
     try:
         response = openai.chat.completions.create(
-            model="gpt-4o-mini",
-            temperature=0.6,
-            messages=[{"role": "user", "content": prompt}],
+            model="gpt-4.1",
+            response_format={"type": "json_object"},   # 🔥 Guarantees valid JSON
+            messages=messages,
+            temperature=0.4,
         )
     except Exception as exc:
         raise HTTPException(status_code=502, detail=f"OpenAI request failed: {exc}")
 
-    model_output = extract_openai_text(response)
+    # Direct JSON — no more parsing issues
+    parsed = response.choices[0].message.parsed
 
-    # Parse JSON payload from model output
-    try:
-        parsed = extract_json_payload(model_output)
-    except Exception:
-        raise HTTPException(
-            status_code=502,
-            detail={
-                "error": "Unable to parse OpenAI response as JSON",
-                "rawResponse": model_output
-            },
-        )
+    # Validate shapes
+    score_val = parsed.get("score", 0)
+    explanation = parsed.get("explanation", "")
 
-    # Validate required fields
-    if "score" not in parsed or "explanation" not in parsed:
-        raise HTTPException(
-            status_code=502,
-            detail={
-                "error": "OpenAI response missing required fields",
-                "rawResponse": model_output
-            }
-        )
-
-    # Clamp score to 0..100
     try:
-        score_val = int(float(parsed["score"]))
+        score_val = int(float(score_val))
         score_val = max(0, min(100, score_val))
     except Exception:
         score_val = 0
@@ -316,6 +214,5 @@ def financial_score(payload: ScoreRequest):
     return {
         "userId": payload.userId,
         "score": score_val,
-        "explanation": parsed["explanation"]
+        "explanation": explanation
     }
-