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

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

@@ -0,0 +1,211 @@
+from dataclasses import asdict
+from enum import Enum
+from typing import Optional, Union
+
+from opentelemetry._events import Event, EventLogger
+from .event_models import (
+    ChoiceEvent,
+    MessageEvent,
+    ToolCall,
+)
+from .utils import (
+    should_emit_events,
+    should_send_prompts,
+)
+from opentelemetry.semconv._incubating.attributes import (
+    gen_ai_attributes as GenAIAttributes,
+)
+
+
+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.ANTHROPIC.value
+}
+"""The attributes to be used for the event."""
+
+
+def emit_input_events(event_logger: Optional[EventLogger], kwargs):
+    if kwargs.get("prompt") is not None:
+        emit_event(
+            MessageEvent(content=kwargs.get("prompt"), role="user"), event_logger
+        )
+
+    elif kwargs.get("messages") is not None:
+        if kwargs.get("system"):
+            emit_event(
+                MessageEvent(content=kwargs.get("system"), role="system"), event_logger
+            )
+        for message in kwargs.get("messages"):
+            emit_event(
+                MessageEvent(content=message.get("content"), role=message.get("role")),
+                event_logger,
+            )
+    if kwargs.get("tools") is not None:
+        emit_event(
+            MessageEvent(content={"tools": kwargs.get("tools")}, role="user"),
+            event_logger,
+        )
+
+
+def emit_response_events(event_logger: Optional[EventLogger], response):
+    if not isinstance(response, dict):
+        response = dict(response)
+
+    if response.get("completion"):
+        emit_event(
+            ChoiceEvent(
+                index=0,
+                message={
+                    "content": response.get("completion"),
+                    "role": response.get("role", "assistant"),
+                },
+                finish_reason=response.get("stop_reason"),
+            ),
+            event_logger,
+        )
+    elif response.get("content"):
+        for i, completion in enumerate(response.get("content")):
+            # Parse message
+            if completion.type == "text":
+                message = {
+                    "content": completion.text,
+                    "role": response.get("role", "assistant"),
+                }
+            elif completion.type == "thinking":
+                message = {
+                    "content": completion.thinking,
+                    "role": response.get("role", "assistant"),
+                }
+            elif completion.type == "tool_use":
+                message = {
+                    "content": None,
+                    "role": response.get("role", "assistant"),
+                }
+            else:
+                message = {
+                    "content": None,
+                    "role": response.get("role", "assistant"),
+                }
+
+            # Parse tool calls
+            if completion.type == "tool_use":
+                tool_calls = [
+                    ToolCall(
+                        id=completion.id,
+                        function={
+                            "name": completion.name,
+                            "arguments": completion.input,
+                        },
+                        type="function",
+                    )
+                ]
+            else:
+                tool_calls = None
+
+            # Emit the event
+            emit_event(
+                ChoiceEvent(
+                    index=i,
+                    message=message,
+                    finish_reason=response.get("stop_reason"),
+                    tool_calls=tool_calls,
+                ),
+                event_logger,
+            )
+
+
+def emit_streaming_response_events(
+    event_logger: Optional[EventLogger], complete_response: dict
+):
+    for message in complete_response.get("events", []):
+        emit_event(
+            ChoiceEvent(
+                index=message.get("index", 0),
+                message={
+                    "content": {
+                        "type": message.get("type"),
+                        "content": message.get("text"),
+                    },
+                    "role": message.get("role", "assistant"),
+                },
+                finish_reason=message.get("finish_reason", "unknown"),
+            ),
+            event_logger,
+        )
+
+
+def emit_event(
+    event: Union[MessageEvent, ChoiceEvent], event_logger: EventLogger
+) -> None:
+    """
+    Emit an event to the OpenTelemetry SDK.
+
+    Args:
+        event: The event to emit.
+    """
+    if not should_emit_events():
+        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/anthropic/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

lmnr/opentelemetry_lib/opentelemetry/instrumentation/anthropic/span_utils.py

