diff --git a/optillm/server.py b/optillm/server.py index c9697369..bfd4e28b 100644 --- a/optillm/server.py +++ b/optillm/server.py @@ -555,6 +555,45 @@ def execute_n_times(n: int, approaches, operation: str, system_prompt: str, init return responses[0], total_tokens return responses, total_tokens +def generate_streaming_completion(completion, model): + """Stream a full chat completion (dict, or list of dicts) as SSE chunks. + + Unlike generate_streaming_response, which only forwards message text, this + keeps tool calls and reasoning, so agents that stream (and call tools) get + the whole answer. Each choice goes out as one chunk. + """ + completions = completion if isinstance(completion, list) else [completion] + response_id = f"chatcmpl-{int(time.time()*1000)}" + created = int(time.time()) + usage = None + for comp in completions: + if not isinstance(comp, dict): + comp = {"choices": [{"message": {"content": str(comp)}}]} + usage = comp.get("usage") or usage + for i, choice in enumerate(comp.get("choices") or []): + message = choice.get("message") or {} + delta = {"role": "assistant", "content": message.get("content") or ""} + for key in ("reasoning", "reasoning_content"): + if message.get(key): + delta[key] = message[key] + if message.get("tool_calls"): + delta["tool_calls"] = [dict(tc, index=j) for j, tc in enumerate(message["tool_calls"])] + chunk = { + "id": comp.get("id") or response_id, + "object": "chat.completion.chunk", + "created": comp.get("created") or created, + "model": comp.get("model") or model, + "choices": [{"index": choice.get("index", i), "delta": delta, + "finish_reason": choice.get("finish_reason") + or ("tool_calls" if delta.get("tool_calls") else "stop")}], + } + yield "data: " + json.dumps(chunk) + "\n\n" + if usage: + yield "data: " + json.dumps({"id": response_id, "object": "chat.completion.chunk", + "created": created, "model": model, "choices": [], + "usage": usage}) + "\n\n" + yield "data: [DONE]\n\n" + def generate_streaming_response(final_response, model): # Generate a unique response ID response_id = f"chatcmpl-{int(time.time()*1000)}" @@ -846,7 +885,7 @@ def proxy(): if stream: if request_id: logger.info(f'Request {request_id}: Completed (streaming response)') - return Response(generate_streaming_response(extract_contents(result), model), content_type='text/event-stream') + return Response(generate_streaming_completion(result, model), content_type='text/event-stream') else : if request_id: logger.info(f'Request {request_id}: Completed') @@ -869,7 +908,7 @@ def proxy(): if stream: if request_id: logger.info(f'Request {request_id}: Completed (streaming response)') - return Response(generate_streaming_response(extract_contents(response), model), content_type='text/event-stream') + return Response(generate_streaming_completion(response, model), content_type='text/event-stream') else: if request_id: logger.info(f'Request {request_id}: Completed') diff --git a/tests/test_streaming_completion.py b/tests/test_streaming_completion.py new file mode 100644 index 00000000..a75ffbad --- /dev/null +++ b/tests/test_streaming_completion.py @@ -0,0 +1,52 @@ +"""Streaming a completed response must keep tool calls and reasoning. + +The proxy and `none` paths get a full chat completion back and, when the client +asked for `stream: true`, re-emit it as SSE. They used to forward only the +message text, so a tool call (which has empty content) reached the client as a +bare `[DONE]`, and agents that stream saw an empty reply. +""" +import json + +from optillm.server import generate_streaming_completion + + +def _chunks(gen): + out = [] + for frame in gen: + assert frame.startswith("data: ") and frame.endswith("\n\n") + body = frame[len("data: "):-2] + out.append(body if body == "[DONE]" else json.loads(body)) + return out + + +def test_tool_calls_survive_streaming(): + completion = { + "id": "abc", "model": "m", "created": 1, + "choices": [{"index": 0, "finish_reason": "tool_calls", "message": { + "role": "assistant", "content": "", + "tool_calls": [{"id": "call_1", "type": "function", + "function": {"name": "get_weather", "arguments": "{\"city\":\"Paris\"}"}}]}}], + "usage": {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8}, + } + chunks = _chunks(generate_streaming_completion(completion, "optillm")) + assert chunks[-1] == "[DONE]" + first = chunks[0]["choices"][0] + assert first["finish_reason"] == "tool_calls" + call = first["delta"]["tool_calls"][0] + assert call["index"] == 0 and call["function"]["name"] == "get_weather" + assert chunks[1]["usage"]["total_tokens"] == 8 + + +def test_text_and_reasoning_survive_streaming(): + completion = {"choices": [{"message": {"content": "hi", "reasoning": "thinking"}, + "finish_reason": "stop"}]} + first = _chunks(generate_streaming_completion(completion, "optillm"))[0]["choices"][0] + assert first["delta"]["content"] == "hi" + assert first["delta"]["reasoning"] == "thinking" + + +def test_a_list_of_completions_streams_every_choice(): + completions = [{"choices": [{"message": {"content": "a"}}]}, + {"choices": [{"message": {"content": "b"}}]}] + chunks = _chunks(generate_streaming_completion(completions, "optillm")) + assert [c["choices"][0]["delta"]["content"] for c in chunks[:-1]] == ["a", "b"]