Skip to content

Commit ad00198

Browse files
committed
WIP vertex tool calls
1 parent bf60da4 commit ad00198

File tree

6 files changed

+489
-9
lines changed

6 files changed

+489
-9
lines changed

instrumentation-genai/opentelemetry-instrumentation-vertexai/CHANGELOG.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,5 +15,5 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
1515
([#3208](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3208))
1616
- VertexAI emit user, system, and assistant events
1717
([#3203](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3203))
18-
- Add Vertex gen AI response span attributes
18+
- Add Vertex gen AI response attributes and `gen_ai.choice` events
1919
([#3227](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3227))

instrumentation-genai/opentelemetry-instrumentation-vertexai/src/opentelemetry/instrumentation/vertexai/events.py

Lines changed: 30 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@
2323
from __future__ import annotations
2424

2525
from dataclasses import asdict, dataclass
26-
from typing import Literal
26+
from typing import Any, Iterable, Literal
2727

2828
from opentelemetry._events import Event
2929
from opentelemetry.semconv._incubating.attributes import gen_ai_attributes
@@ -104,36 +104,58 @@ class ChoiceMessage:
104104
role: str = "assistant"
105105

106106

107+
@dataclass
108+
class ChoiceToolCall:
109+
"""The tool_calls field for a gen_ai.choice event"""
110+
111+
@dataclass
112+
class Function:
113+
name: str
114+
arguments: AnyValue = None
115+
116+
function: Function
117+
id: str
118+
type: Literal["function"] = "function"
119+
120+
107121
FinishReason = Literal[
108122
"content_filter", "error", "length", "stop", "tool_calls"
109123
]
110124

111125

112-
# TODO add tool calls
113-
# https://github.com/open-telemetry/opentelemetry-python-contrib/issues/3216
114126
def choice_event(
115127
*,
116128
finish_reason: FinishReason | str,
117129
index: int,
118130
message: ChoiceMessage,
131+
tool_calls: Iterable[ChoiceToolCall] = (),
119132
) -> Event:
120133
"""Creates a choice event, which describes the Gen AI response message.
121134
https://github.com/open-telemetry/semantic-conventions/blob/v1.28.0/docs/gen-ai/gen-ai-events.md#event-gen_aichoice
122135
"""
123136
body: dict[str, AnyValue] = {
124137
"finish_reason": finish_reason,
125138
"index": index,
126-
"message": asdict(
127-
message,
128-
# filter nulls
129-
dict_factory=lambda kvs: {k: v for (k, v) in kvs if v is not None},
130-
),
139+
"message": _asdict_filter_nulls(message),
131140
}
132141

142+
tool_calls_list = [
143+
_asdict_filter_nulls(tool_call) for tool_call in tool_calls
144+
]
145+
if tool_calls_list:
146+
body["tool_calls"] = tool_calls_list
147+
133148
return Event(
134149
name="gen_ai.choice",
135150
attributes={
136151
gen_ai_attributes.GEN_AI_SYSTEM: gen_ai_attributes.GenAiSystemValues.VERTEX_AI.value,
137152
},
138153
body=body,
139154
)
155+
156+
157+
def _asdict_filter_nulls(instance: Any) -> dict[str, AnyValue]:
158+
return asdict(
159+
instance,
160+
dict_factory=lambda kvs: {k: v for (k, v) in kvs if v is not None},
161+
)

instrumentation-genai/opentelemetry-instrumentation-vertexai/src/opentelemetry/instrumentation/vertexai/utils.py

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,9 +26,12 @@
2626
)
2727
from urllib.parse import urlparse
2828

29+
from google.protobuf import json_format
30+
2931
from opentelemetry._events import Event
3032
from opentelemetry.instrumentation.vertexai.events import (
3133
ChoiceMessage,
34+
ChoiceToolCall,
3235
FinishReason,
3336
assistant_event,
3437
choice_event,
@@ -231,6 +234,17 @@ def response_to_events(
231234
capture_content: bool,
232235
) -> Iterable[Event]:
233236
for candidate in response.candidates:
237+
# NOTE: since function_call appears in content.parts, it will also be in the choice
238+
# event. This is different from OpenAI where the tool calls are outside of content:
239+
# https://platform.openai.com/docs/api-reference/chat/object. I would prefer not to
240+
# filter from choice event to keep indexing obvious.
241+
#
242+
# There is similarly a pair of executable_code and
243+
# code_execution_result which are similar to tool call in that the model is asking for
244+
# you to do something rather than generating content:
245+
# https://github.com/googleapis/googleapis/blob/ae87dc8a3830f37d575e2cff577c9b5a4737176b/google/cloud/aiplatform/v1beta1/content.proto#L123-L128
246+
tool_calls = _extract_tool_calls(candidate)
247+
234248
yield choice_event(
235249
finish_reason=_map_finish_reason(candidate.finish_reason),
236250
index=candidate.index,
@@ -242,6 +256,26 @@ def response_to_events(
242256
parts=candidate.content.parts,
243257
),
244258
),
259+
tool_calls=tool_calls,
260+
)
261+
262+
263+
def _extract_tool_calls(
264+
candidate: content.Candidate | content_v1beta1.Candidate,
265+
) -> Iterable[ChoiceToolCall]:
266+
for part in candidate.content.parts:
267+
if "function_call" not in part:
268+
continue
269+
270+
yield ChoiceToolCall(
271+
# NOTE: Vertex does not have an id but this required
272+
id="",
273+
function=ChoiceToolCall.Function(
274+
name=part.function_call.name,
275+
arguments=json_format.MessageToDict(
276+
part.function_call._pb.args # type: ignore[reportUnknownMemberType]
277+
),
278+
),
245279
)
246280

247281

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,119 @@
1+
interactions:
2+
- request:
3+
body: |-
4+
{
5+
"contents": [
6+
{
7+
"role": "user",
8+
"parts": [
9+
{
10+
"text": "Get weather details in New Delhi and San Francisco?"
11+
}
12+
]
13+
}
14+
],
15+
"tools": [
16+
{
17+
"functionDeclarations": [
18+
{
19+
"name": "get_current_weather",
20+
"description": "Get the current weather in a given location",
21+
"parameters": {
22+
"type": 6,
23+
"properties": {
24+
"location": {
25+
"type": 1,
26+
"description": "The location for which to get the weather. It can be a city name, a city name and state, or a zip code. Examples: 'San Francisco', 'San Francisco, CA', '95616', etc."
27+
}
28+
},
29+
"propertyOrdering": [
30+
"location"
31+
]
32+
}
33+
}
34+
]
35+
}
36+
]
37+
}
38+
headers:
39+
Accept:
40+
- '*/*'
41+
Accept-Encoding:
42+
- gzip, deflate
43+
Connection:
44+
- keep-alive
45+
Content-Length:
46+
- '824'
47+
Content-Type:
48+
- application/json
49+
User-Agent:
50+
- python-requests/2.32.3
51+
method: POST
52+
uri: https://us-central1-aiplatform.googleapis.com/v1/projects/fake-project/locations/us-central1/publishers/google/models/gemini-1.5-flash-002:generateContent?%24alt=json%3Benum-encoding%3Dint
53+
response:
54+
body:
55+
string: |-
56+
{
57+
"candidates": [
58+
{
59+
"content": {
60+
"role": "model",
61+
"parts": [
62+
{
63+
"functionCall": {
64+
"name": "get_current_weather",
65+
"args": {
66+
"location": "New Delhi"
67+
}
68+
}
69+
},
70+
{
71+
"functionCall": {
72+
"name": "get_current_weather",
73+
"args": {
74+
"location": "San Francisco"
75+
}
76+
}
77+
}
78+
]
79+
},
80+
"finishReason": 1,
81+
"avgLogprobs": -0.00018195244774688035
82+
}
83+
],
84+
"usageMetadata": {
85+
"promptTokenCount": 72,
86+
"candidatesTokenCount": 16,
87+
"totalTokenCount": 88,
88+
"promptTokensDetails": [
89+
{
90+
"modality": 1,
91+
"tokenCount": 72
92+
}
93+
],
94+
"candidatesTokensDetails": [
95+
{
96+
"modality": 1,
97+
"tokenCount": 16
98+
}
99+
]
100+
},
101+
"modelVersion": "gemini-1.5-flash-002",
102+
"createTime": "2025-02-04T03:52:01.562155Z",
103+
"responseId": "4Y6hZ-unIrfQnvgPr73XgA0"
104+
}
105+
headers:
106+
Content-Type:
107+
- application/json; charset=UTF-8
108+
Transfer-Encoding:
109+
- chunked
110+
Vary:
111+
- Origin
112+
- X-Origin
113+
- Referer
114+
content-length:
115+
- '1029'
116+
status:
117+
code: 200
118+
message: OK
119+
version: 1
Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,119 @@
1+
interactions:
2+
- request:
3+
body: |-
4+
{
5+
"contents": [
6+
{
7+
"role": "user",
8+
"parts": [
9+
{
10+
"text": "Get weather details in New Delhi and San Francisco?"
11+
}
12+
]
13+
}
14+
],
15+
"tools": [
16+
{
17+
"functionDeclarations": [
18+
{
19+
"name": "get_current_weather",
20+
"description": "Get the current weather in a given location",
21+
"parameters": {
22+
"type": 6,
23+
"properties": {
24+
"location": {
25+
"type": 1,
26+
"description": "The location for which to get the weather. It can be a city name, a city name and state, or a zip code. Examples: 'San Francisco', 'San Francisco, CA', '95616', etc."
27+
}
28+
},
29+
"propertyOrdering": [
30+
"location"
31+
]
32+
}
33+
}
34+
]
35+
}
36+
]
37+
}
38+
headers:
39+
Accept:
40+
- '*/*'
41+
Accept-Encoding:
42+
- gzip, deflate
43+
Connection:
44+
- keep-alive
45+
Content-Length:
46+
- '824'
47+
Content-Type:
48+
- application/json
49+
User-Agent:
50+
- python-requests/2.32.3
51+
method: POST
52+
uri: https://us-central1-aiplatform.googleapis.com/v1/projects/fake-project/locations/us-central1/publishers/google/models/gemini-1.5-flash-002:generateContent?%24alt=json%3Benum-encoding%3Dint
53+
response:
54+
body:
55+
string: |-
56+
{
57+
"candidates": [
58+
{
59+
"content": {
60+
"role": "model",
61+
"parts": [
62+
{
63+
"functionCall": {
64+
"name": "get_current_weather",
65+
"args": {
66+
"location": "New Delhi"
67+
}
68+
}
69+
},
70+
{
71+
"functionCall": {
72+
"name": "get_current_weather",
73+
"args": {
74+
"location": "San Francisco"
75+
}
76+
}
77+
}
78+
]
79+
},
80+
"finishReason": 1,
81+
"avgLogprobs": -0.00018187805835623294
82+
}
83+
],
84+
"usageMetadata": {
85+
"promptTokenCount": 72,
86+
"candidatesTokenCount": 16,
87+
"totalTokenCount": 88,
88+
"promptTokensDetails": [
89+
{
90+
"modality": 1,
91+
"tokenCount": 72
92+
}
93+
],
94+
"candidatesTokensDetails": [
95+
{
96+
"modality": 1,
97+
"tokenCount": 16
98+
}
99+
]
100+
},
101+
"modelVersion": "gemini-1.5-flash-002",
102+
"createTime": "2025-02-04T22:05:24.282269Z",
103+
"responseId": "JI-iZ52dEeWI3NoP39Od0Qg"
104+
}
105+
headers:
106+
Content-Type:
107+
- application/json; charset=UTF-8
108+
Transfer-Encoding:
109+
- chunked
110+
Vary:
111+
- Origin
112+
- X-Origin
113+
- Referer
114+
content-length:
115+
- '1029'
116+
status:
117+
code: 200
118+
message: OK
119+
version: 1

0 commit comments

Comments
 (0)