netra-sdk 0.1.0__py3-none-any.whl

netra/__init__.py

@@ -0,0 +1,148 @@
+import logging
+import threading
+from typing import Any, Dict, Optional, Set
+
+from netra.instrumentation.instruments import NetraInstruments
+
+from .config import Config
+
+# Instrumentor functions
+from .instrumentation import init_instrumentations
+from .session import Session
+from .session_manager import SessionManager
+from .tracer import Tracer
+
+logger = logging.getLogger(__name__)
+
+
+class Netra:
+    """
+    Main SDK class. Call SDK.init(...) at the start of your application
+    to configure OpenTelemetry and enable all built-in LLM + VectorDB instrumentations.
+    """
+
+    _initialized = False
+    # Use RLock so the thread that already owns the lock can re-acquire it safely
+    _init_lock = threading.RLock()
+
+    @classmethod
+    def is_initialized(cls) -> bool:
+        """Thread-safe check if Netra has been initialized.
+
+        Returns:
+            bool: True if Netra has been initialized, False otherwise
+        """
+        with cls._init_lock:
+            return cls._initialized
+
+    @classmethod
+    def init(
+        cls,
+        app_name: Optional[str] = None,
+        headers: Optional[str] = None,
+        disable_batch: Optional[bool] = None,
+        trace_content: Optional[bool] = None,
+        resource_attributes: Optional[Dict[str, Any]] = None,
+        environment: Optional[str] = None,
+        instruments: Optional[Set[NetraInstruments]] = None,
+        block_instruments: Optional[Set[NetraInstruments]] = None,
+    ) -> None:
+        # Acquire lock at the start of the method and hold it throughout
+        # to prevent race conditions during initialization
+        with cls._init_lock:
+            # Check if already initialized while holding the lock
+            if cls._initialized:
+                logger.warning("Netra.init() called more than once; ignoring subsequent calls.")
+                return
+
+            # Build Config
+            cfg = Config(
+                app_name=app_name,
+                headers=headers,
+                disable_batch=disable_batch,
+                trace_content=trace_content,
+                resource_attributes=resource_attributes,
+                environment=environment,
+            )
+
+            # Initialize tracer (OTLP exporter, span processor, resource)
+            Tracer(cfg)
+
+            # Instrument all supported modules
+            #    Pass trace_content flag to instrumentors that can capture prompts/completions
+            init_instrumentations(
+                should_enrich_metrics=True,
+                base64_image_uploader=None,
+                instruments=instruments,
+                block_instruments=block_instruments,
+            )
+
+            cls._initialized = True
+            logger.info("Netra successfully initialized.")
+
+    @classmethod
+    def set_session_id(cls, session_id: str) -> None:
+        """
+        Set session_id context attributes in the current OpenTelemetry context.
+
+        Args:
+            session_id: Session identifier
+        """
+        SessionManager.set_session_context("session_id", session_id)
+
+    @classmethod
+    def set_user_id(cls, user_id: str) -> None:
+        """
+        Set user_id context attributes in the current OpenTelemetry context.
+
+        Args:
+            user_id: User identifier
+        """
+        SessionManager.set_session_context("user_id", user_id)
+
+    @classmethod
+    def set_tenant_id(cls, tenant_id: str) -> None:
+        """
+        Set user_account_id context attributes in the current OpenTelemetry context.
+
+        Args:
+            user_account_id: User account identifier
+        """
+        SessionManager.set_session_context("tenant_id", tenant_id)
+
+    @classmethod
+    def set_custom_attributes(cls, key: str, value: Any) -> None:
+        """
+        Set custom attributes context in the current OpenTelemetry context.
+
+        Args:
+            key: Custom attribute key
+            value: Custom attribute value
+        """
+        SessionManager.set_session_context("custom_attributes", {key: value})
+
+    @classmethod
+    def set_custom_event(cls, event_name: str, attributes: Any) -> None:
+        """
+        Set custom event in the current OpenTelemetry context.
+
+        Args:
+            event_name: Name of the custom event
+            attributes: Attributes of the custom event
+        """
+        SessionManager.set_custom_event(event_name, attributes)
+
+    @classmethod
+    def start_session(
+        cls,
+        name: str,
+        attributes: Optional[Dict[str, str]] = None,
+        module_name: str = "combat_sdk",
+    ) -> Session:
+        """
+        Start a new session.
+        """
+        return Session(name, attributes, module_name)
+
+
+__all__ = ["Netra"]

