opentelemetry-instrumentation-vertexai 0.12.4__tar.gz → 0.49.6__tar.gz

{opentelemetry_instrumentation_vertexai-0.12.4 → opentelemetry_instrumentation_vertexai-0.49.6}/PKG-INFO

@@ -1,20 +1,25 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.4
 Name: opentelemetry-instrumentation-vertexai
-Version: 0.12.4
+Version: 0.49.6
 Summary: OpenTelemetry Vertex AI instrumentation
 License: Apache-2.0
 Author: Gal Kleinman
 Author-email: gal@traceloop.com
-Requires-Python: >=3.8.1,<4
+Requires-Python: >=3.9,<4
 Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
-Requires-Dist: opentelemetry-api (>=1.22.0,<2.0.0)
-Requires-Dist: opentelemetry-instrumentation (==0.43b0)
-Requires-Dist: opentelemetry-semantic-conventions-ai (>=0.0.20,<0.0.21)
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Provides-Extra: instruments
+Requires-Dist: opentelemetry-api (>=1.38.0,<2.0.0)
+Requires-Dist: opentelemetry-instrumentation (>=0.59b0)
+Requires-Dist: opentelemetry-semantic-conventions (>=0.59b0)
+Requires-Dist: opentelemetry-semantic-conventions-ai (>=0.4.13,<0.5.0)
+Project-URL: Repository, https://github.com/traceloop/openllmetry/tree/main/packages/opentelemetry-instrumentation-vertexai
 Description-Content-Type: text/markdown
 
 # OpenTelemetry VertexAI Instrumentation
@@ -31,6 +36,14 @@ This library allows tracing VertexAI prompts and completions sent with the offic
 pip install opentelemetry-instrumentation-vertexai
 ```
 
+## Example usage
+
+```python
+from opentelemetry.instrumentation.vertexai import VertexAIInstrumentor
+
+VertexAIInstrumentor().instrument()
+```
+
 ## Privacy
 
 **By default, this instrumentation logs prompts, completions, and embeddings to span attributes**. This gives you a clear visibility into how your LLM application is working, and can make it easy to debug and evaluate the quality of the outputs.

{opentelemetry_instrumentation_vertexai-0.12.4 → opentelemetry_instrumentation_vertexai-0.49.6}/README.md

@@ -12,6 +12,14 @@ This library allows tracing VertexAI prompts and completions sent with the offic
 pip install opentelemetry-instrumentation-vertexai
 ```
 
+## Example usage
+
+```python
+from opentelemetry.instrumentation.vertexai import VertexAIInstrumentor
+
+VertexAIInstrumentor().instrument()
+```
+
 ## Privacy
 
 **By default, this instrumentation logs prompts, completions, and embeddings to span attributes**. This gives you a clear visibility into how your LLM application is working, and can make it easy to debug and evaluate the quality of the outputs.

opentelemetry_instrumentation_vertexai-0.49.6/opentelemetry/instrumentation/vertexai/__init__.py

@@ -0,0 +1,369 @@
+"""OpenTelemetry Vertex AI instrumentation"""
+
+import logging
+import types
+from typing import Collection
+
+from opentelemetry import context as context_api
+from opentelemetry._logs import get_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._incubating.attributes import (
+    gen_ai_attributes as GenAIAttributes,
+)
+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
+    )
+
+    span.set_status(Status(StatusCode.OK))
+    span.end()
+
+
+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
+
+        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
+
+    name = to_wrap.get("span_name")
+    span = tracer.start_span(
+        name,
+        kind=SpanKind.CLIENT,
+        attributes={
+            GenAIAttributes.GEN_AI_SYSTEM: "Google",
+            SpanAttributes.LLM_REQUEST_TYPE: LLMRequestTypeValues.COMPLETION.value,
+        },
+    )
+
+    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={
+            GenAIAttributes.GEN_AI_SYSTEM: "Google",
+            SpanAttributes.LLM_REQUEST_TYPE: LLMRequestTypeValues.COMPLETION.value,
+        },
+    )
+
+    # 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)
+
+    span.end()
+    return response
+
+
+class VertexAIInstrumentor(BaseInstrumentor):
+    """An instrumentor for VertextAI's client library."""
+
+    def __init__(self, exception_logger=None, use_legacy_attributes=True, upload_base64_image=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
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return _instruments
+
+    def _instrument(self, **kwargs):
+        tracer_provider = kwargs.get("tracer_provider")
+        tracer = get_tracer(__name__, __version__, tracer_provider)
+
+        event_logger = None
+
+        if should_emit_events():
+            logger_provider = kwargs.get("logger_provider")
+            event_logger = get_logger(
+                __name__,
+                __version__,
+                logger_provider=logger_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")
+
+            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)
+                ),
+            )
+
+    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", ""),
+            )

