Skip to content

Commit f3ebc26

Browse files
yeesiancopybara-github
yeesian
authored and
copybara-github
committed
feat: Add Resource to tracer provider in Agent Engine templates
PiperOrigin-RevId: 763962855
1 parent 47ab05a commit f3ebc26

File tree

8 files changed

+310
-24
lines changed

8 files changed

+310
-24
lines changed

‎tests/unit/vertexai/test_generative_models.py

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,13 +21,15 @@
2121
from typing import Dict, Iterable, List, MutableSequence, Optional
2222
from unittest import mock
2323

24+
from google.api_core import operation as ga_operation
2425
import vertexai
2526
from google.cloud.aiplatform import initializer
2627
from google.cloud.aiplatform_v1 import types as types_v1
2728
from google.cloud.aiplatform_v1.services import (
2829
prediction_service as prediction_service_v1,
2930
)
3031
from google.cloud.aiplatform_v1beta1 import types as types_v1beta1
32+
from google.cloud.aiplatform_v1beta1.services import endpoint_service
3133
from vertexai import generative_models
3234
from vertexai.preview import (
3335
generative_models as preview_generative_models,
@@ -48,6 +50,7 @@
4850
)
4951
from vertexai.generative_models import _function_calling_utils
5052
from vertexai.caching import CachedContent
53+
from google.protobuf import field_mask_pb2
5154

5255

5356
_TEST_PROJECT = "test-project"
@@ -1710,6 +1713,115 @@ def test_defs_ref_renaming(self):
17101713
_fix_schema_dict_for_gapic_in_place(actual)
17111714
assert actual == expected
17121715

1716+
@pytest.mark.parametrize(
1717+
"generative_models",
1718+
[preview_generative_models], # Only preview supports set_logging_config
1719+
)
1720+
@mock.patch.object(endpoint_service.EndpointServiceClient, "update_endpoint")
1721+
def test_set_logging_config_for_endpoint(
1722+
self, mock_update_endpoint, generative_models: generative_models
1723+
):
1724+
endpoint_name = (
1725+
f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/endpoints/12345"
1726+
)
1727+
model = generative_models.GenerativeModel(endpoint_name)
1728+
1729+
mock_update_endpoint.return_value = types_v1beta1.Endpoint(name=endpoint_name)
1730+
1731+
enabled = True
1732+
sampling_rate = 0.5
1733+
bigquery_destination = f"bq://{_TEST_PROJECT}.my_dataset.my_table"
1734+
enable_otel_logging = True
1735+
1736+
model.set_request_response_logging_config(
1737+
enabled=enabled,
1738+
sampling_rate=sampling_rate,
1739+
bigquery_destination=bigquery_destination,
1740+
enable_otel_logging=enable_otel_logging,
1741+
)
1742+
1743+
expected_logging_config = types_v1beta1.PredictRequestResponseLoggingConfig(
1744+
enabled=enabled,
1745+
sampling_rate=sampling_rate,
1746+
bigquery_destination=types_v1beta1.BigQueryDestination(
1747+
output_uri=bigquery_destination
1748+
),
1749+
enable_otel_logging=enable_otel_logging,
1750+
)
1751+
expected_endpoint = types_v1beta1.Endpoint(
1752+
name=endpoint_name,
1753+
predict_request_response_logging_config=expected_logging_config,
1754+
)
1755+
expected_update_mask = field_mask_pb2.FieldMask(
1756+
paths=["predict_request_response_logging_config"]
1757+
)
1758+
1759+
mock_update_endpoint.assert_called_once_with(
1760+
types_v1beta1.UpdateEndpointRequest(
1761+
endpoint=expected_endpoint,
1762+
update_mask=expected_update_mask,
1763+
)
1764+
)
1765+
1766+
@pytest.mark.parametrize(
1767+
"generative_models",
1768+
[preview_generative_models], # Only preview supports set_logging_config
1769+
)
1770+
@mock.patch.object(
1771+
endpoint_service.EndpointServiceClient, "set_publisher_model_config"
1772+
)
1773+
def test_set_logging_config_for_publisher_model(
1774+
self, mock_set_publisher_model_config, generative_models: generative_models
1775+
):
1776+
model_name = "gemini-pro"
1777+
model = generative_models.GenerativeModel(model_name)
1778+
full_model_name = f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/publishers/google/models/{model_name}"
1779+
1780+
enabled = False
1781+
sampling_rate = 1.0
1782+
bigquery_destination = f"bq://{_TEST_PROJECT}.another_dataset"
1783+
enable_otel_logging = False
1784+
1785+
mock_operation = mock.Mock(spec=ga_operation.Operation)
1786+
mock_set_publisher_model_config.return_value = mock_operation
1787+
mock_operation.result.return_value = types_v1beta1.PublisherModelConfig(
1788+
logging_config=types_v1beta1.PredictRequestResponseLoggingConfig(
1789+
enabled=enabled,
1790+
sampling_rate=sampling_rate,
1791+
bigquery_destination=types_v1beta1.BigQueryDestination(
1792+
output_uri=bigquery_destination
1793+
),
1794+
enable_otel_logging=enable_otel_logging,
1795+
)
1796+
)
1797+
1798+
model.set_request_response_logging_config(
1799+
enabled=enabled,
1800+
sampling_rate=sampling_rate,
1801+
bigquery_destination=bigquery_destination,
1802+
enable_otel_logging=enable_otel_logging,
1803+
)
1804+
1805+
expected_logging_config = types_v1beta1.PredictRequestResponseLoggingConfig(
1806+
enabled=enabled,
1807+
sampling_rate=sampling_rate,
1808+
bigquery_destination=types_v1beta1.BigQueryDestination(
1809+
output_uri=bigquery_destination
1810+
),
1811+
enable_otel_logging=enable_otel_logging,
1812+
)
1813+
expected_publisher_model_config = types_v1beta1.PublisherModelConfig(
1814+
logging_config=expected_logging_config
1815+
)
1816+
1817+
mock_set_publisher_model_config.assert_called_once_with(
1818+
types_v1beta1.SetPublisherModelConfigRequest(
1819+
name=full_model_name,
1820+
publisher_model_config=expected_publisher_model_config,
1821+
)
1822+
)
1823+
mock_operation.result.assert_called_once()
1824+
17131825

