netra-sdk 0.1.34__py3-none-any.whl → 0.1.35__py3-none-any.whl

netra/decorators.py

@@ -8,7 +8,26 @@ import functools
 import inspect
 import json
 import logging
-from typing import Any, Awaitable, Callable, Dict, Optional, ParamSpec, Tuple, TypeVar, Union, cast
+from typing import (
+    Any,
+    AsyncGenerator,
+    Awaitable,
+    Callable,
+    Dict,
+    Generator,
+    Optional,
+    ParamSpec,
+    Tuple,
+    TypeVar,
+    Union,
+    cast,
+)
+
+try:
+    # Optional import: only present in FastAPI/Starlette environments
+    from starlette.responses import StreamingResponse
+except Exception:  # pragma: no cover - starlette may not be installed in some environments
+    StreamingResponse = None
 
 from opentelemetry import trace
 
@@ -73,6 +92,176 @@ def _add_output_attributes(span: trace.Span, result: Any) -> None:
         span.set_attribute(f"{Config.LIBRARY_NAME}.entity.output_error", str(e))
 
 
+def _is_streaming_response(obj: Any) -> bool:
+    """Return True if obj is a Starlette StreamingResponse instance."""
+    if StreamingResponse is None:
+        return False
+    try:
+        return isinstance(obj, StreamingResponse)
+    except Exception:
+        return False
+
+
+def _is_async_generator(obj: Any) -> bool:
+    return inspect.isasyncgen(obj)
+
+
+def _is_sync_generator(obj: Any) -> bool:
+    return inspect.isgenerator(obj)
+
+
+def _wrap_async_generator_with_span(
+    span: trace.Span,
+    agen: AsyncGenerator[Any, None],
+    span_name: str,
+    entity_type: str,
+) -> AsyncGenerator[Any, None]:
+    """Wrap an async generator so the span remains current for the full iteration and ends afterwards."""
+
+    async def _wrapped() -> AsyncGenerator[Any, None]:
+        # Activate span for the entire iteration
+        with trace.use_span(span, end_on_exit=False):
+            try:
+                async for item in agen:
+                    yield item
+            except Exception as e:
+                try:
+                    span.set_attribute(f"{Config.LIBRARY_NAME}.entity.error", str(e))
+                    span.record_exception(e)
+                finally:
+                    span.end()
+                    # De-register and pop entity at the very end for streaming lifecycle
+                    try:
+                        SessionManager.unregister_span(span_name, span)
+                    except Exception:
+                        logger.exception("Failed to unregister span '%s' from SessionManager", span_name)
+                    SessionManager.pop_entity(entity_type)
+                raise
+            else:
+                # Normal completion
+                span.end()
+                try:
+                    SessionManager.unregister_span(span_name, span)
+                except Exception:
+                    logger.exception("Failed to unregister span '%s' from SessionManager", span_name)
+                SessionManager.pop_entity(entity_type)
+
+    return _wrapped()
+
+
+def _wrap_sync_generator_with_span(
+    span: trace.Span,
+    gen: Generator[Any, None, None],
+    span_name: str,
+    entity_type: str,
+) -> Generator[Any, None, None]:
+    """Wrap a sync generator so the span remains current for the full iteration and ends afterwards."""
+
+    def _wrapped() -> Generator[Any, None, None]:
+        with trace.use_span(span, end_on_exit=False):
+            try:
+                for item in gen:
+                    yield item
+            except Exception as e:
+                try:
+                    span.set_attribute(f"{Config.LIBRARY_NAME}.entity.error", str(e))
+                    span.record_exception(e)
+                finally:
+                    span.end()
+                    try:
+                        SessionManager.unregister_span(span_name, span)
+                    except Exception:
+                        logger.exception("Failed to unregister span '%s' from SessionManager", span_name)
+                    SessionManager.pop_entity(entity_type)
+                raise
+            else:
+                span.end()
+                try:
+                    SessionManager.unregister_span(span_name, span)
+                except Exception:
+                    logger.exception("Failed to unregister span '%s' from SessionManager", span_name)
+                SessionManager.pop_entity(entity_type)
+
+    return _wrapped()
+
+
+def _wrap_streaming_response_with_span(
+    span: trace.Span,
+    resp: Any,
+    span_name: str,
+    entity_type: str,
+) -> Any:
+    """Wrap StreamingResponse.body_iterator with a generator that keeps span current and ends it afterwards."""
+    try:
+        body_iter = getattr(resp, "body_iterator", None)
+        if body_iter is None:
+            return resp
+        # Async iterator
+        if inspect.isasyncgen(body_iter) or hasattr(body_iter, "__aiter__"):
+
+            async def _aiter_wrapper():  # type: ignore[no-untyped-def]
+                with trace.use_span(span, end_on_exit=False):
+                    try:
+                        async for chunk in body_iter:
+                            yield chunk
+                    except Exception as e:
+                        try:
+                            span.set_attribute(f"{Config.LIBRARY_NAME}.entity.error", str(e))
+                            span.record_exception(e)
+                        finally:
+                            span.end()
+                            try:
+                                SessionManager.unregister_span(span_name, span)
+                            except Exception:
+                                logger.exception("Failed to unregister span '%s' from SessionManager", span_name)
+                            SessionManager.pop_entity(entity_type)
+                        raise
+                    else:
+                        span.end()
+                        try:
+                            SessionManager.unregister_span(span_name, span)
+                        except Exception:
+                            logger.exception("Failed to unregister span '%s' from SessionManager", span_name)
+                        SessionManager.pop_entity(entity_type)
+
+            resp.body_iterator = _aiter_wrapper()  # type: ignore[no-untyped-call]
+            return resp
+
+        # Sync iterator
+        if inspect.isgenerator(body_iter) or hasattr(body_iter, "__iter__"):
+
+            def _iter_wrapper():  # type: ignore[no-untyped-def]
+                with trace.use_span(span, end_on_exit=False):
+                    try:
+                        for chunk in body_iter:
+                            yield chunk
+                    except Exception as e:
+                        try:
+                            span.set_attribute(f"{Config.LIBRARY_NAME}.entity.error", str(e))
+                            span.record_exception(e)
+                        finally:
+                            span.end()
+                            try:
+                                SessionManager.unregister_span(span_name, span)
+                            except Exception:
+                                logger.exception("Failed to unregister span '%s' from SessionManager", span_name)
+                            SessionManager.pop_entity(entity_type)
+                        raise
+                    else:
+                        span.end()
+                        try:
+                            SessionManager.unregister_span(span_name, span)
+                        except Exception:
+                            logger.exception("Failed to unregister span '%s' from SessionManager", span_name)
+                        SessionManager.pop_entity(entity_type)
+
+            resp.body_iterator = _iter_wrapper()  # type: ignore[no-untyped-call]
+            return resp
+    except Exception:
+        logger.exception("Failed to wrap StreamingResponse with span '%s'", span_name)
+    return resp
+
+
 def _create_function_wrapper(func: Callable[P, R], entity_type: str, name: Optional[str] = None) -> Callable[P, R]:
     module_name = func.__name__
     is_async = inspect.iscoroutinefunction(func)
