diff --git a/posthog/ai/gemini/_shared.py b/posthog/ai/gemini/_shared.py new file mode 100644 index 000000000..9ac080c8b --- /dev/null +++ b/posthog/ai/gemini/_shared.py @@ -0,0 +1,230 @@ +import os +import uuid +from typing import Any, Dict, Optional + +try: + from google import genai +except ImportError: + raise ModuleNotFoundError( + "Please install the Google Gemini SDK to use this feature: 'pip install google-genai'" + ) + +from ... import setup +from ...client import Client as PostHogClient +from ..types import StreamingEventData, TokenUsage +from ..utils import ( + _capture_ai_event, + capture_streaming_event, + finalize_ai_content, + merge_system_prompt, + with_privacy_mode, +) +from .gemini_converter import ( + extract_gemini_embedding_token_count, + format_gemini_streaming_output, +) + +_GEMINI_BASE_URL = "https://generativelanguage.googleapis.com" + + +def _resolve_posthog_client( + posthog_client: Optional[PostHogClient], +) -> PostHogClient: + client = posthog_client or setup() + if client is None: + raise ValueError("posthog_client is required for PostHog tracking") + return client + + +def _build_gemini_client_args( + *, + api_key: Optional[str], + vertexai: Optional[bool], + credentials: Optional[Any], + project: Optional[str], + location: Optional[str], + debug_config: Optional[Any], + http_options: Optional[Any], +) -> Dict[str, Any]: + """Build provider client arguments while preserving Gemini auth precedence.""" + client_args: Dict[str, Any] = {} + + optional_args = { + "vertexai": vertexai, + "credentials": credentials, + "project": project, + "location": location, + "debug_config": debug_config, + "http_options": http_options, + } + client_args.update( + {name: value for name, value in optional_args.items() if value is not None} + ) + + if vertexai: + if api_key is not None: + client_args["api_key"] = api_key + return client_args + + resolved_api_key = api_key + if resolved_api_key is None: + resolved_api_key = os.environ.get("GOOGLE_API_KEY") or os.environ.get("API_KEY") + if resolved_api_key is None: + raise ValueError( + "API key must be provided either as parameter or via GOOGLE_API_KEY/API_KEY environment variable" + ) + + client_args["api_key"] = resolved_api_key + return client_args + + +class _GeminiModelsPolicy: + """Shared telemetry policy for the explicit sync and async Gemini adapters.""" + + _ph_client: PostHogClient + + def _initialize_policy( + self, + *, + api_key: Optional[str], + vertexai: Optional[bool], + credentials: Optional[Any], + project: Optional[str], + location: Optional[str], + debug_config: Optional[Any], + http_options: Optional[Any], + posthog_client: Optional[PostHogClient], + posthog_distinct_id: Optional[str], + posthog_properties: Optional[Dict[str, Any]], + posthog_privacy_mode: bool, + posthog_groups: Optional[Dict[str, Any]], + ) -> 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, + ) + self._client = genai.Client(**client_args) + self._base_url = _GEMINI_BASE_URL + + def _merge_posthog_params( + self, + call_distinct_id: Optional[str], + call_trace_id: Optional[str], + call_properties: Optional[Dict[str, Any]], + call_privacy_mode: Optional[bool], + call_groups: Optional[Dict[str, Any]], + ): + """Merge call-level PostHog parameters with client defaults.""" + distinct_id = ( + call_distinct_id + if call_distinct_id is not None + else self._default_distinct_id + ) + privacy_mode = ( + call_privacy_mode + if call_privacy_mode is not None + else self._default_privacy_mode + ) + groups = call_groups if call_groups is not None else self._default_groups + + properties = dict(self._default_properties) + if call_properties: + properties.update(call_properties) + + trace_id = call_trace_id if call_trace_id is not None else str(uuid.uuid4()) + return distinct_id, trace_id, properties, privacy_mode, groups + + def _capture_streaming_event( + self, + model: str, + contents, + distinct_id: Optional[str], + trace_id: Optional[str], + properties: Optional[Dict[str, Any]], + privacy_mode: bool, + groups: Optional[Dict[str, Any]], + kwargs: Dict[str, Any], + usage_stats: TokenUsage, + latency: float, + output: Any, + stop_reason: Optional[str] = None, + ) -> None: + formatted_input = merge_system_prompt( + {"contents": contents, **kwargs}, "gemini" + ) + event_data = StreamingEventData( + provider="gemini", + model=model, + base_url=self._base_url, + kwargs=kwargs, + formatted_input=formatted_input, + formatted_output=format_gemini_streaming_output(output), + usage_stats=usage_stats, + latency=latency, + distinct_id=distinct_id, + trace_id=trace_id, + properties=properties, + privacy_mode=privacy_mode, + groups=groups, + stop_reason=stop_reason, + ) + capture_streaming_event(self._ph_client, event_data) + + def _capture_embedding_outcome( + self, + *, + model: str, + contents: Any, + distinct_id: Optional[str], + trace_id: str, + properties: Optional[Dict[str, Any]], + privacy_mode: bool, + groups: Optional[Dict[str, Any]], + response: Any, + error: Optional[Exception], + latency: float, + ) -> None: + input_tokens = extract_gemini_embedding_token_count(response) if response else 0 + event_properties = { + "$ai_provider": "gemini", + "$ai_model": model, + "$ai_input": with_privacy_mode( + self._ph_client, + privacy_mode, + finalize_ai_content(contents, self._ph_client), + ), + "$ai_http_status": ( + getattr(error, "status_code", 0) if error is not None else 200 + ), + "$ai_input_tokens": input_tokens, + "$ai_latency": latency, + "$ai_trace_id": trace_id, + "$ai_base_url": self._base_url, + **(properties or {}), + } + + if error: + event_properties["$ai_is_error"] = True + event_properties["$ai_error"] = str(error) + + if distinct_id is None: + event_properties["$process_person_profile"] = False + + _capture_ai_event( + self._ph_client, + "$ai_embedding", + distinct_id=distinct_id or trace_id, + properties=event_properties, + groups=groups, + ) diff --git a/posthog/ai/gemini/gemini.py b/posthog/ai/gemini/gemini.py index 09c0d7645..d8e977e9f 100644 --- a/posthog/ai/gemini/gemini.py +++ b/posthog/ai/gemini/gemini.py @@ -1,35 +1,26 @@ -import os import time -import uuid from typing import Any, Dict, Optional -from posthog.ai.types import TokenUsage, StreamingEventData -from posthog.ai.utils import merge_system_prompt - -try: - from google import genai -except ImportError: - raise ModuleNotFoundError( - "Please install the Google Gemini SDK to use this feature: 'pip install google-genai'" - ) - -from posthog import setup -from posthog.ai.utils import ( +from ... import setup as setup +from ...client import Client as PostHogClient +from ..types import StreamingEventData as StreamingEventData +from ..types import TokenUsage +from ..utils import ( call_llm_and_track_usage, - _capture_ai_event, - capture_streaming_event, - finalize_ai_content, + capture_streaming_event as capture_streaming_event, + finalize_ai_content as finalize_ai_content, + merge_system_prompt as merge_system_prompt, merge_usage_stats, + with_privacy_mode as with_privacy_mode, ) -from posthog.ai.gemini.gemini_converter import ( - extract_gemini_embedding_token_count, - extract_gemini_usage_from_chunk, +from ._shared import _GeminiModelsPolicy, _resolve_posthog_client +from .gemini_converter import ( extract_gemini_content_from_chunk, + extract_gemini_embedding_token_count as extract_gemini_embedding_token_count, extract_gemini_stop_reason_from_chunk, - format_gemini_streaming_output, + extract_gemini_usage_from_chunk, + format_gemini_streaming_output as format_gemini_streaming_output, ) -from posthog.ai.utils import with_privacy_mode -from posthog.client import Client as PostHogClient class Client: @@ -85,10 +76,7 @@ def __init__( **kwargs: Additional arguments (for future compatibility) """ - self._ph_client = posthog_client or setup() - - if self._ph_client is None: - raise ValueError("posthog_client is required for PostHog tracking") + self._ph_client = _resolve_posthog_client(posthog_client) self.models = Models( api_key=api_key, @@ -107,7 +95,7 @@ def __init__( ) -class Models: +class Models(_GeminiModelsPolicy): """ Models interface that mimics genai.Client().models with PostHog tracking. """ @@ -147,92 +135,20 @@ def __init__( **kwargs: Additional arguments (for future compatibility) """ - self._ph_client = posthog_client or setup() - - if self._ph_client is None: - raise ValueError("posthog_client is required for PostHog tracking") - - # Store default PostHog settings - 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 - - # Build genai.Client arguments - client_args: Dict[str, Any] = {} - - # Add Vertex AI parameters if provided - if vertexai is not None: - client_args["vertexai"] = vertexai - - if credentials is not None: - client_args["credentials"] = credentials - - if project is not None: - client_args["project"] = project - - if location is not None: - client_args["location"] = location - - if debug_config is not None: - client_args["debug_config"] = debug_config - - if http_options is not None: - client_args["http_options"] = http_options - - # Handle API key authentication - if vertexai: - # For Vertex AI, api_key is optional - if api_key is not None: - client_args["api_key"] = api_key - else: - # For non-Vertex AI mode, api_key is required (backwards compatibility) - if api_key is None: - api_key = os.environ.get("GOOGLE_API_KEY") or os.environ.get("API_KEY") - - if api_key is None: - raise ValueError( - "API key must be provided either as parameter or via GOOGLE_API_KEY/API_KEY environment variable" - ) - - client_args["api_key"] = api_key - - self._client = genai.Client(**client_args) - self._base_url = "https://generativelanguage.googleapis.com" - - def _merge_posthog_params( - self, - call_distinct_id: Optional[str], - call_trace_id: Optional[str], - call_properties: Optional[Dict[str, Any]], - call_privacy_mode: Optional[bool], - call_groups: Optional[Dict[str, Any]], - ): - """Merge call-level PostHog parameters with client defaults.""" - - # Use call-level values if provided, otherwise fall back to defaults - distinct_id = ( - call_distinct_id - if call_distinct_id is not None - else self._default_distinct_id - ) - privacy_mode = ( - call_privacy_mode - if call_privacy_mode is not None - else self._default_privacy_mode + self._initialize_policy( + api_key=api_key, + vertexai=vertexai, + credentials=credentials, + project=project, + location=location, + debug_config=debug_config, + http_options=http_options, + posthog_client=posthog_client, + posthog_distinct_id=posthog_distinct_id, + posthog_properties=posthog_properties, + posthog_privacy_mode=posthog_privacy_mode, + posthog_groups=posthog_groups, ) - groups = call_groups if call_groups is not None else self._default_groups - - # Merge properties: default properties + call properties (call properties override) - properties = dict(self._default_properties) - - if call_properties: - properties.update(call_properties) - - if call_trace_id is None: - call_trace_id = str(uuid.uuid4()) - - return distinct_id, call_trace_id, properties, privacy_mode, groups def generate_content( self, @@ -354,50 +270,6 @@ def generator(): return generator() - def _capture_streaming_event( - self, - model: str, - contents, - distinct_id: Optional[str], - trace_id: Optional[str], - properties: Optional[Dict[str, Any]], - privacy_mode: bool, - groups: Optional[Dict[str, Any]], - kwargs: Dict[str, Any], - usage_stats: TokenUsage, - latency: float, - output: Any, - stop_reason: Optional[str] = None, - ): - formatted_input = self._format_input(contents, **kwargs) - - event_data = StreamingEventData( - provider="gemini", - model=model, - base_url=self._base_url, - kwargs=kwargs, - formatted_input=formatted_input, - formatted_output=format_gemini_streaming_output(output), - usage_stats=usage_stats, - latency=latency, - distinct_id=distinct_id, - trace_id=trace_id, - properties=properties, - privacy_mode=privacy_mode, - groups=groups, - stop_reason=stop_reason, - ) - - # Use the common capture function - capture_streaming_event(self._ph_client, event_data) - - def _format_input(self, contents, **kwargs): - """Format input contents for PostHog tracking""" - - # Create kwargs dict with contents for merge_system_prompt - input_kwargs = {"contents": contents, **kwargs} - return merge_system_prompt(input_kwargs, "gemini") - def generate_content_stream( self, model: str, @@ -485,7 +357,6 @@ def embed_content( start_time = time.time() response = None error = None - http_status = 200 try: response = self._client.models.embed_content( @@ -493,44 +364,18 @@ def embed_content( ) except Exception as exc: error = exc - http_status = getattr(exc, "status_code", 0) finally: - end_time = time.time() - latency = end_time - start_time - - input_tokens = ( - extract_gemini_embedding_token_count(response) if response else 0 - ) - - event_properties = { - "$ai_provider": "gemini", - "$ai_model": model, - "$ai_input": with_privacy_mode( - self._ph_client, - privacy_mode, - finalize_ai_content(contents, self._ph_client), - ), - "$ai_http_status": http_status, - "$ai_input_tokens": input_tokens, - "$ai_latency": latency, - "$ai_trace_id": trace_id, - "$ai_base_url": self._base_url, - **(properties or {}), - } - - if error: - event_properties["$ai_is_error"] = True - event_properties["$ai_error"] = str(error) - - if distinct_id is None: - event_properties["$process_person_profile"] = False - - _capture_ai_event( - self._ph_client, - "$ai_embedding", - distinct_id=distinct_id or trace_id, - properties=event_properties, + self._capture_embedding_outcome( + model=model, + contents=contents, + distinct_id=distinct_id, + trace_id=trace_id, + properties=properties, + privacy_mode=privacy_mode, groups=groups, + response=response, + error=error, + latency=time.time() - start_time, ) if error: diff --git a/posthog/ai/gemini/gemini_async.py b/posthog/ai/gemini/gemini_async.py index a7497f4ed..78ee7dbee 100644 --- a/posthog/ai/gemini/gemini_async.py +++ b/posthog/ai/gemini/gemini_async.py @@ -1,36 +1,27 @@ -import os import time -import uuid from typing import Any, Dict, Optional -from posthog.ai.stream import AsyncStreamWrapper -from posthog.ai.types import TokenUsage, StreamingEventData -from posthog.ai.utils import merge_system_prompt - -try: - from google import genai -except ImportError: - raise ModuleNotFoundError( - "Please install the Google Gemini SDK to use this feature: 'pip install google-genai'" - ) - -from posthog import setup -from posthog.ai.utils import ( +from ... import setup as setup +from ...client import Client as PostHogClient +from ..stream import AsyncStreamWrapper +from ..types import StreamingEventData as StreamingEventData +from ..types import TokenUsage +from ..utils import ( call_llm_and_track_usage_async, - _capture_ai_event, - capture_streaming_event, - finalize_ai_content, + capture_streaming_event as capture_streaming_event, + finalize_ai_content as finalize_ai_content, + merge_system_prompt as merge_system_prompt, merge_usage_stats, + with_privacy_mode as with_privacy_mode, ) -from posthog.ai.gemini.gemini_converter import ( - extract_gemini_embedding_token_count, - extract_gemini_usage_from_chunk, +from ._shared import _GeminiModelsPolicy, _resolve_posthog_client +from .gemini_converter import ( extract_gemini_content_from_chunk, + extract_gemini_embedding_token_count as extract_gemini_embedding_token_count, extract_gemini_stop_reason_from_chunk, - format_gemini_streaming_output, + extract_gemini_usage_from_chunk, + format_gemini_streaming_output as format_gemini_streaming_output, ) -from posthog.ai.utils import with_privacy_mode -from posthog.client import Client as PostHogClient class AsyncClient: @@ -86,10 +77,7 @@ def __init__( **kwargs: Additional arguments (for future compatibility) """ - self._ph_client = posthog_client or setup() - - if self._ph_client is None: - raise ValueError("posthog_client is required for PostHog tracking") + self._ph_client = _resolve_posthog_client(posthog_client) self.models = AsyncModels( api_key=api_key, @@ -108,7 +96,7 @@ def __init__( ) -class AsyncModels: +class AsyncModels(_GeminiModelsPolicy): """ Async Models interface that mimics genai.Client().aio.models with PostHog tracking. """ @@ -148,92 +136,20 @@ def __init__( **kwargs: Additional arguments (for future compatibility) """ - self._ph_client = posthog_client or setup() - - if self._ph_client is None: - raise ValueError("posthog_client is required for PostHog tracking") - - # Store default PostHog settings - 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 - - # Build genai.Client arguments - client_args: Dict[str, Any] = {} - - # Add Vertex AI parameters if provided - if vertexai is not None: - client_args["vertexai"] = vertexai - - if credentials is not None: - client_args["credentials"] = credentials - - if project is not None: - client_args["project"] = project - - if location is not None: - client_args["location"] = location - - if debug_config is not None: - client_args["debug_config"] = debug_config - - if http_options is not None: - client_args["http_options"] = http_options - - # Handle API key authentication - if vertexai: - # For Vertex AI, api_key is optional - if api_key is not None: - client_args["api_key"] = api_key - else: - # For non-Vertex AI mode, api_key is required (backwards compatibility) - if api_key is None: - api_key = os.environ.get("GOOGLE_API_KEY") or os.environ.get("API_KEY") - - if api_key is None: - raise ValueError( - "API key must be provided either as parameter or via GOOGLE_API_KEY/API_KEY environment variable" - ) - - client_args["api_key"] = api_key - - self._client = genai.Client(**client_args) - self._base_url = "https://generativelanguage.googleapis.com" - - def _merge_posthog_params( - self, - call_distinct_id: Optional[str], - call_trace_id: Optional[str], - call_properties: Optional[Dict[str, Any]], - call_privacy_mode: Optional[bool], - call_groups: Optional[Dict[str, Any]], - ): - """Merge call-level PostHog parameters with client defaults.""" - - # Use call-level values if provided, otherwise fall back to defaults - distinct_id = ( - call_distinct_id - if call_distinct_id is not None - else self._default_distinct_id - ) - privacy_mode = ( - call_privacy_mode - if call_privacy_mode is not None - else self._default_privacy_mode + self._initialize_policy( + api_key=api_key, + vertexai=vertexai, + credentials=credentials, + project=project, + location=location, + debug_config=debug_config, + http_options=http_options, + posthog_client=posthog_client, + posthog_distinct_id=posthog_distinct_id, + posthog_properties=posthog_properties, + posthog_privacy_mode=posthog_privacy_mode, + posthog_groups=posthog_groups, ) - groups = call_groups if call_groups is not None else self._default_groups - - # Merge properties: default properties + call properties (call properties override) - properties = dict(self._default_properties) - - if call_properties: - properties.update(call_properties) - - if call_trace_id is None: - call_trace_id = str(uuid.uuid4()) - - return distinct_id, call_trace_id, properties, privacy_mode, groups async def generate_content( self, @@ -358,50 +274,6 @@ async def async_generator(): return AsyncStreamWrapper(async_generator(), stream=response) - def _capture_streaming_event( - self, - model: str, - contents, - distinct_id: Optional[str], - trace_id: Optional[str], - properties: Optional[Dict[str, Any]], - privacy_mode: bool, - groups: Optional[Dict[str, Any]], - kwargs: Dict[str, Any], - usage_stats: TokenUsage, - latency: float, - output: Any, - stop_reason: Optional[str] = None, - ): - formatted_input = self._format_input(contents, **kwargs) - - event_data = StreamingEventData( - provider="gemini", - model=model, - base_url=self._base_url, - kwargs=kwargs, - formatted_input=formatted_input, - formatted_output=format_gemini_streaming_output(output), - usage_stats=usage_stats, - latency=latency, - distinct_id=distinct_id, - trace_id=trace_id, - properties=properties, - privacy_mode=privacy_mode, - groups=groups, - stop_reason=stop_reason, - ) - - # Use the common capture function - capture_streaming_event(self._ph_client, event_data) - - def _format_input(self, contents, **kwargs): - """Format input contents for PostHog tracking""" - - # Create kwargs dict with contents for merge_system_prompt - input_kwargs = {"contents": contents, **kwargs} - return merge_system_prompt(input_kwargs, "gemini") - async def generate_content_stream( self, model: str, @@ -489,7 +361,6 @@ async def embed_content( start_time = time.time() response = None error = None - http_status = 200 try: response = await self._client.aio.models.embed_content( @@ -497,44 +368,18 @@ async def embed_content( ) except Exception as exc: error = exc - http_status = getattr(exc, "status_code", 0) finally: - end_time = time.time() - latency = end_time - start_time - - input_tokens = ( - extract_gemini_embedding_token_count(response) if response else 0 - ) - - event_properties = { - "$ai_provider": "gemini", - "$ai_model": model, - "$ai_input": with_privacy_mode( - self._ph_client, - privacy_mode, - finalize_ai_content(contents, self._ph_client), - ), - "$ai_http_status": http_status, - "$ai_input_tokens": input_tokens, - "$ai_latency": latency, - "$ai_trace_id": trace_id, - "$ai_base_url": self._base_url, - **(properties or {}), - } - - if error: - event_properties["$ai_is_error"] = True - event_properties["$ai_error"] = str(error) - - if distinct_id is None: - event_properties["$process_person_profile"] = False - - _capture_ai_event( - self._ph_client, - "$ai_embedding", - distinct_id=distinct_id or trace_id, - properties=event_properties, + self._capture_embedding_outcome( + model=model, + contents=contents, + distinct_id=distinct_id, + trace_id=trace_id, + properties=properties, + privacy_mode=privacy_mode, groups=groups, + response=response, + error=error, + latency=time.time() - start_time, ) if error: diff --git a/posthog/test/ai/gemini/test_gemini_parity.py b/posthog/test/ai/gemini/test_gemini_parity.py new file mode 100644 index 000000000..5218b2b00 --- /dev/null +++ b/posthog/test/ai/gemini/test_gemini_parity.py @@ -0,0 +1,67 @@ +from unittest.mock import MagicMock, patch + +import pytest + +try: + from google import genai as google_genai + + from posthog.ai.gemini import AsyncClient, Client +except ImportError: + pytest.skip("Google Gemini package is not available", allow_module_level=True) + + +@pytest.mark.parametrize("client_class", [Client, AsyncClient]) +def test_sync_and_async_clients_share_api_key_environment_precedence( + client_class, monkeypatch +): + monkeypatch.setenv("GOOGLE_API_KEY", "google-key") + monkeypatch.setenv("API_KEY", "legacy-key") + + with patch.object(google_genai, "Client") as provider_client: + client_class(posthog_client=MagicMock()) + + provider_client.assert_called_once_with(api_key="google-key") + + +@pytest.mark.parametrize("client_class", [Client, AsyncClient]) +def test_sync_and_async_clients_preserve_explicit_empty_api_key( + client_class, monkeypatch +): + monkeypatch.setenv("GOOGLE_API_KEY", "environment-key") + + with patch.object(google_genai, "Client") as provider_client: + client_class(api_key="", posthog_client=MagicMock()) + + provider_client.assert_called_once_with(api_key="") + + +@pytest.mark.parametrize("client_class", [Client, AsyncClient]) +def test_sync_and_async_clients_merge_posthog_defaults_without_mutation(client_class): + default_properties = {"shared": "default", "default-only": True} + + with patch.object(google_genai, "Client"): + client = client_class( + api_key="test-key", + posthog_client=MagicMock(), + posthog_distinct_id="default-id", + posthog_properties=default_properties, + posthog_privacy_mode=True, + posthog_groups={"organization": "default-org"}, + ) + + merged = client.models._merge_posthog_params( + None, + "", + {"shared": "call", "call-only": True}, + False, + None, + ) + + assert merged == ( + "default-id", + "", + {"shared": "call", "default-only": True, "call-only": True}, + False, + {"organization": "default-org"}, + ) + assert default_properties == {"shared": "default", "default-only": True}