Skip to content

Commit 8ee0569

Browse files
nate-marnate-mar
authored
fix(google-genai): fix wrong tool attribute when setting llm input_message attribute (#2260)
1 parent 8403631 commit 8ee0569

File tree

3 files changed

+86
-34
lines changed

3 files changed

+86
-34
lines changed

‎python/instrumentation/openinference-instrumentation-google-genai/src/openinference/instrumentation/google_genai/_request_attributes_extractor.py‎

Lines changed: 50 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
OpenInferenceSpanKindValues,
1818
SpanAttributes,
1919
ToolAttributes,
20+
ToolCallAttributes,
2021
)
2122

2223
__all__ = ("_RequestAttributesExtractor",)
@@ -322,7 +323,7 @@ def _get_attributes_from_message_param(
322323
elif isinstance(input_contents, Content) or isinstance(input_contents, UserContent):
323324
yield from self._get_attributes_from_content(input_contents)
324325
elif isinstance(input_contents, Part):
325-
yield from self._get_attributes_from_part(input_contents)
326+
yield from self._get_attributes_from_part(input_contents, 0)
326327
else:
327328
# TODO: Implement for File, PIL_Image
328329
logger.exception(f"Unexpected input contents type: {type(input_contents)}")
@@ -345,28 +346,54 @@ def _get_attributes_from_content(
345346
yield from self._flatten_parts(parts)
346347

347348
def _get_attributes_from_function_call(
348-
self, function_call: FunctionCall
349+
self, function_call: FunctionCall, tool_call_index: int
349350
) -> Iterator[Tuple[str, AttributeValue]]:
350351
if name := get_attribute(function_call, "name"):
351352
if isinstance(name, str):
352-
yield (MessageAttributes.MESSAGE_FUNCTION_CALL_NAME, name)
353+
yield (
354+
MessageAttributes.MESSAGE_TOOL_CALLS
355+
+ f".{tool_call_index}."
356+
+ ToolCallAttributes.TOOL_CALL_FUNCTION_NAME,
357+
name,
358+
)
353359
if args := get_attribute(function_call, "args"):
354-
yield (MessageAttributes.MESSAGE_FUNCTION_CALL_ARGUMENTS_JSON, safe_json_dumps(args))
360+
yield (
361+
MessageAttributes.MESSAGE_TOOL_CALLS
362+
+ f".{tool_call_index}."
363+
+ ToolCallAttributes.TOOL_CALL_FUNCTION_ARGUMENTS_JSON,
364+
safe_json_dumps(args),
365+
)
366+
367+
if id := get_attribute(function_call, "id"):
368+
yield (
369+
MessageAttributes.MESSAGE_TOOL_CALLS
370+
+ f".{tool_call_index}."
371+
+ ToolCallAttributes.TOOL_CALL_ID,
372+
id,
373+
)
355374

356375
def _get_attributes_from_function_response(
357376
self, function_response: FunctionResponse
358377
) -> Iterator[Tuple[str, AttributeValue]]:
359378
if response := get_attribute(function_response, "response"):
360379
yield (MessageAttributes.MESSAGE_CONTENT, safe_json_dumps(response))
380+
if id := get_attribute(function_response, "id"):
381+
yield (
382+
MessageAttributes.MESSAGE_TOOL_CALL_ID,
383+
id,
384+
)
361385

362386
def _flatten_parts(self, parts: list[Part]) -> Iterator[Tuple[str, AttributeValue]]:
363387
content_values = []
388+
tool_call_index = 0
364389
for part in parts:
365-
for attr, value in self._get_attributes_from_part(part):
366-
if attr in [
367-
MessageAttributes.MESSAGE_FUNCTION_CALL_NAME,
368-
MessageAttributes.MESSAGE_FUNCTION_CALL_ARGUMENTS_JSON,
369-
]:
390+
for attr, value in self._get_attributes_from_part(part, tool_call_index):
391+
if attr.startswith(MessageAttributes.MESSAGE_TOOL_CALLS):
392+
# Increment tool call index if there happens to be multiple tool calls
393+
# across parts
394+
tool_call_index = self._extract_tool_call_index(attr) + 1
395+
yield (attr, value)
396+
elif attr == MessageAttributes.MESSAGE_TOOL_CALL_ID:
370397
yield (attr, value)
371398
elif isinstance(value, str):
372399
# Flatten all other string values into a single message content
@@ -377,15 +404,27 @@ def _flatten_parts(self, parts: list[Part]) -> Iterator[Tuple[str, AttributeValu
377404
if content_values:
378405
yield (MessageAttributes.MESSAGE_CONTENT, "\n\n".join(content_values))
379406

380-
def _get_attributes_from_part(self, part: Part) -> Iterator[Tuple[str, AttributeValue]]:
407+
def _extract_tool_call_index(self, attr: str) -> int:
408+
"""Extract tool call index from message tool call attribute key.
409+
410+
Example: 'message.tool_calls.0.function_name' -> 0
411+
"""
412+
parts = attr.split(".")
413+
if len(parts) >= 3 and parts[2].isdigit():
414+
return int(parts[2])
415+
return 0
416+
417+
def _get_attributes_from_part(
418+
self, part: Part, tool_call_index: int
419+
) -> Iterator[Tuple[str, AttributeValue]]:
381420
# https://github.com/googleapis/python-genai/blob/main/google/genai/types.py#L566
382421
if text := get_attribute(part, "text"):
383422
yield (
384423
MessageAttributes.MESSAGE_CONTENT,
385424
text,
386425
)
387426
elif function_call := get_attribute(part, "function_call"):
388-
yield from self._get_attributes_from_function_call(function_call)
427+
yield from self._get_attributes_from_function_call(function_call, tool_call_index)
389428
elif function_response := get_attribute(part, "function_response"):
390429
yield from self._get_attributes_from_function_response(function_response)
391430
else:

‎python/instrumentation/openinference-instrumentation-google-genai/tests/cassettes/test_generate_content.yaml‎

Lines changed: 13 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,28 +1,30 @@
11
interactions:
22
- request:
33
body: '{"contents": [{"parts": [{"text": "What''s the weather like?"}], "role":
4-
"user"}, {"parts": [{"functionCall": {"args": {"location": "San Francisco"},
5-
"name": "get_weather"}}], "role": "model"}, {"parts": [{"functionResponse":
6-
{"name": "get_weather", "response": {"location": "San Francisco", "temperature":
7-
65, "unit": "fahrenheit", "condition": "foggy", "humidity": "85%"}}}], "role":
8-
"user"}], "systemInstruction": {"parts": [{"text": "You are a helpful assistant
9-
that can answer questions and help with tasks."}], "role": "user"}, "generationConfig":
10-
{}}'
4+
"user"}, {"parts": [{"functionCall": {"id": "call_abc123", "args": {"location":
5+
"San Francisco"}, "name": "get_weather"}}], "role": "model"}, {"parts": [{"functionResponse":
6+
{"id": "call_abc123", "name": "get_weather", "response": {"location": "San Francisco",
7+
"temperature": 65, "unit": "fahrenheit", "condition": "foggy", "humidity": "85%"}}}],
8+
"role": "user"}], "systemInstruction": {"parts": [{"text": "You are a helpful
9+
assistant that can answer questions and help with tasks."}], "role": "user"},
10+
"generationConfig": {}}'
1111
headers: {}
1212
method: POST
1313
uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent
1414
response:
15-
content: "{\n \"candidates\": [\n {\n \"content\": {\n \"parts\":
15+
body:
16+
string: "{\n \"candidates\": [\n {\n \"content\": {\n \"parts\":
1617
[\n {\n \"text\": \"The weather in San Francisco is foggy
1718
with 85% humidity. The temperature is 65 degrees Fahrenheit.\\n\"\n }\n
1819
\ ],\n \"role\": \"model\"\n },\n \"finishReason\":
19-
\"STOP\",\n \"avgLogprobs\": -0.021388312180836994\n }\n ],\n \"usageMetadata\":
20+
\"STOP\",\n \"avgLogprobs\": -0.026007329424222309\n }\n ],\n \"usageMetadata\":
2021
{\n \"promptTokenCount\": 45,\n \"candidatesTokenCount\": 24,\n \"totalTokenCount\":
2122
69,\n \"promptTokensDetails\": [\n {\n \"modality\": \"TEXT\",\n
2223
\ \"tokenCount\": 45\n }\n ],\n \"candidatesTokensDetails\":
2324
[\n {\n \"modality\": \"TEXT\",\n \"tokenCount\": 24\n
2425
\ }\n ]\n },\n \"modelVersion\": \"gemini-2.0-flash\"\n}\n"
2526
headers: {}
26-
http_version: HTTP/1.1
27-
status_code: 200
27+
status:
28+
code: 200
29+
message: OK
2830
version: 1

‎python/instrumentation/openinference-instrumentation-google-genai/tests/test_instrumentation.py‎

Lines changed: 23 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,9 @@
77
from google import genai
88
from google.genai.types import (
99
Content,
10+
FunctionCall,
1011
FunctionDeclaration,
12+
FunctionResponse,
1113
GenerateContentConfig,
1214
Part,
1315
Tool,
@@ -79,21 +81,28 @@ def test_generate_content(
7981
Content(
8082
role="model",
8183
parts=[
82-
Part.from_function_call(name="get_weather", args={"location": "San Francisco"}),
84+
Part(
85+
function_call=FunctionCall(
86+
name="get_weather", args={"location": "San Francisco"}, id="call_abc123"
87+
)
88+
),
8389
],
8490
),
8591
Content(
8692
role="user",
8793
parts=[
88-
Part.from_function_response(
89-
name="get_weather",
90-
response={
91-
"location": "San Francisco",
92-
"temperature": 65,
93-
"unit": "fahrenheit",
94-
"condition": "foggy",
95-
"humidity": "85%",
96-
},
94+
Part(
95+
function_response=FunctionResponse(
96+
name="get_weather",
97+
response={
98+
"location": "San Francisco",
99+
"temperature": 65,
100+
"unit": "fahrenheit",
101+
"condition": "foggy",
102+
"humidity": "85%",
103+
},
104+
id="call_abc123",
105+
)
97106
),
98107
],
99108
),
@@ -123,10 +132,12 @@ def test_generate_content(
123132
f"{SpanAttributes.LLM_INPUT_MESSAGES}.1.{MessageAttributes.MESSAGE_ROLE}": "user",
124133
f"{SpanAttributes.LLM_INPUT_MESSAGES}.1.{MessageAttributes.MESSAGE_CONTENT}": "What's the weather like?",
125134
f"{SpanAttributes.LLM_INPUT_MESSAGES}.2.{MessageAttributes.MESSAGE_ROLE}": "model",
126-
f"{SpanAttributes.LLM_INPUT_MESSAGES}.2.{MessageAttributes.MESSAGE_FUNCTION_CALL_NAME}": "get_weather",
127-
f"{SpanAttributes.LLM_INPUT_MESSAGES}.2.{MessageAttributes.MESSAGE_FUNCTION_CALL_ARGUMENTS_JSON}": json.dumps(
135+
f"{SpanAttributes.LLM_INPUT_MESSAGES}.2.{MessageAttributes.MESSAGE_TOOL_CALLS}.0.{ToolCallAttributes.TOOL_CALL_FUNCTION_NAME}": "get_weather",
136+
f"{SpanAttributes.LLM_INPUT_MESSAGES}.2.{MessageAttributes.MESSAGE_TOOL_CALLS}.0.{ToolCallAttributes.TOOL_CALL_FUNCTION_ARGUMENTS_JSON}": json.dumps(
128137
{"location": "San Francisco"}
129138
),
139+
f"{SpanAttributes.LLM_INPUT_MESSAGES}.2.{MessageAttributes.MESSAGE_TOOL_CALLS}.0.{ToolCallAttributes.TOOL_CALL_ID}": "call_abc123",
140+
f"{SpanAttributes.LLM_INPUT_MESSAGES}.3.{MessageAttributes.MESSAGE_TOOL_CALL_ID}": "call_abc123",
130141
f"{SpanAttributes.LLM_INPUT_MESSAGES}.3.{MessageAttributes.MESSAGE_ROLE}": "user",
131142
f"{SpanAttributes.LLM_INPUT_MESSAGES}.3.{MessageAttributes.MESSAGE_CONTENT}": json.dumps(
132143
{

0 commit comments

Comments
 (0)