message.py

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

    def __repr__(self) -> str:
        return f"Message(origin={self.origin}, message={self.message})"