from dataclasses import dataclass
from typing import List, Literal, Optional

from langchain.schema import Document


@dataclass
class Message:
    """Class for keeping track of a chat message."""

    origin: Literal["human", "ai"]
    message: str
    documents: Optional[List[Document]] = None