netra/anonymizer/__init__.py

@@ -0,0 +1,7 @@
+from .anonymizer import Anonymizer
+from .base import AnonymizationResult
+
+__all__ = [
+    "Anonymizer",
+    "AnonymizationResult",
+]

netra/anonymizer/anonymizer.py

@@ -0,0 +1,79 @@
+"""
+Custom anonymizer for PII data that provides consistent hashing of entities.
+
+This module provides a custom anonymizer that can be used to replace PII entities
+with consistent hash values, allowing for tracking the same entities across multiple
+texts while maintaining privacy.
+"""
+
+from typing import Callable, List, Optional
+
+from presidio_analyzer.recognizer_result import RecognizerResult
+
+from .base import AnonymizationResult, BaseAnonymizer
+from .fp_anonymizer import FormatPreservingEmailAnonymizer
+
+
+class Anonymizer:
+    """
+    Main anonymizer that delegates to different anonymizer classes based on entity type.
+
+    This anonymizer analyzes the entity types and uses appropriate anonymization
+    strategies - format-preserving for email addresses and hash-based for other types.
+    """
+
+    def __init__(self, hash_function: Optional[Callable[[str], str]] = None, cache_size: int = 1000):
+        """
+        Initialize the Anonymizer.
+
+        Args:
+            hash_function: Optional custom hash function that takes a string and returns a hash.
+                           If not provided, a default hash function will be used.
+            cache_size: Maximum number of entities to cache. Uses LRU eviction policy.
+                       Default is 1000. Set to 0 to disable caching.
+        """
+        # Initialize different anonymizer instances
+        self.base_anonymizer = BaseAnonymizer(hash_function=hash_function, cache_size=cache_size)
+        self.email_anonymizer = FormatPreservingEmailAnonymizer()
+
+    def anonymize(self, text: str, analyzer_results: List[RecognizerResult]) -> AnonymizationResult:
+        """
+        Anonymize text by replacing detected entities using appropriate anonymization strategies.
+
+        Args:
+            text: The original text containing PII.
+            analyzer_results: List of RecognizerResult objects from the Presidio analyzer.
+
+        Returns:
+            AnonymizationResult containing the masked text and a mapping of entity hashes to original values.
+        """
+        # Sort results by start index in descending order to avoid offset issues when replacing
+        sorted_results = sorted(analyzer_results, key=lambda x: x.start, reverse=True)
+
+        # Make a copy of the original text that we'll modify
+        masked_text = text
+
+        # Dictionary to store mapping of anonymized values to original entity values
+        entities_map = {}
+
+        # Replace each entity with its anonymized value
+        for result in sorted_results:
+            entity_type = result.entity_type
+            entity_value = text[result.start : result.end]
+
+            # Use appropriate anonymizer based on entity type
+            if entity_type.upper() in ["EMAIL", "EMAIL_ADDRESS"]:
+                # Use format-preserving email anonymization
+                anonymized_value = self.email_anonymizer._anonymize_email(entity_value)
+                placeholder = anonymized_value
+                entities_map[anonymized_value] = entity_value
+            else:
+                # Use base anonymizer for other entity types
+                entity_hash = self.base_anonymizer._get_entity_hash(entity_type, entity_value)
+                placeholder = f"<{entity_hash}>"
+                entities_map[entity_hash] = entity_value
+
+            # Replace the entity in the text with the placeholder
+            masked_text = masked_text[: result.start] + placeholder + masked_text[result.end :]
+
+        return AnonymizationResult(masked_text=masked_text, entities=entities_map)

netra/anonymizer/base.py

