Skip to content

Commit 5dd2a21

Browse files
DylanRussellDylanRussell
committed
Update instrumentation to emit spans in the new format, call the upload hook
1 parent 7e139ac commit 5dd2a21

File tree

12 files changed

+625
-87
lines changed

12 files changed

+625
-87
lines changed

‎cassettes/test_tool_events_no_content‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
interactions:
2+
- request:
3+
body: "{\n \"contents\": [\n {\n \"role\": \"user\",\n \"parts\":
4+
[\n {\n \"text\": \"Get weather details in New Delhi and San
5+
Francisco?\"\n }\n ]\n },\n {\n \"role\": \"model\",\n
6+
\ \"parts\": [\n {\n \"functionCall\": {\n \"name\":
7+
\"get_current_weather\",\n \"args\": {\n \"location\":
8+
\"New Delhi\"\n }\n }\n },\n {\n \"functionCall\":
9+
{\n \"name\": \"get_current_weather\",\n \"args\": {\n
10+
\ \"location\": \"San Francisco\"\n }\n }\n
11+
\ }\n ]\n },\n {\n \"role\": \"user\",\n \"parts\":
12+
[\n {\n \"functionResponse\": {\n \"name\": \"get_current_weather\",\n
13+
\ \"response\": {\n \"content\": \"{\\\"temperature\\\":
14+
35, \\\"unit\\\": \\\"C\\\"}\"\n }\n }\n },\n {\n
15+
\ \"functionResponse\": {\n \"name\": \"get_current_weather\",\n
16+
\ \"response\": {\n \"content\": \"{\\\"temperature\\\":
17+
25, \\\"unit\\\": \\\"C\\\"}\"\n }\n }\n }\n ]\n
18+
\ }\n ],\n \"tools\": [\n {\n \"functionDeclarations\": [\n {\n
19+
\ \"name\": \"get_current_weather\",\n \"description\": \"Get
20+
the current weather in a given location\",\n \"parameters\": {\n \"type\":
21+
6,\n \"properties\": {\n \"location\": {\n \"type\":
22+
1,\n \"description\": \"The location for which to get the weather.
23+
It can be a city name, a city name and state, or a zip code. Examples: 'San
24+
Francisco', 'San Francisco, CA', '95616', etc.\"\n }\n },\n
25+
\ \"propertyOrdering\": [\n \"location\"\n ]\n
26+
\ }\n }\n ]\n }\n ]\n}"
27+
headers:
28+
Accept:
29+
- '*/*'
30+
Accept-Encoding:
31+
- gzip, deflate
32+
Connection:
33+
- keep-alive
34+
Content-Length:
35+
- '1731'
36+
Content-Type:
37+
- application/json
38+
User-Agent:
39+
- python-requests/2.32.3
40+
x-goog-api-client:
41+
- model-builder/1.79.0+top_google_constructor_method+vertexai.generative_models.GenerativeModel.generate_content
42+
gl-python/3.11.9 grpc/1.68.1 gax/2.23.0 gapic/1.79.0+top_google_constructor_method+vertexai.generative_models.GenerativeModel.generate_content
43+
x-goog-request-params:
44+
- model=projects/fake-project/locations/us-central1/publishers/google/models/gemini-2.5-pro
45+
method: POST
46+
uri: https://us-central1-aiplatform.googleapis.com/v1/projects/fake-project/locations/us-central1/publishers/google/models/gemini-2.5-pro:generateContent?%24alt=json%3Benum-encoding%3Dint
47+
response:
48+
body:
49+
string: !!binary |
50+
H4sIAAAAAAAC/91TwU7rMBC89ytWPRcHECdOVG3EqwQFkXJCCBlnSRccb7A3BYT679ghQODAgeM7
51+
OTueHc+OndcRwBi9Zz8+hNdYxNJwibE62N2bvAM1hqCrhI0v8LHFIEABagqBXAU+QuSxBN3KGp2Q
52+
0ULswEQsldoqyJ8bNBI5Z9NIgn3QxkRREH5ANwHLFcUG5gdCYA8chTxstKVfVQtEWIs04TDLStyg
53+
5QZ9UBVzZVEZrjPquPKSBarcDrnsCW8T1bALbHGn8XwfjalxP2kQLW1Ig14up5erf/lytZhNV/n8
54+
g1CiaLKJcdUB0IfWbR7JS9OllNbehm4odFZ6V74xao63bbVwd9yrDpT/54z7WbeTvyeXp5f6MzmP
55+
OrBL3bOLfJ6ubHpS3JwuimKxPP4WMdeaOuL3E4acOl5DqUV//g5f+JrLr15lLLel0tRYLXfsa7XZ
56+
U+cxODIpwwL9hgyqY3ToteCMncSYBiel1/ZOSqIDnR/ePhu2o+F6PUpf29EbabJOVcEDAAA=
57+
headers:
58+
Content-Encoding:
59+
- gzip
60+
Content-Type:
61+
- application/json; charset=UTF-8
62+
Date:
63+
- Tue, 30 Sep 2025 15:26:53 GMT
64+
Server:
65+
- scaffolding on HTTPServer2
66+
Transfer-Encoding:
67+
- chunked
68+
Vary:
69+
- Origin
70+
- X-Origin
71+
- Referer
72+
WWW-Authenticate:
73+
- Bearer realm="https://accounts.google.com/"
74+
X-Content-Type-Options:
75+
- nosniff
76+
X-Frame-Options:
77+
- SAMEORIGIN
78+
X-XSS-Protection:
79+
- '0'
80+
status:
81+
code: 401
82+
message: Unauthorized
83+
version: 1

