Skip to content

Commit 745c296

Browse files
DylanRussellDylanRussell
committed
Address comments
1 parent 6f7aa72 commit 745c296

File tree

5 files changed

+130
-222
lines changed

5 files changed

+130
-222
lines changed

‎instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/shared_test_utils.py‎

Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,122 @@
1+
# Copyright The OpenTelemetry Authors
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from vertexai.generative_models import (
16+
Content,
17+
FunctionDeclaration,
18+
GenerativeModel,
19+
Part,
20+
Tool,
21+
)
22+
23+
24+
def weather_tool() -> Tool:
25+
# Adapted from https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/function-calling#parallel-samples
26+
get_current_weather_func = FunctionDeclaration(
27+
name="get_current_weather",
28+
description="Get the current weather in a given location",
29+
parameters={
30+
"type": "object",
31+
"properties": {
32+
"location": {
33+
"type": "string",
34+
"description": "The location for which to get the weather. "
35+
"It can be a city name, a city name and state, or a zip code. "
36+
"Examples: 'San Francisco', 'San Francisco, CA', '95616', etc.",
37+
},
38+
},
39+
},
40+
)
41+
return Tool(
42+
function_declarations=[get_current_weather_func],
43+
)
44+
45+
46+
def ask_about_weather(generate_content: callable) -> None:
47+
model = GenerativeModel("gemini-2.5-pro", tools=[weather_tool()])
48+
# Model will respond asking for function calls
49+
generate_content(
50+
model,
51+
[
52+
# User asked about weather
53+
Content(
54+
role="user",
55+
parts=[
56+
Part.from_text(
57+
"Get weather details in New Delhi and San Francisco?"
58+
),
59+
],
60+
),
61+
],
62+
)
63+
64+
65+
def ask_about_weather_function_response(
66+
generate_content: callable,
67+
) -> None:
68+
model = GenerativeModel("gemini-2.5-pro", tools=[weather_tool()])
69+
generate_content(
70+
model,
71+
[
72+
# User asked about weather
73+
Content(
74+
role="user",
75+
parts=[
76+
Part.from_text(
77+
"Get weather details in New Delhi and San Francisco?"
78+
),
79+
],
80+
),
81+
# Model requests two function calls
82+
Content(
83+
role="model",
84+
parts=[
85+
Part.from_dict(
86+
{
87+
"function_call": {
88+
"name": "get_current_weather",
89+
"args": {"location": "New Delhi"},
90+
}
91+
},
92+
),
93+
Part.from_dict(
94+
{
95+
"function_call": {
96+
"name": "get_current_weather",
97+
"args": {"location": "San Francisco"},
98+
}
99+
},
100+
),
101+
],
102+
),
103+
# User responds with function responses
104+
Content(
105+
role="user",
106+
parts=[
107+
Part.from_function_response(
108+
name="get_current_weather",
109+
response={
110+
"content": '{"temperature": 35, "unit": "C"}'
111+
},
112+
),
113+
Part.from_function_response(
114+
name="get_current_weather",
115+
response={
116+
"content": '{"temperature": 25, "unit": "C"}'
117+
},
118+
),
119+
],
120+
),
121+
],
122+
)

‎instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/test_chat_completions.py‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -354,7 +354,6 @@ def test_generate_content_all_events(
354354
"You are a clever language model"
355355
),
356356
),
357-
generate_content,
358357
log_exporter,
359358
instrument_with_content,
360359
)
@@ -373,15 +372,13 @@ def test_preview_generate_content_all_input_events(
373372
"You are a clever language model"
374373
),
375374
),
376-
generate_content,
377375
log_exporter,
378376
instrument_with_content,
379377
)
380378

381379

382380
def generate_content_all_input_events(
383381
model: GenerativeModel | PreviewGenerativeModel,
384-
generate_content: callable,
385382
log_exporter: InMemoryLogExporter,
386383
instrument_with_content: VertexAIInstrumentor,
387384
):

