"""Trace context for grouping related API calls."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Optional


@dataclass
class Trace:
    """
    Trace context for grouping related API calls.

    Usage:
        with tracer.trace("user-query-123") as trace:
            response1 = client.chat.completions.create(...)
            trace.set_metadata({"user_id": "user-456"})
    """

    trace_id: str
    name: Optional[str] = None
    metadata: dict[str, Any] = field(default_factory=dict)

    def set_metadata(self, metadata: dict[str, Any]) -> "Trace":
        """
        Set metadata for this trace.

        Args:
            metadata: Dictionary of key-value pairs.

        Returns:
            Self for chaining.
        """
        self.metadata.update(metadata)
        return self
