Still wip on rewrite

1 files changed, 15 insertions, 0 deletions

diff --git a/rag/message.py b/rag/message.py
new file mode 100644
index 0000000..d628982
--- /dev/null
+++ b/rag/message.py
@@ -0,0 +1,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}