opentelemetry-instrumentation-vertexai 0.47.3__py3-none-any.whl → 2.1b0__py3-none-any.whl

opentelemetry/instrumentation/vertexai/__init__.py

@@ -1,366 +1,173 @@
-"""OpenTelemetry Vertex AI instrumentation"""
-
-import logging
-import types
-from typing import Collection
-
-from opentelemetry import context as context_api
-from opentelemetry._events import get_event_logger
-from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
-from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY, unwrap
-from opentelemetry.instrumentation.vertexai.config import Config
-from opentelemetry.instrumentation.vertexai.event_emitter import (
-    emit_prompt_events,
-    emit_response_events,
-)
-from opentelemetry.instrumentation.vertexai.span_utils import (
-    set_input_attributes,
-    set_input_attributes_sync,
-    set_model_input_attributes,
-    set_model_response_attributes,
-    set_response_attributes,
-)
-from opentelemetry.instrumentation.vertexai.utils import dont_throw, should_emit_events
-from opentelemetry.instrumentation.vertexai.version import __version__
-from opentelemetry.semconv_ai import (
-    SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY,
-    LLMRequestTypeValues,
-    SpanAttributes,
-)
-from opentelemetry.trace import SpanKind, get_tracer
-from opentelemetry.trace.status import Status, StatusCode
-from wrapt import wrap_function_wrapper
-
-logger = logging.getLogger(__name__)
-
-_instruments = ("google-cloud-aiplatform >= 1.38.1",)
-
-WRAPPED_METHODS = [
-    {
-        "package": "vertexai.generative_models",
-        "object": "GenerativeModel",
-        "method": "generate_content",
-        "span_name": "vertexai.generate_content",
-        "is_async": False,
-    },
-    {
-        "package": "vertexai.generative_models",
-        "object": "GenerativeModel",
-        "method": "generate_content_async",
-        "span_name": "vertexai.generate_content_async",
-        "is_async": True,
-    },
-    {
-        "package": "vertexai.generative_models",
-        "object": "ChatSession",
-        "method": "send_message",
-        "span_name": "vertexai.send_message",
-        "is_async": False,
-    },
-    {
-        "package": "vertexai.preview.generative_models",
-        "object": "GenerativeModel",
-        "method": "generate_content",
-        "span_name": "vertexai.generate_content",
-        "is_async": False,
-    },
-    {
-        "package": "vertexai.preview.generative_models",
-        "object": "GenerativeModel",
-        "method": "generate_content_async",
-        "span_name": "vertexai.generate_content_async",
-        "is_async": True,
-    },
-    {
-        "package": "vertexai.preview.generative_models",
-        "object": "ChatSession",
-        "method": "send_message",
-        "span_name": "vertexai.send_message",
-        "is_async": False,
-    },
-    {
-        "package": "vertexai.language_models",
-        "object": "TextGenerationModel",
-        "method": "predict",
-        "span_name": "vertexai.predict",
-        "is_async": False,
-    },
-    {
-        "package": "vertexai.language_models",
-        "object": "TextGenerationModel",
-        "method": "predict_async",
-        "span_name": "vertexai.predict_async",
-        "is_async": True,
-    },
-    {
-        "package": "vertexai.language_models",
-        "object": "TextGenerationModel",
-        "method": "predict_streaming",
-        "span_name": "vertexai.predict_streaming",
-        "is_async": False,
-    },
-    {
-        "package": "vertexai.language_models",
-        "object": "TextGenerationModel",
-        "method": "predict_streaming_async",
-        "span_name": "vertexai.predict_streaming_async",
-        "is_async": True,
-    },
-    {
-        "package": "vertexai.language_models",
-        "object": "ChatSession",
-        "method": "send_message",
-        "span_name": "vertexai.send_message",
-        "is_async": False,
-    },
-    {
-        "package": "vertexai.language_models",
-        "object": "ChatSession",
-        "method": "send_message_streaming",
-        "span_name": "vertexai.send_message_streaming",
-        "is_async": False,
-    },
-]
-
-
-def is_streaming_response(response):
-    return isinstance(response, types.GeneratorType)
-
-
-def is_async_streaming_response(response):
-    return isinstance(response, types.AsyncGeneratorType)
-
-
-@dont_throw
-def handle_streaming_response(span, event_logger, llm_model, response, token_usage):
-    set_model_response_attributes(span, llm_model, token_usage)
-    if should_emit_events():
-        emit_response_events(response, event_logger)
-    else:
-        set_response_attributes(span, llm_model, response)
-    if span.is_recording():
-        span.set_status(Status(StatusCode.OK))
-
-
-def _build_from_streaming_response(span, event_logger, response, llm_model):
-    complete_response = ""
-    token_usage = None
-    for item in response:
-        item_to_yield = item
-        complete_response += str(item.text)
-        if item.usage_metadata:
-            token_usage = item.usage_metadata
-
-        yield item_to_yield
-
-    handle_streaming_response(
-        span, event_logger, llm_model, complete_response, token_usage
+# Copyright The OpenTelemetry Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+VertexAI client instrumentation supporting `google-cloud-aiplatform` SDK, it can be enabled by
+using ``VertexAIInstrumentor``.
+
+.. _vertexai: https://pypi.org/project/google-cloud-aiplatform/
+
+Usage
+-----
+
+.. code:: python
+
+    import vertexai
+    from vertexai.generative_models import GenerativeModel
+    from opentelemetry.instrumentation.vertexai import VertexAIInstrumentor
+
+    VertexAIInstrumentor().instrument()
+
+    vertexai.init()
+    model = GenerativeModel("gemini-1.5-flash-002")
+    chat_completion = model.generate_content(
+        "Write a short poem on OpenTelemetry."
     )
 
-    span.set_status(Status(StatusCode.OK))
-    span.end()
+API
+---
+"""
 
+from __future__ import annotations
 
-async def _abuild_from_streaming_response(span, event_logger, response, llm_model):
-    complete_response = ""
-    token_usage = None
-    async for item in response:
-        item_to_yield = item
-        complete_response += str(item.text)
-        if item.usage_metadata:
-            token_usage = item.usage_metadata
+from typing import Any, Collection
 
-        yield item_to_yield
-
-    handle_streaming_response(span, event_logger, llm_model, response, token_usage)
-
-    span.set_status(Status(StatusCode.OK))
-    span.end()
-
-
-@dont_throw
-async def _handle_request(span, event_logger, args, kwargs, llm_model):
-    set_model_input_attributes(span, kwargs, llm_model)
-    if should_emit_events():
-        emit_prompt_events(args, event_logger)
-    else:
-        await set_input_attributes(span, args)
-
-
-def _handle_response(span, event_logger, response, llm_model):
-    set_model_response_attributes(span, llm_model, response.usage_metadata)
-    if should_emit_events():
-        emit_response_events(response, event_logger)
-    else:
-        set_response_attributes(
-            span, llm_model, response.candidates[0].text if response.candidates else ""
-        )
-    if span.is_recording():
-        span.set_status(Status(StatusCode.OK))
-
-
-def _with_tracer_wrapper(func):
-    """Helper for providing tracer for wrapper functions."""
-
-    def _with_tracer(tracer, event_logger, to_wrap):
-        def wrapper(wrapped, instance, args, kwargs):
-            return func(tracer, event_logger, to_wrap, wrapped, instance, args, kwargs)
-
-        return wrapper
-
-    return _with_tracer
-
-
-@_with_tracer_wrapper
-async def _awrap(tracer, event_logger, to_wrap, wrapped, instance, args, kwargs):
-    """Instruments and calls every function defined in TO_WRAP."""
-    if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY) or context_api.get_value(
-        SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY
-    ):
-        return await wrapped(*args, **kwargs)
-
-    llm_model = "unknown"
-    if hasattr(instance, "_model_id"):
-        llm_model = instance._model_id
-    if hasattr(instance, "_model_name"):
-        llm_model = instance._model_name.replace("publishers/google/models/", "")
-    # For ChatSession, try to get model from the parent model object
-    if hasattr(instance, "_model") and hasattr(instance._model, "_model_name"):
-        llm_model = instance._model._model_name.replace("publishers/google/models/", "")
-    elif hasattr(instance, "_model") and hasattr(instance._model, "_model_id"):
-        llm_model = instance._model._model_id
+from wrapt import (
+    wrap_function_wrapper,  # type: ignore[reportUnknownVariableType]
+)
 
-    name = to_wrap.get("span_name")
-    span = tracer.start_span(
-        name,
-        kind=SpanKind.CLIENT,
-        attributes={
-            SpanAttributes.LLM_SYSTEM: "Google",
-            SpanAttributes.LLM_REQUEST_TYPE: LLMRequestTypeValues.COMPLETION.value,
-        },
+from opentelemetry._logs import get_logger
+from opentelemetry.instrumentation._semconv import (
+    _OpenTelemetrySemanticConventionStability,
+    _OpenTelemetryStabilitySignalType,
+    _StabilityMode,
+)
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.instrumentation.utils import unwrap
+from opentelemetry.instrumentation.vertexai.package import _instruments
+from opentelemetry.instrumentation.vertexai.patch import MethodWrappers
+from opentelemetry.instrumentation.vertexai.utils import is_content_enabled
+from opentelemetry.semconv.schemas import Schemas
+from opentelemetry.trace import get_tracer
+from opentelemetry.util.genai.completion_hook import load_completion_hook
+
+
+def _methods_to_wrap(
+    method_wrappers: MethodWrappers,
+):
+    # This import is very slow, do it lazily in case instrument() is not called
+    # pylint: disable=import-outside-toplevel
+    from google.cloud.aiplatform_v1.services.prediction_service import (
+        async_client,
+        client,
     )
-
-    await _handle_request(span, event_logger, args, kwargs, llm_model)
-
-    response = await wrapped(*args, **kwargs)
-
-    if response:
-        if is_streaming_response(response):
-            return _build_from_streaming_response(
-                span, event_logger, response, llm_model
-            )
-        elif is_async_streaming_response(response):
-            return _abuild_from_streaming_response(
-                span, event_logger, response, llm_model
-            )
-        else:
-            _handle_response(span, event_logger, response, llm_model)
-
-    span.end()
-    return response
-
-
-@_with_tracer_wrapper
-def _wrap(tracer, event_logger, to_wrap, wrapped, instance, args, kwargs):
-    """Instruments and calls every function defined in TO_WRAP."""
-    if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY) or context_api.get_value(
-        SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY
-    ):
-        return wrapped(*args, **kwargs)
-
-    llm_model = "unknown"
-    if hasattr(instance, "_model_id"):
-        llm_model = instance._model_id
-    if hasattr(instance, "_model_name"):
-        llm_model = instance._model_name.replace("publishers/google/models/", "")
-    # For ChatSession, try to get model from the parent model object
-    if hasattr(instance, "_model") and hasattr(instance._model, "_model_name"):
-        llm_model = instance._model._model_name.replace("publishers/google/models/", "")
-    elif hasattr(instance, "_model") and hasattr(instance._model, "_model_id"):
-        llm_model = instance._model._model_id
-
-    name = to_wrap.get("span_name")
-    span = tracer.start_span(
-        name,
-        kind=SpanKind.CLIENT,
-        attributes={
-            SpanAttributes.LLM_SYSTEM: "Google",
-            SpanAttributes.LLM_REQUEST_TYPE: LLMRequestTypeValues.COMPLETION.value,
-        },
+    from google.cloud.aiplatform_v1beta1.services.prediction_service import (
+        async_client as async_client_v1beta1,
+    )
+    from google.cloud.aiplatform_v1beta1.services.prediction_service import (
+        client as client_v1beta1,
     )
 
-    # Use sync version for non-async wrapper to avoid image processing for now
-    set_model_input_attributes(span, kwargs, llm_model)
-    if should_emit_events():
-        emit_prompt_events(args, event_logger)
-    else:
-        set_input_attributes_sync(span, args)
-
-    response = wrapped(*args, **kwargs)
-
-    if response:
-        if is_streaming_response(response):
-            return _build_from_streaming_response(
-                span, event_logger, response, llm_model
-            )
-        elif is_async_streaming_response(response):
-            return _abuild_from_streaming_response(
-                span, event_logger, response, llm_model
-            )
-        else:
-            _handle_response(span, event_logger, response, llm_model)
+    for client_class in (
+        client.PredictionServiceClient,
+        client_v1beta1.PredictionServiceClient,
+    ):
+        yield (
+            client_class,
+            client_class.generate_content.__name__,  # type: ignore[reportUnknownMemberType]
+            method_wrappers.generate_content,
+        )
 
-    span.end()
-    return response
+    for client_class in (
+        async_client.PredictionServiceAsyncClient,
+        async_client_v1beta1.PredictionServiceAsyncClient,
+    ):
+        yield (
+            client_class,
+            client_class.generate_content.__name__,  # type: ignore[reportUnknownMemberType]
+            method_wrappers.agenerate_content,
+        )
 
 
 class VertexAIInstrumentor(BaseInstrumentor):
-    """An instrumentor for VertextAI's client library."""
-
-    def __init__(self, exception_logger=None, use_legacy_attributes=True, upload_base64_image=None):
+    def __init__(self) -> None:
         super().__init__()
-        Config.exception_logger = exception_logger
-        Config.use_legacy_attributes = use_legacy_attributes
-        if upload_base64_image:
-            Config.upload_base64_image = upload_base64_image
+        self._methods_to_unwrap: list[tuple[Any, str]] = []
 
     def instrumentation_dependencies(self) -> Collection[str]:
         return _instruments
 
-    def _instrument(self, **kwargs):
+    def _instrument(self, **kwargs: Any):
+        """Enable VertexAI instrumentation."""
+        completion_hook = (
+            kwargs.get("completion_hook") or load_completion_hook()
+        )
+        sem_conv_opt_in_mode = _OpenTelemetrySemanticConventionStability._get_opentelemetry_stability_opt_in_mode(
+            _OpenTelemetryStabilitySignalType.GEN_AI,
+        )
         tracer_provider = kwargs.get("tracer_provider")
-        tracer = get_tracer(__name__, __version__, tracer_provider)
-
-        event_logger = None
-
-        if should_emit_events():
-            event_logger_provider = kwargs.get("event_logger_provider")
-            event_logger = get_event_logger(
-                __name__,
-                __version__,
-                event_logger_provider=event_logger_provider,
+        schema = (
+            Schemas.V1_28_0.value
+            if sem_conv_opt_in_mode == _StabilityMode.DEFAULT
+            else Schemas.V1_36_0.value
+        )
+        tracer = get_tracer(
+            __name__,
+            "",
+            tracer_provider,
+            schema_url=schema,
+        )
+        logger_provider = kwargs.get("logger_provider")
+        logger = get_logger(
+            __name__,
+            "",
+            logger_provider=logger_provider,
+            schema_url=schema,
+        )
+        sem_conv_opt_in_mode = _OpenTelemetrySemanticConventionStability._get_opentelemetry_stability_opt_in_mode(
+            _OpenTelemetryStabilitySignalType.GEN_AI,
+        )
+        if sem_conv_opt_in_mode == _StabilityMode.DEFAULT:
+            # Type checker now knows sem_conv_opt_in_mode is a Literal[_StabilityMode.DEFAULT]
+            method_wrappers = MethodWrappers(
+                tracer,
+                logger,
+                is_content_enabled(sem_conv_opt_in_mode),
+                sem_conv_opt_in_mode,
+                completion_hook,
             )
-
-        for wrapped_method in WRAPPED_METHODS:
-            wrap_package = wrapped_method.get("package")
-            wrap_object = wrapped_method.get("object")
-            wrap_method = wrapped_method.get("method")
-
+        elif sem_conv_opt_in_mode == _StabilityMode.GEN_AI_LATEST_EXPERIMENTAL:
+            # Type checker now knows it's the other literal
+            method_wrappers = MethodWrappers(
+                tracer,
+                logger,
+                is_content_enabled(sem_conv_opt_in_mode),
+                sem_conv_opt_in_mode,
+                completion_hook,
+            )
+        else:
+            raise RuntimeError(f"{sem_conv_opt_in_mode} mode not supported")
+        for client_class, method_name, wrapper in _methods_to_wrap(
+            method_wrappers
+        ):
             wrap_function_wrapper(
-                wrap_package,
-                f"{wrap_object}.{wrap_method}",
-                (
-                    _awrap(tracer, event_logger, wrapped_method)
-                    if wrapped_method.get("is_async")
-                    else _wrap(tracer, event_logger, wrapped_method)
-                ),
+                client_class,
+                name=method_name,
+                wrapper=wrapper,
             )
+            self._methods_to_unwrap.append((client_class, method_name))
 
-    def _uninstrument(self, **kwargs):
-        for wrapped_method in WRAPPED_METHODS:
-            wrap_package = wrapped_method.get("package")
-            wrap_object = wrapped_method.get("object")
-            unwrap(
-                f"{wrap_package}.{wrap_object}",
-                wrapped_method.get("method", ""),
-            )
+    def _uninstrument(self, **kwargs: Any) -> None:
+        for client_class, method_name in self._methods_to_unwrap:
+            unwrap(client_class, method_name)