app/routes.py

import re
import time

from fastapi import APIRouter, Request

from app.config import MAX_NEW_TOKENS, MODEL_NAME
from app.schemas import PredictionResponse, PromptRequest


router = APIRouter()

COMMAND_PATTERN = re.compile(
    r'"command"\s*:\s*"([^"]+)"',
)


@router.get("/")
def root():
    return {
        "status": "running",
    }


@router.get("/health")
def health(request: Request):
    model_loaded = (
        hasattr(request.app.state, "model")
        and hasattr(request.app.state, "tokenizer")
        and request.app.state.model is not None
        and request.app.state.tokenizer is not None
    )

    return {
        "status": "healthy",
        "model_loaded": model_loaded,
        "model_name": MODEL_NAME,
    }


@router.get("/model-info")
def model_info():
    return {
        "model_name": MODEL_NAME,
    }


@router.post("/predict", response_model=PredictionResponse)
def predict(payload: PromptRequest, request: Request):
    import torch

    start_time = time.time()

    tokenizer = request.app.state.tokenizer
    model = request.app.state.model

    messages = [
        {
            "role": "user",
            "content": payload.prompt,
        }
    ]

    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )

    inputs = tokenizer(
        text,
        return_tensors="pt",
    ).to(model.device)

    with torch.inference_mode():
        output = model.generate(
            **inputs,
            max_new_tokens=MAX_NEW_TOKENS,
            do_sample=False,
        )

    prompt_token_count = inputs["input_ids"].shape[1]
    generated_tokens = output[0][prompt_token_count:]

    response = tokenizer.decode(
        generated_tokens,
        skip_special_tokens=True,
    )

    command = None
    match = COMMAND_PATTERN.search(response)

    if match:
        command = match.group(1)

    latency_seconds = round(
        time.time() - start_time,
        3,
    )

    return PredictionResponse(
        prompt=payload.prompt,
        command=command,
        raw_output=response,
        latency_seconds=latency_seconds,
    )