‎cassettes/test_tool_events_no_content.yaml‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
interactions:
2+
- request:
3+
body: "{\n \"contents\": [\n {\n \"role\": \"user\",\n \"parts\":
4+
[\n {\n \"text\": \"Get weather details in New Delhi and San
5+
Francisco?\"\n }\n ]\n },\n {\n \"role\": \"model\",\n
6+
\ \"parts\": [\n {\n \"functionCall\": {\n \"name\":
7+
\"get_current_weather\",\n \"args\": {\n \"location\":
8+
\"New Delhi\"\n }\n }\n },\n {\n \"functionCall\":
9+
{\n \"name\": \"get_current_weather\",\n \"args\": {\n
10+
\ \"location\": \"San Francisco\"\n }\n }\n
11+
\ }\n ]\n },\n {\n \"role\": \"user\",\n \"parts\":
12+
[\n {\n \"functionResponse\": {\n \"name\": \"get_current_weather\",\n
13+
\ \"response\": {\n \"content\": \"{\\\"temperature\\\":
14+
35, \\\"unit\\\": \\\"C\\\"}\"\n }\n }\n },\n {\n
15+
\ \"functionResponse\": {\n \"name\": \"get_current_weather\",\n
16+
\ \"response\": {\n \"content\": \"{\\\"temperature\\\":
17+
25, \\\"unit\\\": \\\"C\\\"}\"\n }\n }\n }\n ]\n
18+
\ }\n ],\n \"tools\": [\n {\n \"functionDeclarations\": [\n {\n
19+
\ \"name\": \"get_current_weather\",\n \"description\": \"Get
20+
the current weather in a given location\",\n \"parameters\": {\n \"type\":
21+
6,\n \"properties\": {\n \"location\": {\n \"type\":
22+
1,\n \"description\": \"The location for which to get the weather.
23+
It can be a city name, a city name and state, or a zip code. Examples: 'San
24+
Francisco', 'San Francisco, CA', '95616', etc.\"\n }\n },\n
25+
\ \"propertyOrdering\": [\n \"location\"\n ]\n
26+
\ }\n }\n ]\n }\n ]\n}"
27+
headers:
28+
Accept:
29+
- '*/*'
30+
Accept-Encoding:
31+
- gzip, deflate
32+
Connection:
33+
- keep-alive
34+
Content-Length:
35+
- '1731'
36+
Content-Type:
37+
- application/json
38+
User-Agent:
39+
- python-requests/2.32.3
40+
x-goog-api-client:
41+
- model-builder/1.79.0+top_google_constructor_method+vertexai.generative_models.GenerativeModel.generate_content
42+
gl-python/3.11.9 grpc/1.68.1 gax/2.23.0 gapic/1.79.0+top_google_constructor_method+vertexai.generative_models.GenerativeModel.generate_content
43+
x-goog-request-params:
44+
- model=projects/fake-project/locations/us-central1/publishers/google/models/gemini-2.5-pro
45+
method: POST
46+
uri: https://us-central1-aiplatform.googleapis.com/v1/projects/fake-project/locations/us-central1/publishers/google/models/gemini-2.5-pro:generateContent?%24alt=json%3Benum-encoding%3Dint
47+
response:
48+
body:
49+
string: !!binary |
50+
H4sIAAAAAAAC/91TwU7rMBC89ytWPRcHECdOVG3EqwQFkXJCCBlnSRccb7A3BYT679ghQODAgeM7
51+
OTueHc+OndcRwBi9Zz8+hNdYxNJwibE62N2bvAM1hqCrhI0v8LHFIEABagqBXAU+QuSxBN3KGp2Q
52+
0ULswEQsldoqyJ8bNBI5Z9NIgn3QxkRREH5ANwHLFcUG5gdCYA8chTxstKVfVQtEWIs04TDLStyg
53+
5QZ9UBVzZVEZrjPquPKSBarcDrnsCW8T1bALbHGn8XwfjalxP2kQLW1Ig14up5erf/lytZhNV/n8
54+
g1CiaLKJcdUB0IfWbR7JS9OllNbehm4odFZ6V74xao63bbVwd9yrDpT/54z7WbeTvyeXp5f6MzmP
55+
OrBL3bOLfJ6ubHpS3JwuimKxPP4WMdeaOuL3E4acOl5DqUV//g5f+JrLr15lLLel0tRYLXfsa7XZ
56+
U+cxODIpwwL9hgyqY3ToteCMncSYBiel1/ZOSqIDnR/ePhu2o+F6PUpf29EbabJOVcEDAAA=
57+
headers:
58+
Content-Encoding:
59+
- gzip
60+
Content-Type:
61+
- application/json; charset=UTF-8
62+
Date:
63+
- Tue, 30 Sep 2025 15:27:34 GMT
64+
Server:
65+
- scaffolding on HTTPServer2
66+
Transfer-Encoding:
67+
- chunked
68+
Vary:
69+
- Origin
70+
- X-Origin
71+
- Referer
72+
WWW-Authenticate:
73+
- Bearer realm="https://accounts.google.com/"
74+
X-Content-Type-Options:
75+
- nosniff
76+
X-Frame-Options:
77+
- SAMEORIGIN
78+
X-XSS-Protection:
79+
- '0'
80+
status:
81+
code: 401
82+
message: Unauthorized
83+
version: 1

