opentelemetry-instrumentation-replicate 0.40.14__py3-none-any.whl → 0.42.0__py3-none-any.whl

opentelemetry/instrumentation/replicate/__init__.py

@@ -1,25 +1,33 @@
 """OpenTelemetry Replicate instrumentation"""
 
 import logging
-import os
 import types
 from typing import Collection
-from opentelemetry.instrumentation.replicate.config import Config
-from opentelemetry.instrumentation.replicate.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._events import get_event_logger
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.instrumentation.replicate.config import Config
+from opentelemetry.instrumentation.replicate.event_emitter import (
+    emit_choice_events,
+    emit_event,
+)
+from opentelemetry.instrumentation.replicate.event_models import MessageEvent
+from opentelemetry.instrumentation.replicate.span_utils import (
+    set_input_attributes,
+    set_model_input_attributes,
+    set_response_attributes,
+)
+from opentelemetry.instrumentation.replicate.utils import dont_throw, should_emit_events
+from opentelemetry.instrumentation.replicate.version import __version__
 from opentelemetry.instrumentation.utils import (
     _SUPPRESS_INSTRUMENTATION_KEY,
     unwrap,
 )
-
-from opentelemetry.semconv_ai import SpanAttributes, LLMRequestTypeValues
-from opentelemetry.instrumentation.replicate.version import __version__
+from opentelemetry.semconv_ai import 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__)
 
@@ -44,68 +52,11 @@ WRAPPED_METHODS = [
 ]
 
 
-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 _set_span_attribute(span, name, value):
-    if value is not None:
-        if value != "":
-            span.set_attribute(name, value)
-    return
-
-
-input_attribute_map = {
-    "prompt": f"{SpanAttributes.LLM_PROMPTS}.0.user",
-    "temperature": SpanAttributes.LLM_REQUEST_TEMPERATURE,
-    "top_p": SpanAttributes.LLM_REQUEST_TOP_P,
-}
-
-
-def _set_input_attributes(span, args, kwargs):
-    if args is not None and len(args) > 0:
-        _set_span_attribute(span, SpanAttributes.LLM_REQUEST_MODEL, args[0])
-    elif kwargs.get("version"):
-        _set_span_attribute(
-            span, SpanAttributes.LLM_REQUEST_MODEL, kwargs.get("version").id
-        )
-    else:
-        _set_span_attribute(span, SpanAttributes.LLM_REQUEST_MODEL, "unknown")
-
-    input_attribute = kwargs.get("input")
-    for key in input_attribute:
-        if key in input_attribute_map:
-            if key == "prompt" and not should_send_prompts():
-                continue
-            _set_span_attribute(
-                span,
-                input_attribute_map.get(key, f"llm.request.{key}"),
-                input_attribute.get(key),
-            )
-    return
-
-
-@dont_throw
-def _set_response_attributes(span, response):
-    if should_send_prompts():
-        if isinstance(response, list):
-            for index, item in enumerate(response):
-                prefix = f"{SpanAttributes.LLM_COMPLETIONS}.{index}"
-                _set_span_attribute(span, f"{prefix}.content", item)
-        elif isinstance(response, str):
-            _set_span_attribute(
-                span, f"{SpanAttributes.LLM_COMPLETIONS}.0.content", response
-            )
-    return
-
-
-def _build_from_streaming_response(span, response):
+def _build_from_streaming_response(span, event_logger, response):
     complete_response = ""
     for item in response:
         item_to_yield = item
@@ -113,32 +64,40 @@ def _build_from_streaming_response(span, response):
 
         yield item_to_yield
 
-    _set_response_attributes(span, complete_response)
+    _handle_response(span, event_logger, complete_response)
 
-    span.set_status(Status(StatusCode.OK))
     span.end()
 
 
 @dont_throw
-def _handle_request(span, args, kwargs):
-    if span.is_recording():
-        _set_input_attributes(span, args, kwargs)
+def _handle_request(span, event_logger, args, kwargs):
+    set_model_input_attributes(span, args, kwargs)
+
+    model_input = kwargs.get("input") or (args[1] if len(args) > 1 else None)
+
+    if should_emit_events() and event_logger:
+        emit_event(MessageEvent(content=model_input.get("prompt")), event_logger)
+    else:
+        set_input_attributes(span, args, kwargs)
 
 
 @dont_throw
