File size: 2,257 Bytes
215f97f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
"""

Training endpoints.



POST /train             — trigger a background training job

GET  /train/status/{job_id} — check training job status

"""

import uuid
from datetime import datetime, timezone

from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException

from app.api.dependencies import verify_internal_service
from app.models.schemas import TrainingJobResponse, ErrorResponse
from app.core.firebase import get_firestore_db
from app.jobs.training_job import run_training_job

router = APIRouter(prefix="/train", tags=["Training"])


@router.post(

    "",

    response_model=TrainingJobResponse,

    dependencies=[Depends(verify_internal_service)],

    summary="Trigger a model training job",

    description=(

        "Creates a background training job that loads labeled documents from Supabase, "

        "trains a new RETVec+CNN model, evaluates it, and stores the resulting model "

        "to Firebase Storage and Firestore. Returns the job ID immediately."

    ),

    responses={

        401: {"model": ErrorResponse, "description": "Unauthorized — Missing or invalid X-Internal-Token header"},

        403: {"model": ErrorResponse, "description": "Forbidden — Client IP banned due to 3 failed token attempts"},

        500: {"model": ErrorResponse, "description": "Internal Server Error — Failed to initialize training record in Firestore"},

    },

)
async def start_training(background_tasks: BackgroundTasks):
    """Start a new training job in the background."""
    db = get_firestore_db()
    job_id = str(uuid.uuid4())

    # Create job record in Firestore (survives service restarts)
    if db is not None:
        try:
            db.collection("training_jobs").document(job_id).set(
                {
                    "jobId": job_id,
                    "status": "queued",
                    "createdAt": datetime.now(timezone.utc),
                }
            )
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"Failed to create job in Firestore: {str(e)}")

    # Launch training as a background task
    background_tasks.add_task(run_training_job, job_id)

    return {"jobId": job_id, "status": "queued"}