‎instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/test_chat_completions_experimental.py‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -358,7 +358,6 @@ def test_generate_content_all_events(
358358
"You are a clever language model"
359359
),
360360
),
361-
generate_content,
362361
log_exporter,
363362
instrument_with_experimental_semconvs,
364363
)
@@ -377,15 +376,13 @@ def test_preview_generate_content_all_input_events(
377376
"You are a clever language model"
378377
),
379378
),
380-
generate_content,
381379
log_exporter,
382380
instrument_with_experimental_semconvs,
383381
)
384382

385383

386384
def generate_content_all_input_events(
387385
model: GenerativeModel | PreviewGenerativeModel,
388-
generate_content: callable,
389386
log_exporter: InMemoryLogExporter,
390387
instrument_with_experimental_semconvs: VertexAIInstrumentor,
391388
):

‎instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/test_function_calling.py‎

Lines changed: 4 additions & 108 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,4 @@
11
import pytest
2-
from vertexai.generative_models import (
3-
Content,
4-
FunctionDeclaration,
5-
GenerativeModel,
6-
Part,
7-
Tool,
8-
)
92

103
from opentelemetry.instrumentation.vertexai import VertexAIInstrumentor
114
from opentelemetry.sdk._logs._internal.export.in_memory_log_exporter import (
@@ -14,6 +7,10 @@
147
from opentelemetry.sdk.trace.export.in_memory_span_exporter import (
158
InMemorySpanExporter,
169
)
10+
from tests.shared_test_utils import (
11+
ask_about_weather,
12+
ask_about_weather_function_response,
13+
)
1714

1815

1916
@pytest.mark.vcr()
@@ -315,104 +312,3 @@ def test_tool_events_no_content(
315312
"index": 0,
316313
"message": {"role": "model"},
317314
}
318-
319-
320-
def weather_tool() -> Tool:
321-
# Adapted from https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/function-calling#parallel-samples
322-
get_current_weather_func = FunctionDeclaration(
323-
name="get_current_weather",
324-
description="Get the current weather in a given location",
325-
parameters={
326-
"type": "object",
327-
"properties": {
328-
"location": {
329-
"type": "string",
330-
"description": "The location for which to get the weather. "
331-
"It can be a city name, a city name and state, or a zip code. "
332-
"Examples: 'San Francisco', 'San Francisco, CA', '95616', etc.",
333-
},
334-
},
335-
},
336-
)
337-
return Tool(
338-
function_declarations=[get_current_weather_func],
339-
)
340-
341-
342-
def ask_about_weather(generate_content: callable) -> None:
343-
model = GenerativeModel("gemini-2.5-pro", tools=[weather_tool()])
344-
# Model will respond asking for function calls
345-
generate_content(
346-
model,
347-
[
348-
# User asked about weather
349-
Content(
350-
role="user",
351-
parts=[
352-
Part.from_text(
353-
"Get weather details in New Delhi and San Francisco?"
354-
),
355-
],
356-
),
357-
],
358-
)
359-
360-
361-
def ask_about_weather_function_response(
362-
generate_content: callable,
363-
) -> None:
364-
model = GenerativeModel("gemini-2.5-pro", tools=[weather_tool()])
365-
generate_content(
366-
model,
367-
[
368-
# User asked about weather
369-
Content(
370-
role="user",
371-
parts=[
372-
Part.from_text(
373-
"Get weather details in New Delhi and San Francisco?"
374-
),
375-
],
376-
),
377-
# Model requests two function calls
378-
Content(
379-
role="model",
380-
parts=[
381-
Part.from_dict(
382-
{
383-
"function_call": {
384-
"name": "get_current_weather",
385-
"args": {"location": "New Delhi"},
386-
}
387-
},
388-
),
389-
Part.from_dict(
390-
{
391-
"function_call": {
392-
"name": "get_current_weather",
393-
"args": {"location": "San Francisco"},
394-
}
395-
},
396-
),
397-
],
398-
),
399-
# User responds with function responses
400-
Content(
401-
role="user",
402-
parts=[
403-
Part.from_function_response(
404-
name="get_current_weather",
405-
response={
406-
"content": '{"temperature": 35, "unit": "C"}'
407-
},
408-
),
409-
Part.from_function_response(
410-
name="get_current_weather",
411-
response={
412-
"content": '{"temperature": 25, "unit": "C"}'
413-
},
414-
),
415-
],
416-
),
417-
],
418-
)

0 commit comments

Comments
 (0)