opentelemetry-instrumentation-vertexai 0.38.7__py3-none-any.whl → 2.0b0__py3-none-any.whl

opentelemetry/instrumentation/vertexai/__init__.py

@@ -1,368 +1,96 @@
-"""OpenTelemetry Vertex AI instrumentation"""
-
-import logging
-import os
-import types
-from typing import Collection
-from opentelemetry.instrumentation.vertexai.config import Config
-from opentelemetry.instrumentation.vertexai.utils import dont_throw
-from wrapt import wrap_function_wrapper
-
-from opentelemetry import context as context_api
-from opentelemetry.trace import get_tracer, SpanKind
-from opentelemetry.trace.status import Status, StatusCode
-
-from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
-from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY, unwrap
-
-from opentelemetry.semconv_ai import (
-    SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY,
-    SpanAttributes,
-    LLMRequestTypeValues,
-)
-from opentelemetry.instrumentation.vertexai.version import __version__
-
-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.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.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 should_send_prompts():
-    return (
-        os.getenv("TRACELOOP_TRACE_CONTENT") or "true"
-    ).lower() == "true" or context_api.get_value("override_enable_content_tracing")
-
-
-def is_streaming_response(response):
-    return isinstance(response, types.GeneratorType)
-
-
-def is_async_streaming_response(response):
-    return isinstance(response, types.AsyncGeneratorType)
-
-
-def _set_span_attribute(span, name, value):
-    if value is not None:
-        if value != "":
-            span.set_attribute(name, value)
-    return
-
-
-def _set_input_attributes(span, args, kwargs, llm_model):
-    if should_send_prompts() and args is not None and len(args) > 0:
-        prompt = ""
-        for arg in args:
-            if isinstance(arg, str):
-                prompt = f"{prompt}{arg}\n"
-            elif isinstance(arg, list):
-                for subarg in arg:
-                    prompt = f"{prompt}{subarg}\n"
-
-        _set_span_attribute(
-            span,
-            f"{SpanAttributes.LLM_PROMPTS}.0.user",
-            prompt,
-        )
-
-    _set_span_attribute(span, SpanAttributes.LLM_REQUEST_MODEL, llm_model)
-    _set_span_attribute(
-        span, f"{SpanAttributes.LLM_PROMPTS}.0.user", kwargs.get("prompt")
-    )
-    _set_span_attribute(
-        span, SpanAttributes.LLM_REQUEST_TEMPERATURE, kwargs.get("temperature")
-    )
-    _set_span_attribute(
-        span, SpanAttributes.LLM_REQUEST_MAX_TOKENS, kwargs.get("max_output_tokens")
-    )
-    _set_span_attribute(span, SpanAttributes.LLM_REQUEST_TOP_P, kwargs.get("top_p"))
-    _set_span_attribute(span, SpanAttributes.LLM_TOP_K, kwargs.get("top_k"))
-    _set_span_attribute(
-        span, SpanAttributes.LLM_PRESENCE_PENALTY, kwargs.get("presence_penalty")
-    )
-    _set_span_attribute(
-        span, SpanAttributes.LLM_FREQUENCY_PENALTY, kwargs.get("frequency_penalty")
-    )
-
-    return
-
-
-@dont_throw
-def _set_response_attributes(span, llm_model, generation_text, token_usage):
-    _set_span_attribute(span, SpanAttributes.LLM_RESPONSE_MODEL, llm_model)
-
-    if token_usage:
-        _set_span_attribute(
-            span,
-            SpanAttributes.LLM_USAGE_TOTAL_TOKENS,
-            token_usage.total_token_count,
-        )
-        _set_span_attribute(
-            span,
-            SpanAttributes.LLM_USAGE_COMPLETION_TOKENS,
-            token_usage.candidates_token_count,
-        )
-        _set_span_attribute(
-            span,
-            SpanAttributes.LLM_USAGE_PROMPT_TOKENS,
-            token_usage.prompt_token_count,
-        )
-
-    _set_span_attribute(span, f"{SpanAttributes.LLM_COMPLETIONS}.0.role", "assistant")
-    _set_span_attribute(
-        span,
-        f"{SpanAttributes.LLM_COMPLETIONS}.0.content",
-        generation_text,
+# 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."
     )
 
