import logging import os import time from collections import defaultdict from functools import lru_cache from fastapi import FastAPI from fastapi import HTTPException, Request from fastapi.middleware.cors import CORSMiddleware import joblib import pandas as pd from pathlib import Path from pydantic import BaseModel, Field, field_validator from starlette.concurrency import run_in_threadpool from urllib.parse import urlsplit logger = logging.getLogger(__name__) app = FastAPI( title="Task Time Predictor API", description="API for predicting task duration categories based on task details.", version="1.0", ) def parse_allowed_origins(): configured_origins = { origin.strip().rstrip("/") for origin in os.getenv("ALLOWED_ORIGINS", "").split(",") if origin.strip() } vercel_urls = { os.getenv("VERCEL_URL", ""), os.getenv("VERCEL_BRANCH_URL", ""), os.getenv("VERCEL_PROJECT_PRODUCTION_URL", ""), } for vercel_url in {url.strip().rstrip("/") for url in vercel_urls if url.strip()}: configured_origins.add( vercel_url if vercel_url.startswith("http") else f"https://{vercel_url}" ) local_origins = { "http://localhost:5173", "http://127.0.0.1:5173", "http://localhost:8000", "http://127.0.0.1:8000", } return configured_origins | local_origins ALLOWED_ORIGINS = parse_allowed_origins() RATE_LIMIT_WINDOW_SECONDS = int(os.getenv("RATE_LIMIT_WINDOW_SECONDS", "60")) RATE_LIMIT_MAX_REQUESTS = int(os.getenv("RATE_LIMIT_MAX_REQUESTS", "20")) RATE_LIMIT_MAX_CLIENTS = int(os.getenv("RATE_LIMIT_MAX_CLIENTS", "2000")) MAX_PREDICT_CONTENT_LENGTH = int(os.getenv("MAX_PREDICT_CONTENT_LENGTH", "8000")) REQUEST_LOG = defaultdict(list) app.add_middleware( CORSMiddleware, allow_origins=sorted(ALLOWED_ORIGINS), allow_credentials=False, allow_methods=["GET", "POST", "OPTIONS"], allow_headers=["Content-Type"], ) PROJECT_ROOT = Path(__file__).resolve().parents[1] MODEL_PATH = PROJECT_ROOT / "models" / "duration_logistic_regression_classifier.joblib" @lru_cache(maxsize=1) def load_model(): if not MODEL_PATH.exists(): raise RuntimeError(f"Model file not found at {MODEL_PATH}") return joblib.load(MODEL_PATH) class TaskInput(BaseModel): summary: str = Field(..., min_length=1, max_length=300) description: str = Field(default="", max_length=3000) issuetype_name: str = Field(default="unknown", max_length=80) priority_name: str = Field(default="unknown", max_length=80) project_key: str = Field(default="unknown", max_length=40) project_category_name: str = Field(default="unknown", max_length=120) created_year: int = Field(default=0, ge=0, le=2100) created_month: int = Field(default=0, ge=0, le=12) labels_count: int = Field(default=0, ge=0, le=100) has_assignee: int = Field(default=0, ge=0, le=1) votes_votes: float = Field(default=0, ge=0, le=1_000_000) watches_watch_count: float = Field(default=0, ge=0, le=1_000_000) @field_validator( "summary", "description", "issuetype_name", "priority_name", "project_key", "project_category_name", mode="before", ) @classmethod def strip_text_fields(cls, value): return value.strip() if isinstance(value, str) else value def origin_from_url(url): if not url: return None parsed_url = urlsplit(url) if not parsed_url.scheme or not parsed_url.netloc: return None return f"{parsed_url.scheme}://{parsed_url.netloc}" def request_origin(request): origin = request.headers.get("origin") if origin: return origin.rstrip("/") return origin_from_url(request.headers.get("referer")) def request_base_origin(request): forwarded_proto = request.headers.get("x-forwarded-proto", "") forwarded_host = request.headers.get("x-forwarded-host", "") proto = forwarded_proto.split(",", 1)[0].strip() or request.url.scheme host = forwarded_host.split(",", 1)[0].strip() or request.headers.get("host") return f"{proto}://{host}".rstrip("/") if host else None def is_allowed_request_origin(request): origin = request_origin(request) return bool(origin and (origin in ALLOWED_ORIGINS or origin == request_base_origin(request))) def client_key(request): forwarded_for = request.headers.get("x-forwarded-for", "") if forwarded_for: return forwarded_for.split(",", 1)[0].strip() return request.client.host if request.client else "unknown" def enforce_predict_security(request): content_length = request.headers.get("content-length") try: content_length_value = int(content_length) if content_length else 0 except ValueError: raise HTTPException(status_code=400, detail="Invalid Content-Length header.") if content_length_value > MAX_PREDICT_CONTENT_LENGTH: raise HTTPException(status_code=413, detail="Prediction request is too large.") if not is_allowed_request_origin(request): raise HTTPException( status_code=403, detail="Use the deployed app or its Swagger docs page to submit predictions.", ) now = time.monotonic() key = client_key(request) if len(REQUEST_LOG) > RATE_LIMIT_MAX_CLIENTS: expired_keys = [ request_key for request_key, timestamps in REQUEST_LOG.items() if not timestamps or now - max(timestamps) >= RATE_LIMIT_WINDOW_SECONDS ] for request_key in expired_keys: REQUEST_LOG.pop(request_key, None) if len(REQUEST_LOG) > RATE_LIMIT_MAX_CLIENTS: REQUEST_LOG.clear() recent_requests = [ timestamp for timestamp in REQUEST_LOG[key] if now - timestamp < RATE_LIMIT_WINDOW_SECONDS ] if len(recent_requests) >= RATE_LIMIT_MAX_REQUESTS: raise HTTPException( status_code=429, detail="Too many prediction requests. Please wait before trying again.", ) recent_requests.append(now) REQUEST_LOG[key] = recent_requests def api_metadata(): return { "name": "Task Time Predictor API", "description": "API for predicting task duration categories based on task details.", "status": "Running", "model": { "type": "Logistic Regression", "classes": ["Short", "Standard", "Long-running"], }, "endpoints": { "root": "/", "health": "/health", "predict": "/predict", }, } @app.get("/") async def root(): return api_metadata() @app.get("/health") async def health(): return api_metadata() def run_prediction(features): model = load_model() prediction = model.predict(features)[0] probabilities = model.predict_proba(features)[0] class_probabilities = { str(class_name): float(probability) for class_name, probability in zip(model.classes_, probabilities) } return str(prediction), class_probabilities @app.post("/predict") async def predict_task_time(task: TaskInput, request: Request): enforce_predict_security(request) summary = task.summary or "" description = task.description or "" summary_word_count = len(summary.split()) description_word_count = len(description.split()) summary_to_description_word_ratio = ( summary_word_count / description_word_count if description_word_count else summary_word_count ) issue_priority = f"{task.issuetype_name}__{task.priority_name}" features = pd.DataFrame([{ "summary_text": summary, "description_text": description, "total_text": f"{summary} {description}", "priority_name": task.priority_name, "issuetype_name": task.issuetype_name, "project_key": task.project_key, "project_category_name": task.project_category_name, "created_year": task.created_year, "created_month": task.created_month, "issue_priority": issue_priority, "summary_char_count": len(summary), "summary_word_count": summary_word_count, "description_char_count": len(description), "description_word_count": description_word_count, "has_description": int(description_word_count > 0), "labels_count": task.labels_count, "has_assignee": task.has_assignee, "votes_votes": task.votes_votes, "watches_watch_count": task.watches_watch_count, "summary_to_description_word_ratio": summary_to_description_word_ratio, }]) try: prediction, class_probabilities = await run_in_threadpool( run_prediction, features, ) except Exception as e: logger.exception("Prediction failed") raise HTTPException( status_code=503, detail="Prediction service is temporarily unavailable.", ) from e return { "duration_category": prediction, "probabilities": class_probabilities, }