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

opentelemetry/instrumentation/vertexai/events.py

@@ -0,0 +1,190 @@
+# 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.
+
+# type: ignore[reportUnknownDeprecated]
+
+"""
+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._logs import LogRecord
+from opentelemetry.semconv._incubating.attributes import gen_ai_attributes
+from opentelemetry.util.types import AnyValue
+
+
+def user_event(
+    *,
+    role: str = "user",
+    content: AnyValue = None,
+) -> LogRecord:
+    """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 LogRecord(
+        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,
+) -> LogRecord:
+    """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 LogRecord(
+        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,
+) -> LogRecord:
+    """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 LogRecord(
+        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,
+) -> LogRecord:
+    """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 LogRecord(
+        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] = (),
+) -> LogRecord:
+    """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 LogRecord(
+        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",)

opentelemetry/instrumentation/vertexai/patch.py

@@ -0,0 +1,371 @@
+# 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.
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+from dataclasses import asdict
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Awaitable,
+    Callable,
+    Literal,
+    MutableSequence,
+    Union,
+    cast,
+    overload,
+)
+
+from opentelemetry._logs import Logger, LogRecord
+from opentelemetry.instrumentation._semconv import (
+    _StabilityMode,
+)
+from opentelemetry.instrumentation.vertexai.utils import (
+    GenerateContentParams,
+    _map_finish_reason,
+    convert_content_to_message_parts,
+    get_genai_request_attributes,
+    get_genai_response_attributes,
+    get_server_attributes,
+    get_span_name,
+    request_to_events,
+    response_to_events,
+)
+from opentelemetry.semconv._incubating.attributes import (
+    gen_ai_attributes as GenAI,
+)
+from opentelemetry.trace import SpanKind, Tracer
+from opentelemetry.util.genai.completion_hook import CompletionHook
+from opentelemetry.util.genai.types import (
+    ContentCapturingMode,
+    InputMessage,
+    OutputMessage,
+)
+from opentelemetry.util.genai.utils import gen_ai_json_dumps
+
+if TYPE_CHECKING:
+    from google.cloud.aiplatform_v1.services.prediction_service import client
+    from google.cloud.aiplatform_v1.types import (
+        content,
+        prediction_service,
+    )
+    from google.cloud.aiplatform_v1beta1.services.prediction_service import (
+        client as client_v1beta1,
+    )
+    from google.cloud.aiplatform_v1beta1.types import (
+        content as content_v1beta1,
+    )
+    from google.cloud.aiplatform_v1beta1.types import (
+        prediction_service as prediction_service_v1beta1,
+    )
+
+
+# Use parameter signature from
+# https://github.com/googleapis/python-aiplatform/blob/v1.76.0/google/cloud/aiplatform_v1/services/prediction_service/client.py#L2088
+# to handle named vs positional args robustly
+def _extract_params(
+    request: prediction_service.GenerateContentRequest
+    | prediction_service_v1beta1.GenerateContentRequest
+    | dict[Any, Any]
+    | None = None,
+    *,
+    model: str | None = None,
+    contents: MutableSequence[content.Content]
+    | MutableSequence[content_v1beta1.Content]
+    | None = None,
+    **_kwargs: Any,
+) -> GenerateContentParams:
+    # Request vs the named parameters are mututally exclusive or the RPC will fail
+    if not request:
+        return GenerateContentParams(
+            model=model or "",
+            contents=contents,
+        )
+
+    if isinstance(request, dict):
+        return GenerateContentParams(**request)
+
+    return GenerateContentParams(
+        model=request.model,
+        contents=request.contents,
+        system_instruction=request.system_instruction,
+        tools=request.tools,
+        tool_config=request.tool_config,
+        labels=request.labels,
+        safety_settings=request.safety_settings,
+        generation_config=request.generation_config,
+    )
+
+
+# For details about GEN_AI_LATEST_EXPERIMENTAL stability mode see
+# https://github.com/open-telemetry/semantic-conventions/blob/v1.37.0/docs/gen-ai/gen-ai-agent-spans.md?plain=1#L18-L37
+class MethodWrappers:
+    @overload
+    def __init__(
+        self,
+        tracer: Tracer,
+        logger: Logger,
+        capture_content: ContentCapturingMode,
+        sem_conv_opt_in_mode: Literal[
+            _StabilityMode.GEN_AI_LATEST_EXPERIMENTAL
+        ],
+        completion_hook: CompletionHook,
+    ) -> None: ...
+
+    @overload
+    def __init__(
+        self,
+        tracer: Tracer,
+        logger: Logger,
+        capture_content: bool,
+        sem_conv_opt_in_mode: Literal[_StabilityMode.DEFAULT],
+        completion_hook: CompletionHook,
+    ) -> None: ...
+
+    def __init__(
+        self,
+        tracer: Tracer,
+        logger: Logger,
+        capture_content: Union[bool, ContentCapturingMode],
+        sem_conv_opt_in_mode: Union[
+            Literal[_StabilityMode.DEFAULT],
+            Literal[_StabilityMode.GEN_AI_LATEST_EXPERIMENTAL],
+        ],
+        completion_hook: CompletionHook,
+    ) -> None:
+        self.tracer = tracer
+        self.logger = logger
+        self.capture_content = capture_content
+        self.sem_conv_opt_in_mode = sem_conv_opt_in_mode
+        self.completion_hook = completion_hook
+
+    @contextmanager
+    def _with_new_instrumentation(
+        self,
+        capture_content: ContentCapturingMode,
+        instance: client.PredictionServiceClient
+        | client_v1beta1.PredictionServiceClient,
+        args: Any,
+        kwargs: Any,
+    ):
+        params = _extract_params(*args, **kwargs)
+        request_attributes = get_genai_request_attributes(True, params)
+        with self.tracer.start_as_current_span(
+            name=f"{GenAI.GenAiOperationNameValues.CHAT.value} {request_attributes.get(GenAI.GEN_AI_REQUEST_MODEL, '')}".strip(),
+            kind=SpanKind.CLIENT,
+        ) as span:
+
+            def handle_response(
+                response: prediction_service.GenerateContentResponse
+                | prediction_service_v1beta1.GenerateContentResponse
+                | None,
+            ) -> None:
+                event = LogRecord(
+                    event_name="gen_ai.client.inference.operation.details",
+                )
+                attributes = (
+                    get_server_attributes(instance.api_endpoint)  # type: ignore[reportUnknownMemberType]
+                    | request_attributes
+                    | get_genai_response_attributes(response)
+                )
+                system_instructions, inputs, outputs = [], [], []
+                if params.system_instruction:
+                    system_instructions = convert_content_to_message_parts(
+                        params.system_instruction
+                    )
+                if params.contents:
+                    inputs = [
+                        InputMessage(
+                            role=content.role,
+                            parts=convert_content_to_message_parts(content),
+                        )
+                        for content in params.contents
+                    ]
+                if response:
+                    outputs = [
+                        OutputMessage(
+                            finish_reason=_map_finish_reason(
+                                candidate.finish_reason
+                            ),
+                            role=candidate.content.role,
+                            parts=convert_content_to_message_parts(
+                                candidate.content
+                            ),
+                        )
+                        for candidate in response.candidates
+                    ]
+                self.completion_hook.on_completion(
+                    inputs=inputs,
+                    outputs=outputs,
+                    system_instruction=system_instructions,
+                    span=span,
+                    log_record=event,
+                )
+                content_attributes = {
+                    k: [asdict(x) for x in v]
+                    for k, v in [
+                        (
+                            GenAI.GEN_AI_SYSTEM_INSTRUCTIONS,
+                            system_instructions,
+                        ),
+                        (GenAI.GEN_AI_INPUT_MESSAGES, inputs),
+                        (GenAI.GEN_AI_OUTPUT_MESSAGES, outputs),
+                    ]
+                    if v
+                }
+                if span.is_recording():
+                    span.set_attributes(attributes)
+                    if capture_content in (
+                        ContentCapturingMode.SPAN_AND_EVENT,
+                        ContentCapturingMode.SPAN_ONLY,
+                    ):
+                        span.set_attributes(
+                            {
+                                k: gen_ai_json_dumps(v)
+                                for k, v in content_attributes.items()
+                            }
+                        )
+                event.attributes = attributes
+                if capture_content in (
+                    ContentCapturingMode.SPAN_AND_EVENT,
+                    ContentCapturingMode.EVENT_ONLY,
+                ):
+                    event.attributes |= content_attributes
+                self.logger.emit(event)
+
+            yield handle_response
+
+    @contextmanager
+    def _with_default_instrumentation(
+        self,
+        capture_content: bool,
+        instance: client.PredictionServiceClient
+        | client_v1beta1.PredictionServiceClient,
+        args: Any,
+        kwargs: Any,
+    ):
+        params = _extract_params(*args, **kwargs)
+        api_endpoint: str = instance.api_endpoint  # type: ignore[reportUnknownMemberType]
+        span_attributes = {
+            **get_genai_request_attributes(False, params),
+            **get_server_attributes(api_endpoint),
+        }
+
+        span_name = get_span_name(span_attributes)
+
+        with self.tracer.start_as_current_span(
+            name=span_name,
+            kind=SpanKind.CLIENT,
+            attributes=span_attributes,
+        ) as span:
+            for event in request_to_events(
+                params=params, capture_content=capture_content
+            ):
+                self.logger.emit(event)
+
+            # TODO: set error.type attribute
+            # https://github.com/open-telemetry/semantic-conventions/blob/main/docs/gen-ai/gen-ai-spans.md
+
+            def handle_response(
+                response: prediction_service.GenerateContentResponse
+                | prediction_service_v1beta1.GenerateContentResponse,
+            ) -> None:
+                if span.is_recording():
+                    # When streaming, this is called multiple times so attributes would be
+                    # overwritten. In practice, it looks the API only returns the interesting
+                    # attributes on the last streamed response. However, I couldn't find
+                    # documentation for this and setting attributes shouldn't be too expensive.
+                    span.set_attributes(
+                        get_genai_response_attributes(response)
+                    )
+
+                for event in response_to_events(
+                    response=response, capture_content=capture_content
+                ):
+                    self.logger.emit(event)
+
+            yield handle_response
+
+    def generate_content(
+        self,
+        wrapped: Callable[
+            ...,
+            prediction_service.GenerateContentResponse
+            | prediction_service_v1beta1.GenerateContentResponse,
+        ],
+        instance: client.PredictionServiceClient
+        | client_v1beta1.PredictionServiceClient,
+        args: Any,
+        kwargs: Any,
+    ) -> (
+        prediction_service.GenerateContentResponse
+        | prediction_service_v1beta1.GenerateContentResponse
+    ):
+        if self.sem_conv_opt_in_mode == _StabilityMode.DEFAULT:
+            capture_content_bool = cast(bool, self.capture_content)
+            with self._with_default_instrumentation(
+                capture_content_bool, instance, args, kwargs
+            ) as handle_response:
+                response = wrapped(*args, **kwargs)
+                handle_response(response)
+                return response
+        else:
+            capture_content = cast(ContentCapturingMode, self.capture_content)
+            with self._with_new_instrumentation(
+                capture_content, instance, args, kwargs
+            ) as handle_response:
+                response = None
+                try:
+                    response = wrapped(*args, **kwargs)
+                    return response
+                finally:
+                    handle_response(response)
+
+    async def agenerate_content(
+        self,
+        wrapped: Callable[
+            ...,
+            Awaitable[
+                prediction_service.GenerateContentResponse
+                | prediction_service_v1beta1.GenerateContentResponse
+            ],
+        ],
+        instance: client.PredictionServiceClient
+        | client_v1beta1.PredictionServiceClient,
+        args: Any,
+        kwargs: Any,
+    ) -> (
+        prediction_service.GenerateContentResponse
+        | prediction_service_v1beta1.GenerateContentResponse
+    ):
+        if self.sem_conv_opt_in_mode == _StabilityMode.DEFAULT:
+            capture_content_bool = cast(bool, self.capture_content)
+            with self._with_default_instrumentation(
+                capture_content_bool, instance, args, kwargs
+            ) as handle_response:
+                response = await wrapped(*args, **kwargs)
+                handle_response(response)
+                return response
+        else:
+            capture_content = cast(ContentCapturingMode, self.capture_content)
+            with self._with_new_instrumentation(
+                capture_content, instance, args, kwargs
+            ) as handle_response:
+                response = None
+                try:
+                    response = await wrapped(*args, **kwargs)
+                    return response
+                finally:
+                    handle_response(response)