Skip to content

fix(ai): preserve Gemini tool call IDs in captured generations #1150

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
marandaneto merged 2 commits into from
Oct 30, 2026
Merged

fix(ai): preserve Gemini tool call IDs in captured generations #1150

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3cbbc2d
fix(ai): preserve Gemini tool call IDs in captured generations
mosafariuk Oct 29, 2026
ea15bb3
Merge remote-tracking branch 'origin/main' into fix/gemini-tool-call-ids
marandaneto Oct 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions .sampo/changesets/gemini-tool-call-ids.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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.
59 changes: 37 additions & 22 deletions posthog/ai/gemini/gemini_converter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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}
Expand All @@ -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))
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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(
Expand Down
122 changes: 122 additions & 0 deletions posthog/test/ai/gemini/test_gemini.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading