Harden OpenAI chat completion streams
This commit is contained in:
wasimysaid
parent
ede6a2bcee
commit
bd1e5eabe5
3 changed files with 164 additions and 22 deletions
|
|
@ -4748,7 +4748,7 @@ async def openai_chat_completions(
|
|||
finally:
|
||||
_tracker.__exit__(None, None, None)
|
||||
|
||||
return StreamingResponse(
|
||||
return _SameTaskStreamingResponse(
|
||||
audio_input_stream(),
|
||||
media_type = "text/event-stream",
|
||||
headers = {
|
||||
|
|
@ -5239,7 +5239,7 @@ async def openai_chat_completions(
|
|||
pass
|
||||
_tracker.__exit__(None, None, None)
|
||||
|
||||
return StreamingResponse(
|
||||
return _SameTaskStreamingResponse(
|
||||
gguf_tool_stream(),
|
||||
media_type = "text/event-stream",
|
||||
headers = {
|
||||
|
|
@ -5393,7 +5393,7 @@ async def openai_chat_completions(
|
|||
finally:
|
||||
_tracker.__exit__(None, None, None)
|
||||
|
||||
return StreamingResponse(
|
||||
return _SameTaskStreamingResponse(
|
||||
gguf_stream_chunks(),
|
||||
media_type = "text/event-stream",
|
||||
headers = {
|
||||
|
|
@ -5842,7 +5842,7 @@ async def openai_chat_completions(
|
|||
_sf_tracker.__exit__(None, None, None)
|
||||
|
||||
if payload.stream:
|
||||
return StreamingResponse(
|
||||
return _SameTaskStreamingResponse(
|
||||
sf_tool_stream(),
|
||||
media_type = "text/event-stream",
|
||||
headers = {
|
||||
|
|
@ -6062,7 +6062,7 @@ async def openai_chat_completions(
|
|||
finally:
|
||||
_tracker.__exit__(None, None, None)
|
||||
|
||||
return StreamingResponse(
|
||||
return _SameTaskStreamingResponse(
|
||||
stream_chunks(),
|
||||
media_type = "text/event-stream",
|
||||
headers = {
|
||||
|
|
@ -9547,7 +9547,7 @@ async def _openai_passthrough_stream(
|
|||
except Exception:
|
||||
pass
|
||||
_tracker.__exit__(None, None, None)
|
||||
return StreamingResponse(
|
||||
return _SameTaskStreamingResponse(
|
||||
iter(()),
|
||||
media_type = "text/event-stream",
|
||||
headers = {
|
||||
|
|
@ -9598,6 +9598,26 @@ async def _openai_passthrough_stream(
|
|||
_await_disconnect_then_close(request, resp, cancel_event)
|
||||
)
|
||||
monitor_done = False
|
||||
saw_finish_reason = False
|
||||
saw_done = False
|
||||
last_chunk_id = completion_id
|
||||
last_chunk_model = model_name
|
||||
last_chunk_created = int(time.time())
|
||||
|
||||
def _synthetic_finish_line() -> str:
|
||||
chunk = ChatCompletionChunk(
|
||||
id = last_chunk_id,
|
||||
created = last_chunk_created,
|
||||
model = last_chunk_model,
|
||||
choices = [
|
||||
ChunkChoice(
|
||||
delta = ChoiceDelta(),
|
||||
finish_reason = "stop",
|
||||
)
|
||||
],
|
||||
)
|
||||
return f"data: {chunk.model_dump_json(exclude_none = True)}"
|
||||
|
||||
try:
|
||||
lines_iter = resp.aiter_lines()
|
||||
async for raw_line in _aiter_llama_stream_items(
|
||||
|
|
@ -9611,6 +9631,37 @@ async def _openai_passthrough_stream(
|
|||
continue
|
||||
if not raw_line.startswith("data: "):
|
||||
continue
|
||||
data_text = raw_line[6:].strip()
|
||||
if data_text == "[DONE]":
|
||||
saw_done = True
|
||||
if not saw_finish_reason and not cancel_event.is_set():
|
||||
finish_line = _synthetic_finish_line()
|
||||
_monitor_openai_sse_line(
|
||||
monitor_id,
|
||||
finish_line,
|
||||
llama_backend.context_length,
|
||||
)
|
||||
yield finish_line + "\n\n"
|
||||
saw_finish_reason = True
|
||||
yield raw_line + "\n\n"
|
||||
monitor_done = True
|
||||
break
|
||||
try:
|
||||
chunk_data = json.loads(data_text)
|
||||
except json.JSONDecodeError:
|
||||
chunk_data = None
|
||||
if isinstance(chunk_data, dict):
|
||||
if isinstance(chunk_data.get("id"), str):
|
||||
last_chunk_id = chunk_data["id"]
|
||||
if isinstance(chunk_data.get("model"), str):
|
||||
last_chunk_model = chunk_data["model"]
|
||||
if isinstance(chunk_data.get("created"), int):
|
||||
last_chunk_created = chunk_data["created"]
|
||||
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
|
||||
|
|
@ -9625,9 +9676,19 @@ async def _openai_passthrough_stream(
|
|||
# Relay verbatim to preserve llama-server's native id,
|
||||
# finish_reason, delta.tool_calls, and usage chunks.
|
||||
yield raw_line + "\n\n"
|
||||
if monitor_event == "done" or raw_line[6:].strip() == "[DONE]":
|
||||
if monitor_event == "done":
|
||||
monitor_done = True
|
||||
break
|
||||
if not saw_done and not saw_finish_reason and not cancel_event.is_set():
|
||||
finish_line = _synthetic_finish_line()
|
||||
_monitor_openai_sse_line(
|
||||
monitor_id,
|
||||
finish_line,
|
||||
llama_backend.context_length,
|
||||
)
|
||||
yield finish_line + "\n\n"
|
||||
yield "data: [DONE]\n\n"
|
||||
monitor_done = True
|
||||
if not monitor_done:
|
||||
api_monitor.finish(
|
||||
monitor_id,
|
||||
|
|
@ -9680,7 +9741,7 @@ async def _openai_passthrough_stream(
|
|||
pass
|
||||
_tracker.__exit__(None, None, None)
|
||||
|
||||
return StreamingResponse(
|
||||
return _SameTaskStreamingResponse(
|
||||
_stream(),
|
||||
media_type = "text/event-stream",
|
||||
headers = {
|
||||
|
|
|
|||
|
|
@ -48,6 +48,7 @@ from routes.inference import (
|
|||
_openai_passthrough_stream,
|
||||
_openai_stream_usage_chunk,
|
||||
_proxy_to_external_provider,
|
||||
_SameTaskStreamingResponse,
|
||||
_set_or_prepend_system_message,
|
||||
openai_completions,
|
||||
openai_embeddings,
|
||||
|
|
@ -2239,6 +2240,7 @@ class TestApiMonitorProviderAndCompletionStreams:
|
|||
"chatcmpl-test",
|
||||
monitor_id = monitor_id,
|
||||
)
|
||||
assert isinstance(response, _SameTaskStreamingResponse)
|
||||
iterator = response.body_iterator
|
||||
first = await anext(iterator)
|
||||
assert "hello" in first
|
||||
|
|
@ -2256,6 +2258,67 @@ class TestApiMonitorProviderAndCompletionStreams:
|
|||
|
||||
asyncio.run(_run())
|
||||
|
||||
def test_passthrough_stream_synthesizes_missing_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":{"content":"hello"}}]}'
|
||||
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,
|
||||
)
|
||||
chunks = [chunk async for chunk in response.body_iterator]
|
||||
body = "".join(chunks)
|
||||
|
||||
assert '"finish_reason":"stop"' in body.replace(" ", "")
|
||||
assert "data: [DONE]" in body
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
def test_passthrough_non_streaming_cancel_finalizes_monitor(self, monkeypatch):
|
||||
async def _run():
|
||||
import routes.inference as inf_mod
|
||||
|
|
|
|||
|
|
@ -146,24 +146,42 @@ def test_async_generators_cleanup_tracker_in_finally():
|
|||
)
|
||||
|
||||
|
||||
def test_streaming_responses_have_no_background_task():
|
||||
top = None
|
||||
for n in ast.walk(_TREE):
|
||||
if isinstance(n, ast.AsyncFunctionDef) and n.name == "openai_chat_completions":
|
||||
top = n
|
||||
break
|
||||
assert top is not None
|
||||
def test_chat_completions_streams_avoid_starlette_task_group():
|
||||
top = _async_function("openai_chat_completions")
|
||||
legacy_calls = []
|
||||
same_task_calls = 0
|
||||
for sub in ast.walk(top):
|
||||
if not (isinstance(sub, ast.Call) and isinstance(sub.func, ast.Name)):
|
||||
continue
|
||||
if sub.func.id != "StreamingResponse":
|
||||
if sub.func.id == "StreamingResponse":
|
||||
legacy_calls.append(sub.lineno)
|
||||
if sub.func.id == "_SameTaskStreamingResponse":
|
||||
same_task_calls += 1
|
||||
assert not legacy_calls, (
|
||||
"Streaming /v1/chat/completions must use _SameTaskStreamingResponse, "
|
||||
"not Starlette's legacy task-group StreamingResponse. Lines: "
|
||||
f"{legacy_calls}"
|
||||
)
|
||||
assert same_task_calls >= 5
|
||||
|
||||
|
||||
def test_openai_passthrough_stream_avoids_starlette_task_group():
|
||||
top = _async_function("_openai_passthrough_stream")
|
||||
legacy_calls = []
|
||||
same_task_calls = 0
|
||||
for sub in ast.walk(top):
|
||||
if not (isinstance(sub, ast.Call) and isinstance(sub.func, ast.Name)):
|
||||
continue
|
||||
kwargs = {kw.arg for kw in sub.keywords if kw.arg}
|
||||
assert "background" not in kwargs, (
|
||||
"StreamingResponse in openai_chat_completions must not pass "
|
||||
"`background=` -- cleanup now lives in the generator's finally "
|
||||
"block; a BackgroundTask would be skipped on abrupt disconnect"
|
||||
)
|
||||
if sub.func.id == "StreamingResponse":
|
||||
legacy_calls.append(sub.lineno)
|
||||
if sub.func.id == "_SameTaskStreamingResponse":
|
||||
same_task_calls += 1
|
||||
assert not legacy_calls, (
|
||||
"OpenAI passthrough streams must use _SameTaskStreamingResponse, "
|
||||
"not Starlette's legacy task-group StreamingResponse. Lines: "
|
||||
f"{legacy_calls}"
|
||||
)
|
||||
assert same_task_calls >= 2
|
||||
|
||||
|
||||
def test_direct_llama_server_streams_install_disconnect_watcher():
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue