netra-sdk 0.1.29__tar.gz → 0.1.31__tar.gz

{netra_sdk-0.1.29 → netra_sdk-0.1.31}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: netra-sdk
-Version: 0.1.29
+Version: 0.1.31
 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
@@ -303,6 +303,7 @@ async def async_span(data):
 - **CrewAI** - Multi-agent AI systems
 - **Pydantic AI** - AI model communication standard
 - **MCP (Model Context Protocol)** - AI model communication standard
+- **LiteLLM** - LLM provider agnostic client
 
 ## 🛡️ Privacy Protection & Security
 

{netra_sdk-0.1.29 → netra_sdk-0.1.31}/README.md

@@ -223,6 +223,7 @@ async def async_span(data):
 - **CrewAI** - Multi-agent AI systems
 - **Pydantic AI** - AI model communication standard
 - **MCP (Model Context Protocol)** - AI model communication standard
+- **LiteLLM** - LLM provider agnostic client
 
 ## 🛡️ Privacy Protection & Security
 

{netra_sdk-0.1.29 → netra_sdk-0.1.31}/netra/__init__.py

@@ -255,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.29 → netra_sdk-0.1.31}/netra/config.py

@@ -25,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,

{netra_sdk-0.1.29 → netra_sdk-0.1.31}/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.29 → netra_sdk-0.1.31}/netra/instrumentation/__init__.py

