fix(llma): fix Gemini

Author: Radu-Raicea

posthog/ai/gemini/gemini_converter.py

@@ -343,7 +343,7 @@ 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, 0 otherwise.
+    Returns 1 if grounding_metadata is present with actual search data, 0 otherwise.
 
     Args:
         response: The response from Gemini API
@@ -359,7 +359,21 @@ def extract_gemini_web_search_count(response: Any) -> int:
                 hasattr(candidate, "grounding_metadata")
                 and candidate.grounding_metadata
             ):
-                return 1
+                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:
@@ -369,6 +383,7 @@ def extract_gemini_web_search_count(response: Any) -> int:
                             function_name = getattr(
                                 part.function_call, "name", ""
                             ).lower()
+
                             if (
                                 "google_search" in function_name
                                 or "grounding" in function_name

posthog/test/ai/gemini/test_gemini.py

@@ -1114,3 +1114,130 @@ async def mock_async_streaming_response():
     assert props["$ai_web_search_count"] == 1
     assert props["$ai_input_tokens"] == 30
     assert props["$ai_output_tokens"] == 15
+
+
+def test_empty_grounding_metadata_no_web_search(mock_client, mock_google_genai_client):
+    """Test that empty grounding_metadata (all null fields) does not count as web search."""
+
+    # Create mock response with empty grounding metadata (all null fields)
+    mock_response = MagicMock()
+
+    # Mock usage metadata
+    mock_usage = MagicMock()
+    mock_usage.prompt_token_count = 10
+    mock_usage.candidates_token_count = 10
+    mock_usage.cached_content_token_count = 0
+    mock_usage.thoughts_token_count = 0
+    mock_response.usage_metadata = mock_usage
+
+    # Mock empty grounding metadata (all fields are None)
+    mock_grounding_metadata = MagicMock()
+    mock_grounding_metadata.web_search_queries = None
+    mock_grounding_metadata.grounding_chunks = None
+    mock_grounding_metadata.grounding_supports = None
+    mock_grounding_metadata.retrieval_metadata = None
+    mock_grounding_metadata.retrieval_queries = None
+    mock_grounding_metadata.search_entry_point = None
+
+    # Mock text part
+    mock_text_part = MagicMock()
+    mock_text_part.text = "Hey there! How can I help you today?"
+    type(mock_text_part).text = mock_text_part.text
+
+    # Mock content with parts
+    mock_content = MagicMock()
+    mock_content.parts = [mock_text_part]
+
+    # Mock candidate with empty grounding metadata
+    mock_candidate = MagicMock()
+    mock_candidate.content = mock_content
+    mock_candidate.grounding_metadata = mock_grounding_metadata
+    type(mock_candidate).grounding_metadata = mock_candidate.grounding_metadata
+
+    mock_response.candidates = [mock_candidate]
+    mock_response.text = "Hey there! How can I help you today?"
+
+    # Mock the generate_content method
+    mock_google_genai_client.models.generate_content.return_value = mock_response
+
+    client = Client(api_key="test-key", posthog_client=mock_client)
+
+    response = client.models.generate_content(
+        model="gemini-2.5-flash",
+        contents="Hello",
+        posthog_distinct_id="test-id",
+    )
+
+    assert response == mock_response
+    assert mock_client.capture.call_count == 1
+
+    call_args = mock_client.capture.call_args[1]
+    props = call_args["properties"]
+
+    # Verify web search count is 0 (not present in properties when 0)
+    assert "$ai_web_search_count" not in props
+    assert props["$ai_input_tokens"] == 10
+    assert props["$ai_output_tokens"] == 10
+
+
+def test_empty_array_grounding_metadata_no_web_search(
+    mock_client, mock_google_genai_client
+):
+    """Test that grounding_metadata with empty arrays does not count as web search."""
+
+    # Create mock response with grounding metadata having empty arrays
+    mock_response = MagicMock()
+
+    # Mock usage metadata
+    mock_usage = MagicMock()
+    mock_usage.prompt_token_count = 15
+    mock_usage.candidates_token_count = 12
+    mock_usage.cached_content_token_count = 0
+    mock_usage.thoughts_token_count = 0
+    mock_response.usage_metadata = mock_usage
+
+    # Mock grounding metadata with empty arrays
+    mock_grounding_metadata = MagicMock()
+    mock_grounding_metadata.web_search_queries = []
+    mock_grounding_metadata.grounding_chunks = []
+    mock_grounding_metadata.grounding_supports = []
+
+    # Mock text part
+    mock_text_part = MagicMock()
+    mock_text_part.text = "I can help with that."
+    type(mock_text_part).text = mock_text_part.text
+
+    # Mock content with parts
+    mock_content = MagicMock()
+    mock_content.parts = [mock_text_part]
+
+    # Mock candidate with grounding metadata containing empty arrays
+    mock_candidate = MagicMock()
+    mock_candidate.content = mock_content
+    mock_candidate.grounding_metadata = mock_grounding_metadata
+    type(mock_candidate).grounding_metadata = mock_candidate.grounding_metadata
+
+    mock_response.candidates = [mock_candidate]
+    mock_response.text = "I can help with that."
+
+    # Mock the generate_content method
+    mock_google_genai_client.models.generate_content.return_value = mock_response
+
+    client = Client(api_key="test-key", posthog_client=mock_client)
+
+    response = client.models.generate_content(
+        model="gemini-2.5-flash",
+        contents="What can you do?",
+        posthog_distinct_id="test-id",
+    )
+
+    assert response == mock_response
+    assert mock_client.capture.call_count == 1
+
+    call_args = mock_client.capture.call_args[1]
+    props = call_args["properties"]
+
+    # Verify web search count is 0 (not present in properties when 0)
+    assert "$ai_web_search_count" not in props
+    assert props["$ai_input_tokens"] == 15
+    assert props["$ai_output_tokens"] == 12