server/stream.py

"""Translate the router's OpenAI-shaped SSE into the events the UI consumes.

The translation is a pure function over decoded frames so it can be tested
against recorded fixtures with no network and no credentials, which is what
lets the test job gate the deploy without ever holding a token.
"""

from __future__ import annotations

import json
from dataclasses import dataclass, field
from typing import Iterator


@dataclass
class Accumulated:
    """Running totals, so a stream that dies mid-flight still reports what it got."""
    content: str = ""
    reasoning: str = ""
    stats: dict = field(default_factory=dict)
    finished: bool = False
    # Tool calls arrive as fragments keyed by index: the name in one chunk, the
    # arguments a character at a time across dozens more. They are only useful
    # once the whole call has landed, so they are assembled here and emitted at
    # the end rather than forwarded piecemeal.
    tool_calls: dict = field(default_factory=dict)
    tools_sent: bool = False
    # Some providers drop the separate reasoning channel and inline the model's
    # thinking as <think>...</think> in the content instead. `thinking` tracks
    # which side of that tag we are on, and `held` is the part of a tag that
    # arrived split across two chunks.
    thinking: bool = False
    held: str = ""


def sse(event: str, data: str) -> str:
    """One SSE frame. Data is split so a newline in the payload cannot end it."""
    body = "".join(f"data: {line}\n" for line in data.split("\n"))
    return f"event: {event}\n{body}\n"


def translate(raw_line: str, acc: Accumulated) -> Iterator[str]:
    """Feed one upstream line; yield zero or more frames for the browser."""
    line = raw_line.strip()
    if not line or not line.startswith("data:"):
        return
    payload = line[5:].strip()

    if payload == "[DONE]":
        # Anything still held back was the start of a tag that never completed,
        # so it was ordinary text after all.
        if acc.held:
            acc.content += acc.held
            yield sse("delta", json.dumps({"content": acc.held}))
            acc.held = ""
        acc.finished = True
        yield from _tool_frame(acc)
        yield sse("stats", json.dumps(acc.stats))
        return

    try:
        chunk = json.loads(payload)
    except json.JSONDecodeError:
        return

    # Some providers report an error inside a 200 stream rather than as a
    # status code, which is otherwise indistinguishable from a silent stall.
    if isinstance(chunk.get("error"), (dict, str)):
        err = chunk["error"]
        msg = err.get("message") if isinstance(err, dict) else str(err)
        yield sse("error", msg or "The provider returned an error.")
        return

    choices = chunk.get("choices") or []
    if choices:
        choice = choices[0]
        delta = choice.get("delta") or {}
        out: dict = {}
        if delta.get("content"):
            body, inline_thought = split_thinking(delta["content"], acc)
            if body:
                out["content"] = body
                acc.content += body
            if inline_thought:
                out["reasoning"] = inline_thought
                acc.reasoning += inline_thought
        # GLM-5.3 streams thinking on its own channel; some providers spell it
        # `reasoning` rather than `reasoning_content`.
        thought = delta.get("reasoning_content") or delta.get("reasoning")
        if thought:
            out["reasoning"] = out.get("reasoning", "") + thought
            acc.reasoning += thought
        if out:
            yield sse("delta", json.dumps(out))
        _collect_tool_calls(delta.get("tool_calls"), acc)
        if choice.get("finish_reason"):
            acc.stats["finishReason"] = choice["finish_reason"]
            if choice["finish_reason"] == "tool_calls":
                yield from _tool_frame(acc)

    if chunk.get("id"):
        acc.stats["requestId"] = chunk["id"]
    if chunk.get("model"):
        acc.stats["model"] = chunk["model"].split("/")[-1]

    usage = chunk.get("usage")
    if isinstance(usage, dict):
        acc.stats.update(_usage(usage))
        yield sse("stats", json.dumps(acc.stats))


THINK_OPEN = "<think>"
THINK_CLOSE = "</think>"