‎instrumentation-genai/opentelemetry-instrumentation-vertexai/src/opentelemetry/instrumentation/vertexai/__init__.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,7 @@
6060
from opentelemetry.instrumentation.vertexai.utils import is_content_enabled
6161
from opentelemetry.semconv.schemas import Schemas
6262
from opentelemetry.trace import get_tracer
63+
from opentelemetry.util.genai.completion_hook import load_completion_hook
6364

6465

6566
def _methods_to_wrap(
@@ -109,6 +110,9 @@ def instrumentation_dependencies(self) -> Collection[str]:
109110

110111
def _instrument(self, **kwargs: Any):
111112
"""Enable VertexAI instrumentation."""
113+
completion_hook = (
114+
kwargs.get("completion_hook") or load_completion_hook()
115+
)
112116
sem_conv_opt_in_mode = _OpenTelemetrySemanticConventionStability._get_opentelemetry_stability_opt_in_mode(
113117
_OpenTelemetryStabilitySignalType.GEN_AI,
114118
)
@@ -141,6 +145,7 @@ def _instrument(self, **kwargs: Any):
141145
event_logger,
142146
is_content_enabled(sem_conv_opt_in_mode),
143147
sem_conv_opt_in_mode,
148+
completion_hook,
144149
)
145150
elif sem_conv_opt_in_mode == _StabilityMode.GEN_AI_LATEST_EXPERIMENTAL:
146151
# Type checker now knows it's the other literal
@@ -149,6 +154,7 @@ def _instrument(self, **kwargs: Any):
149154
event_logger,
150155
is_content_enabled(sem_conv_opt_in_mode),
151156
sem_conv_opt_in_mode,
157+
completion_hook,
152158
)
153159
else:
154160
raise RuntimeError(f"{sem_conv_opt_in_mode} mode not supported")

