"""Admin dataset management endpoints (auth-protected, under /admin). Covers list/view/purge, image→dataset lookup, embed-all, and match-all. The dataset *load* job endpoints live in api/jobs.py (background job + polling for the dashboard's load form). """ from __future__ import annotations from datetime import date from fastapi import APIRouter, Depends, HTTPException, status from sqlalchemy import select from sqlalchemy.orm import Session from ..db import get_db from ..models import Dataset, Picture, User from ..models.base import SubjectType from ..schemas.dataset import ( DatasetDetailOut, DatasetOut, EmbedAllResult, ImageDatasetOut, MatchDatasetResult, PurgeResult, ) from ..security import require_admin from ..services import datasets as ds router = APIRouter(prefix="/admin", tags=["admin-datasets"]) def _get_dataset(db: Session, dataset_id: int) -> Dataset: dataset = db.get(Dataset, dataset_id) if not dataset: raise HTTPException(status.HTTP_404_NOT_FOUND, "Dataset not found") return dataset @router.get("/datasets", response_model=list[DatasetOut]) def list_datasets( _: User = Depends(require_admin), db: Session = Depends(get_db), ) -> list[DatasetOut]: datasets = db.execute(select(Dataset).order_by(Dataset.creation_time.desc())).scalars().all() out: list[DatasetOut] = [] for d in datasets: item = DatasetOut.model_validate(d) item.case_count = ds.case_count(db, d) out.append(item) return out @router.get("/datasets/{dataset_id}", response_model=DatasetDetailOut) def get_dataset( dataset_id: int, _: User = Depends(require_admin), db: Session = Depends(get_db), ) -> DatasetDetailOut: dataset = _get_dataset(db, dataset_id) stats = ds.dataset_stats(db, dataset) base = DatasetOut.model_validate(dataset) base.case_count = stats["case_count"] return DatasetDetailOut(**base.model_dump(), stats=stats) @router.get("/test-match") def test_match( kind: str, dog_id: int, top_k: int = 10, apply_breed_gate: bool = False, _: User = Depends(require_admin), db: Session = Depends(get_db), ) -> dict: """Score one dog against ALL dogs in the opposite category by embedding similarity (testing).""" if kind not in ("known", "unknown"): raise HTTPException(status.HTTP_400_BAD_REQUEST, "kind must be 'known' or 'unknown'") subject_type = SubjectType.known if kind == "known" else SubjectType.unknown query = ds.dog_profile(db, subject_type, dog_id) if query is None: raise HTTPException(status.HTTP_404_NOT_FOUND, "Dog not found") top_k = max(1, min(top_k, 50)) from ..services.matching import rank_dog_against_opposite ranked, considered, query_embedded = rank_dog_against_opposite( db, subject_type=subject_type, subject_id=dog_id, top_k=top_k, apply_breed_gate=apply_breed_gate, ) results = [] for cand_type, cid, score in ranked: prof = ds.dog_profile(db, cand_type, cid) if prof: results.append({**prof, "score": score}) from ..ml import get_embedder e = get_embedder() return { "query": query, "results": results, "model": f"{e.name}/{e.version}", "candidate_count": considered, "query_embedded": query_embedded, } @router.get("/dog/{kind}/{dog_id}") def dog_detail( kind: str, dog_id: int, _: User = Depends(require_admin), db: Session = Depends(get_db), ) -> dict: """A dog's profile plus ALL of its photos (full images) — for the profile viewer.""" if kind not in ("known", "unknown"): raise HTTPException(status.HTTP_400_BAD_REQUEST, "kind must be 'known' or 'unknown'") subject_type = SubjectType.known if kind == "known" else SubjectType.unknown profile = ds.dog_profile(db, subject_type, dog_id) if profile is None: raise HTTPException(status.HTTP_404_NOT_FOUND, "Dog not found") from .helpers import pictures_for photos = [p.model_dump() for p in pictures_for(db, subject_type, dog_id)] return {"profile": profile, "photos": photos} @router.get("/breeds") def list_breeds( _: User = Depends(require_admin), db: Session = Depends(get_db), ) -> dict: """Predicted breed labels (HF softmax, active model) available to filter by.""" return {"model": ds.active_breed_model_in_db(db), "breeds": ds.list_breeds(db)} @router.get("/dogs") def list_all_dogs( kind: str = "all", limit: int = 30, offset: int = 0, breed: str | None = None, breed_k: int = 10, sort: str = "newest", zip: str | None = None, added_from: date | None = None, added_to: date | None = None, _: User = Depends(require_admin), db: Session = Depends(get_db), ) -> dict: """Browse every dog profile, filterable by predicted breed / ZIP / added-date and sortable. ``breed_k`` = 1 matches the top predicted breed; 5–10 matches within the top-K. ``sort`` is ``newest`` (default) or ``oldest``. ``zip`` matches a dog's ZIP by prefix; ``added_from`` / ``added_to`` bound the added date (inclusive). """ kind = kind if kind in ("all", "known", "unknown") else "all" limit = max(1, min(limit, 100)) offset = max(0, offset) breed_k = max(1, min(breed_k, 50)) sort = sort if sort in ("newest", "oldest") else "newest" return ds.list_all_dogs( db, kind, limit, offset, breed=breed or None, breed_k=breed_k, sort=sort, zip_prefix=(zip or None), added_from=added_from, added_to=added_to, ) @router.get("/datasets/{dataset_id}/dogs") def list_dataset_dogs( dataset_id: int, limit: int = 25, offset: int = 0, _: User = Depends(require_admin), db: Session = Depends(get_db), ) -> dict: dataset = _get_dataset(db, dataset_id) limit = max(1, min(limit, 100)) return ds.list_dogs(db, dataset, limit, offset) @router.delete("/datasets/{dataset_id}", response_model=PurgeResult) def delete_dataset( dataset_id: int, _: User = Depends(require_admin), db: Session = Depends(get_db), ) -> PurgeResult: dataset = _get_dataset(db, dataset_id) deleted = ds.purge_dataset(db, dataset) return PurgeResult(dataset_id=dataset_id, deleted=deleted) @router.get("/image/{image_id}/dataset", response_model=ImageDatasetOut) def image_dataset( image_id: int, _: User = Depends(require_admin), db: Session = Depends(get_db), ) -> ImageDatasetOut: pic = db.get(Picture, image_id) if not pic: raise HTTPException(status.HTTP_404_NOT_FOUND, "Image not found") from ..models import KnownDog, UnknownDog if pic.subject_type == SubjectType.known: dog = db.get(KnownDog, pic.subject_id) else: dog = db.get(UnknownDog, pic.subject_id) dataset_id = getattr(dog, "dataset_id", None) if dog else None dataset_name = None if dataset_id: d = db.get(Dataset, dataset_id) dataset_name = d.name if d else None return ImageDatasetOut( image_id=image_id, subject_type=pic.subject_type.value, subject_id=pic.subject_id, dataset_id=dataset_id, dataset_name=dataset_name, ) @router.post("/datasets/{dataset_id}/embed-all", response_model=EmbedAllResult) def embed_all( dataset_id: int, _: User = Depends(require_admin), db: Session = Depends(get_db), ) -> EmbedAllResult: dataset = _get_dataset(db, dataset_id) result = ds.embed_all(db, dataset) return EmbedAllResult(dataset_id=dataset_id, **result) @router.post("/datasets/{dataset_id}/match", response_model=MatchDatasetResult) def match_dataset( dataset_id: int, candidate_dataset_id: int | None = None, _: User = Depends(require_admin), db: Session = Depends(get_db), ) -> MatchDatasetResult: dataset = _get_dataset(db, dataset_id) result = ds.match_dataset(db, dataset, candidate_dataset_id=candidate_dataset_id) return MatchDatasetResult( dataset_id=dataset_id, cases_processed=result["cases_processed"], matches_created=result["matches_created"], candidate_dataset_id=result["candidate_dataset_id"], )