from dataclasses import dataclass
from pathlib import Path
from typing import Optional, Union

from starlette.requests import Request
from starlette.responses import Response
from starlette.websockets import WebSocket

from phoenix.core.evals import Evals
from phoenix.core.model_schema import Model
from phoenix.core.traces import Traces


@dataclass
class Context:
    request: Union[Request, WebSocket]
    response: Optional[Response]
    model: Model
    export_path: Path
    corpus: Optional[Model] = None
    traces: Optional[Traces] = None
    evals: Optional[Evals] = None