17141826
EXPECTED_SCHEMA_FOR_GET_CURRENT_WEATHER = {
17151827
"title": "get_current_weather",

‎vertexai/agent_engines/_utils.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -685,6 +685,20 @@ def _import_opentelemetry_or_warn() -> Optional[types.ModuleType]:
685685
return None
686686

687687

688+
def _import_opentelemetry_sdk_resources_or_warn() -> Optional[types.ModuleType]:
689+
"""Tries to import the opentelemetry.sdk.trace module."""
690+
try:
691+
import opentelemetry.sdk.resources # noqa:F401
692+
693+
return opentelemetry.sdk.resources
694+
except ImportError:
695+
LOGGER.warning(
696+
"Failed to import opentelemetry.sdk.resources. Please call "
697+
"'pip install google-cloud-aiplatform[agent_engines]'."
698+
)
699+
return None
700+
701+
688702
def _import_opentelemetry_sdk_trace_or_warn() -> Optional[types.ModuleType]:
689703
"""Tries to import the opentelemetry.sdk.trace module."""
690704
try:

‎vertexai/agent_engines/templates/ag2.py

Lines changed: 20 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -96,16 +96,23 @@ def _default_instrumentor_builder(project_id: str):
9696
openinference_autogen = _utils._import_openinference_autogen_or_warn()
9797
opentelemetry = _utils._import_opentelemetry_or_warn()
9898
opentelemetry_sdk_trace = _utils._import_opentelemetry_sdk_trace_or_warn()
99+
opentelemetry_sdk_resources = _utils._import_opentelemetry_sdk_resources_or_warn()
99100
if all(
100101
(
101102
cloud_trace_exporter,
102103
cloud_trace_v2,
103104
openinference_autogen,
104105
opentelemetry,
105106
opentelemetry_sdk_trace,
107+
opentelemetry_sdk_resources,
106108
)
107109
):
108110
import google.auth
111+
import os
112+
113+
SERVICE_INSTANCE_ID = opentelemetry_sdk_resources.SERVICE_INSTANCE_ID
114+
SERVICE_NAME = opentelemetry_sdk_resources.SERVICE_NAME
115+
AGENT_ENGINE_ID = os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID", "")
109116

110117
credentials, _ = google.auth.default()
111118
span_exporter = cloud_trace_exporter.CloudTraceSpanExporter(
@@ -119,6 +126,12 @@ def _default_instrumentor_builder(project_id: str):
119126
span_exporter=span_exporter,
120127
)
121128
)
129+
resource = opentelemetry_sdk_resources.Resource.create(
130+
attributes={
131+
SERVICE_NAME: "aiplatform.googleapis.com/ReasoningEngine",
132+
SERVICE_INSTANCE_ID: AGENT_ENGINE_ID,
133+
}
134+
)
122135
tracer_provider: TracerProvider = opentelemetry.trace.get_tracer_provider()
123136
# Get the appropriate tracer provider:
124137
# 1. If _TRACER_PROVIDER is already set, use that.
@@ -127,7 +140,7 @@ def _default_instrumentor_builder(project_id: str):
127140
# 3. As a final fallback, use _PROXY_TRACER_PROVIDER.
128141
# If none of the above is set, we log a warning, and
129142
# create a tracer provider.
130-
if not tracer_provider:
143+
if AGENT_ENGINE_ID or not tracer_provider:
131144
from google.cloud.aiplatform import base
132145

