Skip to content

feat(gemini): add aio surface and files API (Fixes #315) - #933

Open
DanielTobi0 wants to merge 1 commit into
PostHog:mainfrom
DanielTobi0:feat/gemini-aio-and-files
Open

feat(gemini): add aio surface and files API (Fixes #315)#933
DanielTobi0 wants to merge 1 commit into
PostHog:mainfrom
DanielTobi0:feat/gemini-aio-and-files

Conversation

@DanielTobi0

Copy link
Copy Markdown

The Gemini adapter advertised itself as a drop-in replacement for genai.Client, but two surfaces of the real SDK were missing:

  • client.aio.models.* - async generation was only reachable by swapping the class out for AsyncClient, so SDK-shaped code hit an AttributeError on aio.
  • client.files.* - absent entirely, which breaks any multimodal flow that uploads a file before referencing it in contents.

Client now builds the provider client once and shares it across every surface, so models, aio.models and files no longer open separate connections. aio.models is a tracked AsyncModels inheriting the same PostHog defaults; files and aio.files pass straight through to the provider, since uploads are not generations and emit no events.

AsyncClient gains the matching pair: files resolves to the provider's async Files API, and aio aliases its already-async models.

💡 Motivation and Context

Fixes #315.

posthog.ai.gemini is documented as a drop-in replacement for genai.Client, but only client.models existed. Two things followed from that:

  • Async wasn't drop-in. Async generation landed in feat(llma): add Gemini async #375 as a separate AsyncClient class. That works, but it isn't the shape the Google SDK has — anyone copying from Google's docs writes await client.aio.models.generate_content(...) and gets an AttributeError, and switching to PostHog means editing every call site rather than one import.
  • Files were unreachable. There was no files attribute at all, so uploading a PDF/image/video and referencing it in contents — the standard multimodal flow — failed outright with no workaround short of holding a second, untracked genai.Client.

Since Models and AsyncModels each constructed their own genai.Client, simply bolting on a second tracked surface would have doubled the number of provider connections per PostHog client. _initialize_policy now takes an optional provider_client so all surfaces share one.

client.chats is still missing — the same class of gap, but outside what the issue asked for. Happy to add it here if you'd prefer it in one pass.

Docs: there are no in-repo docs covering this adapter, but the posthog.com LLM analytics Gemini page should gain a mention of client.aio and client.files. I haven't opened that PR.

💚 How did you test it?

Five tests added to posthog/test/ai/gemini/test_gemini_parity.py, against a mocked provider client:

Test Covers
test_every_surface_shares_one_provider_client genai.Client is constructed exactly once; models and aio.models hold the same instance (parametrized over Client and AsyncClient)
test_sync_client_exposes_the_provider_files_api files/aio.files are the provider's, and upload() delegates with its arguments intact
test_async_client_exposes_the_async_files_api AsyncClient.files resolves to aio.files, and aio.models aliases models
test_sync_client_aio_models_inherits_posthog_defaults distinct id, properties, privacy mode and groups reach aio.models
test_sync_client_aio_models_tracks_generations await client.aio.models.generate_content(...) awaits the provider and emits one $ai_generation with the right distinct id and model

Full CI-aligned run:

pytest posthog/test/ai/    677 passed, 15 skipped
ruff format --check .      289 files already formatted
ruff check .               All checks passed
mypy | mypy-baseline       Success: no issues found in 227 source files
make public_api_check      passes (snapshot regenerated in this PR)
python -W error -c "import posthog"   clean

Not tested: no live calls against the real Google GenAI API. Everything above runs against mocks, so this verifies wiring, event capture and the public shape — not real upload/download behaviour against Google's Files service.

📝 Checklist

  • I reviewed the submitted code.
  • I added tests to verify the changes.
  • I updated the docs if needed.
  • No breaking change or entry added to the changelog.

If releasing new changes

  • Ran sampo add to generate a changeset file

🤖 Agent context

Autonomy: Human-driven (agent-assisted)

the agent was pointed at #315, asked to read and explain it, then asked to implement and submit the fix.

Decisions worth flagging for review:

  • Passthrough over wrapping for files. A tracked wrapper was considered and rejected — uploads aren't generations, so there's no $ai_* event that fits them, and a plain passthrough covers upload/get/list/delete/download with nothing to keep in sync as Google adds methods.
  • AsyncClient.files resolves to the provider's aio.files, not files. AsyncClient is async end to end, so handing back a blocking Files API would have been a trap.
  • Shared provider client instead of a second one. The first cut had Client.aio build its own AsyncModels, which meant a second genai.Client per PostHog client. Threading an optional provider_client through _initialize_policy was the smaller change and is asserted by a test, so the count can't silently regress.
  • Eager aio, not a lazy property. Once the provider client is shared, constructing AsyncModels up front costs one Python object, so laziness bought nothing and added state.
  • chats left out to keep the diff scoped to what the issue asked for.

The Gemini adapter advertised itself as a drop-in replacement for
genai.Client, but two surfaces of the real SDK were missing:

- `client.aio.models.*` - async generation was only reachable by
  swapping the class out for `AsyncClient`, so SDK-shaped code hit an
  AttributeError on `aio`.
- `client.files.*` - absent entirely, which breaks any multimodal flow
  that uploads a file before referencing it in `contents`.

`Client` now builds the provider client once and shares it across every
surface, so `models`, `aio.models` and `files` no longer open separate
connections. `aio.models` is a tracked `AsyncModels` inheriting the same
PostHog defaults; `files` and `aio.files` pass straight through to the
provider, since uploads are not generations and emit no events.

`AsyncClient` gains the matching pair: `files` resolves to the provider's
async Files API, and `aio` aliases its already-async `models`.
@DanielTobi0
DanielTobi0 requested a review from a team as a code owner September 8, 2026 18:20
@marandaneto
marandaneto requested a review from a team September 9, 2026 08:35
@marandaneto

Copy link
Copy Markdown
Member

@Radu-Raicea you are assigned to #315, so I wonder if you've already cooked something up?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Support for Async in Google GenAI SDK - LLM Observability

2 participants