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

lmnr/opentelemetry_lib/decorators/__init__.py

@@ -6,8 +6,15 @@ import types
 from typing import Any, AsyncGenerator, Callable, Generator, Literal
 
 from opentelemetry import context as context_api
-from opentelemetry.trace import Span
-
+from opentelemetry.trace import Span, Status, StatusCode
+
+from lmnr.opentelemetry_lib.tracing.context import (
+    CONTEXT_SESSION_ID_KEY,
+    CONTEXT_USER_ID_KEY,
+    attach_context,
+    detach_context,
+    get_event_attributes_from_context,
+)
 from lmnr.sdk.utils import get_input_from_func_args, is_method
 from lmnr.opentelemetry_lib import MAX_MANUAL_SPAN_PAYLOAD_SIZE
 from lmnr.opentelemetry_lib.tracing.tracer import get_tracer_with_context
@@ -180,7 +187,21 @@ def observe_base(
 
             span = _setup_span(span_name, span_type, association_properties)
             new_context = wrapper.push_span_context(span)
+            if session_id := association_properties.get("session_id"):
+                new_context = context_api.set_value(
+                    CONTEXT_SESSION_ID_KEY, session_id, new_context
+                )
+            if user_id := association_properties.get("user_id"):
+                new_context = context_api.set_value(
+                    CONTEXT_USER_ID_KEY, user_id, new_context
+                )
+            # Some auto-instrumentations are not under our control, so they
+            # don't have access to our isolated context. We attach the context
+            # to the OTEL global context, so that spans know their parent
+            # span and trace_id.
             ctx_token = context_api.attach(new_context)
+            # update our isolated context too
+            isolated_ctx_token = attach_context(new_context)
 
             _process_input(
                 span, fn, args, kwargs, ignore_input, ignore_inputs, input_formatter
@@ -195,7 +216,7 @@ def observe_base(
             finally:
                 # Always restore global context
                 context_api.detach(ctx_token)
-
+                detach_context(isolated_ctx_token)
             # span will be ended in the generator
             if isinstance(res, types.GeneratorType):
                 return _handle_generator(span, ctx_token, res)
@@ -240,7 +261,21 @@ def async_observe_base(
 
             span = _setup_span(span_name, span_type, association_properties)
             new_context = wrapper.push_span_context(span)
+            if session_id := association_properties.get("session_id"):
+                new_context = context_api.set_value(
+                    CONTEXT_SESSION_ID_KEY, session_id, new_context
+                )
+            if user_id := association_properties.get("user_id"):
+                new_context = context_api.set_value(
+                    CONTEXT_USER_ID_KEY, user_id, new_context
+                )
+            # Some auto-instrumentations are not under our control, so they
+            # don't have access to our isolated context. We attach the context
+            # to the OTEL global context, so that spans know their parent
+            # span and trace_id.
             ctx_token = context_api.attach(new_context)
+            # update our isolated context too
+            isolated_ctx_token = attach_context(new_context)
 
             _process_input(
                 span, fn, args, kwargs, ignore_input, ignore_inputs, input_formatter
@@ -255,6 +290,7 @@ def async_observe_base(
             finally:
                 # Always restore global context
                 context_api.detach(ctx_token)
+                detach_context(isolated_ctx_token)
 
             # span will be ended in the generator
             if isinstance(res, types.AsyncGeneratorType):
@@ -288,4 +324,7 @@ async def _ahandle_generator(span: Span, wrapper: TracerWrapper, res: AsyncGener
 
 def _process_exception(span: Span, e: Exception):
     # Note that this `escaped` is sent as a StringValue("True"), not a boolean.
-    span.record_exception(e, escaped=True)
+    span.record_exception(
+        e, attributes=get_event_attributes_from_context(), escaped=True
+    )
+    span.set_status(Status(StatusCode.ERROR, str(e)))

lmnr/opentelemetry_lib/litellm/__init__.py

@@ -7,6 +7,7 @@ from opentelemetry.trace import SpanKind, Status, StatusCode, Tracer
 from lmnr.opentelemetry_lib.litellm.utils import model_as_dict, set_span_attribute
 from lmnr.opentelemetry_lib.tracing import TracerWrapper
 
+from lmnr.opentelemetry_lib.tracing.context import get_event_attributes_from_context
 from lmnr.opentelemetry_lib.utils.package_check import is_package_installed
 from lmnr.sdk.log import get_default_logger
 
@@ -141,10 +142,12 @@ try:
                 else:
                     span.set_status(Status(StatusCode.ERROR))
                     if isinstance(response_obj, Exception):
-                        span.record_exception(response_obj)
+                        attributes = get_event_attributes_from_context()
+                        span.record_exception(response_obj, attributes=attributes)
 
             except Exception as e:
-                span.record_exception(e)
+                attributes = get_event_attributes_from_context()
+                span.record_exception(e, attributes=attributes)
                 logger.error(f"Error in Laminar LiteLLM instrumentation: {e}")
             finally:
                 span.end(int(end_time.timestamp() * 1e9))

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

@@ -29,6 +29,10 @@ from .utils import (
     shared_metrics_attributes,
     should_emit_events,
 )
+from .streaming import (
+    WrappedAsyncMessageStreamManager,
+    WrappedMessageStreamManager,
+)
 from .version import __version__
 
 from lmnr.opentelemetry_lib.tracing.context import get_current_context
@@ -52,6 +56,7 @@ logger = logging.getLogger(__name__)
 
 _instruments = ("anthropic >= 0.3.11",)
 
+
 WRAPPED_METHODS = [
     {
         "package": "anthropic.resources.completions",
@@ -71,6 +76,15 @@ WRAPPED_METHODS = [
         "method": "stream",
         "span_name": "anthropic.chat",
     },
+    # This method is on an async resource, but is meant to be called as
+    # an async context manager (async with), which we don't need to await;
+    # thus, we wrap it with a sync wrapper
+    {
+        "package": "anthropic.resources.messages",
+        "object": "AsyncMessages",
+        "method": "stream",
+        "span_name": "anthropic.chat",
+    },
 ]
 
 WRAPPED_AMETHODS = [
@@ -86,12 +100,6 @@ WRAPPED_AMETHODS = [
         "method": "create",
         "span_name": "anthropic.chat",
     },
-    {
-        "package": "anthropic.resources.messages",
-        "object": "AsyncMessages",
-        "method": "stream",
-        "span_name": "anthropic.chat",
-    },
 ]
 
 
@@ -99,6 +107,23 @@ def is_streaming_response(response):
     return isinstance(response, Stream) or isinstance(response, AsyncStream)
 
 
+def is_stream_manager(response):
+    """Check if response is a MessageStreamManager or AsyncMessageStreamManager"""
+    try:
+        from anthropic.lib.streaming._messages import (
+            MessageStreamManager,
+            AsyncMessageStreamManager,
+        )
+
+        return isinstance(response, (MessageStreamManager, AsyncMessageStreamManager))
+    except ImportError:
+        # Check by class name as fallback
+        return (
+            response.__class__.__name__ == "MessageStreamManager"
+            or response.__class__.__name__ == "AsyncMessageStreamManager"
+        )
+
+
 @dont_throw
 async def _aset_token_usage(
     span,
@@ -437,6 +462,33 @@ def _wrap(
             event_logger,
             kwargs,
         )
+    elif is_stream_manager(response):
+        if response.__class__.__name__ == "AsyncMessageStreamManager":
+            return WrappedAsyncMessageStreamManager(
+                response,
+                span,
+                instance._client,
+                start_time,
+                token_histogram,
+                choice_counter,
+                duration_histogram,
+                exception_counter,
+                event_logger,
+                kwargs,
+            )
+        else:
+            return WrappedMessageStreamManager(
+                response,
+                span,
+                instance._client,
+                start_time,
+                token_histogram,
+                choice_counter,
+                duration_histogram,
+                exception_counter,
+                event_logger,
+                kwargs,
+            )
     elif response:
         try:
             metric_attributes = shared_metrics_attributes(response)
@@ -532,6 +584,33 @@ async def _awrap(
             event_logger,
             kwargs,
         )
+    elif is_stream_manager(response):
+        if response.__class__.__name__ == "AsyncMessageStreamManager":
+            return WrappedAsyncMessageStreamManager(
+                response,
+                span,
+                instance._client,
+                start_time,
+                token_histogram,
+                choice_counter,
+                duration_histogram,
+                exception_counter,
+                event_logger,
+                kwargs,
+            )
+        else:
+            return WrappedMessageStreamManager(
+                response,
+                span,
+                instance._client,
+                start_time,
+                token_histogram,
+                choice_counter,
+                duration_histogram,
+                exception_counter,
+                event_logger,
+                kwargs,
+            )
     elif response:
         metric_attributes = shared_metrics_attributes(response)
 

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

@@ -113,18 +113,43 @@ async def aset_input_attributes(span, kwargs):
                 )
             for i, message in enumerate(kwargs.get("messages")):
                 prompt_index = i + (1 if has_system_message else 0)
+                content = message.get("content")
+                tool_use_blocks = []
+                other_blocks = []
+                if isinstance(content, list):
+                    for block in content:
+                        if dict(block).get("type") == "tool_use":
+                            tool_use_blocks.append(dict(block))
+                        else:
+                            other_blocks.append(block)
+                    content = other_blocks
                 set_span_attribute(
                     span,
                     f"{SpanAttributes.LLM_PROMPTS}.{prompt_index}.content",
-                    await _dump_content(
-                        message_index=i, span=span, content=message.get("content")
-                    ),
+                    await _dump_content(message_index=i, span=span, content=content),
                 )
                 set_span_attribute(
                     span,
                     f"{SpanAttributes.LLM_PROMPTS}.{prompt_index}.role",
                     message.get("role"),
                 )
+                if tool_use_blocks:
+                    for tool_num, tool_use_block in enumerate(tool_use_blocks):
+                        set_span_attribute(
+                            span,
+                            f"{SpanAttributes.LLM_PROMPTS}.{prompt_index}.tool_calls.{tool_num}.id",
+                            tool_use_block.get("id"),
+                        )
+                        set_span_attribute(
+                            span,
+                            f"{SpanAttributes.LLM_PROMPTS}.{prompt_index}.tool_calls.{tool_num}.name",
+                            tool_use_block.get("name"),
+                        )
+                        set_span_attribute(
+                            span,
+                            f"{SpanAttributes.LLM_PROMPTS}.{prompt_index}.tool_calls.{tool_num}.arguments",
+                            json.dumps(tool_use_block.get("input")),
+                        )
 
         if kwargs.get("tools") is not None:
             for i, tool in enumerate(kwargs.get("tools")):
@@ -160,8 +185,8 @@ def _set_span_completions(span, response):
             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
+            if content_block_type == "text" and hasattr(content, "text"):
+                text += content.text or ""
             elif content_block_type == "thinking":
                 content = dict(content)
                 # override the role to thinking
@@ -242,15 +267,33 @@ def set_streaming_response_attributes(span, complete_response_events):
     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}"
+    index = 0
+    for event in complete_response_events:
+        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"
+        # Thinking is added as a separate completion, so we need to increment the index
+        if event.get("type") == "thinking":
+            index += 1
+        set_span_attribute(span, f"{prefix}.role", role)
+        if event.get("type") == "tool_use":
+            set_span_attribute(
+                span,
+                f"{prefix}.tool_calls.0.id",
+                event.get("id"),
+            )
             set_span_attribute(
-                span, f"{prefix}.finish_reason", event.get("finish_reason")
+                span,
+                f"{prefix}.tool_calls.0.name",
+                event.get("name"),
             )
-            role = "thinking" if event.get("type") == "thinking" else "assistant"
-            set_span_attribute(span, f"{prefix}.role", role)
+            tool_arguments = event.get("input")
+            if tool_arguments is not None:
+                set_span_attribute(
+                    span,
+                    f"{prefix}.tool_calls.0.arguments",
+                    # already stringified
+                    tool_arguments,
+                )
+        else:
             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))

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

@@ -40,15 +40,19 @@ def _process_response_item(item, complete_response):
             complete_response["events"].append(
                 {"index": index, "text": "", "type": item.content_block.type}
             )
-    elif item.type == "content_block_delta" and item.delta.type in [
-        "thinking_delta",
-        "text_delta",
-    ]:
+            if item.content_block.type == "tool_use":
+                complete_response["events"][index]["id"] = item.content_block.id
+                complete_response["events"][index]["name"] = item.content_block.name
+                complete_response["events"][index]["input"] = ""
+
+    elif item.type == "content_block_delta":
         index = item.index
         if item.delta.type == "thinking_delta":
-            complete_response["events"][index]["text"] += item.delta.thinking
+            complete_response["events"][index]["text"] += item.delta.thinking or ""
         elif item.delta.type == "text_delta":
-            complete_response["events"][index]["text"] += item.delta.text
+            complete_response["events"][index]["text"] += item.delta.text or ""
+        elif item.delta.type == "input_json_delta":
+            complete_response["events"][index]["input"] += item.delta.partial_json
     elif item.type == "message_delta":
         for event in complete_response.get("events", []):
             event["finish_reason"] = item.delta.stop_reason
@@ -293,3 +297,99 @@ async def abuild_from_streaming_response(
     if span.is_recording():
         span.set_status(Status(StatusCode.OK))
         span.end()
+
+
+class WrappedMessageStreamManager:
+    """Wrapper for MessageStreamManager that handles instrumentation"""
+
+    def __init__(
+        self,
+        stream_manager,
+        span,
+        instance,
+        start_time,
+        token_histogram,
+        choice_counter,
+        duration_histogram,
+        exception_counter,
+        event_logger,
+        kwargs,
+    ):
+        self._stream_manager = stream_manager
+        self._span = span
+        self._instance = instance
+        self._start_time = start_time
+        self._token_histogram = token_histogram
+        self._choice_counter = choice_counter
+        self._duration_histogram = duration_histogram
+        self._exception_counter = exception_counter
+        self._event_logger = event_logger
+        self._kwargs = kwargs
+
+    def __enter__(self):
+        # Call the original stream manager's __enter__ to get the actual stream
+        stream = self._stream_manager.__enter__()
+        # Return the wrapped stream
+        return build_from_streaming_response(
+            self._span,
+            stream,
+            self._instance,
+            self._start_time,
+            self._token_histogram,
+            self._choice_counter,
+            self._duration_histogram,
+            self._exception_counter,
+            self._event_logger,
+            self._kwargs,
+        )
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        return self._stream_manager.__exit__(exc_type, exc_val, exc_tb)
+
+
+class WrappedAsyncMessageStreamManager:
+    """Wrapper for AsyncMessageStreamManager that handles instrumentation"""
+
+    def __init__(
+        self,
+        stream_manager,
+        span,
+        instance,
+        start_time,
+        token_histogram,
+        choice_counter,
+        duration_histogram,
+        exception_counter,
+        event_logger,
+        kwargs,
+    ):
+        self._stream_manager = stream_manager
+        self._span = span
+        self._instance = instance
+        self._start_time = start_time
+        self._token_histogram = token_histogram
+        self._choice_counter = choice_counter
+        self._duration_histogram = duration_histogram
+        self._exception_counter = exception_counter
+        self._event_logger = event_logger
+        self._kwargs = kwargs
+
+    async def __aenter__(self):
+        # Call the original stream manager's __aenter__ to get the actual stream
+        stream = await self._stream_manager.__aenter__()
+        # Return the wrapped stream
+        return abuild_from_streaming_response(
+            self._span,
+            stream,
+            self._instance,
+            self._start_time,
+            self._token_histogram,
+            self._choice_counter,
+            self._duration_histogram,
+            self._exception_counter,
+            self._event_logger,
+            self._kwargs,
+        )
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        return await self._stream_manager.__aexit__(exc_type, exc_val, exc_tb)

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

@@ -8,7 +8,10 @@ from typing import AsyncGenerator, Callable, Collection, Generator
 
 from google.genai import types
 
-from lmnr.opentelemetry_lib.tracing.context import get_current_context
+from lmnr.opentelemetry_lib.tracing.context import (
+    get_current_context,
+    get_event_attributes_from_context,
+)
 
 from .config import (
     Config,
@@ -491,8 +494,9 @@ def _wrap(tracer: Tracer, to_wrap, wrapped, instance, args, kwargs):
         span.end()
         return response
     except Exception as e:
+        attributes = get_event_attributes_from_context()
         span.set_attribute(ERROR_TYPE, e.__class__.__name__)
-        span.record_exception(e)
+        span.record_exception(e, attributes=attributes)
         span.set_status(Status(StatusCode.ERROR, str(e)))
         span.end()
         raise e
@@ -529,8 +533,9 @@ async def _awrap(tracer: Tracer, to_wrap, wrapped, instance, args, kwargs):
             span.end()
             return response
     except Exception as e:
+        attributes = get_event_attributes_from_context()
         span.set_attribute(ERROR_TYPE, e.__class__.__name__)
-        span.record_exception(e)
+        span.record_exception(e, attributes=attributes)
         span.set_status(Status(StatusCode.ERROR, str(e)))
         span.end()
         raise e

lmnr/opentelemetry_lib/opentelemetry/instrumentation/openai/shared/__init__.py

@@ -395,6 +395,12 @@ def get_token_count_from_string(string: str, model_name: str):
                 f"Failed to get tiktoken encoding for model_name {model_name}, error: {str(ex)}"
             )
             return None
+        except Exception as ex:
+            # Other exceptions in tiktoken
+            logger.warning(
+                f"Failed to get tiktoken encoding for model_name {model_name}, error: {str(ex)}"
+            )
+            return None
 
         tiktoken_encodings[model_name] = encoding
     else: