netra-sdk 0.1.28__tar.gz → 0.1.30__tar.gz

{netra_sdk-0.1.28 → netra_sdk-0.1.30}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: netra-sdk
-Version: 0.1.28
+Version: 0.1.30
 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.28 → netra_sdk-0.1.30}/netra/__init__.py

@@ -1,7 +1,12 @@
+import atexit
 import logging
 import threading
 from typing import Any, Dict, Optional, Set
 
+from opentelemetry import context as context_api
+from opentelemetry import trace
+from opentelemetry.trace import SpanKind
+
 from netra.instrumentation.instruments import InstrumentSet, NetraInstruments
 
 from .config import Config
@@ -27,6 +32,8 @@ class Netra:
     _initialized = False
     # Use RLock so the thread that already owns the lock can re-acquire it safely
     _init_lock = threading.RLock()
+    _root_span = None
+    _root_ctx_token = None
 
     @classmethod
     def is_initialized(cls) -> bool:
@@ -46,6 +53,7 @@ class Netra:
         disable_batch: Optional[bool] = None,
         trace_content: Optional[bool] = None,
         debug_mode: Optional[bool] = None,
+        enable_root_span: Optional[bool] = None,
         resource_attributes: Optional[Dict[str, Any]] = None,
         environment: Optional[str] = None,
         instruments: Optional[Set[NetraInstruments]] = None,
@@ -66,6 +74,7 @@ class Netra:
                 disable_batch=disable_batch,
                 trace_content=trace_content,
                 debug_mode=debug_mode,
+                enable_root_span=enable_root_span,
                 resource_attributes=resource_attributes,
                 environment=environment,
             )
@@ -106,6 +115,61 @@ class Netra:
             cls._initialized = True
             logger.info("Netra successfully initialized.")
 
+            # Create and attach a long-lived root span if enabled
+            if cfg.enable_root_span:
+                tracer = trace.get_tracer("netra.root.span")
+                root_name = f"{Config.LIBRARY_NAME}.root.span"
+                root_span = tracer.start_span(root_name, kind=SpanKind.INTERNAL)
+                # Add useful attributes
+                if cfg.app_name:
+                    root_span.set_attribute("service.name", cfg.app_name)
+                root_span.set_attribute("netra.environment", cfg.environment)
+                root_span.set_attribute("netra.library.version", Config.LIBRARY_VERSION)
+
+                # Attach span to current context so subsequent spans become its children
+                ctx = trace.set_span_in_context(root_span)
+                token = context_api.attach(ctx)
+
+                # Save for potential shutdown/cleanup and session tracking
+                cls._root_span = root_span
+                cls._root_ctx_token = token
+                try:
+                    SessionManager.set_current_span(root_span)
+                except Exception:
+                    pass
+                logger.info("Netra root span created and attached to context.")
+
+                # Ensure cleanup at process exit
+                atexit.register(cls.shutdown)
+
+    @classmethod
+    def shutdown(cls) -> None:
+        """Optional cleanup to end the root span and detach context."""
+        with cls._init_lock:
+            if cls._root_ctx_token is not None:
+                try:
+                    context_api.detach(cls._root_ctx_token)
+                except Exception:
+                    pass
+                finally:
+                    cls._root_ctx_token = None
+            if cls._root_span is not None:
+                try:
+                    cls._root_span.end()
+                except Exception:
+                    pass
+                finally:
+                    cls._root_span = None
+            # Try to flush and shutdown the tracer provider to ensure export
+            try:
+                provider = trace.get_tracer_provider()
+                if hasattr(provider, "force_flush"):
+                    provider.force_flush()
+                if hasattr(provider, "shutdown"):
+                    provider.shutdown()
+            except Exception:
+                pass
+
     @classmethod
     def set_session_id(cls, session_id: str) -> None:
         """
@@ -191,5 +255,30 @@ class Netra:
         """
         return SpanWrapper(name, attributes, module_name)
 
+    @classmethod
+    def set_input(cls, value: Any, span_name: Optional[str] = None) -> None:
+        """
+        Set custom attribute `netra.span.input` on a target span.
+
+        Args:
+            value: Input payload to record (string or JSON-serializable object)
+            span_name: Optional. When provided, sets the attribute on the span registered
+                       with this name. Otherwise sets on the active span.
+        """
+        SessionManager.set_attribute_on_target_span(f"{Config.LIBRARY_NAME}.span.input", value, span_name)
+
+    @classmethod
+    def set_output(cls, value: Any, span_name: Optional[str] = None) -> None:
+        """
+        Set custom attribute `netra.span.output` on a target span.
+
+        Args:
+            value: Output payload to record (string or JSON-serializable object)
+            span_name: Optional. When provided, sets the attribute on the span registered
+                       with this name. Otherwise sets on the active span.
+        """
+        if value:
+            SessionManager.set_attribute_on_target_span(f"{Config.LIBRARY_NAME}.span.output", value, span_name)
+
 
 __all__ = ["Netra", "UsageModel", "ActionModel"]

{netra_sdk-0.1.28 → netra_sdk-0.1.30}/netra/config.py

@@ -17,6 +17,7 @@ class Config:
       - disable_batch:           Whether to disable batch span processor (bool)
       - trace_content:           Whether to capture prompt/completion content (bool)
       - debug_mode:              Whether to enable SDK logging; default False (bool)
+      - enable_root_span:        Whether to create a process root span; default False (bool)
       - resource_attributes:     Custom resource attributes dict (e.g., {'env': 'prod', 'version': '1.0.0'})
     """
 
@@ -24,6 +25,8 @@ class Config:
     SDK_NAME = "netra"
     LIBRARY_NAME = "netra"
     LIBRARY_VERSION = __version__
+    # Maximum length for any attribute value (strings and bytes). Processors should honor this.
+    ATTRIBUTE_MAX_LEN = 1000
 
     def __init__(
         self,
@@ -32,6 +35,7 @@ class Config:
         disable_batch: Optional[bool] = None,
         trace_content: Optional[bool] = None,
         debug_mode: Optional[bool] = None,
+        enable_root_span: Optional[bool] = None,
         resource_attributes: Optional[Dict[str, Any]] = None,
         environment: Optional[str] = None,
     ):
@@ -106,6 +110,13 @@ class Config:
         else:
             self.environment = os.getenv("NETRA_ENV", "local")
 
+        # Enable a long-lived root span for the process? Default False.
+        if enable_root_span is not None:
+            self.enable_root_span = enable_root_span
+        else:
+            env_root = os.getenv("NETRA_ENABLE_ROOT_SPAN")
+            self.enable_root_span = True if (env_root is not None and env_root.lower() in ("1", "true")) else False
+
         # Resource attributes: param override, else parse JSON from env, else empty dict
         if resource_attributes is not None:
             self.resource_attributes = resource_attributes

{netra_sdk-0.1.28 → netra_sdk-0.1.30}/netra/decorators.py

@@ -7,6 +7,7 @@ Decorators can be applied to both functions and classes.
 import functools
 import inspect
 import json
+import logging
 from typing import Any, Awaitable, Callable, Dict, Optional, ParamSpec, Tuple, TypeVar, Union, cast
 
 from opentelemetry import trace
@@ -14,6 +15,8 @@ from opentelemetry import trace
 from .config import Config
 from .session_manager import SessionManager
 
+logger = logging.getLogger(__name__)
+
 P = ParamSpec("P")
 R = TypeVar("R")
 
@@ -84,6 +87,13 @@ def _create_function_wrapper(func: Callable[P, R], entity_type: str, name: Optio
 
             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)
+
                 _add_span_attributes(span, func, args, kwargs, entity_type)
                 try:
                     result = await cast(Awaitable[Any], func(*args, **kwargs))
@@ -93,7 +103,11 @@ def _create_function_wrapper(func: Callable[P, R], entity_type: str, name: Optio
                     span.set_attribute(f"{Config.LIBRARY_NAME}.entity.error", str(e))
                     raise
                 finally:
-                    # Pop entity from stack after function call is done
+                    # Unregister and pop entity from stack after function call is done
+                    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 cast(Callable[P, R], async_wrapper)
@@ -107,6 +121,13 @@ def _create_function_wrapper(func: Callable[P, R], entity_type: str, name: Optio
 
             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)
+
                 _add_span_attributes(span, func, args, kwargs, entity_type)
                 try:
                     result = func(*args, **kwargs)
@@ -116,7 +137,11 @@ def _create_function_wrapper(func: Callable[P, R], entity_type: str, name: Optio
                     span.set_attribute(f"{Config.LIBRARY_NAME}.entity.error", str(e))
                     raise
                 finally:
-                    # Pop entity from stack after function call is done
+                    # Unregister and pop entity from stack after function call is done
+                    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 cast(Callable[P, R], sync_wrapper)

netra_sdk-0.1.30/netra/processors/__init__.py

@@ -0,0 +1,4 @@
+from netra.processors.instrumentation_span_processor import InstrumentationSpanProcessor
+from netra.processors.session_span_processor import SessionSpanProcessor
+
+__all__ = ["SessionSpanProcessor", "InstrumentationSpanProcessor"]

netra_sdk-0.1.30/netra/processors/instrumentation_span_processor.py

@@ -0,0 +1,102 @@
+import logging
+from typing import Any, Callable, Optional
+
+from opentelemetry import context as otel_context
+from opentelemetry import trace
+from opentelemetry.sdk.trace import SpanProcessor
+
+from netra.config import Config
+from netra.instrumentation.instruments import InstrumentSet
+
+logger = logging.getLogger(__name__)
+
+
+ALLOWED_INSTRUMENTATION_NAMES = {member.value for member in InstrumentSet}  # type: ignore[attr-defined]
+
+
+class InstrumentationSpanProcessor(SpanProcessor):  # type: ignore[misc]
+    """Span processor to record this span's instrumentation name and wrap set_attribute.
+
+    - Records raw instrumentation scope name for the span
+    - Wraps span.set_attribute to truncate string values to max 1000 chars (also inside simple lists/dicts of strings)
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+
+    def _detect_raw_instrumentation_name(self, span: trace.Span) -> Optional[str]:
+        """Detect the raw instrumentation name for the span."""
+        scope = getattr(span, "instrumentation_info", None)
+        if scope is not None:
+            name = getattr(scope, "name", None)
+            if isinstance(name, str) and name:
+                # Normalize common pattern like 'opentelemetry.instrumentation.httpx' -> 'httpx'
+                try:
+                    base = name.rsplit(".", 1)[-1].strip()
+                    if base:
+                        return base
+                except Exception:
+                    pass
+                return name
+        return None
+
+    def _truncate_value(self, value: Any) -> Any:
+        """Truncate string values to max 1000 chars (also inside simple lists/dicts of strings)."""
+        try:
+            if isinstance(value, str):
+                return value if len(value) <= Config.ATTRIBUTE_MAX_LEN else value[: Config.ATTRIBUTE_MAX_LEN]
+            if isinstance(value, (bytes, bytearray)):
+                return value[: Config.ATTRIBUTE_MAX_LEN]
+            if isinstance(value, list):
+                # Shallow truncate strings inside lists
+                return [self._truncate_value(v) if isinstance(v, (str, bytes, bytearray)) else v for v in value]
+            if isinstance(value, dict):
+                # Shallow truncate strings inside dicts
+                return {
+                    k: self._truncate_value(v) if isinstance(v, (str, bytes, bytearray)) else v
+                    for k, v in value.items()
+                }
+        except Exception:
+            return value
+        return value
+
+    def on_start(self, span: trace.Span, parent_context: Optional[otel_context.Context] = None) -> None:
+        """Start span and wrap set_attribute."""
+        try:
+            # Wrap set_attribute first so subsequent sets are also processed
+            original_set_attribute: Callable[[str, Any], None] = span.set_attribute
+
+            def wrapped_set_attribute(key: str, value: Any) -> None:
+                try:
+                    # Truncate value(s)
+                    truncated = self._truncate_value(value)
+                    # Forward to original
+                    original_set_attribute(key, truncated)
+                    # Special rule: if model key set, mark span as llm
+                    if key == "gen_ai.request.model":
+                        original_set_attribute(f"{Config.LIBRARY_NAME}.span.type", "llm")
+                except Exception:
+                    # Best-effort; never break span
+                    try:
+                        original_set_attribute(key, value)
+                    except Exception:
+                        pass
+
+            # Monkey patch for this span's lifetime
+            setattr(span, "set_attribute", wrapped_set_attribute)
+
+            # Set this span's instrumentation name
+            name = self._detect_raw_instrumentation_name(span)
+            if name and any(allowed in name for allowed in ALLOWED_INSTRUMENTATION_NAMES):
+                span.set_attribute(f"{Config.LIBRARY_NAME}.instrumentation.name", name)
+        except Exception:
+            pass
+
+    def on_end(self, span: trace.Span) -> None:
+        """End span."""
+
+    def force_flush(self, timeout_millis: int = 30000) -> None:
+        """Force flush span."""
+
+    def shutdown(self) -> None:
+        pass