+API
+---
+"""
 
-def _build_from_streaming_response(span, 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
-
-    _set_response_attributes(span, llm_model, complete_response, token_usage)
-
-    span.set_status(Status(StatusCode.OK))
-    span.end()
-
-
-async def _abuild_from_streaming_response(span, 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
-
-    _set_response_attributes(span, llm_model, complete_response, token_usage)
-
-    span.set_status(Status(StatusCode.OK))
-    span.end()
-
-
-@dont_throw
-def _handle_request(span, args, kwargs, llm_model):
-    if span.is_recording():
-        _set_input_attributes(span, args, kwargs, llm_model)
-
-
-@dont_throw
-def _handle_response(span, response, llm_model):
-    if span.is_recording():
-        _set_response_attributes(
-            span, llm_model, response.candidates[0].text, response.usage_metadata
-        )
-
-        span.set_status(Status(StatusCode.OK))
-
-
-def _with_tracer_wrapper(func):
-    """Helper for providing tracer for wrapper functions."""
-
-    def _with_tracer(tracer, to_wrap):
-        def wrapper(wrapped, instance, args, kwargs):
-            return func(tracer, to_wrap, wrapped, instance, args, kwargs)
-
-        return wrapper
-
-    return _with_tracer
-
-
-@_with_tracer_wrapper
-async def _awrap(tracer, 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/", "")
-
-    name = to_wrap.get("span_name")
-    span = tracer.start_span(
-        name,
-        kind=SpanKind.CLIENT,
-        attributes={
-            SpanAttributes.LLM_SYSTEM: "VertexAI",
-            SpanAttributes.LLM_REQUEST_TYPE: LLMRequestTypeValues.COMPLETION.value,
-        },
-    )
-
-    _handle_request(span, args, kwargs, llm_model)
-
-    response = await wrapped(*args, **kwargs)
-
-    if response:
-        if is_streaming_response(response):
-            return _build_from_streaming_response(span, response, llm_model)
-        elif is_async_streaming_response(response):
-            return _abuild_from_streaming_response(span, response, llm_model)
-        else:
-            _handle_response(span, response, llm_model)
-
-    span.end()
-    return response
-
-
-@_with_tracer_wrapper
-def _wrap(tracer, 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/", "")
-
-    name = to_wrap.get("span_name")
-    span = tracer.start_span(
-        name,
-        kind=SpanKind.CLIENT,
-        attributes={
-            SpanAttributes.LLM_SYSTEM: "VertexAI",
-            SpanAttributes.LLM_REQUEST_TYPE: LLMRequestTypeValues.COMPLETION.value,
-        },
-    )
-
-    _handle_request(span, args, kwargs, llm_model)
-
-    response = wrapped(*args, **kwargs)
-
-    if response:
-        if is_streaming_response(response):
-            return _build_from_streaming_response(span, response, llm_model)
-        elif is_async_streaming_response(response):
-            return _abuild_from_streaming_response(span, response, llm_model)
-        else:
-            _handle_response(span, response, llm_model)
+from wrapt import (
+    wrap_function_wrapper,  # type: ignore[reportUnknownVariableType]
+)
 
-    span.end()
-    return response
+from opentelemetry._events import get_event_logger
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.instrumentation.vertexai.package import _instruments
+from opentelemetry.instrumentation.vertexai.patch import (
+    generate_content_create,
+)
+from opentelemetry.instrumentation.vertexai.utils import is_content_enabled
+from opentelemetry.semconv.schemas import Schemas
+from opentelemetry.trace import get_tracer
 
 
 class VertexAIInstrumentor(BaseInstrumentor):
-    """An instrumentor for VertextAI's client library."""
-
-    def __init__(self, exception_logger=None):
-        super().__init__()
-        Config.exception_logger = exception_logger
-
     def instrumentation_dependencies(self) -> Collection[str]:
         return _instruments
 
