From 3cbbc2dcb38ccd5ce2c5a5197b89f705e15afcde Mon Sep 17 00:00:00 2001 From: mo Date: Thu, 29 Oct 2026 16:07:32 +0000 Subject: [PATCH] fix(ai): preserve Gemini tool call IDs in captured generations Gemini returns IDs on function calls and function responses, but the captured $ai_output_choices and $ai_input dropped them, so tool results could not be matched to their calls. Keep the call ID on output (streaming and non-streaming) and on input history calls, and the result's ID as tool_use_id, matching PostHog/posthog-js#5087. --- .sampo/changesets/gemini-tool-call-ids.md | 5 + posthog/ai/gemini/gemini_converter.py | 59 +++++++---- posthog/test/ai/gemini/test_gemini.py | 122 ++++++++++++++++++++++ 3 files changed, 164 insertions(+), 22 deletions(-) create mode 100644 .sampo/changesets/gemini-tool-call-ids.md diff --git a/.sampo/changesets/gemini-tool-call-ids.md b/.sampo/changesets/gemini-tool-call-ids.md new file mode 100644 index 000000000..e1716fadd --- /dev/null +++ b/.sampo/changesets/gemini-tool-call-ids.md @@ -0,0 +1,5 @@ +--- +pypi/posthog: patch +--- + +Preserve Gemini tool call IDs in captured generations, on output tool calls and on tool calls and results in the input history. diff --git a/posthog/ai/gemini/gemini_converter.py b/posthog/ai/gemini/gemini_converter.py index ae4522746..09ae82d25 100644 --- a/posthog/ai/gemini/gemini_converter.py +++ b/posthog/ai/gemini/gemini_converter.py @@ -40,6 +40,18 @@ def _format_media_payload(payload: Any) -> Dict[str, Any]: return payload +def _function_call_block(name: Any, arguments: Any, call_id: Any) -> Dict[str, Any]: + block: Dict[str, Any] = { + "type": "function", + "function": {"name": name, "arguments": arguments}, + } + # Only a real string ID: Gemini omits it for some calls, and mocked parts + # answer any attribute. + if isinstance(call_id, str): + block["id"] = call_id + return block + + def _format_part(part: Any) -> Optional[FormattedContentItem]: if isinstance(part, str): return {"type": "text", "text": part} @@ -64,20 +76,26 @@ def _format_part(part: Any) -> Optional[FormattedContentItem]: # shape leaves them unrenderable and invisible to evals. if "function_call" in plain: call = to_plain(plain["function_call"]) or {} - return { + block: Dict[str, Any] = { "type": "function", "function": { "name": call.get("name"), "arguments": call.get("args"), }, } + if isinstance(call.get("id"), str): + block["id"] = call["id"] + return block if "function_response" in plain: response = to_plain(plain["function_response"]) or {} - return { + response_block: Dict[str, Any] = { "type": "function", "tool_name": response.get("name") or "", "content": response.get("response"), } + if isinstance(response.get("id"), str): + response_block["tool_use_id"] = response["id"] + return response_block if not plain: return None key = next(iter(plain)) @@ -269,13 +287,11 @@ def format_gemini_response(response: Any) -> List[FormattedMessage]: if hasattr(part, "function_call") and part.function_call: function_call = part.function_call content.append( - { - "type": "function", - "function": { - "name": function_call.name, - "arguments": function_call.args, - }, - } + _function_call_block( + function_call.name, + function_call.args, + getattr(function_call, "id", None), + ) ) continue @@ -622,13 +638,11 @@ def extract_gemini_content_from_chunk( if hasattr(part, "function_call") and part.function_call: function_call = part.function_call blocks.append( - { - "type": "function", - "function": { - "name": function_call.name, - "arguments": function_call.args, - }, - } + _function_call_block( + function_call.name, + function_call.args, + getattr(function_call, "id", None), + ) ) continue @@ -689,12 +703,13 @@ def format_gemini_streaming_output( text_parts = [] # Add the function call - content.append( - { - "type": "function", - "function": item.get("function", {}), - } - ) + function_block: Dict[str, Any] = { + "type": "function", + "function": item.get("function", {}), + } + if isinstance(item.get("id"), str): + function_block["id"] = item["id"] + content.append(function_block) else: if text_parts: content.append( diff --git a/posthog/test/ai/gemini/test_gemini.py b/posthog/test/ai/gemini/test_gemini.py index f2a801e68..b616dcd89 100644 --- a/posthog/test/ai/gemini/test_gemini.py +++ b/posthog/test/ai/gemini/test_gemini.py @@ -797,6 +797,128 @@ def test_function_calls_only_no_content( assert props["$ai_http_status"] == 200 +def _function_call_response(call_id): + from google.genai import types + + return types.GenerateContentResponse( + candidates=[ + types.Candidate( + content=types.Content( + role="model", + parts=[ + types.Part( + function_call=types.FunctionCall( + id=call_id, name="describe_shape", args={"sides": 3} + ) + ) + ], + ), + finish_reason="STOP", + ) + ], + usage_metadata=types.GenerateContentResponseUsageMetadata( + prompt_token_count=10, candidates_token_count=5, total_token_count=15 + ), + ) + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("call_id", ["call_649034", None]) +def test_function_call_ids_in_output_choices( + mock_client, mock_google_genai_client, stream, call_id +): + """Gemini's function-call ID is kept in $ai_output_choices when present.""" + response = _function_call_response(call_id) + client = Client(api_key="test-key", posthog_client=mock_client) + + if stream: + mock_google_genai_client.models.generate_content_stream.return_value = iter( + [response] + ) + list( + client.models.generate_content_stream( + model="gemini-2.5-flash", + contents="Describe a triangle", + posthog_distinct_id="test-id", + ) + ) + else: + mock_google_genai_client.models.generate_content.return_value = response + client.models.generate_content( + model="gemini-2.5-flash", + contents="Describe a triangle", + posthog_distinct_id="test-id", + ) + + expected_call = { + "type": "function", + "function": {"name": "describe_shape", "arguments": {"sides": 3}}, + } + if call_id is not None: + expected_call["id"] = call_id + props = mock_client.capture.call_args[1]["properties"] + assert props["$ai_output_choices"] == [ + {"role": "assistant", "content": [expected_call]} + ] + + +def test_function_call_ids_in_input_history( + mock_client, mock_google_genai_client, mock_gemini_response +): + """IDs on function_call and function_response parts in the input history are + kept, so a tool result can be matched to its call.""" + from google.genai import types + + mock_google_genai_client.models.generate_content.return_value = mock_gemini_response + client = Client(api_key="test-key", posthog_client=mock_client) + client.models.generate_content( + model="gemini-2.5-flash", + contents=[ + types.Content(role="user", parts=[types.Part(text="Describe a triangle")]), + types.Content( + role="model", + parts=[ + types.Part( + function_call=types.FunctionCall( + id="call_649034", name="describe_shape", args={"sides": 3} + ) + ) + ], + ), + types.Content( + role="user", + parts=[ + types.Part( + function_response=types.FunctionResponse( + id="call_649034", + name="describe_shape", + response={"shape": "triangle"}, + ) + ) + ], + ), + ], + posthog_distinct_id="test-id", + ) + + props = mock_client.capture.call_args[1]["properties"] + call_block, result_block = [ + block + for message in props["$ai_input"] + for block in ( + message["content"] if isinstance(message["content"], list) else [] + ) + if block.get("type") == "function" + ] + assert call_block["id"] == "call_649034" + assert call_block["function"] == { + "name": "describe_shape", + "arguments": {"sides": 3}, + } + assert result_block["tool_use_id"] == "call_649034" + assert result_block["tool_name"] == "describe_shape" + + def test_cache_and_reasoning_tokens(mock_client, mock_google_genai_client): """Test that cache and reasoning tokens are properly extracted""" # Create a mock response with cache and reasoning tokens