api.py

import logging
from typing import Annotated

from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, StringConstraints
from outlines import generate

from generate import model, sampler, stream_jsonl_file

logger = logging.getLogger(__name__)


class Status(BaseModel):
    status: Annotated[str, StringConstraints(pattern="ok")]

status_generator = generate.json(model, Status, sampler=sampler)

status_stream = status_generator.stream("status:")
status = "".join(char.strip() for char in status_stream if char.strip())
logger.warning("Model status: " + status)


async def stream_response(filename: str, prompt: str, columns: list[str], seed: int, size: int):
    for chunk in stream_jsonl_file(
        filename=filename,
        prompt=prompt,
        columns=columns,
        seed=seed,
        size=size,
    ):
        yield chunk


async def dummy_stream():
    yield ""


app = FastAPI()

@app.head("/{filename}.jsonl")
@app.get("/{filename}.jsonl")
async def read_item(request: Request, filename: str, prompt: str = "", columns: str = "", seed: int = 42, size: int = 3):
    if request.method == 'GET':
        columns = [field.strip() for field in columns.split(",") if field.strip()]
        content = stream_response(
            filename,
            prompt=prompt,
            columns=columns,
            seed=seed,
            size=size
        )
    else:
        content = dummy_stream()
    response = StreamingResponse(content, media_type="text/jsonlines")
    response.headers["Content-Disposition"] = f"attachment; filename={filename}.jsonl"
    return response