posthoganalytics 6.7.5__py3-none-any.whl → 7.4.3__py3-none-any.whl

posthoganalytics/ai/gemini/gemini_async.py

@@ -0,0 +1,423 @@
+import os
+import time
+import uuid
+from typing import Any, Dict, Optional
+
+from posthoganalytics.ai.types import TokenUsage, StreamingEventData
+from posthoganalytics.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 posthoganalytics import setup
+from posthoganalytics.ai.utils import (
+    call_llm_and_track_usage_async,
+    capture_streaming_event,
+    merge_usage_stats,
+)
+from posthoganalytics.ai.gemini.gemini_converter import (
+    extract_gemini_usage_from_chunk,
+    extract_gemini_content_from_chunk,
+    format_gemini_streaming_output,
+)
+from posthoganalytics.ai.sanitization import sanitize_gemini
+from posthoganalytics.client import Client as PostHogClient
+
+
+class AsyncClient:
+    """
+    An async drop-in replacement for genai.Client that automatically sends LLM usage events to PostHog.
+
+    Usage:
+        client = AsyncClient(
+            api_key="your_api_key",
+            posthog_client=posthog_client,
+            posthog_distinct_id="default_user",  # Optional defaults
+            posthog_properties={"team": "ai"}    # Optional defaults
+        )
+        response = await client.models.generate_content(
+            model="gemini-2.0-flash",
+            contents=["Hello world"],
+            posthog_distinct_id="specific_user"  # Override default
+        )
+    """
+
+    _ph_client: PostHogClient
+
+    def __init__(
+        self,
+        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,
+    ):
+        """
+        Args:
+            api_key: Google AI API key. If not provided, will use GOOGLE_API_KEY or API_KEY environment variable (not required for Vertex AI)
+            vertexai: Whether to use Vertex AI authentication
+            credentials: Vertex AI credentials object
+            project: GCP project ID for Vertex AI
+            location: GCP location for Vertex AI
+            debug_config: Debug configuration for the client
+            http_options: HTTP options for the client
+            posthog_client: PostHog client for tracking usage
+            posthog_distinct_id: Default distinct ID for all calls (can be overridden per call)
+            posthog_properties: Default properties for all calls (can be overridden per call)
+            posthog_privacy_mode: Default privacy mode for all calls (can be overridden per call)
+            posthog_groups: Default groups for all calls (can be overridden per call)
+            **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.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,
+        )
+
+
+class AsyncModels:
+    """
+    Async Models interface that mimics genai.Client().aio.models with PostHog tracking.
+    """
+
+    _ph_client: PostHogClient  # Not None after __init__ validation
+
+    def __init__(
+        self,
+        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,
+    ):
+        """
+        Args:
+            api_key: Google AI API key. If not provided, will use GOOGLE_API_KEY or API_KEY environment variable (not required for Vertex AI)
+            vertexai: Whether to use Vertex AI authentication
+            credentials: Vertex AI credentials object
+            project: GCP project ID for Vertex AI
+            location: GCP location for Vertex AI
+            debug_config: Debug configuration for the client
+            http_options: HTTP options for the client
+            posthog_client: PostHog client for tracking usage
+            posthog_distinct_id: Default distinct ID for all calls
+            posthog_properties: Default properties for all calls
+            posthog_privacy_mode: Default privacy mode for all calls
+            posthog_groups: Default groups for all calls
+            **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
+        )
+        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,
+        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,
+    ):
+        """
+        Generate content using Gemini's API while tracking usage in PostHog.
+
+        This method signature exactly matches genai.Client().aio.models.generate_content()
+        with additional PostHog tracking parameters.
+
+        Args:
+            model: The model to use (e.g., 'gemini-2.0-flash')
+            contents: The input content for generation
+            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 generate_content
+        """
+
+        # Merge PostHog parameters
+        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,
+            )
+        )
+
+        kwargs_with_contents = {"model": model, "contents": contents, **kwargs}
+
+        return await call_llm_and_track_usage_async(
+            distinct_id,
+            self._ph_client,
+            "gemini",
+            trace_id,
+            properties,
+            privacy_mode,
+            groups,
+            self._base_url,
+            self._client.aio.models.generate_content,
+            **kwargs_with_contents,
+        )
+
+    async def _generate_content_streaming(
+        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: Any,
+    ):
+        start_time = time.time()
+        usage_stats: TokenUsage = TokenUsage(input_tokens=0, output_tokens=0)
+        accumulated_content = []
+
+        kwargs_without_stream = {"model": model, "contents": contents, **kwargs}
+        response = await self._client.aio.models.generate_content_stream(
+            **kwargs_without_stream
+        )
+
+        async def async_generator():
+            nonlocal usage_stats
+            nonlocal accumulated_content
+
+            try:
+                async for chunk in response:
+                    # Extract usage stats from chunk
+                    chunk_usage = extract_gemini_usage_from_chunk(chunk)
+
+                    if chunk_usage:
+                        # Gemini reports cumulative totals, not incremental values
+                        merge_usage_stats(usage_stats, chunk_usage, mode="cumulative")
+
+                    # Extract content from chunk (now returns content blocks)
+                    content_block = extract_gemini_content_from_chunk(chunk)
+
+                    if content_block is not None:
+                        accumulated_content.append(content_block)
+
+                    yield chunk
+
+            finally:
+                end_time = time.time()
+                latency = end_time - start_time
+
+                self._capture_streaming_event(
+                    model,
+                    contents,
+                    distinct_id,
+                    trace_id,
+                    properties,
+                    privacy_mode,
+                    groups,
+                    kwargs,
+                    usage_stats,
+                    latency,
+                    accumulated_content,
+                )
+
+        return async_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,
+    ):
+        # Prepare standardized event data
+        formatted_input = self._format_input(contents, **kwargs)
+        sanitized_input = sanitize_gemini(formatted_input)
+
+        event_data = StreamingEventData(
+            provider="gemini",
+            model=model,
+            base_url=self._base_url,
+            kwargs=kwargs,
+            formatted_input=sanitized_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,
+        )
+
+        # 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,
+        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,
+    ):
+        # Merge PostHog parameters
+        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,
+            )
+        )
+
+        return await self._generate_content_streaming(
+            model,
+            contents,
+            distinct_id,
+            trace_id,
+            properties,
+            privacy_mode,
+            groups,
+            **kwargs,
+        )

