Add native Anthropic, OpenAI, and Gemini SDK support#3
Merged
Conversation
- src/lago_agent_sdk/adapters/anthropic_native.py — extract_anthropic_native - src/lago_agent_sdk/wrappers/anthropic.py — wraps messages.create (sync + async, streaming and non-streaming) and messages.stream context manager - Wired into sdk.wrap() dispatch and adapters/__init__.py exports - anthropic = ["anthropic>=0.30"] optional-dep group - 19 new unit tests + 3 live integration tests; 256 unit tests pass - Coverage 80.71% — gate maintained - 9 captured response fixtures from real Anthropic API - README + CHANGELOG updated
The test set max_batch_size == max_buffer_size == 100, which caused the push that brings the buffer to 100 to trigger a wake on the background worker. The worker would take a batch (emptying the buffer) and then race with the remaining 150 pushes to call slow_sender. On CI's slower runners the worker sometimes squeezed in additional batches before slow_sender finally blocked, leaving the buffer with fewer items than the expected sliding window. Setting max_batch_size > max_buffer_size guarantees push() never sets the wake event (buffer can never reach max_batch_size). Combined with a long flush_interval the worker only runs once shutdown() releases the pause in the finally block — fully deterministic. Verified with 5 consecutive runs.
Adapter handles both API shapes with auto-detection:
Chat Completions (client.chat.completions.create):
usage.prompt_tokens -> input
usage.completion_tokens -> output
usage.prompt_tokens_details.cached_tokens -> cache_read
usage.prompt_tokens_details.audio_tokens -> audio_input
usage.completion_tokens_details.reasoning_tokens -> reasoning (o-series)
usage.completion_tokens_details.audio_tokens -> audio_output
count of choices[0].message.tool_calls -> tool_calls
Responses API (client.responses.create):
usage.input_tokens -> input
usage.output_tokens -> output
usage.input_tokens_details.cached_tokens -> cache_read
usage.output_tokens_details.reasoning_tokens -> reasoning
count of output[].type == "function_call" -> tool_calls
Wrapper covers both methods, sync + streaming, on both OpenAI and AsyncOpenAI.
For Chat Completions streaming, auto-injects stream_options.include_usage=true
when missing so the final chunk carries usage data (without that flag, OpenAI
emits no usage on streamed responses).
CanonicalUsage extended with audio_output (mapped to llm_audio_output_tokens)
to capture GPT-4o-audio output usage.
OpenAI is the first provider to actually populate llm_reasoning_tokens
(o-series surfaces reasoning tokens separately; Anthropic/Bedrock fold them
into output_tokens).
Predicted Outputs tokens (accepted/rejected_prediction_tokens) are
intentionally not surfaced -- documented in the adapter docstring as a
v1 gap.
27 new unit tests (18 adapter + 9 wrapper). 5 live integration tests gated
on OPENAI_API_KEY. 10 captured response fixtures from the real OpenAI API.
Total: 283 unit tests passing, ruff + mypy strict clean.
Adapter maps usage_metadata fields to CanonicalUsage: prompt_token_count -> input candidates_token_count -> output cached_content_token_count -> cache_read thoughts_token_count -> reasoning prompt_tokens_details[modality=AUDIO].token_count -> audio_input prompt_tokens_details[modality=IMAGE].token_count -> image_input candidates_tokens_details[modality=AUDIO].token_count -> audio_output count of candidates[0].content.parts[].function_call -> tool_calls Wrapper covers client.models.generate_content + generate_content_stream (sync) and the async variants under client.aio.models. Idempotent via _lago_instrumented sentinel. Detector now returns 'gemini' (was 'google') for google-genai clients -- matches the naming convention used by other providers (bedrock, anthropic, openai, mistral). Semantic note vs OpenAI: Gemini's `thoughts_token_count` is ADDITIVE to `candidates_token_count` (verified by math across all 5 fixtures: input + output + reasoning = total). OpenAI's `reasoning_tokens` is a SUBSET of `completion_tokens`. Documented in adapter docstring + README for customers configuring per-metric billing. Gemini 2.5 emits reasoning tokens by default (no explicit thinking_config needed) -- second provider populating llm_reasoning_tokens. 21 new unit tests (15 adapter + 6 wrapper). 4 live integration tests gated on GEMINI_API_KEY. 5 captured response fixtures (plain, tool use, streaming, thinking, multi-turn). Total: 304 unit tests passing, ruff + mypy strict clean.
CI runs `ruff format --check` which was failing because earlier dev only ran `ruff check` (linter) locally, not the formatter. Auto-formatting restores whitespace consistency in: - src/lago_agent_sdk/adapters/gemini_native.py - src/lago_agent_sdk/wrappers/openai.py - tests/unit/adapters/fixtures/capture_openai.py - tests/unit/adapters/test_gemini_native.py - tests/unit/test_wrapper_gemini.py No functional changes.
vladmarascu
self-requested a review
June 1, 2026 12:09
vladmarascu
reviewed
Jun 1, 2026
vladmarascu
reviewed
Jun 1, 2026
Contributor
|
Good PR @anassg-lago 🙌 Left a few comms with some important bits and also these 2 more generals nits here:
|
vladmarascu
requested changes
Jun 1, 2026
Three bugs in Python async paths and the matching coverage gaps:
OpenAI wrapper
- Responses API + stream=True was injecting stream_options.include_usage,
a parameter the Responses API does not accept. It would raise TypeError
or HTTP 400 on the customer's call. Factory now takes is_responses_api
and only injects on the chat-completions path.
- Streaming usage extraction missed Responses API events: their terminal
response.completed event nests usage under event.response.usage rather
than at the top level. New _extract_stream_usage helper handles both
shapes.
Anthropic wrapper
- _LagoStreamManager.__aexit__ called the sync _emit_final, which in turn
called get_final_message() without await. On AsyncMessageStream the
method is a coroutine; the un-awaited coroutine fell through to {}
and nothing got billed. Added _emit_final_async; __aexit__ awaits it.
Mistral wrapper
- chat.stream_async in mistralai v2 is async def, returning a coroutine
that must be awaited before iterating. The wrapper was iterating the
coroutine directly, raising TypeError on first iteration. One-line
`await` added.
Plus new test coverage for the async paths that previously had none:
- 5 async tests on OpenAI wrapper (chat + responses, with stream variants)
- 3 async tests on Anthropic wrapper (messages.create + messages.stream
context manager)
- 3 async tests on Mistral wrapper (complete_async + stream_async)
- 2 async tests on Gemini wrapper (client.aio.models surface)
Coverage climbs from 80.16% to 89.28%. Total Python unit tests 304 → 318.
Three new live tests, skipped unless their matching API key is set, that
exercise the real provider API end-to-end against a mock Lago server:
- test_live_openai_responses_create_with_stream_emits_to_lago
Hits the Responses API with stream=True. Before the fix the call
crashed because stream_options was injected; after the fix the call
succeeds and usage is extracted from event.response.usage on the
terminal response.completed event.
- test_live_async_anthropic_messages_stream_context_manager_emits
Uses AsyncAnthropic + `async with client.messages.stream(...)`.
Before the fix __aexit__ ran the sync _emit_final, which called
get_final_message() without await; the un-awaited coroutine fell
through as {} and no event was emitted.
- test_live_mistral_chat_stream_async_emits_to_lago
Uses `stream = await client.chat.stream_async(...)`. Before the fix
the wrapper's _stream_async was sync def, so customers got an async
generator back directly instead of the awaitable shape the real
mistralai SDK provides. Fix: make _stream_async async def, preserving
the SDK's API contract.
Also corrected the unit test for stream_async to use the natural
`await client.chat.stream_async(...)` pattern.
All three live tests verified passing against the real OpenAI, Anthropic
and Mistral APIs.
Publish:
.github/workflows/publish.yml triggered on `v*.*.*` tag push. Four-stage
pipeline:
1. test — full CI gate (ruff, mypy, pytest, cov ≥ 80%) on 3.10/3.11/3.12
2. build — verify tag matches pyproject version, build sdist + wheel
3. publish — upload to PyPI via OIDC trusted publishing (no token in repo)
4. release — open a GitHub Release with auto-generated notes + artifacts
Auth via OIDC means no long-lived PYPI_TOKEN secret to leak. One-time
PyPI-side setup documented in RELEASING.md.
Customer-facing release flow becomes one command:
git tag v0.X.Y && git push --tags
Dev-group cleanup:
Consolidated [project.optional-dependencies].dev and
[dependency-groups].dev into a single dev extra. Removed duplicates
(boto3/mistralai were in both bedrock/mistral extras AND the dev extra;
anthropic/openai/google-genai were in their own extras AND the
dependency-groups dev). One source of truth now:
`pip install -e '.[dev]'` gives you everything to develop and test.
User-facing per-provider extras (bedrock, mistral, anthropic, openai,
gemini) are unchanged.
vladmarascu
reviewed
Jun 2, 2026
| | audio_output | `llm_audio_output_tokens` | ✗ | ✗ | ✗ | ✓ (GPT-4o-audio) | ✓ (multimodal AUDIO) | | ||
| | image_input | `llm_image_input_tokens` | ✗ | ✗ | ✗ | ✗ (Phase 3) | ✓ (multimodal IMAGE) | | ||
|
|
||
| **Semantic note on `reasoning`:** |
Contributor
There was a problem hiding this comment.
cache_read, audio_input, and image_input are all subsets of input for both OpenAI (prompt_tokens_details.*) and Gemini (prompt_tokens_details[modality=…]) , same trap as reasoning. A customer metering llm_input_tokens + llm_audio_input_tokens separately and summing them will double-count. Can we add a note here that these are subsets of input, not additive? Thxxx
vladmarascu
reviewed
Jun 2, 2026
| return "mistral" | ||
| if "google" in module and ("genai" in module or "generativeai" in module): | ||
| return "google" | ||
| return "gemini" |
Contributor
There was a problem hiding this comment.
This also matches the legacy google-generativeai SDK, but wrap_gemini_client only handles the new google-genai .models/.aio surface , a legacy client routes here and silently wraps nothing. Either narrow this to genai only, or handle/reject generativeai explicitly.
vladmarascu
reviewed
Jun 2, 2026
| for event in src: | ||
| payload = event.model_dump() if hasattr(event, "model_dump") else event | ||
| if isinstance(payload, dict): | ||
| usage = payload.get("usage") |
Contributor
There was a problem hiding this comment.
Only top-level usage is read, which is present on message_delta. The authoritative input/cache counts arrive on message_start (nested under message.usage) and are ignored --> billing relies on message_delta echoing input_tokens. If the SDK ever drops that, input silently bills 0. Consider merging message_start (input/cache) with message_delta (final output).
Contributor
|
Having a look at the release pipeline and have some suggestions for security:
permissions:
I don't have a lot of exp with release workflows so feel free to take this with a grain of salt or ask for a second pair of eyes! |
vladmarascu
self-requested a review
June 2, 2026 15:06
Bugs
- Anthropic messages.create(stream=True) under-billed input tokens. The
stream wrapper read only the top-level `usage`, which on a basic stream
appears only on message_delta as {output_tokens: N}. The authoritative
input/cache counts arrive nested under message.usage on message_start and
were dropped, so input billed 0. New _merge_stream_usage folds both
locations (message_start input/cache + message_delta cumulative output)
across sync and async paths. Fixtures now use the realistic wire shape
(message_delta carries no input echo), so the stream tests are genuine
regressions.
- Legacy google-generativeai SDK silently emitted nothing. The detector
matched both google-genai and the deprecated google-generativeai, but the
wrapper only instruments the unified Client.models/.aio surface, so a
legacy GenerativeModel wrapped nothing. Detector now returns a distinct
'gemini_legacy' kind and wrap() rejects it with a migrate-to-google-genai
message. ("genai" is not a substring of "generativeai", so no overlap.)
Docs
- README: cache_read / audio_input / image_input are subsets of input for
OpenAI and Gemini, not additive — summing them double-counts.
Publish-workflow hardening
- Least-privilege default (permissions: contents: read); only publish gets
id-token: write, only release gets contents: write.
- All third-party actions pinned to full commit SHAs (version in comment).
- Added `if: startsWith(github.ref, 'refs/tags/v')` to the publish job as
defense-in-depth.
- Added .github/dependabot.yml (github-actions) to keep SHA pins fresh.
- RELEASING.md documents pypi environment protection (required reviewers +
protected-tag restriction) as a REQUIRED setup step.
Gate: ruff + format + mypy clean; 319 unit tests pass; coverage 89.27%.
vladmarascu
approved these changes
Jun 3, 2026
vladmarascu
left a comment
vladmarascu
left a comment
Contributor
There was a problem hiding this comment.
lgtm! ty for addressing 🙌
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds three new native LLM provider integrations to the SDK, bringing total coverage to five (Bedrock + Mistral + Anthropic + OpenAI + Gemini). Plus a canonical-schema extension for audio-output tokens and a flaky-test fix.
Commits in this PR
c23e692test_repeated_overflow_keeps_window_sliding(race condition withmax_batch_size == max_buffer_size)8da7b82messages.createsync + stream,messages.streamcontext manager)0f79c5a6c487abgoogle-genai, sync + stream + async viaclient.aio.models)Provider coverage matrix
Key design decisions
audio_output→llm_audio_output_tokensmetric code. Populated by GPT-4o-audio output and Gemini TTS responses.gemini(wasgoogle) forgoogle-genaiclients — keeps naming consistent withbedrock/anthropic/openai/mistral.stream_options.include_usage=Truewhen missing. Without it, OpenAI emits no usage on streamed responses — silent under-billing.prompt_tokensvsinput_tokens).reasoning_tokensis a SUBSET ofcompletion_tokens; Gemini'sthoughtsTokenCountis ADDITIVE tocandidatesTokenCount. Customers configuring per-metric billing must account for this. Documented in adapter docstrings + README.Tests
Quality gate
ruff check src tests— cleanmypy --strict src— 0 issues in 21 source filespytest tests/unit— 304/304 passKnown gaps (intentional, documented)
accepted_/rejected_prediction_tokens) — not surfaced; would risk double-counting againstcompletion_tokensTest plan
audio_output)kind == "google"consumers (none today)pytest tests/unit -q)