-    def _instrument(self, **kwargs):
+    def _instrument(self, **kwargs: Any):
+        """Enable VertexAI instrumentation."""
         tracer_provider = kwargs.get("tracer_provider")
-        tracer = get_tracer(__name__, __version__, tracer_provider)
-        for wrapped_method in WRAPPED_METHODS:
-            wrap_package = wrapped_method.get("package")
-            wrap_object = wrapped_method.get("object")
-            wrap_method = wrapped_method.get("method")
+        tracer = get_tracer(
+            __name__,
+            "",
+            tracer_provider,
+            schema_url=Schemas.V1_28_0.value,
+        )
+        event_logger_provider = kwargs.get("event_logger_provider")
+        event_logger = get_event_logger(
+            __name__,
+            "",
+            schema_url=Schemas.V1_28_0.value,
+            event_logger_provider=event_logger_provider,
+        )
 
-            wrap_function_wrapper(
-                wrap_package,
-                f"{wrap_object}.{wrap_method}",
-                (
-                    _awrap(tracer, wrapped_method)
-                    if wrapped_method.get("is_async")
-                    else _wrap(tracer, wrapped_method)
-                ),
-            )
+        wrap_function_wrapper(
+            module="google.cloud.aiplatform_v1beta1.services.prediction_service.client",
+            name="PredictionServiceClient.generate_content",
+            wrapper=generate_content_create(
+                tracer, event_logger, is_content_enabled()
+            ),
+        )
+        wrap_function_wrapper(
+            module="google.cloud.aiplatform_v1.services.prediction_service.client",
+            name="PredictionServiceClient.generate_content",
+            wrapper=generate_content_create(
+                tracer, event_logger, is_content_enabled()
+            ),
+        )
 
-    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:
+        """TODO: implemented in later PR"""

opentelemetry/instrumentation/vertexai/events.py

