from typing import TypedDict, Optional, Any, Literal

StepType = Literal["llm_call", "tool_call", "reasoning", "other"]
TraceStatus = Literal["running", "success", "error"]


class TraceOptions(TypedDict, total=False):
    user_id: Optional[str]
    session_id: Optional[str]
    input: Any
    metadata: dict[str, Any]
    prompt_id: Optional[str]
    prompt_version: Optional[int]


class StepOptions(TypedDict, total=False):
    input: Any
    metadata: dict[str, Any]
    model: Optional[str]


class EndStepOptions(TypedDict, total=False):
    output: Any
    status: TraceStatus
    error_message: Optional[str]
    model: Optional[str]
    prompt_tokens: Optional[int]
    completion_tokens: Optional[int]


class EndTraceOptions(TypedDict, total=False):
    output: Any
    status: TraceStatus
    error_message: Optional[str]
    metadata: dict[str, Any]