{netra_sdk-0.1.28 → netra_sdk-0.1.30}/netra/session_manager.py

@@ -28,6 +28,9 @@ class SessionManager:
     _agent_stack: List[str] = []
     _span_stack: List[str] = []
 
+    # Span registry: name -> stack of spans (most-recent last)
+    _spans_by_name: Dict[str, List[trace.Span]] = {}
+
     @classmethod
     def set_current_span(cls, span: Optional[trace.Span]) -> None:
         """
@@ -48,6 +51,49 @@ class SessionManager:
         """
         return cls._current_span
 
+    @classmethod
+    def register_span(cls, name: str, span: trace.Span) -> None:
+        """
+        Register a span under a given name. Supports nested spans with the same name via a stack.
+        """
+        try:
+            stack = cls._spans_by_name.get(name)
+            if stack is None:
+                cls._spans_by_name[name] = [span]
+            else:
+                stack.append(span)
+        except Exception:
+            logger.exception("Failed to register span '%s'", name)
+
+    @classmethod
+    def unregister_span(cls, name: str, span: trace.Span) -> None:
+        """
+        Unregister a span for a given name. Safe if not present.
+        """
+        try:
+            stack = cls._spans_by_name.get(name)
+            if not stack:
+                return
+            # Remove the last matching instance (normal case)
+            for i in range(len(stack) - 1, -1, -1):
+                if stack[i] is span:
+                    stack.pop(i)
+                    break
+            if not stack:
+                cls._spans_by_name.pop(name, None)
+        except Exception:
+            logger.exception("Failed to unregister span '%s'", name)
+
+    @classmethod
+    def get_span_by_name(cls, name: str) -> Optional[trace.Span]:
+        """
+        Get the most recently registered span with the given name.
+        """
+        stack = cls._spans_by_name.get(name)
+        if stack:
+            return stack[-1]
+        return None
+
     @classmethod
     def push_entity(cls, entity_type: str, entity_name: str) -> None:
         """
@@ -190,3 +236,45 @@ class SessionManager:
                     span.add_event(name=name, attributes=attributes, timestamp=timestamp_ns)
         except Exception as e:
             logger.exception(f"Failed to add custom event: {name} - {e}")
+
+    @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.
+
+        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`.
+        """
+        try:
+            # Convert attribute value to a JSON-safe string representation
+            try:
+                if isinstance(attr_value, str):
+                    attr_str = attr_value
+                else:
+                    import json
+
+                    attr_str = json.dumps(attr_value)
+            except Exception:
+                attr_str = str(attr_value)
+
+            # If a target span name is provided, use the registry for explicit lookup
+            if span_name is not None:
+                target = cls.get_span_by_name(span_name)
+                if target is None:
+                    logger.debug("No span found with name '%s' to set attribute %s", span_name, attr_key)
+                    return
+                target.set_attribute(attr_key, attr_str)
+                return
+
+            # Otherwise annotate the active 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()
+            if candidate is None:
+                logger.debug("No active span found to set attribute %s", attr_key)
+                return
+            candidate.set_attribute(attr_key, attr_str)
+        except Exception as e:
+            logger.exception("Failed setting attribute %s: %s", attr_key, e)

