BuddyMath / main.py
dotandru's picture
Force Rebuild: Updated restart trigger in main.py
f4a6b4d
Raw
History Blame
27.2 kB
# main.py - V5.10.1 (RESTART TRIGGER: 2026-03-13 20:45)
from contextlib import asynccontextmanager
from fastapi import FastAPI, UploadFile, File, Form, HTTPException, Request, Header
from fastapi.responses import JSONResponse, HTMLResponse
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from sse_starlette.sse import EventSourceResponse
from fastapi.templating import Jinja2Templates
from typing import Optional, List
import logging
import base64
import json
import io
import sys
import os
import asyncio
from pydantic import BaseModel
from typing import Any
class AskQuestionRequest(BaseModel):
context_data: dict | Any
question: str
student_name: str = "תלמיד"
id_token: Optional[str] = None
# --- HEALTH CHECK : Top-level Dependency Verification ---
# We do this before standard imports to ensure a clear error message
# if the virtual environment is inactive and libraries are missing.
try:
import cv2
import numpy as np
except ModuleNotFoundError as e:
print(f"🔥 [HEALTH-CHECK FAILED] Missing critical dependency: {e}. Are you running inside the .venv?")
sys.exit(1)
from orchestrator import orchestrator, build_standard_response
from quota_system import quota_manager
from quota_system_v2 import quota_manager_v2
from config import IS_PRODUCTION, ENV
from firebase_manager import firebase_manager
if hasattr(sys.stdout, 'reconfigure'):
sys.stdout.reconfigure(encoding='utf-8')
# הגדרת לוגר HamoraServer
try:
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
handlers=[
logging.FileHandler("/tmp/server.log", encoding="utf-8"),
logging.StreamHandler(sys.stdout)
]
)
except PermissionError:
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
handlers=[logging.StreamHandler(sys.stdout)]
)
logger = logging.getLogger("HamoraServer")
# --- INFRA HARDENING: Global Async Exception Handler ---
def custom_async_exception_handler(loop, context):
"""
Catches unhandled asynchronous exceptions to prevent the Event Loop from crashing.
"""
msg = context.get("exception", context["message"])
logger.critical(f"🚨 [ASYNC-CRASH-PREVENTION] Caught unhandled exception in Event Loop: {msg}")
# The loop remains alive. We just log the critical error.
# --- INFRA HARDENING: Health Check Function ---
def verify_system_health():
"""
Verifies execution environment and critical dependencies.
"""
logger.info("🩺 [HEALTH-CHECK] Verifying core dependencies and environment...")
# Check if running in a virtual environment
if sys.prefix == sys.base_prefix:
logger.warning("⚠️ [HEALTH-CHECK] Not running inside a virtual environment (.venv). Proceeding anyway...")
logger.info(f"✅ [HEALTH-CHECK] cv2 version: {cv2.__version__}, numpy: {np.__version__}")
logger.info(f"✅ [HEALTH-CHECK] Environment: {ENV.upper()}, Production Mode: {IS_PRODUCTION}")
# V5.8.1: API Key Validation
if not os.environ.get("GOOGLE_API_KEY"):
logger.error("❌ [HEALTH-CHECK] GOOGLE_API_KEY is missing! Gemini calls will fail.")
else:
logger.info("✅ [HEALTH-CHECK] GOOGLE_API_KEY is detected.")
def _ensure_dev_admin_exists():
"""V5.9.6: Ensures dev-bypass-user exists with admin role in DEV mode."""
if IS_PRODUCTION:
return
try:
db = firebase_manager.get_db()
if not db:
return
uid = "dev-bypass-user"
user_ref = db.collection('users').document(uid)
doc = user_ref.get()
# We always merge to ensure roles are correct even if doc exists
user_ref.set({
'uid': uid,
'name': 'Dev Bypass User (Admin)',
'isAdmin': True,
'role': 'admin',
'status': 'approved',
'tier': 'student_premium',
'monthly_token_budget': 9999999
}, merge=True)
logger.info(f"🛡️ [STARTUP] Dev admin user '{uid}' provisioned/verified.")
except Exception as e:
logger.error(f"❌ [STARTUP] Failed to provision dev admin: {e}")
# --- INFRA HARDENING: Lifespan Context Manager ---
@asynccontextmanager
async def lifespan(app: FastAPI):
# Startup Phase
verify_system_health()
# Register Global Async Exception Handler
loop = asyncio.get_running_loop()
loop.set_exception_handler(custom_async_exception_handler)
logger.info("🛡️ [STARTUP] Global Async Exception Handler registered.")
# V5.9.6: Provision dev admin if needed
_ensure_dev_admin_exists()
yield # Yield control back to FastAPI
# Shutdown Phase
logger.info("🛑 [SHUTDOWN] BuddyMath Server is shutting down cleanly.")
# Application Setup
app = FastAPI(title="BuddyMath Server - OpenCV Engine", lifespan=lifespan)
# Setup Jinja templates
base_dir = os.path.dirname(os.path.abspath(__file__))
templates = Jinja2Templates(directory=os.path.join(base_dir, "templates"))
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# Static files for audio fallback
os.makedirs("/tmp/static", exist_ok=True)
app.mount("/static", StaticFiles(directory="/tmp/static"), name="static")
@app.get("/")
async def root():
return {"status": f"BuddyMath API V5.8.0 ({ENV.upper()})", "engine": "OpenCV Base + Infra Hardening"}
@app.get("/admin/stats")
async def get_admin_stats(request: Request):
"""
V5.9.4: Returns usage and cost statistics from usage.jsonl.
Requires Admin authorization.
"""
auth_header = request.headers.get('Authorization')
if not auth_header or not auth_header.startswith('Bearer '):
return JSONResponse(status_code=401, content={"error": "Unauthorized"})
token = auth_header.split('Bearer ')[1].strip()
# Dev bypass check for stats too
from config import DEV_BYPASS_TOKEN
if not IS_PRODUCTION and token == DEV_BYPASS_TOKEN:
pass # OK
else:
decoded = firebase_manager.verify_token(token)
if not decoded:
return JSONResponse(status_code=401, content={"error": "Invalid token"})
# Check if user is admin
uid = decoded.get('uid')
db = firebase_manager.get_db()
user_doc = db.collection('users').document(uid).get()
user_data = user_doc.to_dict() if user_doc.exists else {}
if user_data.get('role') != 'admin' and not user_data.get('isAdmin'):
return JSONResponse(status_code=403, content={"error": "Admin access required"})
from cost_tracker import LOG_FILE
stats = {
"total_input_tokens": 0,
"total_output_tokens": 0,
"total_cost_usd": 0.0,
"request_count": 0,
"total_users": 0
}
# Bug 1 Fix: Count users from Firestore directly
try:
db = firebase_manager.get_db()
users_ref = db.collection('users')
# Use an aggregation query for efficiency if available, or just count
stats["total_users"] = len(list(users_ref.stream())) # Simple for now, can be optimized
except Exception as e:
logger.error(f"Error counting users: {e}")
if os.path.exists(LOG_FILE):
try:
with open(LOG_FILE, "r", encoding="utf-8") as f:
for line in f:
entry = json.loads(line)
stats["total_input_tokens"] += entry.get("input_tokens", 0)
stats["total_output_tokens"] += entry.get("output_tokens", 0)
stats["request_count"] += 1
# Calculate total cost
from cost_tracker import PRICING
stats["total_cost_usd"] = (stats["total_input_tokens"] / 1e6 * PRICING["input"]) + \
(stats["total_output_tokens"] / 1e6 * PRICING["output"])
except Exception as e:
logger.error(f"Error parsing log file: {e}")
return {"cost_summary": stats}
@app.post("/admin/delete_user/{uid}")
async def delete_user_full(uid: str, request: Request):
"""
V5.9.5: Deletes user from both Firestore and Firebase Auth.
Requires Admin authorization.
"""
auth_header = request.headers.get('Authorization')
if not auth_header or not auth_header.startswith('Bearer '):
return JSONResponse(status_code=401, content={"error": "Unauthorized"})
token = auth_header.split('Bearer ')[1].strip()
# Admin check logic (reused from stats)
from config import DEV_BYPASS_TOKEN
if not IS_PRODUCTION and token == DEV_BYPASS_TOKEN:
pass # OK
else:
decoded = firebase_manager.verify_token(token)
if not decoded or not decoded.get('uid'):
return JSONResponse(status_code=401, content={"error": "Invalid token"})
admin_uid = decoded.get('uid')
db = firebase_manager.get_db()
user_doc = db.collection('users').document(admin_uid).get()
user_data = user_doc.to_dict() if user_doc.exists else {}
if user_data.get('role') != 'admin' and not user_data.get('isAdmin'):
return JSONResponse(status_code=403, content={"error": "Admin access required"})
try:
# 1. Delete from Firestore
db = firebase_manager.get_db()
db.collection('users').document(uid).delete()
# 2. Delete from Firebase Auth
try:
from firebase_admin import auth
auth.delete_user(uid)
logger.info(f"🗑️ [ADMIN] User {uid} deleted from Auth and Firestore.")
except auth.UserNotFoundError:
logger.warning(f"⚠️ [ADMIN] User {uid} not found in Auth, but deleted from Firestore.")
except Exception as e:
logger.error(f"❌ [ADMIN] Auth deletion failed for {uid}: {e}")
# We continue because the DB part is done
return {"status": "success", "message": f"User {uid} removed."}
except Exception as e:
logger.error(f"❌ [ADMIN] Deletion failed: {e}")
return JSONResponse(status_code=500, content={"error": str(e)})
@app.post("/solve_stream")
async def solve_stream(
user: Optional[str] = Form(None),
student_name: Optional[str] = Form(None),
grade: str = Form("י'"),
student_gender: str = Form("M"),
mode: str = Form("solve"),
user_note: Optional[str] = Form(None),
file: UploadFile = File(...)
):
"""
V5.8.0: המורה למתמטיקה - Multipart & OpenCV Base.
מקבל קובץ ישירות מהפלאטר ומפענח אותו עם OpenCV.
"""
final_student_name = student_name or user or "תלמיד"
print(f"🚀 🟢 BIT-LOG: Received Multipart request from {final_student_name}. Grade: {grade}")
# Quota Check
if final_student_name == "dev-bypass-user":
is_allowed, msg, current_usage, limit = True, "Dev Bypass", 0, 999
else:
is_allowed, msg, current_usage, limit = quota_manager.check_limit(final_student_name)
if not is_allowed:
response_content = build_standard_response(
final_answer=f"הגעת למכסה היומית ({limit} שאלות)",
teacher_summary="נא להמתין למחר לקבלת מכסה חדשה.",
logic_error=True,
response_type="error"
)
response_content["error"] = "QUOTA_EXCEEDED"
return JSONResponse(status_code=429, content=response_content)
quota_manager.increment_usage(final_student_name)
try:
# 1. קריאת הבינארי
image_bytes = await file.read()
print(f"📸 [BIT-LOG] Image received. Size: {len(image_bytes)} bytes")
# 2. OpenCV Decoder
nparr = np.frombuffer(image_bytes, np.uint8)
img_cv2 = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
if img_cv2 is None:
print("❌ [BIT-LOG] OpenCV failed to decode image!")
raise HTTPException(status_code=400, detail="Invalid image data")
print(f"✅ [BIT-LOG] OpenCV Matrix Ready: {img_cv2.shape}")
# 3. OCR & Solving Pipeline (Streaming)
print("🚀 [TRACE-MAIN] Initiating streaming orchestrator.solve_problem...")
async def event_generator():
try:
async for event in orchestrator.solve_problem(
problem_text="", # Will be extracted by OCR
grade=grade,
student_name=final_student_name,
student_gender=student_gender,
user_note=user_note,
image_data=image_bytes,
mode=mode,
uid=uid
):
# SSE Protocol: yield a dict with "data" key
yield {
"event": "message",
"id": event.question_id,
"data": event.model_dump_json() # Pydantic v2
}
except Exception as e:
logger.error(f"STREAMING ERROR: {e}")
yield {
"event": "error",
"data": json.dumps({"error": str(e)})
}
return EventSourceResponse(event_generator())
except Exception as e:
logger.exception("CRITICAL FLOW ERROR")
print(f"🔥 [BIT-LOG] CRITICAL ERROR: {str(e)}")
import traceback
traceback.print_exc()
response_content = build_standard_response(
final_answer="שגיאה בפענוח התמונה או התרגיל",
teacher_summary="המורה למתמטיקה מתנצל, אך חלה שגיאה לא צפויה.",
logic_error=True,
response_type="error"
)
return JSONResponse(status_code=500, content=response_content)
@app.post("/v2/solve_stream")
async def solve_stream_v2(
request: Request,
user: Optional[str] = Form(None),
student_name: Optional[str] = Form(None),
grade: str = Form("י'"),
student_gender: str = Form("M"),
mode: str = Form("solve"),
user_note: Optional[str] = Form(None),
files: List[UploadFile] = File(...)
):
"""
V2: Token-based Auth and Firestore Quota Management.
"""
auth_header = request.headers.get('Authorization')
if not auth_header or not auth_header.startswith('Bearer '):
logger.warning("🚨 [V2_ENDPOINT] Missing or invalid Authorization header.")
return JSONResponse(status_code=401, content={"error": "Unauthorized: Missing Token"})
id_token = auth_header.split('Bearer ')[1].strip()
uid = None
from config import IS_PRODUCTION, DEV_BYPASS_TOKEN
# V5.9.3: Log incoming token for debugging
logger.info(f"🔑 [V2_ENDPOINT] Received Token (len={len(id_token)}): {id_token[:5]}...")
# V5.9.1: DEV Auth Bypass (Strict non-prod check)
if not IS_PRODUCTION and id_token == DEV_BYPASS_TOKEN:
logger.info("🛠️ [V2_ENDPOINT] Using DEV Auth Bypass Token.")
uid = "dev-bypass-user"
else:
decoded_token = firebase_manager.verify_token(id_token)
if not decoded_token:
logger.warning("🚨 [V2_ENDPOINT] Invalid or expired Firebase ID token.")
return JSONResponse(status_code=401, content={"error": "Unauthorized: Invalid Token"})
uid = decoded_token.get('uid')
final_student_name = student_name or user or "תלמיד"
print(f"🚀 🟢 [V2] Received request from UID: {uid} ({final_student_name}). Grade: {grade}")
device_id = request.headers.get('Device-ID')
# V5.10.0: Fetch user tier for Digital Binder (History) support
user_tier = "student_basic"
try:
db = firebase_manager.get_db()
user_doc = db.collection('users').document(uid).get()
if user_doc.exists:
user_tier = user_doc.to_dict().get('tier', 'student_basic')
except Exception as e:
logger.error(f"Error fetching user tier: {e}")
# V2 Quota Check (Firestore)
if uid == "dev-bypass-user":
is_allowed, msg, current_usage, limit = True, "Dev Bypass", 0, 999
else:
is_allowed, msg, current_usage, limit = quota_manager_v2.check_limit(uid, device_id=device_id)
if not is_allowed:
response_content = build_standard_response(
final_answer=f"הגעת למכסה היומית ({limit} שאלות)",
teacher_summary="נא להמתין למחר לקבלת מכסה חדשה או לשדרג לפרימיום.",
logic_error=True,
response_type="error"
)
response_content["error"] = "QUOTA_EXCEEDED"
return JSONResponse(status_code=403, content=response_content) # Changed to 403 Forbidden for quota specifically
# Only increment usage if OCR/Solving process starts successfully
try:
# 1. קריאת הבינארי
image_bytes_list = []
for single_file in files:
image_bytes_list.append(await single_file.read())
print(f"📸 [V2-LOG] Received {len(image_bytes_list)} images.")
if not image_bytes_list:
raise HTTPException(status_code=400, detail="No images provided")
# 2. OpenCV Decoder (validate the first image)
nparr = np.frombuffer(image_bytes_list[0], np.uint8)
img_cv2 = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
if img_cv2 is None:
print("❌ [V2-LOG] OpenCV failed to decode first image!")
raise HTTPException(status_code=400, detail="Invalid image data")
print(f"✅ [V2-LOG] OpenCV Matrix Ready: {img_cv2.shape}")
# Increment quota AFTER we are sure the image is valid
quota_manager_v2.increment_usage(uid)
# 3. OCR & Solving Pipeline (Streaming)
print("🚀 [TRACE-V2] Initiating streaming orchestrator.solve_problem with multiple images...")
async def event_generator():
import cost_tracker
cost_tracker.current_request_tokens.set(0)
try:
# orchestrator handles the rest using student_name for pedagogical stuff, but quota is already handled.
async for event in orchestrator.solve_problem(
problem_text="", # Will be extracted by OCR
grade=grade,
student_name=final_student_name,
student_gender=student_gender,
user_note=user_note,
image_data_list=image_bytes_list,
mode=mode,
uid=uid,
tier=user_tier # V5.10.0: Pass tier for history saving
):
# SSE Protocol: yield a dict with "data" key
yield {
"event": "message",
"id": event.question_id,
"data": event.model_dump_json() # Pydantic v2
}
except Exception as e:
logger.error(f"STREAMING ERROR (V2): {e}")
yield {
"event": "error",
"data": json.dumps({"error": str(e)})
}
finally:
total_tokens = cost_tracker.current_request_tokens.get()
if total_tokens > 0:
quota_manager_v2.increment_usage(uid, increment_questions=0, tokens_used=total_tokens)
print(f"🪙 [V2-QUOTA] Deducted {total_tokens} tokens for UID: {uid}")
return EventSourceResponse(event_generator())
except Exception as e:
logger.exception("CRITICAL FLOW ERROR (V2)")
print(f"🔥 [V2-LOG] CRITICAL ERROR: {str(e)}")
import traceback
traceback.print_exc()
response_content = build_standard_response(
final_answer="שגיאה בפענוח התמונה או התרגיל",
teacher_summary="המורה למתמטיקה מתנצל, אך חלה שגיאה לא צפויה.",
logic_error=True,
response_type="error"
)
return JSONResponse(status_code=500, content=response_content)
@app.post("/explain_step")
async def explain_step(request: Request):
data = await request.json()
# Auth Hardening (V3.1)
auth_header = request.headers.get('Authorization')
token = (data.get("id_token") or (auth_header.split('Bearer ')[1] if auth_header and auth_header.startswith('Bearer ') else None))
if not token:
return JSONResponse(status_code=401, content={"error": "Unauthorized: Missing Token"})
decoded_token = firebase_manager.verify_token(token)
if not decoded_token:
return JSONResponse(status_code=401, content={"error": "Unauthorized: Invalid Token"})
uid = decoded_token.get('uid')
# Quota check (Gate only, no increment for simple explanations yet)
is_allowed, msg, _, _ = quota_manager_v2.check_limit(uid)
if not is_allowed:
return JSONResponse(status_code=403, content={"error": "QUOTA_EXCEEDED", "message": msg})
res = await orchestrator.explain_specific_step(data.get("context"), data.get("step_text"), data.get("student_name"))
return JSONResponse(content=res)
@app.post("/ask_question")
async def ask_question(request: Request, ask_req: AskQuestionRequest):
data = ask_req.dict()
# Auth Hardening (V3.1)
auth_header = request.headers.get('Authorization')
token = (data.get("id_token") or (auth_header.split('Bearer ')[1] if auth_header and auth_header.startswith('Bearer ') else None))
if not token:
return JSONResponse(status_code=401, content={"error": "Unauthorized: Missing Token"})
decoded_token = firebase_manager.verify_token(token)
if not decoded_token:
return JSONResponse(status_code=401, content={"error": "Unauthorized: Invalid Token"})
uid = decoded_token.get('uid')
# Quota check
is_allowed, msg, _, _ = quota_manager_v2.check_limit(uid)
if not is_allowed:
return JSONResponse(status_code=403, content={"error": "QUOTA_EXCEEDED", "message": msg})
res = await orchestrator.ask_question(data.get("context_data"), data.get("question"), data.get("student_name"))
return JSONResponse(content=res)
@app.get("/pay", response_class=HTMLResponse)
async def payment_page(request: Request, uid: Optional[str] = None):
"""
Serves the Premium Payment Web Page.
Fetches the student name from Firestore if uid is provided.
"""
student_name = ""
if uid:
try:
db = firebase_manager.get_db()
user_doc = db.collection('users').document(uid).get()
if user_doc.exists:
student_name = user_doc.to_dict().get("student_name", "")
except Exception as e:
logger.error(f"Failed to fetch user for payment page: {e}")
return templates.TemplateResponse("payment.html", {
"request": request,
"uid": uid or "",
"student_name": student_name
})
class UpgradeRequest(BaseModel):
uid: str
parent_email: str
@app.post("/api/upgrade_success")
async def upgrade_success(req: UpgradeRequest):
"""
Mock Webhook for successful payment.
Updates the Firestore user to Premium Tier.
"""
if not req.uid:
raise HTTPException(status_code=400, detail="Missing user ID")
try:
db = firebase_manager.get_db()
user_ref = db.collection('users').document(req.uid)
user_ref.set({
"tier": "parent_premium",
"parent_email": req.parent_email,
"monthly_token_budget": 2800000
}, merge=True)
logger.info(f"🎉 UPGRADED USER {req.uid} to parent_premium")
return {"status": "success", "message": "User upgraded successfully"}
except Exception as e:
logger.error(f"Failed to upgrade user {req.uid}: {e}")
raise HTTPException(status_code=500, detail="Database update failed")
class MigrationRequest(BaseModel):
batch_size: Optional[int] = 50
@app.post("/admin/migrate_to_v2")
async def migrate_to_v2(req: MigrationRequest, x_api_key: str = Header(None)):
"""
V280.2: Admin endpoint to migrate users to V2 Quota system.
"""
from config import FIREBASE_CREDENTIALS_JSON
import os
# Just a basic safety check (in production, use a real secret)
admin_secret = os.environ.get("ADMIN_SECRET_KEY", "buddy_admin_123")
if x_api_key != admin_secret:
raise HTTPException(status_code=401, detail="Unauthorized")
try:
from scripts.migrate_users_to_cloud import migrate_users
import asyncio
# Run migration in background so we don't block
asyncio.create_task(asyncio.to_thread(migrate_users))
return {"status": "success", "message": "Migration started in background."}
except Exception as e:
logger.error(f"Migration error: {e}")
raise HTTPException(status_code=500, detail=str(e))
class ReportRequest(BaseModel):
uid: str
student_name: str
parent_email: str
week_id: str
@app.post("/v2/send_weekly_report")
async def send_weekly_report(req: ReportRequest):
"""
Generates and emails the weekly AI Assessment report to the parent.
"""
try:
import os
from report_generator import report_generator
# 1. Produce HTML
html_content = report_generator.produce_weekly_report(
uid=req.uid,
week_id=req.week_id,
student_name=req.student_name
)
# 2. Generate PDF
pdf_path = f"/tmp/report_{req.uid}_{req.week_id}.pdf"
# Create tmp dir if it doesn't exist (local dev)
os.makedirs(os.path.dirname(pdf_path), exist_ok=True)
report_generator.export_to_pdf(html_content, pdf_path)
# 3. Email PDF
success = report_generator.send_report_email(
parent_email=req.parent_email,
student_name=req.student_name,
pdf_path=pdf_path
)
# Cleanup
try:
if os.path.exists(pdf_path):
os.remove(pdf_path)
except Exception as e:
print(f"Cleanup failed for {pdf_path}: {e}")
if success:
return {"status": "success", "message": f"Report sent to {req.parent_email}"}
return JSONResponse(status_code=500, content={"status": "error", "message": "Failed to send email via SendGrid."})
except Exception as e:
logger.error(f"Failed to generate report: {e}")
import traceback
traceback.print_exc()
return JSONResponse(status_code=500, content={"status": "error", "message": str(e)})
if __name__ == "__main__":
import uvicorn
uvicorn.run("main:app", host="127.0.0.1", port=8000, reload=True)