rag/message.py


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

from dataclasses import dataclass
from typing import Dict


@dataclass
class Message:
    role: str
    content: str
    client: str

    def as_dict(self) -> Dict[str, str]:
        if self.client == "cohere":
            return {"role": self.role, "message": self.content}
        else:
            return {"role": self.role, "content": self.content}