feat(google): instrument aio.models.generate_content_stream for async streaming

[Aftabbs]

Summary

Extends the Google GenAI instrumentation to cover client.aio.models.generate_content_stream — the async iterator variant of the streaming API. Before this PR, only the non-streaming client.models.generate_content path was instrumented; async streaming calls produced no spans.

Changes:

• generate_content.py: adds wrap_generate_content_stream_async using immutable_wrap_async_iterator, with pre_hook / yield_hook / error_hook / finally_hook that mirror the patterns already used for OpenAI and Anthropic async streaming. Also extracts _process_google_stream_chunk and _flush_stream_span helpers (shared logic between this and the analogous sync streaming PR feat(google): instrument generate_content_stream for streaming spans #760).
• wrapper.py: wrap_google_client now calls wrap_generate_content_stream_async in addition to wrap_generate_content_sync, so a single client wrap covers all three Google API paths.
• test_generate_content_stream_async.py: 8 unit tests covering span creation, llm span kind, multi-chunk text accumulation, token-usage recording from the final chunk, error status propagation, method replacement, and passthrough of unchanged chunk objects.

Relationship to PR #760

PR #760 adds wrap_generate_content_stream_sync for client.models.generate_content_stream (sync iterator). This PR targets the async counterpart client.aio.models.generate_content_stream and is independent of #760 at the code level (different API path). If both merge, the _process_google_stream_chunk helper will need a minor reconciliation — happy to rebase onto #760 once it lands if that's easier.

Test plan

pytest src/tests/instrumentation/llm/google/test_generate_content_stream_async.py -v

All 8 tests cover:

• Span is created with name GOOGLE_API_CALL
• Span kind attribute is llm
• Text chunks are accumulated across multiple yields
• Token usage is recorded from the final (usage-bearing) chunk
• Chunks with text=None are skipped in accumulation
• Stream errors set span status to ERROR
• wrap_generate_content_stream_async replaces the original method
• All yielded chunks pass through unchanged

---

[devin-ai-integration]

Devin Review found 1 potential issue.

View 2 additional findings in Devin Review.

[devin-ai-integration]

🔴 Async streaming instrumentation breaks real Google GenAI streaming calls

The async streaming method is wrapped as if it directly returns a stream of chunks (immutable_wrap_async_iterator at src/judgeval/instrumentation/llm/llm_google/generate_content.py:201-207), but the real Google method must first be awaited to obtain the stream, so instrumented streaming calls fail instead of working.

Impact: Any application using the real google-genai client for async streaming will crash once the client is wrapped, instead of getting traced streaming responses.

Coroutine-vs-async-iterator mismatch with the real google-genai API

In google-genai, client.aio.models.generate_content_stream is a coroutine function (async def generate_content_stream(...) at google/genai/models.py:8898 which does return stream_generator() at google/genai/models.py:9188). The documented usage is async for chunk in await client.aio.models.generate_content_stream(...) — the call returns a coroutine that must be awaited to yield the AsyncIterator.

immutable_wrap_async_iterator (src/judgeval/utils/wrappers/immutable_wrap_async_iterator.py:69-84) instead assumes func(*args, **kwargs) is directly an async iterable: it does async for value in func(*args, **kwargs). When func is the real coroutine function this raises TypeError (a coroutine is not async-iterable), which trips error_hook and re-raises.

Worse, the replacement installed via setattr(client.aio.models, "generate_content_stream", wrapped) at src/judgeval/instrumentation/llm/llm_google/generate_content.py:209 is itself an async generator function. Existing caller code async for chunk in await client.aio.models.generate_content_stream(...) now awaits an async-generator object, which raises TypeError: object async_generator can't be used in 'await' expression.

The correct pattern is the one used for OpenAI/Anthropic async streaming (src/judgeval/instrumentation/llm/llm_openai/chat_completions.py:381-501): wrap with mutable_wrap_async, await the coroutine to get the real AsyncIterator, then in a mutate_hook wrap that iterator with immutable_wrap_async_iterator.

The unit tests pass only because they replace generate_content_stream with make_async_stream (src/tests/instrumentation/llm/google/test_generate_content_stream_async.py:27-32), an async generator function that is directly async-iterable and is never awaited — which does not match the real API's coroutine-returning contract.

Prompt for agents

The real google-genai async API client.aio.models.generate_content_stream is a coroutine function (async def) that returns an AsyncIterator only after being awaited; documented usage is `async for chunk in await client.aio.models.generate_content_stream(...)`. The current code in wrap_generate_content_stream_async wraps this coroutine function directly with immutable_wrap_async_iterator, which (a) internally does `async for value in func(...)` on an un-awaited coroutine and (b) installs an async-generator function that callers can no longer `await`. Both break real usage. Follow the existing OpenAI/Anthropic async streaming pattern (see src/judgeval/instrumentation/llm/llm_openai/chat_completions.py, function _wrap_streaming_async): use mutable_wrap_async so the coroutine is awaited to obtain the AsyncIterator, then in a mutate_hook wrap the resulting AsyncIterator with immutable_wrap_async_iterator (via an inner traced_generator). Move the span pre_hook/yield/error/finally logic accordingly, and keep setattr installing an async (awaitable) replacement so `await client.aio.models.generate_content_stream(...)` still returns an async iterator. Also update the tests to use a coroutine-returning stub (an async def that returns an async generator) to match the real API contract rather than a bare async generator function.

---

Was this helpful? React with 👍 or 👎 to provide feedback.