netra-sdk 0.1.3__py3-none-any.whl → 0.1.5__py3-none-any.whl

netra/decorators.py

@@ -12,6 +12,7 @@ from typing import Any, Awaitable, Callable, Dict, Optional, ParamSpec, Tuple, T
 from opentelemetry import trace
 
 from .config import Config
+from .session_manager import SessionManager
 
 P = ParamSpec("P")
 R = TypeVar("R")
@@ -78,6 +79,9 @@ def _create_function_wrapper(func: Callable[P, R], entity_type: str, name: Optio
 
         @functools.wraps(func)
         async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+            # Push entity to stack before span starts so SessionSpanProcessor can capture it
+            SessionManager.push_entity(entity_type, span_name)
+
             tracer = trace.get_tracer(module_name)
             with tracer.start_as_current_span(span_name) as span:
                 _add_span_attributes(span, func, args, kwargs, entity_type)
@@ -87,8 +91,10 @@ def _create_function_wrapper(func: Callable[P, R], entity_type: str, name: Optio
                     return result
                 except Exception as e:
                     span.set_attribute(f"{Config.LIBRARY_NAME}.entity.error", str(e))
-                    span.record_exception(e)
                     raise
+                finally:
+                    # Pop entity from stack after function call is done
+                    SessionManager.pop_entity(entity_type)
 
         return cast(Callable[P, R], async_wrapper)
 
@@ -96,6 +102,9 @@ def _create_function_wrapper(func: Callable[P, R], entity_type: str, name: Optio
 
         @functools.wraps(func)
         def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+            # Push entity to stack before span starts so SessionSpanProcessor can capture it
+            SessionManager.push_entity(entity_type, span_name)
+
             tracer = trace.get_tracer(module_name)
             with tracer.start_as_current_span(span_name) as span:
                 _add_span_attributes(span, func, args, kwargs, entity_type)
@@ -105,8 +114,10 @@ def _create_function_wrapper(func: Callable[P, R], entity_type: str, name: Optio
                     return result
                 except Exception as e:
                     span.set_attribute(f"{Config.LIBRARY_NAME}.entity.error", str(e))
-                    span.record_exception(e)
                     raise
+                finally:
+                    # Pop entity from stack after function call is done
+                    SessionManager.pop_entity(entity_type)
 
         return cast(Callable[P, R], sync_wrapper)
 

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/processors/__init__.py

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

netra/processors/error_detection_processor.py

@@ -0,0 +1,80 @@
+import logging
+from typing import Any, Optional
+
+import httpx
+from opentelemetry.sdk.trace import SpanProcessor
+from opentelemetry.trace import Context, Span, Status, StatusCode
+
+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)
+
+        # Wrap set_status
+        original_set_status = span.set_status
+
+        def wrapped_set_status(status: Status) -> Any:
+            # Check if status code is ERROR
+            if status.status_code == StatusCode.ERROR:
+                event_attributes = {
+                    "has_error": True,
+                }
+                Netra.set_custom_event(event_name="error_detected", attributes=event_attributes)
+
+            return original_set_status(status)
+
+        span.set_attribute = wrapped_set_attribute
+        span.set_status = wrapped_set_status

netra/processors/session_span_processor.py

@@ -42,6 +42,12 @@ class SessionSpanProcessor(SpanProcessor):  # type: ignore[misc]
                     value = baggage.get_baggage(f"custom.{key}", ctx)
                     if value:
                         span.set_attribute(f"{Config.LIBRARY_NAME}.custom.{key}", value)
+
+            # Add entity attributes from SessionManager
+            entity_attributes = SessionManager.get_current_entity_attributes()
+            for attr_key, attr_value in entity_attributes.items():
+                span.set_attribute(attr_key, attr_value)
+
         except Exception as e:
             logger.exception(f"Error setting span attributes: {e}")
 

netra/session.py

@@ -1,11 +1,13 @@
+import json
 import logging
 import time
-from typing import Any, Dict, Literal, Optional
+from typing import Any, Dict, List, Literal, Optional
 
 from opentelemetry import context as context_api
 from opentelemetry import trace
 from opentelemetry.trace import SpanKind, Status, StatusCode
 from opentelemetry.trace.propagation import set_span_in_context
+from pydantic import BaseModel
 
 from netra.config import Config
 
@@ -14,22 +16,22 @@ logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 
+class UsageModel(BaseModel):  # type: ignore[misc]
+    model: str
+    type: str
+    unit_used: Optional[int] = None
+    cost_in_usd: Optional[float] = None
+
+
 class ATTRIBUTE:
     LLM_SYSTEM = "llm_system"
     MODEL = "model"
     PROMPT = "prompt"
     NEGATIVE_PROMPT = "negative_prompt"
-    IMAGE_HEIGHT = "image_height"
-    IMAGE_WIDTH = "image_width"
-    CREDITS = "credits"
-    TOTAL_COST = "total_cost"
-    TOTAL_TOKENS = "total_tokens"
-    PROMPT_TOKENS_COST = "prompt_tokens_cost"
-    PROMPT_TOKENS = "prompt_tokens"
-    COMPLETION_TOKENS_COST = "completion_tokens_cost"
-    COMPLETION_TOKENS = "completion_tokens"
-    CACHED_TOKENS_COST = "cached_tokens_cost"
-    CACHED_TOKENS = "cached_tokens"
+    HEIGHT = "height"
+    WIDTH = "width"
+    OUTPUT_TYPE = "output_type"
+    USAGE = "usage"
     STATUS = "status"
     DURATION_MS = "duration_ms"
     ERROR_MESSAGE = "error_message"
@@ -46,7 +48,7 @@ class Session:
             # External API call
             result = external_api.generate_video(...)
 
-            session.set_tokens("20").set_credits("30")
+            session.set_usage(usage_data)
     """
 
     def __init__(self, name: str, attributes: Optional[Dict[str, str]] = None, module_name: str = "combat_sdk"):
@@ -139,49 +141,23 @@ class Session:
         """Set the negative prompt."""
         return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.NEGATIVE_PROMPT}", negative_prompt)
 
-    def set_image_height(self, height: str) -> "Session":
-        """Set the image height."""
-        return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.IMAGE_HEIGHT}", height)
-
-    def set_image_width(self, width: str) -> "Session":
-        """Set the image width."""
-        return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.IMAGE_WIDTH}", width)
-
-    def set_total_tokens(self, tokens: str) -> "Session":
-        """Set the number of tokens used."""
-        return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.TOTAL_TOKENS}", tokens)
-
-    def set_prompt_tokens_cost(self, cost: str) -> "Session":
-        """Set the number of tokens used."""
-        return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.PROMPT_TOKENS_COST}", cost)
-
-    def set_prompt_tokens(self, tokens: str) -> "Session":
-        """Set the number of tokens used."""
-        return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.PROMPT_TOKENS}", tokens)
-
-    def set_completion_tokens_cost(self, cost: str) -> "Session":
-        """Set the number of tokens used."""
-        return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.COMPLETION_TOKENS_COST}", cost)
-
-    def set_completion_tokens(self, tokens: str) -> "Session":
-        """Set the number of tokens used."""
-        return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.COMPLETION_TOKENS}", tokens)
-
-    def set_cached_tokens_cost(self, cost: str) -> "Session":
-        """Set the number of tokens used."""
-        return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.CACHED_TOKENS_COST}", cost)
+    def set_height(self, height: str) -> "Session":
+        """Set the height."""
+        return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.HEIGHT}", height)
 
-    def set_cached_tokens(self, tokens: str) -> "Session":
-        """Set the number of tokens used."""
-        return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.CACHED_TOKENS}", tokens)
+    def set_width(self, width: str) -> "Session":
+        """Set the width."""
+        return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.WIDTH}", width)
 
-    def set_credits(self, credits: str) -> "Session":
-        """Set the number of credits used."""
-        return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.CREDITS}", credits)
+    def set_output_type(self, output_type: str) -> "Session":
+        """Set the output type."""
+        return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.OUTPUT_TYPE}", output_type)
 
-    def set_total_cost(self, cost: str) -> "Session":
-        """Set the cost of the operation."""
-        return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.TOTAL_COST}", cost)
+    def set_usage(self, usage: List[UsageModel]) -> "Session":
+        """Set the usage data as a JSON string."""
+        usage_dict = [u.model_dump() for u in usage]
+        usage_json = json.dumps(usage_dict)
+        return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.USAGE}", usage_json)
 
     def set_model(self, model: str) -> "Session":
         """Set the model used."""

netra/session_manager.py

@@ -5,7 +5,7 @@ Handles automatic session and user ID management for applications.
 
 import logging
 from datetime import datetime
-from typing import Any, Dict, Optional, Union
+from typing import Any, Dict, List, Optional, Union
 
 from opentelemetry import baggage
 from opentelemetry import context as otel_context
@@ -22,6 +22,11 @@ class SessionManager:
     # Class variable to track the current span
     _current_span: Optional[trace.Span] = None
 
+    # Class variables to track separate entity stacks
+    _workflow_stack: List[str] = []
+    _task_stack: List[str] = []
+    _agent_stack: List[str] = []
+
     @classmethod
     def set_current_span(cls, span: Optional[trace.Span]) -> None:
         """
@@ -42,6 +47,86 @@ class SessionManager:
         """
         return cls._current_span
 
+    @classmethod
+    def push_entity(cls, entity_type: str, entity_name: str) -> None:
+        """
+        Push an entity onto the appropriate entity stack.
+
+        Args:
+            entity_type: Type of entity (workflow, task, agent)
+            entity_name: Name of the entity
+        """
+        if entity_type == "workflow":
+            cls._workflow_stack.append(entity_name)
+        elif entity_type == "task":
+            cls._task_stack.append(entity_name)
+        elif entity_type == "agent":
+            cls._agent_stack.append(entity_name)
+
+    @classmethod
+    def pop_entity(cls, entity_type: str) -> Optional[str]:
+        """
+        Pop the most recent entity from the specified entity stack.
+
+        Args:
+            entity_type: Type of entity (workflow, task, agent)
+
+        Returns:
+            Entity name or None if stack is empty
+        """
+        if entity_type == "workflow" and cls._workflow_stack:
+            return cls._workflow_stack.pop()
+        elif entity_type == "task" and cls._task_stack:
+            return cls._task_stack.pop()
+        elif entity_type == "agent" and cls._agent_stack:
+            return cls._agent_stack.pop()
+        return None
+
+    @classmethod
+    def get_current_entity_attributes(cls) -> Dict[str, str]:
+        """
+        Get current entity attributes for span annotation.
+
+        Returns:
+            Dictionary of entity attributes to add to spans
+        """
+        attributes = {}
+
+        # Add current workflow if exists
+        if cls._workflow_stack:
+            attributes[f"{Config.LIBRARY_NAME}.workflow.name"] = cls._workflow_stack[-1]
+
+        # Add current task if exists
+        if cls._task_stack:
+            attributes[f"{Config.LIBRARY_NAME}.task.name"] = cls._task_stack[-1]
+
+        # Add current agent if exists
+        if cls._agent_stack:
+            attributes[f"{Config.LIBRARY_NAME}.agent.name"] = cls._agent_stack[-1]
+
+        return attributes
+
+    @classmethod
+    def clear_entity_stacks(cls) -> None:
+        """Clear all entity stacks."""
+        cls._workflow_stack.clear()
+        cls._task_stack.clear()
+        cls._agent_stack.clear()
+
+    @classmethod
+    def get_stack_info(cls) -> Dict[str, List[str]]:
+        """
+        Get information about all current stacks.
+
+        Returns:
+            Dictionary containing all stack contents
+        """
+        return {
+            "workflows": cls._workflow_stack.copy(),
+            "tasks": cls._task_stack.copy(),
+            "agents": cls._agent_stack.copy(),
+        }
+
     @staticmethod
     def set_session_context(session_key: str, value: Union[str, Dict[str, str]]) -> None:
         """

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/version.py

@@ -1 +1 @@
-__version__ = "0.1.3"
+__version__ = "0.1.5"

{netra_sdk-0.1.3.dist-info → netra_sdk-0.1.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: netra-sdk
-Version: 0.1.3
+Version: 0.1.5
 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
@@ -416,17 +416,18 @@ Netra.init(app_name="My App")
 with Netra.start_session("video_generation_task") as session:
     # Set attributes before the API call
     session.set_prompt("A cat playing piano")
-    session.set_image_height("1024")
-    session.set_image_width("1024")
+    session.set_height("1024")
+    session.set_width("1024")
+    session.set_output_type("image")
     session.set_model("stable-diffusion-xl")
 
     # Make your external API call
     result = external_api.generate_video(...)
 
     # Set post-call attributes
-    session.set_tokens("1250")
+    session.set_completion_tokens("1250")
     session.set_credits("30")
-    session.set_cost("0.15")
+    session.set_completion_tokens_cost("0.15")
 
     # Add events during session
     session.add_event("processing_completed", {"step": "rendering"})
@@ -662,42 +663,7 @@ pre-commit install --hook-type pre-push
 
 ## 🤝 Contributing
 
-We welcome contributions! Please follow these guidelines:
-
-### Commit Message Format
-
-We use [Conventional Commits](https://www.conventionalcommits.org/) for commit messages:
-
-```
-<type>[optional scope]: <description>
-
-[optional body]
-
-[optional footer(s)]
-```
-
-**Types:**
-- **feat**: A new feature
-- **fix**: A bug fix
-- **docs**: Documentation only changes
-- **style**: Changes that do not affect the meaning of the code
-- **refactor**: A code change that neither fixes a bug nor adds a feature
-- **perf**: A code change that improves performance
-- **test**: Adding missing tests or correcting existing tests
-- **chore**: Changes to the build process or auxiliary tools
-
-**Examples:**
-```
-feat: add support for Claude AI instrumentation
-fix(pii): resolve masking issue with nested objects
-docs: update installation instructions
-```
-
-**Scope** can be used to specify the area of change (e.g., `pii`, `instrumentation`, `decorators`).
-
-**Body** should include the motivation for the change and contrast with previous behavior.
-
-**Footer** can be used for "BREAKING CHANGE:" or issue references.
+We welcome contributions! Please see our [Contributing Guide](CONTRIBUTING.md) for detailed information on how to contribute to the project, including development setup, testing, and our commit message format.
 
 ---
 

{netra_sdk-0.1.3.dist-info → netra_sdk-0.1.5.dist-info}/RECORD

@@ -4,12 +4,12 @@ netra/anonymizer/anonymizer.py,sha256=1VeYAsFpF_tYDlqJF-Q82-ZXGOR4YWBqrKUsRw3qOr
 netra/anonymizer/base.py,sha256=ytPxHCUD2OXlEY6fNTuMmwImNdIjgj294I41FIgoXpU,5946
 netra/anonymizer/fp_anonymizer.py,sha256=_6svIYmE0eejdIMkhKBUWCNjGtGimtrGtbLvPSOp8W4,6493
 netra/config.py,sha256=6aUG-BWWDOyptnTlEjOpB41n9xLmuo-XBAgs5GE5Tbg,4761
-netra/decorators.py,sha256=WX3EAvcZ_b-wUHO6UFXN07_J4H9shmrLSYPidYvTL5w,6257
+netra/decorators.py,sha256=V_WpZ2IgW2Y7B_WnSXmKUGGhkM5Cra2TwONddmJpPaI,6837
 netra/exceptions/__init__.py,sha256=uDgcBxmC4WhdS7HRYQk_TtJyxH1s1o6wZmcsnSHLAcM,174
 netra/exceptions/injection.py,sha256=ke4eUXRYUFJkMZgdSyPPkPt5PdxToTI6xLEBI0hTWUQ,1332
 netra/exceptions/pii.py,sha256=MT4p_x-zH3VtYudTSxw1Z9qQZADJDspq64WrYqSWlZc,2438
 netra/input_scanner.py,sha256=bzP3s7YudGHQrIbUgQGrcIBEJ6CmOewzuYNSu75cVXM,4988
-netra/instrumentation/__init__.py,sha256=s0PnnEZgxf-Yd0N6_AmpGYjQh_bINAIAUsX8nnSD9oU,39066
+netra/instrumentation/__init__.py,sha256=s-sXykQZ4CKUHLqHRR7buOrkN9hXGTZpNALRZkdIHB0,38757
 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
@@ -28,15 +28,15 @@ netra/instrumentation/mistralai/version.py,sha256=d6593s-XBNvVxri9lr2qLUDZQ3Zk3-
 netra/instrumentation/weaviate/__init__.py,sha256=EOlpWxobOLHYKqo_kMct_7nu26x1hr8qkeG5_h99wtg,4330
 netra/instrumentation/weaviate/version.py,sha256=PiCZHjonujPbnIn0KmD3Yl68hrjPRG_oKe5vJF3mmG8,24
 netra/pii.py,sha256=S7GnVzoNJEzKiUWnqN9bOCKPeNLsriztgB2E6Rx-yJU,27023
-netra/processors/__init__.py,sha256=f8Ck-uCdDwNgEvm3hm6I9wFSVlnRYjmSfrEIQkdBaZQ,218
-netra/processors/session_span_processor.py,sha256=uv-KL5lwil3C3wQGdYWiYQMHLBsXrt8hTy_ql6kUWXE,2171
-netra/processors/span_aggregation_processor.py,sha256=eOBAHTi5pgBCrnNthWGby0IY6b7iSUPAIFkDDmwLKOY,15579
+netra/processors/__init__.py,sha256=G16VumYTpgV4jsWrKNFSgm6xMQAsZ2Rrux25UVeo5YQ,215
+netra/processors/error_detection_processor.py,sha256=vlFQ8-Zgoj83htk_ciVhlT41RfU2GiGmDHGBzjKgtYQ,2701
+netra/processors/session_span_processor.py,sha256=qcsBl-LnILWefsftI8NQhXDGb94OWPc8LvzhVA0JS_c,2432
 netra/scanner.py,sha256=wqjMZnEbVvrGMiUSI352grUyHpkk94oBfHfMiXPhpGU,3866
-netra/session.py,sha256=Lsd7yps2YtjN7P10HsGN3bRZxufoeUnGxAH6Us8_l-Y,8734
-netra/session_manager.py,sha256=UusP3MRZlLeU4NtBVlXQ_sCgRg-LGleVdYPq5MwLvi8,3555
-netra/tracer.py,sha256=WZDBIpPp3oL8MkKBhlw0gxAXXSd1FY3bicLUKt2GNxc,3540
-netra/version.py,sha256=XEqb2aiIn8fzGE68Mph4ck1FtQqsR_am0wRWvrYPffQ,22
-netra_sdk-0.1.3.dist-info/LICENCE,sha256=8B_UoZ-BAl0AqiHAHUETCgd3I2B9yYJ1WEQtVb_qFMA,11359
-netra_sdk-0.1.3.dist-info/METADATA,sha256=7tXAkE-43Y10ksq8X6-VozfhZwkzuqmRnRqqiYyOaog,24349
-netra_sdk-0.1.3.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-netra_sdk-0.1.3.dist-info/RECORD,,
+netra/session.py,sha256=o1wXrPzMauqj_3P-iNBHVlcIR7zcKcbsmkrcHjQMKuY,7263
+netra/session_manager.py,sha256=EVcnWcSj4NdkH--HmqHx0mmzivQiM4GCyFLu6lwi33M,6252
+netra/tracer.py,sha256=9jAKdIHXbaZ6WV_p8I1syQiMdqXVCXMhpEhCBsbbci8,3538
+netra/version.py,sha256=rPSfWgIeq2YWVPyESOAwCBt8vftsTpIkuLAGDEzyRQc,22
+netra_sdk-0.1.5.dist-info/LICENCE,sha256=8B_UoZ-BAl0AqiHAHUETCgd3I2B9yYJ1WEQtVb_qFMA,11359
+netra_sdk-0.1.5.dist-info/METADATA,sha256=-82nUeU3ogSqDhIpMzYAUkjgS5emC5yy_jA_WTkpsPU,23518
+netra_sdk-0.1.5.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+netra_sdk-0.1.5.dist-info/RECORD,,

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