Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions tests/unit/vertex_adk/test_agent_engine_templates_adk.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
import os
import cloudpickle
import sys
import re
from unittest import mock
from typing import Optional

Expand Down Expand Up @@ -769,8 +770,17 @@ def test_tracing_setup(
otlp_span_exporter_mock.assert_called_once_with(
session=mock.ANY,
endpoint="https://telemetry.googleapis.com/v1/traces",
headers=mock.ANY,
)

user_agent = otlp_span_exporter_mock.call_args.kwargs["headers"]["User-Agent"]
assert (
re.fullmatch(
r"Vertex-Agent-Engine\/[\d\.]+ OTel-OTLP-Exporter-Python\/[\d\.]+",
user_agent,
)
is not None
)
assert (
trace_provider_mock.call_args.kwargs["resource"].attributes
== expected_attributes
Expand Down
10 changes: 10 additions & 0 deletions tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
import importlib
import json
import os
import re
from unittest import mock
from typing import Optional

Expand Down Expand Up @@ -723,8 +724,17 @@ def test_tracing_setup(
otlp_span_exporter_mock.assert_called_once_with(
session=mock.ANY,
endpoint="https://telemetry.googleapis.com/v1/traces",
headers=mock.ANY,
)

user_agent = otlp_span_exporter_mock.call_args.kwargs["headers"]["User-Agent"]
assert (
re.fullmatch(
r"Vertex-Agent-Engine\/[\d\.]+ OTel-OTLP-Exporter-Python\/[\d\.]+",
user_agent,
)
is not None
)
assert (
trace_provider_mock.call_args.kwargs["resource"].attributes
== expected_attributes
Expand Down
7 changes: 7 additions & 0 deletions vertexai/agent_engines/templates/adk.py
Original file line number Diff line number Diff line change
Expand Up @@ -357,8 +357,10 @@ def _detect_cloud_resource_id(project_id: str) -> Optional[str]:

if enable_tracing:
try:
import opentelemetry.exporter.otlp.proto.http.version
import opentelemetry.exporter.otlp.proto.http.trace_exporter
import google.auth.transport.requests
from google.cloud.aiplatform import version as aip_version
except (ImportError, AttributeError):
return _warn_missing_dependency(
"opentelemetry-exporter-otlp-proto-http", needed_for_tracing=True
Expand All @@ -367,12 +369,17 @@ def _detect_cloud_resource_id(project_id: str) -> Optional[str]:
import google.auth

credentials, _ = google.auth.default()
vertex_sdk_version = aip_version.__version__
otlp_http_version = opentelemetry.exporter.otlp.proto.http.version.__version__
user_agent = f"Vertex-Agent-Engine/{vertex_sdk_version} OTel-OTLP-Exporter-Python/{otlp_http_version}"

span_exporter = (
opentelemetry.exporter.otlp.proto.http.trace_exporter.OTLPSpanExporter(
session=google.auth.transport.requests.AuthorizedSession(
credentials=credentials
),
endpoint="https://telemetry.googleapis.com/v1/traces",
headers={"User-Agent": user_agent},
)
)
span_processor = opentelemetry.sdk.trace.export.BatchSpanProcessor(
Expand Down
7 changes: 7 additions & 0 deletions vertexai/preview/reasoning_engines/templates/adk.py
Original file line number Diff line number Diff line change
Expand Up @@ -359,8 +359,10 @@ def _detect_cloud_resource_id(project_id: str) -> Optional[str]:

if enable_tracing:
try:
import opentelemetry.exporter.otlp.proto.http.version
import opentelemetry.exporter.otlp.proto.http.trace_exporter
import google.auth.transport.requests
from google.cloud.aiplatform import version as aip_version
except (ImportError, AttributeError):
return _warn_missing_dependency(
"opentelemetry-exporter-otlp-proto-http", needed_for_tracing=True
Expand All @@ -369,12 +371,17 @@ def _detect_cloud_resource_id(project_id: str) -> Optional[str]:
import google.auth

credentials, _ = google.auth.default()
vertex_sdk_version = aip_version.__version__
otlp_http_version = opentelemetry.exporter.otlp.proto.http.version.__version__
user_agent = f"Vertex-Agent-Engine/{vertex_sdk_version} OTel-OTLP-Exporter-Python/{otlp_http_version}"

span_exporter = (
opentelemetry.exporter.otlp.proto.http.trace_exporter.OTLPSpanExporter(
session=google.auth.transport.requests.AuthorizedSession(
credentials=credentials
),
endpoint="https://telemetry.googleapis.com/v1/traces",
headers={"User-Agent": user_agent},
)
)
span_processor = opentelemetry.sdk.trace.export.BatchSpanProcessor(
Expand Down