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

.sampo/changesets/gemini-aio-and-files.md

@@ -0,0 +1,5 @@
+---
+pypi/posthog: minor
+---
+
+The Gemini adapter now covers two surfaces of `genai.Client` it previously lacked. `Client.aio.models` reaches the tracked async models adapter, so `await client.aio.models.generate_content(...)` works without swapping the class out for `AsyncClient`, and `client.files` (plus `client.aio.files`, and `AsyncClient.files` for the async Files API) passes through to the provider, so multimodal flows that upload a file before referencing it in `contents` no longer fail. Every surface of one client shares a single provider client instead of opening its own.

posthog/ai/gemini/_shared.py

@@ -78,6 +78,44 @@ def _build_gemini_client_args(
     return client_args
 
 
+def _build_gemini_client(
+    *,
+    api_key: Optional[str],
+    vertexai: Optional[bool],
+    credentials: Optional[Any],
+    project: Optional[str],
+    location: Optional[str],
+    debug_config: Optional[Any],
+    http_options: Optional[Any],
+) -> Any:
+    """Construct the provider client shared by every surface of one adapter."""
+    return genai.Client(
+        **_build_gemini_client_args(
+            api_key=api_key,
+            vertexai=vertexai,
+            credentials=credentials,
+            project=project,
+            location=location,
+            debug_config=debug_config,
+            http_options=http_options,
+        )
+    )
+
+
+class _GeminiAioNamespace:
+    """
+    Mirrors ``genai.Client().aio``: the async surface of a single client.
+
+    Pairs the tracked async ``models`` adapter with the provider's async
+    ``files`` API so ``client.aio.models`` and ``client.aio.files`` resolve the
+    same way they do on the real SDK.
+    """
+
+    def __init__(self, models: Any, files: Any):
+        self.models = models
+        self.files = files
+
+
 class _GeminiModelsPolicy:
     """Shared telemetry policy for the explicit sync and async Gemini adapters."""
 
@@ -98,23 +136,29 @@ def _initialize_policy(
         posthog_properties: Optional[Dict[str, Any]],
         posthog_privacy_mode: bool,
         posthog_groups: Optional[Dict[str, Any]],
+        provider_client: Optional[Any] = None,
     ) -> None:
         self._ph_client = _resolve_posthog_client(posthog_client)
         self._default_distinct_id = posthog_distinct_id
         self._default_properties = posthog_properties or {}
         self._default_privacy_mode = posthog_privacy_mode
         self._default_groups = posthog_groups
 
-        client_args = _build_gemini_client_args(
-            api_key=api_key,
-            vertexai=vertexai,
-            credentials=credentials,
-            project=project,
-            location=location,
-            debug_config=debug_config,
-            http_options=http_options,
+        # Every surface of one PostHog client (sync models, aio models, files)
+        # shares a single provider client rather than opening its own.
+        self._client = (
+            provider_client
+            if provider_client is not None
+            else _build_gemini_client(
+                api_key=api_key,
+                vertexai=vertexai,
+                credentials=credentials,
+                project=project,
+                location=location,
+                debug_config=debug_config,
+                http_options=http_options,
+            )
         )
-        self._client = genai.Client(**client_args)
         self._base_url = _GEMINI_BASE_URL
 
     def _merge_posthog_params(

posthog/ai/gemini/gemini.py

@@ -13,7 +13,13 @@
     merge_usage_stats,
     with_privacy_mode as with_privacy_mode,
 )
-from ._shared import _GeminiModelsPolicy, _resolve_posthog_client
+from ._shared import (
+    _GeminiAioNamespace,
+    _GeminiModelsPolicy,
+    _build_gemini_client,
+    _resolve_posthog_client,
+)
+from .gemini_async import AsyncModels
 from .gemini_converter import (
     extract_gemini_content_from_chunk,
     extract_gemini_embedding_token_count as extract_gemini_embedding_token_count,
@@ -39,6 +45,16 @@ class Client:
             contents=["Hello world"],
             posthog_distinct_id="specific_user"  # Override default
         )
+
+    The async surface lives under ``aio``, exactly as it does on genai.Client:
+
+        response = await client.aio.models.generate_content(
+            model="gemini-2.0-flash",
+            contents=["Hello world"],
+        )
+
+    ``files`` (and ``aio.files``) pass straight through to the provider. Uploads
+    are not generations, so they emit no PostHog events.
     """
 
     _ph_client: PostHogClient
@@ -78,21 +94,40 @@ def __init__(
 
         self._ph_client = _resolve_posthog_client(posthog_client)
 
-        self.models = Models(
+        # Built once here and shared by every surface, so a client that uses both
+        # `models` and `aio.models` still opens a single provider client.
+        self._provider_client = _build_gemini_client(
             api_key=api_key,
             vertexai=vertexai,
             credentials=credentials,
             project=project,
             location=location,
             debug_config=debug_config,
             http_options=http_options,
+        )
+
+        self.models = Models(
+            provider_client=self._provider_client,
             posthog_client=self._ph_client,
             posthog_distinct_id=posthog_distinct_id,
             posthog_properties=posthog_properties,
             posthog_privacy_mode=posthog_privacy_mode,
             posthog_groups=posthog_groups,
             **kwargs,
         )
+        self.files = self._provider_client.files
+        self.aio = _GeminiAioNamespace(
+            models=AsyncModels(
+                provider_client=self._provider_client,
+                posthog_client=self._ph_client,
+                posthog_distinct_id=posthog_distinct_id,
+                posthog_properties=posthog_properties,
+                posthog_privacy_mode=posthog_privacy_mode,
+                posthog_groups=posthog_groups,
+                **kwargs,
+            ),
+            files=self._provider_client.aio.files,
+        )
 
 
 class Models(_GeminiModelsPolicy):
@@ -116,6 +151,7 @@ def __init__(
         posthog_properties: Optional[Dict[str, Any]] = None,
         posthog_privacy_mode: bool = False,
         posthog_groups: Optional[Dict[str, Any]] = None,
+        provider_client: Optional[Any] = None,
         **kwargs,
     ):
         """
@@ -132,6 +168,8 @@ def __init__(
             posthog_properties: Default properties for all calls
             posthog_privacy_mode: Default privacy mode for all calls
             posthog_groups: Default groups for all calls
+            provider_client: An already-built genai.Client to reuse instead of
+                constructing one from the connection arguments above
             **kwargs: Additional arguments (for future compatibility)
         """
 
@@ -148,6 +186,7 @@ def __init__(
             posthog_properties=posthog_properties,
             posthog_privacy_mode=posthog_privacy_mode,
             posthog_groups=posthog_groups,
+            provider_client=provider_client,
         )
 
     def generate_content(

posthog/ai/gemini/gemini_async.py

@@ -14,7 +14,12 @@
     merge_usage_stats,
     with_privacy_mode as with_privacy_mode,
 )
-from ._shared import _GeminiModelsPolicy, _resolve_posthog_client
+from ._shared import (
+    _GeminiAioNamespace,
+    _GeminiModelsPolicy,
+    _build_gemini_client,
+    _resolve_posthog_client,
+)
 from .gemini_converter import (
     extract_gemini_content_from_chunk,
     extract_gemini_embedding_token_count as extract_gemini_embedding_token_count,
@@ -40,6 +45,11 @@ class AsyncClient:
             contents=["Hello world"],
             posthog_distinct_id="specific_user"  # Override default
         )
+
+    ``models`` is already the async surface, and ``aio.models`` is an alias for
+    it so code copied from the Google SDK docs keeps working. ``files`` (and
+    ``aio.files``) pass straight through to the provider's async Files API;
+    uploads are not generations, so they emit no PostHog events.
     """
 
     _ph_client: PostHogClient
@@ -79,21 +89,29 @@ def __init__(
 
         self._ph_client = _resolve_posthog_client(posthog_client)
 
-        self.models = AsyncModels(
+        # Built once here and shared by every surface, so `models`, `aio.models`
+        # and `files` all go through a single provider client.
+        self._provider_client = _build_gemini_client(
             api_key=api_key,
             vertexai=vertexai,
             credentials=credentials,
             project=project,
             location=location,
             debug_config=debug_config,
             http_options=http_options,
+        )
+
+        self.models = AsyncModels(
+            provider_client=self._provider_client,
             posthog_client=self._ph_client,
             posthog_distinct_id=posthog_distinct_id,
             posthog_properties=posthog_properties,
             posthog_privacy_mode=posthog_privacy_mode,
             posthog_groups=posthog_groups,
             **kwargs,
         )
+        self.files = self._provider_client.aio.files
+        self.aio = _GeminiAioNamespace(models=self.models, files=self.files)
 
 
 class AsyncModels(_GeminiModelsPolicy):
@@ -117,6 +135,7 @@ def __init__(
         posthog_properties: Optional[Dict[str, Any]] = None,
         posthog_privacy_mode: bool = False,
         posthog_groups: Optional[Dict[str, Any]] = None,
+        provider_client: Optional[Any] = None,
         **kwargs,
     ):
         """
@@ -133,6 +152,8 @@ def __init__(
             posthog_properties: Default properties for all calls
             posthog_privacy_mode: Default privacy mode for all calls
             posthog_groups: Default groups for all calls
+            provider_client: An already-built genai.Client to reuse instead of
+                constructing one from the connection arguments above
             **kwargs: Additional arguments (for future compatibility)
         """
 
@@ -149,6 +170,7 @@ def __init__(
             posthog_properties=posthog_properties,
             posthog_privacy_mode=posthog_privacy_mode,
             posthog_groups=posthog_groups,
+            provider_client=provider_client,
         )
 
     async def generate_content(

posthog/test/ai/gemini/test_gemini_parity.py

@@ -1,4 +1,4 @@
-from unittest.mock import MagicMock, patch
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 
@@ -65,3 +65,121 @@ def test_sync_and_async_clients_merge_posthog_defaults_without_mutation(client_c
         {"organization": "default-org"},
     )
     assert default_properties == {"shared": "default", "default-only": True}
+
+
+@pytest.fixture
+def provider_client_class():
+    with patch.object(google_genai, "Client") as patched:
+        patched.return_value = MagicMock()
+        yield patched
+
+
+@pytest.fixture
+def provider_client(provider_client_class):
+    """A stand-in for the underlying google-genai client, with both surfaces."""
+    return provider_client_class.return_value
+
+
+@pytest.fixture
+def gemini_response():
+    response = MagicMock()
+    response.text = "Test response from Gemini"
+
+    usage = MagicMock()
+    usage.prompt_token_count = 20
+    usage.candidates_token_count = 10
+    usage.cached_content_token_count = 0
+    usage.thoughts_token_count = 0
+    response.usage_metadata = usage
+
+    part = MagicMock()
+    part.text = "Test response from Gemini"
+    content = MagicMock()
+    content.parts = [part]
+    candidate = MagicMock()
+    candidate.content = content
+    response.candidates = [candidate]
+
+    return response
+
+
+@pytest.mark.parametrize("client_class", [Client, AsyncClient])
+def test_every_surface_shares_one_provider_client(
+    client_class, provider_client_class, provider_client
+):
+    """`models`, `aio.models` and `files` must not open separate connections."""
+    client = client_class(api_key="test-key", posthog_client=MagicMock())
+
+    provider_client_class.assert_called_once()
+    assert client.models._client is provider_client
+    assert client.aio.models._client is provider_client
+
+
+def test_sync_client_exposes_the_provider_files_api(provider_client):
+    client = Client(api_key="test-key", posthog_client=MagicMock())
+
+    assert client.files is provider_client.files
+    assert client.aio.files is provider_client.aio.files
+
+    uploaded = client.files.upload(file="notes.pdf")
+
+    provider_client.files.upload.assert_called_once_with(file="notes.pdf")
+    assert uploaded is provider_client.files.upload.return_value
+
+
+def test_async_client_exposes_the_async_files_api(provider_client):
+    """AsyncClient is async end to end, so its files API is the aio one."""
+    client = AsyncClient(api_key="test-key", posthog_client=MagicMock())
+
+    assert client.files is provider_client.aio.files
+    assert client.aio.files is provider_client.aio.files
+    # `models` is already async; `aio.models` is an alias so SDK-shaped code works.
+    assert client.aio.models is client.models
+
+
+def test_sync_client_aio_models_inherits_posthog_defaults(provider_client):
+    client = Client(
+        api_key="test-key",
+        posthog_client=MagicMock(),
+        posthog_distinct_id="default-id",
+        posthog_properties={"team": "ai"},
+        posthog_privacy_mode=True,
+        posthog_groups={"organization": "default-org"},
+    )
+
+    assert client.aio.models._merge_posthog_params(None, "trace", None, None, None) == (
+        "default-id",
+        "trace",
+        {"team": "ai"},
+        True,
+        {"organization": "default-org"},
+    )
+
+
+@pytest.mark.asyncio
+async def test_sync_client_aio_models_tracks_generations(
+    provider_client, gemini_response
+):
+    """`client.aio.models.generate_content` is the drop-in async entry point."""
+    provider_client.aio.models.generate_content = AsyncMock(
+        return_value=gemini_response
+    )
+    posthog_client = MagicMock()
+    posthog_client.privacy_mode = False
+
+    client = Client(api_key="test-key", posthog_client=posthog_client)
+
+    response = await client.aio.models.generate_content(
+        model="gemini-2.0-flash",
+        contents=["Tell me a fun fact about hedgehogs"],
+        posthog_distinct_id="test-id",
+    )
+
+    assert response is gemini_response
+    provider_client.aio.models.generate_content.assert_awaited_once()
+
+    assert posthog_client.capture.call_count == 1
+    call_args = posthog_client.capture.call_args[1]
+    assert call_args["distinct_id"] == "test-id"
+    assert call_args["event"] == "$ai_generation"
+    assert call_args["properties"]["$ai_model"] == "gemini-2.0-flash"