diff --git a/pycodeloop/providers/generic.py b/pycodeloop/providers/generic.py index 754e5cd..7bcafe9 100644 --- a/pycodeloop/providers/generic.py +++ b/pycodeloop/providers/generic.py @@ -121,6 +121,7 @@ class _ConnectionSnapshot: request_builder: RequestBuilder response_parser: ResponseParser supports_openai_sse: bool + include_usage_in_stream: bool class GenericProvider(Provider): @@ -196,6 +197,7 @@ def __init__( repetition_repeats: int = _REPETITION_REPEATS, context_window: int | None = None, supports_openai_sse: bool = True, + include_usage_in_stream: bool = True, **kwargs, ) -> None: super().__init__(model=model, api_key=api_key, **kwargs) @@ -211,6 +213,7 @@ def __init__( self.repetition_repeats = repetition_repeats self.context_window = context_window self._supports_openai_sse = supports_openai_sse + self._include_usage_in_stream = include_usage_in_stream self._config_path: Path | None = None self._lock = threading.Lock() @@ -258,6 +261,7 @@ def _build_from_json(cls, path: str | Path) -> GenericProvider: timeout=data.get("timeout", 60.0), context_window=data.get("context_window"), supports_openai_sse=response_shape != "anthropic", + include_usage_in_stream=data.get("include_usage_in_stream", True), ) def reload(self) -> None: @@ -281,6 +285,7 @@ def reload(self) -> None: self.timeout = fresh.timeout self.context_window = fresh.context_window self._supports_openai_sse = fresh._supports_openai_sse + self._include_usage_in_stream = fresh._include_usage_in_stream @staticmethod def _default_request( @@ -308,6 +313,7 @@ def _snapshot_locked(self) -> _ConnectionSnapshot: request_builder=self.request_builder, response_parser=self.response_parser, supports_openai_sse=self._supports_openai_sse, + include_usage_in_stream=self._include_usage_in_stream, ) def _headers(self, config: _ConnectionSnapshot) -> dict[str, str]: @@ -381,6 +387,12 @@ def _stream( config: _ConnectionSnapshot, ) -> ProviderResponse: body = {**body, "stream": True} + if config.include_usage_in_stream: + existing_stream_options = body.get("stream_options") or {} + body["stream_options"] = { + "include_usage": True, + **existing_stream_options, + } text = "" pending: dict[int, dict] = {} stop_reason: str | None = None diff --git a/tests/providers/test_generic.py b/tests/providers/test_generic.py index 903846b..2037e2d 100644 --- a/tests/providers/test_generic.py +++ b/tests/providers/test_generic.py @@ -338,6 +338,123 @@ def test_streaming_flags_a_connection_dropped_mid_response(self): self.assertEqual(result.text, "cut off mid") self.assertEqual(result.stop_reason, "connection_lost") + def test_streaming_requests_usage_and_captures_it_from_final_chunk(self): + """Regression: streaming previously sent `stream: True` with no + `stream_options.include_usage`, so OpenAI-compatible servers that + only report usage when asked (e.g. Ollama's /v1/chat/completions) + never sent a usage chunk and every streamed response reported + 0/0 tokens.""" + path = self._write_config( + {"url": "http://fake/v1/chat/completions", "model": "my-model"} + ) + provider = GenericProvider.from_json(path) + + chunks = [ + {"choices": [{"delta": {"content": "hi"}}]}, + {"choices": [{"delta": {}, "finish_reason": "stop"}]}, + { + "choices": [], + "usage": {"prompt_tokens": 12, "completion_tokens": 3}, + }, + ] + sse_body = ( + "".join(f"data: {json.dumps(c)}\n" for c in chunks) + + "data: [DONE]\n" + ).encode() + + captured_requests = [] + + def fake_urlopen(request, timeout=None): + captured_requests.append(json.loads(request.data)) + return _FakeResponse(sse_body) + + with mock.patch( + "pycodeloop.providers.generic.urllib.request.urlopen", + side_effect=fake_urlopen, + ): + result = provider.complete("sys", [], [], on_delta=lambda _: None) + + self.assertEqual( + captured_requests[0]["stream_options"], {"include_usage": True} + ) + self.assertEqual(result.usage.input_tokens, 12) + self.assertEqual(result.usage.output_tokens, 3) + + def test_streaming_merges_include_usage_into_callers_stream_options( + self, + ): + """A caller opting out via params.stream_options.include_usage + (e.g. a provider that rejects the field) must not be silently + overwritten, and sibling flags must survive the merge.""" + path = self._write_config( + { + "url": "http://fake/v1/chat/completions", + "model": "my-model", + "request": { + "params": { + "stream_options": { + "include_usage": False, + "include_intermediary_tokens": True, + } + } + }, + } + ) + provider = GenericProvider.from_json(path) + + sse_body = ( + b'data: {"choices": [{"delta": {}, "finish_reason": "stop"}]}\n' + b"data: [DONE]\n" + ) + + captured_requests = [] + + def fake_urlopen(request, timeout=None): + captured_requests.append(json.loads(request.data)) + return _FakeResponse(sse_body) + + with mock.patch( + "pycodeloop.providers.generic.urllib.request.urlopen", + side_effect=fake_urlopen, + ): + provider.complete("sys", [], [], on_delta=lambda _: None) + + self.assertEqual( + captured_requests[0]["stream_options"], + {"include_usage": False, "include_intermediary_tokens": True}, + ) + + def test_include_usage_in_stream_false_omits_stream_options(self): + """Strict OpenAI-compatible endpoints that 400 on unknown fields + can opt out entirely via config.""" + path = self._write_config( + { + "url": "http://fake/v1/chat/completions", + "model": "my-model", + "include_usage_in_stream": False, + } + ) + provider = GenericProvider.from_json(path) + + sse_body = ( + b'data: {"choices": [{"delta": {}, "finish_reason": "stop"}]}\n' + b"data: [DONE]\n" + ) + + captured_requests = [] + + def fake_urlopen(request, timeout=None): + captured_requests.append(json.loads(request.data)) + return _FakeResponse(sse_body) + + with mock.patch( + "pycodeloop.providers.generic.urllib.request.urlopen", + side_effect=fake_urlopen, + ): + provider.complete("sys", [], [], on_delta=lambda _: None) + + self.assertNotIn("stream_options", captured_requests[0]) + def test_streaming_cuts_a_looping_response_short(self): path = self._write_config( {"url": "http://fake/v1/chat/completions", "model": "my-model"}