‎instrumentation-genai/opentelemetry-instrumentation-vertexai/src/opentelemetry/instrumentation/vertexai/patch.py‎

Lines changed: 77 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,9 @@
1414

1515
from __future__ import annotations
1616

17+
import json
1718
from contextlib import contextmanager
19+
from dataclasses import asdict
1820
from typing import (
1921
TYPE_CHECKING,
2022
Any,
@@ -27,13 +29,14 @@
2729
overload,
2830
)
2931

30-
from opentelemetry._events import EventLogger
32+
from opentelemetry._events import Event, EventLogger
3133
from opentelemetry.instrumentation._semconv import (
3234
_StabilityMode,
3335
)
3436
from opentelemetry.instrumentation.vertexai.utils import (
3537
GenerateContentParams,
36-
create_operation_details_event,
38+
convert_content_to_message,
39+
convert_response_to_output_messages,
3740
get_genai_request_attributes,
3841
get_genai_response_attributes,
3942
get_server_attributes,
@@ -45,7 +48,11 @@
4548
gen_ai_attributes as GenAI,
4649
)
4750
from opentelemetry.trace import SpanKind, Tracer
48-
from opentelemetry.util.genai.types import ContentCapturingMode
51+
from opentelemetry.util.genai.completion_hook import CompletionHook
52+
from opentelemetry.util.genai.types import (
53+
ContentCapturingMode,
54+
Text,
55+
)
4956

5057
if TYPE_CHECKING:
5158
from google.cloud.aiplatform_v1.services.prediction_service import client
@@ -113,6 +120,7 @@ def __init__(
113120
sem_conv_opt_in_mode: Literal[
114121
_StabilityMode.GEN_AI_LATEST_EXPERIMENTAL
115122
],
123+
completion_hook: CompletionHook,
116124
) -> None: ...
117125

118126
@overload
@@ -122,6 +130,7 @@ def __init__(
122130
event_logger: EventLogger,
123131
capture_content: bool,
124132
sem_conv_opt_in_mode: Literal[_StabilityMode.DEFAULT],
133+
completion_hook: CompletionHook,
125134
) -> None: ...
126135

127136
def __init__(
@@ -133,11 +142,13 @@ def __init__(
133142
Literal[_StabilityMode.DEFAULT],
134143
Literal[_StabilityMode.GEN_AI_LATEST_EXPERIMENTAL],
135144
],
145+
completion_hook: CompletionHook,
136146
) -> None:
137147
self.tracer = tracer
138148
self.event_logger = event_logger
139149
self.capture_content = capture_content
140150
self.sem_conv_opt_in_mode = sem_conv_opt_in_mode
151+
self.completion_hook = completion_hook
141152

