evalvault 1.63.1__py3-none-any.whl → 1.65.0__py3-none-any.whl

evalvault/adapters/outbound/tracker/phoenix_adapter.py

@@ -9,6 +9,13 @@ from datetime import datetime
 from typing import TYPE_CHECKING, Any
 
 from evalvault.adapters.outbound.tracer.open_rag_trace_helpers import serialize_json
+from evalvault.adapters.outbound.tracker.log_sanitizer import (
+    MAX_CONTEXT_CHARS,
+    MAX_LOG_CHARS,
+    sanitize_payload,
+    sanitize_text,
+    sanitize_text_list,
+)
 from evalvault.domain.entities import (
     EvaluationRun,
     GenerationData,
@@ -19,8 +26,7 @@ from evalvault.domain.entities import (
 from evalvault.ports.outbound.tracker_port import TrackerPort
 
 if TYPE_CHECKING:
-    from opentelemetry.sdk.trace import Span, TracerProvider
-    from opentelemetry.trace import Tracer
+    from opentelemetry.sdk.trace import TracerProvider
 
 
 class PhoenixAdapter(TrackerPort):
@@ -55,9 +61,10 @@ class PhoenixAdapter(TrackerPort):
         """
         self._endpoint = endpoint
         self._service_name = service_name
-        self._tracer: Tracer | None = None
+        self._tracer: Any | None = None
         self._tracer_provider: TracerProvider | None = None
-        self._active_spans: dict[str, Span] = {}
+        self._active_spans: dict[str, Any] = {}
+        self._tracer_any: Any | None = None
         self._initialized = False
 
     def _ensure_initialized(self) -> None:
@@ -83,7 +90,8 @@ class PhoenixAdapter(TrackerPort):
                 provider = get_tracer_provider()
                 if provider:
                     self._tracer_provider = provider
-                    self._tracer = trace.get_tracer(__name__)
+                    self._tracer_any = trace.get_tracer(__name__)
+                    self._tracer = self._tracer_any
                     self._initialized = True
                     return
 
@@ -102,7 +110,8 @@ class PhoenixAdapter(TrackerPort):
             trace.set_tracer_provider(self._tracer_provider)
 
             # Get tracer
-            self._tracer = trace.get_tracer(__name__)
+            self._tracer_any = trace.get_tracer(__name__)
+            self._tracer = self._tracer_any
             self._initialized = True
 
         except ImportError as e:
@@ -127,7 +136,12 @@ class PhoenixAdapter(TrackerPort):
         self._ensure_initialized()
 
         # Start a new span as root
-        span = self._tracer.start_span(name)
+        tracer = self._tracer_any
+        if tracer is None:
+            tracer = self._tracer
+        if tracer is None:
+            raise RuntimeError("Phoenix tracer is not initialized")
+        span = tracer.start_span(name)
         trace_id = str(uuid.uuid4())
 
         # Set metadata as span attributes
@@ -166,14 +180,21 @@ class PhoenixAdapter(TrackerPort):
 
         from opentelemetry import trace
 
+        tracer = self._tracer_any
+        if tracer is None:
+            tracer = self._tracer
+        if tracer is None:
+            raise RuntimeError("Phoenix tracer is not initialized")
         parent_span = self._active_spans[trace_id]
         context = trace.set_span_in_context(parent_span)
 
-        with self._tracer.start_span(name, context=context) as span:
+        with tracer.start_span(name, context=context) as span:
             if input_data is not None:
-                span.set_attribute("input", json.dumps(input_data, default=str))
+                safe_input = sanitize_payload(input_data, max_chars=MAX_LOG_CHARS)
+                span.set_attribute("input", json.dumps(safe_input, default=str))
             if output_data is not None:
-                span.set_attribute("output", json.dumps(output_data, default=str))
+                safe_output = sanitize_payload(output_data, max_chars=MAX_LOG_CHARS)
+                span.set_attribute("output", json.dumps(safe_output, default=str))
 
     def log_score(
         self,
@@ -270,7 +291,7 @@ class PhoenixAdapter(TrackerPort):
             passed_count = sum(
                 1
                 for r in run.results
-                if r.get_metric(metric_name) and r.get_metric(metric_name).passed
+                if (metric := r.get_metric(metric_name)) and metric.passed is True
             )
             avg_score = run.get_avg_score(metric_name)
             threshold = run.thresholds.get(metric_name, 0.7)
@@ -360,20 +381,33 @@ class PhoenixAdapter(TrackerPort):
         """
         from opentelemetry import trace
 
+        tracer = self._tracer_any
+        if tracer is None:
+            tracer = self._tracer
+        if tracer is None:
+            raise RuntimeError("Phoenix tracer is not initialized")
         parent_span = self._active_spans[trace_id]
         context = trace.set_span_in_context(parent_span)
 
-        with self._tracer.start_span(
+        with tracer.start_span(
             f"test-case-{result.test_case_id}",
             context=context,
         ) as span:
             # Input data
-            span.set_attribute("input.question", result.question or "")
-            span.set_attribute("input.answer", result.answer or "")
+            safe_question = sanitize_text(result.question, max_chars=MAX_LOG_CHARS) or ""
+            safe_answer = sanitize_text(result.answer, max_chars=MAX_LOG_CHARS) or ""
+            span.set_attribute("input.question", safe_question)
+            span.set_attribute("input.answer", safe_answer)
             if result.contexts:
-                span.set_attribute("input.contexts", json.dumps(result.contexts))
+                safe_contexts = sanitize_text_list(
+                    result.contexts,
+                    max_chars=MAX_CONTEXT_CHARS,
+                )
+                span.set_attribute("input.contexts", json.dumps(safe_contexts))
             if result.ground_truth:
-                span.set_attribute("input.ground_truth", result.ground_truth)
+                safe_ground_truth = sanitize_text(result.ground_truth, max_chars=MAX_LOG_CHARS)
+                if safe_ground_truth:
+                    span.set_attribute("input.ground_truth", safe_ground_truth)
 
             # Metrics
             span.set_attribute("output.all_passed", result.all_passed)
@@ -461,15 +495,22 @@ class PhoenixAdapter(TrackerPort):
         parent_span = self._active_spans[trace_id]
         context = trace.set_span_in_context(parent_span)
 
-        with self._tracer.start_span("retrieval", context=context) as span:
+        tracer = self._tracer_any
+        if tracer is None:
+            tracer = self._tracer
+        if tracer is None:
+            raise RuntimeError("Phoenix tracer is not initialized")
+        with tracer.start_span("retrieval", context=context) as span:
             # Set retrieval attributes
             for key, value in data.to_span_attributes().items():
                 span.set_attribute(key, value)
 
             # Set query
             if data.query:
-                span.set_attribute("retrieval.query", data.query)
-                span.set_attribute("input.value", data.query)
+                safe_query = sanitize_text(data.query, max_chars=MAX_LOG_CHARS)
+                if safe_query:
+                    span.set_attribute("retrieval.query", safe_query)
+                    span.set_attribute("input.value", safe_query)
 
             span.set_attribute("spec.version", "0.1")
             span.set_attribute("rag.module", "retrieve")
@@ -495,11 +536,14 @@ class PhoenixAdapter(TrackerPort):
                     event_attrs["doc.rerank_rank"] = doc.rerank_rank
                 if doc.chunk_id:
                     event_attrs["doc.chunk_id"] = doc.chunk_id
-                preview = doc.content[:200] if doc.content else ""
-                if preview:
-                    event_attrs["doc.preview"] = preview
+                safe_preview = (
+                    sanitize_text(doc.content, max_chars=MAX_CONTEXT_CHARS) if doc.content else ""
+                )
+                if safe_preview:
+                    event_attrs["doc.preview"] = safe_preview
                 if doc.metadata:
-                    event_attrs["doc.metadata"] = json.dumps(doc.metadata, default=str)
+                    safe_metadata = sanitize_payload(doc.metadata, max_chars=MAX_LOG_CHARS)
+                    event_attrs["doc.metadata"] = json.dumps(safe_metadata, default=str)
                 span.add_event(f"retrieved_doc_{i}", attributes=event_attrs)
 
     def log_generation(
@@ -538,15 +582,19 @@ class PhoenixAdapter(TrackerPort):
         parent_span = self._active_spans[trace_id]
         context = trace.set_span_in_context(parent_span)
 
-        with self._tracer.start_span("generation", context=context) as span:
+        tracer = self._tracer_any
+        if tracer is None:
+            tracer = self._tracer
+        if tracer is None:
+            raise RuntimeError("Phoenix tracer is not initialized")
+        with tracer.start_span("generation", context=context) as span:
             # Set generation attributes
             for key, value in data.to_span_attributes().items():
                 span.set_attribute(key, value)
 
             # Set prompt/response (truncate if too long)
-            max_len = 10000
-            prompt = data.prompt[:max_len] if data.prompt else ""
-            response = data.response[:max_len] if data.response else ""
+            prompt = sanitize_text(data.prompt, max_chars=MAX_LOG_CHARS) or ""
+            response = sanitize_text(data.response, max_chars=MAX_LOG_CHARS) or ""
             if prompt:
                 span.set_attribute("generation.prompt", prompt)
                 span.set_attribute("input.value", prompt)
@@ -559,24 +607,28 @@ class PhoenixAdapter(TrackerPort):
 
             # Set prompt template if available
             if data.prompt_template:
-                span.set_attribute("generation.prompt_template", data.prompt_template[:max_len])
+                safe_template = sanitize_text(data.prompt_template, max_chars=MAX_LOG_CHARS)
+                if safe_template:
+                    span.set_attribute("generation.prompt_template", safe_template)
 
     def log_rag_trace(self, data: RAGTraceData) -> str:
         """Log a full RAG trace (retrieval + generation) to Phoenix."""
 
         self._ensure_initialized()
         metadata = {"event_type": "rag_trace", "total_time_ms": data.total_time_ms}
-        if data.query:
-            metadata["query"] = data.query
+        safe_query = sanitize_text(data.query, max_chars=MAX_LOG_CHARS)
+        if safe_query:
+            metadata["query"] = safe_query
         if data.metadata:
-            metadata.update(data.metadata)
+            safe_metadata = sanitize_payload(data.metadata, max_chars=MAX_LOG_CHARS)
+            metadata.update(safe_metadata)
 
         should_end = False
         trace_id = data.trace_id
         if trace_id and trace_id in self._active_spans:
             span = self._active_spans[trace_id]
         else:
-            trace_name = f"rag-trace-{(data.query or 'run')[:12]}"
+            trace_name = f"rag-trace-{(safe_query or 'run')[:12]}"
             trace_id = self.start_trace(trace_name, metadata=metadata)
             span = self._active_spans[trace_id]
             should_end = True
@@ -589,12 +641,13 @@ class PhoenixAdapter(TrackerPort):
         if data.generation:
             self.log_generation(trace_id, data.generation)
         if data.final_answer:
-            preview = data.final_answer[:1000]
-            span.set_attribute("rag.final_answer", preview)
-            span.set_attribute("output.value", preview)
+            preview = sanitize_text(data.final_answer, max_chars=MAX_LOG_CHARS)
+            if preview:
+                span.set_attribute("rag.final_answer", preview)
+                span.set_attribute("output.value", preview)
 
-        if data.query:
-            span.set_attribute("input.value", data.query)
+        if safe_query:
+            span.set_attribute("input.value", safe_query)
 
         span.set_attribute("spec.version", "0.1")
         span.set_attribute("rag.module", "custom.pipeline")

evalvault/config/secret_manager.py

@@ -0,0 +1,118 @@
+from __future__ import annotations
+
+import base64
+import os
+from dataclasses import dataclass
+from typing import Protocol
+
+SECRET_REF_PREFIX = "secret://"
+
+
+class SecretProvider(Protocol):
+    def get_secret(self, name: str) -> str: ...
+
+
+class SecretProviderError(RuntimeError):
+    pass
+
+
+@dataclass
+class EnvSecretProvider:
+    def get_secret(self, name: str) -> str:
+        value = os.environ.get(name)
+        if value is None:
+            raise SecretProviderError(f"Missing secret in environment: {name}")
+        return value
+
+
+@dataclass
+class AwsSecretsManagerProvider:
+    region_name: str | None = None
+
+    def get_secret(self, name: str) -> str:
+        try:
+            import boto3  # type: ignore
+        except ImportError as exc:
+            raise SecretProviderError("boto3 is required for AWS Secrets Manager") from exc
+        client = boto3.client("secretsmanager", region_name=self.region_name)
+        response = client.get_secret_value(SecretId=name)
+        if "SecretString" in response and response["SecretString"] is not None:
+            return response["SecretString"]
+        secret_binary = response.get("SecretBinary")
+        if secret_binary is None:
+            raise SecretProviderError("Empty secret value returned from AWS Secrets Manager")
+        return base64.b64decode(secret_binary).decode("utf-8")
+
+
+@dataclass
+class GcpSecretManagerProvider:
+    def get_secret(self, name: str) -> str:
+        try:
+            from google.cloud import secretmanager  # type: ignore
+        except ImportError as exc:
+            raise SecretProviderError(
+                "google-cloud-secret-manager is required for GCP Secret Manager"
+            ) from exc
+        client = secretmanager.SecretManagerServiceClient()
+        response = client.access_secret_version(request={"name": name})
+        return response.payload.data.decode("utf-8")
+
+
+@dataclass
+class VaultSecretProvider:
+    def get_secret(self, name: str) -> str:
+        try:
+            import hvac  # type: ignore
+        except ImportError as exc:
+            raise SecretProviderError("hvac is required for Vault secret access") from exc
+        client = hvac.Client()
+        if not client.is_authenticated():
+            raise SecretProviderError("Vault client authentication failed")
+        response = client.secrets.kv.v2.read_secret_version(path=name)
+        data = response.get("data", {}).get("data", {})
+        if not data:
+            raise SecretProviderError("Vault secret payload is empty")
+        if "value" in data:
+            return str(data["value"])
+        if len(data) == 1:
+            return str(next(iter(data.values())))
+        raise SecretProviderError("Vault secret has multiple keys; specify 'value' key")
+
+
+def is_secret_reference(value: str | None) -> bool:
+    return bool(value) and value.startswith(SECRET_REF_PREFIX)
+
+
+def parse_secret_reference(value: str) -> str:
+    return value.removeprefix(SECRET_REF_PREFIX).strip()
+
+
+def build_secret_provider(provider_name: str | None) -> SecretProvider:
+    provider = (provider_name or "").strip().lower()
+    if not provider:
+        raise SecretProviderError("Secret provider is not configured.")
+    if provider == "env":
+        return EnvSecretProvider()
+    if provider in {"aws", "aws-secrets-manager", "secretsmanager"}:
+        return AwsSecretsManagerProvider(region_name=os.environ.get("AWS_REGION"))
+    if provider in {"gcp", "gcp-secret-manager", "secretmanager"}:
+        return GcpSecretManagerProvider()
+    if provider in {"vault", "hashicorp-vault"}:
+        return VaultSecretProvider()
+    raise SecretProviderError(f"Unknown secret provider: {provider_name}")
+
+
+def resolve_secret_reference(
+    value: str,
+    provider: SecretProvider,
+    cache: dict[str, str] | None = None,
+) -> str:
+    secret_name = parse_secret_reference(value)
+    if not secret_name:
+        raise SecretProviderError("Secret reference must include a name.")
+    if cache is not None and secret_name in cache:
+        return cache[secret_name]
+    secret_value = provider.get_secret(secret_name)
+    if cache is not None:
+        cache[secret_name] = secret_value
+    return secret_value

evalvault/config/settings.py

@@ -3,9 +3,16 @@
 from pathlib import Path
 from typing import Any
 
-from pydantic import Field
+from pydantic import Field, PrivateAttr
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
+from evalvault.config.secret_manager import (
+    SecretProviderError,
+    build_secret_provider,
+    is_secret_reference,
+    resolve_secret_reference,
+)
+
 
 def _detect_repo_root(start: Path, max_depth: int = 6) -> Path | None:
     current = start
@@ -38,6 +45,75 @@ def _ensure_http_scheme(url_value: str) -> str:
     return f"http://{value}"
 
 
+def is_production_profile(profile_name: str | None) -> bool:
+    return (profile_name or "").strip().lower() == "prod"
+
+
+def _parse_cors_origins(cors_origins: str | None) -> list[str]:
+    if not cors_origins:
+        return []
+    return [origin.strip() for origin in cors_origins.split(",") if origin.strip()]
+
+
+SECRET_REFERENCE_FIELDS = (
+    "api_auth_tokens",
+    "knowledge_read_tokens",
+    "knowledge_write_tokens",
+    "openai_api_key",
+    "anthropic_api_key",
+    "azure_api_key",
+    "vllm_api_key",
+    "langfuse_public_key",
+    "langfuse_secret_key",
+    "phoenix_api_token",
+    "postgres_password",
+    "postgres_connection_string",
+)
+
+
+def _validate_production_settings(settings: "Settings") -> None:
+    if not is_production_profile(settings.evalvault_profile):
+        return
+
+    missing: list[str] = []
+
+    if not settings.api_auth_tokens:
+        missing.append("API_AUTH_TOKENS")
+
+    if settings.llm_provider == "openai" and not settings.openai_api_key:
+        missing.append("OPENAI_API_KEY")
+
+    if settings.tracker_provider == "langfuse":
+        if not settings.langfuse_public_key:
+            missing.append("LANGFUSE_PUBLIC_KEY")
+        if not settings.langfuse_secret_key:
+            missing.append("LANGFUSE_SECRET_KEY")
+
+    if settings.tracker_provider == "mlflow" and not settings.mlflow_tracking_uri:
+        missing.append("MLFLOW_TRACKING_URI")
+
+    if (
+        settings.postgres_connection_string is None
+        and settings.postgres_host
+        and not settings.postgres_password
+    ):
+        missing.append("POSTGRES_PASSWORD")
+
+    cors_origins = _parse_cors_origins(settings.cors_origins)
+    if not cors_origins:
+        missing.append("CORS_ORIGINS")
+    else:
+        localhost_origins = {"localhost", "127.0.0.1"}
+        for origin in cors_origins:
+            if any(host in origin for host in localhost_origins):
+                raise ValueError("Production profile forbids localhost in CORS_ORIGINS.")
+
+    if missing:
+        raise ValueError(
+            "Missing required settings for prod profile: " + ", ".join(sorted(set(missing)))
+        )
+
+
 class Settings(BaseSettings):
     """Application configuration settings."""
 
@@ -48,6 +124,8 @@ class Settings(BaseSettings):
         extra="ignore",
     )
 
+    _secret_cache: dict[str, str] = PrivateAttr(default_factory=dict)
+
     # Profile Configuration (YAML 기반 모델 프로필)
     evalvault_profile: str | None = Field(
         default=None,
@@ -58,6 +136,45 @@ class Settings(BaseSettings):
         default="http://localhost:5173,http://127.0.0.1:5173",
         description="Comma-separated list of allowed CORS origins.",
     )
+    secret_provider: str | None = Field(
+        default=None,
+        description="Secret provider name for secret:// references (env/aws/gcp/vault).",
+    )
+    secret_cache_enabled: bool = Field(
+        default=True,
+        description="Cache resolved secret references in memory.",
+    )
+    api_auth_tokens: str | None = Field(
+        default=None,
+        description=(
+            "Comma-separated list of API bearer tokens for FastAPI auth. "
+            "Leave empty to disable authentication."
+        ),
+    )
+    knowledge_read_tokens: str | None = Field(
+        default=None,
+        description="Comma-separated read tokens for knowledge endpoints.",
+    )
+    knowledge_write_tokens: str | None = Field(
+        default=None,
+        description="Comma-separated write tokens for knowledge endpoints.",
+    )
+    rate_limit_enabled: bool = Field(
+        default=False,
+        description="Enable API rate limiting for /api routes.",
+    )
+    rate_limit_requests: int = Field(
+        default=120,
+        description="Max requests allowed within rate_limit_window_seconds.",
+    )
+    rate_limit_window_seconds: int = Field(
+        default=60,
+        description="Window size for rate limit checks in seconds.",
+    )
+    rate_limit_block_threshold: int = Field(
+        default=10,
+        description="Log suspicious activity after this many rate limit blocks.",
+    )
     evalvault_db_path: str = Field(
         default="data/db/evalvault.db",
         description="SQLite database path for API/CLI storage.",
@@ -71,6 +188,26 @@ class Settings(BaseSettings):
         self.evalvault_db_path = _resolve_storage_path(self.evalvault_db_path)
         self.evalvault_memory_db_path = _resolve_storage_path(self.evalvault_memory_db_path)
         self.ollama_base_url = _ensure_http_scheme(self.ollama_base_url)
+        self._resolve_secret_references()
+
+    def _resolve_secret_references(self) -> None:
+        secret_values = [
+            value
+            for value in (getattr(self, field, None) for field in SECRET_REFERENCE_FIELDS)
+            if isinstance(value, str)
+        ]
+        if not any(is_secret_reference(value) for value in secret_values):
+            return
+        try:
+            provider = build_secret_provider(self.secret_provider)
+        except SecretProviderError as exc:
+            raise ValueError(str(exc)) from exc
+        cache = self._secret_cache if self.secret_cache_enabled else None
+        for field in SECRET_REFERENCE_FIELDS:
+            value = getattr(self, field, None)
+            if isinstance(value, str) and is_secret_reference(value):
+                resolved = resolve_secret_reference(value, provider, cache)
+                setattr(self, field, resolved)
 
     # LLM Provider Selection
     llm_provider: str = Field(
@@ -314,6 +451,8 @@ def get_settings() -> Settings:
         if _settings.evalvault_profile:
             _settings = apply_profile(_settings, _settings.evalvault_profile)
 
+        _validate_production_settings(_settings)
+
     return _settings
 
 
@@ -346,6 +485,7 @@ def apply_runtime_overrides(overrides: dict[str, object]) -> Settings:
     updated = Settings.model_validate(payload)
     if updated.evalvault_profile:
         updated = apply_profile(updated, updated.evalvault_profile)
+    _validate_production_settings(updated)
     for key, value in updated.model_dump().items():
         setattr(settings, key, value)
 

evalvault/domain/entities/__init__.py

@@ -34,6 +34,12 @@ from evalvault.domain.entities.improvement import (
     RAGComponent,
     RAGImprovementGuide,
 )
+from evalvault.domain.entities.judge_calibration import (
+    JudgeCalibrationCase,
+    JudgeCalibrationMetric,
+    JudgeCalibrationResult,
+    JudgeCalibrationSummary,
+)
 from evalvault.domain.entities.kg import EntityModel, RelationModel
 from evalvault.domain.entities.method import MethodInput, MethodInputDataset, MethodOutput
 from evalvault.domain.entities.prompt import Prompt, PromptSet, PromptSetBundle, PromptSetItem
@@ -104,6 +110,10 @@ __all__ = [
     "PatternType",
     "RAGComponent",
     "RAGImprovementGuide",
+    "JudgeCalibrationCase",
+    "JudgeCalibrationMetric",
+    "JudgeCalibrationResult",
+    "JudgeCalibrationSummary",
     # KG
     "EntityModel",
     "RelationModel",

evalvault/domain/entities/judge_calibration.py

@@ -0,0 +1,50 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+
+@dataclass
+class JudgeCalibrationCase:
+    test_case_id: str
+    raw_score: float
+    calibrated_score: float
+    label: float | None = None
+    label_source: str | None = None
+
+
+@dataclass
+class JudgeCalibrationMetric:
+    metric: str
+    method: str
+    sample_count: int
+    label_count: int
+    mae: float | None
+    pearson: float | None
+    spearman: float | None
+    temperature: float | None = None
+    parameters: dict[str, float | None] = field(default_factory=dict)
+    gate_passed: bool | None = None
+    warning: str | None = None
+
+
+@dataclass
+class JudgeCalibrationSummary:
+    run_id: str
+    labels_source: str
+    method: str
+    metrics: list[str]
+    holdout_ratio: float
+    seed: int
+    total_labels: int
+    total_samples: int
+    gate_passed: bool
+    gate_threshold: float | None = None
+    notes: list[str] = field(default_factory=list)
+
+
+@dataclass
+class JudgeCalibrationResult:
+    summary: JudgeCalibrationSummary
+    metrics: list[JudgeCalibrationMetric] = field(default_factory=list)
+    case_results: dict[str, list[JudgeCalibrationCase]] = field(default_factory=dict)
+    warnings: list[str] = field(default_factory=list)

evalvault/domain/entities/stage.py

@@ -4,7 +4,7 @@ from __future__ import annotations
 
 from dataclasses import dataclass, field
 from datetime import datetime
-from typing import Any
+from typing import Any, Literal, overload
 from uuid import uuid4
 
 REQUIRED_STAGE_TYPES: tuple[str, ...] = ("system_prompt", "input", "retrieval", "output")
@@ -82,8 +82,8 @@ class StageEvent:
             duration_ms=_optional_float(payload.get("duration_ms")),
             input_ref=input_ref,
             output_ref=output_ref,
-            attributes=_ensure_dict(payload.get("attributes")),
-            metadata=_ensure_dict(payload.get("metadata")),
+            attributes=_ensure_dict(payload.get("attributes"), allow_none=False),
+            metadata=_ensure_dict(payload.get("metadata"), allow_none=False),
             trace_id=_optional_str(payload.get("trace_id") or trace_payload.get("trace_id")),
             span_id=_optional_str(payload.get("span_id") or trace_payload.get("span_id")),
         )
@@ -187,6 +187,14 @@ def _parse_datetime(value: Any) -> datetime | None:
     raise ValueError("Invalid datetime value")
 
 
+@overload
+def _ensure_dict(value: None, *, allow_none: Literal[True]) -> None: ...
+
+
+@overload
+def _ensure_dict(value: Any, *, allow_none: Literal[False] = False) -> dict[str, Any]: ...
+
+
 def _ensure_dict(value: Any, *, allow_none: bool = False) -> dict[str, Any] | None:
     if value is None:
         return None if allow_none else {}