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

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

@@ -1,6 +1,7 @@
 import copy
 import json
 import logging
+import threading
 import time
 from functools import singledispatch
 from typing import List, Optional, Union
@@ -39,7 +40,10 @@ from ..utils import (
     should_emit_events,
     should_send_prompts,
 )
-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 opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
 from opentelemetry.metrics import Counter, Histogram
 from opentelemetry.semconv.attributes.error_attributes import ERROR_TYPE
@@ -111,7 +115,8 @@ def chat_wrapper(
             exception_counter.add(1, attributes=attributes)
 
         span.set_attribute(ERROR_TYPE, e.__class__.__name__)
-        span.record_exception(e)
+        attributes = get_event_attributes_from_context()
+        span.record_exception(e, attributes=attributes)
         span.set_status(Status(StatusCode.ERROR, str(e)))
         span.end()
 
@@ -211,7 +216,8 @@ async def achat_wrapper(
             exception_counter.add(1, attributes=attributes)
 
         span.set_attribute(ERROR_TYPE, e.__class__.__name__)
-        span.record_exception(e)
+        attributes = get_event_attributes_from_context()
+        span.record_exception(e, attributes=attributes)
         span.set_status(Status(StatusCode.ERROR, str(e)))
         span.end()
 
@@ -296,6 +302,7 @@ def _handle_response(
     choice_counter=None,
     duration_histogram=None,
     duration=None,
+    is_streaming: bool = False,
 ):
     if is_openai_v1():
         response_dict = model_as_dict(response)
@@ -310,6 +317,7 @@ def _handle_response(
         duration_histogram,
         response_dict,
         duration,
+        is_streaming,
     )
 
     # span attributes
@@ -327,13 +335,19 @@ def _handle_response(
 
 
 def _set_chat_metrics(
-    instance, token_counter, choice_counter, duration_histogram, response_dict, duration
+    instance,
+    token_counter,
+    choice_counter,
+    duration_histogram,
+    response_dict,
+    duration,
+    is_streaming: bool = False,
 ):
     shared_attributes = metric_shared_attributes(
         response_model=response_dict.get("model") or None,
         operation="chat",
         server_address=_get_openai_base_url(instance),
-        is_streaming=False,
+        is_streaming=is_streaming,
     )
 
     # token metrics
@@ -520,11 +534,9 @@ def _set_completions(span, choices):
 def _set_streaming_token_metrics(
     request_kwargs, complete_response, span, token_counter, shared_attributes
 ):
-    # use tiktoken calculate token usage
     if not should_record_stream_token_usage():
         return
 
-    # kwargs={'model': 'gpt-3.5', 'messages': [{'role': 'user', 'content': '...'}], 'stream': True}
     prompt_usage = -1
     completion_usage = -1
 
@@ -621,11 +633,35 @@ class ChatStream(ObjectProxy):
         self._time_of_first_token = self._start_time
         self._complete_response = {"choices": [], "model": ""}
 
+        # Cleanup state tracking to prevent duplicate operations
+        self._cleanup_completed = False
+        self._cleanup_lock = threading.Lock()
+
+    def __del__(self):
+        """Cleanup when object is garbage collected"""
+        if hasattr(self, "_cleanup_completed") and not self._cleanup_completed:
+            self._ensure_cleanup()
+
     def __enter__(self):
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        self.__wrapped__.__exit__(exc_type, exc_val, exc_tb)
+        cleanup_exception = None
+        try:
+            self._ensure_cleanup()
+        except Exception as e:
+            cleanup_exception = e
+            # Don't re-raise to avoid masking original exception
+
+        result = self.__wrapped__.__exit__(exc_type, exc_val, exc_tb)
+
+        if cleanup_exception:
+            # Log cleanup exception but don't affect context manager behavior
+            logger.debug(
+                "Error during ChatStream cleanup in __exit__: %s", cleanup_exception
+            )
+
+        return result
 
     async def __aenter__(self):
         return self
@@ -645,7 +681,12 @@ class ChatStream(ObjectProxy):
         except Exception as e:
             if isinstance(e, StopIteration):
                 self._process_complete_response()
-            raise e
+            else:
+                # Handle cleanup for other exceptions during stream iteration
+                self._ensure_cleanup()
+                if self._span and self._span.is_recording():
+                    self._span.set_status(Status(StatusCode.ERROR, str(e)))
+            raise
         else:
             self._process_item(chunk)
             return chunk
@@ -656,7 +697,12 @@ class ChatStream(ObjectProxy):
         except Exception as e:
             if isinstance(e, StopAsyncIteration):
                 self._process_complete_response()
-            raise e
+            else:
+                # Handle cleanup for other exceptions during stream iteration
+                self._ensure_cleanup()
+                if self._span and self._span.is_recording():
+                    self._span.set_status(Status(StatusCode.ERROR, str(e)))
+            raise
         else:
             self._process_item(chunk)
             return chunk
@@ -727,6 +773,82 @@ class ChatStream(ObjectProxy):
 
         self._span.set_status(Status(StatusCode.OK))
         self._span.end()
+        self._cleanup_completed = True
+
+    @dont_throw
+    def _ensure_cleanup(self):
+        """Thread-safe cleanup method that handles different cleanup scenarios"""
+        with self._cleanup_lock:
+            if self._cleanup_completed:
+                logger.debug("ChatStream cleanup already completed, skipping")
+                return
+
+            try:
+                logger.debug("Starting ChatStream cleanup")
+
+                # Set span status and close it
+                if self._span and self._span.is_recording():
+                    self._span.set_status(Status(StatusCode.OK))
+                    self._span.end()
+                    logger.debug("ChatStream span closed successfully")
+
+                # Calculate partial metrics based on available data
+                self._record_partial_metrics()
+
+                self._cleanup_completed = True
+                logger.debug("ChatStream cleanup completed successfully")
+
+            except Exception as e:
+                # Log cleanup errors but don't propagate to avoid masking original issues
+                logger.debug("Error during ChatStream cleanup: %s", str(e))
+
+                # Still try to close the span even if metrics recording failed
+                try:
+                    if self._span and self._span.is_recording():
+                        self._span.set_status(
+                            Status(StatusCode.ERROR, "Cleanup failed")
+                        )
+                        self._span.end()
+                    self._cleanup_completed = True
+                except Exception:
+                    # Final fallback - just mark as completed to prevent infinite loops
+                    self._cleanup_completed = True
+
+    @dont_throw
+    def _record_partial_metrics(self):
+        """Record metrics based on available partial data"""
+        # Always record duration if we have start time
+        if (
+            self._start_time
+            and isinstance(self._start_time, (float, int))
+            and self._duration_histogram
+        ):
+            duration = time.time() - self._start_time
+            self._duration_histogram.record(
+                duration, attributes=self._shared_attributes()
+            )
+
+        # Record basic span attributes even without complete response
+        if self._span and self._span.is_recording():
+            _set_response_attributes(self._span, self._complete_response)
+
+        # Record partial token metrics if we have any data
+        if self._complete_response.get("choices") or self._request_kwargs:
+            _set_streaming_token_metrics(
+                self._request_kwargs,
+                self._complete_response,
+                self._span,
+                self._token_counter,
+                self._shared_attributes(),
+            )
+
+        # Record choice metrics if we have any choices processed
+        if self._choice_counter and self._complete_response.get("choices"):
+            _set_choice_counter_metrics(
+                self._choice_counter,
+                self._complete_response.get("choices"),
+                self._shared_attributes(),
+            )
 
 
 # Backward compatibility with OpenAI v0
@@ -975,6 +1097,13 @@ def _accumulate_stream_items(item, complete_response):
     complete_response["model"] = item.get("model")
     complete_response["id"] = item.get("id")
 
+    # capture usage information from the last stream chunks
+    if item.get("usage"):
+        complete_response["usage"] = item.get("usage")
+    elif item.get("choices") and item["choices"][0].get("usage"):
+        # Some LLM providers like moonshot mistakenly place token usage information within choices[0], handle this.
+        complete_response["usage"] = item["choices"][0].get("usage")
+
     # prompt filter results
     if item.get("prompt_filter_results"):
         complete_response["prompt_filter_results"] = item.get("prompt_filter_results")

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

@@ -27,7 +27,10 @@ from ..utils import (
     should_emit_events,
     should_send_prompts,
 )
-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 opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
 from opentelemetry.semconv.attributes.error_attributes import ERROR_TYPE
 from opentelemetry.semconv_ai import (
@@ -65,7 +68,8 @@ def completion_wrapper(tracer, wrapped, instance, args, kwargs):
         response = wrapped(*args, **kwargs)
     except Exception as e:
         span.set_attribute(ERROR_TYPE, e.__class__.__name__)
-        span.record_exception(e)
+        attributes = get_event_attributes_from_context()
+        span.record_exception(e, attributes=attributes)
         span.set_status(Status(StatusCode.ERROR, str(e)))
         span.end()
         raise
@@ -100,7 +104,8 @@ async def acompletion_wrapper(tracer, wrapped, instance, args, kwargs):
         response = await wrapped(*args, **kwargs)
     except Exception as e:
         span.set_attribute(ERROR_TYPE, e.__class__.__name__)
-        span.record_exception(e)
+        attributes = get_event_attributes_from_context()
+        span.record_exception(e, attributes=attributes)
         span.set_status(Status(StatusCode.ERROR, str(e)))
         span.end()
         raise

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

@@ -3,6 +3,8 @@ import time
 from collections.abc import Iterable
 
 from opentelemetry import context as context_api
+
+from lmnr.opentelemetry_lib.tracing.context import get_event_attributes_from_context
 from ..shared import (
     OPENAI_LLM_USAGE_TOKEN_TYPES,
     _get_openai_base_url,
@@ -91,7 +93,8 @@ def embeddings_wrapper(
                 exception_counter.add(1, attributes=attributes)
 
             span.set_attribute(ERROR_TYPE, e.__class__.__name__)
-            span.record_exception(e)
+            attributes = get_event_attributes_from_context()
+            span.record_exception(e, attributes=attributes)
             span.set_status(Status(StatusCode.ERROR, str(e)))
             span.end()
 
@@ -156,7 +159,8 @@ async def aembeddings_wrapper(
                 exception_counter.add(1, attributes=attributes)
 
             span.set_attribute(ERROR_TYPE, e.__class__.__name__)
-            span.record_exception(e)
+            attributes = get_event_attributes_from_context()
+            span.record_exception(e, attributes=attributes)
             span.set_status(Status(StatusCode.ERROR, str(e)))
             span.end()
 

lmnr/opentelemetry_lib/opentelemetry/instrumentation/openai/v1/assistant_wrappers.py

@@ -17,7 +17,10 @@ from ..utils import (
     dont_throw,
     should_emit_events,
 )
-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 opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
 from opentelemetry.semconv.attributes.error_attributes import ERROR_TYPE
 from opentelemetry.semconv_ai import LLMRequestTypeValues, SpanAttributes
@@ -132,7 +135,7 @@ def messages_list_wrapper(tracer, wrapped, instance, args, kwargs):
 
     if exception := run.get("exception"):
         span.set_attribute(ERROR_TYPE, exception.__class__.__name__)
-        span.record_exception(exception)
+        span.record_exception(exception, attributes=get_event_attributes_from_context())
         span.set_status(Status(StatusCode.ERROR, str(exception)))
         span.end(run.get("end_time"))
 
@@ -316,7 +319,7 @@ def runs_create_and_stream_wrapper(tracer, wrapped, instance, args, kwargs):
         return response
     except Exception as e:
         span.set_attribute(ERROR_TYPE, e.__class__.__name__)
-        span.record_exception(e)
+        span.record_exception(e, attributes=get_event_attributes_from_context())
         span.set_status(Status(StatusCode.ERROR, str(e)))
         span.end()
         raise

lmnr/opentelemetry_lib/opentelemetry/instrumentation/openai/v1/event_handler_wrapper.py

@@ -1,3 +1,4 @@
+from lmnr.opentelemetry_lib.tracing.context import get_event_attributes_from_context
 from ..shared import _set_span_attribute
 from ..shared.event_emitter import emit_event
 from ..shared.event_models import ChoiceEvent
@@ -69,7 +70,9 @@ class EventHandlerWrapper(AssistantEventHandler):
     @override
     def on_exception(self, exception: Exception):
         self._span.set_attribute(ERROR_TYPE, exception.__class__.__name__)
-        self._span.record_exception(exception)
+        self._span.record_exception(
+            exception, attributes=get_event_attributes_from_context()
+        )
         self._span.set_status(Status(StatusCode.ERROR, str(exception)))
         self._original_handler.on_exception(exception)
 

lmnr/opentelemetry_lib/opentelemetry/instrumentation/openai/v1/responses_wrappers.py

@@ -36,7 +36,10 @@ except ImportError:
     ResponseOutputMessageParam = Dict[str, Any]
     RESPONSES_AVAILABLE = False
 
-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 openai._legacy_response import LegacyAPIResponse
 from opentelemetry import context as context_api
 from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
@@ -433,7 +436,7 @@ def responses_get_or_create_wrapper(tracer: Tracer, wrapped, instance, args, kwa
             context=get_current_context(),
         )
         span.set_attribute(ERROR_TYPE, e.__class__.__name__)
-        span.record_exception(e)
+        span.record_exception(e, attributes=get_event_attributes_from_context())
         span.set_status(StatusCode.ERROR, str(e))
         if traced_data:
             set_data_attributes(traced_data, span)
@@ -529,7 +532,7 @@ async def async_responses_get_or_create_wrapper(
             context=get_current_context(),
         )
         span.set_attribute(ERROR_TYPE, e.__class__.__name__)
-        span.record_exception(e)
+        span.record_exception(e, attributes=get_event_attributes_from_context())
         span.set_status(StatusCode.ERROR, str(e))
         if traced_data:
             set_data_attributes(traced_data, span)
@@ -597,7 +600,10 @@ def responses_cancel_wrapper(tracer: Tracer, wrapped, instance, args, kwargs):
             record_exception=True,
             context=get_current_context(),
         )
-        span.record_exception(Exception("Response cancelled"))
+        span.record_exception(
+            Exception("Response cancelled"),
+            attributes=get_event_attributes_from_context(),
+        )
         set_data_attributes(existing_data, span)
         span.end()
     return response
@@ -624,7 +630,10 @@ async def async_responses_cancel_wrapper(
             record_exception=True,
             context=get_current_context(),
         )
-        span.record_exception(Exception("Response cancelled"))
+        span.record_exception(
+            Exception("Response cancelled"),
+            attributes=get_event_attributes_from_context(),
+        )
         set_data_attributes(existing_data, span)
         span.end()
     return response

lmnr/opentelemetry_lib/tracing/context.py

@@ -2,7 +2,9 @@ import threading
 
 from abc import ABC, abstractmethod
 from contextvars import ContextVar
-from opentelemetry.context import Context, Token
+from opentelemetry.context import Context, Token, create_key, get_value
+
+from lmnr.opentelemetry_lib.tracing.attributes import SESSION_ID, USER_ID
 
 
 class _IsolatedRuntimeContext(ABC):
@@ -107,3 +109,18 @@ def attach_context(context: Context) -> Token[Context]:
 def detach_context(token: Token[Context]) -> None:
     """Detach a context from the isolated runtime context."""
     _ISOLATED_RUNTIME_CONTEXT.detach(token)
+
+
+CONTEXT_USER_ID_KEY = create_key(f"lmnr.{USER_ID}")
+CONTEXT_SESSION_ID_KEY = create_key(f"lmnr.{SESSION_ID}")
+
+
+def get_event_attributes_from_context(context: Context | None = None) -> dict[str, str]:
+    """Get the event attributes from the context."""
+    context = context or get_current_context()
+    attributes = {}
+    if session_id := get_value(CONTEXT_SESSION_ID_KEY, context):
+        attributes["lmnr.event.session_id"] = session_id
+    if user_id := get_value(CONTEXT_USER_ID_KEY, context):
+        attributes["lmnr.event.user_id"] = user_id
+    return attributes

lmnr/sdk/browser/pw_utils.py

@@ -33,7 +33,7 @@ except ImportError as e:
 logger = logging.getLogger(__name__)
 
 current_dir = os.path.dirname(os.path.abspath(__file__))
-with open(os.path.join(current_dir, "rrweb", "rrweb.umd.min.cjs"), "r") as f:
+with open(os.path.join(current_dir, "recorder", "record.umd.min.cjs"), "r") as f:
     RRWEB_CONTENT = f"() => {{ {f.read()} }}"
 
 INJECT_PLACEHOLDER = """
@@ -358,14 +358,6 @@ INJECT_PLACEHOLDER = """
 
     setInterval(sendBatchIfReady, BATCH_TIMEOUT);
 
-    // Add heartbeat events
-    setInterval(() => {
-        window.lmnrRrweb.record.addCustomEvent('heartbeat', {
-            title: document.title,
-            url: document.URL,
-        })
-    }, HEARTBEAT_INTERVAL);
-
     async function bufferToBase64(buffer) {
         const base64url = await new Promise(r => {
             const reader = new FileReader()
@@ -397,63 +389,22 @@ INJECT_PLACEHOLDER = """
         collectFonts: true,
         recordCrossOriginIframes: true
     });
-}
-"""
-
-
-async def send_events_async(
-    page: Page, session_id: str, trace_id: str, client: AsyncLaminarClient
-):
-    """Fetch events from the page and send them to the server"""
-    try:
-        # Check if function exists first
-        events = await page.evaluate(
-            """
-        () => {
-            if (typeof window.lmnrGetAndClearEvents !== 'function') {
-                return [];
-            }
-            return window.lmnrGetAndClearEvents();
-        }
-        """
-        )
-
-        if not events or len(events) == 0:
-            return
-
-        await client._browser_events.send(session_id, trace_id, events)
-    except Exception as e:
-        if "Page.evaluate: Target page, context or browser has been closed" not in str(
-            e
-        ):
-            logger.debug(f"Could not send events: {e}")
 
+    function heartbeat() {
+        // Add heartbeat events
+        setInterval(() => {
+            window.lmnrRrweb.record.addCustomEvent('heartbeat', {
+                title: document.title,
+                    url: document.URL,
+                })
+            }, HEARTBEAT_INTERVAL
+        );
+    }
 
-def send_events_sync(
-    page: SyncPage, session_id: str, trace_id: str, client: LaminarClient
-):
-    """Synchronous version of send_events"""
-    try:
-        events = page.evaluate(
-            """
-        () => {
-            if (typeof window.lmnrGetAndClearEvents !== 'function') {
-                return [];
-            }
-            return window.lmnrGetAndClearEvents();
-        }
-        """
-        )
-        if not events or len(events) == 0:
-            return
-
-        client._browser_events.send(session_id, trace_id, events)
+    heartbeat();
 
-    except Exception as e:
-        if "Page.evaluate: Target page, context or browser has been closed" not in str(
-            e
-        ):
-            logger.debug(f"Could not send events: {e}")
+}
+"""
 
 
 def inject_session_recorder_sync(page: SyncPage):
@@ -483,10 +434,10 @@ def inject_session_recorder_sync(page: SyncPage):
             ):
                 return
 
-        try:
-            page.evaluate(INJECT_PLACEHOLDER)
-        except Exception as e:
-            logger.debug(f"Failed to inject session recorder: {e}")
+            try:
+                page.evaluate(INJECT_PLACEHOLDER)
+            except Exception as e:
+                logger.debug(f"Failed to inject session recorder: {e}")
 
     except Exception as e:
         logger.error(f"Error during session recorder injection: {e}")
@@ -519,10 +470,10 @@ async def inject_session_recorder_async(page: Page):
             ):
                 return
 
-        try:
-            await page.evaluate(INJECT_PLACEHOLDER)
-        except Exception as e:
-            logger.debug(f"Failed to inject session recorder placeholder: {e}")
+            try:
+                await page.evaluate(INJECT_PLACEHOLDER)
+            except Exception as e:
+                logger.debug(f"Failed to inject session recorder placeholder: {e}")
 
     except Exception as e:
         logger.error(f"Error during session recorder injection: {e}")
@@ -536,30 +487,6 @@ def start_recording_events_sync(page: SyncPage, session_id: str, client: Laminar
     trace_id = format(span.get_span_context().trace_id, "032x")
     span.set_attribute("lmnr.internal.has_browser_session", True)
 
-    try:
-        if page.evaluate("""() => typeof window.lmnrSendEvents !== 'undefined'"""):
-            return
-    except Exception:
-        pass
-
-    def on_load():
-        try:
-            inject_session_recorder_sync(page)
-        except Exception as e:
-            logger.error(f"Error in on_load handler: {e}")
-
-    def on_close():
-        try:
-            send_events_sync(page, session_id, trace_id, client)
-        except Exception:
-            pass
-
-    page.on("load", on_load)
-    page.on("close", on_close)
-
-    inject_session_recorder_sync(page)
-
-    # Expose function to browser so it can call us when events are ready
     def send_events_from_browser(events):
         try:
             if events and len(events) > 0:
@@ -572,6 +499,16 @@ def start_recording_events_sync(page: SyncPage, session_id: str, client: Laminar
     except Exception as e:
         logger.debug(f"Could not expose function: {e}")
 
+    inject_session_recorder_sync(page)
+
+    def on_load(p):
+        try:
+            inject_session_recorder_sync(p)
+        except Exception as e:
+            logger.error(f"Error in on_load handler: {e}")
+
+    page.on("domcontentloaded", on_load)
+
 
 @observe(name="playwright.page", ignore_input=True, ignore_output=True)
 async def start_recording_events_async(
@@ -581,33 +518,7 @@ async def start_recording_events_async(
     span = trace.get_current_span(ctx)
     trace_id = format(span.get_span_context().trace_id, "032x")
     span.set_attribute("lmnr.internal.has_browser_session", True)
-
-    try:
-        if await page.evaluate(
-            """() => typeof window.lmnrSendEvents !== 'undefined'"""
-        ):
-            return
-    except Exception:
-        pass
-
-    async def on_load(p):
-        try:
-            await inject_session_recorder_async(p)
-        except Exception as e:
-            logger.error(f"Error in on_load handler: {e}")
-
-    async def on_close(p):
-        try:
-            # Send any remaining events before closing
-            await send_events_async(p, session_id, trace_id, client)
-        except Exception:
-            pass
-
-    page.on("load", on_load)
-    page.on("close", on_close)
-
-    await inject_session_recorder_async(page)
-
+    
     async def send_events_from_browser(events):
         try:
             if events and len(events) > 0:
@@ -620,6 +531,16 @@ async def start_recording_events_async(
     except Exception as e:
         logger.debug(f"Could not expose function: {e}")
 
+    await inject_session_recorder_async(page)
+    
+    async def on_load(p):
+        try:
+            await inject_session_recorder_async(p)
+        except Exception as e:
+            logger.error(f"Error in on_load handler: {e}")
+
+    page.on("domcontentloaded", on_load)
+
 
 def take_full_snapshot(page: Page):
     return page.evaluate(