posthoganalytics/ai/gemini/gemini_converter.py

@@ -29,35 +29,76 @@ class GeminiMessage(TypedDict, total=False):
     text: str
 
 
-def _extract_text_from_parts(parts: List[Any]) -> str:
+def _format_parts_as_content_blocks(parts: List[Any]) -> List[FormattedContentItem]:
     """
-    Extract and concatenate text from a parts array.
+    Format Gemini parts array into structured content blocks.
+
+    Preserves structure for multimodal content (text + images) instead of
+    concatenating everything into a string.
 
     Args:
-        parts: List of parts that may contain text content
+        parts: List of parts that may contain text, inline_data, etc.
 
     Returns:
-        Concatenated text from all parts
+        List of formatted content blocks
     """
-
-    content_parts = []
+    content_blocks: List[FormattedContentItem] = []
 
     for part in parts:
+        # Handle dict with text field
         if isinstance(part, dict) and "text" in part:
-            content_parts.append(part["text"])
+            content_blocks.append({"type": "text", "text": part["text"]})
 
+        # Handle string parts
         elif isinstance(part, str):
-            content_parts.append(part)
+            content_blocks.append({"type": "text", "text": part})
+
+        # Handle dict with inline_data (images, documents, etc.)
+        elif isinstance(part, dict) and "inline_data" in part:
+            inline_data = part["inline_data"]
+            mime_type = inline_data.get("mime_type", "")
+            content_type = "image" if mime_type.startswith("image/") else "document"
+
+            content_blocks.append(
+                {
+                    "type": content_type,
+                    "inline_data": inline_data,
+                }
+            )
 