142153
@contextmanager
143154
def _with_new_instrumentation(
@@ -149,11 +160,9 @@ def _with_new_instrumentation(
149160
kwargs: Any,
150161
):
151162
params = _extract_params(*args, **kwargs)
152-
api_endpoint: str = instance.api_endpoint # type: ignore[reportUnknownMemberType]
153163
request_attributes = get_genai_request_attributes(True, params)
154-
server_attributes = get_server_attributes(api_endpoint)
155164
with self.tracer.start_as_current_span(
156-
name=f"{GenAI.GenAiOperationNameValues.CHAT.value} {request_attributes.get(GenAI.GEN_AI_REQUEST_MODEL, '')}",
165+
name=f"{GenAI.GenAiOperationNameValues.CHAT.value} {request_attributes.get(GenAI.GEN_AI_REQUEST_MODEL, '')}".strip(),
157166
kind=SpanKind.CLIENT,
158167
) as span:
159168

@@ -162,30 +171,69 @@ def handle_response(
162171
| prediction_service_v1beta1.GenerateContentResponse
163172
| None,
164173
) -> None:
165-
response_attributes = (
166-
{}
167-
if not response
168-
else get_genai_response_attributes(response)
174+
attributes = (
175+
get_server_attributes(instance.api_endpoint) # type: ignore[reportUnknownMemberType]
176+
| request_attributes
177+
| get_genai_response_attributes(response)
169178
)
170-
if span.is_recording() and response:
171-
# When streaming, this is called multiple times so attributes would be
172-
# overwritten. In practice, it looks the API only returns the interesting
173-
# attributes on the last streamed response. However, I couldn't find
174-
# documentation for this and setting attributes shouldn't be too expensive.
175-
span.set_attributes(
176-
**response_attributes,
177-
**server_attributes,
178-
**request_attributes,
179-
)
180-
# event = Event(name="gen_ai.client.inference.operation.details")
181-
182-
self.event_logger.emit(
183-
create_operation_details_event(
184-
api_endpoint=api_endpoint,
185-
params=params,
186-
capture_content=capture_content,
187-
response=response,
188-
)
179+
system_instructions, inputs, outputs = [], [], []
180+
if params.system_instruction:
181+
system_instructions = [
182+
Text(
183+
content="\n".join(
184+
part.text
185+
for part in params.system_instruction.parts
186+
)
187+
)
188+
]
189+
if params.contents:
190+
inputs = [
191+
convert_content_to_message(content)
192+
for content in params.contents
193+
]
194+
if response:
195+
outputs = convert_response_to_output_messages(response)
196+
content = {
197+
k: [asdict(x) for x in v]
198+
for k, v in [
199+
(
200+
GenAI.GEN_AI_SYSTEM_INSTRUCTIONS,
201+
system_instructions,
202+
),
203+
(GenAI.GEN_AI_INPUT_MESSAGES, inputs),
204+
(GenAI.GEN_AI_OUTPUT_MESSAGES, outputs),
205+
]
206+
if v
207+
}
208+
if span.is_recording():
209+
span.set_attributes(attributes)
210+
if capture_content in frozenset(
211+
[
212+
ContentCapturingMode.SPAN_AND_EVENT,
213+
ContentCapturingMode.SPAN_ONLY,
214+
]
215+
):
216+
span.set_attributes(
217+
{k: json.dumps(v) for k, v in content.items()}
218+
)
219+
event = Event(
220+
name="gen_ai.client.inference.operation.details",
221+
)
222+
event.attributes = attributes
223+
if capture_content in frozenset(
224+
[
225+
ContentCapturingMode.SPAN_AND_EVENT,
226+
ContentCapturingMode.EVENT_ONLY,
227+
]
228+
):
229+
event.attributes |= content
230+
self.event_logger.emit(event)
231+
self.completion_hook.on_completion(
232+
inputs=inputs,
233+
outputs=outputs,
234+
system_instruction=system_instructions,
235+
span=span,
236+
log_record=event,
189237
)
190238

191239
yield handle_response

0 commit comments

Comments
 (0)