Spaces:
Sleeping
Sleeping
| """ | |
| PensionSight — FastAPI Backend | |
| Run: uvicorn main:app --reload --port 8000 | |
| """ | |
| import os | |
| import json | |
| import re | |
| from typing import Optional, List | |
| from fastapi import FastAPI, HTTPException | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from pydantic import BaseModel | |
| import google.generativeai as genai | |
| import uvicorn | |
| from dotenv import load_dotenv | |
| load_dotenv() | |
| from pension_engine import ( | |
| SubscriberProfile, | |
| PensionCalculationEngine, | |
| NPS_SCHEMES, | |
| SCENARIO_RETURNS, | |
| ) | |
| # ── App Setup ───────────────────────────────────────────────────────────────── | |
| app = FastAPI(title="PensionSight API", version="1.0.0") | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=["*"], | |
| allow_methods=["*"], | |
| allow_headers=["*"], | |
| ) | |
| engine = PensionCalculationEngine() | |
| API_KEY = os.getenv("GEMINI_API_KEY") | |
| if API_KEY: | |
| genai.configure(api_key=API_KEY) | |
| # ── Pydantic Models ─────────────────────────────────────────────────────────── | |
| class ProfileInput(BaseModel): | |
| current_age: int | |
| retirement_age: int | |
| monthly_contribution: float | |
| existing_corpus: float = 0 | |
| sector: str = "non_government" | |
| scheme: str = "LC50" | |
| annuity_percent: float = 0.40 | |
| annuity_rate: float = 0.06 | |
| annual_contribution_increase: float = 0 | |
| employer_contribution: float = 0 | |
| is_gig_worker: bool = False | |
| monthly_incomes: List[float] = [] | |
| deferral_age: Optional[int] = None | |
| desired_monthly_pension: Optional[float] = None | |
| class ReversePlanInput(BaseModel): | |
| desired_monthly_pension: float | |
| current_age: int | |
| retirement_age: int | |
| scheme: str = "LC50" | |
| scenario: str = "realistic" | |
| annuity_percent: float = 0.40 | |
| existing_corpus: float = 0 | |
| annual_step_up: float = 0 | |
| class ChatMessage(BaseModel): | |
| role: str | |
| content: str | |
| class ChatRequest(BaseModel): | |
| messages: List[ChatMessage] | |
| session_id: Optional[str] = None | |
| # ── Helpers ─────────────────────────────────────────────────────────────────── | |
| def fmt_inr(amount: float) -> str: | |
| if amount >= 1e7: | |
| return f"₹{amount/1e7:.2f} Cr" | |
| elif amount >= 1e5: | |
| return f"₹{amount/1e5:.2f} Lakh" | |
| return f"₹{amount:,.0f}" | |
| def safe_float(v, d=0.0): return float(v) if v is not None else d | |
| def safe_int(v, d=0): return int(v) if v is not None else d | |
| def profile_from_input(p: ProfileInput) -> SubscriberProfile: | |
| return SubscriberProfile( | |
| current_age=p.current_age, | |
| retirement_age=p.retirement_age, | |
| monthly_contribution=p.monthly_contribution, | |
| existing_corpus=p.existing_corpus, | |
| sector=p.sector, | |
| scheme=p.scheme, | |
| annuity_percent=p.annuity_percent, | |
| annuity_rate=p.annuity_rate, | |
| annual_contribution_increase=p.annual_contribution_increase, | |
| employer_contribution=p.employer_contribution, | |
| is_gig_worker=p.is_gig_worker, | |
| monthly_incomes=p.monthly_incomes, | |
| deferral_age=p.deferral_age, | |
| desired_monthly_pension=p.desired_monthly_pension, | |
| ) | |
| def scenario_to_dict(r, profile: SubscriberProfile) -> dict: | |
| ap = int(profile.annuity_percent * 100) | |
| lp = 100 - ap | |
| return { | |
| "scenario": r.scenario, | |
| "label": r.label, | |
| "blended_return": r.blended_return, | |
| "projected_corpus": r.projected_corpus, | |
| "real_corpus": r.real_corpus, | |
| "lumpsum": r.lumpsum_withdrawal, | |
| "lumpsum_pct": lp, | |
| "annuity_corpus": r.annuity_corpus, | |
| "annuity_pct": ap, | |
| "monthly_pension": r.monthly_pension, | |
| "real_monthly_pension": r.real_monthly_pension, | |
| "total_contributions": r.total_contributions, | |
| "wealth_gained": r.wealth_gained, | |
| "investment_years": r.investment_years, | |
| } | |
| # ── Routes ──────────────────────────────────────────────────────────────────── | |
| def root(): | |
| return {"status": "PensionSight API running", "version": "1.0.0"} | |
| def project(profile_input: ProfileInput): | |
| try: | |
| profile = profile_from_input(profile_input) | |
| results = engine.project_all_scenarios(profile) | |
| nudges = engine.generate_nudges(profile) | |
| return { | |
| "success": True, | |
| "profile": { | |
| "current_age": profile.current_age, | |
| "retirement_age": profile.retirement_age, | |
| "investment_years": profile.investment_years, | |
| "monthly_sip": profile.total_monthly_contribution, | |
| "scheme": NPS_SCHEMES[profile.scheme]["name"], | |
| "scheme_key": profile.scheme, | |
| "annual_step_up": profile.annual_contribution_increase, | |
| "annuity_percent": profile.annuity_percent, | |
| }, | |
| "scenarios": { | |
| key: scenario_to_dict(r, profile) | |
| for key, r in results.items() | |
| }, | |
| "nudges": [ | |
| { | |
| "type": n.nudge_type, | |
| "message": n.message, | |
| "impact": n.impact_rupees, | |
| "impact_fmt": fmt_inr(n.impact_rupees), | |
| "current": n.current_value, | |
| "improved": n.improved_value, | |
| } | |
| for n in nudges | |
| ], | |
| } | |
| except Exception as e: | |
| raise HTTPException(status_code=400, detail=str(e)) | |
| def reverse_plan(inp: ReversePlanInput): | |
| try: | |
| result = engine.reverse_plan( | |
| desired_monthly_pension=inp.desired_monthly_pension, | |
| current_age=inp.current_age, | |
| retirement_age=inp.retirement_age, | |
| scheme=inp.scheme, | |
| scenario=inp.scenario, | |
| annuity_percent=inp.annuity_percent, | |
| existing_corpus=inp.existing_corpus, | |
| annual_step_up=inp.annual_step_up, | |
| ) | |
| return { | |
| "success": True, | |
| "desired_monthly_pension": result.desired_monthly_pension, | |
| "desired_fmt": fmt_inr(result.desired_monthly_pension), | |
| "required_corpus": result.required_corpus, | |
| "required_corpus_fmt": fmt_inr(result.required_corpus), | |
| "required_monthly_sip": result.required_monthly_sip, | |
| "required_sip_fmt": fmt_inr(result.required_monthly_sip), | |
| "scenario": result.scenario, | |
| "investment_years": result.investment_years, | |
| "current_age": result.current_age, | |
| "retirement_age": result.retirement_age, | |
| } | |
| except Exception as e: | |
| raise HTTPException(status_code=400, detail=str(e)) | |
| def get_schemes(): | |
| return { | |
| key: { | |
| "name": val["name"], | |
| "max_equity": int(val["max_equity"] * 100), | |
| } | |
| for key, val in NPS_SCHEMES.items() | |
| } | |
| # ── Gemini Chat ─────────────────────────────────────────────────────────────── | |
| SYSTEM_INSTRUCTION = """ | |
| You are PensionSight, India's smartest NPS retirement planning assistant for PFRDA Innovate4NPS 2026. | |
| YOUR MOST IMPORTANT RULE: | |
| You MUST NEVER calculate, estimate, or guess any financial numbers (corpus, pension, SIP amounts). | |
| When you have all required data, emit a tool_call JSON block. The engine calculates — you explain. | |
| TOOL CALL FORMAT: | |
| ```tool_call | |
| { | |
| "tool": "project", | |
| "profile": { | |
| "current_age": 20, | |
| "retirement_age": 60, | |
| "monthly_contribution": 20000, | |
| "existing_corpus": 0, | |
| "sector": "non_government", | |
| "scheme": "LC75", | |
| "annuity_percent": 0.40, | |
| "annuity_rate": 0.06, | |
| "annual_contribution_increase": 5, | |
| "employer_contribution": 0, | |
| "is_gig_worker": true, | |
| "monthly_incomes": [80000, 80000, 80000], | |
| "deferral_age": null, | |
| "desired_monthly_pension": null | |
| } | |
| } | |
| ``` | |
| For reverse planning: | |
| ```tool_call | |
| { | |
| "tool": "reverse_plan", | |
| "desired_monthly_pension": 40000, | |
| "current_age": 28, | |
| "retirement_age": 60, | |
| "scheme": "LC50", | |
| "annuity_percent": 0.40, | |
| "existing_corpus": 0, | |
| "annual_step_up": 5 | |
| } | |
| ``` | |
| DATA TO COLLECT (1-2 questions at a time): | |
| 1. current_age (18–74) | |
| 2. retirement_age (60–75) | |
| 3. monthly_contribution (≥500) | |
| 4. sector (government/non_government; gig workers = non_government) | |
| 5. existing_corpus (0 if new) | |
| 6. scheme (LC25/LC50/LC75; recommend LC75 if age<35) | |
| 7. annuity_percent (0.40–1.0; min 40%) | |
| 8. annual_contribution_increase (0–20%; suggest 5%) | |
| 9. employer_contribution (0 for gig/self-employed) | |
| 10. is_gig_worker (true/false); if true → monthly_incomes list | |
| NPS 2025 RULES: | |
| - Min retirement: 60, max deferral: 75 | |
| - Min annuity: 40% (max lumpsum 60%, tax-free) | |
| - Government: employer 14% of basic (Budget 2025) | |
| - Schemes: LC25 (low), LC50 (moderate), LC75 (high equity), AGGRESSIVE, ACTIVE | |
| TONE: Warm, concise, no jargon. Max 2 questions per reply. Use ₹ symbol. | |
| After projection results, offer: (1) reverse planning, (2) nudges, (3) annuity split change. | |
| """ | |
| GEMINI_MODEL = None | |
| if API_KEY: | |
| GEMINI_MODEL = genai.GenerativeModel( | |
| model_name="gemini-2.5-flash-lite", | |
| system_instruction=SYSTEM_INSTRUCTION | |
| ) | |
| def extract_tool_call(text: str): | |
| for pattern in [ | |
| r"```tool_call\s*([\s\S]*?)```", | |
| r"```json\s*([\s\S]*?)```", | |
| r"```\s*(\{[\s\S]*?\})\s*```", | |
| ]: | |
| m = re.search(pattern, text) | |
| if m: | |
| try: | |
| d = json.loads(m.group(1).strip()) | |
| if "tool" in d: | |
| return d | |
| except Exception: | |
| pass | |
| return None | |
| def strip_code_blocks(text: str) -> str: | |
| return re.sub(r"```[\w]*\s*[\s\S]*?```", "", text).strip() | |
| def run_tool(tj: dict) -> dict: | |
| tool = tj.get("tool") | |
| if tool == "project": | |
| p = tj["profile"] | |
| inp = ProfileInput( | |
| current_age=safe_int(p.get("current_age"), 30), | |
| retirement_age=safe_int(p.get("retirement_age"), 60), | |
| monthly_contribution=safe_float(p.get("monthly_contribution"), 1000), | |
| existing_corpus=safe_float(p.get("existing_corpus")), | |
| sector=p.get("sector", "non_government"), | |
| scheme=p.get("scheme", "LC50"), | |
| annuity_percent=safe_float(p.get("annuity_percent"), 0.40), | |
| annuity_rate=safe_float(p.get("annuity_rate"), 0.06), | |
| annual_contribution_increase=safe_float(p.get("annual_contribution_increase")), | |
| employer_contribution=safe_float(p.get("employer_contribution")), | |
| is_gig_worker=bool(p.get("is_gig_worker", False)), | |
| monthly_incomes=list(p.get("monthly_incomes") or []), | |
| deferral_age=safe_int(p["deferral_age"]) if p.get("deferral_age") else None, | |
| ) | |
| return {"type": "project", "data": project(inp)} | |
| elif tool == "reverse_plan": | |
| inp = ReversePlanInput( | |
| desired_monthly_pension=safe_float(tj["desired_monthly_pension"]), | |
| current_age=safe_int(tj["current_age"]), | |
| retirement_age=safe_int(tj["retirement_age"]), | |
| scheme=tj.get("scheme", "LC50"), | |
| annuity_percent=safe_float(tj.get("annuity_percent"), 0.40), | |
| existing_corpus=safe_float(tj.get("existing_corpus")), | |
| annual_step_up=safe_float(tj.get("annual_step_up")), | |
| ) | |
| return {"type": "reverse_plan", "data": reverse_plan(inp)} | |
| return {"type": "error", "data": {"message": f"Unknown tool: {tool}"}} | |
| def chat(req: ChatRequest): | |
| if not GEMINI_MODEL: | |
| raise HTTPException(status_code=500, detail="GEMINI_API_KEY not configured") | |
| try: | |
| # Build Gemini history (exclude last user message) | |
| history = [] | |
| messages = req.messages | |
| for msg in messages[:-1]: | |
| history.append({"role": msg.role if msg.role != "assistant" else "model", | |
| "parts": [msg.content]}) | |
| chat_session = GEMINI_MODEL.start_chat(history=history) | |
| last_msg = messages[-1].content | |
| response = chat_session.send_message(last_msg) | |
| bot_text = response.text | |
| tool_json = extract_tool_call(bot_text) | |
| display_text = strip_code_blocks(bot_text) | |
| tool_result = None | |
| follow_up = None | |
| if tool_json: | |
| tool_result = run_tool(tool_json) | |
| # Get Gemini to interpret the result | |
| result_summary = json.dumps(tool_result["data"], indent=2)[:3000] | |
| fu_resp = chat_session.send_message( | |
| f"[TOOL RESULT]\n{result_summary}\n\n" | |
| "Summarise the realistic scenario for the user in 3-4 sentences. " | |
| "Then offer: (1) reverse planning, (2) nudges, (3) change annuity split." | |
| ) | |
| follow_up = fu_resp.text | |
| return { | |
| "success": True, | |
| "text": display_text, | |
| "tool_result": tool_result, | |
| "follow_up": follow_up, | |
| } | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| if __name__ == "__main__": | |
| uvicorn.run("app:app", host="0.0.0.0", port=7860) | |