-def _handle_response(span, response):
-    if span.is_recording():
-        _set_response_attributes(span, response)
+def _handle_response(span, event_logger, response):
+    if should_emit_events() and event_logger:
+        emit_choice_events(response, event_logger)
+    else:
+        set_response_attributes(span, response)
 
+    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, to_wrap):
+    def _with_tracer(tracer, event_logger, to_wrap):
         def wrapper(wrapped, instance, args, kwargs):
-            return func(tracer, to_wrap, wrapped, instance, args, kwargs)
+            return func(tracer, event_logger, to_wrap, wrapped, instance, args, kwargs)
 
         return wrapper
 
@@ -146,7 +105,15 @@ def _with_tracer_wrapper(func):
 
 
 @_with_tracer_wrapper
-def _wrap(tracer, to_wrap, wrapped, instance, args, kwargs):
+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):
         return wrapped(*args, **kwargs)
@@ -161,15 +128,15 @@ def _wrap(tracer, to_wrap, wrapped, instance, args, kwargs):
         },
     )
 
-    _handle_request(span, args, kwargs)
+    _handle_request(span, event_logger, args, kwargs)
 
     response = wrapped(*args, **kwargs)
 
     if response:
         if is_streaming_response(response):
-            return _build_from_streaming_response(span, response)
+            return _build_from_streaming_response(span, event_logger, response)
         else:
-            _handle_response(span, response)
+            _handle_response(span, event_logger, response)
 
     span.end()
     return response
@@ -178,9 +145,10 @@ def _wrap(tracer, to_wrap, wrapped, instance, args, kwargs):
 class ReplicateInstrumentor(BaseInstrumentor):
     """An instrumentor for Replicate's client library."""
 
-    def __init__(self, exception_logger=None):
+    def __init__(self, exception_logger=None, use_legacy_attributes=True):
         super().__init__()
         Config.exception_logger = exception_logger
+        Config.use_legacy_attributes = use_legacy_attributes
 
     def instrumentation_dependencies(self) -> Collection[str]:
         return _instruments
@@ -188,13 +156,23 @@ class ReplicateInstrumentor(BaseInstrumentor):
     def _instrument(self, **kwargs):
         tracer_provider = kwargs.get("tracer_provider")
         tracer = get_tracer(__name__, __version__, tracer_provider)
+
+        event_logger = None
+        if not Config.use_legacy_attributes:
+            event_logger_provider = kwargs.get("event_logger_provider")
+            event_logger = get_event_logger(
+                __name__, __version__, event_logger_provider=event_logger_provider
+            )
+
         for wrapper_method in WRAPPED_METHODS:
             wrap_function_wrapper(
                 wrapper_method.get("module"),
                 wrapper_method.get("method"),
-                _wrap(tracer, wrapper_method),
+                _wrap(tracer, event_logger, wrapper_method),
             )
 
     def _uninstrument(self, **kwargs):
+        import replicate
+
         for wrapper_method in WRAPPED_METHODS:
-            unwrap(wrapper_method.get("module"), wrapper_method.get("method", ""))
+            unwrap(replicate, wrapper_method.get("method", ""))

opentelemetry/instrumentation/replicate/config.py

@@ -1,2 +1,3 @@
 class Config:
     exception_logger = None
+    use_legacy_attributes = True

opentelemetry/instrumentation/replicate/event_emitter.py

@@ -0,0 +1,134 @@
+from dataclasses import asdict
+from enum import Enum
+from typing import Union
+
+from opentelemetry._events import Event, EventLogger
+from opentelemetry.instrumentation.replicate.event_models import (
+    ChoiceEvent,
+    MessageEvent,
+)
+from opentelemetry.instrumentation.replicate.utils import (
+    dont_throw,
+    should_emit_events,
+    should_send_prompts,
+)
+from opentelemetry.semconv._incubating.attributes import (
+    gen_ai_attributes as GenAIAttributes,
+)
+
+from replicate.prediction import Prediction
+
+
+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: "replicate"}
+"""The attributes to be used for the event."""
+
+
+@dont_throw
+def emit_choice_events(
+    response: Union[str, list, Prediction], event_logger: Union[EventLogger, None]
+):
+    # Handle replicate.run responses
+    if isinstance(response, list):
+        for i, generation in enumerate(response):
+            emit_event(
+                ChoiceEvent(
+                    index=i, message={"content": generation, "role": "assistant"}
+                ),
+                event_logger,
+            )
+    # Handle replicate.predictions.create responses
+    elif isinstance(response, Prediction):
+        emit_event(
+            ChoiceEvent(
+                index=0, message={"content": response.output, "role": "assistant"}
+            ),
+            event_logger,
+        )
+    # Handle replicate.stream responses built from _build_from_streaming_response
+    elif isinstance(response, str):
+        emit_event(
+            ChoiceEvent(index=0, message={"content": response, "role": "assistant"}),
+            event_logger,
+        )
+    else:
+        raise ValueError(
+            "It wasn't possible to emit the choice events due to an unsupported response type"
+        )
+
+
+def emit_event(
+    event: Union[MessageEvent, ChoiceEvent], event_logger: Union[EventLogger, None]
+) -> 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: EventLogger) -> 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)
+
+    event_logger.emit(Event(name=name, body=body, attributes=EVENT_ATTRIBUTES))
+
+
+def _emit_choice_event(event: ChoiceEvent, event_logger: EventLogger) -> 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)
+
+    event_logger.emit(
+        Event(name="gen_ai.choice", body=body, attributes=EVENT_ATTRIBUTES)
+    )

