diff --git a/.sampo/changesets/gemini-aio-and-files.md b/.sampo/changesets/gemini-aio-and-files.md new file mode 100644 index 000000000..ba29c361d --- /dev/null +++ b/.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. diff --git a/posthog/ai/gemini/_shared.py b/posthog/ai/gemini/_shared.py index 95ebf019d..fc3635589 100644 --- a/posthog/ai/gemini/_shared.py +++ b/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,6 +136,7 @@ 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 @@ -105,16 +144,21 @@ def _initialize_policy( 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( diff --git a/posthog/ai/gemini/gemini.py b/posthog/ai/gemini/gemini.py index da76cb3db..ff74bd270 100644 --- a/posthog/ai/gemini/gemini.py +++ b/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,7 +94,9 @@ 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, @@ -86,6 +104,10 @@ def __init__( 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, @@ -93,6 +115,19 @@ def __init__( 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( diff --git a/posthog/ai/gemini/gemini_async.py b/posthog/ai/gemini/gemini_async.py index bf05a68d8..edede3dca 100644 --- a/posthog/ai/gemini/gemini_async.py +++ b/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,7 +89,9 @@ 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, @@ -87,6 +99,10 @@ def __init__( 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, @@ -94,6 +110,8 @@ def __init__( 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( diff --git a/posthog/test/ai/gemini/test_gemini_parity.py b/posthog/test/ai/gemini/test_gemini_parity.py index 5218b2b00..de9f9b82f 100644 --- a/posthog/test/ai/gemini/test_gemini_parity.py +++ b/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" diff --git a/references/public_api_snapshot.txt b/references/public_api_snapshot.txt index 06d7394ee..4334edcfb 100644 --- a/references/public_api_snapshot.txt +++ b/references/public_api_snapshot.txt @@ -98,6 +98,7 @@ alias posthog.ai.gemini.Client -> posthog.ai.gemini.gemini.Client alias posthog.ai.gemini.extract_gemini_tools -> posthog.ai.gemini.gemini_converter.extract_gemini_tools alias posthog.ai.gemini.format_gemini_input -> posthog.ai.gemini.gemini_converter.format_gemini_input alias posthog.ai.gemini.format_gemini_response -> posthog.ai.gemini.gemini_converter.format_gemini_response +alias posthog.ai.gemini.gemini.AsyncModels -> posthog.ai.gemini.gemini_async.AsyncModels alias posthog.ai.gemini.gemini.PostHogClient -> posthog.client.Client alias posthog.ai.gemini.gemini.StreamingEventData -> posthog.ai.types.StreamingEventData alias posthog.ai.gemini.gemini.TokenUsage -> posthog.ai.types.TokenUsage @@ -376,8 +377,12 @@ attribute posthog.ai.claude_agent_sdk.client.log = logging.getLogger('posthog') attribute posthog.ai.claude_agent_sdk.processor.log = logging.getLogger('posthog') attribute posthog.ai.gateway.POSTHOG_AI_GATEWAY_HOSTS = ['gateway.posthog.com', 'gateway.us.posthog.com', 'gateway.eu.posthog.com', 'ai-gateway.us.posthog.com', 'ai-gateway.eu.posthog.com'] attribute posthog.ai.gateway.log = logging.getLogger('posthog') -attribute posthog.ai.gemini.gemini.Client.models = Models(api_key=api_key, vertexai=vertexai, credentials=credentials, project=project, location=location, debug_config=debug_config, http_options=http_options, 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) -attribute posthog.ai.gemini.gemini_async.AsyncClient.models = AsyncModels(api_key=api_key, vertexai=vertexai, credentials=credentials, project=project, location=location, debug_config=debug_config, http_options=http_options, 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) +attribute posthog.ai.gemini.gemini.Client.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)) +attribute posthog.ai.gemini.gemini.Client.files = self._provider_client.files +attribute posthog.ai.gemini.gemini.Client.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) +attribute posthog.ai.gemini.gemini_async.AsyncClient.aio = _GeminiAioNamespace(models=(self.models), files=(self.files)) +attribute posthog.ai.gemini.gemini_async.AsyncClient.files = self._provider_client.aio.files +attribute posthog.ai.gemini.gemini_async.AsyncClient.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) attribute posthog.ai.gemini.gemini_converter.GeminiMessage.content: Union[str, List[Any]] attribute posthog.ai.gemini.gemini_converter.GeminiMessage.parts: List[Union[GeminiPart, Dict[str, Any]]] attribute posthog.ai.gemini.gemini_converter.GeminiMessage.role: str @@ -908,9 +913,9 @@ class posthog.ai.anthropic.anthropic_providers.AsyncAnthropicVertex(posthog_clie class posthog.ai.claude_agent_sdk.client.PostHogClaudeSDKClient(options: Optional[ClaudeAgentOptions] = None, transport: Any = None, *, posthog_client: Optional[Client] = None, posthog_distinct_id: Optional[Union[str, Callable[[ResultMessage], Optional[str]]]] = None, posthog_trace_id: Optional[str] = None, posthog_properties: Optional[Dict[str, Any]] = None, posthog_privacy_mode: bool = False, posthog_groups: Optional[Dict[str, Any]] = None) class posthog.ai.claude_agent_sdk.processor.PostHogClaudeAgentProcessor(client: Optional[Client] = None, distinct_id: Optional[Union[str, Callable[[ResultMessage], Optional[str]]]] = None, privacy_mode: bool = False, groups: Optional[Dict[str, Any]] = None, properties: Optional[Dict[str, Any]] = None) class posthog.ai.gemini.gemini.Client(api_key: Optional[str] = None, vertexai: Optional[bool] = None, credentials: Optional[Any] = None, project: Optional[str] = None, location: Optional[str] = None, debug_config: Optional[Any] = None, http_options: Optional[Any] = None, posthog_client: Optional[PostHogClient] = None, posthog_distinct_id: Optional[str] = None, posthog_properties: Optional[Dict[str, Any]] = None, posthog_privacy_mode: bool = False, posthog_groups: Optional[Dict[str, Any]] = None, **kwargs) -class posthog.ai.gemini.gemini.Models(api_key: Optional[str] = None, vertexai: Optional[bool] = None, credentials: Optional[Any] = None, project: Optional[str] = None, location: Optional[str] = None, debug_config: Optional[Any] = None, http_options: Optional[Any] = None, posthog_client: Optional[PostHogClient] = None, posthog_distinct_id: Optional[str] = None, posthog_properties: Optional[Dict[str, Any]] = None, posthog_privacy_mode: bool = False, posthog_groups: Optional[Dict[str, Any]] = None, **kwargs) +class posthog.ai.gemini.gemini.Models(api_key: Optional[str] = None, vertexai: Optional[bool] = None, credentials: Optional[Any] = None, project: Optional[str] = None, location: Optional[str] = None, debug_config: Optional[Any] = None, http_options: Optional[Any] = None, posthog_client: Optional[PostHogClient] = None, posthog_distinct_id: Optional[str] = None, 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) class posthog.ai.gemini.gemini_async.AsyncClient(api_key: Optional[str] = None, vertexai: Optional[bool] = None, credentials: Optional[Any] = None, project: Optional[str] = None, location: Optional[str] = None, debug_config: Optional[Any] = None, http_options: Optional[Any] = None, posthog_client: Optional[PostHogClient] = None, posthog_distinct_id: Optional[str] = None, posthog_properties: Optional[Dict[str, Any]] = None, posthog_privacy_mode: bool = False, posthog_groups: Optional[Dict[str, Any]] = None, **kwargs) -class posthog.ai.gemini.gemini_async.AsyncModels(api_key: Optional[str] = None, vertexai: Optional[bool] = None, credentials: Optional[Any] = None, project: Optional[str] = None, location: Optional[str] = None, debug_config: Optional[Any] = None, http_options: Optional[Any] = None, posthog_client: Optional[PostHogClient] = None, posthog_distinct_id: Optional[str] = None, posthog_properties: Optional[Dict[str, Any]] = None, posthog_privacy_mode: bool = False, posthog_groups: Optional[Dict[str, Any]] = None, **kwargs) +class posthog.ai.gemini.gemini_async.AsyncModels(api_key: Optional[str] = None, vertexai: Optional[bool] = None, credentials: Optional[Any] = None, project: Optional[str] = None, location: Optional[str] = None, debug_config: Optional[Any] = None, http_options: Optional[Any] = None, posthog_client: Optional[PostHogClient] = None, posthog_distinct_id: Optional[str] = None, 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) class posthog.ai.gemini.gemini_converter.GeminiMessage class posthog.ai.gemini.gemini_converter.GeminiPart class posthog.ai.langchain.callbacks.CallbackHandler(client: Optional[Client] = None, *, distinct_id: Optional[Union[str, int, UUID]] = None, trace_id: Optional[Union[str, int, float, UUID]] = None, properties: Optional[Dict[str, Any]] = None, privacy_mode: bool = False, groups: Optional[Dict[str, Any]] = None)