def split_thinking(text: str, acc: Accumulated) -> tuple[str, str]:
    """Route one content delta into (answer, thinking).

    A model whose provider drops the reasoning channel writes its thinking into
    the content as ``<think>…</think>``, and the reader gets the scratchpad and
    a stray closing tag in the middle of the answer.

    The tags straddle chunk boundaries constantly -- ``<thi`` then ``nk>`` -- so
    a trailing fragment that could still become one is held back until the next
    chunk decides. Stateful across the stream, pure in everything else.
    """
    buf = acc.held + text
    acc.held = ""
    answer: list[str] = []
    thinking: list[str] = []

    while buf:
        tag = THINK_CLOSE if acc.thinking else THINK_OPEN
        into = thinking if acc.thinking else answer
        at = buf.find(tag)
        if at >= 0:
            into.append(buf[:at])
            buf = buf[at + len(tag):]
            acc.thinking = not acc.thinking
            continue
        keep = _partial_tag(buf, tag)
        into.append(buf[:len(buf) - keep] if keep else buf)
        acc.held = buf[len(buf) - keep:] if keep else ""
        break

    return "".join(answer), "".join(thinking)


def _partial_tag(text: str, tag: str) -> int:
    """Length of the trailing fragment that could still grow into ``tag``."""
    for n in range(min(len(tag) - 1, len(text)), 0, -1):
        if text.endswith(tag[:n]):
            return n
    return 0


def _collect_tool_calls(chunks, acc: Accumulated) -> None:
    """Merge one chunk's worth of tool-call fragments into the accumulator."""
    for tc in chunks or []:
        if not isinstance(tc, dict):
            continue
        # A provider that omits `index` is streaming one call at a time.
        index = tc.get("index", 0)
        slot = acc.tool_calls.setdefault(index, {"id": "", "name": "", "arguments": ""})
        if tc.get("id"):
            slot["id"] = str(tc["id"])
        fn = tc.get("function") or {}
        if fn.get("name"):
            slot["name"] = str(fn["name"])
        if fn.get("arguments"):
            slot["arguments"] += str(fn["arguments"])


def _tool_frame(acc: Accumulated) -> Iterator[str]:
    """The assembled calls, once. Emitted at `tool_calls` or at [DONE]."""
    if acc.tools_sent or not acc.tool_calls:
        return
    acc.tools_sent = True
    calls = []
    for index in sorted(acc.tool_calls):
        call = acc.tool_calls[index]
        # A fragment that never carried a name is an incomplete call; forwarding
        # it would have the client invoke a tool called "".
        if not call["name"]:
            continue
        # Some providers omit the id entirely. The round trip only needs it to
        # be unique within the turn, and the model is handed back whatever we
        # send here, so an index-derived one is as good as theirs.
        call["id"] = call["id"] or f"call_{index}_{call['name']}"
        calls.append(call)
    if calls:
        yield sse("tools", json.dumps(calls))


def _usage(u: dict) -> dict:
    """Flatten the usage block, including the fields most clients drop."""
    out: dict = {}
    for src, dst in (
        ("prompt_tokens", "promptTokens"),
        ("completion_tokens", "completionTokens"),
    ):
        if u.get(src) is not None:
            out[dst] = u[src]

    prompt_detail = u.get("prompt_tokens_details") or {}
    if prompt_detail.get("cached_tokens") is not None:
        out["cachedTokens"] = prompt_detail["cached_tokens"]

    detail = u.get("completion_tokens_details") or {}
    for src, dst in (
        ("reasoning_tokens", "reasoningTokens"),
        ("accepted_prediction_tokens", "acceptedTokens"),
        ("rejected_prediction_tokens", "rejectedTokens"),
    ):
        if detail.get(src) is not None:
            out[dst] = detail[src]
    return out


MAX_TOOLS = 32
_ROLES = ("system", "user", "assistant", "tool")


