From 5395193e2ddc90991fb7abf052f17bb2adf50225 Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Thu, 13 Aug 2026 01:44:57 -0700 Subject: [PATCH] chore: Migrate all templates from Cloud Trace API to Telemetry API and drop trace_v2 helpers. PiperOrigin-RevId: 963935484 --- agentplatform/_genai/_agent_engines_utils.py | 28 ----- agentplatform/agent_engines/templates/ag2.py | 38 ++++-- .../agent_engines/templates/langchain.py | 38 ++++-- .../agent_engines/templates/langgraph.py | 40 ++++-- .../agent_engines/templates/llama_index.py | 39 +++--- setup.py | 4 - .../frameworks/test_frameworks_ag2.py | 97 ++++++++++++-- .../frameworks/test_frameworks_langchain.py | 102 +++++++++++++-- .../frameworks/test_frameworks_langgraph.py | 96 ++++++++++++-- .../frameworks/test_frameworks_llama_index.py | 118 +++++++++++++++--- .../test_agent_engine_templates_ag2.py | 97 ++++++++++++-- .../test_reasoning_engine_templates_ag2.py | 94 ++++++++++++-- .../test_agent_engine_templates_langchain.py | 98 +++++++++++++-- .../test_agent_engine_templates_langgraph.py | 96 ++++++++++++-- ...st_reasoning_engine_templates_langchain.py | 95 ++++++++++++-- ...st_reasoning_engine_templates_langgraph.py | 93 ++++++++++++-- ..._reasoning_engine_templates_llama_index.py | 95 ++++++++++++-- vertexai/agent_engines/_utils.py | 28 ----- vertexai/agent_engines/templates/ag2.py | 38 ++++-- vertexai/agent_engines/templates/langchain.py | 38 ++++-- vertexai/agent_engines/templates/langgraph.py | 38 ++++-- .../reasoning_engines/templates/ag2.py | 37 ++++-- .../reasoning_engines/templates/langchain.py | 37 ++++-- .../reasoning_engines/templates/langgraph.py | 37 ++++-- .../templates/llama_index.py | 37 ++++-- vertexai/reasoning_engines/_utils.py | 28 ----- 26 files changed, 1248 insertions(+), 338 deletions(-) diff --git a/agentplatform/_genai/_agent_engines_utils.py b/agentplatform/_genai/_agent_engines_utils.py index f03e8f32d5..41f9d180ec 100644 --- a/agentplatform/_genai/_agent_engines_utils.py +++ b/agentplatform/_genai/_agent_engines_utils.py @@ -2256,34 +2256,6 @@ def _import_opentelemetry_sdk_trace_or_warn() -> Optional[types.ModuleType]: return None -def _import_cloud_trace_v2_or_warn() -> Optional[types.ModuleType]: - """Tries to import the google.cloud.trace_v2 module.""" - try: - import google.cloud.trace_v2 - - return google.cloud.trace_v2 - except ImportError: - logger.warning( - "google-cloud-trace is not installed. Please call " - "'pip install google-cloud-aiplatform[agent_engines]'." - ) - return None - - -def _import_cloud_trace_exporter_or_warn() -> Optional[types.ModuleType]: - """Tries to import the opentelemetry.exporter.cloud_trace module.""" - try: - import opentelemetry.exporter.cloud_trace - - return opentelemetry.exporter.cloud_trace - except ImportError: - logger.warning( - "opentelemetry-exporter-gcp-trace is not installed. Please " - "call 'pip install google-cloud-aiplatform[agent_engines]'." - ) - return None - - def _import_openinference_langchain_or_warn() -> Optional[types.ModuleType]: """Tries to import the openinference.instrumentation.langchain module.""" try: diff --git a/agentplatform/agent_engines/templates/ag2.py b/agentplatform/agent_engines/templates/ag2.py index ee12e49d8f..c6d7f8c3a0 100644 --- a/agentplatform/agent_engines/templates/ag2.py +++ b/agentplatform/agent_engines/templates/ag2.py @@ -91,8 +91,6 @@ def _default_runnable_builder( def _default_instrumentor_builder(project_id: str): from agentplatform._genai import _agent_engines_utils - cloud_trace_exporter = _agent_engines_utils._import_cloud_trace_exporter_or_warn() - cloud_trace_v2 = _agent_engines_utils._import_cloud_trace_v2_or_warn() openinference_autogen = _agent_engines_utils._import_openinference_autogen_or_warn() opentelemetry = _agent_engines_utils._import_opentelemetry_or_warn() opentelemetry_sdk_trace = ( @@ -100,21 +98,37 @@ def _default_instrumentor_builder(project_id: str): ) if all( ( - cloud_trace_exporter, - cloud_trace_v2, openinference_autogen, opentelemetry, opentelemetry_sdk_trace, ) ): - import google.auth + try: + import opentelemetry.exporter.otlp.proto.http.trace_exporter + import opentelemetry.sdk.resources + import google.auth + import google.auth.transport.requests + except (ImportError, AttributeError): + from google.cloud.aiplatform import base + + _LOGGER = base.Logger(__name__) + _LOGGER.warning( + "enable_tracing=True but proceeding with tracing disabled " + "because opentelemetry-exporter-otlp-proto-http is not installed." + ) + return None credentials, _ = google.auth.default() - span_exporter = cloud_trace_exporter.CloudTraceSpanExporter( - project_id=project_id, - client=cloud_trace_v2.TraceServiceClient( - credentials=credentials.with_quota_project(project_id), - ), + resource = opentelemetry.sdk.resources.Resource.create( + attributes={"gcp.project_id": project_id}, + ).merge(opentelemetry.sdk.resources.OTELResourceDetector().detect()) + 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", + ) ) span_processor: SpanProcessor = ( opentelemetry_sdk_trace.export.SimpleSpanProcessor( @@ -139,13 +153,13 @@ def _default_instrumentor_builder(project_id: str): "OTEL_PYTHON_TRACER_PROVIDER, _TRACER_PROVIDER, " "or _PROXY_TRACER_PROVIDER." ) - tracer_provider = opentelemetry_sdk_trace.TracerProvider() + tracer_provider = opentelemetry_sdk_trace.TracerProvider(resource=resource) opentelemetry.trace.set_tracer_provider(tracer_provider) # Avoids AttributeError: # 'ProxyTracerProvider' and 'NoOpTracerProvider' objects has no # attribute 'add_span_processor'. if _agent_engines_utils.is_noop_or_proxy_tracer_provider(tracer_provider): - tracer_provider = opentelemetry_sdk_trace.TracerProvider() + tracer_provider = opentelemetry_sdk_trace.TracerProvider(resource=resource) opentelemetry.trace.set_tracer_provider(tracer_provider) # Avoids OpenTelemetry client already exists error. _override_active_span_processor( diff --git a/agentplatform/agent_engines/templates/langchain.py b/agentplatform/agent_engines/templates/langchain.py index 17751524ba..bf8947b230 100644 --- a/agentplatform/agent_engines/templates/langchain.py +++ b/agentplatform/agent_engines/templates/langchain.py @@ -194,8 +194,6 @@ def _default_runnable_builder( def _default_instrumentor_builder(project_id: str): from agentplatform._genai import _agent_engines_utils - cloud_trace_exporter = _agent_engines_utils._import_cloud_trace_exporter_or_warn() - cloud_trace_v2 = _agent_engines_utils._import_cloud_trace_v2_or_warn() openinference_langchain = ( _agent_engines_utils._import_openinference_langchain_or_warn() ) @@ -205,21 +203,37 @@ def _default_instrumentor_builder(project_id: str): ) if all( ( - cloud_trace_exporter, - cloud_trace_v2, openinference_langchain, opentelemetry, opentelemetry_sdk_trace, ) ): - import google.auth + try: + import opentelemetry.exporter.otlp.proto.http.trace_exporter + import opentelemetry.sdk.resources + import google.auth + import google.auth.transport.requests + except (ImportError, AttributeError): + from google.cloud.aiplatform import base + + _LOGGER = base.Logger(__name__) + _LOGGER.warning( + "enable_tracing=True but proceeding with tracing disabled " + "because opentelemetry-exporter-otlp-proto-http is not installed." + ) + return None credentials, _ = google.auth.default() - span_exporter = cloud_trace_exporter.CloudTraceSpanExporter( - project_id=project_id, - client=cloud_trace_v2.TraceServiceClient( - credentials=credentials.with_quota_project(project_id), - ), + resource = opentelemetry.sdk.resources.Resource.create( + attributes={"gcp.project_id": project_id}, + ).merge(opentelemetry.sdk.resources.OTELResourceDetector().detect()) + 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", + ) ) span_processor: SpanProcessor = ( opentelemetry_sdk_trace.export.SimpleSpanProcessor( @@ -244,13 +258,13 @@ def _default_instrumentor_builder(project_id: str): "OTEL_PYTHON_TRACER_PROVIDER, _TRACER_PROVIDER, " "or _PROXY_TRACER_PROVIDER." ) - tracer_provider = opentelemetry_sdk_trace.TracerProvider() + tracer_provider = opentelemetry_sdk_trace.TracerProvider(resource=resource) opentelemetry.trace.set_tracer_provider(tracer_provider) # Avoids AttributeError: # 'ProxyTracerProvider' and 'NoOpTracerProvider' objects has no # attribute 'add_span_processor'. if _agent_engines_utils.is_noop_or_proxy_tracer_provider(tracer_provider): - tracer_provider = opentelemetry_sdk_trace.TracerProvider() + tracer_provider = opentelemetry_sdk_trace.TracerProvider(resource=resource) opentelemetry.trace.set_tracer_provider(tracer_provider) # Avoids OpenTelemetry client already exists error. _override_active_span_processor( diff --git a/agentplatform/agent_engines/templates/langgraph.py b/agentplatform/agent_engines/templates/langgraph.py index 0ebdb690cc..3afc19c499 100644 --- a/agentplatform/agent_engines/templates/langgraph.py +++ b/agentplatform/agent_engines/templates/langgraph.py @@ -170,8 +170,6 @@ def _default_runnable_builder( def _default_instrumentor_builder(project_id: str): from agentplatform._genai import _agent_engines_utils - cloud_trace_exporter = _agent_engines_utils._import_cloud_trace_exporter_or_warn() - cloud_trace_v2 = _agent_engines_utils._import_cloud_trace_v2_or_warn() openinference_langchain = ( _agent_engines_utils._import_openinference_langchain_or_warn() ) @@ -181,21 +179,38 @@ def _default_instrumentor_builder(project_id: str): ) if all( ( - cloud_trace_exporter, - cloud_trace_v2, openinference_langchain, opentelemetry, opentelemetry_sdk_trace, ) ): - import google.auth + try: + import opentelemetry.exporter.otlp.proto.http.trace_exporter + import opentelemetry.sdk.resources + import google.auth + import google.auth.transport.requests + except (ImportError, AttributeError): + from google.cloud.aiplatform import base + + _LOGGER = base.Logger(__name__) + _LOGGER.warning( + "enable_tracing=True but proceeding with tracing disabled " + "because opentelemetry-exporter-otlp-proto-http is not installed." + ) + return None credentials, _ = google.auth.default() - span_exporter = cloud_trace_exporter.CloudTraceSpanExporter( - project_id=project_id, - client=cloud_trace_v2.TraceServiceClient( - credentials=credentials.with_quota_project(project_id), - ), + + resource = opentelemetry.sdk.resources.Resource.create( + attributes={"gcp.project_id": project_id}, + ).merge(opentelemetry.sdk.resources.OTELResourceDetector().detect()) + 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", + ) ) span_processor: SpanProcessor = ( opentelemetry_sdk_trace.export.SimpleSpanProcessor( @@ -219,13 +234,14 @@ def _default_instrumentor_builder(project_id: str): "OTEL_PYTHON_TRACER_PROVIDER, _TRACER_PROVIDER, " "or _PROXY_TRACER_PROVIDER." ) - tracer_provider = opentelemetry_sdk_trace.TracerProvider() + tracer_provider = opentelemetry_sdk_trace.TracerProvider(resource=resource) opentelemetry.trace.set_tracer_provider(tracer_provider) + # Avoids AttributeError: # 'ProxyTracerProvider' and 'NoOpTracerProvider' objects has no # attribute 'add_span_processor'. if _agent_engines_utils.is_noop_or_proxy_tracer_provider(tracer_provider): - tracer_provider = opentelemetry_sdk_trace.TracerProvider() + tracer_provider = opentelemetry_sdk_trace.TracerProvider(resource=resource) opentelemetry.trace.set_tracer_provider(tracer_provider) # Avoids OpenTelemetry client already exists error. _override_active_span_processor( diff --git a/agentplatform/agent_engines/templates/llama_index.py b/agentplatform/agent_engines/templates/llama_index.py index d9147fccef..d0c5ed258e 100644 --- a/agentplatform/agent_engines/templates/llama_index.py +++ b/agentplatform/agent_engines/templates/llama_index.py @@ -388,10 +388,6 @@ def set_up(self): _agent_engines_utils, ) - cloud_trace_exporter = ( - _agent_engines_utils._import_cloud_trace_exporter_or_warn() - ) - cloud_trace_v2 = _agent_engines_utils._import_cloud_trace_v2_or_warn() openinference_llama_index = ( _agent_engines_utils._import_openinference_llama_index_or_warn() ) @@ -399,23 +395,34 @@ def set_up(self): opentelemetry_sdk_trace = ( _agent_engines_utils._import_opentelemetry_sdk_trace_or_warn() ) + try: + import opentelemetry.exporter.otlp.proto.http.trace_exporter + import opentelemetry.sdk.resources + import google.auth + import google.auth.transport.requests + + _otlp_span_exporter_module = ( + opentelemetry.exporter.otlp.proto.http.trace_exporter + ) + except (ImportError, AttributeError): + _otlp_span_exporter_module = None if all( ( - cloud_trace_exporter, - cloud_trace_v2, + _otlp_span_exporter_module, openinference_llama_index, opentelemetry, opentelemetry_sdk_trace, ) ): - import google.auth - credentials, _ = google.auth.default() - span_exporter = cloud_trace_exporter.CloudTraceSpanExporter( - project_id=self._project, - client=cloud_trace_v2.TraceServiceClient( - credentials=credentials.with_quota_project(self._project), + resource = opentelemetry.sdk.resources.Resource.create( + attributes={"gcp.project_id": self._project}, + ).merge(opentelemetry.sdk.resources.OTELResourceDetector().detect()) + span_exporter = _otlp_span_exporter_module.OTLPSpanExporter( + session=google.auth.transport.requests.AuthorizedSession( + credentials=credentials ), + endpoint="https://telemetry.googleapis.com/v1/traces", ) span_processor: SpanProcessor = ( opentelemetry_sdk_trace.export.SimpleSpanProcessor( @@ -442,7 +449,9 @@ def set_up(self): "OTEL_PYTHON_TRACER_PROVIDER, _TRACER_PROVIDER, " "or _PROXY_TRACER_PROVIDER." ) - tracer_provider = opentelemetry_sdk_trace.TracerProvider() + tracer_provider = opentelemetry_sdk_trace.TracerProvider( + resource=resource + ) opentelemetry.trace.set_tracer_provider(tracer_provider) # Avoids AttributeError: # 'ProxyTracerProvider' and 'NoOpTracerProvider' objects has no @@ -450,7 +459,9 @@ def set_up(self): if _agent_engines_utils.is_noop_or_proxy_tracer_provider( tracer_provider ): - tracer_provider = opentelemetry_sdk_trace.TracerProvider() + tracer_provider = opentelemetry_sdk_trace.TracerProvider( + resource=resource + ) opentelemetry.trace.set_tracer_provider(tracer_provider) # Avoids OpenTelemetry client already exists error. _override_active_span_processor( diff --git a/setup.py b/setup.py index 4df6ddf59d..6c65920a57 100644 --- a/setup.py +++ b/setup.py @@ -150,10 +150,8 @@ reasoning_engine_extra_require = [ "cloudpickle >= 3.0, < 4.0", - "google-cloud-trace < 2", "opentelemetry-sdk < 2", "opentelemetry-exporter-gcp-logging >= 1.11.0a0, < 2.0.0", - "opentelemetry-exporter-gcp-trace < 2", "opentelemetry-exporter-otlp-proto-http < 2", "opentelemetry-instrumentation-google-genai>=0.3b0, <1.0.0", # TODO(b/538550724): update to stable version of @@ -167,11 +165,9 @@ agent_engines_extra_require = [ "packaging >= 24.0", "cloudpickle >= 3.0, < 4.0", - "google-cloud-trace < 2", "google-cloud-logging < 4", "opentelemetry-sdk < 2", "opentelemetry-exporter-gcp-logging >= 1.11.0a0, < 2.0.0", - "opentelemetry-exporter-gcp-trace < 2", "opentelemetry-exporter-otlp-proto-http < 2", "pydantic >= 2.11.1, < 3", "typing_extensions", diff --git a/tests/unit/agentplatform/frameworks/test_frameworks_ag2.py b/tests/unit/agentplatform/frameworks/test_frameworks_ag2.py index b8e2337f51..1b5fbd2eb3 100644 --- a/tests/unit/agentplatform/frameworks/test_frameworks_ag2.py +++ b/tests/unit/agentplatform/frameworks/test_frameworks_ag2.py @@ -101,15 +101,6 @@ def to_json_serializable_autogen_object_mock(): yield to_json_serializable_autogen_object_mock -@pytest.fixture -def cloud_trace_exporter_mock(): - with mock.patch.object( - _agent_engines_utils, - "_import_cloud_trace_exporter_or_warn", - ) as cloud_trace_exporter_mock: - yield cloud_trace_exporter_mock - - @pytest.fixture def tracer_provider_mock(): with mock.patch("opentelemetry.sdk.trace.TracerProvider") as tracer_provider_mock: @@ -124,6 +115,46 @@ def simple_span_processor_mock(): yield simple_span_processor_mock +@pytest.fixture +def otlp_span_exporter_mock(): + import opentelemetry.exporter.otlp.proto.http.trace_exporter + + with mock.patch.object( + opentelemetry.exporter.otlp.proto.http.trace_exporter, + "OTLPSpanExporter", + ) as otlp_span_exporter_mock: + yield otlp_span_exporter_mock + + +@pytest.fixture +def resource_create_mock(): + import opentelemetry.sdk.resources + + with mock.patch.object( + opentelemetry.sdk.resources.Resource, "create" + ) as resource_create_mock: + yield resource_create_mock + + +@pytest.fixture +def otel_resource_detector_mock(): + import opentelemetry.sdk.resources + + with mock.patch.object( + opentelemetry.sdk.resources, "OTELResourceDetector" + ) as otel_resource_detector_mock: + yield otel_resource_detector_mock + + +@pytest.fixture +def is_noop_or_proxy_tracer_provider_mock(): + with mock.patch.object( + _agent_engines_utils, "is_noop_or_proxy_tracer_provider" + ) as is_noop_or_proxy_tracer_provider_mock: + is_noop_or_proxy_tracer_provider_mock.return_value = True + yield is_noop_or_proxy_tracer_provider_mock + + @pytest.fixture def autogen_instrumentor_mock(): with mock.patch.object( @@ -257,7 +288,6 @@ def test_query(self, to_json_serializable_autogen_object_mock): def test_enable_tracing( self, caplog, - cloud_trace_exporter_mock, tracer_provider_mock, simple_span_processor_mock, autogen_instrumentor_mock, @@ -285,6 +315,53 @@ def test_enable_tracing_warning(self, caplog, autogen_instrumentor_none_mock): # agent.set_up() # assert "enable_tracing=True but proceeding with tracing disabled" in caplog.text + def test_tracing_setup( + self, + tracer_provider_mock, + simple_span_processor_mock, + otlp_span_exporter_mock, + resource_create_mock, + otel_resource_detector_mock, + is_noop_or_proxy_tracer_provider_mock, + autogen_instrumentor_mock, + ): + agent = agent_engines.AG2Agent( + model=_TEST_MODEL, + runnable_name=_TEST_RUNNABLE_NAME, + runnable_builder=lambda **kwargs: kwargs, + enable_tracing=True, + ) + assert agent._tmpl_attrs.get("instrumentor") is None + + agent.set_up() + + otlp_span_exporter_mock.assert_called_once_with( + session=mock.ANY, + endpoint="https://telemetry.googleapis.com/v1/traces", + ) + resource_create_mock.assert_called_with( + attributes={"gcp.project_id": _TEST_PROJECT}, + ) + tracer_provider_mock.assert_called_with( + resource=resource_create_mock.return_value.merge.return_value, + ) + simple_span_processor_mock.assert_called_once_with( + span_exporter=otlp_span_exporter_mock.return_value, + ) + tracer_provider_mock.return_value.add_span_processor.assert_called_once_with( + simple_span_processor_mock.return_value, + ) + autogen_instrumentor_class = ( + autogen_instrumentor_mock.return_value.AutogenInstrumentor + ) + autogen_instrumentor_class.assert_called_once_with() + autogen_instrumentor_class.return_value.uninstrument.assert_called_once_with() + autogen_instrumentor_class.return_value.instrument.assert_called_once_with() + assert ( + agent._tmpl_attrs.get("instrumentor") + is autogen_instrumentor_class.return_value + ) + def _return_input_no_typing(input_): """Returns input back to user.""" diff --git a/tests/unit/agentplatform/frameworks/test_frameworks_langchain.py b/tests/unit/agentplatform/frameworks/test_frameworks_langchain.py index 8d50518aff..2ac9fe375c 100644 --- a/tests/unit/agentplatform/frameworks/test_frameworks_langchain.py +++ b/tests/unit/agentplatform/frameworks/test_frameworks_langchain.py @@ -87,15 +87,6 @@ def langchain_dump_mock(): yield langchain_dump_mock -@pytest.fixture -def cloud_trace_exporter_mock(): - with mock.patch.object( - _agent_engines_utils, - "_import_cloud_trace_exporter_or_warn", - ) as cloud_trace_exporter_mock: - yield cloud_trace_exporter_mock - - @pytest.fixture def tracer_provider_mock(): with mock.patch("opentelemetry.sdk.trace.TracerProvider") as tracer_provider_mock: @@ -110,6 +101,46 @@ def simple_span_processor_mock(): yield simple_span_processor_mock +@pytest.fixture +def otlp_span_exporter_mock(): + import opentelemetry.exporter.otlp.proto.http.trace_exporter + + with mock.patch.object( + opentelemetry.exporter.otlp.proto.http.trace_exporter, + "OTLPSpanExporter", + ) as otlp_span_exporter_mock: + yield otlp_span_exporter_mock + + +@pytest.fixture +def resource_create_mock(): + import opentelemetry.sdk.resources + + with mock.patch.object( + opentelemetry.sdk.resources.Resource, "create" + ) as resource_create_mock: + yield resource_create_mock + + +@pytest.fixture +def otel_resource_detector_mock(): + import opentelemetry.sdk.resources + + with mock.patch.object( + opentelemetry.sdk.resources, "OTELResourceDetector" + ) as otel_resource_detector_mock: + yield otel_resource_detector_mock + + +@pytest.fixture +def is_noop_or_proxy_tracer_provider_mock(): + with mock.patch.object( + _agent_engines_utils, "is_noop_or_proxy_tracer_provider" + ) as is_noop_or_proxy_tracer_provider_mock: + is_noop_or_proxy_tracer_provider_mock.return_value = True + yield is_noop_or_proxy_tracer_provider_mock + + @pytest.fixture def langchain_instrumentor_mock(): with mock.patch.object( @@ -235,7 +266,6 @@ def test_stream_query(self, langchain_dump_mock): def test_enable_tracing( self, caplog, - cloud_trace_exporter_mock, tracer_provider_mock, simple_span_processor_mock, langchain_instrumentor_mock, @@ -268,6 +298,54 @@ def test_enable_tracing_warning(self, caplog, langchain_instrumentor_none_mock): # agent.set_up() # assert "enable_tracing=True but proceeding with tracing disabled" in caplog.text + def test_tracing_setup( + self, + tracer_provider_mock, + simple_span_processor_mock, + otlp_span_exporter_mock, + resource_create_mock, + otel_resource_detector_mock, + is_noop_or_proxy_tracer_provider_mock, + langchain_instrumentor_mock, + ): + agent = agent_engines.LangchainAgent( + model=_TEST_MODEL, + prompt=self.prompt, + output_parser=self.output_parser, + runnable_builder=lambda **kwargs: kwargs, + enable_tracing=True, + ) + assert agent._tmpl_attrs.get("instrumentor") is None + + agent.set_up() + + otlp_span_exporter_mock.assert_called_once_with( + session=mock.ANY, + endpoint="https://telemetry.googleapis.com/v1/traces", + ) + resource_create_mock.assert_called_with( + attributes={"gcp.project_id": _TEST_PROJECT}, + ) + tracer_provider_mock.assert_called_with( + resource=resource_create_mock.return_value.merge.return_value, + ) + simple_span_processor_mock.assert_called_once_with( + span_exporter=otlp_span_exporter_mock.return_value, + ) + tracer_provider_mock.return_value.add_span_processor.assert_called_once_with( + simple_span_processor_mock.return_value, + ) + langchain_instrumentor_class = ( + langchain_instrumentor_mock.return_value.LangChainInstrumentor + ) + langchain_instrumentor_class.assert_called_once_with() + langchain_instrumentor_class.return_value.uninstrument.assert_called_once_with() + langchain_instrumentor_class.return_value.instrument.assert_called_once_with() + assert ( + agent._tmpl_attrs.get("instrumentor") + is langchain_instrumentor_class.return_value + ) + def _return_input_no_typing(input_): """Returns input back to user.""" @@ -284,7 +362,9 @@ def test_raise_untyped_input_args(self, agentplatform_init_mock): class TestSystemInstructionAndPromptRaisesErrors: - def test_raise_both_system_instruction_and_prompt_error(self, agentplatform_init_mock): + def test_raise_both_system_instruction_and_prompt_error( + self, agentplatform_init_mock + ): with pytest.raises( ValueError, match=r"Only one of `prompt` or `system_instruction` should be specified.", diff --git a/tests/unit/agentplatform/frameworks/test_frameworks_langgraph.py b/tests/unit/agentplatform/frameworks/test_frameworks_langgraph.py index 82fa32d53b..70948abb63 100644 --- a/tests/unit/agentplatform/frameworks/test_frameworks_langgraph.py +++ b/tests/unit/agentplatform/frameworks/test_frameworks_langgraph.py @@ -98,15 +98,6 @@ def langchain_dump_mock(): yield langchain_dump_mock -@pytest.fixture -def cloud_trace_exporter_mock(): - with mock.patch.object( - _agent_engines_utils, - "_import_cloud_trace_exporter_or_warn", - ) as cloud_trace_exporter_mock: - yield cloud_trace_exporter_mock - - @pytest.fixture def tracer_provider_mock(): with mock.patch("opentelemetry.sdk.trace.TracerProvider") as tracer_provider_mock: @@ -121,6 +112,46 @@ def simple_span_processor_mock(): yield simple_span_processor_mock +@pytest.fixture +def otlp_span_exporter_mock(): + import opentelemetry.exporter.otlp.proto.http.trace_exporter + + with mock.patch.object( + opentelemetry.exporter.otlp.proto.http.trace_exporter, + "OTLPSpanExporter", + ) as otlp_span_exporter_mock: + yield otlp_span_exporter_mock + + +@pytest.fixture +def resource_create_mock(): + import opentelemetry.sdk.resources + + with mock.patch.object( + opentelemetry.sdk.resources.Resource, "create" + ) as resource_create_mock: + yield resource_create_mock + + +@pytest.fixture +def otel_resource_detector_mock(): + import opentelemetry.sdk.resources + + with mock.patch.object( + opentelemetry.sdk.resources, "OTELResourceDetector" + ) as otel_resource_detector_mock: + yield otel_resource_detector_mock + + +@pytest.fixture +def is_noop_or_proxy_tracer_provider_mock(): + with mock.patch.object( + _agent_engines_utils, "is_noop_or_proxy_tracer_provider" + ) as is_noop_or_proxy_tracer_provider_mock: + is_noop_or_proxy_tracer_provider_mock.return_value = True + yield is_noop_or_proxy_tracer_provider_mock + + @pytest.fixture def langchain_instrumentor_mock(): with mock.patch.object( @@ -233,7 +264,6 @@ def test_stream_query(self, langchain_dump_mock): def test_enable_tracing( self, caplog, - cloud_trace_exporter_mock, tracer_provider_mock, simple_span_processor_mock, langchain_instrumentor_mock, @@ -256,6 +286,52 @@ def test_enable_tracing_warning(self, caplog, langchain_instrumentor_none_mock): # agent.set_up() # assert "enable_tracing=True but proceeding with tracing disabled" in caplog.text + def test_tracing_setup( + self, + tracer_provider_mock, + simple_span_processor_mock, + otlp_span_exporter_mock, + resource_create_mock, + otel_resource_detector_mock, + is_noop_or_proxy_tracer_provider_mock, + langchain_instrumentor_mock, + ): + agent = agent_engines.LanggraphAgent( + model=_TEST_MODEL, + runnable_builder=lambda **kwargs: kwargs, + enable_tracing=True, + ) + assert agent._tmpl_attrs.get("instrumentor") is None + + agent.set_up() + + otlp_span_exporter_mock.assert_called_once_with( + session=mock.ANY, + endpoint="https://telemetry.googleapis.com/v1/traces", + ) + resource_create_mock.assert_called_with( + attributes={"gcp.project_id": _TEST_PROJECT}, + ) + tracer_provider_mock.assert_called_with( + resource=resource_create_mock.return_value.merge.return_value, + ) + simple_span_processor_mock.assert_called_once_with( + span_exporter=otlp_span_exporter_mock.return_value, + ) + tracer_provider_mock.return_value.add_span_processor.assert_called_once_with( + simple_span_processor_mock.return_value, + ) + langchain_instrumentor_class = ( + langchain_instrumentor_mock.return_value.LangChainInstrumentor + ) + langchain_instrumentor_class.assert_called_once_with() + langchain_instrumentor_class.return_value.uninstrument.assert_called_once_with() + langchain_instrumentor_class.return_value.instrument.assert_called_once_with() + assert ( + agent._tmpl_attrs.get("instrumentor") + is langchain_instrumentor_class.return_value + ) + def test_get_state_history_empty(self): agent = agent_engines.LanggraphAgent(model=_TEST_MODEL) agent._tmpl_attrs["runnable"] = mock.Mock() diff --git a/tests/unit/agentplatform/frameworks/test_frameworks_llama_index.py b/tests/unit/agentplatform/frameworks/test_frameworks_llama_index.py index aaef28749b..6302c852e4 100644 --- a/tests/unit/agentplatform/frameworks/test_frameworks_llama_index.py +++ b/tests/unit/agentplatform/frameworks/test_frameworks_llama_index.py @@ -22,7 +22,9 @@ from agentplatform.agent_engines.templates import ( llama_index, ) -from agentplatform._genai.agent_engines import _agent_engines_utils +from agentplatform._genai.agent_engines import ( + _agent_engines_utils, +) from llama_index.core import prompts from llama_index.core.base.llms import types @@ -68,15 +70,6 @@ def model_builder_mock(): yield model_builder_mock -@pytest.fixture -def cloud_trace_exporter_mock(): - with mock.patch.object( - _agent_engines_utils, - "_import_cloud_trace_exporter_or_warn", - ) as cloud_trace_exporter_mock: - yield cloud_trace_exporter_mock - - @pytest.fixture def tracer_provider_mock(): with mock.patch("opentelemetry.sdk.trace.TracerProvider") as tracer_provider_mock: @@ -91,6 +84,46 @@ def simple_span_processor_mock(): yield simple_span_processor_mock +@pytest.fixture +def otlp_span_exporter_mock(): + import opentelemetry.exporter.otlp.proto.http.trace_exporter + + with mock.patch.object( + opentelemetry.exporter.otlp.proto.http.trace_exporter, + "OTLPSpanExporter", + ) as otlp_span_exporter_mock: + yield otlp_span_exporter_mock + + +@pytest.fixture +def resource_create_mock(): + import opentelemetry.sdk.resources + + with mock.patch.object( + opentelemetry.sdk.resources.Resource, "create" + ) as resource_create_mock: + yield resource_create_mock + + +@pytest.fixture +def otel_resource_detector_mock(): + import opentelemetry.sdk.resources + + with mock.patch.object( + opentelemetry.sdk.resources, "OTELResourceDetector" + ) as otel_resource_detector_mock: + yield otel_resource_detector_mock + + +@pytest.fixture +def is_noop_or_proxy_tracer_provider_mock(): + with mock.patch.object( + _agent_engines_utils, "is_noop_or_proxy_tracer_provider" + ) as is_noop_or_proxy_tracer_provider_mock: + is_noop_or_proxy_tracer_provider_mock.return_value = True + yield is_noop_or_proxy_tracer_provider_mock + + @pytest.fixture def llama_index_instrumentor_mock(): with mock.patch.object( @@ -228,7 +261,6 @@ def test_query_with_batch_input(self, json_loads_mock, nest_asyncio_apply_mock): def test_enable_tracing( self, caplog, - cloud_trace_exporter_mock, tracer_provider_mock, simple_span_processor_mock, llama_index_instrumentor_mock, @@ -259,6 +291,54 @@ def test_enable_tracing_warning(self, caplog, llama_index_instrumentor_none_mock # agent.set_up() # assert "enable_tracing=True but proceeding with tracing disabled" in caplog.text + def test_tracing_setup( + self, + tracer_provider_mock, + simple_span_processor_mock, + otlp_span_exporter_mock, + resource_create_mock, + otel_resource_detector_mock, + is_noop_or_proxy_tracer_provider_mock, + llama_index_instrumentor_mock, + ): + agent = llama_index.LlamaIndexQueryPipelineAgent( + model=_TEST_MODEL, + prompt=self.prompt, + model_builder=lambda **kwargs: kwargs, + runnable_builder=lambda **kwargs: kwargs, + enable_tracing=True, + ) + assert agent._instrumentor is None + + agent.set_up() + + otlp_span_exporter_mock.assert_called_once_with( + session=mock.ANY, + endpoint="https://telemetry.googleapis.com/v1/traces", + ) + resource_create_mock.assert_called_with( + attributes={"gcp.project_id": _TEST_PROJECT}, + ) + tracer_provider_mock.assert_called_with( + resource=resource_create_mock.return_value.merge.return_value, + ) + simple_span_processor_mock.assert_called_once_with( + span_exporter=otlp_span_exporter_mock.return_value, + ) + tracer_provider_mock.return_value.add_span_processor.assert_called_once_with( + simple_span_processor_mock.return_value, + ) + llama_index_instrumentor_class = ( + llama_index_instrumentor_mock.return_value.LlamaIndexInstrumentor + ) + llama_index_instrumentor_class.assert_called_once_with() + llama_index_instrumentor_class.return_value.uninstrument.assert_called_once_with() + llama_index_instrumentor_class.return_value.instrument.assert_called_once_with() + assert agent._instrumentor is llama_index_instrumentor_class.return_value + llama_index_instrumentor_class.assert_called_once_with() + llama_index_instrumentor_class.return_value.instrument.assert_called_once_with() + assert agent._instrumentor is llama_index_instrumentor_class.return_value + class TestToJsonSerializableLlamaIndexObject: """Tests for `_utils.to_json_serializable_llama_index_object`.""" @@ -285,7 +365,9 @@ def test_llama_index_response(self): "source_nodes": ['{"name": "model1"}', '{"name": "model2"}'], "metadata": {"key": "value"}, } - got = _agent_engines_utils.to_json_serializable_llama_index_object(mock_response) + got = _agent_engines_utils.to_json_serializable_llama_index_object( + mock_response + ) assert got == want def test_llama_index_chat_response(self): @@ -298,7 +380,9 @@ def test_llama_index_chat_response(self): ) want = {"content": "chat message"} - got = _agent_engines_utils.to_json_serializable_llama_index_object(mock_chat_response) + got = _agent_engines_utils.to_json_serializable_llama_index_object( + mock_chat_response + ) assert got == want def test_llama_index_base_model(self): @@ -308,7 +392,9 @@ def test_llama_index_base_model(self): mock_base_model.model_dump_json = lambda: '{"name": "test_model"}' want = {"name": "test_model"} - got = _agent_engines_utils.to_json_serializable_llama_index_object(mock_base_model) + got = _agent_engines_utils.to_json_serializable_llama_index_object( + mock_base_model + ) assert got == want def test_sequence_of_llama_index_base_model(self): @@ -323,7 +409,9 @@ def test_sequence_of_llama_index_base_model(self): mock_base_model_list = [mock_base_model1, mock_base_model2] want = [{"name": "test_model1"}, {"name": "test_model2"}] - got = _agent_engines_utils.to_json_serializable_llama_index_object(mock_base_model_list) + got = _agent_engines_utils.to_json_serializable_llama_index_object( + mock_base_model_list + ) assert got == want def test_sequence_of_mixed_types(self): diff --git a/tests/unit/vertex_ag2/test_agent_engine_templates_ag2.py b/tests/unit/vertex_ag2/test_agent_engine_templates_ag2.py index a495488e83..69d95a63fb 100644 --- a/tests/unit/vertex_ag2/test_agent_engine_templates_ag2.py +++ b/tests/unit/vertex_ag2/test_agent_engine_templates_ag2.py @@ -101,15 +101,6 @@ def to_json_serializable_autogen_object_mock(): yield to_json_serializable_autogen_object_mock -@pytest.fixture -def cloud_trace_exporter_mock(): - with mock.patch.object( - _utils, - "_import_cloud_trace_exporter_or_warn", - ) as cloud_trace_exporter_mock: - yield cloud_trace_exporter_mock - - @pytest.fixture def tracer_provider_mock(): with mock.patch("opentelemetry.sdk.trace.TracerProvider") as tracer_provider_mock: @@ -124,6 +115,46 @@ def simple_span_processor_mock(): yield simple_span_processor_mock +@pytest.fixture +def otlp_span_exporter_mock(): + import opentelemetry.exporter.otlp.proto.http.trace_exporter + + with mock.patch.object( + opentelemetry.exporter.otlp.proto.http.trace_exporter, + "OTLPSpanExporter", + ) as otlp_span_exporter_mock: + yield otlp_span_exporter_mock + + +@pytest.fixture +def resource_create_mock(): + import opentelemetry.sdk.resources + + with mock.patch.object( + opentelemetry.sdk.resources.Resource, "create" + ) as resource_create_mock: + yield resource_create_mock + + +@pytest.fixture +def otel_resource_detector_mock(): + import opentelemetry.sdk.resources + + with mock.patch.object( + opentelemetry.sdk.resources, "OTELResourceDetector" + ) as otel_resource_detector_mock: + yield otel_resource_detector_mock + + +@pytest.fixture +def is_noop_or_proxy_tracer_provider_mock(): + with mock.patch.object( + _utils, "is_noop_or_proxy_tracer_provider" + ) as is_noop_or_proxy_tracer_provider_mock: + is_noop_or_proxy_tracer_provider_mock.return_value = True + yield is_noop_or_proxy_tracer_provider_mock + + @pytest.fixture def autogen_instrumentor_mock(): with mock.patch.object( @@ -257,7 +288,6 @@ def test_query(self, to_json_serializable_autogen_object_mock): def test_enable_tracing( self, caplog, - cloud_trace_exporter_mock, tracer_provider_mock, simple_span_processor_mock, autogen_instrumentor_mock, @@ -285,6 +315,53 @@ def test_enable_tracing_warning(self, caplog, autogen_instrumentor_none_mock): # agent.set_up() # assert "enable_tracing=True but proceeding with tracing disabled" in caplog.text + def test_tracing_setup( + self, + tracer_provider_mock, + simple_span_processor_mock, + otlp_span_exporter_mock, + resource_create_mock, + otel_resource_detector_mock, + is_noop_or_proxy_tracer_provider_mock, + autogen_instrumentor_mock, + ): + agent = agent_engines.AG2Agent( + model=_TEST_MODEL, + runnable_name=_TEST_RUNNABLE_NAME, + runnable_builder=lambda **kwargs: kwargs, + enable_tracing=True, + ) + assert agent._tmpl_attrs.get("instrumentor") is None + + agent.set_up() + + otlp_span_exporter_mock.assert_called_once_with( + session=mock.ANY, + endpoint="https://telemetry.googleapis.com/v1/traces", + ) + resource_create_mock.assert_called_with( + attributes={"gcp.project_id": _TEST_PROJECT}, + ) + tracer_provider_mock.assert_called_with( + resource=resource_create_mock.return_value.merge.return_value, + ) + simple_span_processor_mock.assert_called_once_with( + span_exporter=otlp_span_exporter_mock.return_value, + ) + tracer_provider_mock.return_value.add_span_processor.assert_called_once_with( + simple_span_processor_mock.return_value, + ) + autogen_instrumentor_class = ( + autogen_instrumentor_mock.return_value.AutogenInstrumentor + ) + autogen_instrumentor_class.assert_called_once_with() + autogen_instrumentor_class.return_value.uninstrument.assert_called_once_with() + autogen_instrumentor_class.return_value.instrument.assert_called_once_with() + assert ( + agent._tmpl_attrs.get("instrumentor") + is autogen_instrumentor_class.return_value + ) + def _return_input_no_typing(input_): """Returns input back to user.""" diff --git a/tests/unit/vertex_ag2/test_reasoning_engine_templates_ag2.py b/tests/unit/vertex_ag2/test_reasoning_engine_templates_ag2.py index 62145f6c94..e6f284be1d 100644 --- a/tests/unit/vertex_ag2/test_reasoning_engine_templates_ag2.py +++ b/tests/unit/vertex_ag2/test_reasoning_engine_templates_ag2.py @@ -81,15 +81,6 @@ def dataclasses_asdict_mock(): yield dataclasses_asdict_mock -@pytest.fixture -def cloud_trace_exporter_mock(): - with mock.patch.object( - _utils, - "_import_cloud_trace_exporter_or_warn", - ) as cloud_trace_exporter_mock: - yield cloud_trace_exporter_mock - - @pytest.fixture def tracer_provider_mock(): with mock.patch("opentelemetry.sdk.trace.TracerProvider") as tracer_provider_mock: @@ -104,6 +95,46 @@ def simple_span_processor_mock(): yield simple_span_processor_mock +@pytest.fixture +def otlp_span_exporter_mock(): + import opentelemetry.exporter.otlp.proto.http.trace_exporter + + with mock.patch.object( + opentelemetry.exporter.otlp.proto.http.trace_exporter, + "OTLPSpanExporter", + ) as otlp_span_exporter_mock: + yield otlp_span_exporter_mock + + +@pytest.fixture +def resource_create_mock(): + import opentelemetry.sdk.resources + + with mock.patch.object( + opentelemetry.sdk.resources.Resource, "create" + ) as resource_create_mock: + yield resource_create_mock + + +@pytest.fixture +def otel_resource_detector_mock(): + import opentelemetry.sdk.resources + + with mock.patch.object( + opentelemetry.sdk.resources, "OTELResourceDetector" + ) as otel_resource_detector_mock: + yield otel_resource_detector_mock + + +@pytest.fixture +def is_noop_or_proxy_tracer_provider_mock(): + with mock.patch.object( + _utils, "is_noop_or_proxy_tracer_provider" + ) as is_noop_or_proxy_tracer_provider_mock: + is_noop_or_proxy_tracer_provider_mock.return_value = True + yield is_noop_or_proxy_tracer_provider_mock + + @pytest.fixture def autogen_instrumentor_mock(): with mock.patch.object( @@ -223,7 +254,6 @@ def test_query(self, dataclasses_asdict_mock): def test_enable_tracing( self, caplog, - cloud_trace_exporter_mock, tracer_provider_mock, simple_span_processor_mock, autogen_instrumentor_mock, @@ -251,6 +281,50 @@ def test_enable_tracing_warning(self, caplog, autogen_instrumentor_none_mock): # agent.set_up() # assert "enable_tracing=True but proceeding with tracing disabled" in caplog.text + def test_tracing_setup( + self, + tracer_provider_mock, + simple_span_processor_mock, + otlp_span_exporter_mock, + resource_create_mock, + otel_resource_detector_mock, + is_noop_or_proxy_tracer_provider_mock, + autogen_instrumentor_mock, + ): + agent = reasoning_engines.AG2Agent( + model=_TEST_MODEL, + runnable_name=_TEST_RUNNABLE_NAME, + runnable_builder=lambda **kwargs: kwargs, + enable_tracing=True, + ) + assert agent._instrumentor is None + + agent.set_up() + + otlp_span_exporter_mock.assert_called_once_with( + session=mock.ANY, + endpoint="https://telemetry.googleapis.com/v1/traces", + ) + resource_create_mock.assert_called_with( + attributes={"gcp.project_id": _TEST_PROJECT}, + ) + tracer_provider_mock.assert_called_with( + resource=resource_create_mock.return_value.merge.return_value, + ) + simple_span_processor_mock.assert_called_once_with( + span_exporter=otlp_span_exporter_mock.return_value, + ) + tracer_provider_mock.return_value.add_span_processor.assert_called_once_with( + simple_span_processor_mock.return_value, + ) + autogen_instrumentor_class = ( + autogen_instrumentor_mock.return_value.AutogenInstrumentor + ) + autogen_instrumentor_class.assert_called_once_with() + autogen_instrumentor_class.return_value.uninstrument.assert_called_once_with() + autogen_instrumentor_class.return_value.instrument.assert_called_once_with() + assert agent._instrumentor is autogen_instrumentor_class.return_value + def _return_input_no_typing(input_): """Returns input back to user.""" diff --git a/tests/unit/vertex_langchain/test_agent_engine_templates_langchain.py b/tests/unit/vertex_langchain/test_agent_engine_templates_langchain.py index 46eafd1c34..ea7faa4697 100644 --- a/tests/unit/vertex_langchain/test_agent_engine_templates_langchain.py +++ b/tests/unit/vertex_langchain/test_agent_engine_templates_langchain.py @@ -88,15 +88,6 @@ def langchain_dump_mock(): yield langchain_dump_mock -@pytest.fixture -def cloud_trace_exporter_mock(): - with mock.patch.object( - _utils, - "_import_cloud_trace_exporter_or_warn", - ) as cloud_trace_exporter_mock: - yield cloud_trace_exporter_mock - - @pytest.fixture def tracer_provider_mock(): with mock.patch("opentelemetry.sdk.trace.TracerProvider") as tracer_provider_mock: @@ -111,6 +102,46 @@ def simple_span_processor_mock(): yield simple_span_processor_mock +@pytest.fixture +def otlp_span_exporter_mock(): + import opentelemetry.exporter.otlp.proto.http.trace_exporter + + with mock.patch.object( + opentelemetry.exporter.otlp.proto.http.trace_exporter, + "OTLPSpanExporter", + ) as otlp_span_exporter_mock: + yield otlp_span_exporter_mock + + +@pytest.fixture +def resource_create_mock(): + import opentelemetry.sdk.resources + + with mock.patch.object( + opentelemetry.sdk.resources.Resource, "create" + ) as resource_create_mock: + yield resource_create_mock + + +@pytest.fixture +def otel_resource_detector_mock(): + import opentelemetry.sdk.resources + + with mock.patch.object( + opentelemetry.sdk.resources, "OTELResourceDetector" + ) as otel_resource_detector_mock: + yield otel_resource_detector_mock + + +@pytest.fixture +def is_noop_or_proxy_tracer_provider_mock(): + with mock.patch.object( + _utils, "is_noop_or_proxy_tracer_provider" + ) as is_noop_or_proxy_tracer_provider_mock: + is_noop_or_proxy_tracer_provider_mock.return_value = True + yield is_noop_or_proxy_tracer_provider_mock + + @pytest.fixture def langchain_instrumentor_mock(): with mock.patch.object( @@ -237,7 +268,6 @@ def test_stream_query(self, langchain_dump_mock): def test_enable_tracing( self, caplog, - cloud_trace_exporter_mock, tracer_provider_mock, simple_span_processor_mock, langchain_instrumentor_mock, @@ -270,6 +300,54 @@ def test_enable_tracing_warning(self, caplog, langchain_instrumentor_none_mock): # agent.set_up() # assert "enable_tracing=True but proceeding with tracing disabled" in caplog.text + def test_tracing_setup( + self, + tracer_provider_mock, + simple_span_processor_mock, + otlp_span_exporter_mock, + resource_create_mock, + otel_resource_detector_mock, + is_noop_or_proxy_tracer_provider_mock, + langchain_instrumentor_mock, + ): + agent = agent_engines.LangchainAgent( + model=_TEST_MODEL, + prompt=self.prompt, + output_parser=self.output_parser, + runnable_builder=lambda **kwargs: kwargs, + enable_tracing=True, + ) + assert agent._tmpl_attrs.get("instrumentor") is None + + agent.set_up() + + otlp_span_exporter_mock.assert_called_once_with( + session=mock.ANY, + endpoint="https://telemetry.googleapis.com/v1/traces", + ) + resource_create_mock.assert_called_with( + attributes={"gcp.project_id": _TEST_PROJECT}, + ) + tracer_provider_mock.assert_called_with( + resource=resource_create_mock.return_value.merge.return_value, + ) + simple_span_processor_mock.assert_called_once_with( + span_exporter=otlp_span_exporter_mock.return_value, + ) + tracer_provider_mock.return_value.add_span_processor.assert_called_once_with( + simple_span_processor_mock.return_value, + ) + langchain_instrumentor_class = ( + langchain_instrumentor_mock.return_value.LangChainInstrumentor + ) + langchain_instrumentor_class.assert_called_once_with() + langchain_instrumentor_class.return_value.uninstrument.assert_called_once_with() + langchain_instrumentor_class.return_value.instrument.assert_called_once_with() + assert ( + agent._tmpl_attrs.get("instrumentor") + is langchain_instrumentor_class.return_value + ) + def _return_input_no_typing(input_): """Returns input back to user.""" diff --git a/tests/unit/vertex_langchain/test_agent_engine_templates_langgraph.py b/tests/unit/vertex_langchain/test_agent_engine_templates_langgraph.py index 8e9dea13e5..f569cf8f15 100644 --- a/tests/unit/vertex_langchain/test_agent_engine_templates_langgraph.py +++ b/tests/unit/vertex_langchain/test_agent_engine_templates_langgraph.py @@ -98,15 +98,6 @@ def langchain_dump_mock(): yield langchain_dump_mock -@pytest.fixture -def cloud_trace_exporter_mock(): - with mock.patch.object( - _utils, - "_import_cloud_trace_exporter_or_warn", - ) as cloud_trace_exporter_mock: - yield cloud_trace_exporter_mock - - @pytest.fixture def tracer_provider_mock(): with mock.patch("opentelemetry.sdk.trace.TracerProvider") as tracer_provider_mock: @@ -121,6 +112,46 @@ def simple_span_processor_mock(): yield simple_span_processor_mock +@pytest.fixture +def otlp_span_exporter_mock(): + import opentelemetry.exporter.otlp.proto.http.trace_exporter + + with mock.patch.object( + opentelemetry.exporter.otlp.proto.http.trace_exporter, + "OTLPSpanExporter", + ) as otlp_span_exporter_mock: + yield otlp_span_exporter_mock + + +@pytest.fixture +def resource_create_mock(): + import opentelemetry.sdk.resources + + with mock.patch.object( + opentelemetry.sdk.resources.Resource, "create" + ) as resource_create_mock: + yield resource_create_mock + + +@pytest.fixture +def otel_resource_detector_mock(): + import opentelemetry.sdk.resources + + with mock.patch.object( + opentelemetry.sdk.resources, "OTELResourceDetector" + ) as otel_resource_detector_mock: + yield otel_resource_detector_mock + + +@pytest.fixture +def is_noop_or_proxy_tracer_provider_mock(): + with mock.patch.object( + _utils, "is_noop_or_proxy_tracer_provider" + ) as is_noop_or_proxy_tracer_provider_mock: + is_noop_or_proxy_tracer_provider_mock.return_value = True + yield is_noop_or_proxy_tracer_provider_mock + + @pytest.fixture def langchain_instrumentor_mock(): with mock.patch.object( @@ -233,7 +264,6 @@ def test_stream_query(self, langchain_dump_mock): def test_enable_tracing( self, caplog, - cloud_trace_exporter_mock, tracer_provider_mock, simple_span_processor_mock, langchain_instrumentor_mock, @@ -256,6 +286,52 @@ def test_enable_tracing_warning(self, caplog, langchain_instrumentor_none_mock): # agent.set_up() # assert "enable_tracing=True but proceeding with tracing disabled" in caplog.text + def test_tracing_setup( + self, + tracer_provider_mock, + simple_span_processor_mock, + otlp_span_exporter_mock, + resource_create_mock, + otel_resource_detector_mock, + is_noop_or_proxy_tracer_provider_mock, + langchain_instrumentor_mock, + ): + agent = agent_engines.LanggraphAgent( + model=_TEST_MODEL, + runnable_builder=lambda **kwargs: kwargs, + enable_tracing=True, + ) + assert agent._tmpl_attrs.get("instrumentor") is None + + agent.set_up() + + otlp_span_exporter_mock.assert_called_once_with( + session=mock.ANY, + endpoint="https://telemetry.googleapis.com/v1/traces", + ) + resource_create_mock.assert_called_with( + attributes={"gcp.project_id": _TEST_PROJECT}, + ) + tracer_provider_mock.assert_called_with( + resource=resource_create_mock.return_value.merge.return_value, + ) + simple_span_processor_mock.assert_called_once_with( + span_exporter=otlp_span_exporter_mock.return_value, + ) + tracer_provider_mock.return_value.add_span_processor.assert_called_once_with( + simple_span_processor_mock.return_value, + ) + langchain_instrumentor_class = ( + langchain_instrumentor_mock.return_value.LangChainInstrumentor + ) + langchain_instrumentor_class.assert_called_once_with() + langchain_instrumentor_class.return_value.uninstrument.assert_called_once_with() + langchain_instrumentor_class.return_value.instrument.assert_called_once_with() + assert ( + agent._tmpl_attrs.get("instrumentor") + is langchain_instrumentor_class.return_value + ) + def test_get_state_history_empty(self): agent = agent_engines.LanggraphAgent(model=_TEST_MODEL) agent._tmpl_attrs["runnable"] = mock.Mock() diff --git a/tests/unit/vertex_langchain/test_reasoning_engine_templates_langchain.py b/tests/unit/vertex_langchain/test_reasoning_engine_templates_langchain.py index 760b84b76b..42b3fe3d9a 100644 --- a/tests/unit/vertex_langchain/test_reasoning_engine_templates_langchain.py +++ b/tests/unit/vertex_langchain/test_reasoning_engine_templates_langchain.py @@ -88,15 +88,6 @@ def langchain_dump_mock(): yield langchain_dump_mock -@pytest.fixture -def cloud_trace_exporter_mock(): - with mock.patch.object( - _utils, - "_import_cloud_trace_exporter_or_warn", - ) as cloud_trace_exporter_mock: - yield cloud_trace_exporter_mock - - @pytest.fixture def tracer_provider_mock(): with mock.patch("opentelemetry.sdk.trace.TracerProvider") as tracer_provider_mock: @@ -111,6 +102,46 @@ def simple_span_processor_mock(): yield simple_span_processor_mock +@pytest.fixture +def otlp_span_exporter_mock(): + import opentelemetry.exporter.otlp.proto.http.trace_exporter + + with mock.patch.object( + opentelemetry.exporter.otlp.proto.http.trace_exporter, + "OTLPSpanExporter", + ) as otlp_span_exporter_mock: + yield otlp_span_exporter_mock + + +@pytest.fixture +def resource_create_mock(): + import opentelemetry.sdk.resources + + with mock.patch.object( + opentelemetry.sdk.resources.Resource, "create" + ) as resource_create_mock: + yield resource_create_mock + + +@pytest.fixture +def otel_resource_detector_mock(): + import opentelemetry.sdk.resources + + with mock.patch.object( + opentelemetry.sdk.resources, "OTELResourceDetector" + ) as otel_resource_detector_mock: + yield otel_resource_detector_mock + + +@pytest.fixture +def is_noop_or_proxy_tracer_provider_mock(): + with mock.patch.object( + _utils, "is_noop_or_proxy_tracer_provider" + ) as is_noop_or_proxy_tracer_provider_mock: + is_noop_or_proxy_tracer_provider_mock.return_value = True + yield is_noop_or_proxy_tracer_provider_mock + + @pytest.fixture def langchain_instrumentor_mock(): with mock.patch.object( @@ -237,7 +268,6 @@ def test_stream_query(self, langchain_dump_mock): def test_enable_tracing( self, caplog, - cloud_trace_exporter_mock, tracer_provider_mock, simple_span_processor_mock, langchain_instrumentor_mock, @@ -270,6 +300,51 @@ def test_enable_tracing_warning(self, caplog, langchain_instrumentor_none_mock): # agent.set_up() # assert "enable_tracing=True but proceeding with tracing disabled" in caplog.text + def test_tracing_setup( + self, + tracer_provider_mock, + simple_span_processor_mock, + otlp_span_exporter_mock, + resource_create_mock, + otel_resource_detector_mock, + is_noop_or_proxy_tracer_provider_mock, + langchain_instrumentor_mock, + ): + agent = reasoning_engines.LangchainAgent( + model=_TEST_MODEL, + prompt=self.prompt, + output_parser=self.output_parser, + runnable_builder=lambda **kwargs: kwargs, + enable_tracing=True, + ) + assert agent._instrumentor is None + + agent.set_up() + + otlp_span_exporter_mock.assert_called_once_with( + session=mock.ANY, + endpoint="https://telemetry.googleapis.com/v1/traces", + ) + resource_create_mock.assert_called_with( + attributes={"gcp.project_id": _TEST_PROJECT}, + ) + tracer_provider_mock.assert_called_with( + resource=resource_create_mock.return_value.merge.return_value, + ) + simple_span_processor_mock.assert_called_once_with( + span_exporter=otlp_span_exporter_mock.return_value, + ) + tracer_provider_mock.return_value.add_span_processor.assert_called_once_with( + simple_span_processor_mock.return_value, + ) + langchain_instrumentor_class = ( + langchain_instrumentor_mock.return_value.LangChainInstrumentor + ) + langchain_instrumentor_class.assert_called_once_with() + langchain_instrumentor_class.return_value.uninstrument.assert_called_once_with() + langchain_instrumentor_class.return_value.instrument.assert_called_once_with() + assert agent._instrumentor is langchain_instrumentor_class.return_value + def _return_input_no_typing(input_): """Returns input back to user.""" diff --git a/tests/unit/vertex_langchain/test_reasoning_engine_templates_langgraph.py b/tests/unit/vertex_langchain/test_reasoning_engine_templates_langgraph.py index bd5865d2f4..063090bd90 100644 --- a/tests/unit/vertex_langchain/test_reasoning_engine_templates_langgraph.py +++ b/tests/unit/vertex_langchain/test_reasoning_engine_templates_langgraph.py @@ -98,15 +98,6 @@ def langchain_dump_mock(): yield langchain_dump_mock -@pytest.fixture -def cloud_trace_exporter_mock(): - with mock.patch.object( - _utils, - "_import_cloud_trace_exporter_or_warn", - ) as cloud_trace_exporter_mock: - yield cloud_trace_exporter_mock - - @pytest.fixture def tracer_provider_mock(): with mock.patch("opentelemetry.sdk.trace.TracerProvider") as tracer_provider_mock: @@ -121,6 +112,46 @@ def simple_span_processor_mock(): yield simple_span_processor_mock +@pytest.fixture +def otlp_span_exporter_mock(): + import opentelemetry.exporter.otlp.proto.http.trace_exporter + + with mock.patch.object( + opentelemetry.exporter.otlp.proto.http.trace_exporter, + "OTLPSpanExporter", + ) as otlp_span_exporter_mock: + yield otlp_span_exporter_mock + + +@pytest.fixture +def resource_create_mock(): + import opentelemetry.sdk.resources + + with mock.patch.object( + opentelemetry.sdk.resources.Resource, "create" + ) as resource_create_mock: + yield resource_create_mock + + +@pytest.fixture +def otel_resource_detector_mock(): + import opentelemetry.sdk.resources + + with mock.patch.object( + opentelemetry.sdk.resources, "OTELResourceDetector" + ) as otel_resource_detector_mock: + yield otel_resource_detector_mock + + +@pytest.fixture +def is_noop_or_proxy_tracer_provider_mock(): + with mock.patch.object( + _utils, "is_noop_or_proxy_tracer_provider" + ) as is_noop_or_proxy_tracer_provider_mock: + is_noop_or_proxy_tracer_provider_mock.return_value = True + yield is_noop_or_proxy_tracer_provider_mock + + @pytest.fixture def langchain_instrumentor_mock(): with mock.patch.object( @@ -233,7 +264,6 @@ def test_stream_query(self, langchain_dump_mock): def test_enable_tracing( self, caplog, - cloud_trace_exporter_mock, tracer_provider_mock, simple_span_processor_mock, langchain_instrumentor_mock, @@ -256,6 +286,49 @@ def test_enable_tracing_warning(self, caplog, langchain_instrumentor_none_mock): # agent.set_up() # assert "enable_tracing=True but proceeding with tracing disabled" in caplog.text + def test_tracing_setup( + self, + tracer_provider_mock, + simple_span_processor_mock, + otlp_span_exporter_mock, + resource_create_mock, + otel_resource_detector_mock, + is_noop_or_proxy_tracer_provider_mock, + langchain_instrumentor_mock, + ): + agent = reasoning_engines.LanggraphAgent( + model=_TEST_MODEL, + runnable_builder=lambda **kwargs: kwargs, + enable_tracing=True, + ) + assert agent._instrumentor is None + + agent.set_up() + + otlp_span_exporter_mock.assert_called_once_with( + session=mock.ANY, + endpoint="https://telemetry.googleapis.com/v1/traces", + ) + resource_create_mock.assert_called_with( + attributes={"gcp.project_id": _TEST_PROJECT}, + ) + tracer_provider_mock.assert_called_with( + resource=resource_create_mock.return_value.merge.return_value, + ) + simple_span_processor_mock.assert_called_once_with( + span_exporter=otlp_span_exporter_mock.return_value, + ) + tracer_provider_mock.return_value.add_span_processor.assert_called_once_with( + simple_span_processor_mock.return_value, + ) + langchain_instrumentor_class = ( + langchain_instrumentor_mock.return_value.LangChainInstrumentor + ) + langchain_instrumentor_class.assert_called_once_with() + langchain_instrumentor_class.return_value.uninstrument.assert_called_once_with() + langchain_instrumentor_class.return_value.instrument.assert_called_once_with() + assert agent._instrumentor is langchain_instrumentor_class.return_value + def test_get_state_history_empty(self): agent = reasoning_engines.LanggraphAgent(model=_TEST_MODEL) agent._runnable = mock.Mock() diff --git a/tests/unit/vertex_llama_index/test_reasoning_engine_templates_llama_index.py b/tests/unit/vertex_llama_index/test_reasoning_engine_templates_llama_index.py index 414e91d387..ef3a900cdd 100644 --- a/tests/unit/vertex_llama_index/test_reasoning_engine_templates_llama_index.py +++ b/tests/unit/vertex_llama_index/test_reasoning_engine_templates_llama_index.py @@ -68,15 +68,6 @@ def model_builder_mock(): yield model_builder_mock -@pytest.fixture -def cloud_trace_exporter_mock(): - with mock.patch.object( - _utils, - "_import_cloud_trace_exporter_or_warn", - ) as cloud_trace_exporter_mock: - yield cloud_trace_exporter_mock - - @pytest.fixture def tracer_provider_mock(): with mock.patch("opentelemetry.sdk.trace.TracerProvider") as tracer_provider_mock: @@ -91,6 +82,46 @@ def simple_span_processor_mock(): yield simple_span_processor_mock +@pytest.fixture +def otlp_span_exporter_mock(): + import opentelemetry.exporter.otlp.proto.http.trace_exporter + + with mock.patch.object( + opentelemetry.exporter.otlp.proto.http.trace_exporter, + "OTLPSpanExporter", + ) as otlp_span_exporter_mock: + yield otlp_span_exporter_mock + + +@pytest.fixture +def resource_create_mock(): + import opentelemetry.sdk.resources + + with mock.patch.object( + opentelemetry.sdk.resources.Resource, "create" + ) as resource_create_mock: + yield resource_create_mock + + +@pytest.fixture +def otel_resource_detector_mock(): + import opentelemetry.sdk.resources + + with mock.patch.object( + opentelemetry.sdk.resources, "OTELResourceDetector" + ) as otel_resource_detector_mock: + yield otel_resource_detector_mock + + +@pytest.fixture +def is_noop_or_proxy_tracer_provider_mock(): + with mock.patch.object( + _utils, "is_noop_or_proxy_tracer_provider" + ) as is_noop_or_proxy_tracer_provider_mock: + is_noop_or_proxy_tracer_provider_mock.return_value = True + yield is_noop_or_proxy_tracer_provider_mock + + @pytest.fixture def llama_index_instrumentor_mock(): with mock.patch.object( @@ -228,7 +259,6 @@ def test_query_with_batch_input(self, json_loads_mock, nest_asyncio_apply_mock): def test_enable_tracing( self, caplog, - cloud_trace_exporter_mock, tracer_provider_mock, simple_span_processor_mock, llama_index_instrumentor_mock, @@ -259,6 +289,51 @@ def test_enable_tracing_warning(self, caplog, llama_index_instrumentor_none_mock # agent.set_up() # assert "enable_tracing=True but proceeding with tracing disabled" in caplog.text + def test_tracing_setup( + self, + tracer_provider_mock, + simple_span_processor_mock, + otlp_span_exporter_mock, + resource_create_mock, + otel_resource_detector_mock, + is_noop_or_proxy_tracer_provider_mock, + llama_index_instrumentor_mock, + ): + agent = llama_index.LlamaIndexQueryPipelineAgent( + model=_TEST_MODEL, + prompt=self.prompt, + model_builder=lambda **kwargs: kwargs, + runnable_builder=lambda **kwargs: kwargs, + enable_tracing=True, + ) + assert agent._instrumentor is None + + agent.set_up() + + otlp_span_exporter_mock.assert_called_once_with( + session=mock.ANY, + endpoint="https://telemetry.googleapis.com/v1/traces", + ) + resource_create_mock.assert_called_with( + attributes={"gcp.project_id": _TEST_PROJECT}, + ) + tracer_provider_mock.assert_called_with( + resource=resource_create_mock.return_value.merge.return_value, + ) + simple_span_processor_mock.assert_called_once_with( + span_exporter=otlp_span_exporter_mock.return_value, + ) + tracer_provider_mock.return_value.add_span_processor.assert_called_once_with( + simple_span_processor_mock.return_value, + ) + llama_index_instrumentor_class = ( + llama_index_instrumentor_mock.return_value.LlamaIndexInstrumentor + ) + llama_index_instrumentor_class.assert_called_once_with() + llama_index_instrumentor_class.return_value.uninstrument.assert_called_once_with() + llama_index_instrumentor_class.return_value.instrument.assert_called_once_with() + assert agent._instrumentor is llama_index_instrumentor_class.return_value + class TestToJsonSerializableLlamaIndexObject: """Tests for `_utils.to_json_serializable_llama_index_object`.""" diff --git a/vertexai/agent_engines/_utils.py b/vertexai/agent_engines/_utils.py index 751beef1d0..e2ce024465 100644 --- a/vertexai/agent_engines/_utils.py +++ b/vertexai/agent_engines/_utils.py @@ -774,34 +774,6 @@ def _import_opentelemetry_sdk_trace_or_warn() -> Optional[types.ModuleType]: return None -def _import_cloud_trace_v2_or_warn() -> Optional[types.ModuleType]: - """Tries to import the google.cloud.trace_v2 module.""" - try: - import google.cloud.trace_v2 - - return google.cloud.trace_v2 - except ImportError: - LOGGER.warning( - "google-cloud-trace is not installed. Please call " - "'pip install google-cloud-aiplatform[agent_engines]'." - ) - return None - - -def _import_cloud_trace_exporter_or_warn() -> Optional[types.ModuleType]: - """Tries to import the opentelemetry.exporter.cloud_trace module.""" - try: - import opentelemetry.exporter.cloud_trace # noqa:F401 - - return opentelemetry.exporter.cloud_trace - except ImportError: - LOGGER.warning( - "opentelemetry-exporter-gcp-trace is not installed. Please " - "call 'pip install google-cloud-aiplatform[agent_engines]'." - ) - return None - - def _import_openinference_langchain_or_warn() -> Optional[types.ModuleType]: """Tries to import the openinference.instrumentation.langchain module.""" try: diff --git a/vertexai/agent_engines/templates/ag2.py b/vertexai/agent_engines/templates/ag2.py index a7261a7d5f..f7087a0b79 100644 --- a/vertexai/agent_engines/templates/ag2.py +++ b/vertexai/agent_engines/templates/ag2.py @@ -91,28 +91,42 @@ def _default_runnable_builder( def _default_instrumentor_builder(project_id: str): from vertexai.agent_engines import _utils - cloud_trace_exporter = _utils._import_cloud_trace_exporter_or_warn() - cloud_trace_v2 = _utils._import_cloud_trace_v2_or_warn() openinference_autogen = _utils._import_openinference_autogen_or_warn() opentelemetry = _utils._import_opentelemetry_or_warn() opentelemetry_sdk_trace = _utils._import_opentelemetry_sdk_trace_or_warn() if all( ( - cloud_trace_exporter, - cloud_trace_v2, openinference_autogen, opentelemetry, opentelemetry_sdk_trace, ) ): - import google.auth + try: + import opentelemetry.exporter.otlp.proto.http.trace_exporter + import opentelemetry.sdk.resources + import google.auth + import google.auth.transport.requests + except (ImportError, AttributeError): + from google.cloud.aiplatform import base + + _LOGGER = base.Logger(__name__) + _LOGGER.warning( + "enable_tracing=True but proceeding with tracing disabled " + "because opentelemetry-exporter-otlp-proto-http is not installed." + ) + return None credentials, _ = google.auth.default() - span_exporter = cloud_trace_exporter.CloudTraceSpanExporter( - project_id=project_id, - client=cloud_trace_v2.TraceServiceClient( - credentials=credentials.with_quota_project(project_id), - ), + resource = opentelemetry.sdk.resources.Resource.create( + attributes={"gcp.project_id": project_id}, + ).merge(opentelemetry.sdk.resources.OTELResourceDetector().detect()) + 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", + ) ) span_processor: SpanProcessor = ( opentelemetry_sdk_trace.export.SimpleSpanProcessor( @@ -137,13 +151,13 @@ def _default_instrumentor_builder(project_id: str): "OTEL_PYTHON_TRACER_PROVIDER, _TRACER_PROVIDER, " "or _PROXY_TRACER_PROVIDER." ) - tracer_provider = opentelemetry_sdk_trace.TracerProvider() + tracer_provider = opentelemetry_sdk_trace.TracerProvider(resource=resource) opentelemetry.trace.set_tracer_provider(tracer_provider) # Avoids AttributeError: # 'ProxyTracerProvider' and 'NoOpTracerProvider' objects has no # attribute 'add_span_processor'. if _utils.is_noop_or_proxy_tracer_provider(tracer_provider): - tracer_provider = opentelemetry_sdk_trace.TracerProvider() + tracer_provider = opentelemetry_sdk_trace.TracerProvider(resource=resource) opentelemetry.trace.set_tracer_provider(tracer_provider) # Avoids OpenTelemetry client already exists error. _override_active_span_processor( diff --git a/vertexai/agent_engines/templates/langchain.py b/vertexai/agent_engines/templates/langchain.py index 4bcf40f55b..9a4dc8bc35 100644 --- a/vertexai/agent_engines/templates/langchain.py +++ b/vertexai/agent_engines/templates/langchain.py @@ -194,28 +194,42 @@ def _default_runnable_builder( def _default_instrumentor_builder(project_id: str): from vertexai.agent_engines import _utils - cloud_trace_exporter = _utils._import_cloud_trace_exporter_or_warn() - cloud_trace_v2 = _utils._import_cloud_trace_v2_or_warn() openinference_langchain = _utils._import_openinference_langchain_or_warn() opentelemetry = _utils._import_opentelemetry_or_warn() opentelemetry_sdk_trace = _utils._import_opentelemetry_sdk_trace_or_warn() if all( ( - cloud_trace_exporter, - cloud_trace_v2, openinference_langchain, opentelemetry, opentelemetry_sdk_trace, ) ): - import google.auth + try: + import opentelemetry.exporter.otlp.proto.http.trace_exporter + import opentelemetry.sdk.resources + import google.auth + import google.auth.transport.requests + except (ImportError, AttributeError): + from google.cloud.aiplatform import base + + _LOGGER = base.Logger(__name__) + _LOGGER.warning( + "enable_tracing=True but proceeding with tracing disabled " + "because opentelemetry-exporter-otlp-proto-http is not installed." + ) + return None credentials, _ = google.auth.default() - span_exporter = cloud_trace_exporter.CloudTraceSpanExporter( - project_id=project_id, - client=cloud_trace_v2.TraceServiceClient( - credentials=credentials.with_quota_project(project_id), - ), + resource = opentelemetry.sdk.resources.Resource.create( + attributes={"gcp.project_id": project_id}, + ).merge(opentelemetry.sdk.resources.OTELResourceDetector().detect()) + 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", + ) ) span_processor: SpanProcessor = ( opentelemetry_sdk_trace.export.SimpleSpanProcessor( @@ -240,13 +254,13 @@ def _default_instrumentor_builder(project_id: str): "OTEL_PYTHON_TRACER_PROVIDER, _TRACER_PROVIDER, " "or _PROXY_TRACER_PROVIDER." ) - tracer_provider = opentelemetry_sdk_trace.TracerProvider() + tracer_provider = opentelemetry_sdk_trace.TracerProvider(resource=resource) opentelemetry.trace.set_tracer_provider(tracer_provider) # Avoids AttributeError: # 'ProxyTracerProvider' and 'NoOpTracerProvider' objects has no # attribute 'add_span_processor'. if _utils.is_noop_or_proxy_tracer_provider(tracer_provider): - tracer_provider = opentelemetry_sdk_trace.TracerProvider() + tracer_provider = opentelemetry_sdk_trace.TracerProvider(resource=resource) opentelemetry.trace.set_tracer_provider(tracer_provider) # Avoids OpenTelemetry client already exists error. _override_active_span_processor( diff --git a/vertexai/agent_engines/templates/langgraph.py b/vertexai/agent_engines/templates/langgraph.py index 63c63c83aa..bd9c5d32c0 100644 --- a/vertexai/agent_engines/templates/langgraph.py +++ b/vertexai/agent_engines/templates/langgraph.py @@ -170,28 +170,42 @@ def _default_runnable_builder( def _default_instrumentor_builder(project_id: str): from vertexai.agent_engines import _utils - cloud_trace_exporter = _utils._import_cloud_trace_exporter_or_warn() - cloud_trace_v2 = _utils._import_cloud_trace_v2_or_warn() openinference_langchain = _utils._import_openinference_langchain_or_warn() opentelemetry = _utils._import_opentelemetry_or_warn() opentelemetry_sdk_trace = _utils._import_opentelemetry_sdk_trace_or_warn() if all( ( - cloud_trace_exporter, - cloud_trace_v2, openinference_langchain, opentelemetry, opentelemetry_sdk_trace, ) ): - import google.auth + try: + import opentelemetry.exporter.otlp.proto.http.trace_exporter + import opentelemetry.sdk.resources + import google.auth + import google.auth.transport.requests + except (ImportError, AttributeError): + from google.cloud.aiplatform import base + + _LOGGER = base.Logger(__name__) + _LOGGER.warning( + "enable_tracing=True but proceeding with tracing disabled " + "because opentelemetry-exporter-otlp-proto-http is not installed." + ) + return None credentials, _ = google.auth.default() - span_exporter = cloud_trace_exporter.CloudTraceSpanExporter( - project_id=project_id, - client=cloud_trace_v2.TraceServiceClient( - credentials=credentials.with_quota_project(project_id), - ), + resource = opentelemetry.sdk.resources.Resource.create( + attributes={"gcp.project_id": project_id}, + ).merge(opentelemetry.sdk.resources.OTELResourceDetector().detect()) + 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", + ) ) span_processor: SpanProcessor = ( opentelemetry_sdk_trace.export.SimpleSpanProcessor( @@ -215,13 +229,13 @@ def _default_instrumentor_builder(project_id: str): "OTEL_PYTHON_TRACER_PROVIDER, _TRACER_PROVIDER, " "or _PROXY_TRACER_PROVIDER." ) - tracer_provider = opentelemetry_sdk_trace.TracerProvider() + tracer_provider = opentelemetry_sdk_trace.TracerProvider(resource=resource) opentelemetry.trace.set_tracer_provider(tracer_provider) # Avoids AttributeError: # 'ProxyTracerProvider' and 'NoOpTracerProvider' objects has no # attribute 'add_span_processor'. if _utils.is_noop_or_proxy_tracer_provider(tracer_provider): - tracer_provider = opentelemetry_sdk_trace.TracerProvider() + tracer_provider = opentelemetry_sdk_trace.TracerProvider(resource=resource) opentelemetry.trace.set_tracer_provider(tracer_provider) # Avoids OpenTelemetry client already exists error. _override_active_span_processor( diff --git a/vertexai/preview/reasoning_engines/templates/ag2.py b/vertexai/preview/reasoning_engines/templates/ag2.py index 3e194449f9..b5c3cdfb89 100644 --- a/vertexai/preview/reasoning_engines/templates/ag2.py +++ b/vertexai/preview/reasoning_engines/templates/ag2.py @@ -299,28 +299,37 @@ def set_up(self): if self._enable_tracing: from vertexai.reasoning_engines import _utils - cloud_trace_exporter = _utils._import_cloud_trace_exporter_or_warn() - cloud_trace_v2 = _utils._import_cloud_trace_v2_or_warn() openinference_autogen = _utils._import_openinference_autogen_or_warn() opentelemetry = _utils._import_opentelemetry_or_warn() opentelemetry_sdk_trace = _utils._import_opentelemetry_sdk_trace_or_warn() + try: + import opentelemetry.exporter.otlp.proto.http.trace_exporter + import opentelemetry.sdk.resources + import google.auth + import google.auth.transport.requests + + _otlp_span_exporter_module = ( + opentelemetry.exporter.otlp.proto.http.trace_exporter + ) + except (ImportError, AttributeError): + _otlp_span_exporter_module = None if all( ( - cloud_trace_exporter, - cloud_trace_v2, + _otlp_span_exporter_module, openinference_autogen, opentelemetry, opentelemetry_sdk_trace, ) ): - import google.auth - credentials, _ = google.auth.default() - span_exporter = cloud_trace_exporter.CloudTraceSpanExporter( - project_id=self._project, - client=cloud_trace_v2.TraceServiceClient( - credentials=credentials.with_quota_project(self._project), + resource = opentelemetry.sdk.resources.Resource.create( + attributes={"gcp.project_id": self._project}, + ).merge(opentelemetry.sdk.resources.OTELResourceDetector().detect()) + span_exporter = _otlp_span_exporter_module.OTLPSpanExporter( + session=google.auth.transport.requests.AuthorizedSession( + credentials=credentials ), + endpoint="https://telemetry.googleapis.com/v1/traces", ) span_processor: SpanProcessor = ( opentelemetry_sdk_trace.export.SimpleSpanProcessor( @@ -347,13 +356,17 @@ def set_up(self): "OTEL_PYTHON_TRACER_PROVIDER, _TRACER_PROVIDER, " "or _PROXY_TRACER_PROVIDER." ) - tracer_provider = opentelemetry_sdk_trace.TracerProvider() + tracer_provider = opentelemetry_sdk_trace.TracerProvider( + resource=resource + ) opentelemetry.trace.set_tracer_provider(tracer_provider) # Avoids AttributeError: # 'ProxyTracerProvider' and 'NoOpTracerProvider' objects has no # attribute 'add_span_processor'. if _utils.is_noop_or_proxy_tracer_provider(tracer_provider): - tracer_provider = opentelemetry_sdk_trace.TracerProvider() + tracer_provider = opentelemetry_sdk_trace.TracerProvider( + resource=resource + ) opentelemetry.trace.set_tracer_provider(tracer_provider) # Avoids OpenTelemetry client already exists error. _override_active_span_processor( diff --git a/vertexai/preview/reasoning_engines/templates/langchain.py b/vertexai/preview/reasoning_engines/templates/langchain.py index cb39fb172e..f678c67c31 100644 --- a/vertexai/preview/reasoning_engines/templates/langchain.py +++ b/vertexai/preview/reasoning_engines/templates/langchain.py @@ -490,28 +490,37 @@ def set_up(self): if self._enable_tracing: from vertexai.reasoning_engines import _utils - cloud_trace_exporter = _utils._import_cloud_trace_exporter_or_warn() - cloud_trace_v2 = _utils._import_cloud_trace_v2_or_warn() openinference_langchain = _utils._import_openinference_langchain_or_warn() opentelemetry = _utils._import_opentelemetry_or_warn() opentelemetry_sdk_trace = _utils._import_opentelemetry_sdk_trace_or_warn() + try: + import opentelemetry.exporter.otlp.proto.http.trace_exporter + import opentelemetry.sdk.resources + import google.auth + import google.auth.transport.requests + + _otlp_span_exporter_module = ( + opentelemetry.exporter.otlp.proto.http.trace_exporter + ) + except (ImportError, AttributeError): + _otlp_span_exporter_module = None if all( ( - cloud_trace_exporter, - cloud_trace_v2, + _otlp_span_exporter_module, openinference_langchain, opentelemetry, opentelemetry_sdk_trace, ) ): - import google.auth - credentials, _ = google.auth.default() - span_exporter = cloud_trace_exporter.CloudTraceSpanExporter( - project_id=self._project, - client=cloud_trace_v2.TraceServiceClient( - credentials=credentials.with_quota_project(self._project), + resource = opentelemetry.sdk.resources.Resource.create( + attributes={"gcp.project_id": self._project}, + ).merge(opentelemetry.sdk.resources.OTELResourceDetector().detect()) + span_exporter = _otlp_span_exporter_module.OTLPSpanExporter( + session=google.auth.transport.requests.AuthorizedSession( + credentials=credentials ), + endpoint="https://telemetry.googleapis.com/v1/traces", ) span_processor: SpanProcessor = ( opentelemetry_sdk_trace.export.SimpleSpanProcessor( @@ -538,13 +547,17 @@ def set_up(self): "OTEL_PYTHON_TRACER_PROVIDER, _TRACER_PROVIDER, " "or _PROXY_TRACER_PROVIDER." ) - tracer_provider = opentelemetry_sdk_trace.TracerProvider() + tracer_provider = opentelemetry_sdk_trace.TracerProvider( + resource=resource + ) opentelemetry.trace.set_tracer_provider(tracer_provider) # Avoids AttributeError: # 'ProxyTracerProvider' and 'NoOpTracerProvider' objects has no # attribute 'add_span_processor'. if _utils.is_noop_or_proxy_tracer_provider(tracer_provider): - tracer_provider = opentelemetry_sdk_trace.TracerProvider() + tracer_provider = opentelemetry_sdk_trace.TracerProvider( + resource=resource + ) opentelemetry.trace.set_tracer_provider(tracer_provider) # Avoids OpenTelemetry client already exists error. _override_active_span_processor( diff --git a/vertexai/preview/reasoning_engines/templates/langgraph.py b/vertexai/preview/reasoning_engines/templates/langgraph.py index a3553fae2c..288890a7e4 100644 --- a/vertexai/preview/reasoning_engines/templates/langgraph.py +++ b/vertexai/preview/reasoning_engines/templates/langgraph.py @@ -418,28 +418,37 @@ def set_up(self): if self._enable_tracing: from vertexai.reasoning_engines import _utils - cloud_trace_exporter = _utils._import_cloud_trace_exporter_or_warn() - cloud_trace_v2 = _utils._import_cloud_trace_v2_or_warn() openinference_langchain = _utils._import_openinference_langchain_or_warn() opentelemetry = _utils._import_opentelemetry_or_warn() opentelemetry_sdk_trace = _utils._import_opentelemetry_sdk_trace_or_warn() + try: + import opentelemetry.exporter.otlp.proto.http.trace_exporter + import opentelemetry.sdk.resources + import google.auth + import google.auth.transport.requests + + _otlp_span_exporter_module = ( + opentelemetry.exporter.otlp.proto.http.trace_exporter + ) + except (ImportError, AttributeError): + _otlp_span_exporter_module = None if all( ( - cloud_trace_exporter, - cloud_trace_v2, + _otlp_span_exporter_module, openinference_langchain, opentelemetry, opentelemetry_sdk_trace, ) ): - import google.auth - credentials, _ = google.auth.default() - span_exporter = cloud_trace_exporter.CloudTraceSpanExporter( - project_id=self._project, - client=cloud_trace_v2.TraceServiceClient( - credentials=credentials.with_quota_project(self._project), + resource = opentelemetry.sdk.resources.Resource.create( + attributes={"gcp.project_id": self._project}, + ).merge(opentelemetry.sdk.resources.OTELResourceDetector().detect()) + span_exporter = _otlp_span_exporter_module.OTLPSpanExporter( + session=google.auth.transport.requests.AuthorizedSession( + credentials=credentials ), + endpoint="https://telemetry.googleapis.com/v1/traces", ) span_processor: SpanProcessor = ( opentelemetry_sdk_trace.export.SimpleSpanProcessor( @@ -465,13 +474,17 @@ def set_up(self): "OTEL_PYTHON_TRACER_PROVIDER, _TRACER_PROVIDER, " "or _PROXY_TRACER_PROVIDER." ) - tracer_provider = opentelemetry_sdk_trace.TracerProvider() + tracer_provider = opentelemetry_sdk_trace.TracerProvider( + resource=resource + ) opentelemetry.trace.set_tracer_provider(tracer_provider) # Avoids AttributeError: # 'ProxyTracerProvider' and 'NoOpTracerProvider' objects has no # attribute 'add_span_processor'. if _utils.is_noop_or_proxy_tracer_provider(tracer_provider): - tracer_provider = opentelemetry_sdk_trace.TracerProvider() + tracer_provider = opentelemetry_sdk_trace.TracerProvider( + resource=resource + ) opentelemetry.trace.set_tracer_provider(tracer_provider) # Avoids OpenTelemetry client already exists error. _override_active_span_processor( diff --git a/vertexai/preview/reasoning_engines/templates/llama_index.py b/vertexai/preview/reasoning_engines/templates/llama_index.py index 93acd90073..418eb665b8 100644 --- a/vertexai/preview/reasoning_engines/templates/llama_index.py +++ b/vertexai/preview/reasoning_engines/templates/llama_index.py @@ -386,30 +386,39 @@ def set_up(self): if self._enable_tracing: from vertexai.reasoning_engines import _utils - cloud_trace_exporter = _utils._import_cloud_trace_exporter_or_warn() - cloud_trace_v2 = _utils._import_cloud_trace_v2_or_warn() openinference_llama_index = ( _utils._import_openinference_llama_index_or_warn() ) opentelemetry = _utils._import_opentelemetry_or_warn() opentelemetry_sdk_trace = _utils._import_opentelemetry_sdk_trace_or_warn() + try: + import opentelemetry.exporter.otlp.proto.http.trace_exporter + import opentelemetry.sdk.resources + import google.auth + import google.auth.transport.requests + + _otlp_span_exporter_module = ( + opentelemetry.exporter.otlp.proto.http.trace_exporter + ) + except (ImportError, AttributeError): + _otlp_span_exporter_module = None if all( ( - cloud_trace_exporter, - cloud_trace_v2, + _otlp_span_exporter_module, openinference_llama_index, opentelemetry, opentelemetry_sdk_trace, ) ): - import google.auth - credentials, _ = google.auth.default() - span_exporter = cloud_trace_exporter.CloudTraceSpanExporter( - project_id=self._project, - client=cloud_trace_v2.TraceServiceClient( - credentials=credentials.with_quota_project(self._project), + resource = opentelemetry.sdk.resources.Resource.create( + attributes={"gcp.project_id": self._project}, + ).merge(opentelemetry.sdk.resources.OTELResourceDetector().detect()) + span_exporter = _otlp_span_exporter_module.OTLPSpanExporter( + session=google.auth.transport.requests.AuthorizedSession( + credentials=credentials ), + endpoint="https://telemetry.googleapis.com/v1/traces", ) span_processor: SpanProcessor = ( opentelemetry_sdk_trace.export.SimpleSpanProcessor( @@ -436,13 +445,17 @@ def set_up(self): "OTEL_PYTHON_TRACER_PROVIDER, _TRACER_PROVIDER, " "or _PROXY_TRACER_PROVIDER." ) - tracer_provider = opentelemetry_sdk_trace.TracerProvider() + tracer_provider = opentelemetry_sdk_trace.TracerProvider( + resource=resource + ) opentelemetry.trace.set_tracer_provider(tracer_provider) # Avoids AttributeError: # 'ProxyTracerProvider' and 'NoOpTracerProvider' objects has no # attribute 'add_span_processor'. if _utils.is_noop_or_proxy_tracer_provider(tracer_provider): - tracer_provider = opentelemetry_sdk_trace.TracerProvider() + tracer_provider = opentelemetry_sdk_trace.TracerProvider( + resource=resource + ) opentelemetry.trace.set_tracer_provider(tracer_provider) # Avoids OpenTelemetry client already exists error. _override_active_span_processor( diff --git a/vertexai/reasoning_engines/_utils.py b/vertexai/reasoning_engines/_utils.py index 82943b5457..9048317c73 100644 --- a/vertexai/reasoning_engines/_utils.py +++ b/vertexai/reasoning_engines/_utils.py @@ -427,34 +427,6 @@ def _import_opentelemetry_sdk_trace_or_warn() -> Optional[types.ModuleType]: return None -def _import_cloud_trace_v2_or_warn() -> Optional[types.ModuleType]: - """Tries to import the google.cloud.trace_v2 module.""" - try: - import google.cloud.trace_v2 - - return google.cloud.trace_v2 - except ImportError: - _LOGGER.warning( - "google-cloud-trace is not installed. Please call " - "'pip install google-cloud-aiplatform[agent_engines]'." - ) - return None - - -def _import_cloud_trace_exporter_or_warn() -> Optional[types.ModuleType]: - """Tries to import the opentelemetry.exporter.cloud_trace module.""" - try: - import opentelemetry.exporter.cloud_trace # noqa:F401 - - return opentelemetry.exporter.cloud_trace - except ImportError: - _LOGGER.warning( - "opentelemetry-exporter-gcp-trace is not installed. Please " - "call 'pip install google-cloud-aiplatform[langchain]'." - ) - return None - - def _import_openinference_langchain_or_warn() -> Optional[types.ModuleType]: """Tries to import the openinference.instrumentation.langchain module.""" try: