lmnr 0.6.20__py3-none-any.whl → 0.7.0__py3-none-any.whl

lmnr/opentelemetry_lib/opentelemetry/instrumentation/groq/__init__.py

@@ -0,0 +1,488 @@
+"""OpenTelemetry Groq instrumentation"""
+
+import logging
+import os
+import time
+from typing import Callable, Collection, Union
+
+from opentelemetry import context as context_api
+from opentelemetry._events import EventLogger, get_event_logger
+from .config import Config
+from .event_emitter import (
+    emit_choice_events,
+    emit_message_events,
+    emit_streaming_response_events,
+)
+from .span_utils import (
+    set_input_attributes,
+    set_model_input_attributes,
+    set_model_response_attributes,
+    set_model_streaming_response_attributes,
+    set_response_attributes,
+    set_streaming_response_attributes,
+)
+from .utils import (
+    error_metrics_attributes,
+    shared_metrics_attributes,
+    should_emit_events,
+)
+from .version import __version__
+from lmnr.opentelemetry_lib.tracing.context import get_current_context
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY, unwrap
+from opentelemetry.metrics import Counter, Histogram, Meter, get_meter
+from opentelemetry.semconv_ai import (
+    SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY,
+    LLMRequestTypeValues,
+    Meters,
+    SpanAttributes,
+)
+from opentelemetry.trace import SpanKind, Tracer, get_tracer
+from opentelemetry.trace.status import Status, StatusCode
+from wrapt import wrap_function_wrapper
+
+from groq._streaming import AsyncStream, Stream
+
+logger = logging.getLogger(__name__)
+
+_instruments = ("groq >= 0.9.0",)
+
+
+WRAPPED_METHODS = [
+    {
+        "package": "groq.resources.chat.completions",
+        "object": "Completions",
+        "method": "create",
+        "span_name": "groq.chat",
+    },
+]
+WRAPPED_AMETHODS = [
+    {
+        "package": "groq.resources.chat.completions",
+        "object": "AsyncCompletions",
+        "method": "create",
+        "span_name": "groq.chat",
+    },
+]
+
+
+def is_streaming_response(response):
+    return isinstance(response, Stream) or isinstance(response, AsyncStream)
+
+
+def _with_chat_telemetry_wrapper(func):
+    """Helper for providing tracer for wrapper functions. Includes metric collectors."""
+
+    def _with_chat_telemetry(
+        tracer,
+        token_histogram,
+        choice_counter,
+        duration_histogram,
+        event_logger,
+        to_wrap,
+    ):
+        def wrapper(wrapped, instance, args, kwargs):
+            return func(
+                tracer,
+                token_histogram,
+                choice_counter,
+                duration_histogram,
+                event_logger,
+                to_wrap,
+                wrapped,
+                instance,
+                args,
+                kwargs,
+            )
+
+        return wrapper
+
+    return _with_chat_telemetry
+
+
+def _create_metrics(meter: Meter):
+    token_histogram = meter.create_histogram(
+        name=Meters.LLM_TOKEN_USAGE,
+        unit="token",
+        description="Measures number of input and output tokens used",
+    )
+
+    choice_counter = meter.create_counter(
+        name=Meters.LLM_GENERATION_CHOICES,
+        unit="choice",
+        description="Number of choices returned by chat completions call",
+    )
+
+    duration_histogram = meter.create_histogram(
+        name=Meters.LLM_OPERATION_DURATION,
+        unit="s",
+        description="GenAI operation duration",
+    )
+
+    return token_histogram, choice_counter, duration_histogram
+
+
+def _process_streaming_chunk(chunk):
+    """Extract content, finish_reason and usage from a streaming chunk."""
+    if not chunk.choices:
+        return None, None, None
+
+    delta = chunk.choices[0].delta
+    content = delta.content if hasattr(delta, "content") else None
+    finish_reason = chunk.choices[0].finish_reason
+
+    # Extract usage from x_groq if present in the final chunk
+    usage = None
+    if hasattr(chunk, "x_groq") and chunk.x_groq and chunk.x_groq.usage:
+        usage = chunk.x_groq.usage
+
+    return content, finish_reason, usage
+
+
+def _handle_streaming_response(
+    span, accumulated_content, finish_reason, usage, event_logger
+):
+    set_model_streaming_response_attributes(span, usage)
+    if should_emit_events() and event_logger:
+        emit_streaming_response_events(accumulated_content, finish_reason, event_logger)
+    else:
+        set_streaming_response_attributes(
+            span, accumulated_content, finish_reason, usage
+        )
+
+
+def _create_stream_processor(response, span, event_logger):
+    """Create a generator that processes a stream while collecting telemetry."""
+    accumulated_content = ""
+    finish_reason = None
+    usage = None
+
+    for chunk in response:
+        content, chunk_finish_reason, chunk_usage = _process_streaming_chunk(chunk)
+        if content:
+            accumulated_content += content
+        if chunk_finish_reason:
+            finish_reason = chunk_finish_reason
+        if chunk_usage:
+            usage = chunk_usage
+        yield chunk
+
+    _handle_streaming_response(
+        span, accumulated_content, finish_reason, usage, event_logger
+    )
+
+    if span.is_recording():
+        span.set_status(Status(StatusCode.OK))
+
+    span.end()
+
+
+async def _create_async_stream_processor(response, span, event_logger):
+    """Create an async generator that processes a stream while collecting telemetry."""
+    accumulated_content = ""
+    finish_reason = None
+    usage = None
+
+    async for chunk in response:
+        content, chunk_finish_reason, chunk_usage = _process_streaming_chunk(chunk)
+        if content:
+            accumulated_content += content
+        if chunk_finish_reason:
+            finish_reason = chunk_finish_reason
+        if chunk_usage:
+            usage = chunk_usage
+        yield chunk
+
+    _handle_streaming_response(
+        span, accumulated_content, finish_reason, usage, event_logger
+    )
+
+    if span.is_recording():
+        span.set_status(Status(StatusCode.OK))
+
+    span.end()
+
+
+def _handle_input(span, kwargs, event_logger):
+    set_model_input_attributes(span, kwargs)
+    if should_emit_events() and event_logger:
+        emit_message_events(kwargs, event_logger)
+    else:
+        set_input_attributes(span, kwargs)
+
+
+def _handle_response(span, response, token_histogram, event_logger):
+    set_model_response_attributes(span, response, token_histogram)
+    if should_emit_events() and event_logger:
+        emit_choice_events(response, event_logger)
+    else:
+        set_response_attributes(span, response)
+
+
+@_with_chat_telemetry_wrapper
+def _wrap(
+    tracer: Tracer,
+    token_histogram: Histogram,
+    choice_counter: Counter,
+    duration_histogram: Histogram,
+    event_logger: Union[EventLogger, None],
+    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)
+
+    name = to_wrap.get("span_name")
+    span = tracer.start_span(
+        name,
+        kind=SpanKind.CLIENT,
+        attributes={
+            SpanAttributes.LLM_SYSTEM: "Groq",
+            SpanAttributes.LLM_REQUEST_TYPE: LLMRequestTypeValues.COMPLETION.value,
+        },
+        context=get_current_context(),
+    )
+
+    _handle_input(span, kwargs, event_logger)
+
+    start_time = time.time()
+    try:
+        response = wrapped(*args, **kwargs)
+    except Exception as e:  # pylint: disable=broad-except
+        end_time = time.time()
+        attributes = error_metrics_attributes(e)
+
+        if duration_histogram:
+            duration = end_time - start_time
+            duration_histogram.record(duration, attributes=attributes)
+
+        raise e
+
+    end_time = time.time()
+
+    if is_streaming_response(response):
+        try:
+            return _create_stream_processor(response, span, event_logger)
+        except Exception as ex:
+            logger.warning(
+                "Failed to process streaming response for groq span, error: %s",
+                str(ex),
+            )
+            span.set_status(Status(StatusCode.ERROR))
+            span.end()
+            raise
+    elif response:
+        try:
+            metric_attributes = shared_metrics_attributes(response)
+
+            if duration_histogram:
+                duration = time.time() - start_time
+                duration_histogram.record(
+                    duration,
+                    attributes=metric_attributes,
+                )
+
+            _handle_response(span, response, token_histogram, event_logger)
+
+        except Exception as ex:  # pylint: disable=broad-except
+            logger.warning(
+                "Failed to set response attributes for groq span, error: %s",
+                str(ex),
+            )
+
+        if span.is_recording():
+            span.set_status(Status(StatusCode.OK))
+    span.end()
+    return response
+
+
+@_with_chat_telemetry_wrapper
+async def _awrap(
+    tracer,
+    token_histogram: Histogram,
+    choice_counter: Counter,
+    duration_histogram: Histogram,
+    event_logger: Union[EventLogger, None],
+    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)
+
+    name = to_wrap.get("span_name")
+    span = tracer.start_span(
+        name,
+        kind=SpanKind.CLIENT,
+        attributes={
+            SpanAttributes.LLM_SYSTEM: "Groq",
+            SpanAttributes.LLM_REQUEST_TYPE: LLMRequestTypeValues.COMPLETION.value,
+        },
+        context=get_current_context(),
+    )
+
+    _handle_input(span, kwargs, event_logger)
+
+    start_time = time.time()
+
+    try:
+        response = await wrapped(*args, **kwargs)
+    except Exception as e:  # pylint: disable=broad-except
+        end_time = time.time()
+        attributes = error_metrics_attributes(e)
+
+        if duration_histogram:
+            duration = end_time - start_time
+            duration_histogram.record(duration, attributes=attributes)
+
+        raise e
+
+    end_time = time.time()
+
+    if is_streaming_response(response):
+        try:
+            return await _create_async_stream_processor(response, span, event_logger)
+        except Exception as ex:
+            logger.warning(
+                "Failed to process streaming response for groq span, error: %s",
+                str(ex),
+            )
+            span.set_status(Status(StatusCode.ERROR))
+            span.end()
+            raise
+    elif response:
+        metric_attributes = shared_metrics_attributes(response)
+
+        if duration_histogram:
+            duration = time.time() - start_time
+            duration_histogram.record(
+                duration,
+                attributes=metric_attributes,
+            )
+
+        _handle_response(span, response, token_histogram, event_logger)
+
+        if span.is_recording():
+            span.set_status(Status(StatusCode.OK))
+    span.end()
+    return response
+
+
+def is_metrics_enabled() -> bool:
+    return (os.getenv("TRACELOOP_METRICS_ENABLED") or "true").lower() == "true"
+
+
+class GroqInstrumentor(BaseInstrumentor):
+    """An instrumentor for Groq's client library."""
+
+    def __init__(
+        self,
+        enrich_token_usage: bool = False,
+        exception_logger=None,
+        use_legacy_attributes: bool = True,
+        get_common_metrics_attributes: Callable[[], dict] = lambda: {},
+    ):
+        super().__init__()
+        Config.exception_logger = exception_logger
+        Config.enrich_token_usage = enrich_token_usage
+        Config.get_common_metrics_attributes = get_common_metrics_attributes
+        Config.use_legacy_attributes = use_legacy_attributes
+
+    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)
+
+        # meter and counters are inited here
+        meter_provider = kwargs.get("meter_provider")
+        meter = get_meter(__name__, __version__, meter_provider)
+
+        if is_metrics_enabled():
+            (
+                token_histogram,
+                choice_counter,
+                duration_histogram,
+            ) = _create_metrics(meter)
+        else:
+            (
+                token_histogram,
+                choice_counter,
+                duration_histogram,
+            ) = (None, None, None)
+
+        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 wrapped_method in WRAPPED_METHODS:
+            wrap_package = wrapped_method.get("package")
+            wrap_object = wrapped_method.get("object")
+            wrap_method = wrapped_method.get("method")
+
+            try:
+                wrap_function_wrapper(
+                    wrap_package,
+                    f"{wrap_object}.{wrap_method}",
+                    _wrap(
+                        tracer,
+                        token_histogram,
+                        choice_counter,
+                        duration_histogram,
+                        event_logger,
+                        wrapped_method,
+                    ),
+                )
+            except ModuleNotFoundError:
+                pass  # that's ok, we don't want to fail if some methods do not exist
+
+        for wrapped_method in WRAPPED_AMETHODS:
+            wrap_package = wrapped_method.get("package")
+            wrap_object = wrapped_method.get("object")
+            wrap_method = wrapped_method.get("method")
+            try:
+                wrap_function_wrapper(
+                    wrap_package,
+                    f"{wrap_object}.{wrap_method}",
+                    _awrap(
+                        tracer,
+                        token_histogram,
+                        choice_counter,
+                        duration_histogram,
+                        event_logger,
+                        wrapped_method,
+                    ),
+                )
+            except ModuleNotFoundError:
+                pass  # that's ok, we don't want to fail if some methods do not exist
+
+    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"),
+            )
+        for wrapped_method in WRAPPED_AMETHODS:
+            wrap_package = wrapped_method.get("package")
+            wrap_object = wrapped_method.get("object")
+            unwrap(
+                f"{wrap_package}.{wrap_object}",
+                wrapped_method.get("method"),
+            )