opentelemetry_instrumentation_vertexai-0.49.6/opentelemetry/instrumentation/vertexai/config.py

@@ -0,0 +1,9 @@
+from typing import Callable
+
+
+class Config:
+    exception_logger = None
+    use_legacy_attributes = True
+    upload_base64_image: Callable[[str, str, str, str], str] = (
+        lambda trace_id, span_id, image_name, base64_string: str
+    )

opentelemetry_instrumentation_vertexai-0.49.6/opentelemetry/instrumentation/vertexai/event_emitter.py

@@ -0,0 +1,173 @@
+from dataclasses import asdict
+from enum import Enum
+from typing import Union
+
+from opentelemetry._logs import LogRecord
+from opentelemetry.instrumentation.vertexai.event_models import (
+    ChoiceEvent,
+    MessageEvent,
+)
+from opentelemetry.instrumentation.vertexai.utils import (
+    dont_throw,
+    should_emit_events,
+    should_send_prompts,
+)
+from opentelemetry.semconv._incubating.attributes import (
+    gen_ai_attributes as GenAIAttributes,
+)
+
+from vertexai.generative_models import GenerationResponse
+
+
+class Roles(Enum):
+    USER = "user"
+    ASSISTANT = "assistant"
+    SYSTEM = "system"
+    TOOL = "tool"
+
+
+VALID_MESSAGE_ROLES = {role.value for role in Roles}
+"""The valid roles for naming the message event."""
+
+EVENT_ATTRIBUTES = {
+    GenAIAttributes.GEN_AI_SYSTEM: GenAIAttributes.GenAiSystemValues.VERTEX_AI.value
+}
+"""The attributes to be used for the event."""
+
+
+def _parse_vertex_finish_reason(reason):
+    if reason is None:
+        return "unknown"
+
+    finish_reason_map = {
+        0: "unspecified",
+        1: "stop",
+        2: "max_tokens",
+        3: "safety",
+        4: "recitation",
+        5: "other",
+        6: "blocklist",
+        7: "prohibited_content",
+        8: "spii",
+        9: "malformed_function_call",
+    }
+
+    if hasattr(reason, "value"):
+        reason_value = reason.value
+    else:
+        reason_value = reason
+
+    return finish_reason_map.get(reason_value, "unknown")
+
+
+@dont_throw
+def emit_prompt_events(args, event_logger):
+    prompt = ""
+    if args is not None and len(args) > 0:
+        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"
+    emit_event(MessageEvent(content=prompt, role=Roles.USER.value), event_logger)
+
+
+def emit_response_events(response, event_logger):
+    if isinstance(response, str):
+        emit_event(
+            ChoiceEvent(
+                index=0,
+                message={"content": response, "role": Roles.ASSISTANT.value},
+                finish_reason="unknown",
+            ),
+            event_logger,
+        )
+    elif isinstance(response, GenerationResponse):
+        for candidate in response.candidates:
+            emit_event(
+                ChoiceEvent(
+                    index=candidate.index,
+                    message={
+                        "content": candidate.text,
+                        "role": Roles.ASSISTANT.value,
+                    },
+                    finish_reason=_parse_vertex_finish_reason(candidate.finish_reason),
+                ),
+                event_logger,
+            )
+
+
+def emit_event(event: Union[MessageEvent, ChoiceEvent], event_logger) -> None:
+    """
+    Emit an event to the OpenTelemetry SDK.
+
+    Args:
+        event: The event to emit.
+    """
+    if not should_emit_events() or event_logger is None:
+        return
+
+    if isinstance(event, MessageEvent):
+        _emit_message_event(event, event_logger)
+    elif isinstance(event, ChoiceEvent):
+        _emit_choice_event(event, event_logger)
+    else:
+        raise TypeError("Unsupported event type")
+
+
+def _emit_message_event(event: MessageEvent, event_logger) -> None:
+    body = asdict(event)
+
+    if event.role in VALID_MESSAGE_ROLES:
+        name = "gen_ai.{}.message".format(event.role)
+        # According to the semantic conventions, the role is conditionally required if available
+        # and not equal to the "role" in the message name. So, remove the role from the body if
+        # it is the same as the in the event name.
+        body.pop("role", None)
+    else:
+        name = "gen_ai.user.message"
+
+    # According to the semantic conventions, only the assistant role has tool call
+    if event.role != Roles.ASSISTANT.value and event.tool_calls is not None:
+        del body["tool_calls"]
+    elif event.tool_calls is None:
+        del body["tool_calls"]
+
+    if not should_send_prompts():
+        del body["content"]
+        if body.get("tool_calls") is not None:
+            for tool_call in body["tool_calls"]:
+                tool_call["function"].pop("arguments", None)
+
+    log_record = LogRecord(
+        body=body,
+        attributes=EVENT_ATTRIBUTES,
+        event_name=name
+    )
+    event_logger.emit(log_record)
+
+
+def _emit_choice_event(event: ChoiceEvent, event_logger) -> None:
+    body = asdict(event)
+    if event.message["role"] == Roles.ASSISTANT.value:
+        # According to the semantic conventions, the role is conditionally required if available
+        # and not equal to "assistant", so remove the role from the body if it is "assistant".
+        body["message"].pop("role", None)
+
+    if event.tool_calls is None:
+        del body["tool_calls"]
+
+    if not should_send_prompts():
+        body["message"].pop("content", None)
+        if body.get("tool_calls") is not None:
+            for tool_call in body["tool_calls"]:
+                tool_call["function"].pop("arguments", None)
+
+    log_record = LogRecord(
+        body=body,
+        attributes=EVENT_ATTRIBUTES,
+        event_name="gen_ai.choice"
+
+    )
+    event_logger.emit(log_record)

opentelemetry_instrumentation_vertexai-0.49.6/opentelemetry/instrumentation/vertexai/event_models.py

@@ -0,0 +1,41 @@
+from dataclasses import dataclass
+from typing import Any, List, Literal, Optional, TypedDict
+
+
+class _FunctionToolCall(TypedDict):
+    function_name: str
+    arguments: Optional[dict[str, Any]]
+
+
+class ToolCall(TypedDict):
+    """Represents a tool call in the AI model."""
+
+    id: str
+    function: _FunctionToolCall
+    type: Literal["function"]
+
+
+class CompletionMessage(TypedDict):
+    """Represents a message in the AI model."""
+
+    content: Any
+    role: str = "assistant"
+
+
+@dataclass
+class MessageEvent:
+    """Represents an input event for the AI model."""
+
+    content: Any
+    role: str = "user"
+    tool_calls: Optional[List[ToolCall]] = None
+
+
+@dataclass
+class ChoiceEvent:
+    """Represents a completion event for the AI model."""
+
+    index: int
+    message: CompletionMessage
+    finish_reason: str = "unknown"
+    tool_calls: Optional[List[ToolCall]] = None