@@ -0,0 +1,159 @@
+"""
+Base anonymizer class for PII data anonymization.
+
+This module provides the base anonymizer class that contains the core anonymization
+logic that can be extended by specific anonymizer implementations.
+"""
+
+import hashlib
+from collections import OrderedDict
+from dataclasses import dataclass
+from typing import Callable, Dict, List, Optional
+
+from presidio_analyzer.recognizer_result import RecognizerResult
+
+
+@dataclass
+class AnonymizationResult:
+    """
+    Result of anonymization containing the masked text and entity mappings.
+
+    Attributes:
+        masked_text: The text with PII entities replaced by hash placeholders.
+        entities: Dictionary mapping entity hashes to their original values.
+    """
+
+    masked_text: str
+    entities: Dict[str, str]
+
+
+class BaseAnonymizer:
+    """
+    Base anonymizer that replaces entities with consistent hash values.
+
+    This base anonymizer provides the core anonymization logic that can be
+    extended by specific anonymizer implementations for different entity types.
+    """
+
+    def __init__(self, hash_function: Optional[Callable[[str], str]] = None, cache_size: int = 1000):
+        """
+        Initialize the BaseAnonymizer.
+
+        Args:
+            hash_function: Optional custom hash function that takes a string and returns a hash.
+                           If not provided, a default hash function will be used.
+            cache_size: Maximum number of entities to cache. Uses LRU eviction policy.
+                       Default is 1000. Set to 0 to disable caching.
+        """
+        self.hash_function = hash_function or self._default_hash_function
+        self.cache_size = cache_size
+
+        # Initialize LRU cache for entity hashes
+        if cache_size > 0:
+            self._entity_hash_cache: Optional[OrderedDict[str, str]] = OrderedDict()
+        else:
+            self._entity_hash_cache = None
+
+    def _default_hash_function(self, value: str) -> str:
+        """
+        Default hash function using SHA-256.
+
+        Args:
+            value: The string to hash.
+
+        Returns:
+            A hexadecimal hash string.
+        """
+        return hashlib.sha256(value.encode()).hexdigest()[:8]
+
+    def _get_entity_hash(self, entity_type: str, entity_value: str) -> str:
+        """
+        Get a consistent hash for an entity value, creating one if it doesn't exist.
+        Uses LRU cache with configurable size to balance performance and memory usage.
+
+        Args:
+            entity_type: The type of entity (e.g., 'EMAIL', 'PHONE', etc.)
+            entity_value: The original value of the entity.
+
+        Returns:
+            A hash string for the entity.
+        """
+        # Skip caching if cache_size is 0
+        if self.cache_size == 0:
+            entity_hash = f"{entity_type}_{self.hash_function(entity_value)}"
+            return entity_hash
+
+        # Create a composite key for the entity cache
+        cache_key = f"{entity_type}:{entity_value}"
+
+        # Check if entity exists in cache and move to end (mark as recently used)
+        if self._entity_hash_cache is not None and cache_key in self._entity_hash_cache:
+            # Move to end to mark as recently used
+            self._entity_hash_cache.move_to_end(cache_key)
+            return self._entity_hash_cache[cache_key]
+
+        # Generate a new hash for this entity
+        entity_hash = f"{entity_type}_{self.hash_function(entity_value)}"
+
+        # Add to cache if cache is enabled
+        if self._entity_hash_cache is not None:
+            self._entity_hash_cache[cache_key] = entity_hash
+
+            # Evict oldest entry if cache exceeds size limit
+            if len(self._entity_hash_cache) > self.cache_size:
+                # Remove the least recently used item (first item)
+                self._entity_hash_cache.popitem(last=False)
+
+        return entity_hash
+
+    def anonymize_entity(self, entity_type: str, entity_value: str) -> str:
+        """
+        Anonymize a single entity value.
+
+        Args:
+            entity_type: The type of entity (e.g., 'EMAIL', 'PHONE', etc.)
+            entity_value: The original value of the entity.
+
+        Returns:
+            The anonymized entity value.
+        """
+        # Get or create hash for this entity
+        entity_hash = self._get_entity_hash(entity_type, entity_value)
+        return f"<{entity_hash}>"
+
+    def anonymize(self, text: str, analyzer_results: List[RecognizerResult]) -> AnonymizationResult:
+        """
+        Anonymize text by replacing detected entities with hash values.
+
+        Args:
+            text: The original text containing PII.
+            analyzer_results: List of RecognizerResult objects from the Presidio analyzer.
+
+        Returns:
+            AnonymizationResult containing the masked text and a mapping of entity hashes to original values.
+        """
+        # Sort results by start index in descending order to avoid offset issues when replacing
+        sorted_results = sorted(analyzer_results, key=lambda x: x.start, reverse=True)
+
+        # Make a copy of the original text that we'll modify
+        masked_text = text
+
+        # Dictionary to store mapping of hash values to original entity values
+        entities_map: Dict[str, str] = {}
+
+        # Replace each entity with its hash
+        for result in sorted_results:
+            entity_type = result.entity_type
+            entity_value = text[result.start : result.end]
+
+            # Get or create hash for this entity
+            entity_hash = self._get_entity_hash(entity_type, entity_value)
+
+            # Replace the entity in the text with the hash placeholder
+            placeholder = f"<{entity_hash}>"
+            masked_text = masked_text[: result.start] + placeholder + masked_text[result.end :]
+
+            # Store the mapping of hash to original value
+            entities_map[entity_hash] = entity_value
+
+        return AnonymizationResult(masked_text=masked_text, entities=entities_map)