+        # Handle object with text attribute
         elif hasattr(part, "text"):
-            # Get the text attribute value
             text_value = getattr(part, "text", "")
-            content_parts.append(text_value if text_value else str(part))
-
-        else:
-            content_parts.append(str(part))
+            if text_value:
+                content_blocks.append({"type": "text", "text": text_value})
+
+        # Handle object with inline_data attribute
+        elif hasattr(part, "inline_data"):
+            inline_data = part.inline_data
+            # Convert to dict if needed
+            if hasattr(inline_data, "mime_type") and hasattr(inline_data, "data"):
+                # Determine type based on mime_type
+                mime_type = inline_data.mime_type
+                content_type = "image" if mime_type.startswith("image/") else "document"
+
+                content_blocks.append(
+                    {
+                        "type": content_type,
+                        "inline_data": {
+                            "mime_type": mime_type,
+                            "data": inline_data.data,
+                        },
+                    }
+                )
+            else:
+                content_blocks.append(
+                    {
+                        "type": "image",
+                        "inline_data": inline_data,
+                    }
+                )
 
-    return "".join(content_parts)
+    return content_blocks
 
 
 def _format_dict_message(item: Dict[str, Any]) -> FormattedMessage:
@@ -73,16 +114,17 @@ def _format_dict_message(item: Dict[str, Any]) -> FormattedMessage:
 
     # Handle dict format with parts array (Gemini-specific format)
     if "parts" in item and isinstance(item["parts"], list):
-        content = _extract_text_from_parts(item["parts"])
-        return {"role": item.get("role", "user"), "content": content}
+        content_blocks = _format_parts_as_content_blocks(item["parts"])
+        return {"role": item.get("role", "user"), "content": content_blocks}
 
     # Handle dict with content field
     if "content" in item:
         content = item["content"]
 
         if isinstance(content, list):
-            # If content is a list, extract text from it
-            content = _extract_text_from_parts(content)
+            # If content is a list, format it as content blocks
+            content_blocks = _format_parts_as_content_blocks(content)
+            return {"role": item.get("role", "user"), "content": content_blocks}
 
         elif not isinstance(content, str):
             content = str(content)
@@ -110,14 +152,14 @@ def _format_object_message(item: Any) -> FormattedMessage:
 
     # Handle object with parts attribute
     if hasattr(item, "parts") and hasattr(item.parts, "__iter__"):
-        content = _extract_text_from_parts(item.parts)
+        content_blocks = _format_parts_as_content_blocks(list(item.parts))
         role = getattr(item, "role", "user") if hasattr(item, "role") else "user"
 
         # Ensure role is a string
         if not isinstance(role, str):
             role = "user"
 
-        return {"role": role, "content": content}
+        return {"role": role, "content": content_blocks}
 
     # Handle object with text attribute
     if hasattr(item, "text"):
@@ -140,7 +182,8 @@ def _format_object_message(item: Any) -> FormattedMessage:
         content = item.content
 
         if isinstance(content, list):
-            content = _extract_text_from_parts(content)
+            content_blocks = _format_parts_as_content_blocks(content)
+            return {"role": role, "content": content_blocks}
 
         elif not isinstance(content, str):
             content = str(content)
@@ -193,6 +236,29 @@ def format_gemini_response(response: Any) -> List[FormattedMessage]:
                                 }
                             )
 
