| from __future__ import annotations |
|
|
| import json |
| import os |
| import secrets |
| import threading |
| import time |
| import uuid |
| from contextlib import asynccontextmanager |
| from typing import Any, Literal |
|
|
| from fastapi import Depends, FastAPI, Header, HTTPException |
| from fastapi.concurrency import run_in_threadpool |
| from fastapi.responses import StreamingResponse |
| from pydantic import BaseModel, Field |
| from tokenizers import Tokenizer |
|
|
| from nayara.artifacts import materialize_checkpoint |
| from nayara.generate import load_model, stream_generate |
|
|
|
|
| class CompletionRequest(BaseModel): |
| model: str = "nayara-v1" |
| prompt: str |
| max_tokens: int = Field(192, ge=1, le=512) |
| temperature: float = Field(0.2, ge=0, le=2) |
| top_k: int = Field(40, ge=1, le=32000) |
| stream: bool = False |
|
|
|
|
| class ChatMessage(BaseModel): |
| role: Literal["system", "user", "assistant"] |
| content: str |
|
|
|
|
| class ChatRequest(BaseModel): |
| model: str = "nayara-v1" |
| messages: list[ChatMessage] |
| max_tokens: int = Field(192, ge=1, le=512) |
| temperature: float = Field(0.2, ge=0, le=2) |
| top_k: int = Field(40, ge=1, le=32000) |
| stream: bool = False |
|
|
|
|
| runtime: dict[str, Any] = {"ready": False} |
| generation_lock = threading.Lock() |
| MODEL_ID = os.environ.get("NAYARA_MODEL_ID", "nayara-v1") |
|
|
|
|
| def _authorize(authorization: str | None = Header(default=None)) -> None: |
| expected = os.environ.get("NAYARA_API_KEY") |
| if not expected: |
| return |
| supplied = authorization.removeprefix("Bearer ") if authorization else "" |
| if not secrets.compare_digest(supplied, expected): |
| raise HTTPException(status_code=401, detail="invalid bearer token") |
|
|
|
|
| def _chat_prompt(messages: list[ChatMessage]) -> str: |
| if not messages or not any(message.role == "user" for message in messages): |
| raise HTTPException(status_code=422, detail="messages must contain at least one user message") |
| parts: list[str] = [] |
| for message in messages: |
| if message.role == "system": |
| parts.append(f"<|user|>\nSystem guidance: {message.content.strip()}") |
| else: |
| parts.append(f"<|{message.role}|>\n{message.content.strip()}") |
| parts.append("<|assistant|>\n") |
| return "\n".join(parts) |
|
|
|
|
| def _generate(prompt: str, request: CompletionRequest | ChatRequest): |
| return stream_generate( |
| runtime["model"], runtime["tokenizer"], prompt, |
| max_new_tokens=request.max_tokens, temperature=request.temperature, top_k=request.top_k, |
| ) |
|
|
|
|
| def _complete(prompt: str, request: CompletionRequest | ChatRequest) -> str: |
| final = "" |
| with generation_lock: |
| for final in _generate(prompt, request): |
| pass |
| return final |
|
|
|
|
| def _sse(prompt: str, request: CompletionRequest | ChatRequest, chat: bool): |
| request_id = f"nayara-{uuid.uuid4().hex}" |
| previous = "" |
| with generation_lock: |
| for partial in _generate(prompt, request): |
| delta, previous = partial[len(previous) :], partial |
| if chat: |
| payload = {"id": request_id, "object": "chat.completion.chunk", "model": MODEL_ID, "choices": [{"index": 0, "delta": {"content": delta}, "finish_reason": None}]} |
| else: |
| payload = {"id": request_id, "object": "text_completion", "model": MODEL_ID, "choices": [{"index": 0, "text": delta, "finish_reason": None}]} |
| yield f"data: {json.dumps(payload)}\n\n" |
| yield "data: [DONE]\n\n" |
|
|
|
|
| @asynccontextmanager |
| async def lifespan(_: FastAPI): |
| if os.environ.get("NAYARA_REQUIRE_API_KEY") == "1" and not os.environ.get("NAYARA_API_KEY"): |
| raise RuntimeError("NAYARA_REQUIRE_API_KEY=1 but NAYARA_API_KEY is unset") |
| checkpoint = await run_in_threadpool(materialize_checkpoint) |
| threads = int(os.environ.get("NAYARA_THREADS", max(1, (os.cpu_count() or 4) // 2))) |
| runtime["tokenizer"] = Tokenizer.from_file(str(checkpoint / "tokenizer.json")) |
| runtime["model"] = await run_in_threadpool(load_model, checkpoint, True, threads) |
| runtime.update({"ready": True, "checkpoint": str(checkpoint), "loaded_at": time.time()}) |
| yield |
| runtime.clear() |
| runtime["ready"] = False |
|
|
|
|
| app = FastAPI(title="NAYARA API", version="1.0.0", lifespan=lifespan) |
|
|
|
|
| @app.get("/healthz") |
| def health() -> dict[str, str]: |
| return {"status": "ok"} |
|
|
|
|
| @app.get("/readyz") |
| def ready() -> dict[str, Any]: |
| if not runtime.get("ready"): |
| raise HTTPException(status_code=503, detail="model is not loaded") |
| return {"status": "ready", "model": MODEL_ID, "loaded_at": runtime["loaded_at"]} |
|
|
|
|
| @app.get("/v1/models", dependencies=[Depends(_authorize)]) |
| def models() -> dict[str, Any]: |
| return {"object": "list", "data": [{"id": MODEL_ID, "object": "model", "owned_by": "nayara"}]} |
|
|
|
|
| @app.post("/v1/completions", dependencies=[Depends(_authorize)]) |
| async def completions(request: CompletionRequest): |
| prompt = request.prompt |
| if request.stream: |
| return StreamingResponse(_sse(prompt, request, False), media_type="text/event-stream") |
| text = await run_in_threadpool(_complete, prompt, request) |
| return {"id": f"nayara-{uuid.uuid4().hex}", "object": "text_completion", "model": MODEL_ID, "choices": [{"index": 0, "text": text, "finish_reason": "stop"}]} |
|
|
|
|
| @app.post("/v1/chat/completions", dependencies=[Depends(_authorize)]) |
| async def chat_completions(request: ChatRequest): |
| prompt = _chat_prompt(request.messages) |
| if request.stream: |
| return StreamingResponse(_sse(prompt, request, True), media_type="text/event-stream") |
| text = await run_in_threadpool(_complete, prompt, request) |
| return {"id": f"nayara-{uuid.uuid4().hex}", "object": "chat.completion", "model": MODEL_ID, "choices": [{"index": 0, "message": {"role": "assistant", "content": text}, "finish_reason": "stop"}]} |
|
|