@@ -0,0 +1,188 @@
+# 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.
+
+"""
+Factories for event types described in
+https://github.com/open-telemetry/semantic-conventions/blob/main/docs/gen-ai/gen-ai-events.md#system-event.
+
+Hopefully this code can be autogenerated by Weaver once Gen AI semantic conventions are
+schematized in YAML and the Weaver tool supports it.
+"""
+
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass
+from typing import Any, Iterable, Literal
+
+from opentelemetry._events import Event
+from opentelemetry.semconv._incubating.attributes import gen_ai_attributes
+from opentelemetry.util.types import AnyValue
+
+
+def user_event(
+    *,
+    role: str = "user",
+    content: AnyValue = None,
+) -> Event:
+    """Creates a User event
+    https://github.com/open-telemetry/semantic-conventions/blob/v1.28.0/docs/gen-ai/gen-ai-events.md#user-event
+    """
+    body: dict[str, AnyValue] = {
+        "role": role,
+    }
+    if content is not None:
+        body["content"] = content
+    return Event(
+        name="gen_ai.user.message",
+        attributes={
+            gen_ai_attributes.GEN_AI_SYSTEM: gen_ai_attributes.GenAiSystemValues.VERTEX_AI.value,
+        },
+        body=body,
+    )
+
+
+def assistant_event(
+    *,
+    role: str = "assistant",
+    content: AnyValue = None,
+) -> Event:
+    """Creates an Assistant event
+    https://github.com/open-telemetry/semantic-conventions/blob/v1.28.0/docs/gen-ai/gen-ai-events.md#assistant-event
+    """
+    body: dict[str, AnyValue] = {
+        "role": role,
+    }
+    if content is not None:
+        body["content"] = content
+    return Event(
+        name="gen_ai.assistant.message",
+        attributes={
+            gen_ai_attributes.GEN_AI_SYSTEM: gen_ai_attributes.GenAiSystemValues.VERTEX_AI.value,
+        },
+        body=body,
+    )
+
+
+def system_event(
+    *,
+    role: str = "system",
+    content: AnyValue = None,
+) -> Event:
+    """Creates a System event
+    https://github.com/open-telemetry/semantic-conventions/blob/v1.28.0/docs/gen-ai/gen-ai-events.md#system-event
+    """
+    body: dict[str, AnyValue] = {
+        "role": role,
+    }
+    if content is not None:
+        body["content"] = content
+    return Event(
+        name="gen_ai.system.message",
+        attributes={
+            gen_ai_attributes.GEN_AI_SYSTEM: gen_ai_attributes.GenAiSystemValues.VERTEX_AI.value,
+        },
+        body=body,
+    )
+
+
+def tool_event(
+    *,
+    role: str | None,
+    id_: str,
+    content: AnyValue = None,
+) -> Event:
+    """Creates a Tool message event
+    https://github.com/open-telemetry/semantic-conventions/blob/v1.28.0/docs/gen-ai/gen-ai-events.md#event-gen_aitoolmessage
+    """
+    if not role:
+        role = "tool"
+
+    body: dict[str, AnyValue] = {
+        "role": role,
+        "id": id_,
+    }
+    if content is not None:
+        body["content"] = content
+    return Event(
+        name="gen_ai.tool.message",
+        attributes={
+            gen_ai_attributes.GEN_AI_SYSTEM: gen_ai_attributes.GenAiSystemValues.VERTEX_AI.value,
+        },
+        body=body,
+    )
+
+
+@dataclass
+class ChoiceMessage:
+    """The message field for a gen_ai.choice event"""
+
+    content: AnyValue = None
+    role: str = "assistant"
+
+
+@dataclass
+class ChoiceToolCall:
+    """The tool_calls field for a gen_ai.choice event"""
+
+    @dataclass
+    class Function:
+        name: str
+        arguments: AnyValue = None
+
+    function: Function
+    id: str
+    type: Literal["function"] = "function"
+
+
+FinishReason = Literal[
+    "content_filter", "error", "length", "stop", "tool_calls"
+]
+
+
+def choice_event(
+    *,
+    finish_reason: FinishReason | str,
+    index: int,
+    message: ChoiceMessage,
+    tool_calls: Iterable[ChoiceToolCall] = (),
+) -> Event:
+    """Creates a choice event, which describes the Gen AI response message.
+    https://github.com/open-telemetry/semantic-conventions/blob/v1.28.0/docs/gen-ai/gen-ai-events.md#event-gen_aichoice
+    """
+    body: dict[str, AnyValue] = {
+        "finish_reason": finish_reason,
+        "index": index,
+        "message": _asdict_filter_nulls(message),
+    }
+
+    tool_calls_list = [
+        _asdict_filter_nulls(tool_call) for tool_call in tool_calls
+    ]
+    if tool_calls_list:
+        body["tool_calls"] = tool_calls_list
+
+    return Event(
+        name="gen_ai.choice",
+        attributes={
+            gen_ai_attributes.GEN_AI_SYSTEM: gen_ai_attributes.GenAiSystemValues.VERTEX_AI.value,
+        },
+        body=body,
+    )
+
+
+def _asdict_filter_nulls(instance: Any) -> dict[str, AnyValue]:
+    return asdict(
+        instance,
+        dict_factory=lambda kvs: {k: v for (k, v) in kvs if v is not None},
+    )

opentelemetry/instrumentation/vertexai/package.py

@@ -0,0 +1,16 @@
+# 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.
+
+
+_instruments = ("google-cloud-aiplatform >= 1.64",)