diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 3917e38e36..36347d4c13 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1,8 +1,13 @@ """ Pydantic schemas for Inference API """ +from __future__ import annotations + +import time +import uuid +from typing import Literal, Optional, List + from pydantic import BaseModel, Field -from typing import Optional, List class LoadRequest(BaseModel): @@ -52,3 +57,91 @@ class InferenceStatusResponse(BaseModel): is_vision: bool = Field(False, description="Whether the active model is a vision model") loading: List[str] = Field(default_factory=list, description="Models currently being loaded") loaded: List[str] = Field(default_factory=list, description="Models currently loaded") + + +# ===================================================================== +# OpenAI-Compatible Chat Completions Models +# ===================================================================== + + +class ChatMessage(BaseModel): + """A single message in the conversation.""" + role: Literal["system", "user", "assistant"] = Field(..., description="Message role") + content: str = Field(..., description="Message content") + + +class ChatCompletionRequest(BaseModel): + """ + OpenAI-compatible chat completion request. + + Extensions (non-OpenAI fields) are marked with 'x-unsloth'. + """ + model: str = Field("default", description="Model identifier (informational; the active model is used)") + messages: list[ChatMessage] = Field(..., description="Conversation messages") + stream: bool = Field(True, description="Whether to stream the response via SSE") + temperature: float = Field(0.7, ge=0.0, le=2.0) + top_p: float = Field(0.9, ge=0.0, le=1.0) + max_tokens: Optional[int] = Field(512, ge=1, le=4096, description="Maximum tokens to generate") + + # ── Unsloth extensions (ignored by standard OpenAI clients) ── + top_k: int = Field(40, ge=1, le=100, description="[x-unsloth] Top-k sampling") + repetition_penalty: float = Field(1.1, ge=1.0, le=2.0, description="[x-unsloth] Repetition penalty") + image_base64: Optional[str] = Field(None, description="[x-unsloth] Base64-encoded image for vision models") + + +# ── Streaming response chunks ──────────────────────────────────── + + +class ChoiceDelta(BaseModel): + """Delta content for a streaming chunk.""" + role: Optional[str] = None + content: Optional[str] = None + + +class ChunkChoice(BaseModel): + """A single choice in a streaming chunk.""" + index: int = 0 + delta: ChoiceDelta + finish_reason: Optional[Literal["stop", "length"]] = None + + +class ChatCompletionChunk(BaseModel): + """A single SSE chunk in OpenAI streaming format.""" + id: str = Field(default_factory=lambda: f"chatcmpl-{uuid.uuid4().hex[:12]}") + object: Literal["chat.completion.chunk"] = "chat.completion.chunk" + created: int = Field(default_factory=lambda: int(time.time())) + model: str = "default" + choices: list[ChunkChoice] + + +# ── Non-streaming response ─────────────────────────────────────── + + +class CompletionMessage(BaseModel): + """The assistant's complete response message.""" + role: Literal["assistant"] = "assistant" + content: str + + +class CompletionChoice(BaseModel): + """A single choice in a non-streaming response.""" + index: int = 0 + message: CompletionMessage + finish_reason: Literal["stop", "length"] = "stop" + + +class CompletionUsage(BaseModel): + """Token usage statistics (approximate).""" + prompt_tokens: int = 0 + completion_tokens: int = 0 + total_tokens: int = 0 + + +class ChatCompletion(BaseModel): + """Non-streaming chat completion response.""" + id: str = Field(default_factory=lambda: f"chatcmpl-{uuid.uuid4().hex[:12]}") + object: Literal["chat.completion"] = "chat.completion" + created: int = Field(default_factory=lambda: int(time.time())) + model: str = "default" + choices: list[CompletionChoice] + usage: CompletionUsage = Field(default_factory=CompletionUsage) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 17d2011205..e98165d83e 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -2,13 +2,17 @@ Inference API routes for model loading and text generation. """ import sys +import time +import uuid from pathlib import Path from fastapi import APIRouter, HTTPException -from fastapi.responses import StreamingResponse +from fastapi.responses import StreamingResponse, JSONResponse from typing import Optional import json import logging + + # Add backend directory to path backend_path = Path(__file__).parent.parent.parent if str(backend_path) not in sys.path: @@ -32,6 +36,13 @@ from models.inference import ( LoadResponse, UnloadResponse, InferenceStatusResponse, + ChatCompletionRequest, + ChatCompletionChunk, + ChatCompletion, + ChunkChoice, + ChoiceDelta, + CompletionChoice, + CompletionMessage, ) router = APIRouter() @@ -222,3 +233,174 @@ async def get_status(): status_code=500, detail=f"Failed to get status: {str(e)}" ) + + +# ===================================================================== +# OpenAI-Compatible Chat Completions (/chat/completions) +# ===================================================================== + + + + +@router.post("/chat/completions") +async def openai_chat_completions(request: ChatCompletionRequest): + """ + OpenAI-compatible chat completions endpoint. + + Streaming (default): returns SSE chunks matching OpenAI's format. + Non-streaming: returns a single ChatCompletion JSON object. + """ + backend = get_inference_backend() + + if not backend.active_model_name: + raise HTTPException( + status_code=400, + detail="No model loaded. Call POST /inference/load first.", + ) + + # ── Extract system prompt from messages ─────────────────────── + system_prompt = "You are a helpful AI assistant." + chat_messages: list[dict] = [] + + for msg in request.messages: + if msg.role == "system": + system_prompt = msg.content + else: + chat_messages.append({"role": msg.role, "content": msg.content}) + + # If no non-system messages were provided, error out + if not chat_messages: + raise HTTPException( + status_code=400, + detail="At least one non-system message is required.", + ) + + # ── Decode image if provided (vision models) ────────────────── + image = None + if request.image_base64: + try: + import base64 + from PIL import Image + from io import BytesIO + + model_info = backend.models.get(backend.active_model_name, {}) + if not model_info.get("is_vision"): + raise HTTPException( + status_code=400, + detail="Image provided but current model is text-only.", + ) + + image_data = base64.b64decode(request.image_base64) + image = Image.open(BytesIO(image_data)) + image = backend.resize_image(image) + + except HTTPException: + raise + except Exception as e: + raise HTTPException(status_code=400, detail=f"Failed to decode image: {e}") + + # ── Shared generation kwargs ────────────────────────────────── + gen_kwargs = dict( + messages=chat_messages, + system_prompt=system_prompt, + image=image, + temperature=request.temperature, + top_p=request.top_p, + top_k=request.top_k, + max_new_tokens=request.max_tokens or 512, + repetition_penalty=request.repetition_penalty, + ) + + model_name = backend.active_model_name or request.model + completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}" + created = int(time.time()) + + # ── Streaming response ──────────────────────────────────────── + if request.stream: + async def stream_chunks(): + try: + # First chunk: send the role + first_chunk = ChatCompletionChunk( + id=completion_id, + created=created, + model=model_name, + choices=[ChunkChoice( + delta=ChoiceDelta(role="assistant"), + finish_reason=None, + )], + ) + yield f"data: {first_chunk.model_dump_json(exclude_none=True)}\n\n" + + # Content chunks — generate_chat_response yields cumulative + # text, so we diff to get incremental deltas. + prev_text = "" + for cumulative in backend.generate_chat_response(**gen_kwargs): + new_text = cumulative[len(prev_text):] + prev_text = cumulative + if not new_text: + continue + chunk = ChatCompletionChunk( + id=completion_id, + created=created, + model=model_name, + choices=[ChunkChoice( + delta=ChoiceDelta(content=new_text), + finish_reason=None, + )], + ) + yield f"data: {chunk.model_dump_json(exclude_none=True)}\n\n" + + # Final chunk: finish_reason = stop + final_chunk = ChatCompletionChunk( + id=completion_id, + created=created, + model=model_name, + choices=[ChunkChoice( + delta=ChoiceDelta(), + finish_reason="stop", + )], + ) + yield f"data: {final_chunk.model_dump_json(exclude_none=True)}\n\n" + yield "data: [DONE]\n\n" + + except Exception as e: + backend.reset_generation_state() + logger.error(f"Error during OpenAI streaming: {e}", exc_info=True) + error_chunk = { + "error": {"message": str(e), "type": "server_error"}, + } + yield f"data: {json.dumps(error_chunk)}\n\n" + + return StreamingResponse( + stream_chunks(), + media_type="text/event-stream", + headers={ + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "X-Accel-Buffering": "no", + }, + ) + + # ── Non-streaming response ──────────────────────────────────── + else: + try: + full_text = "" + for token in backend.generate_chat_response(**gen_kwargs): + full_text = token # generate_stream yields cumulative text + + response = ChatCompletion( + id=completion_id, + created=created, + model=model_name, + choices=[CompletionChoice( + message=CompletionMessage(content=full_text), + finish_reason="stop", + )], + ) + return JSONResponse(content=response.model_dump()) + + except Exception as e: + backend.reset_generation_state() + logger.error(f"Error during OpenAI completion: {e}", exc_info=True) + raise HTTPException(status_code=500, detail=str(e)) + diff --git a/studio/tests/test_openai_completions.py b/studio/tests/test_openai_completions.py new file mode 100644 index 0000000000..31bbc6d51e --- /dev/null +++ b/studio/tests/test_openai_completions.py @@ -0,0 +1,282 @@ +""" +Tests for the OpenAI-compatible /chat/completions endpoint. + +Validates: + - Streaming: SSE chunk format matches OpenAI spec + - Non-streaming: single JSON ChatCompletion response + - System prompt extraction from messages array + - Request validation (no messages, missing model, etc.) + - Response headers for proxy compatibility + +All tests mock the inference backend and bypass auth. +""" +import sys +import json +from pathlib import Path +from unittest.mock import MagicMock, patch, PropertyMock + +import pytest + +# ── Path setup ──────────────────────────────────────────────────── +_backend_root = Path(__file__).resolve().parent.parent / "backend" +if str(_backend_root) not in sys.path: + sys.path.insert(0, str(_backend_root)) + +from fastapi.testclient import TestClient +from main import app + + +# ── Fixtures ────────────────────────────────────────────────────── + +def _make_mock_backend(*, tokens: list[str] | None = None, active_model: str = "test-model"): + """Build a mock InferenceBackend that yields preset tokens.""" + backend = MagicMock() + backend.active_model_name = active_model + backend.models = {active_model: {"is_vision": False}} + + def fake_generate(**kwargs): + for t in (tokens or ["Hello", "Hello world", "Hello world!"]): + yield t + + backend.generate_chat_response = MagicMock(side_effect=fake_generate) + backend.reset_generation_state = MagicMock() + return backend + + +def _parse_sse_data(raw: str) -> list[dict | str]: + """Extract `data:` payloads from raw SSE text. Returns dicts or raw strings.""" + results = [] + for line in raw.split("\n"): + if line.startswith("data: "): + payload = line[len("data: "):] + if payload == "[DONE]": + results.append("[DONE]") + else: + try: + results.append(json.loads(payload)) + except json.JSONDecodeError: + results.append(payload) + return results + + +@pytest.fixture() +def client(): + yield TestClient(app) + + +# ===================================================================== +# Streaming tests +# ===================================================================== + + +class TestStreamingChunkFormat: + """Each SSE chunk must match the OpenAI chat.completion.chunk schema.""" + + def test_chunks_have_required_fields(self, client: TestClient): + mock_backend = _make_mock_backend(tokens=["Hi"]) + with patch("routes.inference.get_inference_backend", return_value=mock_backend): + resp = client.post( + "/api/inference/chat/completions", + json={ + "messages": [{"role": "user", "content": "Hello"}], + "stream": True, + }, + ) + + assert resp.status_code == 200 + chunks = _parse_sse_data(resp.text) + + # Filter to actual chunk dicts (not [DONE]) + json_chunks = [c for c in chunks if isinstance(c, dict) and "choices" in c] + assert len(json_chunks) >= 2 # role chunk + content chunk(s) + final + + for chunk in json_chunks: + assert "id" in chunk + assert chunk["object"] == "chat.completion.chunk" + assert "created" in chunk + assert "model" in chunk + assert len(chunk["choices"]) == 1 + assert "delta" in chunk["choices"][0] + + def test_first_chunk_has_role(self, client: TestClient): + mock_backend = _make_mock_backend(tokens=["Hi"]) + with patch("routes.inference.get_inference_backend", return_value=mock_backend): + resp = client.post( + "/api/inference/chat/completions", + json={"messages": [{"role": "user", "content": "Hello"}]}, + ) + + chunks = [c for c in _parse_sse_data(resp.text) if isinstance(c, dict) and "choices" in c] + first = chunks[0] + assert first["choices"][0]["delta"].get("role") == "assistant" + + def test_last_chunk_has_stop_finish_reason(self, client: TestClient): + mock_backend = _make_mock_backend(tokens=["Done"]) + with patch("routes.inference.get_inference_backend", return_value=mock_backend): + resp = client.post( + "/api/inference/chat/completions", + json={"messages": [{"role": "user", "content": "Hello"}]}, + ) + + chunks = [c for c in _parse_sse_data(resp.text) if isinstance(c, dict) and "choices" in c] + last = chunks[-1] + assert last["choices"][0]["finish_reason"] == "stop" + # Delta should be empty on the final chunk + assert last["choices"][0]["delta"].get("content") is None + + def test_stream_ends_with_done(self, client: TestClient): + mock_backend = _make_mock_backend(tokens=["x"]) + with patch("routes.inference.get_inference_backend", return_value=mock_backend): + resp = client.post( + "/api/inference/chat/completions", + json={"messages": [{"role": "user", "content": "Hello"}]}, + ) + + all_data = _parse_sse_data(resp.text) + assert all_data[-1] == "[DONE]" + + def test_consistent_id_across_chunks(self, client: TestClient): + mock_backend = _make_mock_backend(tokens=["a", "b", "c"]) + with patch("routes.inference.get_inference_backend", return_value=mock_backend): + resp = client.post( + "/api/inference/chat/completions", + json={"messages": [{"role": "user", "content": "Hello"}]}, + ) + + chunks = [c for c in _parse_sse_data(resp.text) if isinstance(c, dict) and "choices" in c] + ids = set(c["id"] for c in chunks) + assert len(ids) == 1, "All chunks should share the same completion ID" + + +class TestStreamingHeaders: + """Verify response headers for SSE proxy compatibility.""" + + def test_headers(self, client: TestClient): + mock_backend = _make_mock_backend(tokens=["x"]) + with patch("routes.inference.get_inference_backend", return_value=mock_backend): + resp = client.post( + "/api/inference/chat/completions", + json={"messages": [{"role": "user", "content": "Hello"}]}, + ) + + assert resp.headers["content-type"].startswith("text/event-stream") + assert resp.headers.get("cache-control") == "no-cache" + assert resp.headers.get("x-accel-buffering") == "no" + + +# ===================================================================== +# Non-streaming tests +# ===================================================================== + + +class TestNonStreaming: + """When stream=false, return a single ChatCompletion JSON object.""" + + def test_returns_json_object(self, client: TestClient): + mock_backend = _make_mock_backend(tokens=["Full response text"]) + with patch("routes.inference.get_inference_backend", return_value=mock_backend): + resp = client.post( + "/api/inference/chat/completions", + json={ + "messages": [{"role": "user", "content": "Hello"}], + "stream": False, + }, + ) + + assert resp.status_code == 200 + body = resp.json() + assert body["object"] == "chat.completion" + assert body["choices"][0]["message"]["role"] == "assistant" + assert body["choices"][0]["message"]["content"] == "Full response text" + assert body["choices"][0]["finish_reason"] == "stop" + + def test_non_streaming_has_model(self, client: TestClient): + mock_backend = _make_mock_backend(tokens=["x"], active_model="my-model") + with patch("routes.inference.get_inference_backend", return_value=mock_backend): + resp = client.post( + "/api/inference/chat/completions", + json={ + "messages": [{"role": "user", "content": "Hi"}], + "stream": False, + }, + ) + + body = resp.json() + assert body["model"] == "my-model" + + +# ===================================================================== +# System prompt extraction +# ===================================================================== + + +class TestSystemPromptExtraction: + """System messages should be extracted and passed as system_prompt.""" + + def test_system_message_extracted(self, client: TestClient): + mock_backend = _make_mock_backend(tokens=["ok"]) + with patch("routes.inference.get_inference_backend", return_value=mock_backend): + client.post( + "/api/inference/chat/completions", + json={ + "messages": [ + {"role": "system", "content": "You are a pirate."}, + {"role": "user", "content": "Hello"}, + ], + "stream": False, + }, + ) + + # Check that generate_chat_response was called with the correct system_prompt + call_kwargs = mock_backend.generate_chat_response.call_args[1] + assert call_kwargs["system_prompt"] == "You are a pirate." + # System message should NOT be in the chat_messages list + assert all(m["role"] != "system" for m in call_kwargs["messages"]) + + def test_default_system_prompt_when_none(self, client: TestClient): + mock_backend = _make_mock_backend(tokens=["ok"]) + with patch("routes.inference.get_inference_backend", return_value=mock_backend): + client.post( + "/api/inference/chat/completions", + json={ + "messages": [{"role": "user", "content": "Hello"}], + "stream": False, + }, + ) + + call_kwargs = mock_backend.generate_chat_response.call_args[1] + assert call_kwargs["system_prompt"] == "You are a helpful AI assistant." + + +# ===================================================================== +# Error handling +# ===================================================================== + + +class TestErrorHandling: + """Validate error responses for bad requests.""" + + def test_no_model_loaded(self, client: TestClient): + mock_backend = _make_mock_backend() + mock_backend.active_model_name = None + with patch("routes.inference.get_inference_backend", return_value=mock_backend): + resp = client.post( + "/api/inference/chat/completions", + json={"messages": [{"role": "user", "content": "Hi"}]}, + ) + + assert resp.status_code == 400 + assert "No model loaded" in resp.json()["detail"] + + def test_only_system_messages_rejected(self, client: TestClient): + mock_backend = _make_mock_backend() + with patch("routes.inference.get_inference_backend", return_value=mock_backend): + resp = client.post( + "/api/inference/chat/completions", + json={ + "messages": [{"role": "system", "content": "You are a bot."}], + }, + ) + + assert resp.status_code == 400 + assert "non-system message" in resp.json()["detail"]