netra/anonymizer/fp_anonymizer.py

@@ -0,0 +1,182 @@
+import hashlib
+import random
+import re
+from typing import Dict, Optional
+
+
+class FormatPreservingEmailAnonymizer:
+    def __init__(self, preserve_length: bool = True, preserve_structure: bool = True):
+        """
+        Initialize the email anonymizer.
+
+        Args:
+            preserve_length: Whether to preserve the length of original parts
+            preserve_structure: Whether to preserve dots, hyphens in the structure
+        """
+        self.preserve_length = preserve_length
+        self.preserve_structure = preserve_structure
+        self.email_cache: Dict[str, str] = {}
+        self.part_cache: Dict[str, str] = {}  # Cache for individual parts
+
+        # Character sets for replacement
+        self.alphanumeric = "abcdefghijklmnopqrstuvwxyz0123456789"
+        self.letters = "abcdefghijklmnopqrstuvwxyz"
+
+    def _get_deterministic_random(self, seed: str) -> random.Random:
+        """Create a deterministic random generator from a seed."""
+        # Use hash of the seed as random seed for consistency
+        hash_int = int(hashlib.md5(seed.encode()).hexdigest()[:8], 16)
+        return random.Random(hash_int)
+
+    def _preserve_structure_replace(self, text: str, seed: str) -> str:
+        """
+        Replace text while preserving structure (length, special chars, case pattern).
+        """
+        if text in self.part_cache:
+            return self.part_cache[text]
+
+        rng = self._get_deterministic_random(seed)
+        result = []
+
+        for char in text:
+            if char.isalpha():
+                # Preserve case pattern
+                new_char = rng.choice(self.letters)
+                result.append(new_char.upper() if char.isupper() else new_char)
+            elif char.isdigit():
+                result.append(str(rng.randint(0, 9)))
+            else:
+                # Keep special characters (dots, hyphens, etc.)
+                result.append(char)
+
+        anonymized = "".join(result)
+        self.part_cache[text] = anonymized
+        return anonymized
+
+    def _simple_hash_replace(self, text: str, target_length: Optional[int] = None) -> str:
+        """Simple hash replacement with optional length preservation."""
+        if target_length is None:
+            target_length = len(text)
+
+        hash_val = hashlib.md5(text.encode()).hexdigest()
+
+        # Create a mix of letters and numbers that looks more natural
+        result = []
+        for i in range(target_length):
+            if i < len(hash_val):
+                char = hash_val[i]
+                if char.isdigit():
+                    result.append(char)
+                else:
+                    # Convert hex chars to letters
+                    result.append(chr(ord("a") + (ord(char) - ord("a")) % 26))
+            else:
+                result.append("x")
+
+        return "".join(result)
+
+    def _anonymize_email(self, email: str) -> str:
+        """
+        Anonymize a single email while preserving format and structure.
+        """
+        if email in self.email_cache:
+            return self.email_cache[email]
+
+        # Split email into local part and domain
+        local_part, domain = email.split("@", 1)
+
+        if self.preserve_structure:
+            # Preserve the structure (dots, hyphens, length, case pattern)
+            local_anonymized = self._preserve_structure_replace(local_part, f"local_{local_part}")
+            domain_anonymized = self._preserve_structure_replace(domain, f"domain_{domain}")
+        else:
+            # Simple length-preserving hash
+            local_length = len(local_part) if self.preserve_length else 8
+            domain_length = len(domain) if self.preserve_length else 8
+
+            local_anonymized = self._simple_hash_replace(local_part, local_length)
+            domain_anonymized = self._simple_hash_replace(domain, domain_length)
+
+        anonymized_email = f"{local_anonymized}@{domain_anonymized}"
+        self.email_cache[email] = anonymized_email
+
+        return anonymized_email
+
+    def anonymize_text(self, text: str) -> str:
+        """
+        Anonymize all emails in the given text while preserving format.
+        """
+        email_pattern = r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b"
+
+        def replace_email(match: re.Match[str]) -> str:
+            email = match.group(0)
+            return self._anonymize_email(email)
+
+        return re.sub(email_pattern, replace_email, text)
+
+    def get_mapping(self) -> Dict[str, str]:
+        """Return the mapping of original emails to anonymized versions."""
+        return self.email_cache.copy()
+
+
+# Example usage and comparison
+if __name__ == "__main__":
+    print("Format-Preserving Email Anonymization - Structure Preserving:\n")
+
+    # Structure-preserving anonymizer
+    anonymizer1 = FormatPreservingEmailAnonymizer(preserve_structure=True)
+
+    test_emails = [
+        "john@gmail.com",
+        "john@gmail.com",
+        "john@outlook.com",
+        "joe@outlook.com",
+        "user.name@company.co.uk",
+        "test-email@sub.example.org",
+        "Admin123@BigCorp.net",
+        "a@b.co",
+    ]
+
+    print("Structure-Preserving Mode:")
+    print("=" * 50)
+    for email in test_emails:
+        anonymized = anonymizer1._anonymize_email(email)
+        print(f"{email:25} -> {anonymized}")
+
+    print("\n" + "=" * 50)
+    print("Length-Preserving Mode:")
+    print("=" * 50)
+
+    # Length-preserving but simpler anonymizer
+    anonymizer2 = FormatPreservingEmailAnonymizer(preserve_length=True, preserve_structure=False)
+
+    for email in test_emails:
+        anonymized = anonymizer2._anonymize_email(email)
+        print(f"{email:25} -> {anonymized}")
+
+    # Test with full text
+    print("\n" + "=" * 70)
+    print("Full Text Anonymization Examples:")
+    print("=" * 70)
+
+    test_texts = [
+        "Hi, my name is John and my email is john@gmail.com",
+        "Contact: support@company.com or admin@BigCorp.net",
+        "Emails: user.name@test.co.uk, simple@domain.org",
+    ]
+
+    for text in test_texts:
+        anonymized = anonymizer1.anonymize_text(text)
+        print(f"Original:   {text}")
+        print(f"Anonymized: {anonymized}")
+        print()
+
+    # Consistency test
+    print("Consistency Test:")
+    print("-" * 30)
+    email = "john@gmail.com"
+    result1 = anonymizer1._anonymize_email(email)
+    result2 = anonymizer1._anonymize_email(email)
+    print(f"First call:  {email} -> {result1}")
+    print(f"Second call: {email} -> {result2}")
+    print(f"Consistent: {'✓' if result1 == result2 else '✗'}")