opentelemetry/instrumentation/replicate/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

opentelemetry/instrumentation/replicate/span_utils.py

@@ -0,0 +1,61 @@
+from opentelemetry.instrumentation.replicate.utils import (
+    dont_throw,
+    should_send_prompts,
+)
+from opentelemetry.semconv_ai import SpanAttributes
+
+
+def _set_span_attribute(span, name, value):
+    if value is not None:
+        if value != "":
+            span.set_attribute(name, value)
+    return
+
+
+@dont_throw
+def set_input_attributes(span, args, kwargs):
+    if not span.is_recording():
+        return
+
+    input_attribute = kwargs.get("input")
+    if should_send_prompts():
+        _set_span_attribute(
+            span, f"{SpanAttributes.LLM_PROMPTS}.0.user", input_attribute.get("prompt")
+        )
+
+
+@dont_throw
+def set_model_input_attributes(span, args, kwargs):
+    if not span.is_recording():
+        return
+
+    if args is not None and len(args) > 0:
+        _set_span_attribute(span, SpanAttributes.LLM_REQUEST_MODEL, args[0])
+    elif kwargs.get("version"):
+        _set_span_attribute(
+            span, SpanAttributes.LLM_REQUEST_MODEL, kwargs.get("version").id
+        )
+    else:
+        _set_span_attribute(span, SpanAttributes.LLM_REQUEST_MODEL, "unknown")
+
+    input_attribute = kwargs.get("input")
+
+    _set_span_attribute(
+        span, SpanAttributes.LLM_REQUEST_TEMPERATURE, input_attribute.get("temperature")
+    )
+    _set_span_attribute(
+        span, SpanAttributes.LLM_REQUEST_TOP_P, input_attribute.get("top_p")
+    )
+
+
+@dont_throw
+def set_response_attributes(span, response):
+    if should_send_prompts():
+        if isinstance(response, list):
+            for index, item in enumerate(response):
+                prefix = f"{SpanAttributes.LLM_COMPLETIONS}.{index}"
+                _set_span_attribute(span, f"{prefix}.content", item)
+        elif isinstance(response, str):
+            _set_span_attribute(
+                span, f"{SpanAttributes.LLM_COMPLETIONS}.0.content", response
+            )

opentelemetry/instrumentation/replicate/utils.py

@@ -1,7 +1,18 @@
 import logging
+import os
 import traceback
+
+from opentelemetry import context as context_api
 from opentelemetry.instrumentation.replicate.config import Config
 
