Skip to content

Commit b5ecace

Browse files
aabmassaabmass
committed
tmp
1 parent 60c60f7 commit b5ecace

15 files changed

+1402
-147
lines changed

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

Lines changed: 16 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -39,14 +39,14 @@
3939
---
4040
"""
4141

42+
import itertools
4243
from typing import Any, Collection
4344

4445
from wrapt import (
4546
wrap_function_wrapper, # type: ignore[reportUnknownVariableType]
4647
)
4748

4849
from opentelemetry._events import get_event_logger
49-
from opentelemetry.instrumentation.utils import unwrap
5050
from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
5151
from opentelemetry.instrumentation.utils import unwrap
5252
from opentelemetry.instrumentation.vertexai.package import _instruments
@@ -56,7 +56,7 @@
5656
from opentelemetry.trace import get_tracer
5757

5858

59-
def _client_classes():
59+
def _methods_to_wrap():
6060
# This import is very slow, do it lazily in case instrument() is not called
6161

6262
# pylint: disable=import-outside-toplevel
@@ -67,9 +67,12 @@ def _client_classes():
6767
client as client_v1beta1,
6868
)
6969

70-
return (
71-
client.PredictionServiceClient,
72-
client_v1beta1.PredictionServiceClient,
70+
return itertools.product(
71+
(
72+
client.PredictionServiceClient,
73+
client_v1beta1.PredictionServiceClient,
74+
),
75+
("generate_content", "stream_generate_content"),
7376
)
7477

7578

@@ -94,32 +97,16 @@ def _instrument(self, **kwargs: Any):
9497
event_logger_provider=event_logger_provider,
9598
)
9699

97-
for client_class in _client_classes():
100+
patched_methods = PatchedMethods(
101+
tracer, event_logger, is_content_enabled()
102+
)
103+
for client_class, method_name in _methods_to_wrap():
98104
wrap_function_wrapper(
99105
client_class,
100-
name="generate_content",
101-
wrapper=generate_content_create(
102-
tracer, event_logger, is_content_enabled()
103-
),
104-
)
105-
106-
for module in (
107-
"google.cloud.aiplatform_v1.services.prediction_service.client",
108-
"google.cloud.aiplatform_v1beta1.services.prediction_service.client",
109-
):
110-
# non streaming
111-
wrap_function_wrapper(
112-
module=module,
113-
name="PredictionServiceClient.generate_content",
114-
wrapper=patched_methods.generate_content,
115-
)
116-
# streaming
117-
wrap_function_wrapper(
118-
module=module,
119-
name="PredictionServiceClient.stream_generate_content",
120-
wrapper=patched_methods.stream_generate_content,
106+
name=method_name,
107+
wrapper=getattr(patched_methods, method_name),
121108
)
122109

123110
def _uninstrument(self, **kwargs: Any) -> None:
124-
for client_class in _client_classes():
125-
unwrap(client_class, "generate_content")
111+
for client_class, method_name in _methods_to_wrap():
112+
unwrap(client_class, method_name)

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -161,10 +161,11 @@ def choice_event(
161161
https://github.com/open-telemetry/semantic-conventions/blob/v1.28.0/docs/gen-ai/gen-ai-events.md#event-gen_aichoice
162162
"""
163163
body: dict[str, AnyValue] = {
164-
"finish_reason": finish_reason,
165164
"index": index,
166165
"message": _asdict_filter_nulls(message),
167166
}
167+
if finish_reason:
168+
body["finish_reason"] = finish_reason
168169

169170
tool_calls_list = [
170171
_asdict_filter_nulls(tool_call) for tool_call in tool_calls

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

Lines changed: 7 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -19,17 +19,10 @@
1919
TYPE_CHECKING,
2020
Any,
2121
Callable,
22-
Generator,
2322
Iterable,
2423
MutableSequence,
2524
)
2625

