diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py index 009a33f0d..91fd82965 100644 --- a/python/freetoken/server/openai_api.py +++ b/python/freetoken/server/openai_api.py @@ -24,6 +24,7 @@ from .request_logger import log_request from .generation import ( DEFAULT_MAX_OUTPUT_TOKENS, + KEEPALIVE, ContentDelta, GenDone, GenerationError, @@ -40,8 +41,12 @@ resolve_sampling, timings_block, submit_generation, + with_keepalive, ) +KEEPALIVE_INTERVAL_S = 15.0 + + #: The wire superset plus "off", DeepSeek's disable synonym that #: effort_toggle_kwargs has always honored. _ACCEPTED_EFFORTS = (*KNOWN_REASONING_EFFORTS, "off") @@ -266,10 +271,12 @@ async def stream_chat_completion_chunks( """Format generate_events() into the OpenAI chat.completion.chunk SSE stream.""" if spec is None: spec = chat_request_to_genspec(req, {}) + created = int(time.time()) yield _sse( _chat_chunk( req, uid, + created, [{"delta": {"role": "assistant", "content": ""}, "index": 0, "finish_reason": None}], ) ) @@ -293,121 +300,143 @@ def emit(chunk: dict[str, Any]) -> bytes: tool_calls_sent = 0 open_tool: dict[str, Any] | None = None - events = generate_events(uid, spec, state, source="/v1/chat/completions", usage=usage) - while True: - try: - ev = await events.__anext__() - except StopAsyncIteration: - break - except GenerationError as exc: - # Request failed before producing output — emit an error chunk + [DONE] so the - # client gets a terminal signal instead of a stalled stream. - yield _sse( - {"error": {"message": str(exc), "type": "invalid_request_error", "code": exc.code}} - ) - break - if isinstance(ev, ReasoningDelta): - yield emit( - _chat_chunk( - req, - uid, - [{"delta": {"reasoning_content": ev.text}, "index": 0, "finish_reason": None}], + output_started = False + events = with_keepalive( + generate_events(uid, spec, state, source="/v1/chat/completions", usage=usage), + KEEPALIVE_INTERVAL_S, + ) + try: + while True: + try: + ev = await events.__anext__() + except StopAsyncIteration: + break + except GenerationError as exc: + # Request failed before producing output — emit an error chunk + [DONE] so the + # client gets a terminal signal instead of a stalled stream. + yield _sse( + {"error": {"message": str(exc), "type": "invalid_request_error", "code": exc.code}} ) - ) - elif isinstance(ev, ContentDelta): - yield emit( - _chat_chunk( - req, - uid, - [{"delta": {"content": ev.text}, "index": 0, "finish_reason": None}], + break + if ev is KEEPALIVE: + if not output_started: + # Event-based idle timers ignore SSE comments during queue/prefill silence. + # A chunk like any other after the role chunk: emit() adds the usage so + # far and, once an ack has delivered one, its timings snapshot. + yield emit(_chat_chunk( + req, uid, created, [{"delta": {}, "index": 0, "finish_reason": None}], + )) + continue + if isinstance(ev, (ContentDelta, ReasoningDelta)): + output_started = output_started or bool(ev.text) + else: + output_started = True + if isinstance(ev, ReasoningDelta): + yield emit( + _chat_chunk( + req, + uid, + created, + [{"delta": {"reasoning_content": ev.text}, "index": 0, "finish_reason": None}], + ) ) - ) - elif isinstance(ev, ToolCallStart): - open_tool = { - "index": tool_calls_sent, - "ordinal": ev.tool_index, - "sent": "", - "stable": ev.args_prefix_stable, - } - yield emit( - _chat_chunk( - req, uid, - [{ - "delta": {"tool_calls": [{ - "index": open_tool["index"], - "id": _tool_call_id(ev.name, open_tool["index"]), - "type": "function", - "function": {"name": ev.name, "arguments": ""}, - }]}, - "index": 0, "finish_reason": None, - }], + elif isinstance(ev, ContentDelta): + yield emit( + _chat_chunk( + req, + uid, + created, + [{"delta": {"content": ev.text}, "index": 0, "finish_reason": None}], + ) ) - ) - elif isinstance(ev, ToolCallArgsDelta): - # Clients concatenate argument fragments, so only prefix-stable - # fragments stream; otherwise the full arguments arrive at close. - if open_tool is not None and open_tool["stable"] and ev.fragment: - open_tool["sent"] += ev.fragment + elif isinstance(ev, ToolCallStart): + open_tool = { + "index": tool_calls_sent, + "ordinal": ev.tool_index, + "sent": "", + "stable": ev.args_prefix_stable, + } yield emit( _chat_chunk( - req, uid, + req, uid, created, [{ "delta": {"tool_calls": [{ "index": open_tool["index"], - "function": {"arguments": ev.fragment}, + "id": _tool_call_id(ev.name, open_tool["index"]), + "type": "function", + "function": {"name": ev.name, "arguments": ""}, }]}, "index": 0, "finish_reason": None, }], ) ) - elif isinstance(ev, ToolCallsDelta): - for call in ev.calls: - if open_tool is not None and open_tool["ordinal"] == call.tool_index: - # Close of a ToolCallStart-opened call: send whatever of the - # final (authoritative) arguments wasn't streamed yet. - final = call.parameters or "" - remainder = ( - final[len(open_tool["sent"]):] - if final.startswith(open_tool["sent"]) - else final if not open_tool["sent"] else "" + elif isinstance(ev, ToolCallArgsDelta): + # Clients concatenate argument fragments, so only prefix-stable + # fragments stream; otherwise the full arguments arrive at close. + if open_tool is not None and open_tool["stable"] and ev.fragment: + open_tool["sent"] += ev.fragment + yield emit( + _chat_chunk( + req, uid, created, + [{ + "delta": {"tool_calls": [{ + "index": open_tool["index"], + "function": {"arguments": ev.fragment}, + }]}, + "index": 0, "finish_reason": None, + }], + ) ) - if remainder: + elif isinstance(ev, ToolCallsDelta): + for call in ev.calls: + if open_tool is not None and open_tool["ordinal"] == call.tool_index: + # Close of a ToolCallStart-opened call: send whatever of the + # final (authoritative) arguments wasn't streamed yet. + final = call.parameters or "" + remainder = ( + final[len(open_tool["sent"]):] + if final.startswith(open_tool["sent"]) + else final if not open_tool["sent"] else "" + ) + if remainder: + yield emit( + _chat_chunk( + req, uid, created, + [{ + "delta": {"tool_calls": [{ + "index": open_tool["index"], + "function": {"arguments": remainder}, + }]}, + "index": 0, "finish_reason": None, + }], + ) + ) + open_tool = None + tool_calls_sent += 1 + continue + # Standalone complete call (buffered fallback path). + for delta in _tool_call_deltas([call], start_index=tool_calls_sent): yield emit( _chat_chunk( - req, uid, - [{ - "delta": {"tool_calls": [{ - "index": open_tool["index"], - "function": {"arguments": remainder}, - }]}, - "index": 0, "finish_reason": None, - }], + req, uid, created, + [{"delta": {"tool_calls": [delta]}, "index": 0, "finish_reason": None}], ) ) - open_tool = None tool_calls_sent += 1 - continue - # Standalone complete call (buffered fallback path). - for delta in _tool_call_deltas([call], start_index=tool_calls_sent): - yield emit( - _chat_chunk( - req, uid, - [{"delta": {"tool_calls": [delta]}, "index": 0, "finish_reason": None}], - ) - ) - tool_calls_sent += 1 - elif isinstance(ev, GenDone): - timings = ev.timings - done = _chat_chunk(req, uid, [{"delta": {}, "index": 0, "finish_reason": ev.finish_reason}]) - if not include_usage and (t := timings_block(state, timings)) is not None: - done["timings"] = t - yield emit(done) + elif isinstance(ev, GenDone): + timings = ev.timings + done = _chat_chunk(req, uid, created, [{"delta": {}, "index": 0, "finish_reason": ev.finish_reason}]) + if not include_usage and (t := timings_block(state, timings)) is not None: + done["timings"] = t + yield emit(done) + finally: + await events.aclose() if include_usage: final: dict[str, Any] = { "id": f"chatcmpl-{uid}", "object": "chat.completion.chunk", - "created": int(time.time()), + "created": created, "model": req.model, "choices": [], "usage": _reported_usage(state, usage), @@ -687,11 +716,13 @@ def _tool_call_id(name: str | None, index: int) -> str: return f"call_{prefix}_{index}_{uuid.uuid4().hex[:8]}" -def _chat_chunk(req: ChatCompletionRequest, uid: int, choices: list[dict[str, Any]]) -> dict[str, Any]: +def _chat_chunk( + req: ChatCompletionRequest, uid: int, created: int, choices: list[dict[str, Any]], +) -> dict[str, Any]: return { "id": f"chatcmpl-{uid}", "object": "chat.completion.chunk", - "created": int(time.time()), + "created": created, "model": req.model, "choices": choices, } diff --git a/tests/server/test_streaming_model_matrix.py b/tests/server/test_streaming_model_matrix.py index d6c84e1ce..bedbe9240 100644 --- a/tests/server/test_streaming_model_matrix.py +++ b/tests/server/test_streaming_model_matrix.py @@ -671,6 +671,150 @@ async def run(): assert out[:-1] and all(ev is KEEPALIVE for ev in out[:-1]) +@pytest.mark.parametrize("include_usage", [False, True]) +def test_chat_keepalive_only_before_content(monkeypatch, include_usage): + from openai.types.chat import ChatCompletionChunk + from freetoken.server import openai_api as O + + interval = 0.02 + monkeypatch.setattr(O, "KEEPALIVE_INTERVAL_S", interval, raising=False) + now = 100 + monkeypatch.setattr(O.time, "time", lambda: now) + + async def run(): + nonlocal now + prefill = asyncio.Event() + decode = asyncio.Event() + state = FakeState([], tool_call_parser="llama3") + + async def wait_for_ack(uid): + await prefill.wait() + yield UserReply(uid=uid, incremental_output="hi", finished=False, + prompt_tokens_delta=2, completion_tokens_delta=1) + await decode.wait() + yield UserReply(uid=uid, incremental_output="!", finished=True, + finish_reason="length", completion_tokens_delta=1) + + monkeypatch.setattr(state, "wait_for_ack", wait_for_ack) + req = ChatCompletionRequest(model="m", messages=[], stream=True, + stream_options={"include_usage": include_usage}) + stream = stream_chat_completion_chunks(42, req, state, _spec(tools=None)) + try: + role = json.loads((await anext(stream)).decode().removeprefix("data: ")) + now = 101 + for _ in range(2): + frame = await asyncio.wait_for(anext(stream), timeout=1) + chunk = ChatCompletionChunk.model_validate_json(frame.decode().removeprefix("data: ")) + assert (chunk.id, chunk.model, chunk.created, chunk.object) == ( + role["id"], role["model"], role["created"], "chat.completion.chunk", + ) + assert json.loads(frame.decode().removeprefix("data: "))["choices"] == [ + {"index": 0, "delta": {}, "finish_reason": None}, + ] + assert chunk.usage is None + + now = 102 + prefill.set() + content = json.loads((await asyncio.wait_for(anext(stream), 1)).decode().removeprefix("data: ")) + assert content["choices"][0]["delta"] == {"content": "hi"} + pending = asyncio.create_task(anext(stream)) + try: + done, _ = await asyncio.wait({pending}, timeout=interval * 3) + assert not done, "heartbeat after the first content delta" + decode.set() + frames = [await asyncio.wait_for(pending, 1)] + frames.extend([frame async for frame in stream]) + finally: + pending.cancel() + await asyncio.gather(pending, return_exceptions=True) + + assert frames[-1] == b"data: [DONE]\n\n" + chunks = [json.loads(frame.decode().removeprefix("data: ")) for frame in frames[:-1]] + assert len(chunks) == (3 if include_usage else 2) + assert chunks[0]["choices"][0]["delta"] == {"content": "!"} + assert chunks[1]["choices"] == [{"delta": {}, "index": 0, "finish_reason": "length"}] + if include_usage: + assert chunks[2]["choices"] == [] + assert chunks[2]["usage"] == {"prompt_tokens": 2, "completion_tokens": 2, "total_tokens": 4} + assert all(chunk["created"] == role["created"] for chunk in [content, *chunks]) + finally: + await stream.aclose() + + asyncio.run(run()) + + +@pytest.mark.parametrize("asked", [True, False]) +def test_chat_keepalive_carries_running_usage_and_timings_when_requested(monkeypatch, asked): + """A keep-alive is one more chunk after the role chunk: with continuous_usage_stats it + carries the usage so far, and with timings_per_token the snapshot of the ack behind it, + so never a block during queue/prefill, where nothing is measured yet.""" + from freetoken.server import openai_api as O + + monkeypatch.setattr(O, "KEEPALIVE_INTERVAL_S", 0.02, raising=False) + + def snapshot(predicted_n, predicted_ms): + return dict(queue_ms=0.0, cache_n=0, prompt_n=5, prompt_ms=100.0, predicted_n=predicted_n, + predicted_ms=predicted_ms, draft_n=0, draft_n_accepted=0) + + async def run(): + prefill = asyncio.Event() + decode = asyncio.Event() + state = FakeState([], tool_call_parser="llama3") + state.config.enable_timings_report = True + + async def wait_for_ack(uid): + await prefill.wait() + # a committed token that reaches no delta yet: the stream stays silent after an ack + yield UserReply(uid=uid, incremental_output="", finished=False, prompt_tokens_delta=5, + completion_tokens_delta=1, timings=snapshot(1, 0.001) if asked else None) + await decode.wait() + yield UserReply(uid=uid, incremental_output="hi", finished=True, finish_reason="stop", + completion_tokens_delta=1, timings=snapshot(2, 40.0)) + + monkeypatch.setattr(state, "wait_for_ack", wait_for_ack) + req = ChatCompletionRequest( + model="m", messages=[], stream=True, timings_per_token=asked, + stream_options={"include_usage": True, "continuous_usage_stats": asked}, + ) + stream = stream_chat_completion_chunks(42, req, state, _spec(tools=None)) + + async def chunk(): + frame = await asyncio.wait_for(anext(stream), timeout=1) + return json.loads(frame.decode().removeprefix("data: ")) + + try: + role = await chunk() + assert "usage" not in role and "timings" not in role + queued = await chunk() # before the first ack + prefill.set() + held = await chunk() # after an ack that produced no delta + for beat in (queued, held): + assert beat["choices"] == [{"index": 0, "delta": {}, "finish_reason": None}] + if asked: + assert queued["usage"] == {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} + assert "timings" not in queued + assert held["usage"] == {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6} + assert held["timings"]["predicted_n"] == 1 + else: + assert all("usage" not in beat and "timings" not in beat for beat in (queued, held)) + + decode.set() + content = await chunk() + assert content["choices"][0]["delta"] == {"content": "hi"} + assert ("usage" in content, "timings" in content) == (asked, asked) + frames = [frame async for frame in stream] + assert frames[-1] == b"data: [DONE]\n\n" + final = json.loads(frames[-2].decode().removeprefix("data: ")) + # the keep-alives leave the terminal usage chunk as it was + assert final["choices"] == [] + assert final["usage"] == {"prompt_tokens": 5, "completion_tokens": 2, "total_tokens": 7} + assert final["timings"]["predicted_n"] == 2 + finally: + await stream.aclose() + + asyncio.run(run()) + + # --------------------------------------------------------------------------- # # Review regressions: empty args, same-chunk trailing text, non-stream parity # --------------------------------------------------------------------------- #