def normalise_messages(messages: list) -> list:
    """Turn the client's messages into the shape the router expects.

    The browser speaks the same camelCase it uses everywhere else; the OpenAI
    wire format is snake_case and nests a tool call inside a `function` object.
    Translating here rather than in the client keeps one JSON shape in the UI
    and one on the wire, and gives the round trip something to test.

    An assistant message that asked for tools must keep its `tool_calls`, and
    every one of them must be answered by a `tool` message: providers reject
    the pair broken apart.
    """
    out: list = []
    for raw in messages:
        if not isinstance(raw, dict):
            continue
        role = raw.get("role")
        if role not in _ROLES:
            continue
        msg: dict = {"role": role, "content": str(raw.get("content") or "")}

        if role == "tool":
            call_id = raw.get("toolCallId") or raw.get("tool_call_id")
            if not call_id:
                continue
            msg["tool_call_id"] = str(call_id)
            if raw.get("name"):
                msg["name"] = str(raw["name"])

        calls = raw.get("toolCalls") or raw.get("tool_calls")
        if role == "assistant" and isinstance(calls, list) and calls:
            msg["tool_calls"] = [
                {
                    "id": str(c.get("id") or ""),
                    "type": "function",
                    "function": {
                        "name": str(c.get("name") or ""),
                        "arguments": _json_object(c.get("arguments")),
                    },
                }
                for c in calls
                if isinstance(c, dict) and c.get("name")
            ]
        out.append(msg)
    return out


def _json_object(value) -> str:
    """A tool call's arguments, as something the router will accept.

    The router rejects the whole request with a 400 when any call's arguments
    are not valid JSON -- so a model that wrote a JavaScript object literal
    kills the conversation on the *next* turn, when its own call is echoed
    back, and the error names neither the call nor the turn it came from.

    The client repairs what it can before it gets here. This is the guarantee:
    what still cannot be parsed becomes an empty object, and the `tool` message
    answering it already says why.
    """
    text = value if isinstance(value, str) else json.dumps(value or {})
    text = text.strip() or "{}"
    try:
        parsed = json.loads(text)
    except (ValueError, TypeError):
        return "{}"
    # An array or a scalar is valid JSON and still not an argument list.
    return text if isinstance(parsed, dict) else "{}"


def normalise_tools(tools) -> list:
    """Accept only well-formed function tools, and not many of them."""
    if not isinstance(tools, list):
        return []
    out = []
    for t in tools[:MAX_TOOLS]:
        if not isinstance(t, dict):
            continue
        fn = t.get("function")
        if not isinstance(fn, dict) or not isinstance(fn.get("name"), str) or not fn["name"]:
            continue
        out.append({
            "type": "function",
            "function": {
                "name": fn["name"],
                "description": str(fn.get("description") or ""),
                "parameters": fn.get("parameters") if isinstance(fn.get("parameters"), dict) else {},
            },
        })
    return out


def build_payload(body: dict, model: str, reasoning: bool) -> dict:
    """The request we send upstream, from the request the browser sent us."""
    payload: dict = {
        "model": model,
        "messages": normalise_messages(body["messages"]),
        "stream": True,
        # Per-chunk usage is what feeds the live token readouts. Without it the
        # totals only arrive at the end, and the inspector sits empty.
        "stream_options": {"include_usage": True},
        "max_tokens": int(body.get("maxTokens") or 4096),
    }
    if body.get("temperature") is not None:
        payload["temperature"] = float(body["temperature"])

    tools = normalise_tools(body.get("tools"))
    if tools:
        payload["tools"] = tools
        payload["tool_choice"] = "auto"

    if reasoning:
        effort = body.get("reasoningEffort")
        # The model card: the parameter defaults to `max` and only `low` and
        # `high` need passing. Sending `max` explicitly is a no-op, so don't.
        if effort in ("low", "high"):
            payload["reasoning_effort"] = effort
        # Also from the card: chat scenarios should set this, or the model
        # carries prior turns' thinking back into context.
        payload["chat_template_kwargs"] = {"clear_thinking": True}
    return payload