blob: d207596447aa87540598f62e934ab6344343dc23 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
|
from dataclasses import dataclass
from typing import Dict, List
from loguru import logger as log
@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}
@dataclass
class Messages:
messages: List[Message]
def __add__(self, message: Message):
self.messages.append(message)
def __len__(self):
return len(self.messages)
def reset(self):
log.debug("Resetting messages...")
self.messages = []
def content(self) -> List[str]:
return [m.content for m in self.messages]
def rerank(self, rankings: List[int]):
self.messages = [self.messages[r] for r in rankings]
|
