feat(server): add buffered streaming responses

Accept text-only OpenAI content blocks and return Chat or Completions results through SSE framing when stream=true. Preserve family-owned generation behavior by buffering the native result until the current Task contract completes.

Document direct Pi configuration and keep tool calling and incremental token delivery outside this MVP.

Signed-off-by: Vivian Chen <140748220+xuanzic@users.noreply.github.com>

Author: xuanzic

apps/server/python/tests/test_app.py

@@ -128,7 +128,13 @@ def test_completion_and_chat_keep_model_semantics_in_worker(caplog: Any) -> None
                 "model": "test/model",
                 "messages": [
                     {"role": "system", "content": "Be brief"},
-                    {"role": "user", "content": "Capital?"},
+                    {
+                        "role": "user",
+                        "content": [
+                            {"type": "text", "text": "Cap"},
+                            {"type": "text", "text": "ital?"},
+                        ],
+                    },
                 ],
                 "max_completion_tokens": 8,
             },
@@ -138,6 +144,7 @@ def test_completion_and_chat_keep_model_semantics_in_worker(caplog: Any) -> None
         assert config["use_chat_template"] is True
         assert config["system_prompt"] == "Be brief"
         assert config["max_new_tokens"] == 8
+        assert registry.requests[-1][1]["prompt"] == "Capital?"
 
     request_logs = [
         json.loads(record.message)
@@ -153,6 +160,86 @@ def test_completion_and_chat_keep_model_semantics_in_worker(caplog: Any) -> None
     assert "Capital?" not in json.dumps(request_logs)
 
 
+def test_streaming_chat_uses_openai_sse_shape() -> None:
+    registry = FakeRegistry()
+    with make_client(registry) as client:
+        response = client.post(
+            "/v1/chat/completions",
+            json={
+                "model": "test/model",
+                "messages": [
+                    {
+                        "role": "user",
+                        "content": [{"type": "text", "text": "Capital?"}],
+                    }
+                ],
+                "stream": True,
+                "stream_options": {"include_usage": True},
+            },
+        )
+
+    assert response.status_code == 200
+    assert response.headers["content-type"].startswith("text/event-stream")
+    events = [event for event in response.text.split("\n\n") if event]
+    assert events[-1] == "data: [DONE]"
+    chunks = [json.loads(event.removeprefix("data: ")) for event in events[:-1]]
+    assert chunks[0]["object"] == "chat.completion.chunk"
+    assert chunks[0]["choices"][0]["delta"] == {
+        "role": "assistant",
+        "content": "Paris",
+    }
+    assert chunks[1]["choices"][0]["delta"] == {}
+    assert chunks[1]["choices"][0]["finish_reason"] is None
+    assert chunks[2]["choices"] == []
+    assert chunks[2]["usage"]["completion_tokens"] == 1
+    assert registry.requests[-1][1]["prompt"] == "Capital?"
+
+
+def test_streaming_completion_and_stream_options_validation() -> None:
+    registry = FakeRegistry()
+    with make_client(registry) as client:
+        streamed = client.post(
+            "/v1/completions",
+            json={"model": "test/model", "prompt": "Capital?", "stream": True},
+        )
+        invalid = client.post(
+            "/v1/completions",
+            json={
+                "model": "test/model",
+                "prompt": "Capital?",
+                "stream_options": {"include_usage": True},
+            },
+        )
+
+    assert streamed.status_code == 200
+    assert '"object":"text_completion"' in streamed.text
+    assert '"text":"Paris"' in streamed.text
+    assert streamed.text.endswith("data: [DONE]\n\n")
+    assert invalid.status_code == 400
+    assert invalid.json()["error"]["param"] == "stream_options"
+
+
+def test_chat_rejects_non_text_content_blocks() -> None:
+    registry = FakeRegistry()
+    with make_client(registry) as client:
+        response = client.post(
+            "/v1/chat/completions",
+            json={
+                "model": "test/model",
+                "messages": [
+                    {
+                        "role": "user",
+                        "content": [{"type": "image_url", "image_url": {"url": "file:///tmp/x"}}],
+                    }
+                ],
+            },
+        )
+
+    assert response.status_code == 400
+    assert response.json()["error"]["param"] in {"str", "type", "text"}
+    assert registry.requests == []
+
+
 def test_validation_overload_and_metrics_are_explicit() -> None:
     registry = FakeRegistry()
     with make_client(registry) as client:

apps/server/python/trtmc_server/app.py

@@ -16,7 +16,7 @@
 
 from fastapi import FastAPI, Request
 from fastapi.exceptions import RequestValidationError
-from fastapi.responses import JSONResponse, PlainTextResponse
+from fastapi.responses import JSONResponse, PlainTextResponse, StreamingResponse
 from starlette.types import ASGIApp, Message, Receive, Scope, Send
 
 from .errors import (
@@ -128,6 +128,84 @@ def error_response(
     )
 
 
+def _sse_data(payload: dict[str, Any] | str) -> bytes:
+    data = payload if isinstance(payload, str) else json.dumps(payload, separators=(",", ":"))
+    return f"data: {data}\n\n".encode("utf-8")
+
+
+def _streaming_completion_response(
+    *,
+    response_id: str,
+    created: int,
+    model: str,
+    text: str,
+    completion_tokens: int,
+    chat: bool,
+    include_usage: bool,
+    request_id: str,
+) -> StreamingResponse:
+    if chat:
+        object_name = "chat.completion.chunk"
+        content_choice: dict[str, Any] = {
+            "index": 0,
+            "delta": {"role": "assistant", "content": text},
+            "logprobs": None,
+            "finish_reason": None,
+        }
+        terminal_choice: dict[str, Any] = {
+            "index": 0,
+            "delta": {},
+            "logprobs": None,
+            "finish_reason": None,
+        }
+    else:
+        object_name = "text_completion"
+        content_choice = {
+            "index": 0,
+            "text": text,
+            "logprobs": None,
+            "finish_reason": None,
+        }
+        terminal_choice = {
+            "index": 0,
+            "text": "",
+            "logprobs": None,
+            "finish_reason": None,
+        }
+
+    def chunk(choices: list[dict[str, Any]]) -> dict[str, Any]:
+        return {
+            "id": response_id,
+            "object": object_name,
+            "created": created,
+            "model": model,
+            "choices": choices,
+        }
+
+    async def events() -> AsyncIterator[bytes]:
+        yield _sse_data(chunk([content_choice]))
+        yield _sse_data(chunk([terminal_choice]))
+        if include_usage:
+            usage_chunk = chunk([])
+            usage_chunk["usage"] = {
+                "prompt_tokens": 0,
+                "completion_tokens": completion_tokens,
+                "total_tokens": completion_tokens,
+            }
+            yield _sse_data(usage_chunk)
+        yield _sse_data("[DONE]")
+
+    return StreamingResponse(
+        events(),
+        media_type="text/event-stream",
+        headers={
+            "Cache-Control": "no-cache",
+            "X-Accel-Buffering": "no",
+            "X-Request-ID": request_id,
+        },
+    )
+
+
 def create_app(registry: ModelRegistry, config: ServerConfig) -> FastAPI:
     metrics = Metrics()
 
@@ -238,10 +316,13 @@ async def execute(
         if request.n != 1:
             metrics.reject(route, 400)
             return error_response(400, "unsupported_parameter", "n must be 1", param="n")
-        if request.stream:
+        if request.stream_options is not None and not request.stream:
             metrics.reject(route, 400)
             return error_response(
-                400, "streaming_not_supported", "streaming is not available", param="stream"
+                400,
+                "invalid_request",
+                "stream_options requires stream=true",
+                param="stream_options",
             )
         if request.stop is not None:
             metrics.reject(route, 400)
@@ -386,11 +467,25 @@ async def execute(
             }
             object_name = "text_completion"
             response_id = f"cmpl-{uuid.uuid4().hex}"
+        created = int(time.time())
+        if request.stream:
+            return _streaming_completion_response(
+                response_id=response_id,
+                created=created,
+                model=request.model,
+                text=text,
+                completion_tokens=completion_tokens,
+                chat=chat,
+                include_usage=(
+                    request.stream_options is not None and request.stream_options.include_usage
+                ),
+                request_id=request_id,
+            )
         return JSONResponse(
             content={
                 "id": response_id,
                 "object": object_name,
-                "created": int(time.time()),
+                "created": created,
                 "model": request.model,
                 "choices": [choice],
                 "usage": {

apps/server/python/trtmc_server/protocol.py

@@ -9,7 +9,7 @@
 from typing import Any
 
 from .errors import WorkerProtocolError, WorkerRemoteError, WorkerRequestTooLargeError
-from .schemas import ChatCompletionRequest, GenerationRequest
+from .schemas import ChatCompletionRequest, GenerationRequest, TextContentPart
 
 
 def generation_config(request: GenerationRequest, max_tokens: int) -> dict[str, Any]:
@@ -24,16 +24,22 @@ def generation_config(request: GenerationRequest, max_tokens: int) -> dict[str,
 def chat_prompt(request: ChatCompletionRequest) -> tuple[str, str]:
     messages = request.messages
     if len(messages) == 1 and messages[0].role == "user":
-        return messages[0].content, ""
+        return _message_text(messages[0].content), ""
     if (
         len(messages) == 2
         and messages[0].role == "system"
         and messages[1].role == "user"
     ):
-        return messages[1].content, messages[0].content
+        return _message_text(messages[1].content), _message_text(messages[0].content)
     raise ValueError("messages must be one user message with an optional preceding system message")
 
 
+def _message_text(content: str | list[TextContentPart]) -> str:
+    if isinstance(content, str):
+        return content
+    return "".join(part.text for part in content)
+
+
 def extract_result(result: Any) -> tuple[str, int, dict[str, float]]:
     if not isinstance(result, Mapping) or not isinstance(result.get("text"), str):
         raise WorkerProtocolError("worker generate result does not contain text")

apps/server/python/trtmc_server/schemas.py

@@ -5,7 +5,7 @@
 
 from __future__ import annotations
 
-from typing import Any
+from typing import Any, Literal
 
 from pydantic import BaseModel, ConfigDict, Field
 
@@ -14,9 +14,18 @@ class StrictRequest(BaseModel):
     model_config = ConfigDict(extra="forbid")
 
 
+class TextContentPart(StrictRequest):
+    type: Literal["text"]
+    text: str
+
+
 class ChatMessage(StrictRequest):
     role: str
-    content: str
+    content: str | list[TextContentPart]
+
+
+class StreamOptions(StrictRequest):
+    include_usage: bool = False
 
 
 class GenerationRequest(StrictRequest):
@@ -30,6 +39,7 @@ class GenerationRequest(StrictRequest):
     enable_thinking: bool | None = None
     n: int = Field(default=1, ge=1)
     stream: bool = False
+    stream_options: StreamOptions | None = None
     stop: Any | None = None
 
 

website/docs/user-guides/serve-text-generation.md

@@ -3,7 +3,7 @@ title: Serve Text Generation
 description: Expose text-generation bundles through the local hybrid server.
 ---
 
-`trtmc-server` exposes a deliberately small, non-streaming subset of the
+`trtmc-server` exposes a deliberately small subset of the
 OpenAI Completions and Chat Completions protocols. It is a local evaluation
 server, not a production or distributed serving system.
 
@@ -96,18 +96,58 @@ support that placement.
 | Route | Supported behavior |
 | --- | --- |
 | `POST /v1/completions` | One string `prompt`; `model` is required. |
-| `POST /v1/chat/completions` | One text `user` message and optional preceding `system` message. |
+| `POST /v1/chat/completions` | One text `user` message and optional preceding `system` message; content may be a string or text-only content blocks. |
 | `GET /v1/models` | Configured text models. |
 | `GET /health/live` | Control-plane process liveness. |
 | `GET /health/ready` | Worker readiness and admission state. |
 | `GET /metrics` | Prometheus text metrics for admission and inference. |
 
 Generation accepts `temperature`, `top_p`, `min_p`, `top_k`, `seed`,
-`enable_thinking`, `n`, and `stream`. Completions accept `max_tokens`;
+`enable_thinking`, `n`, `stream`, and `stream_options.include_usage`.
+Completions accept `max_tokens`;
 chat accepts either `max_tokens` or `max_completion_tokens`. The MVP
-requires `n=1` and `stream=false`. It rejects unknown fields, prompt arrays,
-multi-turn chat, structured content, stop sequences, tools, log probabilities,
-and streaming instead of silently ignoring them.
+requires `n=1`. It rejects unknown fields, prompt arrays, non-text content
+blocks, multi-turn chat, stop sequences, tools, and log probabilities instead
+of silently ignoring them.
+
+Streaming responses use the OpenAI-compatible server-sent event framing and
+terminate with `data: [DONE]`. The native text task currently returns a complete
+generation, so the MVP buffers inference before emitting the content chunk;
+streaming provides client compatibility but does not yet reduce time to first
+token. Incremental token delivery requires a future extension to the family and
+runtime text-generation contract.
+
+For example, a Pi custom provider can use the local endpoint without an
+authentication header:
+
+```json
+{
+  "providers": {
+    "trtmc": {
+      "baseUrl": "http://127.0.0.1:8000/v1",
+      "api": "openai-completions",
+      "apiKey": "EMPTY",
+      "authHeader": false,
+      "compat": {
+        "supportsStore": false,
+        "supportsDeveloperRole": false,
+        "supportsReasoningEffort": false,
+        "supportsFinishReason": false,
+        "maxTokensField": "max_tokens"
+      },
+      "models": [{"id": "Qwen/Qwen3-0.6B"}]
+    }
+  }
+}
+```
+
+Pi sends text-only content blocks and consumes the buffered SSE response. Run
+Pi with `--no-tools`; tool definitions and tool-result messages remain outside
+the MVP protocol.
+
+```bash
+pi --provider trtmc --model Qwen/Qwen3-0.6B --no-tools
+```
 
 Model-specific chat templates, tokenization, sampling, stopping, engine
 composition, and validation remain inside the family selected by the bundle.