ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Address OpenAI stream review issues

Browse source
This commit is contained in:
wasimysaid 2026-06-19 19:40:29 +02:00
commit b53015be85
4 changed files with 163 additions and 16 deletions
Download patch file Download diff file Expand all files Collapse all files

92
studio/backend/routes/inference.py
View file

@ -1303,6 +1303,16 @@ async def _await_disconnect_then_close(request, resp, cancel_event) -> None:
return
async def _await_disconnect_then_cancel(request, cancel_event) -> None:
"""Set ``cancel_event`` when a same-task local stream disconnects."""
try:
while not await request.is_disconnected():
await asyncio.sleep(0.1)
cancel_event.set()
except asyncio.CancelledError:
return
# Centralized local/server tool nudge. Keep render_html guidance gated to turns
# where the canvas tool is actually present in the tool schema; otherwise
# small local models can hallucinate a missing tool call instead of following
@ -5095,6 +5105,9 @@ async def openai_chat_completions(
async def gguf_tool_stream():
gen = None
disconnect_watcher = asyncio.create_task(
_await_disconnect_then_cancel(request, cancel_event)
)
try:
first_chunk = ChatCompletionChunk(
id = completion_id,
@ -5232,6 +5245,11 @@ async def openai_chat_completions(
error_chunk = _openai_stream_error_chunk(e)
yield f"data: {json.dumps(error_chunk)}\n\n"
finally:
disconnect_watcher.cancel()
try:
await disconnect_watcher
except (asyncio.CancelledError, Exception):
pass
if gen is not None:
try:
gen.close()
@ -5283,6 +5301,9 @@ async def openai_chat_completions(
_tracker.__enter__()
async def gguf_stream_chunks():
disconnect_watcher = asyncio.create_task(
_await_disconnect_then_cancel(request, cancel_event)
)
try:
# First chunk: role
first_chunk = ChatCompletionChunk(
@ -5391,6 +5412,11 @@ async def openai_chat_completions(
error_chunk = _openai_stream_error_chunk(e)
yield f"data: {json.dumps(error_chunk)}\n\n"
finally:
disconnect_watcher.cancel()
try:
await disconnect_watcher
except (asyncio.CancelledError, Exception):
pass
_tracker.__exit__(None, None, None)
return _SameTaskStreamingResponse(
@ -5710,6 +5736,9 @@ async def openai_chat_completions(
async def sf_tool_stream():
gen = None
disconnect_watcher = asyncio.create_task(
_await_disconnect_then_cancel(request, cancel_event)
)
try:
first_chunk = ChatCompletionChunk(
id = completion_id,
@ -5834,6 +5863,11 @@ async def openai_chat_completions(
}
yield f"data: {json.dumps(error_chunk)}\n\n"
finally:
disconnect_watcher.cancel()
try:
await disconnect_watcher
except (asyncio.CancelledError, Exception):
pass
if gen is not None:
try:
gen.close()
@ -5952,6 +5986,9 @@ async def openai_chat_completions(
_tracker.__enter__()
async def stream_chunks():
disconnect_watcher = asyncio.create_task(
_await_disconnect_then_cancel(request, cancel_event)
)
try:
first_chunk = ChatCompletionChunk(
id = completion_id,
@ -6060,6 +6097,11 @@ async def openai_chat_completions(
}
yield f"data: {json.dumps(error_chunk)}\n\n"
finally:
disconnect_watcher.cancel()
try:
await disconnect_watcher
except (asyncio.CancelledError, Exception):
pass
_tracker.__exit__(None, None, None)
return _SameTaskStreamingResponse(
@ -7115,8 +7157,6 @@ async def _responses_non_streaming(
# the model produced content, so clients expecting a pure tool-call turn
# (finish_reason="tool_calls") don't see a spurious empty message item.
output_items: list[dict] = []
if reasoning_text and not text and not tool_calls:
text = reasoning_text
if reasoning_text:
output_items.append(_responses_reasoning_output_item(reasoning_text))
if text:
@ -7421,6 +7461,7 @@ async def _responses_stream(
client = httpx.AsyncClient(timeout = _llama_streaming_generation_timeout())
resp = None
lines_iter = None
disconnect_watcher = None
disconnect_event = threading.Event()
try:
req = client.build_request(
@ -7482,6 +7523,9 @@ async def _responses_stream(
return
lines_iter = resp.aiter_lines()
disconnect_watcher = asyncio.create_task(
_await_disconnect_then_close(request, resp, disconnect_event)
)
async for raw_line in _aiter_llama_stream_items(
lines_iter,
cancel_event = disconnect_event,
@ -7643,6 +7687,12 @@ async def _responses_stream(
)
return
finally:
if disconnect_watcher is not None:
disconnect_watcher.cancel()
try:
await disconnect_watcher
except (asyncio.CancelledError, Exception):
pass
if lines_iter is not None:
try:
await lines_iter.aclose()
@ -9600,11 +9650,13 @@ async def _openai_passthrough_stream(
monitor_done = False
saw_finish_reason = False
saw_done = False
saw_tool_call_delta = False
last_chunk_id = completion_id
last_chunk_model = model_name
last_chunk_created = int(time.time())
def _synthetic_finish_line() -> str:
finish_reason = "tool_calls" if saw_tool_call_delta else "stop"
chunk = ChatCompletionChunk(
id = last_chunk_id,
created = last_chunk_created,
@ -9612,7 +9664,7 @@ async def _openai_passthrough_stream(
choices = [
ChunkChoice(
delta = ChoiceDelta(),
finish_reason = "stop",
finish_reason = finish_reason,
)
],
)
@ -9643,9 +9695,21 @@ async def _openai_passthrough_stream(
)
yield finish_line + "\n\n"
saw_finish_reason = True
_monitor_openai_sse_line(
monitor_id,
raw_line,
llama_backend.context_length,
)
yield raw_line + "\n\n"
monitor_done = True
break
# Honor parallel_tool_calls=false (best-effort): drop tool_call
# deltas with index>=1 so only the first call streams. Only
# lines carrying tool_calls are reparsed; everything else is
# relayed byte-for-byte.
if payload.parallel_tool_calls is False and '"tool_calls"' in raw_line:
raw_line = _cap_parallel_tool_calls_sse_line(raw_line)
data_text = raw_line[6:].strip()
try:
chunk_data = json.loads(data_text)
except json.JSONDecodeError:
@ -9660,14 +9724,12 @@ async def _openai_passthrough_stream(
choices = chunk_data.get("choices")
if isinstance(choices, list) and choices:
choice = choices[0]
if isinstance(choice, dict) and choice.get("finish_reason"):
saw_finish_reason = True
# Honor parallel_tool_calls=false (best-effort): drop tool_call
# deltas with index>=1 so only the first call streams. Only
# lines carrying tool_calls are reparsed; everything else is
# relayed byte-for-byte.
if payload.parallel_tool_calls is False and '"tool_calls"' in raw_line:
raw_line = _cap_parallel_tool_calls_sse_line(raw_line)
if isinstance(choice, dict):
if choice.get("finish_reason"):
saw_finish_reason = True
delta = choice.get("delta")
if isinstance(delta, dict) and delta.get("tool_calls"):
saw_tool_call_delta = True
monitor_event = _monitor_openai_sse_line(
monitor_id,
raw_line,
@ -9687,7 +9749,13 @@ async def _openai_passthrough_stream(
llama_backend.context_length,
)
yield finish_line + "\n\n"
yield "data: [DONE]\n\n"
done_line = "data: [DONE]"
_monitor_openai_sse_line(
monitor_id,
done_line,
llama_backend.context_length,
)
yield done_line + "\n\n"
monitor_done = True
if not monitor_done:
api_monitor.finish(

74
studio/backend/tests/test_openai_tool_passthrough.py
View file

@ -2316,6 +2316,75 @@ class TestApiMonitorProviderAndCompletionStreams:
assert '"finish_reason":"stop"' in body.replace(" ", "")
assert "data: [DONE]" in body
assert monitor.active_count() == 0
asyncio.run(_run())
def test_passthrough_stream_synthesizes_tool_call_finish_reason(self, monkeypatch):
async def _run():
import routes.inference as inf_mod
class Request:
async def is_disconnected(self):
return False
async def fake_send(*_args, **_kwargs):
return httpx.Response(200, content = b"")
async def fake_items(*_args, **_kwargs):
yield (
'data: {"id":"upstream","created":123,"model":"gguf",'
'"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,'
'"id":"call_1","type":"function","function":{"name":"lookup",'
'"arguments":"{}"}}]}}]}'
)
yield "data: [DONE]"
monitor = ApiMonitor(max_entries = 3)
monkeypatch.setattr(inf_mod, "api_monitor", monitor)
monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send)
monkeypatch.setattr(inf_mod, "_aiter_llama_stream_items", fake_items)
monitor_id = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "gguf",
prompt = "hi",
)
payload = ChatCompletionRequest(
model = "default",
messages = [ChatMessage(role = "user", content = "hi")],
stream = True,
tools = [
{
"type": "function",
"function": {
"name": "lookup",
"parameters": {"type": "object", "properties": {}},
},
}
],
)
response = await _openai_passthrough_stream(
Request(),
threading.Event(),
SimpleNamespace(
base_url = "http://llama.test",
context_length = 4096,
_request_reasoning_kwargs = lambda *_args, **_kwargs: None,
),
payload,
"gguf",
"chatcmpl-test",
monitor_id = monitor_id,
)
body = "".join([chunk async for chunk in response.body_iterator])
compact = body.replace(" ", "")
assert '"finish_reason":"tool_calls"' in compact
assert '"finish_reason":"stop"' not in compact
assert "data: [DONE]" in body
assert monitor.active_count() == 0
asyncio.run(_run())
@ -2434,7 +2503,10 @@ class TestApiMonitorProviderAndCompletionStreams:
async for chunk in response.body_iterator:
chunks.append(chunk)
assert chunks == ['data: {"choices":[{"delta":{"content":"hello"}}]}\n\n']
assert chunks[0] == 'data: {"choices":[{"delta":{"content":"hello"}}]}\n\n'
compact = "".join(chunks).replace(" ", "")
assert '"finish_reason":"stop"' in compact
assert chunks[-1] == "data: [DONE]\n\n"
[entry] = monitor.snapshot()
assert entry["status"] == "completed"
assert entry["reply"] == "hello"

5
studio/backend/tests/test_responses_tool_passthrough.py
View file

@ -992,7 +992,7 @@ class TestResponsesNonStreamingAdapter:
assert [item["type"] for item in body["output"]] == ["message"]
assert body["output"][0]["content"][0]["text"] == "33"
def test_reasoning_only_is_also_visible_message_text(self, monkeypatch):
def test_reasoning_only_stays_out_of_visible_message_text(self, monkeypatch):
payload = ResponsesRequest(input = "hi", reasoning = {"effort": "high"})
body = self._run_with_message(
monkeypatch,
@ -1000,9 +1000,8 @@ class TestResponsesNonStreamingAdapter:
payload = payload,
)
assert [item["type"] for item in body["output"]] == ["reasoning", "message"]
assert [item["type"] for item in body["output"]] == ["reasoning"]
assert body["output"][0]["content"][0]["text"] == "plan"
assert body["output"][1]["content"][0]["text"] == "plan"
# =====================================================================

8
tests/studio/test_stream_cancel_registration_timing.py
View file

@ -184,6 +184,14 @@ def test_openai_passthrough_stream_avoids_starlette_task_group():
assert same_task_calls >= 2
def test_local_chat_streams_install_same_task_disconnect_watcher():
top = _async_function("openai_chat_completions")
assert _calls_name(top, "_await_disconnect_then_cancel"), (
"Local same-task streams must watch request disconnects themselves; "
"do not restore Starlette's task-group StreamingResponse for this."
)
def test_direct_llama_server_streams_install_disconnect_watcher():
required = {
"openai_completions",