netra-sdk 0.1.30__py3-none-any.whl → 0.1.31__py3-none-any.whl

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/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/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

netra/instrumentation/litellm/version.py

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

netra/instrumentation/litellm/wrappers.py

@@ -0,0 +1,557 @@
+import logging
+import time
+from collections.abc import Awaitable
+from typing import Any, AsyncIterator, Callable, Dict, Iterator, Tuple
+
+from opentelemetry import context as context_api
+from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
+from opentelemetry.semconv_ai import (
+    SpanAttributes,
+)
+from opentelemetry.trace import Span, SpanKind, Tracer
+from opentelemetry.trace.status import Status, StatusCode
+from wrapt import ObjectProxy
+
+logger = logging.getLogger(__name__)
+
+COMPLETION_SPAN_NAME = "litellm.completion"
+EMBEDDING_SPAN_NAME = "litellm.embedding"
+IMAGE_GENERATION_SPAN_NAME = "litellm.image_generation"
+
+
+def should_suppress_instrumentation() -> bool:
+    """Check if instrumentation should be suppressed"""
+    return context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY) is True
+
+
+def is_streaming_response(response: Any) -> bool:
+    """Check if response is a streaming response"""
+    return hasattr(response, "__iter__") and not isinstance(response, (str, bytes, dict))
+
+
+def model_as_dict(obj: Any) -> Dict[str, Any]:
+    """Convert LiteLLM model object to dictionary"""
+    if hasattr(obj, "model_dump"):
+        result = obj.model_dump()
+        return result if isinstance(result, dict) else {}
+    elif hasattr(obj, "to_dict"):
+        result = obj.to_dict()
+        return result if isinstance(result, dict) else {}
+    elif isinstance(obj, dict):
+        return obj
+    else:
+        return {}
+
+
+def set_request_attributes(span: Span, kwargs: Dict[str, Any], operation_type: str) -> None:
+    """Set request attributes on span"""
+    if not span.is_recording():
+        return
+
+    # Set operation type
+    span.set_attribute(f"{SpanAttributes.LLM_REQUEST_TYPE}", operation_type)
+    span.set_attribute(f"{SpanAttributes.LLM_SYSTEM}", "LiteLLM")
+
+    # Common attributes
+    if kwargs.get("model"):
+        span.set_attribute(f"{SpanAttributes.LLM_REQUEST_MODEL}", kwargs["model"])
+
+    if kwargs.get("temperature") is not None:
+        span.set_attribute(f"{SpanAttributes.LLM_REQUEST_TEMPERATURE}", kwargs["temperature"])
+
+    if kwargs.get("max_tokens") is not None:
+        span.set_attribute(f"{SpanAttributes.LLM_REQUEST_MAX_TOKENS}", kwargs["max_tokens"])
+
+    if kwargs.get("stream") is not None:
+        span.set_attribute("gen_ai.stream", kwargs["stream"])
+
+    # Chat completion specific attributes
+    if operation_type == "chat" and kwargs.get("messages"):
+        messages = kwargs["messages"]
+        if isinstance(messages, list) and len(messages) > 0:
+            for index, message in enumerate(messages):
+                if isinstance(message, dict):
+                    span.set_attribute(f"{SpanAttributes.LLM_PROMPTS}.{index}.role", message.get("role", "user"))
+                    span.set_attribute(f"{SpanAttributes.LLM_PROMPTS}.{index}.content", str(message.get("content", "")))
+
+    # Embedding specific attributes
+    if operation_type == "embedding" and kwargs.get("input"):
+        input_data = kwargs["input"]
+        if isinstance(input_data, str):
+            span.set_attribute(f"{SpanAttributes.LLM_PROMPTS}.0.content", input_data)
+        elif isinstance(input_data, list):
+            for index, text in enumerate(input_data):
+                if isinstance(text, str):
+                    span.set_attribute(f"{SpanAttributes.LLM_PROMPTS}.{index}.content", text)
+
+    # Image generation specific attributes
+    if operation_type == "image_generation":
+        if kwargs.get("prompt"):
+            span.set_attribute("gen_ai.prompt", kwargs["prompt"])
+        if kwargs.get("n"):
+            span.set_attribute("gen_ai.request.n", kwargs["n"])
+        if kwargs.get("size"):
+            span.set_attribute("gen_ai.request.size", kwargs["size"])
+        if kwargs.get("quality"):
+            span.set_attribute("gen_ai.request.quality", kwargs["quality"])
+        if kwargs.get("style"):
+            span.set_attribute("gen_ai.request.style", kwargs["style"])
+
+
+def set_response_attributes(span: Span, response_dict: Dict[str, Any], operation_type: str) -> None:
+    """Set response attributes on span"""
+    if not span.is_recording():
+        return
+
+    if response_dict.get("model"):
+        span.set_attribute(f"{SpanAttributes.LLM_RESPONSE_MODEL}", response_dict["model"])
+
+    if response_dict.get("id"):
+        span.set_attribute("gen_ai.response.id", response_dict["id"])
+
+    # Usage information
+    usage = response_dict.get("usage", {})
+    if usage:
+        if usage.get("prompt_tokens"):
+            span.set_attribute(f"{SpanAttributes.LLM_USAGE_PROMPT_TOKENS}", usage["prompt_tokens"])
+        if usage.get("completion_tokens"):
+            span.set_attribute(f"{SpanAttributes.LLM_USAGE_COMPLETION_TOKENS}", usage["completion_tokens"])
+        if usage.get("cache_read_input_tokens"):
+            span.set_attribute(f"{SpanAttributes.LLM_USAGE_CACHE_READ_INPUT_TOKENS}", usage["cache_read_input_tokens"])
+        if usage.get("cache_creation_input_tokens"):
+            span.set_attribute("gen_ai.usage.cache_creation_input_tokens", usage["cache_creation_input_tokens"])
+        if usage.get("total_tokens"):
+            span.set_attribute(f"{SpanAttributes.LLM_USAGE_TOTAL_TOKENS}", usage["total_tokens"])
+
+    # Chat completion response content
+    if operation_type == "chat":
+        choices = response_dict.get("choices", [])
+        for index, choice in enumerate(choices):
+            if choice.get("message", {}).get("role"):
+                span.set_attribute(f"{SpanAttributes.LLM_COMPLETIONS}.{index}.role", choice["message"]["role"])
+            if choice.get("message", {}).get("content"):
+                span.set_attribute(f"{SpanAttributes.LLM_COMPLETIONS}.{index}.content", choice["message"]["content"])
+            if choice.get("finish_reason"):
+                span.set_attribute(f"{SpanAttributes.LLM_COMPLETIONS}.{index}.finish_reason", choice["finish_reason"])
+
+    # Embedding response content
+    elif operation_type == "embedding":
+        data = response_dict.get("data", [])
+        for index, embedding_data in enumerate(data):
+            if embedding_data.get("index") is not None:
+                span.set_attribute(f"gen_ai.response.embeddings.{index}.index", embedding_data["index"])
+            if embedding_data.get("embedding"):
+                # Don't log the actual embedding vector, just its dimensions
+                embedding_vector = embedding_data["embedding"]
+                if isinstance(embedding_vector, list):
+                    span.set_attribute(f"gen_ai.response.embeddings.{index}.dimensions", len(embedding_vector))
+
+    # Image generation response content
+    elif operation_type == "image_generation":
+        data = response_dict.get("data", [])
+        for index, image_data in enumerate(data):
+            if image_data.get("url"):
+                span.set_attribute(f"gen_ai.response.images.{index}.url", image_data["url"])
+            if image_data.get("b64_json"):
+                span.set_attribute(f"gen_ai.response.images.{index}.has_b64_json", True)
+            if image_data.get("revised_prompt"):
+                span.set_attribute(f"gen_ai.response.images.{index}.revised_prompt", image_data["revised_prompt"])
+
+
+def completion_wrapper(tracer: Tracer) -> Callable[..., Any]:
+    """Wrapper for LiteLLM completion function"""
+
+    def wrapper(wrapped: Callable[..., Any], instance: Any, args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> Any:
+        logger.debug(f"LiteLLM completion wrapper called with model: {kwargs.get('model')}")
+
+        if should_suppress_instrumentation():
+            logger.debug("LiteLLM instrumentation suppressed")
+            return wrapped(*args, **kwargs)
+
+        # Check if streaming
+        is_streaming = kwargs.get("stream", False)
+
+        if is_streaming:
+            # Use start_span for streaming - returns span directly
+            span = tracer.start_span(
+                COMPLETION_SPAN_NAME, kind=SpanKind.CLIENT, attributes={"llm.request.type": "chat"}
+            )
+
+            set_request_attributes(span, kwargs, "chat")
+
+            try:
+                start_time = time.time()
+                response = wrapped(*args, **kwargs)
+
+                return StreamingWrapper(span=span, response=response, start_time=start_time, request_kwargs=kwargs)
+            except Exception as e:
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+                span.record_exception(e)
+                span.end()
+                raise
+        else:
+            # Use start_as_current_span for non-streaming - returns context manager
+            with tracer.start_as_current_span(
+                COMPLETION_SPAN_NAME, kind=SpanKind.CLIENT, attributes={"llm.request.type": "chat"}
+            ) as span:
+                set_request_attributes(span, kwargs, "chat")
+
+                try:
+                    start_time = time.time()
+                    response = wrapped(*args, **kwargs)
+                    end_time = time.time()
+
+                    response_dict = model_as_dict(response)
+                    set_response_attributes(span, response_dict, "chat")
+
+                    span.set_attribute("llm.response.duration", end_time - start_time)
+                    span.set_status(Status(StatusCode.OK))
+
+                    return response
+                except Exception as e:
+                    span.set_status(Status(StatusCode.ERROR, str(e)))
+                    raise
+
+    return wrapper
+
+
+def acompletion_wrapper(tracer: Tracer) -> Callable[..., Awaitable[Any]]:
+    """Async wrapper for LiteLLM acompletion function"""
+
+    async def wrapper(
+        wrapped: Callable[..., Awaitable[Any]], instance: Any, args: Tuple[Any, ...], kwargs: Dict[str, Any]
+    ) -> Any:
+        if should_suppress_instrumentation():
+            return await wrapped(*args, **kwargs)
+
+        # Check if streaming
+        is_streaming = kwargs.get("stream", False)
+
+        if is_streaming:
+            # Use start_span for streaming - returns span directly
+            span = tracer.start_span(
+                COMPLETION_SPAN_NAME, kind=SpanKind.CLIENT, attributes={"llm.request.type": "chat"}
+            )
+
+            set_request_attributes(span, kwargs, "chat")
+
+            try:
+                start_time = time.time()
+                response = await wrapped(*args, **kwargs)
+
+                return AsyncStreamingWrapper(span=span, response=response, start_time=start_time, request_kwargs=kwargs)
+            except Exception as e:
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+                span.record_exception(e)
+                span.end()
+                raise
+        else:
+            # Use start_as_current_span for non-streaming - returns context manager
+            with tracer.start_as_current_span(
+                COMPLETION_SPAN_NAME, kind=SpanKind.CLIENT, attributes={"llm.request.type": "chat"}
+            ) as span:
+                set_request_attributes(span, kwargs, "chat")
+
+                try:
+                    start_time = time.time()
+                    response = await wrapped(*args, **kwargs)
+                    end_time = time.time()
+
+                    response_dict = model_as_dict(response)
+                    set_response_attributes(span, response_dict, "chat")
+
+                    span.set_attribute("llm.response.duration", end_time - start_time)
+                    span.set_status(Status(StatusCode.OK))
+
+                    return response
+                except Exception as e:
+                    span.set_status(Status(StatusCode.ERROR, str(e)))
+                    raise
+
+    return wrapper
+
+
+def embedding_wrapper(tracer: Tracer) -> Callable[..., Any]:
+    """Wrapper for LiteLLM embedding function"""
+
+    def wrapper(wrapped: Callable[..., Any], instance: Any, args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> Any:
+        if should_suppress_instrumentation():
+            return wrapped(*args, **kwargs)
+
+        # Embeddings are never streaming, always use start_as_current_span
+        with tracer.start_as_current_span(
+            EMBEDDING_SPAN_NAME, kind=SpanKind.CLIENT, attributes={"llm.request.type": "embedding"}
+        ) as span:
+            set_request_attributes(span, kwargs, "embedding")
+
+            try:
+                start_time = time.time()
+                response = wrapped(*args, **kwargs)
+                end_time = time.time()
+
+                response_dict = model_as_dict(response)
+                set_response_attributes(span, response_dict, "embedding")
+
+                span.set_attribute("llm.response.duration", end_time - start_time)
+                span.set_status(Status(StatusCode.OK))
+
+                return response
+            except Exception as e:
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+                raise
+
+    return wrapper
+
+
+def aembedding_wrapper(tracer: Tracer) -> Callable[..., Awaitable[Any]]:
+    """Async wrapper for LiteLLM aembedding function"""
+
+    async def wrapper(
+        wrapped: Callable[..., Awaitable[Any]], instance: Any, args: Tuple[Any, ...], kwargs: Dict[str, Any]
+    ) -> Any:
+        if should_suppress_instrumentation():
+            return await wrapped(*args, **kwargs)
+
+        # Embeddings are never streaming, always use start_as_current_span
+        with tracer.start_as_current_span(
+            EMBEDDING_SPAN_NAME, kind=SpanKind.CLIENT, attributes={"llm.request.type": "embedding"}
+        ) as span:
+            set_request_attributes(span, kwargs, "embedding")
+
+            try:
+                start_time = time.time()
+                response = await wrapped(*args, **kwargs)
+                end_time = time.time()
+
+                response_dict = model_as_dict(response)
+                set_response_attributes(span, response_dict, "embedding")
+
+                span.set_attribute("llm.response.duration", end_time - start_time)
+                span.set_status(Status(StatusCode.OK))
+
+                return response
+            except Exception as e:
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+                raise
+
+    return wrapper
+
+
+def image_generation_wrapper(tracer: Tracer) -> Callable[..., Any]:
+    """Wrapper for LiteLLM image_generation function"""
+
+    def wrapper(wrapped: Callable[..., Any], instance: Any, args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> Any:
+        if should_suppress_instrumentation():
+            return wrapped(*args, **kwargs)
+
+        # Image generation is never streaming, always use start_as_current_span
+        with tracer.start_as_current_span(
+            IMAGE_GENERATION_SPAN_NAME, kind=SpanKind.CLIENT, attributes={"llm.request.type": "image_generation"}
+        ) as span:
+            set_request_attributes(span, kwargs, "image_generation")
+
+            try:
+                start_time = time.time()
+                response = wrapped(*args, **kwargs)
+                end_time = time.time()
+
+                response_dict = model_as_dict(response)
+                set_response_attributes(span, response_dict, "image_generation")
+
+                span.set_attribute("llm.response.duration", end_time - start_time)
+                span.set_status(Status(StatusCode.OK))
+
+                return response
+            except Exception as e:
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+                raise
+
+    return wrapper
+
+
+def aimage_generation_wrapper(tracer: Tracer) -> Callable[..., Awaitable[Any]]:
+    """Async wrapper for LiteLLM aimage_generation function"""
+
+    async def wrapper(
+        wrapped: Callable[..., Awaitable[Any]], instance: Any, args: Tuple[Any, ...], kwargs: Dict[str, Any]
+    ) -> Any:
+        if should_suppress_instrumentation():
+            return await wrapped(*args, **kwargs)
+
+        # Image generation is never streaming, always use start_as_current_span
+        with tracer.start_as_current_span(
+            IMAGE_GENERATION_SPAN_NAME, kind=SpanKind.CLIENT, attributes={"llm.request.type": "image_generation"}
+        ) as span:
+            set_request_attributes(span, kwargs, "image_generation")
+
+            try:
+                start_time = time.time()
+                response = await wrapped(*args, **kwargs)
+                end_time = time.time()
+
+                response_dict = model_as_dict(response)
+                set_response_attributes(span, response_dict, "image_generation")
+
+                span.set_attribute("llm.response.duration", end_time - start_time)
+                span.set_status(Status(StatusCode.OK))
+
+                return response
+            except Exception as e:
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+                raise
+
+    return wrapper
+
+
+class StreamingWrapper(ObjectProxy):  # type: ignore[misc]
+    """Wrapper for streaming responses"""
+
+    def __init__(self, span: Span, response: Iterator[Any], start_time: float, request_kwargs: Dict[str, Any]) -> None:
+        super().__init__(response)
+        self._span = span
+        self._start_time = start_time
+        self._request_kwargs = request_kwargs
+        self._complete_response: Dict[str, Any] = {"choices": [], "model": ""}
+        self._content_parts: list[str] = []
+
+    def __iter__(self) -> Iterator[Any]:
+        return self
+
+    def __next__(self) -> Any:
+        try:
+            chunk = self.__wrapped__.__next__()
+            self._process_chunk(chunk)
+            return chunk
+        except StopIteration:
+            self._finalize_span()
+            raise
+
+    def _process_chunk(self, chunk: Any) -> None:
+        """Process streaming chunk"""
+        chunk_dict = model_as_dict(chunk)
+
+        # Accumulate response data
+        if chunk_dict.get("model"):
+            self._complete_response["model"] = chunk_dict["model"]
+
+        # Accumulate usage information from chunks
+        if chunk_dict.get("usage"):
+            self._complete_response["usage"] = chunk_dict["usage"]
+
+        # Collect content from delta
+        choices = chunk_dict.get("choices", [])
+        for choice in choices:
+            delta = choice.get("delta", {})
+            if delta.get("content"):
+                self._content_parts.append(delta["content"])
+
+            # Collect finish_reason from choices
+            if choice.get("finish_reason"):
+                if "choices" not in self._complete_response:
+                    self._complete_response["choices"] = []
+                # Ensure we have enough choice entries
+                while len(self._complete_response["choices"]) <= len(choices) - 1:
+                    self._complete_response["choices"].append(
+                        {"message": {"role": "assistant", "content": ""}, "finish_reason": None}
+                    )
+
+                choice_index = choice.get("index", 0)
+                if choice_index < len(self._complete_response["choices"]):
+                    self._complete_response["choices"][choice_index]["finish_reason"] = choice["finish_reason"]
+
+        # Add chunk event
+        self._span.add_event("llm.content.completion.chunk")
+
+    def _finalize_span(self) -> None:
+        """Finalize span when streaming is complete"""
+        end_time = time.time()
+        duration = end_time - self._start_time
+
+        # Set accumulated content
+        if self._content_parts:
+            full_content = "".join(self._content_parts)
+            self._span.set_attribute(f"{SpanAttributes.LLM_COMPLETIONS}.0.content", full_content)
+            self._span.set_attribute(f"{SpanAttributes.LLM_COMPLETIONS}.0.role", "assistant")
+
+        set_response_attributes(self._span, self._complete_response, "chat")
+        self._span.set_attribute("llm.response.duration", duration)
+        self._span.set_status(Status(StatusCode.OK))
+        self._span.end()
+
+
+class AsyncStreamingWrapper(ObjectProxy):  # type: ignore[misc]
+    """Async wrapper for streaming responses"""
+
+    def __init__(
+        self, span: Span, response: AsyncIterator[Any], start_time: float, request_kwargs: Dict[str, Any]
+    ) -> None:
+        super().__init__(response)
+        self._span = span
+        self._start_time = start_time
+        self._request_kwargs = request_kwargs
+        self._complete_response: Dict[str, Any] = {"choices": [], "model": ""}
+        self._content_parts: list[str] = []
+
+    def __aiter__(self) -> AsyncIterator[Any]:
+        return self
+
+    async def __anext__(self) -> Any:
+        try:
+            chunk = await self.__wrapped__.__anext__()
+            self._process_chunk(chunk)
+            return chunk
+        except StopAsyncIteration:
+            self._finalize_span()
+            raise
+
+    def _process_chunk(self, chunk: Any) -> None:
+        """Process streaming chunk"""
+        chunk_dict = model_as_dict(chunk)
+
+        # Accumulate response data
+        if chunk_dict.get("model"):
+            self._complete_response["model"] = chunk_dict["model"]
+
+        # Accumulate usage information from chunks
+        if chunk_dict.get("usage"):
+            self._complete_response["usage"] = chunk_dict["usage"]
+
+        # Collect content from delta
+        choices = chunk_dict.get("choices", [])
+        for choice in choices:
+            delta = choice.get("delta", {})
+            if delta.get("content"):
+                self._content_parts.append(delta["content"])
+
+            # Collect finish_reason from choices
+            if choice.get("finish_reason"):
+                if "choices" not in self._complete_response:
+                    self._complete_response["choices"] = []
+                # Ensure we have enough choice entries
+                while len(self._complete_response["choices"]) <= len(choices) - 1:
+                    self._complete_response["choices"].append(
+                        {"message": {"role": "assistant", "content": ""}, "finish_reason": None}
+                    )
+
+                choice_index = choice.get("index", 0)
+                if choice_index < len(self._complete_response["choices"]):
+                    self._complete_response["choices"][choice_index]["finish_reason"] = choice["finish_reason"]
+
+        # Add chunk event
+        self._span.add_event("llm.content.completion.chunk")
+
+    def _finalize_span(self) -> None:
+        """Finalize span when streaming is complete"""
+        end_time = time.time()
+        duration = end_time - self._start_time
+
+        # Set accumulated content
+        if self._content_parts:
+            full_content = "".join(self._content_parts)
+            self._span.set_attribute(f"{SpanAttributes.LLM_COMPLETIONS}.0.content", full_content)
+            self._span.set_attribute(f"{SpanAttributes.LLM_COMPLETIONS}.0.role", "assistant")
+
+        set_response_attributes(self._span, self._complete_response, "chat")
+        self._span.set_attribute("llm.response.duration", duration)
+        self._span.set_status(Status(StatusCode.OK))
+        self._span.end()

netra/version.py

@@ -1 +1 @@
-__version__ = "0.1.30"
+__version__ = "0.1.31"

{netra_sdk-0.1.30.dist-info → netra_sdk-0.1.31.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: netra-sdk
-Version: 0.1.30
+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.30.dist-info → netra_sdk-0.1.31.dist-info}/RECORD

@@ -9,7 +9,7 @@ netra/exceptions/__init__.py,sha256=uDgcBxmC4WhdS7HRYQk_TtJyxH1s1o6wZmcsnSHLAcM,
 netra/exceptions/injection.py,sha256=ke4eUXRYUFJkMZgdSyPPkPt5PdxToTI6xLEBI0hTWUQ,1332
 netra/exceptions/pii.py,sha256=MT4p_x-zH3VtYudTSxw1Z9qQZADJDspq64WrYqSWlZc,2438
 netra/input_scanner.py,sha256=At6N9gNY8cR0O6S8x3K6swWBV3P1a_9O-XBNM_pcKz4,5348
-netra/instrumentation/__init__.py,sha256=pJOkAO1h7rdM_VwvZ_FZZ-zq8PCmLOzW4jvr_CwWYCI,40792
+netra/instrumentation/__init__.py,sha256=HdG3n5TxPRUNlOxsqjlvwDmBcnm3UtYx1OecLhnLeQM,41578
 netra/instrumentation/aiohttp/__init__.py,sha256=M1kuF0R3gKY5rlbhEC1AR13UWHelmfokluL2yFysKWc,14398
 netra/instrumentation/aiohttp/version.py,sha256=Zy-0Aukx-HS_Mo3NKPWg-hlUoWKDzS0w58gLoVtJec8,24
 netra/instrumentation/cohere/__init__.py,sha256=3XwmCAZwZiMkHdNN3YvcBOLsNCx80ymbU31TyMzv1IY,17685
@@ -22,7 +22,10 @@ netra/instrumentation/google_genai/utils.py,sha256=2OeSN5jUaMKF4x5zWiW65R1LB_a44
 netra/instrumentation/google_genai/version.py,sha256=Hww1duZrC8kYK7ThBSQVyz0HNOb0ys_o8Pln-wVQ1hI,23
 netra/instrumentation/httpx/__init__.py,sha256=w1su_eQP_w5ZJHq0Lf-4miF5zM4OOW0ItmRp0wi85Ew,19388
 netra/instrumentation/httpx/version.py,sha256=ZRQKbgDaGz_yuLk-cUKuk6ZBKCSRKZC8nQd041NRNXk,23
-netra/instrumentation/instruments.py,sha256=JJF8J2O2Xd3w3k33ZYxpFNrwWgl_veRNxV6QUFCsFn0,4301
+netra/instrumentation/instruments.py,sha256=O6MI_BO-5EBkVqI-dr5eqhYnk8mP5QEpI0RWJ7Fe3FQ,4349
+netra/instrumentation/litellm/__init__.py,sha256=H9FsdEq-CL39zbl_dLm8D43-D1vAjoNqFTBpbmZsVXs,6740
+netra/instrumentation/litellm/version.py,sha256=J-j-u0itpEFT6irdmWmixQqYMadNl1X91TxUmoiLHMI,22
+netra/instrumentation/litellm/wrappers.py,sha256=H_UG0et6PUmj6CQagvNzbs_WodNTMruzzGOHhedmTko,22840
 netra/instrumentation/mistralai/__init__.py,sha256=RE0b-rS6iXdoynJMFKHL9s97eYo5HghrJa013fR4ZhI,18910
 netra/instrumentation/mistralai/config.py,sha256=XCyo3mk30qkvqyCqeTrKwROahu0gcOEwmbDLOo53J5k,121
 netra/instrumentation/mistralai/utils.py,sha256=nhdIer5gJFxuGwg8FCT222hggDHeMQDhJctnDSwLqcc,894
@@ -44,8 +47,8 @@ netra/scanner.py,sha256=kyDpeZiscCPb6pjuhS-sfsVj-dviBFRepdUWh0sLoEY,11554
 netra/session_manager.py,sha256=AoQa-k4dFcq7PeOD8G8DNzhLzL1JrHUW6b_y8mRyTQo,10255
 netra/span_wrapper.py,sha256=lGuV1F4Q5I_swIoIof5myzOQCFmGFdtrpgfQt7dTTus,8105
 netra/tracer.py,sha256=YiuijB_5DBOLVgE39Lj3thWVmUqHLcqbdFVB0HGovW0,3543
-netra/version.py,sha256=2GUJJyX8g8EAXKUqyj7DGVzG-jNXOGaqVSWilvGYuX8,23
-netra_sdk-0.1.30.dist-info/LICENCE,sha256=8B_UoZ-BAl0AqiHAHUETCgd3I2B9yYJ1WEQtVb_qFMA,11359
-netra_sdk-0.1.30.dist-info/METADATA,sha256=9quLP0B1XfOdcQWU4vsIY44Gf2CwWiF_er2-Wa6wKvM,28151
-netra_sdk-0.1.30.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-netra_sdk-0.1.30.dist-info/RECORD,,
+netra/version.py,sha256=i-fDEsQ0iAiPKXFaj9eERDqcxl3BqNnavaCEqpNxmVI,23
+netra_sdk-0.1.31.dist-info/LICENCE,sha256=8B_UoZ-BAl0AqiHAHUETCgd3I2B9yYJ1WEQtVb_qFMA,11359
+netra_sdk-0.1.31.dist-info/METADATA,sha256=VvltGCy_nbt-TRB91KiH_hu6YC4CceY_uQW-UYXT7NE,28196
+netra_sdk-0.1.31.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+netra_sdk-0.1.31.dist-info/RECORD,,