import json import os from contextlib import asynccontextmanager from pathlib import Path import structlog from dotenv import load_dotenv from fastapi import Depends, FastAPI, Header, HTTPException, Request, APIRouter from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import StreamingResponse from pydantic import BaseModel, Field from slowapi import Limiter, _rate_limit_exceeded_handler from slowapi.util import get_remote_address from slowapi.errors import RateLimitExceeded from typing import Annotated from ai_prompts import TEMPLATES from utils import ( call_llm, call_llm_messages, call_llm_stream, get_cached_response, set_cached_response, make_cache_key, ) from database import ( init_db, GUEST_DAILY_LIMIT, USER_DAILY_LIMIT, get_usage_today, increment_usage_today, ) from auth import get_optional_user from routers.auth import router as auth_router _BACKEND_DIR = Path(__file__).resolve().parent load_dotenv(dotenv_path=_BACKEND_DIR / ".env") # ── Sentry (optional) ───────────────────────────────────────────────────────── _SENTRY_DSN = os.getenv("SENTRY_DSN") if _SENTRY_DSN: import sentry_sdk from sentry_sdk.integrations.fastapi import FastApiIntegration from sentry_sdk.integrations.starlette import StarletteIntegration sentry_sdk.init( dsn=_SENTRY_DSN, integrations=[StarletteIntegration(), FastApiIntegration()], traces_sample_rate=float(os.getenv("SENTRY_TRACES_SAMPLE_RATE", "0.2")), send_default_pii=False, ) _REQUIRED_ENV = ["PROVIDER"] _PROVIDER_KEYS = { "openai": "OPENAI_API_KEY", "claude": "ANTHROPIC_API_KEY", "anthropic": "ANTHROPIC_API_KEY", "groq": "GROQ_API_KEY", } def _check_env() -> None: if os.getenv("SKIP_ENV_CHECK"): return missing = [k for k in _REQUIRED_ENV if not os.getenv(k)] if missing: raise RuntimeError(f"Missing required env vars: {', '.join(missing)}") provider = (os.getenv("PROVIDER") or "").lower() key_var = _PROVIDER_KEYS.get(provider) if key_var and not os.getenv(key_var): raise RuntimeError(f"PROVIDER={provider} requires {key_var} to be set") _check_env() logger = structlog.get_logger("pseudogen.app") @asynccontextmanager async def lifespan(app: FastAPI): init_db() logger.info("app.started") yield logger.info("app.stopped") # ── Rate limiter (Redis-backed when REDIS_URL is set) ───────────────────────── _redis_url = os.getenv("REDIS_URL") limiter = Limiter( key_func=get_remote_address, **{"storage_uri": _redis_url} if _redis_url else {}, ) app = FastAPI(title="Pseudogen API", lifespan=lifespan, docs_url=None, redoc_url=None) app.state.limiter = limiter app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler) # CORS: credentials require explicit origins (can't mix * with allow_credentials=True) _cors_origins_env = os.getenv("CORS_ORIGINS", "*").strip() if _cors_origins_env == "*": _allow_origins = ["*"] _allow_credentials = False else: _allow_origins = [o.strip() for o in _cors_origins_env.split(",") if o.strip()] _allow_credentials = True app.add_middleware( CORSMiddleware, allow_origins=_allow_origins, allow_credentials=_allow_credentials, allow_methods=["*"], allow_headers=["*"], ) MAX_INPUT_LEN = 4000 class MessageItem(BaseModel): role: Annotated[str, Field(pattern="^(user|assistant|system)$")] content: Annotated[str, Field(min_length=1, max_length=8000)] class GenerateRequest(BaseModel): problem_description: Annotated[str, Field(min_length=1, max_length=MAX_INPUT_LEN)] style: Annotated[str, Field(pattern="^(Academic|Developer-Friendly|English-Like|Step-by-Step)$")] detail: Annotated[str, Field(pattern="^(Concise|Detailed)$")] context: list[MessageItem] | None = None class SummarizeRequest(BaseModel): text: Annotated[str, Field(min_length=1, max_length=2000)] _STYLE_SYSTEM = { "Academic": ( "You generate Academic pseudocode using uppercase keywords " "(BEGIN, END, IF, ELSE, WHILE, FOR, FUNCTION, RETURN) with formal, concise logical flow. " "Output Markdown formatted pseudocode only." ), "Developer-Friendly": ( "You generate Developer-Friendly pseudocode with code-like syntax " "(Function, If, Else, While, For, Return), clear indentation, and comments where needed. " "Output Markdown formatted pseudocode only." ), "English-Like": ( "You convert problems into plain English steps with no programming syntax. " "Output numbered or bulleted Markdown steps only." ), "Step-by-Step": ( "You generate beginner-friendly pseudocode using simple English keywords " "(FUNCTION, IF, ELSE, WHILE, FOR, RETURN). " "Output Markdown formatted pseudocode only." ), } app.include_router(auth_router) v1_router = APIRouter(prefix="/v1", tags=["v1"]) @app.get("/") async def root(): return {"service": "Pseudogen API", "version": "1"} @app.get("/health") async def health(): return {"status": "ok"} @app.get("/usage") async def usage( request: Request, user: dict | None = Depends(get_optional_user), x_session_id: str | None = Header(default=None), ): if user: identifier = f"user:{user['id']}" limit = USER_DAILY_LIMIT is_guest = False else: identifier = f"ip:{_get_client_ip(request)}" limit = GUEST_DAILY_LIMIT is_guest = True used = get_usage_today(identifier) return {"used": used, "limit": limit, "remaining": max(0, limit - used), "is_guest": is_guest} @app.post("/summarize") @limiter.limit("60/minute") async def summarize_title(request: Request, req: SummarizeRequest): prompt = ( "Write a 4-6 word title for this programming problem. " "Title case. No punctuation. No quotes. No explanation. Just the title:\n\n" + req.text[:500] ) try: title = call_llm(prompt) title = title.strip().split("\n")[0][:60] return {"title": title} except Exception: logger.exception("summarize.failed") raise HTTPException(status_code=502, detail="Summarization failed") def _get_client_ip(request: Request) -> str: forwarded = request.headers.get("X-Forwarded-For") if forwarded: return forwarded.split(",")[0].strip() return request.client.host if request.client else "unknown" def _resolve_identity(request: Request, user: dict | None, x_session_id: str | None): if user: return f"user:{user['id']}", USER_DAILY_LIMIT, False # Track guests by IP — session UUID is trivially bypassed by opening incognito ip = _get_client_ip(request) return f"ip:{ip}", GUEST_DAILY_LIMIT, True def _build_messages(req: GenerateRequest) -> list | None: if not req.context: return None system_msg = ( f"{_STYLE_SYSTEM.get(req.style, 'You generate pseudocode.')} " f"Detail level: {req.detail}. " "When asked to modify or improve, update the pseudocode accordingly." ) context = req.context[-10:] return [ {"role": "system", "content": system_msg}, *[{"role": m.role, "content": m.content} for m in context], {"role": "user", "content": req.problem_description}, ] # ── Main endpoint — SSE streaming ───────────────────────────────────────────── @app.post("/generate-pseudocode") @limiter.limit("30/minute") async def generate( request: Request, req: GenerateRequest, user: dict | None = Depends(get_optional_user), x_session_id: str | None = Header(default=None), ): identifier, limit, is_guest = _resolve_identity(request, user, x_session_id) used = get_usage_today(identifier) if used >= limit: raise HTTPException( status_code=429, detail=( f"You've used all {limit} free prompts for today. Create a free account to get {USER_DAILY_LIMIT} per day." if is_guest else f"Daily limit of {limit} prompts reached. Resets at midnight UTC." ), ) # Increment before streaming to prevent quota abuse via cancel new_count = increment_usage_today(identifier) remaining = max(0, limit - new_count) messages = _build_messages(req) def _sse(): try: if messages: token_stream = call_llm_stream(messages) else: template = TEMPLATES.get(req.style) if template is None: yield f"data: {json.dumps({'error': 'Unknown style'})}\n\n" return prompt = template.format(user_input=req.problem_description, detail=req.detail) token_stream = call_llm_stream([{"role": "user", "content": prompt}]) for token in token_stream: yield f"data: {json.dumps({'token': token})}\n\n" yield f"data: {json.dumps({'usage': {'used': new_count, 'limit': limit, 'remaining': remaining, 'is_guest': is_guest}})}\n\n" yield "data: [DONE]\n\n" except Exception as exc: logger.error("generate.stream.error", error=str(exc)) yield f"data: {json.dumps({'error': 'Generation failed. Please try again.'})}\n\n" return StreamingResponse(_sse(), media_type="text/event-stream") # ── v1 endpoint — non-streaming with Redis cache ────────────────────────────── @v1_router.post("/generate-pseudocode") @limiter.limit("30/minute") async def generate_v1( request: Request, req: GenerateRequest, user: dict | None = Depends(get_optional_user), x_session_id: str | None = Header(default=None), ): identifier, limit, is_guest = _resolve_identity(request, user, x_session_id) cache_key = make_cache_key(req.problem_description, req.style, req.detail) cached = get_cached_response(cache_key) if cached: used = get_usage_today(identifier) return { "markdown": cached, "used": used, "limit": limit, "remaining": max(0, limit - used), "is_guest": is_guest, "cached": True, } used = get_usage_today(identifier) if used >= limit: raise HTTPException( status_code=429, detail=( f"You've used all {limit} free prompts for today. Create a free account to get {USER_DAILY_LIMIT} per day." if is_guest else f"Daily limit of {limit} prompts reached. Resets at midnight UTC." ), ) try: messages = _build_messages(req) if messages: response_text = call_llm_messages(messages) else: template = TEMPLATES.get(req.style) if template is None: raise HTTPException(status_code=400, detail="Unknown style") prompt = template.format(user_input=req.problem_description, detail=req.detail) response_text = call_llm(prompt) except HTTPException: raise except Exception: logger.exception("generate_v1.failed") raise HTTPException(status_code=502, detail="Failed to generate pseudocode. Please try again.") set_cached_response(cache_key, response_text) new_count = increment_usage_today(identifier) remaining = max(0, limit - new_count) return { "markdown": response_text, "used": new_count, "limit": limit, "remaining": remaining, "is_guest": is_guest, "cached": False, } app.include_router(v1_router)