+                        elif hasattr(part, "inline_data") and part.inline_data:
+                            # Handle audio/media inline data
+                            import base64
+
+                            inline_data = part.inline_data
+                            mime_type = getattr(inline_data, "mime_type", "audio/pcm")
+                            raw_data = getattr(inline_data, "data", b"")
+
+                            # Encode binary data as base64 string for JSON serialization
+                            if isinstance(raw_data, bytes):
+                                data = base64.b64encode(raw_data).decode("utf-8")
+                            else:
+                                # Already a string (base64)
+                                data = raw_data
+
+                            content.append(
+                                {
+                                    "type": "audio",
+                                    "mime_type": mime_type,
+                                    "data": data,
+                                }
+                            )
+
                 if content:
                     output.append(
                         {
@@ -338,6 +404,61 @@ def format_gemini_input(contents: Any) -> List[FormattedMessage]:
     return [_format_object_message(contents)]
 
 
+def extract_gemini_web_search_count(response: Any) -> int:
+    """
+    Extract web search count from Gemini response.
+
+    Gemini bills per request that uses grounding, not per query.
+    Returns 1 if grounding_metadata is present with actual search data, 0 otherwise.
+
+    Args:
+        response: The response from Gemini API
+
+    Returns:
+        1 if web search/grounding was used, 0 otherwise
+    """
+
+    # Check for grounding_metadata in candidates
+    if hasattr(response, "candidates"):
+        for candidate in response.candidates:
+            if (
+                hasattr(candidate, "grounding_metadata")
+                and candidate.grounding_metadata
+            ):
+                grounding_metadata = candidate.grounding_metadata
+
+                # Check if web_search_queries exists and is non-empty
+                if hasattr(grounding_metadata, "web_search_queries"):
+                    queries = grounding_metadata.web_search_queries
+
+                    if queries is not None and len(queries) > 0:
+                        return 1
+
+                # Check if grounding_chunks exists and is non-empty
+                if hasattr(grounding_metadata, "grounding_chunks"):
+                    chunks = grounding_metadata.grounding_chunks
+
+                    if chunks is not None and len(chunks) > 0:
+                        return 1
+
+            # Also check for google_search or grounding in function call names
+            if hasattr(candidate, "content") and candidate.content:
+                if hasattr(candidate.content, "parts") and candidate.content.parts:
+                    for part in candidate.content.parts:
+                        if hasattr(part, "function_call") and part.function_call:
+                            function_name = getattr(
+                                part.function_call, "name", ""
+                            ).lower()
+
+                            if (
+                                "google_search" in function_name
+                                or "grounding" in function_name
+                            ):
+                                return 1
+
+    return 0
+
+
 def _extract_usage_from_metadata(metadata: Any) -> TokenUsage:
     """
     Common logic to extract usage from Gemini metadata.
@@ -382,7 +503,14 @@ def extract_gemini_usage_from_response(response: Any) -> TokenUsage:
     if not hasattr(response, "usage_metadata") or not response.usage_metadata:
         return TokenUsage(input_tokens=0, output_tokens=0)
 
-    return _extract_usage_from_metadata(response.usage_metadata)
+    usage = _extract_usage_from_metadata(response.usage_metadata)
+
+    # Add web search count if present
+    web_search_count = extract_gemini_web_search_count(response)
+    if web_search_count > 0:
+        usage["web_search_count"] = web_search_count
+
+    return usage
 
 
 def extract_gemini_usage_from_chunk(chunk: Any) -> TokenUsage:
@@ -398,11 +526,19 @@ def extract_gemini_usage_from_chunk(chunk: Any) -> TokenUsage:
 
     usage: TokenUsage = TokenUsage()
 
+    # Extract web search count from the chunk before checking for usage_metadata
+    # Web search indicators can appear on any chunk, not just those with usage data
+    web_search_count = extract_gemini_web_search_count(chunk)
+    if web_search_count > 0:
+        usage["web_search_count"] = web_search_count
+
     if not hasattr(chunk, "usage_metadata") or not chunk.usage_metadata:
         return usage
 
-    # Use the shared helper to extract usage
-    usage = _extract_usage_from_metadata(chunk.usage_metadata)
+    usage_from_metadata = _extract_usage_from_metadata(chunk.usage_metadata)
+
+    # Merge the usage from metadata with any web search count we found
+    usage.update(usage_from_metadata)
 
     return usage