src/api.py

"""Cortex-0 — API FastAPI (Phase 12).

Expose le moteur multimédia maison (`src/inference.CortexEngine`) :

* ``GET /health``            : état du service + si un checkpoint entraîné est chargé ;
* ``GET /capabilities``      : tailles / modalités supportées ;
* ``POST /chat``             : texte -> texte (JSON) ;
* ``POST /vision``           : image (upload) -> légende ;
* ``POST /audio``            : WAV (upload) -> transcription ;
* ``POST /video``            : frames + audio optionnel -> description ;
* ``POST /memory/ingest``    : ingérer un cours (RAG, mode A) ;
* ``POST /memory/query``     : interroger la mémoire d'un cours (réponse extractive).

Aucun secret en dur ; le checkpoint est choisi via ``CORTEX_CHECKPOINT``. Aucun poids
pré-entraîné n'est téléchargé : seuls nos propres checkpoints sont chargés.
"""

from __future__ import annotations

import os
from pathlib import Path

from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from pydantic import BaseModel, Field

from src.inference import CortexEngine
from src.memory.rag import LessonMemory

app = FastAPI(title="Cortex-0", version="0.1.0",
              description="IA multimodale from scratch (texte/image/audio/vidéo).")

_engine: CortexEngine | None = None
_memory = LessonMemory(dim=512, max_words=40)


def get_engine() -> CortexEngine:
    global _engine
    if _engine is None:
        ckpt = os.environ.get("CORTEX_CHECKPOINT") or None
        _engine = CortexEngine(checkpoint=ckpt)
    return _engine


# ---------------------------------------------------------------------------
# Schémas
# ---------------------------------------------------------------------------


class ChatRequest(BaseModel):
    prompt: str = Field(..., min_length=1)
    max_new_tokens: int = Field(40, ge=1, le=512)


class IngestRequest(BaseModel):
    text: str = Field(..., min_length=1)
    source: str = "manual"


class QueryRequest(BaseModel):
    question: str = Field(..., min_length=1)
    k: int = Field(3, ge=1, le=10)


# ---------------------------------------------------------------------------
# Routes
# ---------------------------------------------------------------------------


@app.get("/health")
def health() -> dict:
    eng = get_engine()
    return {"status": "ok", "trained": eng.trained, "device": str(eng.device),
            "vocab_size": eng.tokenizer.vocab_size}


@app.get("/capabilities")
def capabilities() -> dict:
    eng = get_engine()
    return {
        "modalities": ["text", "image", "audio", "video"],
        "image_tokens": eng.n_image_tokens,
        "audio_tokens": eng.n_audio_tokens,
        "video_tokens": eng.n_video_tokens,
        "from_scratch": True,
        "external_weights": False,
    }


@app.post("/chat")
def chat(req: ChatRequest) -> dict:
    return get_engine().chat(req.prompt, max_new_tokens=req.max_new_tokens)


@app.post("/vision")
async def vision(file: UploadFile = File(...),
                 max_new_tokens: int = Form(20)) -> dict:
    data = await file.read()
    if not data:
        raise HTTPException(status_code=400, detail="image vide")
    try:
        return get_engine().describe_image(data, max_new_tokens=max_new_tokens)
    except TypeError as exc:
        raise HTTPException(status_code=415, detail=str(exc)) from exc


@app.post("/audio")
async def audio(file: UploadFile = File(...),
                max_new_tokens: int = Form(20)) -> dict:
    data = await file.read()
    if not data:
        raise HTTPException(status_code=400, detail="audio vide")
    try:
        return get_engine().transcribe_audio(data, max_new_tokens=max_new_tokens)
    except TypeError as exc:
        raise HTTPException(status_code=415, detail=str(exc)) from exc


@app.post("/video")
async def video(frames: UploadFile = File(...),
                audio_file: UploadFile | None = File(None),
                max_new_tokens: int = Form(24)) -> dict:
    """``frames`` = archive numpy ``.npy`` de forme ``(T, H, W, 3)``."""
    import io

    import numpy as np

    raw = await frames.read()
    try:
        arr = np.load(io.BytesIO(raw), allow_pickle=False)
    except Exception as exc:  # fichier invalide
        raise HTTPException(status_code=400, detail=f"frames .npy invalides : {exc}") from exc
    aud = await audio_file.read() if audio_file is not None else None
    return get_engine().describe_video(arr, audio=aud, max_new_tokens=max_new_tokens)


@app.post("/memory/ingest")
def memory_ingest(req: IngestRequest) -> dict:
    _memory.add_text(req.text, source=req.source)
    return {"passages": len(_memory.passages), "source": req.source}


@app.post("/memory/query")
def memory_query(req: QueryRequest) -> dict:
    return _memory.answer(req.question, k=req.k)


def main() -> None:
    import uvicorn

    ckpt = os.environ.get("CORTEX_CHECKPOINT")
    if ckpt and Path(ckpt).is_file():
        get_engine()
    uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", "8000")))


if __name__ == "__main__":
    main()