traccia 0.1.2__py3-none-any.whl → 0.1.6__py3-none-any.whl

traccia/config.py

@@ -0,0 +1,699 @@
+"""Configuration management with Pydantic models and validation."""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Any, Dict, Optional, Literal
+from pydantic import BaseModel, Field, field_validator, model_validator, HttpUrl
+
+from traccia.errors import ConfigError, ValidationError
+
+
+# Environment variable mapping
+ENV_VAR_MAPPING = {
+    # Tracing config
+    "api_key": ["TRACCIA_API_KEY", "AGENT_DASHBOARD_API_KEY"],
+    "endpoint": ["TRACCIA_ENDPOINT", "AGENT_DASHBOARD_ENDPOINT"],
+    "sample_rate": ["TRACCIA_SAMPLE_RATE", "AGENT_DASHBOARD_SAMPLE_RATE"],
+    "auto_start_trace": ["TRACCIA_AUTO_START_TRACE", "AGENT_DASHBOARD_AUTO_START_TRACE"],
+    "auto_trace_name": ["TRACCIA_AUTO_TRACE_NAME"],
+    "use_otlp": ["TRACCIA_USE_OTLP"],
+    "service_name": ["TRACCIA_SERVICE_NAME"],
+    
+    # Exporter config
+    "enable_console": ["TRACCIA_ENABLE_CONSOLE", "AGENT_DASHBOARD_ENABLE_CONSOLE_EXPORTER"],
+    "enable_file": ["TRACCIA_ENABLE_FILE", "AGENT_DASHBOARD_ENABLE_FILE_EXPORTER"],
+    "file_exporter_path": ["TRACCIA_FILE_PATH"],
+    "reset_trace_file": ["TRACCIA_RESET_TRACE_FILE"],
+    
+    # Instrumentation config
+    "enable_patching": ["TRACCIA_ENABLE_PATCHING", "AGENT_DASHBOARD_ENABLE_PATCHING"],
+    "enable_token_counting": ["TRACCIA_ENABLE_TOKEN_COUNTING", "AGENT_DASHBOARD_ENABLE_TOKEN_COUNTING"],
+    "enable_costs": ["TRACCIA_ENABLE_COSTS", "AGENT_DASHBOARD_ENABLE_COSTS"],
+    "auto_instrument_tools": ["TRACCIA_AUTO_INSTRUMENT_TOOLS"],
+    "max_tool_spans": ["TRACCIA_MAX_TOOL_SPANS"],
+    "max_span_depth": ["TRACCIA_MAX_SPAN_DEPTH"],
+    "openai_agents": ["TRACCIA_OPENAI_AGENTS"],
+    
+    # Rate limiting & Batching
+    "max_spans_per_second": ["TRACCIA_MAX_SPANS_PER_SECOND"],
+    "max_queue_size": ["TRACCIA_MAX_QUEUE_SIZE"],
+    "max_block_ms": ["TRACCIA_MAX_BLOCK_MS"],
+    "max_export_batch_size": ["TRACCIA_MAX_EXPORT_BATCH_SIZE"],
+    "schedule_delay_millis": ["TRACCIA_SCHEDULE_DELAY_MILLIS"],
+    
+    # Runtime metadata
+    "session_id": ["TRACCIA_SESSION_ID"],
+    "user_id": ["TRACCIA_USER_ID"],
+    "tenant_id": ["TRACCIA_TENANT_ID"],
+    "project_id": ["TRACCIA_PROJECT_ID"],
+    "agent_id": ["TRACCIA_AGENT_ID", "AGENT_DASHBOARD_AGENT_ID"],
+    
+    # Logging
+    "debug": ["TRACCIA_DEBUG"],
+    "enable_span_logging": ["TRACCIA_ENABLE_SPAN_LOGGING"],
+    
+    # Advanced
+    "attr_truncation_limit": ["TRACCIA_ATTR_TRUNCATION_LIMIT"],
+}
+
+
+class TracingConfig(BaseModel):
+    """Tracing configuration section."""
+    
+    api_key: Optional[str] = Field(
+        default=None,
+        description="API key for authentication (required for SaaS, optional for open-source)"
+    )
+    endpoint: Optional[str] = Field(
+        default=None,
+        description="Endpoint URL for trace ingestion"
+    )
+    sample_rate: float = Field(
+        default=1.0,
+        ge=0.0,
+        le=1.0,
+        description="Sampling rate (0.0 to 1.0)"
+    )
+    auto_start_trace: bool = Field(
+        default=True,
+        description="Automatically start a root trace on init"
+    )
+    auto_trace_name: str = Field(
+        default="root",
+        description="Name for the auto-started root trace"
+    )
+    use_otlp: bool = Field(
+        default=True,
+        description="Use OTLP exporter (set to false for console/file exporters)"
+    )
+    service_name: Optional[str] = Field(
+        default=None,
+        description="Service name for the application"
+    )
+
+
+class ExporterConfig(BaseModel):
+    """Exporter configuration section."""
+    
+    enable_console: bool = Field(
+        default=False,
+        description="Enable console exporter for debugging"
+    )
+    enable_file: bool = Field(
+        default=False,
+        description="Enable file exporter to write traces to local file"
+    )
+    file_exporter_path: str = Field(
+        default="traces.jsonl",
+        description="File path for file exporter"
+    )
+    reset_trace_file: bool = Field(
+        default=False,
+        description="Reset/clear trace file on initialization"
+    )
+    
+    @model_validator(mode='after')
+    def check_single_exporter(self) -> 'ExporterConfig':
+        """Ensure only one exporter is enabled at a time."""
+        enabled_count = sum([
+            self.enable_console,
+            self.enable_file,
+        ])
+        if enabled_count > 1:
+            raise ValidationError(
+                "Only one exporter can be enabled at a time. "
+                "Choose either console or file exporter (OTLP is controlled by use_otlp in tracing section).",
+                details={
+                    "enable_console": self.enable_console,
+                    "enable_file": self.enable_file,
+                }
+            )
+        return self
+
+
+class InstrumentationConfig(BaseModel):
+    """Instrumentation configuration section."""
+    
+    enable_patching: bool = Field(
+        default=True,
+        description="Auto-patch popular libraries (OpenAI, Anthropic, requests)"
+    )
+    enable_token_counting: bool = Field(
+        default=True,
+        description="Count tokens for LLM calls"
+    )
+    enable_costs: bool = Field(
+        default=True,
+        description="Calculate costs for LLM calls"
+    )
+    auto_instrument_tools: bool = Field(
+        default=False,
+        description="Automatically instrument tool calls"
+    )
+    max_tool_spans: int = Field(
+        default=100,
+        gt=0,
+        description="Maximum number of tool spans to create"
+    )
+    max_span_depth: int = Field(
+        default=10,
+        gt=0,
+        description="Maximum depth of nested spans"
+    )
+    openai_agents: bool = Field(
+        default=True,
+        description="Auto-install OpenAI Agents SDK integration when available"
+    )
+
+
+class RateLimitConfig(BaseModel):
+    """Rate limiting and batching configuration section."""
+    
+    max_spans_per_second: Optional[float] = Field(
+        default=None,
+        gt=0,
+        description="Maximum spans per second (None = unlimited)"
+    )
+    max_queue_size: int = Field(
+        default=5000,
+        gt=0,
+        description="Maximum queue size for buffered spans"
+    )
+    max_block_ms: int = Field(
+        default=100,
+        ge=0,
+        description="Maximum milliseconds to block before dropping spans"
+    )
+    max_export_batch_size: int = Field(
+        default=512,
+        gt=0,
+        description="Maximum number of spans in a single export batch"
+    )
+    schedule_delay_millis: int = Field(
+        default=5000,
+        gt=0,
+        description="Delay in milliseconds between export batches"
+    )
+
+
+class LoggingConfig(BaseModel):
+    """Logging configuration section."""
+    
+    debug: bool = Field(
+        default=False,
+        description="Enable debug logging"
+    )
+    enable_span_logging: bool = Field(
+        default=False,
+        description="Enable span-level logging for debugging"
+    )
+
+
+class RuntimeConfig(BaseModel):
+    """Runtime metadata configuration section."""
+    
+    session_id: Optional[str] = Field(
+        default=None,
+        description="Session identifier for grouping traces"
+    )
+    user_id: Optional[str] = Field(
+        default=None,
+        description="User identifier for the current session"
+    )
+    tenant_id: Optional[str] = Field(
+        default=None,
+        description="Tenant/organization identifier"
+    )
+    project_id: Optional[str] = Field(
+        default=None,
+        description="Project identifier"
+    )
+    agent_id: Optional[str] = Field(
+        default=None,
+        description="Agent identifier for the current session"
+    )
+
+
+class AdvancedConfig(BaseModel):
+    """Advanced configuration options."""
+    
+    attr_truncation_limit: Optional[int] = Field(
+        default=None,
+        gt=0,
+        description="Maximum length for attribute values (None = no limit)"
+    )
+
+
+class TracciaConfig(BaseModel):
+    """
+    Complete Traccia SDK configuration.
+    
+    This model validates and merges configuration from multiple sources:
+    1. Config file (traccia.toml)
+    2. Environment variables
+    3. Explicit parameters
+    """
+    
+    tracing: TracingConfig = Field(default_factory=TracingConfig)
+    exporters: ExporterConfig = Field(default_factory=ExporterConfig)
+    instrumentation: InstrumentationConfig = Field(default_factory=InstrumentationConfig)
+    rate_limiting: RateLimitConfig = Field(default_factory=RateLimitConfig)
+    runtime: RuntimeConfig = Field(default_factory=RuntimeConfig)
+    logging: LoggingConfig = Field(default_factory=LoggingConfig)
+    advanced: AdvancedConfig = Field(default_factory=AdvancedConfig)
+    
+    @model_validator(mode='after')
+    def validate_complete_config(self) -> 'TracciaConfig':
+        """Validate the complete configuration for conflicts."""
+        # If OTLP is disabled, at least one other exporter must be enabled
+        if not self.tracing.use_otlp:
+            if not (self.exporters.enable_console or self.exporters.enable_file):
+                raise ConfigError(
+                    "When use_otlp is false, you must enable either console or file exporter.",
+                    details={
+                        "use_otlp": self.tracing.use_otlp,
+                        "enable_console": self.exporters.enable_console,
+                        "enable_file": self.exporters.enable_file,
+                    }
+                )
+        
+        return self
+    
+    def to_flat_dict(self) -> Dict[str, Any]:
+        """Convert to flat dictionary for backward compatibility."""
+        return {
+            # Tracing
+            "api_key": self.tracing.api_key,
+            "endpoint": self.tracing.endpoint,
+            "sample_rate": self.tracing.sample_rate,
+            "auto_start_trace": self.tracing.auto_start_trace,
+            "auto_trace_name": self.tracing.auto_trace_name,
+            "use_otlp": self.tracing.use_otlp,
+            "service_name": self.tracing.service_name,
+            # Exporters
+            "enable_console": self.exporters.enable_console,
+            "enable_file": self.exporters.enable_file,
+            "file_exporter_path": self.exporters.file_exporter_path,
+            "reset_trace_file": self.exporters.reset_trace_file,
+            # Instrumentation
+            "enable_patching": self.instrumentation.enable_patching,
+            "enable_token_counting": self.instrumentation.enable_token_counting,
+            "enable_costs": self.instrumentation.enable_costs,
+            "auto_instrument_tools": self.instrumentation.auto_instrument_tools,
+            "max_tool_spans": self.instrumentation.max_tool_spans,
+            "max_span_depth": self.instrumentation.max_span_depth,
+            "openai_agents": self.instrumentation.openai_agents,
+            # Rate limiting & Batching
+            "max_spans_per_second": self.rate_limiting.max_spans_per_second,
+            "max_queue_size": self.rate_limiting.max_queue_size,
+            "max_block_ms": self.rate_limiting.max_block_ms,
+            "max_export_batch_size": self.rate_limiting.max_export_batch_size,
+            "schedule_delay_millis": self.rate_limiting.schedule_delay_millis,
+            # Runtime
+            "session_id": self.runtime.session_id,
+            "user_id": self.runtime.user_id,
+            "tenant_id": self.runtime.tenant_id,
+            "project_id": self.runtime.project_id,
+            "agent_id": self.runtime.agent_id,
+            # Logging
+            "debug": self.logging.debug,
+            "enable_span_logging": self.logging.enable_span_logging,
+            # Advanced
+            "attr_truncation_limit": self.advanced.attr_truncation_limit,
+        }
+
+
+def load_dotenv(path: str = ".env") -> None:
+    """Minimal .env loader (no external dependency)."""
+    if not os.path.exists(path):
+        return
+    try:
+        with open(path, "r", encoding="utf-8") as f:
+            for line in f:
+                line = line.strip()
+                if not line or line.startswith("#"):
+                    continue
+                if "=" not in line:
+                    continue
+                key, value = line.split("=", 1)
+                key, value = key.strip(), value.strip().strip("\"'")
+                if key and key not in os.environ:
+                    os.environ[key] = value
+    except Exception:
+        # Fail silently; this loader is best-effort.
+        return
+
+
+def find_config_file() -> Optional[str]:
+    """
+    Find traccia.toml config file in standard locations.
+    
+    Lookup order:
+    1. ./traccia.toml (current directory)
+    2. ~/.traccia/config.toml (user home)
+    
+    Returns:
+        Path to config file if found, None otherwise
+    """
+    # Check current directory
+    cwd_config = Path.cwd() / "traccia.toml"
+    if cwd_config.exists():
+        return str(cwd_config)
+    
+    # Check user home directory
+    home_config = Path.home() / ".traccia" / "config.toml"
+    if home_config.exists():
+        return str(home_config)
+    
+    return None
+
+
+def load_toml_config(path: str) -> Dict[str, Any]:
+    """
+    Load configuration from a TOML file.
+    
+    Args:
+        path: Path to the TOML config file
+        
+    Returns:
+        Dictionary with nested config structure
+    """
+    if not os.path.exists(path):
+        return {}
+    
+    try:
+        # Try to import tomli for Python 3.11+, fall back to toml
+        try:
+            import tomli as toml_lib
+            with open(path, "rb") as f:
+                data = toml_lib.load(f)
+        except ImportError:
+            try:
+                import toml as toml_lib  # type: ignore
+                with open(path, "r", encoding="utf-8") as f:
+                    data = toml_lib.load(f)
+            except ImportError:
+                raise ConfigError(
+                    "No TOML library available. Install tomli or toml: pip install tomli"
+                )
+        
+        return data
+    
+    except Exception as e:
+        if isinstance(e, ConfigError):
+            raise
+        raise ConfigError(f"Failed to load config file: {e}")
+
+
+def get_env_value(config_key: str) -> Optional[str]:
+    """
+    Get environment variable value for a config key.
+    
+    Tries multiple environment variable names in order of preference.
+    """
+    env_vars = ENV_VAR_MAPPING.get(config_key, [])
+    for env_var in env_vars:
+        value = os.getenv(env_var)
+        if value is not None:
+            return value
+    return None
+
+
+def load_config_from_env(flat: bool = False) -> Dict[str, Any]:
+    """
+    Load configuration from environment variables.
+    
+    Args:
+        flat: If True, return flat dictionary for backward compatibility
+    
+    Returns:
+        Dictionary of config values from environment (nested structure by default)
+    """
+    env_config = {
+        "tracing": {},
+        "exporters": {},
+        "instrumentation": {},
+        "rate_limiting": {},
+        "runtime": {},
+        "logging": {},
+        "advanced": {},
+    }
+    
+    # Tracing section
+    for key in ["api_key", "endpoint", "sample_rate", "auto_start_trace", "auto_trace_name", "use_otlp", "service_name"]:
+        value = get_env_value(key)
+        if value is not None:
+            # Convert string to appropriate type
+            if key in ["auto_start_trace", "use_otlp"]:
+                env_config["tracing"][key] = value.lower() in ("true", "1", "yes")
+            elif key == "sample_rate":
+                try:
+                    env_config["tracing"][key] = float(value)
+                except ValueError:
+                    raise ConfigError(f"Invalid sample_rate value: {value}. Must be a float between 0.0 and 1.0.")
+            else:
+                env_config["tracing"][key] = value
+    
+    # Exporters section
+    for key in ["enable_console", "enable_file", "file_exporter_path", "reset_trace_file"]:
+        value = get_env_value(key)
+        if value is not None:
+            if key in ["enable_console", "enable_file", "reset_trace_file"]:
+                env_config["exporters"][key] = value.lower() in ("true", "1", "yes")
+            else:
+                env_config["exporters"][key] = value
+    
+    # Instrumentation section
+    for key in ["enable_patching", "enable_token_counting", "enable_costs", "auto_instrument_tools", "openai_agents"]:
+        value = get_env_value(key)
+        if value is not None:
+            env_config["instrumentation"][key] = value.lower() in ("true", "1", "yes")
+    
+    for key in ["max_tool_spans", "max_span_depth"]:
+        value = get_env_value(key)
+        if value is not None:
+            try:
+                env_config["instrumentation"][key] = int(value)
+            except ValueError:
+                raise ConfigError(f"Invalid {key} value: {value}. Must be an integer.")
+    
+    # Rate limiting section
+    for key in ["max_spans_per_second"]:
+        value = get_env_value(key)
+        if value is not None:
+            try:
+                env_config["rate_limiting"][key] = float(value) if value else None
+            except ValueError:
+                raise ConfigError(f"Invalid {key} value: {value}. Must be a number.")
+    
+    for key in ["max_queue_size", "max_block_ms", "max_export_batch_size", "schedule_delay_millis"]:
+        value = get_env_value(key)
+        if value is not None:
+            try:
+                env_config["rate_limiting"][key] = int(value)
+            except ValueError:
+                raise ConfigError(f"Invalid {key} value: {value}. Must be a number.")
+    
+    # Runtime section
+    for key in ["session_id", "user_id", "tenant_id", "project_id", "agent_id"]:
+        value = get_env_value(key)
+        if value is not None:
+            env_config["runtime"][key] = value
+    
+    # Logging section
+    for key in ["debug", "enable_span_logging"]:
+        value = get_env_value(key)
+        if value is not None:
+            env_config["logging"][key] = value.lower() in ("true", "1", "yes")
+    
+    # Advanced section
+    value = get_env_value("attr_truncation_limit")
+    if value is not None:
+        try:
+            env_config["advanced"]["attr_truncation_limit"] = int(value)
+        except ValueError:
+            raise ConfigError(f"Invalid attr_truncation_limit value: {value}. Must be an integer.")
+    
+    # Remove empty sections
+    nested_result = {k: v for k, v in env_config.items() if v}
+    
+    # Flatten if requested (for backward compatibility)
+    if flat:
+        flat_result = {}
+        for section, values in nested_result.items():
+            flat_result.update(values)
+        return flat_result
+    
+    return nested_result
+
+
+def merge_configs(base: Dict[str, Any], override: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Deep merge two config dictionaries.
+    
+    Args:
+        base: Base configuration
+        override: Override configuration (takes precedence)
+        
+    Returns:
+        Merged configuration
+    """
+    result = base.copy()
+    for key, value in override.items():
+        if key in result and isinstance(result[key], dict) and isinstance(value, dict):
+            result[key] = merge_configs(result[key], value)
+        else:
+            result[key] = value
+    return result
+
+
+def load_config(
+    config_file: Optional[str] = None,
+    overrides: Optional[Dict[str, Any]] = None
+) -> TracciaConfig:
+    """
+    Load and validate Traccia configuration from multiple sources.
+    
+    Priority (highest to lowest):
+    1. Explicit overrides (passed as parameters)
+    2. Environment variables
+    3. Config file (./traccia.toml or ~/.traccia/config.toml)
+    4. Defaults
+    
+    Args:
+        config_file: Optional explicit path to config file
+        overrides: Optional dict of explicit parameter overrides
+        
+    Returns:
+        Validated TracciaConfig instance
+        
+    Raises:
+        ConfigError: If configuration is invalid or conflicting
+    """
+    merged_config: Dict[str, Any] = {}
+    
+    # 1. Load from config file (lowest priority)
+    if config_file:
+        file_config = load_toml_config(config_file)
+        merged_config = merge_configs(merged_config, file_config)
+    else:
+        # Try to find config file automatically
+        found_config = find_config_file()
+        if found_config:
+            file_config = load_toml_config(found_config)
+            merged_config = merge_configs(merged_config, file_config)
+    
+    # 2. Override with environment variables (medium priority)
+    env_config = load_config_from_env()
+    merged_config = merge_configs(merged_config, env_config)
+    
+    # 3. Override with explicit parameters (highest priority)
+    if overrides:
+        merged_config = merge_configs(merged_config, overrides)
+    
+    # 4. Create and validate Pydantic model
+    try:
+        return TracciaConfig(**merged_config)
+    except Exception as e:
+        raise ConfigError(f"Configuration validation failed: {e}")
+
+
+def validate_config(
+    config_file: Optional[str] = None,
+    overrides: Optional[Dict[str, Any]] = None
+) -> tuple[bool, str, Optional[TracciaConfig]]:
+    """
+    Validate configuration without loading it.
+    
+    Used by the `traccia doctor` CLI command.
+    
+    Args:
+        config_file: Optional explicit path to config file
+        overrides: Optional dict of explicit parameter overrides
+        
+    Returns:
+        Tuple of (is_valid, message, config_or_none)
+    """
+    try:
+        config = load_config(config_file=config_file, overrides=overrides)
+        return True, "Configuration is valid", config
+    except ConfigError as e:
+        return False, f"Configuration error: {e}", None
+    except Exception as e:
+        return False, f"Unexpected error: {e}", None
+
+
+# Backward compatibility functions
+def load_config_with_priority(
+    config_file: Optional[str] = None,
+    overrides: Optional[Dict[str, Any]] = None
+) -> Dict[str, Any]:
+    """
+    Legacy function for backward compatibility.
+    
+    Returns flattened config dictionary instead of Pydantic model.
+    Accepts flat overrides and converts them to nested format.
+    """
+    # Convert flat overrides to nested format
+    nested_overrides = None
+    if overrides:
+        nested_overrides = {
+            "tracing": {}, 
+            "exporters": {}, 
+            "instrumentation": {}, 
+            "rate_limiting": {},
+            "runtime": {},
+            "logging": {},
+            "advanced": {}
+        }
+        
+        # Map flat keys to nested structure
+        flat_to_nested = {
+            # Tracing
+            "api_key": ("tracing", "api_key"),
+            "endpoint": ("tracing", "endpoint"),
+            "sample_rate": ("tracing", "sample_rate"),
+            "auto_start_trace": ("tracing", "auto_start_trace"),
+            "auto_trace_name": ("tracing", "auto_trace_name"),
+            "use_otlp": ("tracing", "use_otlp"),
+            "service_name": ("tracing", "service_name"),
+            # Exporters
+            "enable_console": ("exporters", "enable_console"),
+            "enable_file": ("exporters", "enable_file"),
+            "file_exporter_path": ("exporters", "file_exporter_path"),
+            "reset_trace_file": ("exporters", "reset_trace_file"),
+            # Instrumentation
+            "enable_patching": ("instrumentation", "enable_patching"),
+            "enable_token_counting": ("instrumentation", "enable_token_counting"),
+            "enable_costs": ("instrumentation", "enable_costs"),
+            "auto_instrument_tools": ("instrumentation", "auto_instrument_tools"),
+            "max_tool_spans": ("instrumentation", "max_tool_spans"),
+            "max_span_depth": ("instrumentation", "max_span_depth"),
+            # Rate limiting & Batching
+            "max_spans_per_second": ("rate_limiting", "max_spans_per_second"),
+            "max_queue_size": ("rate_limiting", "max_queue_size"),
+            "max_block_ms": ("rate_limiting", "max_block_ms"),
+            "max_export_batch_size": ("rate_limiting", "max_export_batch_size"),
+            "schedule_delay_millis": ("rate_limiting", "schedule_delay_millis"),
+            # Runtime
+            "session_id": ("runtime", "session_id"),
+            "user_id": ("runtime", "user_id"),
+            "tenant_id": ("runtime", "tenant_id"),
+            "project_id": ("runtime", "project_id"),
+            "agent_id": ("runtime", "agent_id"),
+            # Logging
+            "debug": ("logging", "debug"),
+            "enable_span_logging": ("logging", "enable_span_logging"),
+            # Advanced
+            "attr_truncation_limit": ("advanced", "attr_truncation_limit"),
+        }
+        
+        for flat_key, value in overrides.items():
+            if flat_key in flat_to_nested:
+                section, nested_key = flat_to_nested[flat_key]
+                nested_overrides[section][nested_key] = value
+    
+    config = load_config(config_file=config_file, overrides=nested_overrides)
+    return config.to_flat_dict()