27-
from google.cloud.aiplatform_v1.types.prediction_service import (
28-
GenerateContentResponse,
29-
)
30-
from google.cloud.aiplatform_v1beta1.types.prediction_service import (
31-
GenerateContentResponse,
32-
)
3326
from opentelemetry._events import EventLogger
3427
from opentelemetry.instrumentation.vertexai.utils import (
3528
GenerateContentParams,
@@ -164,7 +157,10 @@ def generate_content(
164157
| client_v1beta1.PredictionServiceClient,
165158
args: Any,
166159
kwargs: Any,
167-
) -> GenerateContentResponse | GenerateContentResponse:
160+
) -> (
161+
prediction_service.GenerateContentResponse
162+
| prediction_service_v1beta1.GenerateContentResponse
163+
):
168164
with self._start_as_current_span(
169165
instance, args, kwargs
170166
) as handle_response:
@@ -183,8 +179,9 @@ def stream_generate_content(
183179
| client_v1beta1.PredictionServiceClient,
184180
args: Any,
185181
kwargs: Any,
186-
) -> Generator[
187-
GenerateContentResponse | GenerateContentResponse, Any, None
182+
) -> Iterable[
183+
prediction_service.GenerateContentResponse
184+
| prediction_service_v1beta1.GenerateContentResponse,
188185
]:
189186
print("stream_generate_content() starting ctxmanager")
190187
with self._start_as_current_span(

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

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -330,10 +330,9 @@ def _map_finish_reason(
330330
| content_v1beta1.Candidate.FinishReason,
331331
) -> FinishReason | str:
332332
EnumType = type(finish_reason) # pylint: disable=invalid-name
333-
if (
334-
finish_reason is EnumType.FINISH_REASON_UNSPECIFIED
335-
or finish_reason is EnumType.OTHER
336-
):
333+
if finish_reason is EnumType.FINISH_REASON_UNSPECIFIED:
334+
return ""
335+
if finish_reason is EnumType.OTHER:
337336
return "error"
338337
if finish_reason is EnumType.STOP:
339338
return "stop"

‎instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/cassettes/test_stream_function_call_choice.yaml‎

Lines changed: 120 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,120 @@
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:streamGenerateContent?%24alt=json%3Benum-encoding%3Dint
53+
response:
54+
body:
55+
string: |-
56+
[
57+
{
58+
"candidates": [
59+
{
60+
"content": {
61+
"role": "model",
62+
"parts": [
63+
{
64+
"functionCall": {
65+
"name": "get_current_weather",
66+
"args": {
67+
"location": "New Delhi"
68+
}
69+
}
70+
},
71+
{
72+
"functionCall": {
73+
"name": "get_current_weather",
74+
"args": {
75+
"location": "San Francisco"
76+
}
77+
}
78+
}
79+
]
80+
},
81+
"finishReason": 1
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-03-05T04:44:12.226326Z",
103+
"responseId": "nNbHZ5boDZeTmecP49qwuQU"
104+
}
105+
]
106+
headers:
107+
Content-Type:
108+
- application/json; charset=UTF-8
109+
Transfer-Encoding:
110+
- chunked
111+
Vary:
112+
- Origin
113+
- X-Origin
114+
- Referer
115+
content-length:
116+
- '985'
117+
status:
118+
code: 200
119+
message: OK
120+
version: 1

‎instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/cassettes/test_stream_function_call_choice_no_content.yaml‎

Lines changed: 120 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,120 @@
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:streamGenerateContent?%24alt=json%3Benum-encoding%3Dint
53+
response:
54+
body:
55+
string: |-
56+
[
57+
{
58+
"candidates": [
59+
{
60+
"content": {
61+
"role": "model",
62+
"parts": [
63+
{
64+
"functionCall": {
65+
"name": "get_current_weather",
66+
"args": {
67+
"location": "New Delhi"
68+
}
69+
}
70+
},
71+
{
72+
"functionCall": {
73+
"name": "get_current_weather",
74+
"args": {
75+
"location": "San Francisco"
76+
}
77+
}
78+
}
79+
]
80+
},
81+
"finishReason": 1
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-03-05T04:46:18.094334Z",
103+
"responseId": "GtfHZ_7gBe2Om9IPrJa3MQ"
104+
}
105+
]
106+
headers:
107+
Content-Type:
108+
- application/json; charset=UTF-8
109+
Transfer-Encoding:
110+
- chunked
111+
Vary:
112+
- Origin
113+
- X-Origin
114+
- Referer
115+
content-length:
116+
- '984'
117+
status:
118+
code: 200
119+
message: OK
120+
version: 1

‎instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/cassettes/test_generate_content_streaming.yaml‎ renamed to ‎instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/cassettes/test_stream_generate_content.yaml‎

File renamed without changes.

0 commit comments

Comments
 (0)