@@ -93,6 +93,10 @@ def init_instrumentations(
     if CustomInstruments.MISTRALAI in netra_custom_instruments:
         init_mistral_instrumentor()
 
+    # Initialize LiteLLM instrumentation.
+    if CustomInstruments.LITELLM in netra_custom_instruments:
+        init_litellm_instrumentation()
+
     # Initialize OpenAI instrumentation.
     if CustomInstruments.OPENAI in netra_custom_instruments:
         init_openai_instrumentation()
@@ -435,6 +439,26 @@ def init_mistral_instrumentor() -> bool:
         return False
 
 
+def init_litellm_instrumentation() -> bool:
+    """Initialize LiteLLM instrumentation.
+
+    Returns:
+        bool: True if initialization was successful, False otherwise.
+    """
+    try:
+        if is_package_installed("litellm"):
+            from netra.instrumentation.litellm import LiteLLMInstrumentor
+
+            instrumentor = LiteLLMInstrumentor()
+            if not instrumentor.is_instrumented_by_opentelemetry:
+                instrumentor.instrument()
+        return True
+    except Exception as e:
+        logging.error(f"Error initializing LiteLLM instrumentor: {e}")
+        Telemetry().log_exception(e)
+        return False
+
+
 def init_openai_instrumentation() -> bool:
     """Initialize OpenAI instrumentation.
 

{netra_sdk-0.1.29 → netra_sdk-0.1.31}/netra/instrumentation/instruments.py

@@ -8,6 +8,7 @@ class CustomInstruments(Enum):
     AIOHTTP = "aiohttp"
     COHEREAI = "cohere_ai"
     HTTPX = "httpx"
+    LITELLM = "litellm"
     MISTRALAI = "mistral_ai"
     OPENAI = "openai"
     PYDANTIC_AI = "pydantic_ai"
@@ -127,6 +128,7 @@ class InstrumentSet(Enum):
     KAFKA_PYTHON = "kafka_python"
     LANCEDB = "lancedb"
     LANGCHAIN = "langchain"
+    LITELLM = "litellm"
     LLAMA_INDEX = "llama_index"
     LOGGING = "logging"
     MARQO = "marqo"

netra_sdk-0.1.31/netra/instrumentation/litellm/__init__.py

@@ -0,0 +1,161 @@
+import logging
+import time
+from typing import Any, Collection, Dict, Optional
+
+from opentelemetry import context as context_api
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY, unwrap
+from opentelemetry.trace import SpanKind, Tracer, get_tracer
+from opentelemetry.trace.status import Status, StatusCode
+from wrapt import wrap_function_wrapper
+
+from netra.instrumentation.litellm.version import __version__
+from netra.instrumentation.litellm.wrappers import (
+    acompletion_wrapper,
+    aembedding_wrapper,
+    aimage_generation_wrapper,
+    completion_wrapper,
+    embedding_wrapper,
+    image_generation_wrapper,
+)
+
+logger = logging.getLogger(__name__)
+
+_instruments = ("litellm >= 1.0.0",)
+
+
+class LiteLLMInstrumentor(BaseInstrumentor):  # type: ignore[misc]
+    """
+    Custom LiteLLM instrumentor for Netra SDK with enhanced support for:
+    - completion() and acompletion() methods
+    - embedding() and aembedding() methods
+    - image_generation() and aimage_generation() methods
+    - Proper streaming/non-streaming span handling
+    - Integration with Netra tracing
+    """
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return _instruments
+
+    def _instrument(self, **kwargs):  # type: ignore[no-untyped-def]
+        """Instrument LiteLLM methods"""
+        tracer_provider = kwargs.get("tracer_provider")
+        tracer = get_tracer(__name__, __version__, tracer_provider)
+
+        logger.debug("Starting LiteLLM instrumentation...")
+
+        # Force import litellm to ensure it's available for wrapping
+        try:
+            import litellm
+        except ImportError as e:
+            logger.error(f"Failed to import litellm: {e}")
+            return
+
+        # Store original functions for uninstrumentation
+        self._original_completion = getattr(litellm, "completion", None)
+        self._original_acompletion = getattr(litellm, "acompletion", None)
+        self._original_embedding = getattr(litellm, "embedding", None)
+        self._original_aembedding = getattr(litellm, "aembedding", None)
+        self._original_image_generation = getattr(litellm, "image_generation", None)
+        self._original_aimage_generation = getattr(litellm, "aimage_generation", None)
+
+        # Chat completions - use direct monkey patching with proper function wrapping
+        if self._original_completion:
+            try:
+
+                def instrumented_completion(*args, **kwargs):  # type: ignore[no-untyped-def]
+                    wrapper = completion_wrapper(tracer)
+                    return wrapper(self._original_completion, None, args, kwargs)
+
+                litellm.completion = instrumented_completion
+            except Exception as e:
+                logger.error(f"Failed to monkey-patch litellm.completion: {e}")
+
+        if self._original_acompletion:
+            try:
+
+                async def instrumented_acompletion(*args, **kwargs):  # type: ignore[no-untyped-def]
+                    wrapper = acompletion_wrapper(tracer)
+                    return await wrapper(self._original_acompletion, None, args, kwargs)
+
+                litellm.acompletion = instrumented_acompletion
+            except Exception as e:
+                logger.error(f"Failed to monkey-patch litellm.acompletion: {e}")
+
+        # Embeddings
+        if self._original_embedding:
+            try:
+
+                def instrumented_embedding(*args, **kwargs):  # type: ignore[no-untyped-def]
+                    wrapper = embedding_wrapper(tracer)
+                    return wrapper(self._original_embedding, None, args, kwargs)
+
+                litellm.embedding = instrumented_embedding
+            except Exception as e:
+                logger.error(f"Failed to monkey-patch litellm.embedding: {e}")
+
+        if self._original_aembedding:
+            try:
+
+                async def instrumented_aembedding(*args, **kwargs):  # type: ignore[no-untyped-def]
+                    wrapper = aembedding_wrapper(tracer)
+                    return await wrapper(self._original_aembedding, None, args, kwargs)
+
+                litellm.aembedding = instrumented_aembedding
+            except Exception as e:
+                logger.error(f"Failed to monkey-patch litellm.aembedding: {e}")
+
+        # Image generation
+        if self._original_image_generation:
+            try:
+
+                def instrumented_image_generation(*args, **kwargs):  # type: ignore[no-untyped-def]
+                    wrapper = image_generation_wrapper(tracer)
+                    return wrapper(self._original_image_generation, None, args, kwargs)
+
+                litellm.image_generation = instrumented_image_generation
+            except Exception as e:
+                logger.error(f"Failed to monkey-patch litellm.image_generation: {e}")
+
+        if self._original_aimage_generation:
+            try:
+
+                async def instrumented_aimage_generation(*args, **kwargs):  # type: ignore[no-untyped-def]
+                    wrapper = aimage_generation_wrapper(tracer)
+                    return await wrapper(self._original_aimage_generation, None, args, kwargs)
+
+                litellm.aimage_generation = instrumented_aimage_generation
+            except Exception as e:
+                logger.error(f"Failed to monkey-patch litellm.aimage_generation: {e}")
+
+    def _uninstrument(self, **kwargs):  # type: ignore[no-untyped-def]
+        """Uninstrument LiteLLM methods"""
+        try:
+            import litellm
+
+            # Restore original functions
+            if hasattr(self, "_original_completion") and self._original_completion:
+                litellm.completion = self._original_completion
+
+            if hasattr(self, "_original_acompletion") and self._original_acompletion:
+                litellm.acompletion = self._original_acompletion
+
+            if hasattr(self, "_original_embedding") and self._original_embedding:
+                litellm.embedding = self._original_embedding
+
+            if hasattr(self, "_original_aembedding") and self._original_aembedding:
+                litellm.aembedding = self._original_aembedding
+
+            if hasattr(self, "_original_image_generation") and self._original_image_generation:
+                litellm.image_generation = self._original_image_generation
+
+            if hasattr(self, "_original_aimage_generation") and self._original_aimage_generation:
+                litellm.aimage_generation = self._original_aimage_generation
+
+        except ImportError:
+            pass
+
+
+def should_suppress_instrumentation() -> bool:
+    """Check if instrumentation should be suppressed"""
+    return context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY) is True