feat: add Gemini embed_content tracking (#498)

* feat: add gemini embed_content tracking support

* chore: add sampo changeset for gemini embed_content

* chore: fix black formatting

Author: carlos-marchal-ph

.sampo/changesets/gemini-embed-content.md

@@ -0,0 +1,5 @@
+---
+pypi/posthog: minor
+---
+
+Add Gemini `embed_content` tracking support for both sync and async clients

posthog/ai/gemini/gemini.py

@@ -20,11 +20,13 @@
     merge_usage_stats,
 )
 from posthog.ai.gemini.gemini_converter import (
+    extract_gemini_embedding_token_count,
     extract_gemini_usage_from_chunk,
     extract_gemini_content_from_chunk,
     extract_gemini_stop_reason_from_chunk,
     format_gemini_streaming_output,
 )
+from posthog.ai.utils import with_privacy_mode
 from posthog.ai.sanitization import sanitize_gemini
 from posthog.client import Client as PostHogClient
 
@@ -429,3 +431,88 @@ def generate_content_stream(
             groups,
             **kwargs,
         )
+
+    def embed_content(
+        self,
+        model: str,
+        contents,
+        posthog_distinct_id: Optional[str] = None,
+        posthog_trace_id: Optional[str] = None,
+        posthog_properties: Optional[Dict[str, Any]] = None,
+        posthog_privacy_mode: Optional[bool] = None,
+        posthog_groups: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ):
+        """
+        Create embeddings using Gemini's API while tracking usage in PostHog.
+
+        Args:
+            model: The model to use (e.g., 'gemini-embedding-001')
+            contents: The input content for embedding
+            posthog_distinct_id: ID to associate with the usage event (overrides client default)
+            posthog_trace_id: Trace UUID for linking events (auto-generated if not provided)
+            posthog_properties: Extra properties to include in the event (merged with client defaults)
+            posthog_privacy_mode: Whether to redact sensitive information (overrides client default)
+            posthog_groups: Group analytics properties (overrides client default)
+            **kwargs: Arguments passed to Gemini's embed_content (e.g., config)
+        """
+        distinct_id, trace_id, properties, privacy_mode, groups = (
+            self._merge_posthog_params(
+                posthog_distinct_id,
+                posthog_trace_id,
+                posthog_properties,
+                posthog_privacy_mode,
+                posthog_groups,
+            )
+        )
+
+        start_time = time.time()
+        response = None
+        error = None
+        http_status = 200
+
+        try:
+            response = self._client.models.embed_content(
+                model=model, contents=contents, **kwargs
+            )
+        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, contents),
+                "$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
+
+            self._ph_client.capture(
+                distinct_id=distinct_id or trace_id,
+                event="$ai_embedding",
+                properties=event_properties,
+                groups=groups,
+            )
+
+        if error:
+            raise error
+
+        return response

posthog/ai/gemini/gemini_async.py

@@ -20,11 +20,13 @@
     merge_usage_stats,
 )
 from posthog.ai.gemini.gemini_converter import (
+    extract_gemini_embedding_token_count,
     extract_gemini_usage_from_chunk,
     extract_gemini_content_from_chunk,
     extract_gemini_stop_reason_from_chunk,
     format_gemini_streaming_output,
 )
+from posthog.ai.utils import with_privacy_mode
 from posthog.ai.sanitization import sanitize_gemini
 from posthog.client import Client as PostHogClient
 
@@ -432,3 +434,88 @@ async def generate_content_stream(
             groups,
             **kwargs,
         )
+
+    async def embed_content(
+        self,
+        model: str,
+        contents,
+        posthog_distinct_id: Optional[str] = None,
+        posthog_trace_id: Optional[str] = None,
+        posthog_properties: Optional[Dict[str, Any]] = None,
+        posthog_privacy_mode: Optional[bool] = None,
+        posthog_groups: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ):
+        """
+        Create embeddings using Gemini's API while tracking usage in PostHog.
+
+        Args:
+            model: The model to use (e.g., 'gemini-embedding-001')
+            contents: The input content for embedding
+            posthog_distinct_id: ID to associate with the usage event (overrides client default)
+            posthog_trace_id: Trace UUID for linking events (auto-generated if not provided)
+            posthog_properties: Extra properties to include in the event (merged with client defaults)
+            posthog_privacy_mode: Whether to redact sensitive information (overrides client default)
+            posthog_groups: Group analytics properties (overrides client default)
+            **kwargs: Arguments passed to Gemini's embed_content (e.g., config)
+        """
+        distinct_id, trace_id, properties, privacy_mode, groups = (
+            self._merge_posthog_params(
+                posthog_distinct_id,
+                posthog_trace_id,
+                posthog_properties,
+                posthog_privacy_mode,
+                posthog_groups,
+            )
+        )
+
+        start_time = time.time()
+        response = None
+        error = None
+        http_status = 200
+
+        try:
+            response = await self._client.aio.models.embed_content(
+                model=model, contents=contents, **kwargs
+            )
+        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, contents),
+                "$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
+
+            self._ph_client.capture(
+                distinct_id=distinct_id or trace_id,
+                event="$ai_embedding",
+                properties=event_properties,
+                groups=groups,
+            )
+
+        if error:
+            raise error
+
+        return response

posthog/ai/gemini/gemini_converter.py

@@ -675,3 +675,19 @@ def format_gemini_streaming_output(
 
     # Fallback for empty or unexpected input
     return [{"role": "assistant", "content": [{"type": "text", "text": ""}]}]
+
+
+def extract_gemini_embedding_token_count(response) -> int:
+    """
+    Extract total token count from a Gemini embed_content response.
+    Token counts are only available per-embedding via Vertex AI's statistics.token_count.
+    Returns 0 if no token counts are available.
+    """
+    total = 0
+    if hasattr(response, "embeddings") and response.embeddings:
+        for embedding in response.embeddings:
+            if hasattr(embedding, "statistics") and embedding.statistics:
+                token_count = getattr(embedding.statistics, "token_count", None)
+                if token_count is not None:
+                    total += int(token_count)
+    return total