feat: add OpenAI-compatible POST /chat/completions endpoint with streaming and non-streaming support

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)

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))
+

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"]