from __future__ import annotations from typing import Any, Literal from fastapi import Depends, FastAPI, Header, HTTPException, status from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel, Field from data_studio.config import ConfigError, StudioSettings from data_studio.auth import create_or_update_user, generate_temporary_password, get_user from data_studio.db import create_client, ensure_indexes, get_database from data_studio.exports import export_clean_dataset from data_studio.benchmarks import assert_no_exact_benchmark_matches, register_protected_benchmark from data_studio.governance import SourceRegistryV1, SourceState from data_studio.governed_importer import import_governed_record from data_studio.governed_sources import get_registered_source, register_source, transition_source from data_studio.importer import import_source_samples from data_studio.rubrics import RubricV1, register_rubric from data_studio.review import ( ReviewConflict, ReviewValidationError, apply_review_action, claim_next_sample_with_enrichment, queue_overview, release_claim, ) from data_studio.security import create_session, delete_session, get_session_user, has_permission from data_studio.sources import list_sources, seed_default_sources from data_studio.utils import utc_now_iso settings = StudioSettings.from_env() client = create_client(settings) db = get_database(client, settings) ensure_indexes(db) app = FastAPI(title="Rachana Data Studio API", version="1.0.0") app.add_middleware( CORSMiddleware, allow_origins=["http://localhost:3000", "http://127.0.0.1:3000"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) @app.on_event("startup") def startup_seed_sources() -> None: ensure_indexes(db) seed_default_sources(db) class LoginRequest(BaseModel): username: str password: str class ReviewRequest(BaseModel): action: Literal["accept", "edit", "reject", "skip"] expected_version: int cleaned_payload: dict[str, Any] tags: list[str] = Field(default_factory=list) quality_tags: list[str] = Field(default_factory=list) task_tags: list[str] = Field(default_factory=list) notes: str = "" confidence: int = Field(default=4, ge=1, le=5) reason: str = "" class ImportRequest(BaseModel): dataset_key: str max_samples: int = Field(default=100, ge=1, le=50_000) class ExportRequest(BaseModel): version: str = Field(min_length=1, max_length=80, pattern=r"^[A-Za-z0-9._-]+$") repo_id: str | None = None class RegisterSourceRequest(BaseModel): source: SourceRegistryV1 class TransitionSourceRequest(BaseModel): target: SourceState reason: str = Field(min_length=1) class RegisterBenchmarkRequest(BaseModel): benchmark_id: str = Field(min_length=1) exact_revision: str = Field(min_length=1) task: str = Field(min_length=1) exact_payloads: list[dict[str, Any]] class RegisterRubricRequest(BaseModel): rubric: RubricV1 class GovernedRecordRequest(BaseModel): source_id: str = Field(min_length=1) source_revision: str = Field(min_length=1) source_native_id: str = Field(min_length=1) raw_payload: dict[str, Any] normalized_payload: dict[str, Any] class CreateUserRequest(BaseModel): username: str = Field(min_length=1) password: str = Field(min_length=8) display_name: str | None = None role: Literal["reviewer", "manager", "admin"] = "reviewer" update_existing: bool = False inactive: bool = False class ResetPasswordRequest(BaseModel): temporary_password: str | None = Field(default=None, min_length=8) class ResetStudioRequest(BaseModel): confirm: Literal["RESET"] class ResetReviewPoolRequest(BaseModel): confirm: Literal["RESET_REVIEW_POOL"] def bearer_token(authorization: str = Header(default="")) -> str: scheme, _, token = authorization.partition(" ") if scheme.lower() != "bearer" or not token: raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Authentication required.") return token def current_user(token: str = Depends(bearer_token)) -> dict[str, Any]: user = get_session_user(db, token) if user is None: raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Session expired.") return user def require_permission(permission: str): def dependency(user: dict[str, Any] = Depends(current_user)) -> dict[str, Any]: if not has_permission(user, permission): raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail=f"Missing permission: {permission}") return user return dependency @app.get("/api/health") def health() -> dict[str, Any]: return {"status": "ok", **settings.deployment_metadata()} @app.post("/api/auth/login") def login(payload: LoginRequest) -> dict[str, Any]: result = create_session(db, payload.username, payload.password) if result is None: raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid username or password.") token, user = result return {"token": token, "user": user} @app.post("/api/auth/logout") def logout(token: str = Depends(bearer_token)) -> dict[str, bool]: delete_session(db, token) return {"ok": True} @app.get("/api/me") def me(user: dict[str, Any] = Depends(current_user)) -> dict[str, Any]: return user @app.get("/api/overview") def overview(user: dict[str, Any] = Depends(current_user)) -> dict[str, Any]: return {**queue_overview(db, settings), "source_configs": list_sources(db), "user": user} @app.post("/api/review/claim/{queue_name}") def claim( queue_name: str, dataset_key: str | None = None, user: dict[str, Any] = Depends(require_permission("review")), ) -> dict[str, Any]: return {"sample": claim_next_sample_with_enrichment(db, settings, queue_name, user["username"], dataset_key)} @app.post("/api/review/{sample_id}/release") def release(sample_id: str, user: dict[str, Any] = Depends(require_permission("review"))) -> dict[str, bool]: release_claim(db, sample_id, user["username"]) return {"ok": True} @app.post("/api/review/{sample_id}") def review( sample_id: str, payload: ReviewRequest, user: dict[str, Any] = Depends(require_permission("review")), ) -> dict[str, Any]: try: resolved_action = apply_review_action( db=db, settings=settings, sample_id=sample_id, reviewer=user["username"], expected_version=payload.expected_version, action=payload.action, cleaned_payload=payload.cleaned_payload, tags=payload.tags, quality_tags=payload.quality_tags, task_tags=payload.task_tags, notes=payload.notes, confidence=payload.confidence, reason=payload.reason, ) except ReviewConflict as exc: raise HTTPException(status_code=status.HTTP_409_CONFLICT, detail=str(exc)) from exc except ReviewValidationError as exc: raise HTTPException(status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, detail=str(exc)) from exc return resolved_action @app.post("/api/admin/sources/seed") def seed_sources(user: dict[str, Any] = Depends(require_permission("manage_sources"))) -> dict[str, int]: return {"inserted": seed_default_sources(db)} @app.post("/api/admin/reset") def reset_studio( payload: ResetStudioRequest, user: dict[str, Any] = Depends(require_permission("manage_users")), ) -> dict[str, Any]: collections = [ "samples", "review_events", "import_jobs", "import_checkpoints", "exports", "dataset_versions", "sessions", "source_configs", ] cleared = {} for name in collections: result = db[name].delete_many({}) cleared[name] = result.deleted_count inserted_sources = seed_default_sources(db) return {"cleared": cleared, "seeded_sources": inserted_sources, "confirmed": payload.confirm} @app.post("/api/admin/review/reset") def reset_review_pool( payload: ResetReviewPoolRequest, user: dict[str, Any] = Depends(require_permission("manage_users")), ) -> dict[str, Any]: reviewed_statuses = ["accepted", "edited", "rejected", "skipped"] now = utc_now_iso() reset_counts = {"samples": 0, "review_events": 0} reviewed_samples = list(db["samples"].find({"status": {"$in": reviewed_statuses}})) for sample in reviewed_samples: sample_type = sample["sample_type"] review_reset = { "current_text": None, "current_pair": None, "current_transliteration": None, "tags": [], "quality_tags": [], "task_tags": [], "notes": "", "last_reviewed_by": None, "last_reviewed_at": None, "edit_count": 0, "token_stats": None, } if sample_type == "document": review_reset["current_text"] = sample["payload"].get("text", "") elif sample_type == "translation_pair": review_reset["current_pair"] = { "source_text": sample["payload"].get("source_text", ""), "target_text": sample["payload"].get("target_text", ""), } elif sample_type == "transliteration_pair": review_reset["current_transliteration"] = { "native_text": sample["payload"].get("native_text", ""), "latin_text": sample["payload"].get("latin_text", ""), } db["samples"].update_one( {"sample_id": sample["sample_id"]}, { "$set": { "status": "pending", "review": review_reset, "updated_at": now, "metrics.clean_char_count": None, "metrics.clean_token_count": None, "metrics.tokenizer_version": None, "export.eligible": False, "governance.review_accepted": False, "governance.training_export_eligible": False, "live_sync.repo_id": None, "live_sync.path_in_repo": None, "live_sync.status": "pending", "live_sync.last_synced_at": None, "live_sync.last_error": None, }, "$unset": {"lease": ""}, "$inc": {"version": 1}, }, ) reset_counts["samples"] += 1 deleted_events = db["review_events"].delete_many({"sample_id": {"$in": [sample["sample_id"] for sample in reviewed_samples]}}) reset_counts["review_events"] = deleted_events.deleted_count return { "confirmed": payload.confirm, "reset_counts": reset_counts, "metadata": now, } @app.get("/api/admin/users") def list_users(user: dict[str, Any] = Depends(require_permission("manage_users"))) -> dict[str, Any]: users = [] for row in db["users"].find().sort("username", 1): users.append( { "username": row["username"], "display_name": row.get("display_name") or row["username"], "role": row.get("role", "reviewer"), "is_active": bool(row.get("is_active", True)), "last_login_at": row.get("last_login_at"), } ) return {"users": users} @app.post("/api/admin/users") def create_user( payload: CreateUserRequest, user: dict[str, Any] = Depends(require_permission("manage_users")), ) -> dict[str, Any]: try: created = create_or_update_user( db=db, username=payload.username, password=payload.password, display_name=payload.display_name, role=payload.role, is_active=not payload.inactive, update_existing=payload.update_existing, ) except ValueError as exc: raise HTTPException(status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, detail=str(exc)) from exc return {"user": created} @app.post("/api/admin/users/{username}/reset-password") def reset_user_password( username: str, payload: ResetPasswordRequest, user: dict[str, Any] = Depends(require_permission("manage_users")), ) -> dict[str, Any]: existing = get_user(db, username) if existing is None: raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="User not found.") temporary_password = payload.temporary_password or generate_temporary_password() try: created = create_or_update_user( db=db, username=username, password=temporary_password, display_name=existing.get("display_name"), role=existing.get("role", "reviewer"), is_active=bool(existing.get("is_active", True)), update_existing=True, ) except ValueError as exc: raise HTTPException(status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, detail=str(exc)) from exc return {"user": created, "temporary_password": temporary_password} @app.post("/api/admin/import") def run_import( payload: ImportRequest, user: dict[str, Any] = Depends(require_permission("import")), ) -> dict[str, Any]: try: settings.require_legacy_import_export_enabled() except ConfigError as exc: raise HTTPException(status_code=status.HTTP_423_LOCKED, detail=str(exc)) from exc return import_source_samples(db, settings, payload.dataset_key, payload.max_samples) @app.post("/api/admin/export") def run_export( payload: ExportRequest, user: dict[str, Any] = Depends(require_permission("export")), ) -> dict[str, Any]: try: settings.require_legacy_import_export_enabled() except ConfigError as exc: raise HTTPException(status_code=status.HTTP_423_LOCKED, detail=str(exc)) from exc return export_clean_dataset(db, settings, payload.version, user["username"], payload.repo_id) @app.post("/api/governance/sources") def create_governed_source( payload: RegisterSourceRequest, user: dict[str, Any] = Depends(require_permission("manage_sources")), ) -> dict[str, str]: register_source(db, payload.source, user["username"]) return {"source_id": payload.source.source_id, "revision": payload.source.hf_revision} @app.post("/api/governance/sources/{source_id}/{revision}/transition") def change_governed_source_state( source_id: str, revision: str, payload: TransitionSourceRequest, user: dict[str, Any] = Depends(require_permission("approve_sources")), ) -> dict[str, str]: transition_source(db, source_id, revision, payload.target, user["username"], payload.reason) return {"source_id": source_id, "revision": revision, "state": payload.target.value} @app.post("/api/governance/benchmarks") def create_protected_benchmark( payload: RegisterBenchmarkRequest, user: dict[str, Any] = Depends(require_permission("manage_benchmarks")), ) -> dict[str, str]: register_protected_benchmark( db, payload.benchmark_id, payload.exact_revision, payload.task, payload.exact_payloads, user["username"], ) return {"benchmark_id": payload.benchmark_id, "revision": payload.exact_revision} @app.post("/api/governance/rubrics") def create_rubric( payload: RegisterRubricRequest, user: dict[str, Any] = Depends(require_permission("manage_rubrics")), ) -> dict[str, str]: register_rubric(db, payload.rubric, user["username"]) return {"rubric_id": payload.rubric.rubric_id, "rubric_version": payload.rubric.rubric_version} @app.post("/api/governance/records") def create_governed_record( payload: GovernedRecordRequest, user: dict[str, Any] = Depends(require_permission("governed_import")), ) -> dict[str, Any]: source = get_registered_source(db, payload.source_id, payload.source_revision) if source is None: raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Registered source revision not found.") assert_no_exact_benchmark_matches(db, [payload.raw_payload, payload.normalized_payload]) record_id, inserted = import_governed_record( db, source, payload.source_native_id, payload.raw_payload, payload.normalized_payload, ) return {"record_id": record_id, "inserted": inserted}