netra-sdk 0.1.3__tar.gz → 0.1.4__tar.gz

{netra_sdk-0.1.3 → netra_sdk-0.1.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: netra-sdk
-Version: 0.1.3
+Version: 0.1.4
 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.3 → netra_sdk-0.1.4}/netra/instrumentation/__init__.py

@@ -281,23 +281,12 @@ def init_fastapi_instrumentation() -> bool:
         bool: True if initialization was successful, False otherwise.
     """
     try:
-        if not is_package_installed("fastapi"):
-            return True
-        from fastapi import FastAPI
+        if is_package_installed("fastapi"):
+            from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
 
-        original_init = FastAPI.__init__
-
-        def _patched_init(self: FastAPI, *args: Any, **kwargs: Any) -> None:
-            original_init(self, *args, **kwargs)
-
-            try:
-                from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
-
-                FastAPIInstrumentor().instrument_app(self)
-            except Exception as e:
-                logging.warning(f"Failed to auto-instrument FastAPI: {e}")
-
-        FastAPI.__init__ = _patched_init
+            instrumentor = FastAPIInstrumentor()
+            if not instrumentor.is_instrumented_by_opentelemetry:
+                instrumentor.instrument()
         return True
     except Exception as e:
         logging.error(f"Error initializing FastAPI instrumentor: {e}")

netra_sdk-0.1.4/netra/processors/__init__.py

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

netra_sdk-0.1.4/netra/processors/error_detection_processor.py

@@ -0,0 +1,66 @@
+import logging
+from typing import Any, Optional
+
+import httpx
+from opentelemetry.sdk.trace import SpanProcessor
+from opentelemetry.trace import Context, Span
+
+from netra import Netra
+
+logger = logging.getLogger(__name__)
+
+
+class ErrorDetectionProcessor(SpanProcessor):  # type: ignore[misc]
+    """
+    OpenTelemetry span processor that monitors for error attributes in spans and creates custom events.
+    """
+
+    def __init__(self) -> None:
+        pass
+
+    def on_start(self, span: Span, parent_context: Optional[Context] = None) -> None:
+        """Called when a span starts."""
+        span_id = self._get_span_id(span)
+        if not span_id:
+            return
+
+        # Wrap span methods to capture data
+        self._wrap_span_methods(span, span_id)
+
+    def on_end(self, span: Span) -> None:
+        """Called when a span ends."""
+
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        """Force flush any pending data."""
+        return True
+
+    def shutdown(self) -> bool:
+        """Shutdown the processor."""
+        return True
+
+    def _get_span_id(self, span: Span) -> Optional[str]:
+        """Get a unique identifier for the span."""
+        try:
+            span_context = span.get_span_context()
+            return f"{span_context.trace_id:032x}-{span_context.span_id:016x}"
+        except Exception:
+            return None
+
+    def _status_code_processing(self, status_code: int) -> None:
+        if httpx.codes.is_error(status_code):
+            event_attributes = {"has_error": True, "status_code": status_code}
+            Netra.set_custom_event(event_name="error_detected", attributes=event_attributes)
+
+    def _wrap_span_methods(self, span: Span, span_id: str) -> Any:
+        """Wrap span methods to capture attributes and events."""
+        # Wrap set_attribute
+        original_set_attribute = span.set_attribute
+
+        def wrapped_set_attribute(key: str, value: Any) -> Any:
+            # Status code processing
+            if key == "http.status_code":
+                self._status_code_processing(value)
+
+            return original_set_attribute(key, value)
+
+        span.set_attribute = wrapped_set_attribute

{netra_sdk-0.1.3 → netra_sdk-0.1.4}/netra/tracer.py

@@ -66,10 +66,10 @@ class Tracer:
                 headers=self.cfg.headers,
             )
         # Add span processors for session span processing and data aggregation processing
-        from netra.processors import SessionSpanProcessor, SpanAggregationProcessor
+        from netra.processors import ErrorDetectionProcessor, SessionSpanProcessor
 
         provider.add_span_processor(SessionSpanProcessor())
-        provider.add_span_processor(SpanAggregationProcessor())
+        provider.add_span_processor(ErrorDetectionProcessor())
 
         # Install appropriate span processor
         if self.cfg.disable_batch:

netra_sdk-0.1.4/netra/version.py

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

{netra_sdk-0.1.3 → netra_sdk-0.1.4}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [project]
 name = "netra-sdk"
-version = "0.1.3"
+version = "0.1.4"
 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.3/netra/processors/__init__.py

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

netra_sdk-0.1.3/netra/processors/span_aggregation_processor.py

@@ -1,365 +0,0 @@
-"""
-Span aggregation utilities for Combat SDK.
-Handles aggregation of child span data into parent spans.
-"""
-
-import json
-import logging
-from collections import defaultdict
-from typing import Any, Dict, Optional, Set
-
-import httpx
-from opentelemetry import trace
-from opentelemetry.sdk.trace import SpanProcessor
-from opentelemetry.trace import Context, Span
-
-from netra import Netra
-from netra.config import Config
-
-logger = logging.getLogger(__name__)
-
-
-class SpanAggregationData:
-    """Holds aggregated data for a span."""
-
-    def __init__(self) -> None:
-        self.tokens: Dict[str, Dict[str, int]] = defaultdict(lambda: defaultdict(int))
-        self.models: Set[str] = set()
-        self.has_pii: bool = False
-        self.pii_entities: Set[str] = set()
-        self.pii_actions: Dict[str, Set[str]] = defaultdict(set)
-        self.has_violation: bool = False
-        self.violations: Set[str] = set()
-        self.violation_actions: Dict[str, Set[str]] = defaultdict(set)
-        self.has_error: bool = False
-        self.status_codes: Set[int] = set()
-
-    def merge_from_other(self, other: "SpanAggregationData") -> None:
-        """Merge data from another SpanAggregationData instance."""
-        # Merge error data
-        if other.has_error:
-            self.has_error = True
-            self.status_codes.update(other.status_codes)
-
-        # Merge tokens - take the maximum values for each model
-        for model, token_data in other.tokens.items():
-            if model not in self.tokens:
-                self.tokens[model] = {}
-            for token_type, count in token_data.items():
-                self.tokens[model][token_type] = max(self.tokens[model].get(token_type, 0), count)
-
-        # Merge models
-        self.models.update(other.models)
-
-        # Merge PII data
-        if other.has_pii:
-            self.has_pii = True
-        self.pii_entities.update(other.pii_entities)
-        for action, entities in other.pii_actions.items():
-            self.pii_actions[action].update(entities)
-
-        # Merge violation data
-        if other.has_violation:
-            self.has_violation = True
-        self.violations.update(other.violations)
-        for action, violations in other.violation_actions.items():
-            self.violation_actions[action].update(violations)
-
-    def to_attributes(self) -> Dict[str, str]:
-        """Convert aggregated data to span attributes."""
-        attributes = {}
-
-        # Error Data
-        attributes["has_error"] = str(self.has_error).lower()
-        if self.has_error:
-            attributes["status_codes"] = json.dumps(list(self.status_codes))
-
-        # Token usage by model
-        if self.tokens:
-            tokens_dict = {}
-            for model, usage in self.tokens.items():
-                tokens_dict[model] = dict(usage)
-            attributes["tokens"] = json.dumps(tokens_dict)
-
-        # Models used
-        if self.models:
-            attributes["models"] = json.dumps(sorted(list(self.models)))
-
-        # PII information
-        attributes["has_pii"] = str(self.has_pii).lower()
-        if self.pii_entities:
-            attributes["pii_entities"] = json.dumps(sorted(list(self.pii_entities)))
-        if self.pii_actions:
-            pii_actions_dict = {}
-            for action, entities in self.pii_actions.items():
-                pii_actions_dict[action] = sorted(list(entities))
-            attributes["pii_actions"] = json.dumps(pii_actions_dict)
-
-        # Violation information
-        attributes["has_violation"] = str(self.has_violation).lower()
-        if self.violations:
-            attributes["violations"] = json.dumps(sorted(list(self.violations)))
-        if self.violation_actions:
-            violation_actions_dict = {}
-            for action, violations in self.violation_actions.items():
-                violation_actions_dict[action] = sorted(list(violations))
-            attributes["violation_actions"] = json.dumps(violation_actions_dict)
-
-        return attributes
-
-
-class SpanAggregationProcessor(SpanProcessor):  # type: ignore[misc]
-    """
-    OpenTelemetry span processor that aggregates data from child spans into parent spans.
-    """
-
-    def __init__(self) -> None:
-        self._span_data: Dict[str, SpanAggregationData] = {}
-        self._span_hierarchy: Dict[str, Optional[str]] = {}  # child_id -> parent_id
-        self._root_spans: Set[str] = set()
-        self._captured_data: Dict[str, Dict[str, Any]] = {}  # span_id -> {attributes, events}
-        self._active_spans: Dict[str, Span] = {}  # span_id -> original span reference
-
-    def on_start(self, span: Span, parent_context: Optional[Context] = None) -> None:
-        """Called when a span starts."""
-        span_id = self._get_span_id(span)
-        if not span_id:
-            return
-
-        # Store the original span for later use
-        self._active_spans[span_id] = span
-
-        # Initialize aggregation data
-        self._span_data[span_id] = SpanAggregationData()
-        self._captured_data[span_id] = {"attributes": {}, "events": []}
-
-        # Check if this is a root span (no parent)
-        if span.parent is None:
-            self._root_spans.add(span_id)
-        else:
-            # Track parent-child relationship - span.parent is a SpanContext, not a Span
-            try:
-                parent_span_context = span.parent
-                if parent_span_context and parent_span_context.span_id:
-                    parent_span_id = f"{parent_span_context.trace_id:032x}-{parent_span_context.span_id:016x}"
-                    self._span_hierarchy[span_id] = parent_span_id
-                else:
-                    logger.warning(f"DEBUG: Parent span context is invalid for child {span_id}")
-            except Exception as e:
-                logger.warning(f"DEBUG: Could not get parent span ID for child {span_id}: {e}")
-
-        # Wrap span methods to capture data
-        self._wrap_span_methods(span, span_id)
-
-    def on_end(self, span: Span) -> None:
-        """Called when a span ends."""
-        span_id = self._get_span_id(span)
-        if not span_id or span_id not in self._span_data:
-            return
-
-        try:
-            # Process this span's captured data
-            captured = self._captured_data.get(span_id, {})
-            self._process_attributes(self._span_data[span_id], captured.get("attributes", {}))
-
-            # Set aggregated attributes on this span
-            original_span = self._active_spans.get(span_id)
-            if original_span and original_span.is_recording():
-                self._set_span_attributes(original_span, self._span_data[span_id])
-
-            # Handle parent-child aggregation for any remaining data
-            self._aggregate_to_all_parents(span_id)
-
-        except Exception as e:
-            logger.error(f"Error during span aggregation for span {span_id}: {e}")
-            # Even if there's an error, try to do basic aggregation
-            try:
-                original_span = self._active_spans.get(span_id)
-                if original_span and original_span.is_recording():
-                    self._set_span_attributes(original_span, self._span_data[span_id])
-            except Exception as inner_e:
-                logger.error(f"Failed to set basic aggregation attributes: {inner_e}")
-
-        # Clean up
-        self._span_data.pop(span_id, None)
-        self._captured_data.pop(span_id, None)
-        self._active_spans.pop(span_id, None)
-        self._root_spans.discard(span_id)
-        self._span_hierarchy.pop(span_id, None)
-
-    def _wrap_span_methods(self, span: Span, span_id: str) -> Any:
-        """Wrap span methods to capture attributes and events."""
-        # Wrap set_attribute
-        original_set_attribute = span.set_attribute
-
-        def wrapped_set_attribute(key: str, value: Any) -> Any:
-            # Status code processing
-            if key == "http.status_code":
-                self._status_code_processing(value)
-
-            # Capture the all the attribute data
-            self._captured_data[span_id]["attributes"][key] = value
-            return original_set_attribute(key, value)
-
-        span.set_attribute = wrapped_set_attribute
-
-        # Wrap add_event
-        original_add_event = span.add_event
-
-        def wrapped_add_event(name: str, attributes: Dict[str, Any] = {}, timestamp: int = 0) -> Any:
-            # Only process PII and violation events
-            if name == "pii_detected" and attributes:
-                self._process_pii_event(self._span_data[span_id], attributes)
-                if span.is_recording():
-                    self._set_span_attributes(span, self._span_data[span_id])
-                    # Immediately aggregate to parent spans
-                    self._aggregate_to_all_parents(span_id)
-            elif name == "violation_detected" and attributes:
-                self._process_violation_event(self._span_data[span_id], attributes)
-                if span.is_recording():
-                    self._set_span_attributes(span, self._span_data[span_id])
-                    # Immediately aggregate to parent spans
-                    self._aggregate_to_all_parents(span_id)
-
-            # Check if span is still recording before adding event
-            if not span.is_recording():
-                logger.debug(f"Attempted to add event to ended span {span_id}")
-                return None
-            return original_add_event(name, attributes, timestamp)
-
-        span.add_event = wrapped_add_event
-
-    def _process_attributes(self, data: SpanAggregationData, attributes: Dict[str, Any]) -> None:
-        """Process span attributes for aggregation."""
-        # Extract status code for error identification
-        status_code = attributes.get("http.status_code", 200)
-        if httpx.codes.is_error(status_code):
-            data.has_error = True
-            data.status_codes.update([status_code])
-
-        # Extract model information
-        model = attributes.get("gen_ai.request.model") or attributes.get("gen_ai.response.model")
-        if model:
-            data.models.add(model)
-            # Extract token usage
-            token_fields = {
-                "prompt_tokens": attributes.get("gen_ai.usage.prompt_tokens", 0),
-                "completion_tokens": attributes.get("gen_ai.usage.completion_tokens", 0),
-                "total_tokens": attributes.get("llm.usage.total_tokens", 0),
-                "cache_read_input_tokens": attributes.get("gen_ai.usage.cache_read_input_tokens", 0),
-            }
-
-            # Initialize token fields if they don't exist
-            if model not in data.tokens:
-                data.tokens[model] = {}
-
-            # Add token values
-            for field, value in token_fields.items():
-                if isinstance(value, (int, str)):
-                    current_value = data.tokens[model].get(field, 0)
-                    data.tokens[model][field] = current_value + int(value)
-
-    def _process_pii_event(self, data: SpanAggregationData, attrs: Dict[str, Any]) -> None:
-        """Process pii_detected event."""
-        if attrs.get("has_pii"):
-            data.has_pii = True
-
-            # Extract entities from pii_entities field
-            entity_counts_str = attrs.get("pii_entities")
-            if entity_counts_str:
-                try:
-                    entity_counts = (
-                        json.loads(entity_counts_str) if isinstance(entity_counts_str, str) else entity_counts_str
-                    )
-                    if isinstance(entity_counts, dict):
-                        entities = set(entity_counts.keys())
-                        data.pii_entities.update(entities)
-
-                        # Determine action
-                        if attrs.get("is_blocked"):
-                            data.pii_actions["BLOCK"].update(entities)
-                        elif attrs.get("is_masked"):
-                            data.pii_actions["MASK"].update(entities)
-                        else:
-                            data.pii_actions["FLAG"].update(entities)
-                except (json.JSONDecodeError, TypeError):
-                    logger.error(f"Failed to parse pii_entities: {entity_counts_str}")
-
-    def _process_violation_event(self, data: SpanAggregationData, attrs: Dict[str, Any]) -> None:
-        """Process violation_detected event."""
-        if attrs.get("has_violation"):
-            data.has_violation = True
-            violations = attrs.get("violations", [])
-            if violations:
-                data.violations.update(violations)
-                # Set action based on is_blocked flag
-                action = "BLOCK" if attrs.get("is_blocked") else "FLAG"
-                data.violation_actions[action].update(violations)
-
-    def _aggregate_to_all_parents(self, child_span_id: str) -> None:
-        """Aggregate data from child span to all its parent spans in the hierarchy."""
-        if child_span_id not in self._span_data:
-            return
-
-        child_data = self._span_data[child_span_id]
-        current_span_id = child_span_id
-
-        # Traverse up the parent hierarchy
-        while True:
-            parent_id = self._span_hierarchy.get(current_span_id)
-            if not parent_id or parent_id not in self._span_data:
-                break
-
-            # Merge child data into parent
-            self._span_data[parent_id].merge_from_other(child_data)
-
-            # Update parent span attributes if it's still active and recording
-            parent_span = self._active_spans.get(parent_id)
-            if parent_span and parent_span.is_recording():
-                self._set_span_attributes(parent_span, self._span_data[parent_id])
-
-            # Move up to the next parent
-            current_span_id = parent_id
-
-    def _set_span_attributes(self, span: Span, data: SpanAggregationData) -> None:
-        """Set aggregated attributes on the given span."""
-        try:
-            aggregated_attrs = data.to_attributes()
-            # Set all aggregated attributes under a single 'aggregator' key as a JSON object
-            span.set_attribute(f"{Config.LIBRARY_NAME}.aggregated_attributes", json.dumps(aggregated_attrs))
-        except Exception as e:
-            logger.error(f"Failed to set aggregated attributes: {e}")
-
-    def _get_span_id(self, span: Span) -> Optional[str]:
-        """Get a unique identifier for the span."""
-        try:
-            span_context = span.get_span_context()
-            return f"{span_context.trace_id:032x}-{span_context.span_id:016x}"
-        except Exception:
-            return None
-
-    def _get_span_id_from_context(self, context: Context) -> Optional[str]:
-        """Extract span ID from context."""
-        if context:
-            span_context = trace.get_current_span(context).get_span_context()
-            if span_context and span_context.span_id:
-                return f"{span_context.trace_id:032x}-{span_context.span_id:016x}"
-        return None
-
-    def _status_code_processing(self, status_code: int) -> None:
-        if httpx.codes.is_error(status_code):
-            event_attributes = {"has_error": True, "status_code": status_code}
-            Netra.set_custom_event(event_name="error_detected", attributes=event_attributes)
-
-    def force_flush(self, timeout_millis: int = 30000) -> bool:
-        """Force flush any pending data."""
-        return True
-
-    def shutdown(self) -> bool:
-        """Shutdown the processor."""
-        self._span_data.clear()
-        self._span_hierarchy.clear()
-        self._root_spans.clear()
-        self._captured_data.clear()
-        self._active_spans.clear()
-        return True

netra_sdk-0.1.3/netra/version.py

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