@@ -82,33 +271,50 @@ def _create_function_wrapper(func: Callable[P, R], entity_type: str, name: Optio
 
         @functools.wraps(func)
         async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
-            # Push entity to stack before span starts so SessionSpanProcessor can capture it
+            # Push entity before span starts so processors can capture it
             SessionManager.push_entity(entity_type, span_name)
 
             tracer = trace.get_tracer(module_name)
-            with tracer.start_as_current_span(span_name) as span:
-                # Register the span by name for cross-context attribute setting
-                try:
-                    SessionManager.register_span(span_name, span)
-                    SessionManager.set_current_span(span)
-                except Exception:
-                    logger.exception("Failed to register span '%s' with SessionManager", span_name)
-
+            span = tracer.start_span(span_name)
+            # Register and activate span
+            try:
+                SessionManager.register_span(span_name, span)
+                SessionManager.set_current_span(span)
+            except Exception:
+                logger.exception("Failed to register span '%s' with SessionManager", span_name)
+
+            with trace.use_span(span, end_on_exit=False):
                 _add_span_attributes(span, func, args, kwargs, entity_type)
                 try:
                     result = await cast(Awaitable[Any], func(*args, **kwargs))
-                    _add_output_attributes(span, result)
-                    return result
                 except Exception as e:
                     span.set_attribute(f"{Config.LIBRARY_NAME}.entity.error", str(e))
-                    raise
-                finally:
-                    # Unregister and pop entity from stack after function call is done
+                    span.record_exception(e)
+                    span.end()
                     try:
                         SessionManager.unregister_span(span_name, span)
                     except Exception:
                         logger.exception("Failed to unregister span '%s' from SessionManager", span_name)
                     SessionManager.pop_entity(entity_type)
+                    raise
+
+            # If result is streaming, defer span end to when stream completes
+            if _is_streaming_response(result):
+                return _wrap_streaming_response_with_span(span, result, span_name, entity_type)
+            if _is_async_generator(result):
+                return _wrap_async_generator_with_span(span, result, span_name, entity_type)
+            if _is_sync_generator(result):
+                return _wrap_sync_generator_with_span(span, result, span_name, entity_type)
+
+            # Non-streaming: finalize now
+            _add_output_attributes(span, result)
+            span.end()
+            try:
+                SessionManager.unregister_span(span_name, span)
+            except Exception:
+                logger.exception("Failed to unregister span '%s' from SessionManager", span_name)
+            SessionManager.pop_entity(entity_type)
+            return result
 
         return cast(Callable[P, R], async_wrapper)
 
@@ -116,33 +322,50 @@ def _create_function_wrapper(func: Callable[P, R], entity_type: str, name: Optio
 
         @functools.wraps(func)
         def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
-            # Push entity to stack before span starts so SessionSpanProcessor can capture it
+            # Push entity before span starts so processors can capture it
             SessionManager.push_entity(entity_type, span_name)
 
             tracer = trace.get_tracer(module_name)
-            with tracer.start_as_current_span(span_name) as span:
-                # Register the span by name for cross-context attribute setting
-                try:
-                    SessionManager.register_span(span_name, span)
-                    SessionManager.set_current_span(span)
-                except Exception:
-                    logger.exception("Failed to register span '%s' with SessionManager", span_name)
-
+            span = tracer.start_span(span_name)
+            # Register and activate span
+            try:
+                SessionManager.register_span(span_name, span)
+                SessionManager.set_current_span(span)
+            except Exception:
+                logger.exception("Failed to register span '%s' with SessionManager", span_name)
+
+            with trace.use_span(span, end_on_exit=False):
                 _add_span_attributes(span, func, args, kwargs, entity_type)
                 try:
                     result = func(*args, **kwargs)
-                    _add_output_attributes(span, result)
-                    return result
                 except Exception as e:
                     span.set_attribute(f"{Config.LIBRARY_NAME}.entity.error", str(e))
-                    raise
-                finally:
-                    # Unregister and pop entity from stack after function call is done
+                    span.record_exception(e)
+                    span.end()
                     try:
                         SessionManager.unregister_span(span_name, span)
                     except Exception:
                         logger.exception("Failed to unregister span '%s' from SessionManager", span_name)
                     SessionManager.pop_entity(entity_type)
+                    raise
+
+            # If result is streaming, defer span end to when stream completes
+            if _is_streaming_response(result):
+                return _wrap_streaming_response_with_span(span, result, span_name, entity_type)
+            if _is_async_generator(result):
+                return _wrap_async_generator_with_span(span, result, span_name, entity_type)  # type: ignore[arg-type]
+            if _is_sync_generator(result):
+                return _wrap_sync_generator_with_span(span, result, span_name, entity_type)  # type: ignore[arg-type]
+
+            # Non-streaming: finalize now
+            _add_output_attributes(span, result)
+            span.end()
+            try:
+                SessionManager.unregister_span(span_name, span)
+            except Exception:
+                logger.exception("Failed to unregister span '%s' from SessionManager", span_name)
+            SessionManager.pop_entity(entity_type)
+            return result
 
         return cast(Callable[P, R], sync_wrapper)
 

netra/session_manager.py

@@ -31,6 +31,10 @@ class SessionManager:
     # Span registry: name -> stack of spans (most-recent last)
     _spans_by_name: Dict[str, List[trace.Span]] = {}
 
+    # Global stack of active spans in creation order (oldest first, newest last)
+    # Maintained for spans registered via SessionManager (e.g., SpanWrapper)
+    _active_spans: List[trace.Span] = []
+
     @classmethod
     def set_current_span(cls, span: Optional[trace.Span]) -> None:
         """
@@ -62,6 +66,8 @@ class SessionManager:
                 cls._spans_by_name[name] = [span]
             else:
                 stack.append(span)
+            # Track globally as active
+            cls._active_spans.append(span)
         except Exception:
             logger.exception("Failed to register span '%s'", name)
 
@@ -81,6 +87,11 @@ class SessionManager:
                     break
             if not stack:
                 cls._spans_by_name.pop(name, None)
+            # Also remove from global active list (remove last matching instance)
+            for i in range(len(cls._active_spans) - 1, -1, -1):
+                if cls._active_spans[i] is span:
+                    cls._active_spans.pop(i)
+                    break
         except Exception:
             logger.exception("Failed to unregister span '%s'", name)
 
@@ -240,12 +251,13 @@ class SessionManager:
     @classmethod
     def set_attribute_on_target_span(cls, attr_key: str, attr_value: Any, span_name: Optional[str] = None) -> None:
         """
-        Best-effort setter to annotate the active span with the provided attribute.
+        Best-effort setter to annotate a target span with the provided attribute.
 
-        If span_name is provided, we look up that span via SessionManager's registry and set
-        the attribute on that span explicitly. Otherwise, we annotate the active span.
-        We first try the OpenTelemetry current span; if that's invalid, we fall back to
-        the SDK-managed current span from `SessionManager`.
+        Behavior:
+        - If span_name is provided, set the attribute on the span registered with that name.
+        - If no span_name is provided, attempt to set the attribute on the SDK root span
+          (created when Netra.init(enable_root_span=True)). If the root span is unavailable,
+          fall back to the currently active span (OTel current span or SDK-managed current span).
         """
         try:
             # Convert attribute value to a JSON-safe string representation
@@ -268,10 +280,41 @@ class SessionManager:
                 target.set_attribute(attr_key, attr_str)
                 return
 
-            # Otherwise annotate the active span
+            # Otherwise, attempt to set on the root-most span in the current trace
+            candidate = None
+
+            # Determine current trace_id from the active/current span
             current_span = trace.get_current_span()
             has_valid_current = getattr(current_span, "is_recording", None) is not None and current_span.is_recording()
-            candidate = current_span if has_valid_current else cls.get_current_span()
+            base_span = current_span if has_valid_current else cls.get_current_span()
+            trace_id: Optional[int] = None
+            try:
+                if base_span is not None and hasattr(base_span, "get_span_context"):
+                    sc = base_span.get_span_context()
+                    trace_id = getattr(sc, "trace_id", None)
+            except Exception:
+                trace_id = None
+
+            # Find the earliest active span in this process that belongs to the same trace
+            if trace_id is not None:
+                try:
+                    for s in cls._active_spans:
+                        if s is None:
+                            continue
+                        if not getattr(s, "is_recording", lambda: False)():
+                            continue
+                        sc = getattr(s, "get_span_context", lambda: None)()
+                        if sc is None:
+                            continue
+                        if getattr(sc, "trace_id", None) == trace_id:
+                            candidate = s
+                            break
+                except Exception:
+                    candidate = None
+
+            # Fallback to the current active span if no root-most could be found
+            if candidate is None:
+                candidate = base_span
             if candidate is None:
                 logger.debug("No active span found to set attribute %s", attr_key)
                 return

netra/version.py

@@ -1 +1 @@
-__version__ = "0.1.34"
+__version__ = "0.1.35"

{netra_sdk-0.1.34.dist-info → netra_sdk-0.1.35.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: netra-sdk
-Version: 0.1.34
+Version: 0.1.35
 Summary: A Python SDK for AI application observability that provides OpenTelemetry-based monitoring, tracing, and PII protection for LLM and vector database applications. Enables easy instrumentation, session tracking, and privacy-focused data collection for AI systems in production environments.
 License: Apache-2.0
 Keywords: netra,tracing,observability,sdk,ai,llm,vector,database

{netra_sdk-0.1.34.dist-info → netra_sdk-0.1.35.dist-info}/RECORD

@@ -4,7 +4,7 @@ netra/anonymizer/anonymizer.py,sha256=IcrYkdwWrFauGWUeAW-0RwrSUM8VSZCFNtoywZhvIq
 netra/anonymizer/base.py,sha256=ytPxHCUD2OXlEY6fNTuMmwImNdIjgj294I41FIgoXpU,5946
 netra/anonymizer/fp_anonymizer.py,sha256=_6svIYmE0eejdIMkhKBUWCNjGtGimtrGtbLvPSOp8W4,6493
 netra/config.py,sha256=51m8R0NoOrw58gMV7arOniEuFdJ7EIu3PNdFtIQ5xfg,6893
-netra/decorators.py,sha256=yuQP02sdvTRIYkv-myNcP8q7dmPq3ME1AJZxJtryayI,8720
+netra/decorators.py,sha256=qZFHrwdj10FsTFqggo3XjdGB12aMxsrrDMMmslDqZ-0,17424
 netra/exceptions/__init__.py,sha256=uDgcBxmC4WhdS7HRYQk_TtJyxH1s1o6wZmcsnSHLAcM,174
 netra/exceptions/injection.py,sha256=ke4eUXRYUFJkMZgdSyPPkPt5PdxToTI6xLEBI0hTWUQ,1332
 netra/exceptions/pii.py,sha256=MT4p_x-zH3VtYudTSxw1Z9qQZADJDspq64WrYqSWlZc,2438
@@ -45,11 +45,11 @@ netra/processors/instrumentation_span_processor.py,sha256=Ef5FTr8O5FLHcIkBAW3ueU
 netra/processors/scrubbing_span_processor.py,sha256=dJ86Ncmjvmrhm_uAdGTwcGvRpZbVVWqD9AOFwEMWHZY,6701
 netra/processors/session_span_processor.py,sha256=qcsBl-LnILWefsftI8NQhXDGb94OWPc8LvzhVA0JS_c,2432
 netra/scanner.py,sha256=kyDpeZiscCPb6pjuhS-sfsVj-dviBFRepdUWh0sLoEY,11554
-netra/session_manager.py,sha256=AoQa-k4dFcq7PeOD8G8DNzhLzL1JrHUW6b_y8mRyTQo,10255
+netra/session_manager.py,sha256=pyhtWYwFkrieNfrVFxSqxDsErPZzxPIt94zqvABLTAI,12229
 netra/span_wrapper.py,sha256=IygQX78xQRlL_Z1MfKfUbv0okihx92qNClnRlYFtRNc,8004
 netra/tracer.py,sha256=FJO8Cine-WL9K_4wn6RVjQOgX6c1JCp_8QowUbRSVHk,7718
-netra/version.py,sha256=79r5jd-MqbhXLbIBDVBqUJvhvcucjkaId96r46KF18I,23
-netra_sdk-0.1.34.dist-info/LICENCE,sha256=8B_UoZ-BAl0AqiHAHUETCgd3I2B9yYJ1WEQtVb_qFMA,11359
-netra_sdk-0.1.34.dist-info/METADATA,sha256=Wmf0VwjOEEmgtJcNuhmsFlgDlAT2e9RxXbY73TQl-CU,28210
-netra_sdk-0.1.34.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-netra_sdk-0.1.34.dist-info/RECORD,,
+netra/version.py,sha256=ABqgofsnbWf7823vTBbZNQ81eKQbWwrVToAU6T3z-6s,23
+netra_sdk-0.1.35.dist-info/LICENCE,sha256=8B_UoZ-BAl0AqiHAHUETCgd3I2B9yYJ1WEQtVb_qFMA,11359
+netra_sdk-0.1.35.dist-info/METADATA,sha256=Aewd-yCWdZ32uabTT_io4cnfSmrciVuugAxVgDy22nw,28210
+netra_sdk-0.1.35.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+netra_sdk-0.1.35.dist-info/RECORD,,