Files
invoice-master-poc-v2/packages/backend/backend/web/api/v1/admin/training/pool.py
Yaojia Wang ad5ed46b4c WIP
2026-02-11 23:40:38 +01:00

160 lines
5.0 KiB
Python

"""Fine-Tune Pool Endpoints."""
import logging
from typing import Annotated
from fastapi import APIRouter, HTTPException, Query
from backend.web.core.auth import AdminTokenDep, FineTunePoolRepoDep, DocumentRepoDep
from backend.web.schemas.admin.pool import (
PoolAddRequest,
PoolEntryItem,
PoolEntryResponse,
PoolListResponse,
PoolStatsResponse,
)
from ._utils import _validate_uuid
logger = logging.getLogger(__name__)
def register_pool_routes(router: APIRouter) -> None:
"""Register fine-tune pool endpoints on the router."""
@router.post(
"/pool",
response_model=PoolEntryResponse,
summary="Add document to fine-tune pool",
description="Add a labeled document to the fine-tune pool for future fine-tuning.",
)
async def add_to_pool(
request: PoolAddRequest,
admin_token: AdminTokenDep,
pool: FineTunePoolRepoDep,
docs: DocumentRepoDep,
) -> PoolEntryResponse:
"""Add a document to the fine-tune pool."""
_validate_uuid(request.document_id, "document_id")
# Verify document exists
doc = docs.get(request.document_id)
if not doc:
raise HTTPException(status_code=404, detail="Document not found")
# Check if already in pool
existing = pool.get_by_document(request.document_id)
if existing:
raise HTTPException(
status_code=409,
detail=f"Document already in fine-tune pool (entry_id: {existing.entry_id})",
)
entry = pool.add_document(
document_id=request.document_id,
added_by=admin_token,
reason=request.reason,
)
return PoolEntryResponse(
entry_id=str(entry.entry_id),
message="Document added to fine-tune pool",
)
@router.get(
"/pool",
response_model=PoolListResponse,
summary="List fine-tune pool entries",
)
async def list_pool_entries(
admin_token: AdminTokenDep,
pool: FineTunePoolRepoDep,
verified_only: Annotated[bool, Query(description="Filter to verified only")] = False,
limit: Annotated[int, Query(ge=1, le=100)] = 20,
offset: Annotated[int, Query(ge=0)] = 0,
) -> PoolListResponse:
"""List entries in the fine-tune pool."""
entries, total = pool.get_paginated(
verified_only=verified_only,
limit=limit,
offset=offset,
)
return PoolListResponse(
total=total,
limit=limit,
offset=offset,
entries=[
PoolEntryItem(
entry_id=str(e.entry_id),
document_id=str(e.document_id),
added_by=e.added_by,
reason=e.reason,
is_verified=e.is_verified,
verified_at=e.verified_at,
verified_by=e.verified_by,
created_at=e.created_at,
)
for e in entries
],
)
@router.get(
"/pool/stats",
response_model=PoolStatsResponse,
summary="Get fine-tune pool statistics",
)
async def get_pool_stats(
admin_token: AdminTokenDep,
pool: FineTunePoolRepoDep,
) -> PoolStatsResponse:
"""Get statistics about the fine-tune pool."""
total = pool.get_pool_count(verified_only=False)
verified = pool.get_pool_count(verified_only=True)
return PoolStatsResponse(
total_entries=total,
verified_entries=verified,
unverified_entries=total - verified,
is_ready=verified >= 50,
)
@router.post(
"/pool/{entry_id}/verify",
response_model=PoolEntryResponse,
summary="Verify a pool entry",
description="Mark a pool entry as verified (human-reviewed).",
)
async def verify_pool_entry(
entry_id: str,
admin_token: AdminTokenDep,
pool: FineTunePoolRepoDep,
) -> PoolEntryResponse:
"""Mark a pool entry as verified."""
_validate_uuid(entry_id, "entry_id")
entry = pool.verify_entry(entry_id, verified_by=admin_token)
if not entry:
raise HTTPException(status_code=404, detail="Pool entry not found")
return PoolEntryResponse(
entry_id=str(entry.entry_id),
message="Pool entry verified",
)
@router.delete(
"/pool/{entry_id}",
summary="Remove from fine-tune pool",
)
async def remove_from_pool(
entry_id: str,
admin_token: AdminTokenDep,
pool: FineTunePoolRepoDep,
) -> dict:
"""Remove a document from the fine-tune pool."""
_validate_uuid(entry_id, "entry_id")
success = pool.remove_entry(entry_id)
if not success:
raise HTTPException(status_code=404, detail="Pool entry not found")
return {"message": "Entry removed from fine-tune pool"}