+TRACELOOP_TRACE_CONTENT = "TRACELOOP_TRACE_CONTENT"
+
+
+def should_send_prompts():
+    return (
+        os.getenv(TRACELOOP_TRACE_CONTENT) or "true"
+    ).lower() == "true" or context_api.get_value("override_enable_content_tracing")
+
 
 def dont_throw(func):
     """
@@ -26,3 +37,12 @@ def dont_throw(func):
                 Config.exception_logger(e)
 
     return wrapper
+
+
+def should_emit_events() -> bool:
+    """
+    Checks if the instrumentation isn't using the legacy attributes
+    and if the event logger is not None.
+    """
+
+    return not Config.use_legacy_attributes

opentelemetry/instrumentation/replicate/version.py

@@ -1 +1 @@
-__version__ = "0.40.14"
+__version__ = "0.42.0"

{opentelemetry_instrumentation_replicate-0.40.14.dist-info → opentelemetry_instrumentation_replicate-0.42.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: opentelemetry-instrumentation-replicate
-Version: 0.40.14
+Version: 0.42.0
 Summary: OpenTelemetry Replicate instrumentation
 License: Apache-2.0
 Author: Kartik Prajapati
@@ -17,7 +17,7 @@ Provides-Extra: instruments
 Requires-Dist: opentelemetry-api (>=1.28.0,<2.0.0)
 Requires-Dist: opentelemetry-instrumentation (>=0.50b0)
 Requires-Dist: opentelemetry-semantic-conventions (>=0.50b0)
-Requires-Dist: opentelemetry-semantic-conventions-ai (==0.4.9)
+Requires-Dist: opentelemetry-semantic-conventions-ai (==0.4.11)
 Project-URL: Repository, https://github.com/traceloop/openllmetry/tree/main/packages/opentelemetry-instrumentation-replicate
 Description-Content-Type: text/markdown
 

opentelemetry_instrumentation_replicate-0.42.0.dist-info/RECORD

@@ -0,0 +1,11 @@
+opentelemetry/instrumentation/replicate/__init__.py,sha256=efxMBZKhkhfCmAQDEFLeHQqE1JKZmtQW19h44FRj0JQ,5262
+opentelemetry/instrumentation/replicate/config.py,sha256=aaFMaCpJS9eAxkHIP88wvlc8B1tKspysmXdULfhWGSA,75
+opentelemetry/instrumentation/replicate/event_emitter.py,sha256=1TMYvK9IVPnb6IIDaD_Uumay2LCoGxPHqwkDFqJ35R4,4374
+opentelemetry/instrumentation/replicate/event_models.py,sha256=PCfCGxrrArwZqR-4wFcXrhwQq0sBMAxmSrpC4PUMtaM,876
+opentelemetry/instrumentation/replicate/span_utils.py,sha256=6--CZZvb7eAMUyGlKRRZ5B0tF4YgUgRwRDN9NMqKZT8,1839
+opentelemetry/instrumentation/replicate/utils.py,sha256=wtv-bT3CojjX6ACbDtvehZDsINnDHM8s881JOAz0tAo,1307
+opentelemetry/instrumentation/replicate/version.py,sha256=fXDp5M3ZbGNxa2CR646QIbRo8oBi15GhGQT_3rvssIE,23
+opentelemetry_instrumentation_replicate-0.42.0.dist-info/METADATA,sha256=haOfO4RkLJwH4CMj_zvjAOkTpXTlkzmFz8jR_-bI0lA,2201
+opentelemetry_instrumentation_replicate-0.42.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+opentelemetry_instrumentation_replicate-0.42.0.dist-info/entry_points.txt,sha256=YDOSAhxrkimcGYNOJLwqT1VoPS9YM-mXd1DT4hBk0Oc,102
+opentelemetry_instrumentation_replicate-0.42.0.dist-info/RECORD,,

opentelemetry_instrumentation_replicate-0.40.14.dist-info/RECORD

@@ -1,8 +0,0 @@
-opentelemetry/instrumentation/replicate/__init__.py,sha256=RA_x-p8zpGT0MqW26iOau4FTXcTJoRHyo6rX7WpwtpE,5808
-opentelemetry/instrumentation/replicate/config.py,sha256=CtypZov_ytI9nSrfN9lWnjcufbAR9sfkXRA0OstDEUw,42
-opentelemetry/instrumentation/replicate/utils.py,sha256=E4Ur_vkGqrfsNpFJ4oU72GQudgyH0-NFsbOYjLfEkIo,809
-opentelemetry/instrumentation/replicate/version.py,sha256=TzmqqRPz5JsMF0vCMChofQC_r_x0W9P-JB4K5rRCvtE,24
-opentelemetry_instrumentation_replicate-0.40.14.dist-info/METADATA,sha256=ut4Jwr1A5PpR89G8rorz9A3oQQFhEjuhBLutxsyzDnY,2201
-opentelemetry_instrumentation_replicate-0.40.14.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-opentelemetry_instrumentation_replicate-0.40.14.dist-info/entry_points.txt,sha256=YDOSAhxrkimcGYNOJLwqT1VoPS9YM-mXd1DT4hBk0Oc,102
-opentelemetry_instrumentation_replicate-0.40.14.dist-info/RECORD,,