Skip to content
Open
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
230 changes: 148 additions & 82 deletions agentplatform/agent_engines/templates/adk.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,6 @@
import enum
import os
import queue
import sys
import threading
from typing import (
Any,
Expand Down Expand Up @@ -133,6 +132,14 @@
# rather than inherit AuthorizedSession's 120s default.
_TELEMETRY_API_CHECK_TIMEOUT_SECONDS = 5.0

_DEFAULT_TELEMETRY_LOGS_ENDPOINT = "https://telemetry.googleapis.com/v1/logs"

_GCP_LOG_NAME = "gcp.log_name"
_EVENT_NAME = "event.name"
_GCP_RESOURCE_TYPE = "gcp.resource_type"
_LOCATION = "location"
_REASONING_ENGINE_ID = "reasoning_engine_id"


class _MtlsEndpoint(enum.Enum):
"""Enum for the mTLS endpoint setting."""
Expand Down Expand Up @@ -357,8 +364,8 @@ def _warn_missing_dependency(
)
MISSING_LOGGING_IMPORT_ERROR_MESSAGE = (
"proceeding with logging disabled because not all packages (i.e."
" `google-cloud-logging`, `opentelemetry-sdk`,"
" `opentelemetry-exporter-gcp-logging`) for tracing have been installed"
" `opentelemetry-sdk`, `opentelemetry-exporter-otlp-proto-http`)"
" for logging have been installed"
)

if needed_for_tracing and enable_tracing:
Expand All @@ -367,15 +374,6 @@ def _warn_missing_dependency(
_warn(MISSING_LOGGING_IMPORT_ERROR_MESSAGE)
return None

def _detect_cloud_resource_id(project_id: str) -> Optional[str]:
location = os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_LOCATION", "") or os.getenv(
"GOOGLE_CLOUD_LOCATION", ""
)
agent_engine_id = os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_ID")
if all(v is not None for v in (location, agent_engine_id)):
return f"//aiplatform.googleapis.com/projects/{project_id}/locations/{location}/reasoningEngines/{agent_engine_id}"
return None

try:
import opentelemetry
import opentelemetry.trace
Expand All @@ -396,30 +394,7 @@ def _detect_cloud_resource_id(project_id: str) -> Optional[str]:
"opentelemetry-sdk", needed_for_tracing=True, needed_for_logging=True
)

import uuid

# Provide a set of resource attributes but allow to override them with env
# variables like OTEL_RESOURCE_ATTRIBUTES and OTEL_SERVICE_NAME.
cloud_resource_id = _detect_cloud_resource_id(project_id)
resource = opentelemetry.sdk.resources.Resource.create(
attributes={
"gcp.project_id": project_id,
"cloud.account.id": project_id,
"cloud.provider": "gcp",
"cloud.platform": "gcp.agent_engine",
"service.name": os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_ID", ""),
"service.instance.id": f"{uuid.uuid4().hex}-{os.getpid()}",
"cloud.region": (
os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_LOCATION", "")
or os.getenv("GOOGLE_CLOUD_LOCATION", "")
),
}
| (
{"cloud.resource_id": cloud_resource_id}
if cloud_resource_id is not None
else {}
)
).merge(opentelemetry.sdk.resources.OTELResourceDetector().detect())
resource = _create_otel_resource(project_id)

if enable_tracing:
try:
Expand All @@ -435,12 +410,7 @@ 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}"
f" OTel-OTLP-Exporter-Python/{otlp_http_version}"
)
user_agent = _get_user_agent()

session = requests_auth.AuthorizedSession(credentials=credentials)

Expand Down Expand Up @@ -500,54 +470,47 @@ def _detect_cloud_resource_id(project_id: str) -> Optional[str]:

