feat(gemini): add aio surface and files API (Fixes #315) - #933
Open
DanielTobi0 wants to merge 1 commit into
Open
feat(gemini): add aio surface and files API (Fixes #315)#933DanielTobi0 wants to merge 1 commit into
DanielTobi0 wants to merge 1 commit into
Conversation
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`.
Member
|
@Radu-Raicea you are assigned to #315, so I wonder if you've already cooked something up? |
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
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.
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 forAsyncClient, so SDK-shaped code hit an AttributeError onaio.client.files.*- absent entirely, which breaks any multimodal flow that uploads a file before referencing it incontents.Clientnow builds the provider client once and shares it across every surface, somodels,aio.modelsandfilesno longer open separate connections.aio.modelsis a trackedAsyncModelsinheriting the same PostHog defaults;filesandaio.filespass straight through to the provider, since uploads are not generations and emit no events.AsyncClientgains the matching pair:filesresolves to the provider's async Files API, andaioaliases its already-asyncmodels.💡 Motivation and Context
Fixes #315.
posthog.ai.geminiis documented as a drop-in replacement forgenai.Client, but onlyclient.modelsexisted. Two things followed from that:AsyncClientclass. That works, but it isn't the shape the Google SDK has — anyone copying from Google's docs writesawait client.aio.models.generate_content(...)and gets anAttributeError, and switching to PostHog means editing every call site rather than one import.filesattribute at all, so uploading a PDF/image/video and referencing it incontents— the standard multimodal flow — failed outright with no workaround short of holding a second, untrackedgenai.Client.Since
ModelsandAsyncModelseach constructed their owngenai.Client, simply bolting on a second tracked surface would have doubled the number of provider connections per PostHog client._initialize_policynow takes an optionalprovider_clientso all surfaces share one.client.chatsis 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.aioandclient.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_every_surface_shares_one_provider_clientgenai.Clientis constructed exactly once;modelsandaio.modelshold the same instance (parametrized overClientandAsyncClient)test_sync_client_exposes_the_provider_files_apifiles/aio.filesare the provider's, andupload()delegates with its arguments intacttest_async_client_exposes_the_async_files_apiAsyncClient.filesresolves toaio.files, andaio.modelsaliasesmodelstest_sync_client_aio_models_inherits_posthog_defaultsaio.modelstest_sync_client_aio_models_tracks_generationsawait client.aio.models.generate_content(...)awaits the provider and emits one$ai_generationwith the right distinct id and modelFull CI-aligned run:
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
If releasing new changes
sampo addto 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:
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 coversupload/get/list/delete/downloadwith nothing to keep in sync as Google adds methods.AsyncClient.filesresolves to the provider'saio.files, notfiles.AsyncClientis async end to end, so handing back a blocking Files API would have been a trap.Client.aiobuild its ownAsyncModels, which meant a secondgenai.Clientper PostHog client. Threading an optionalprovider_clientthrough_initialize_policywas the smaller change and is asserted by a test, so the count can't silently regress.aio, not a lazy property. Once the provider client is shared, constructingAsyncModelsup front costs one Python object, so laziness bought nothing and added state.chatsleft out to keep the diff scoped to what the issue asked for.