133146
_LOGGER = base.Logger(__name__)
@@ -137,13 +150,17 @@ def _default_instrumentor_builder(project_id: str):
137150
"OTEL_PYTHON_TRACER_PROVIDER, _TRACER_PROVIDER, "
138151
"or _PROXY_TRACER_PROVIDER."
139152
)
140-
tracer_provider = opentelemetry_sdk_trace.TracerProvider()
153+
tracer_provider = opentelemetry_sdk_trace.TracerProvider(
154+
resource=resource,
155+
)
141156
opentelemetry.trace.set_tracer_provider(tracer_provider)
142157
# Avoids AttributeError:
143158
# 'ProxyTracerProvider' and 'NoOpTracerProvider' objects has no
144159
# attribute 'add_span_processor'.
145160
if _utils.is_noop_or_proxy_tracer_provider(tracer_provider):
146-
tracer_provider = opentelemetry_sdk_trace.TracerProvider()
161+
tracer_provider = opentelemetry_sdk_trace.TracerProvider(
162+
resource=resource,
163+
)
147164
opentelemetry.trace.set_tracer_provider(tracer_provider)
148165
# Avoids OpenTelemetry client already exists error.
149166
_override_active_span_processor(

‎vertexai/agent_engines/templates/langchain.py

Lines changed: 20 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -175,16 +175,23 @@ def _default_instrumentor_builder(project_id: str):
175175
openinference_langchain = _utils._import_openinference_langchain_or_warn()
176176
opentelemetry = _utils._import_opentelemetry_or_warn()
177177
opentelemetry_sdk_trace = _utils._import_opentelemetry_sdk_trace_or_warn()
178+
opentelemetry_sdk_resources = _utils._import_opentelemetry_sdk_resources_or_warn()
178179
if all(
179180
(
180181
cloud_trace_exporter,
181182
cloud_trace_v2,
182183
openinference_langchain,
183184
opentelemetry,
184185
opentelemetry_sdk_trace,
186+
opentelemetry_sdk_resources,
185187
)
186188
):
187189
import google.auth
190+
import os
191+
192+
SERVICE_INSTANCE_ID = opentelemetry_sdk_resources.SERVICE_INSTANCE_ID
193+
SERVICE_NAME = opentelemetry_sdk_resources.SERVICE_NAME
194+
AGENT_ENGINE_ID = os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID", "")
188195

189196
credentials, _ = google.auth.default()
190197
span_exporter = cloud_trace_exporter.CloudTraceSpanExporter(
@@ -198,6 +205,12 @@ def _default_instrumentor_builder(project_id: str):
198205
span_exporter=span_exporter,
199206
)
200207
)
208+
resource = opentelemetry_sdk_resources.Resource.create(
209+
attributes={
210+
SERVICE_NAME: "aiplatform.googleapis.com/ReasoningEngine",
211+
SERVICE_INSTANCE_ID: AGENT_ENGINE_ID,
212+
}
213+
)
201214
tracer_provider: TracerProvider = opentelemetry.trace.get_tracer_provider()
202215
# Get the appropriate tracer provider:
203216
# 1. If _TRACER_PROVIDER is already set, use that.
@@ -206,7 +219,7 @@ def _default_instrumentor_builder(project_id: str):
206219
# 3. As a final fallback, use _PROXY_TRACER_PROVIDER.
207220
# If none of the above is set, we log a warning, and
208221
# create a tracer provider.
209-
if not tracer_provider:
222+
if AGENT_ENGINE_ID or not tracer_provider:
210223
from google.cloud.aiplatform import base
211224