if enable_logging:
try:
import opentelemetry.exporter.cloud_logging
import opentelemetry.exporter.otlp.proto.http._log_exporter
import google.auth.transport.requests
except (ImportError, AttributeError):
return _warn_missing_dependency(
"opentelemetry-exporter-gcp-logging", needed_for_logging=True
"opentelemetry-exporter-otlp-proto-http", needed_for_logging=True
)

class _SimpleLogRecordProcessor(
opentelemetry.sdk._logs.export.SimpleLogRecordProcessor
):
import google.auth

def force_flush(
self, timeout_millis: int = 30000
) -> bool: # pylint: disable=no-self-use
sys.stdout.flush()
sys.stderr.flush()
return True

logger_provider = opentelemetry.sdk._logs.LoggerProvider(resource=resource)
# Use the legacy log processor when experimental semconv is enabled.
# Exporting JSON logs to stdout is bugged; Agent Engine fails to
# correctly parse the `gen_ai.client.inference.operation.details`
# messages.
# TODO: b/480102541 - Unify both branches once the regression is fixed.
if "gen_ai_latest_experimental" in os.getenv(
"OTEL_SEMCONV_STABILITY_OPT_IN", ""
).split(","):
logger_provider.add_log_record_processor(
opentelemetry.sdk._logs.export.BatchLogRecordProcessor(
opentelemetry.exporter.cloud_logging.CloudLoggingExporter(
project_id=project_id,
default_log_name=os.getenv(
"GCP_DEFAULT_LOG_NAME", "adk-on-agent-engine"
),
),
)
credentials, _ = google.auth.default()
session = requests_auth.AuthorizedSession(credentials=credentials)

if _use_client_cert_effective():
client_cert_source = (
mtls.default_client_cert_source()
if mtls.has_default_client_cert_source()
else None
)
session.configure_mtls_channel()
endpoint = _get_logs_api_endpoint(client_cert_source)
else:
logger_provider.add_log_record_processor(
_SimpleLogRecordProcessor(
opentelemetry.exporter.cloud_logging.CloudLoggingExporter(
project_id=project_id,
default_log_name=os.getenv(
"GCP_DEFAULT_LOG_NAME", "adk-on-agent-engine"
),
structured_json_file=sys.stdout,
),
)
endpoint = _DEFAULT_TELEMETRY_LOGS_ENDPOINT

# One processor serves stable and experimental semconv. The stdout
# branch experimental records used to need is gone along with the
# Cloud Logging exporter (b/480102541).
logger_provider = opentelemetry.sdk._logs.LoggerProvider(
resource=_create_otel_resource(project_id, "logs")
)
logger_provider.add_log_record_processor(
_named_batch_log_record_processor(
opentelemetry.exporter.otlp.proto.http._log_exporter.OTLPLogExporter(
session=session,
endpoint=endpoint,
headers={"User-Agent": _get_user_agent()},
),
default_log_name=os.getenv(
"GCP_DEFAULT_LOG_NAME", "adk-on-agent-engine"
),
)
)

opentelemetry._logs.set_logger_provider(logger_provider=logger_provider)

Expand Down Expand Up @@ -637,6 +600,109 @@ def _warn_if_telemetry_api_disabled():
_warn(_TELEMETRY_API_DISABLED_WARNING % (project, project))


def _get_user_agent() -> str:
"""Returns the User-Agent to send on OTLP exports."""
from google.cloud.aiplatform import version as aip_version

user_agent = f"Vertex-Agent-Engine/{aip_version.__version__}"
try:
import opentelemetry.exporter.otlp.proto.http.version

user_agent += (
" OTel-OTLP-Exporter-Python/"
f"{opentelemetry.exporter.otlp.proto.http.version.__version__}"
)
except (ImportError, AttributeError):
pass
return user_agent


def _get_logs_api_endpoint(client_cert_source: bytes | None = None) -> str:
"""Returns the logs endpoint matching _get_api_endpoint's mTLS decision.

Args:
client_cert_source (bytes | None): The client certificate source.

Returns:
str: The logs API endpoint to be used.
"""
return _get_api_endpoint(client_cert_source).replace("/v1/traces", "/v1/logs")


def _create_otel_resource(project_id: str, for_signal: str = "unspecified"):
"""Returns the OTel resource describing the Agent Engine deployment.

Args:
project_id: Project to which to send telemetry.
for_signal: The signal the resource is for. `logs` adds the
MonitoredResource hints Cloud Logging needs, which must not be set
on the resource traces and metrics share.

Returns:
The resource to set on the provider for `for_signal`.
"""
import os
import uuid

import opentelemetry.sdk.resources

location = os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_LOCATION", "") or os.getenv(
"GOOGLE_CLOUD_LOCATION", ""
)
agent_engine_id = os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_ID", "")
attributes = {
"gcp.project_id": project_id,
"cloud.account.id": project_id,
"cloud.provider": "gcp",
"cloud.platform": "gcp.agent_engine",
"service.name": agent_engine_id,
"service.instance.id": f"{uuid.uuid4().hex}-{os.getpid()}",
"cloud.region": location,
}
if location and agent_engine_id:
attributes["cloud.resource_id"] = (
f"//aiplatform.googleapis.com/projects/{project_id}"
f"/locations/{location}/reasoningEngines/{agent_engine_id}"
)
if for_signal == "logs":
# Cloud Logging otherwise detects resource as `generic_task`
attributes[_GCP_RESOURCE_TYPE] = "aiplatform.googleapis.com/ReasoningEngine"
attributes[_LOCATION] = location
attributes[_REASONING_ENGINE_ID] = agent_engine_id

# Provide a set of resource attributes but allow to override them with env
# variables like OTEL_RESOURCE_ATTRIBUTES and OTEL_SERVICE_NAME.
return opentelemetry.sdk.resources.Resource.create(attributes=attributes).merge(
opentelemetry.sdk.resources.OTELResourceDetector().detect()
)


def _named_batch_log_record_processor(exporter, *, default_log_name: str):
"""Returns a batch processor that keeps log names and labels stable.

Args:
exporter: The OTLP log exporter to wrap.
default_log_name (str): Log name for records that carry none.

Returns:
The configured log record processor.
"""
import opentelemetry.sdk._logs.export

class _Processor(opentelemetry.sdk._logs.export.BatchLogRecordProcessor):
def on_emit(self, log_record) -> None:
record = log_record.log_record
attributes = dict(record.attributes or {})
if record.event_name:
attributes.setdefault(_EVENT_NAME, record.event_name)
elif _GCP_LOG_NAME not in attributes:
attributes[_GCP_LOG_NAME] = default_log_name
record.attributes = attributes
super().on_emit(log_record)

return _Processor(exporter)


def _get_api_endpoint(client_cert_source: bytes | None = None) -> str:
"""Returns API endpoint based on mTLS configuration and cert availability.

Expand Down
3 changes: 0 additions & 3 deletions setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -151,7 +151,6 @@
reasoning_engine_extra_require = [
"cloudpickle >= 3.0, < 4.0",
"opentelemetry-sdk < 2",
"opentelemetry-exporter-gcp-logging >= 1.11.0a0, < 2.0.0",
"opentelemetry-exporter-otlp-proto-http < 2",
"opentelemetry-instrumentation-google-genai>=0.3b0, <1.0.0",
# TODO(b/538550724): update to stable version of
Expand All @@ -165,9 +164,7 @@
agent_engines_extra_require = [
"packaging >= 24.0",
"cloudpickle >= 3.0, < 4.0",
"google-cloud-logging < 4",
"opentelemetry-sdk < 2",
"opentelemetry-exporter-gcp-logging >= 1.11.0a0, < 2.0.0",
"opentelemetry-exporter-otlp-proto-http < 2",
"pydantic >= 2.11.1, < 3",
"typing_extensions",
Expand Down
Loading
Loading