refactor(streaming): extract chunk_creator dispatch so basedpyright can analyze it

[mateo-berri]

Relevant issues

Part of an effort to make basedpyright actually analyze the codebase's hottest functions. Across litellm/, basedpyright bails out with "Code is too complex to analyze" on exactly three functions, leaving their entire bodies unchecked: completion (litellm/main.py, ~3,500 lines), exception_type (exception_mapping_utils.py, ~2,275 lines), and CustomStreamWrapper.chunk_creator (this PR). This is the first of three, scoped on its own.

Linear ticket

N/A

Pre-Submission checklist

• I have added meaningful tests: a new block in tests/test_litellm/litellm_core_utils/test_streaming_handler.py drives the extracted _dispatch_provider_chunk directly across the legacy provider branches (vllm, petals, palm, cached_response, legacy vertex_ai, registered custom providers and text-completion-codestral) and asserts the tagged-union contract plus the content, finish_reason and fake-stream slicing side effects, on top of the existing mapped suite that already exercises chunk_creator across providers
• My PR passes all unit tests on make test-unit
• My PR's scope is as isolated as possible; it only touches chunk_creator and its new helper in one file
• I have requested a Greptile review and received a Confidence Score of at least 4/5 (5/5)

What and why

CustomStreamWrapper.chunk_creator packed a ~20-branch provider if/elif dispatch plus heavy post-processing into a single body, which pushed it past basedpyright's code-flow complexity ceiling. basedpyright responded with "Code is too complex to analyze" and skipped the whole function, so every type error in one of the hottest streaming paths was invisible and unguarded.

This extracts the provider dispatch into _dispatch_provider_chunk, which returns a tagged union (_ProviderChunkParsed | _ProviderChunkEarlyReturn) so the original early-return and StopIteration semantics are preserved exactly. chunk_creator now sits well under the ceiling and basedpyright type-checks both functions.

Restoring analysis surfaced pre-existing latent type noise in two legacy branches: dynamic proto attribute access in the vertex_ai path, and Optional subscripting of completion_stream in the petals/palm fake-streaming paths. Both are cast to Any at the point of dynamic access, matching the intent of the dead # type: ignore comments they already carried, so no basedpyright budget ceiling moves and no unrelated drift is absorbed.

Screenshots / Proof of Fix

Before, on litellm_internal_staging, basedpyright refuses to analyze the function:

$ uv run basedpyright litellm/litellm_core_utils/streaming_handler.py 2>&1 | grep "too complex"
  litellm/litellm_core_utils/streaming_handler.py:1148:9 - error: Code is too complex to analyze; reduce complexity by refactoring into subroutines or reducing conditional code paths (reportGeneralTypeIssues)

After this PR, the bailout is gone and the per-rule budget gate passes on the unchanged budget (no ceilings moved):

$ uv run basedpyright litellm/litellm_core_utils/streaming_handler.py 2>&1 | grep -c "too complex"
0

$ make lint-basedpyright
OK: every rule is within its basedpyright ceiling (150422 errors total)

$ git diff --stat basedpyright-code-budget.json
(no changes)

Behavioral proof that streaming is unchanged, run against a live proxy hitting real provider APIs and spending real money. The refactored dispatch routes by provider, so this streams through three distinct branches (native Anthropic, Azure AI, OpenAI) and checks each one returns incremental delta.content, a terminal finish_reason, a usage chunk, then data: [DONE].

1. Start the proxy

python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log

2. Stream through the native Anthropic branch

curl -sN http://localhost:4000/v1/chat/completions \
  -H "Authorization: Bearer sk-1234" \
  -H "Content-Type: application/json" \
  -d '{"model":"anthropic-haiku-4-5","stream":true,"stream_options":{"include_usage":true},
       "messages":[{"role":"user","content":"Count from 1 to 8, one number per line."}]}'