212225
_LOGGER = base.Logger(__name__)
@@ -216,13 +229,17 @@ def _default_instrumentor_builder(project_id: str):
216229
"OTEL_PYTHON_TRACER_PROVIDER, _TRACER_PROVIDER, "
217230
"or _PROXY_TRACER_PROVIDER."
218231
)
219-
tracer_provider = opentelemetry_sdk_trace.TracerProvider()
232+
tracer_provider = opentelemetry_sdk_trace.TracerProvider(
233+
resource=resource,
234+
)
220235
opentelemetry.trace.set_tracer_provider(tracer_provider)
221236
# Avoids AttributeError:
222237
# 'ProxyTracerProvider' and 'NoOpTracerProvider' objects has no
223238
# attribute 'add_span_processor'.
224239
if _utils.is_noop_or_proxy_tracer_provider(tracer_provider):
225-
tracer_provider = opentelemetry_sdk_trace.TracerProvider()
240+
tracer_provider = opentelemetry_sdk_trace.TracerProvider(
241+
resource=resource,
242+
)
226243
opentelemetry.trace.set_tracer_provider(tracer_provider)
227244
# Avoids OpenTelemetry client already exists error.
228245
_override_active_span_processor(

‎vertexai/agent_engines/templates/langgraph.py

Lines changed: 20 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -166,16 +166,23 @@ def _default_instrumentor_builder(project_id: str):
166166
openinference_langchain = _utils._import_openinference_langchain_or_warn()
167167
opentelemetry = _utils._import_opentelemetry_or_warn()
168168
opentelemetry_sdk_trace = _utils._import_opentelemetry_sdk_trace_or_warn()
169+
opentelemetry_sdk_resources = _utils._import_opentelemetry_sdk_resources_or_warn()
169170
if all(
170171
(
171172
cloud_trace_exporter,
172173
cloud_trace_v2,
173174
openinference_langchain,
174175
opentelemetry,
175176
opentelemetry_sdk_trace,
177+
opentelemetry_sdk_resources,
176178
)
177179
):
178180
import google.auth
181+
import os
182+
183+
SERVICE_INSTANCE_ID = opentelemetry_sdk_resources.SERVICE_INSTANCE_ID
184+
SERVICE_NAME = opentelemetry_sdk_resources.SERVICE_NAME
185+
AGENT_ENGINE_ID = os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID", "")
179186

180187
credentials, _ = google.auth.default()
181188
span_exporter = cloud_trace_exporter.CloudTraceSpanExporter(
@@ -189,6 +196,12 @@ def _default_instrumentor_builder(project_id: str):
189196
span_exporter=span_exporter,
190197
)
191198
)
199+
resource = opentelemetry_sdk_resources.Resource.create(
200+
attributes={
201+
SERVICE_NAME: "aiplatform.googleapis.com/ReasoningEngine",
202+
SERVICE_INSTANCE_ID: AGENT_ENGINE_ID,
203+
}
204+
)
192205
tracer_provider: TracerProvider = opentelemetry.trace.get_tracer_provider()
193206
# Get the appropriate tracer provider:
194207
# 1. If _TRACER_PROVIDER is already set, use that.
@@ -197,7 +210,7 @@ def _default_instrumentor_builder(project_id: str):
197210
# 3. As a final fallback, use _PROXY_TRACER_PROVIDER.
198211
# If none of the above is set, we log a warning, and
199212
# create a tracer provider.
200-
if not tracer_provider:
213+
if AGENT_ENGINE_ID or not tracer_provider:
201214
from google.cloud.aiplatform import base
202215

203216
base.Logger(__name__).warning(
@@ -206,13 +219,17 @@ def _default_instrumentor_builder(project_id: str):
206219
"OTEL_PYTHON_TRACER_PROVIDER, _TRACER_PROVIDER, "
207220
"or _PROXY_TRACER_PROVIDER."
208221
)
209-
tracer_provider = opentelemetry_sdk_trace.TracerProvider()
222+
tracer_provider = opentelemetry_sdk_trace.TracerProvider(
223+
resource=resource,
224+
)
210225
opentelemetry.trace.set_tracer_provider(tracer_provider)
211226
# Avoids AttributeError:
212227
# 'ProxyTracerProvider' and 'NoOpTracerProvider' objects has no
213228
# attribute 'add_span_processor'.
214229
if _utils.is_noop_or_proxy_tracer_provider(tracer_provider):
215-
tracer_provider = opentelemetry_sdk_trace.TracerProvider()
230+
tracer_provider = opentelemetry_sdk_trace.TracerProvider(
231+
resource=resource,
232+
)
216233
opentelemetry.trace.set_tracer_provider(tracer_provider)
217234
# Avoids OpenTelemetry client already exists error.
218235
_override_active_span_processor(

0 commit comments

Comments
 (0)