{netra_sdk-0.1.28 → netra_sdk-0.1.30}/netra/span_wrapper.py

@@ -11,6 +11,7 @@ from opentelemetry.trace.propagation import set_span_in_context
 from pydantic import BaseModel
 
 from netra.config import Config
+from netra.session_manager import SessionManager
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -84,6 +85,14 @@ class SpanWrapper:
         ctx = set_span_in_context(self.span)
         self.context_token = context_api.attach(ctx)
 
+        # Register with SessionManager for name-based lookup
+        try:
+            SessionManager.register_span(self.name, self.span)
+            # Optionally set as current span for SDK consumers that rely on it
+            SessionManager.set_current_span(self.span)
+        except Exception:
+            logger.exception("Failed to register span '%s' with SessionManager", self.name)
+
         logger.info(f"Started span wrapper: {self.name}")
         return self
 
@@ -120,6 +129,11 @@ class SpanWrapper:
 
         # End OpenTelemetry span and detach context
         if self.span:
+            # Unregister from SessionManager before ending span
+            try:
+                SessionManager.unregister_span(self.name, self.span)
+            except Exception:
+                logger.exception("Failed to unregister span '%s' from SessionManager", self.name)
             self.span.end()
         if self.context_token:
             context_api.detach(self.context_token)

{netra_sdk-0.1.28 → netra_sdk-0.1.30}/netra/tracer.py

@@ -65,9 +65,10 @@ class Tracer:
                 endpoint=self._format_endpoint(self.cfg.otlp_endpoint),
                 headers=self.cfg.headers,
             )
-        # Add span processors for session span processing and data aggregation processing
-        from netra.processors import SessionSpanProcessor
+        # Add span processors: first instrumentation wrapper, then session processor
+        from netra.processors import InstrumentationSpanProcessor, SessionSpanProcessor
 
+        provider.add_span_processor(InstrumentationSpanProcessor())
         provider.add_span_processor(SessionSpanProcessor())
 
         # Install appropriate span processor

netra_sdk-0.1.30/netra/version.py

@@ -0,0 +1 @@
+__version__ = "0.1.30"

{netra_sdk-0.1.28 → netra_sdk-0.1.30}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [project]
 name = "netra-sdk"
-version = "0.1.28"
+version = "0.1.30"
 description = "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."
 authors = [
     {name = "Sooraj Thomas",email = "sooraj@keyvalue.systems"}

netra_sdk-0.1.28/netra/processors/__init__.py

@@ -1,3 +0,0 @@
-from netra.processors.session_span_processor import SessionSpanProcessor
-
-__all__ = ["SessionSpanProcessor"]

netra_sdk-0.1.28/netra/version.py

@@ -1 +0,0 @@
-__version__ = "0.1.28"