diff --git a/examples/gateways/truefoundry/client_example.py b/examples/gateways/truefoundry/client_example.py index bf88e34a5..64314b4aa 100755 --- a/examples/gateways/truefoundry/client_example.py +++ b/examples/gateways/truefoundry/client_example.py @@ -80,17 +80,27 @@ def agent(user_prompt: str) -> str: def run_stream(client: OpenAI, model: str, prompt: str) -> str: - """One streamed completion. The LLM span ends when the stream is fully read.""" + """Return a complete answer, closing the stream on success or failure.""" parts = [] - stream = client.chat.completions.create( + finish_reason = None + with client.chat.completions.create( model=model, messages=[{"role": "user", "content": prompt}], stream=True, stream_options={"include_usage": True}, - ) - for chunk in stream: - if chunk.choices and chunk.choices[0].delta.content: - parts.append(chunk.choices[0].delta.content) + ) as stream: + for chunk in stream: + for choice in chunk.choices: + if choice.index != 0: + continue + if choice.delta.content: + parts.append(choice.delta.content) + if choice.finish_reason is not None: + finish_reason = choice.finish_reason + # EOF and [DONE] alone do not prove the model finished its answer. + # Token limits and content filtering also produce partial answers. + if finish_reason != "stop": + raise RuntimeError("The gateway did not return a complete answer") return "".join(parts) @@ -119,7 +129,7 @@ def main(argv: list) -> int: code = 0 try: print(run_stream(client, model, prompt) if stream else run_agent(client, tracer, model, prompt)) - except Exception as exc: # the failed gateway call is already recorded as an ERROR span + except Exception as exc: print(f"gateway call failed: {type(exc).__name__}", file=sys.stderr) code = 1 finally: diff --git a/examples/gateways/truefoundry/tests/test_stream_completion.py b/examples/gateways/truefoundry/tests/test_stream_completion.py new file mode 100644 index 000000000..55ca6803a --- /dev/null +++ b/examples/gateways/truefoundry/tests/test_stream_completion.py @@ -0,0 +1,85 @@ +"""The normal application must not print a partial answer as success.""" + +import json +from unittest.mock import Mock + +import pytest + +import client_example +from test_client_tracing import MODEL, httpx + + +@pytest.mark.parametrize( + "ending", ["eof", "done", "usage", "empty", "stop", "length", "content_filter"] +) +def test_stream_main_requires_complete_answer(monkeypatch, capsys, ending): + chunk = { + "id": "synthetic-stream", + "object": "chat.completion.chunk", + "created": 1, + "model": MODEL, + "choices": [ + {"index": 0, "delta": {"content": "answer text"}, "finish_reason": None} + ], + } + payload = "" if ending == "empty" else "data: " + json.dumps(chunk) + "\n\n" + if ending in {"stop", "length", "content_filter"}: + chunk["choices"] = [{"index": 0, "delta": {}, "finish_reason": ending}] + payload += "data: " + json.dumps(chunk) + "\n\n" + if ending in {"usage", "stop"}: + chunk["choices"] = [] + chunk["usage"] = {"prompt_tokens": 2, "completion_tokens": 2, "total_tokens": 4} + payload += "data: " + json.dumps(chunk) + "\n\n" + if ending in {"done", "stop"}: + payload += "data: [DONE]\n\n" + responses = [] + + def handler(request): + response = httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + stream=httpx.ByteStream(payload.encode()), + ) + responses.append(response) + return response + + provider = Mock() + with client_example.OpenAI( + api_key="synthetic", + max_retries=0, + http_client=httpx.Client(transport=httpx.MockTransport(handler)), + ) as client: + monkeypatch.setattr(client_example, "build", lambda: (provider, client, None)) + code = client_example.main(["client_example.py", "--stream", "question"]) + output = capsys.readouterr() + assert code == (0 if ending == "stop" else 1) + assert output.out.strip() == ("answer text" if ending == "stop" else "") + assert bool(output.err) == (ending != "stop") + assert responses[0].is_closed + provider.force_flush.assert_called_once_with() + provider.shutdown.assert_called_once_with() + + +@pytest.mark.parametrize( + "error", [RuntimeError("synthetic failure"), KeyboardInterrupt()] +) +def test_stream_cleanup_preserves_iteration_error(error): + class InterruptedStream: + closed = False + + def __enter__(self): + return self + + def __exit__(self, *args): + self.closed = True + + def __iter__(self): + raise error + + stream = InterruptedStream() + client = Mock() + client.chat.completions.create.return_value = stream + with pytest.raises(type(error)) as caught: + client_example.run_stream(client, MODEL, "question") + assert caught.value is error + assert stream.closed