From aac821b148c6c0d35b940609dc9b0ddcb053b28e Mon Sep 17 00:00:00 2001
From: Gustaf Rydholm <gustaf.rydholm@gmail.com>
Date: Wed, 19 Jun 2024 02:07:06 +0200
Subject: Still wip on rewrite

---
 rag/message.py | 15 +++++++++++++++
 1 file changed, 15 insertions(+)
 create mode 100644 rag/message.py

(limited to 'rag/message.py')

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}
-- 
cgit v1.2.3-70-g09d2