data: {"id":"chatcmpl-eddf671e...","object":"chat.completion.chunk","model":"anthropic-haiku-4-5","choices":[{"index":0,"delta":{"role":"assistant","content":"1\n2\n3\n4"}}]}
data: {"id":"chatcmpl-eddf671e...","object":"chat.completion.chunk","model":"anthropic-haiku-4-5","choices":[{"index":0,"delta":{"content":"\n5\n6\n7\n8"}}]}
data: {"id":"chatcmpl-eddf671e...","object":"chat.completion.chunk","model":"anthropic-haiku-4-5","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}
data: {"id":"chatcmpl-eddf671e...","object":"chat.completion.chunk","model":"anthropic-haiku-4-5","choices":[{"index":0,"delta":{}}],"usage":{"completion_tokens":19,"prompt_tokens":21,"total_tokens":40}}
data: [DONE]

3. Repeat through the Azure AI and OpenAI branches (same request shape, model set to azure-haiku-4-5 then gpt-5.5). Both return the same incremental-delta then finish_reason then usage then [DONE] structure.

4. The proxy computes a real, non-zero cost for each stream and writes spend to the DB, which is what proves real money was billed (from litellm.log):

model_group=anthropic-haiku-4-5  provider=anthropic  prompt=21 completion=19  spend=$0.000116
model_group=azure-haiku-4-5      provider=azure_ai   prompt=15 completion=7   spend=$0.00005
model_group=gpt-5.5              provider=openai     prompt=14 completion=18  spend=$0.00061

Type

🧹 Refactoring

Changes

Extracts the provider dispatch out of CustomStreamWrapper.chunk_creator into a new typed helper _dispatch_provider_chunk returning a tagged union, restoring basedpyright analysis of both functions, and casts two dynamic legacy branches to Any so the newly-visible errors do not move any budget ceiling. No behavior change.

A follow-up commit adds the test coverage described above and switches the seven Dict[str, Any] annotations the refactor introduced to the builtin dict[str, Any] generics, which keeps the UP006 strict-rule budget within its ceiling without moving any baseline

Note: two pre-existing tests in the mapped file (test_gemini_legacy_vertex_stop_finish_reason_normalised, test_gemini_legacy_vertex_tool_calls_finish_reason_with_stop_enum) fail in deterministic order on litellm_internal_staging as well; they depend on proto being present in sys.modules from an earlier test. They are unrelated to this change and left for a separate fix.

---

Note

Medium Risk
Touches core streaming dispatch for many LLM providers; behavior is intended to be identical but regressions would affect all streamed completions. Risk is mitigated by broad new unit tests and an explicit no-behavior-change refactor.

Overview
Pulls the large provider-specific if/elif chain out of CustomStreamWrapper.chunk_creator into _dispatch_provider_chunk, so basedpyright can type-check the hot streaming path instead of bailing with “too complex to analyze.”

The helper returns a small tagged union (_ProviderChunkParsed vs _ProviderChunkEarlyReturn) so existing behaviors stay the same: pass-through chunks for registered custom providers, None when only a finish reason arrives, and StopIteration after the stream has finished.

chunk_creator now only runs dispatch and keeps the shared post-processing (tool calls, return_processed_chunk_logic, etc.). cast(Any, …) on legacy vertex_ai and petals/palm dynamic access addresses type noise the refactor surfaced, without changing runtime logic.

Tests drive _dispatch_provider_chunk directly across vllm, fake-stream providers, cached_response, legacy vertex, codestral, triton, ai21, and custom-provider early-return cases.

^{Reviewed by Cursor Bugbot for commit f6958e1. Bugbot is set up for automated code reviews on this repo. Configure here.}

[codecov]

Codecov Report

❌ Patch coverage is 78.09524% with 46 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
litellm/litellm_core_utils/streaming_handler.py	78.09%	46 Missing ⚠️

📢 Thoughts on this report? Let us know!

[greptile-apps]

Greptile Summary

This PR extracts the large provider dispatch block from CustomStreamWrapper.chunk_creator into a new private helper _dispatch_provider_chunk that returns a typed _ProviderChunkParsed | _ProviderChunkEarlyReturn union, restoring basedpyright analysis on both methods without changing any runtime behavior.

• The tagged-union design cleanly preserves the original early-return and StopIteration semantics, and model_response mutations inside the helper remain visible to chunk_creator via Python's pass-by-reference semantics.
• The two cast(Any, ...) additions in the vertex_ai and petals/palm branches silence newly-visible basedpyright noise at the exact points of dynamic attribute access, matching the intent of the dead # type: ignore comments they replace.
• Comprehensive unit tests are added that directly exercise each legacy branch of _dispatch_provider_chunk using pure mock objects with no real network calls.

Confidence Score: 5/5

Safe to merge — this is a pure structural refactor with no behavior changes to the streaming hot path.

The extraction is mechanically equivalent to the original: StopIteration propagates through chunk_creator's existing except StopIteration: raise StopIteration clause, early returns are faithfully wrapped and unwrapped via the tagged union, and model_response mutations inside the helper remain visible to the caller via Python's reference semantics. The GChunk symbol was already a module-level import before this PR (referenced at module scope on line 63), so removing the redundant local import is safe. New tests cover every legacy branch of the extracted helper using pure mocks. The two cast(Any, ...) additions in the legacy vertex_ai and petals/palm branches are scoped to known dynamic-access points and do not widen any existing type guarantees.

No files require special attention.

Important Files Changed

Filename	Overview
litellm/litellm_core_utils/streaming_handler.py	Extracts provider dispatch into _dispatch_provider_chunk returning a typed tagged union; GChunk moved from a local import to the already-existing module-level import; StopIteration propagation and early-return semantics are preserved exactly.
tests/test_litellm/litellm_core_utils/test_streaming_handler.py	Adds 14 new unit tests driving _dispatch_provider_chunk directly across every legacy provider branch; all tests use mock objects with no real network calls, satisfying the no-network rule for this test directory.

_{Reviews (3): Last reviewed commit: "test(streaming): cover _dispatch_provide..." | Re-trigger Greptile}

[mateo-berri]

@greptileai

[mateo-berri]

@greptileai

[mateo-berri]

bugbot run

[cursor]

✅ Bugbot reviewed your changes and found no new issues!

Comment @cursor review or bugbot run to trigger another review on this PR

^{Reviewed by Cursor Bugbot for commit f6958e1. Configure here.}