lmnr/opentelemetry_lib/opentelemetry/instrumentation/groq/config.py

@@ -0,0 +1,8 @@
+from typing import Callable
+
+
+class Config:
+    enrich_token_usage = False
+    exception_logger = None
+    get_common_metrics_attributes: Callable[[], dict] = lambda: {}
+    use_legacy_attributes = True

lmnr/opentelemetry_lib/opentelemetry/instrumentation/groq/event_emitter.py

@@ -0,0 +1,143 @@
+from dataclasses import asdict
+from enum import Enum
+from typing import Union
+
+from opentelemetry._events import Event, EventLogger
+from .event_models import ChoiceEvent, MessageEvent
+from .utils import (
+    dont_throw,
+    should_emit_events,
+    should_send_prompts,
+)
+from opentelemetry.semconv._incubating.attributes import (
+    gen_ai_attributes as GenAIAttributes,
+)
+
+from groq.types.chat.chat_completion import ChatCompletion
+
+
+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 = {
+    # Should be GenAIAttributes.GenAiSystemValues.GROQ.value but it's not defined in the opentelemetry-semconv package
+    GenAIAttributes.GEN_AI_SYSTEM: "groq"
+}
+"""The attributes to be used for the event."""
+
+
+@dont_throw
+def emit_message_events(kwargs: dict, event_logger):
+    for message in kwargs.get("messages", []):
+        emit_event(
+            MessageEvent(
+                content=message.get("content"), role=message.get("role", "unknown")
+            ),
+            event_logger=event_logger,
+        )
+
+
+@dont_throw
+def emit_choice_events(response: ChatCompletion, event_logger):
+    for choice in response.choices:
+        emit_event(
+            ChoiceEvent(
+                index=choice.index,
+                message={
+                    "content": choice.message.content,
+                    "role": choice.message.role or "unknown",
+                },
+                finish_reason=choice.finish_reason,
+            ),
+            event_logger=event_logger,
+        )
+
+
+@dont_throw
+def emit_streaming_response_events(
+    accumulated_content: str, finish_reason: Union[str, None], event_logger
+):
+    """Emit events for streaming response."""
+    emit_event(
+        ChoiceEvent(
+            index=0,
+            message={"content": accumulated_content, "role": "assistant"},
+            finish_reason=finish_reason or "unknown",
+        ),
+        event_logger,
+    )
+
+
+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)
+    )

lmnr/opentelemetry_lib/opentelemetry/instrumentation/groq/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