@@ -0,0 +1,256 @@
+import json
+import logging
+from typing import Any, Dict
+
+from .config import Config
+from .utils import (
+    JSONEncoder,
+    dont_throw,
+    model_as_dict,
+    should_send_prompts,
+)
+from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import (
+    GEN_AI_RESPONSE_ID,
+)
+from opentelemetry.semconv_ai import SpanAttributes
+
+logger = logging.getLogger(__name__)
+
+
+def _is_base64_image(item: Dict[str, Any]) -> bool:
+    if not isinstance(item, dict):
+        return False
+
+    if not isinstance(item.get("source"), dict):
+        return False
+
+    if item.get("type") != "image" or item["source"].get("type") != "base64":
+        return False
+
+    return True
+
+
+async def _process_image_item(item, trace_id, span_id, message_index, content_index):
+    if not Config.upload_base64_image:
+        return item
+
+    image_format = item.get("source").get("media_type").split("/")[1]
+    image_name = f"message_{message_index}_content_{content_index}.{image_format}"
+    base64_string = item.get("source").get("data")
+    url = await Config.upload_base64_image(trace_id, span_id, image_name, base64_string)
+
+    return {"type": "image_url", "image_url": {"url": url}}
+
+
+async def _dump_content(message_index, content, span):
+    if isinstance(content, str):
+        return content
+    elif isinstance(content, list):
+        # If the content is a list of text blocks, concatenate them.
+        # This is more commonly used in prompt caching.
+        if all([model_as_dict(item).get("type") == "text" for item in content]):
+            return "".join([model_as_dict(item).get("text") for item in content])
+
+        content = [
+            (
+                await _process_image_item(
+                    model_as_dict(item),
+                    span.context.trace_id,
+                    span.context.span_id,
+                    message_index,
+                    j,
+                )
+                if _is_base64_image(model_as_dict(item))
+                else model_as_dict(item)
+            )
+            for j, item in enumerate(content)
+        ]
+
+        return json.dumps(content, cls=JSONEncoder)
+
+
+@dont_throw
+async def aset_input_attributes(span, kwargs):
+    from .utils import set_span_attribute
+
+    set_span_attribute(span, SpanAttributes.LLM_REQUEST_MODEL, kwargs.get("model"))
+    set_span_attribute(
+        span, SpanAttributes.LLM_REQUEST_MAX_TOKENS, kwargs.get("max_tokens_to_sample")
+    )
+    set_span_attribute(
+        span, SpanAttributes.LLM_REQUEST_TEMPERATURE, kwargs.get("temperature")
+    )
+    set_span_attribute(span, SpanAttributes.LLM_REQUEST_TOP_P, kwargs.get("top_p"))
+    set_span_attribute(
+        span, SpanAttributes.LLM_FREQUENCY_PENALTY, kwargs.get("frequency_penalty")
+    )
+    set_span_attribute(
+        span, SpanAttributes.LLM_PRESENCE_PENALTY, kwargs.get("presence_penalty")
+    )
+    set_span_attribute(span, SpanAttributes.LLM_IS_STREAMING, kwargs.get("stream"))
+
+    if should_send_prompts():
+        if kwargs.get("prompt") is not None:
+            set_span_attribute(
+                span, f"{SpanAttributes.LLM_PROMPTS}.0.user", kwargs.get("prompt")
+            )
+
+        elif kwargs.get("messages") is not None:
+            has_system_message = False
+            if kwargs.get("system"):
+                has_system_message = True
+                set_span_attribute(
+                    span,
+                    f"{SpanAttributes.LLM_PROMPTS}.0.content",
+                    await _dump_content(
+                        message_index=0, span=span, content=kwargs.get("system")
+                    ),
+                )
+                set_span_attribute(
+                    span,
+                    f"{SpanAttributes.LLM_PROMPTS}.0.role",
+                    "system",
+                )
+            for i, message in enumerate(kwargs.get("messages")):
+                prompt_index = i + (1 if has_system_message else 0)
+                set_span_attribute(
+                    span,
+                    f"{SpanAttributes.LLM_PROMPTS}.{prompt_index}.content",
+                    await _dump_content(
+                        message_index=i, span=span, content=message.get("content")
+                    ),
+                )
+                set_span_attribute(
+                    span,
+                    f"{SpanAttributes.LLM_PROMPTS}.{prompt_index}.role",
+                    message.get("role"),
+                )
+
+        if kwargs.get("tools") is not None:
+            for i, tool in enumerate(kwargs.get("tools")):
+                prefix = f"{SpanAttributes.LLM_REQUEST_FUNCTIONS}.{i}"
+                set_span_attribute(span, f"{prefix}.name", tool.get("name"))
+                set_span_attribute(
+                    span, f"{prefix}.description", tool.get("description")
+                )
+                input_schema = tool.get("input_schema")
+                if input_schema is not None:
+                    set_span_attribute(
+                        span, f"{prefix}.input_schema", json.dumps(input_schema)
+                    )
+
+
+def _set_span_completions(span, response):
+    if not should_send_prompts():
+        return
+    from .utils import set_span_attribute
+
+    index = 0
+    prefix = f"{SpanAttributes.LLM_COMPLETIONS}.{index}"
+    set_span_attribute(span, f"{prefix}.finish_reason", response.get("stop_reason"))
+    if response.get("role"):
+        set_span_attribute(span, f"{prefix}.role", response.get("role"))
+
+    if response.get("completion"):
+        set_span_attribute(span, f"{prefix}.content", response.get("completion"))
+    elif response.get("content"):
+        tool_call_index = 0
+        text = ""
+        for content in response.get("content"):
+            content_block_type = content.type
+            # usually, Antrhopic responds with just one text block,
+            # but the API allows for multiple text blocks, so concatenate them
+            if content_block_type == "text":
+                text += content.text
+            elif content_block_type == "thinking":
+                content = dict(content)
+                # override the role to thinking
+                set_span_attribute(
+                    span,
+                    f"{prefix}.role",
+                    "thinking",
+                )
+                set_span_attribute(
+                    span,
+                    f"{prefix}.content",
+                    content.get("thinking"),
+                )
+                # increment the index for subsequent content blocks
+                index += 1
+                prefix = f"{SpanAttributes.LLM_COMPLETIONS}.{index}"
+                # set the role to the original role on the next completions
+                set_span_attribute(
+                    span,
+                    f"{prefix}.role",
+                    response.get("role"),
+                )
+            elif content_block_type == "tool_use":
+                content = dict(content)
+                set_span_attribute(
+                    span,
+                    f"{prefix}.tool_calls.{tool_call_index}.id",
+                    content.get("id"),
+                )
+                set_span_attribute(
+                    span,
+                    f"{prefix}.tool_calls.{tool_call_index}.name",
+                    content.get("name"),
+                )
+                tool_arguments = content.get("input")
+                if tool_arguments is not None:
+                    set_span_attribute(
+                        span,
+                        f"{prefix}.tool_calls.{tool_call_index}.arguments",
+                        json.dumps(tool_arguments),
+                    )
+                tool_call_index += 1
+        set_span_attribute(span, f"{prefix}.content", text)
+
+
+@dont_throw
+def set_response_attributes(span, response):
+    from .utils import set_span_attribute
+
+    if not isinstance(response, dict):
+        response = response.__dict__
+    set_span_attribute(span, SpanAttributes.LLM_RESPONSE_MODEL, response.get("model"))
+    set_span_attribute(span, GEN_AI_RESPONSE_ID, response.get("id"))
+
+    if response.get("usage"):
+        prompt_tokens = response.get("usage").input_tokens
+        completion_tokens = response.get("usage").output_tokens
+        set_span_attribute(span, SpanAttributes.LLM_USAGE_PROMPT_TOKENS, prompt_tokens)
+        set_span_attribute(
+            span, SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, completion_tokens
+        )
+        set_span_attribute(
+            span,
+            SpanAttributes.LLM_USAGE_TOTAL_TOKENS,
+            prompt_tokens + completion_tokens,
+        )
+
+    _set_span_completions(span, response)
+
+
+@dont_throw
+def set_streaming_response_attributes(span, complete_response_events):
+    if not should_send_prompts():
+        return
+
+    from .utils import set_span_attribute
+
+    if not span.is_recording() or not complete_response_events:
+        return
+
+    try:
+        for event in complete_response_events:
+            index = event.get("index")
+            prefix = f"{SpanAttributes.LLM_COMPLETIONS}.{index}"
+            set_span_attribute(
+                span, f"{prefix}.finish_reason", event.get("finish_reason")
+            )
+            role = "thinking" if event.get("type") == "thinking" else "assistant"
+            set_span_attribute(span, f"{prefix}.role", role)
+            set_span_attribute(span, f"{prefix}.content", event.get("text"))
+    except Exception as e:
+        logger.warning("Failed to set completion attributes, error: %s", str(e))