netra/config.py

@@ -0,0 +1,111 @@
+import json
+import os
+from typing import Any, Dict, Optional
+
+from opentelemetry.util.re import parse_env_headers
+
+from netra.version import __version__
+
+
+class Config:
+    """
+    Holds configuration options for the tracer:
+      - app_name:                Logical name for this service
+      - otlp_endpoint:           URL for OTLP collector
+      - api_key:                 API key for the collector (sent as Bearer token)
+      - headers:                 Additional headers (W3C Correlation-Context format)
+      - disable_batch:           Whether to disable batch span processor (bool)
+      - trace_content:           Whether to capture prompt/completion content (bool)
+      - resource_attributes:     Custom resource attributes dict (e.g., {'env': 'prod', 'version': '1.0.0'})
+    """
+
+    # SDK Constants
+    SDK_NAME = "netra"
+    LIBRARY_NAME = "netra"
+    LIBRARY_VERSION = __version__
+
+    def __init__(
+        self,
+        app_name: Optional[str] = None,
+        headers: Optional[str] = None,
+        disable_batch: Optional[bool] = None,
+        trace_content: Optional[bool] = None,
+        resource_attributes: Optional[Dict[str, Any]] = None,
+        environment: Optional[str] = None,
+    ):
+        # Application name: from param, else env
+        self.app_name = (
+            app_name or os.getenv("OTEL_SERVICE_NAME") or os.getenv("NETRA_APP_NAME") or "llm_tracing_service"
+        )
+
+        # OTLP endpoint: if explicit param, else OTEL_EXPORTER_OTLP_ENDPOINT
+        self.otlp_endpoint = os.getenv("NETRA_OTLP_ENDPOINT") or os.getenv("OTEL_EXPORTER_OTLP_ENDPOINT")
+
+        # API key: if explicit param, else env NETRA_API_KEY
+        self.api_key = os.getenv("NETRA_API_KEY")
+        self.headers = {}
+
+        # Custom headers: comma-separated W3C format (if provided, overrides API key)
+        headers = headers or os.getenv("NETRA_HEADERS")
+
+        if isinstance(headers, str):
+            self.headers = parse_env_headers(headers)
+
+        if self.otlp_endpoint == "https://api.dev.getcombat.ai" and not self.api_key:
+            print("Error: Missing Netra API key, go to https://app.dev.getcombat.ai/api-key to create one")
+            print("Set the NETRA_API_KEY environment variable to the key")
+            return
+
+        # Handle API key authentication based on OTLP endpoint
+        if self.api_key and self.otlp_endpoint:
+            # For Netra endpoints, use x-api-key header
+            if "getcombat" in self.otlp_endpoint.lower():
+                if not self.headers:
+                    self.headers = {"x-api-key": self.api_key}
+                elif "x-api-key" not in self.headers:
+                    self.headers = {**self.headers, "x-api-key": self.api_key}
+            # For other endpoints, set up basic auth
+            else:
+                if not self.headers:
+                    self.headers = {"Authorization": f"Bearer {self.api_key}"}
+                elif "Authorization" not in self.headers:
+                    self.headers = {**self.headers, "Authorization": f"Bearer {self.api_key}"}
+
+        # Disable batch span processor?
+        if disable_batch is not None:
+            self.disable_batch = disable_batch
+        else:
+            # Environment var can be "true"/"false"
+            env_db = os.getenv("NETRA_DISABLE_BATCH")
+            self.disable_batch = True if (env_db is not None and env_db.lower() in ("1", "true")) else False
+
+        # Trace content (prompts/completions)? Default true unless env says false
+        if trace_content is not None:
+            self.trace_content = trace_content
+        else:
+            env_tc = os.getenv("NETRA_TRACE_CONTENT")
+            self.trace_content = False if (env_tc is not None and env_tc.lower() in ("0", "false")) else True
+
+        # 7. Environment: param override, else env
+        if environment is not None:
+            self.environment = environment
+        else:
+            self.environment = os.getenv("NETRA_ENV", "local")
+
+        # Resource attributes: param override, else parse JSON from env, else empty dict
+        if resource_attributes is not None:
+            self.resource_attributes = resource_attributes
+        else:
+            # Expecting something like: {"env":"prod","version":"1.0.0"}
+            env_ra = os.getenv("NETRA_RESOURCE_ATTRS")
+            if env_ra:
+                try:
+                    self.resource_attributes = json.loads(env_ra)
+                except (json.JSONDecodeError, ValueError) as e:
+                    import logging
+
+                    logger = logging.getLogger(__name__)
+                    logger.warning(f"Failed to parse NETRA_RESOURCE_ATTRS: {e}")
+                    self.resource_attributes = {}
+            else:
+                self.resource_attributes = {}