Stream tool calls and reasoning, not just message text

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')

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"]