Add streaming cancellation support

Closed
OT
7mo ago
+3345-300

Summary

Add cancel() method to StreamedRunResult that allows users to stop streaming early and close the underlying HTTP connection. Cancelled responses are marked with incomplete=True.

Notes

  • Base StreamedResponse.cancel(): Sets the _cancelled flag but doesn't raise NotImplementedError. All 8 relevant providers (OpenAI, Anthropic, Google, Groq, Mistral, HuggingFace, Bedrock) implement proper stream closure. The base implementation still stops iteration via flag checks.

  • Usage tracking: It seems to me that different providers send usage data at different times; some with each chunk, some only at stream end. Cancelled streams retain whatever usage arrived before cancellation. This is documented in the new "Cancelling Streams" section.

  • group_by_temporal complexity: Added noqa: C901 due to the try/finally block needed for proper async generator cleanup during cancellation. Open to refactoring suggestions.

  • Type ignores for Google/HuggingFace: Both google.genai.types.AsyncGenerateContentResponse and huggingface_hub.AsyncInferenceClient.chat.completions.create are typed as AsyncIterator but are AsyncGenerator at runtime, hence the implementation uses aclose() with # type: ignore[attr-defined].

Happy to iterate and/or add more in-depth tests if deemed necessary.

Pre-Review Checklist

  • Any AI generated code has been reviewed line-by-line by the human PR author, who stands by it.
  • No breaking changes in accordance with the version policy.
  • Linting and type checking pass per make format and make typecheck.
  • PR title is fit for the release changelog.

Pre-Merge Checklist

  • New tests for any fix or new behavior, maintaining 100% coverage.
  • Updated documentation for new features and behaviors, including docstrings for API docs.