-
Notifications
You must be signed in to change notification settings - Fork 4.6k
/
datamodel.py
317 lines (254 loc) · 8.61 KB
/
datamodel.py
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
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
import uuid
from datetime import datetime
from typing import Any, Callable, Dict, List, Literal, Optional, Union
from pydantic.dataclasses import dataclass
from dataclasses import asdict, field
@dataclass
class Message(object):
user_id: str
role: str
content: str
root_msg_id: Optional[str] = None
msg_id: Optional[str] = None
timestamp: Optional[str] = None
personalize: Optional[bool] = False
ra: Optional[str] = None
code: Optional[str] = None
metadata: Optional[Any] = None
session_id: Optional[str] = None
def __post_init__(self):
if self.msg_id is None:
self.msg_id = str(uuid.uuid4())
if self.timestamp is None:
self.timestamp = datetime.now().isoformat()
def dict(self):
result = asdict(self)
return result
@dataclass
class Skill(object):
title: str
content: str
file_name: Optional[str] = None
id: Optional[str] = None
description: Optional[str] = None
timestamp: Optional[str] = None
user_id: Optional[str] = None
def __post_init__(self):
if self.id is None:
self.id = str(uuid.uuid4())
if self.timestamp is None:
self.timestamp = datetime.now().isoformat()
if self.user_id is None:
self.user_id = "default"
def dict(self):
result = asdict(self)
return result
# web api data models
# autogenflow data models
@dataclass
class Model:
"""Data model for Model Config item in LLMConfig for AutoGen"""
model: str
api_key: Optional[str] = None
base_url: Optional[str] = None
api_type: Optional[str] = None
api_version: Optional[str] = None
id: Optional[str] = None
timestamp: Optional[str] = None
user_id: Optional[str] = None
description: Optional[str] = None
def dict(self):
result = asdict(self)
return result
def __post_init__(self):
if self.id is None:
self.id = str(uuid.uuid4())
if self.timestamp is None:
self.timestamp = datetime.now().isoformat()
if self.user_id is None:
self.user_id = "default"
@dataclass
class LLMConfig:
"""Data model for LLM Config for AutoGen"""
config_list: List[Any] = field(default_factory=list)
temperature: float = 0
cache_seed: Optional[Union[int, None]] = None
timeout: Optional[int] = None
max_tokens: Optional[int] = None
extra_body: Optional[dict] = None
def dict(self):
result = asdict(self)
result["config_list"] = [c.dict() for c in self.config_list]
return result
@dataclass
class AgentConfig:
"""Data model for Agent Config for AutoGen"""
name: str
llm_config: Optional[Union[LLMConfig, bool]] = False
human_input_mode: str = "NEVER"
max_consecutive_auto_reply: int = 10
system_message: Optional[str] = None
is_termination_msg: Optional[Union[bool, str, Callable]] = None
code_execution_config: Optional[Union[bool, str, Dict[str, Any]]] = None
default_auto_reply: Optional[str] = ""
description: Optional[str] = None
def dict(self):
result = asdict(self)
if isinstance(result["llm_config"], LLMConfig):
result["llm_config"] = result["llm_config"].dict()
return result
@dataclass
class AgentFlowSpec:
"""Data model to help flow load agents from config"""
type: Literal["assistant", "userproxy"]
config: AgentConfig
id: Optional[str] = None
timestamp: Optional[str] = None
user_id: Optional[str] = None
skills: Optional[Union[None, List[Skill]]] = None
def __post_init__(self):
if self.timestamp is None:
self.timestamp = datetime.now().isoformat()
if self.id is None:
self.id = str(uuid.uuid4())
if self.user_id is None:
self.user_id = "default"
def dict(self):
result = asdict(self)
return result
@dataclass
class GroupChatConfig:
"""Data model for GroupChat Config for AutoGen"""
agents: List[AgentFlowSpec] = field(default_factory=list)
admin_name: str = "Admin"
messages: List[Dict] = field(default_factory=list)
max_round: Optional[int] = 10
admin_name: Optional[str] = "Admin"
speaker_selection_method: Optional[str] = "auto"
# TODO: match the new group chat default and support transition spec
allow_repeat_speaker: Optional[Union[bool, List[AgentConfig]]] = True
def dict(self):
result = asdict(self)
result["agents"] = [a.dict() for a in self.agents]
return result
@dataclass
class GroupChatFlowSpec:
"""Data model to help flow load agents from config"""
type: Literal["groupchat"]
config: AgentConfig = field(default_factory=AgentConfig)
groupchat_config: Optional[GroupChatConfig] = field(default_factory=GroupChatConfig)
id: Optional[str] = None
timestamp: Optional[str] = None
user_id: Optional[str] = None
skills: Optional[Union[None, List[Skill]]] = None
def __post_init__(self):
if self.timestamp is None:
self.timestamp = datetime.now().isoformat()
if self.id is None:
self.id = str(uuid.uuid4())
if self.user_id is None:
self.user_id = "default"
def dict(self):
result = asdict(self)
# result["config"] = self.config.dict()
# result["groupchat_config"] = self.groupchat_config.dict()
return result
@dataclass
class AgentWorkFlowConfig:
"""Data model for Flow Config for AutoGen"""
name: str
description: str
sender: AgentFlowSpec
receiver: Union[AgentFlowSpec, GroupChatFlowSpec]
type: Literal["twoagents", "groupchat"] = "twoagents"
id: Optional[str] = None
user_id: Optional[str] = None
timestamp: Optional[str] = None
# how the agent message summary is generated. last: only last message is used, none: no summary, llm: use llm to generate summary
summary_method: Optional[Literal["last", "none", "llm"]] = "last"
def init_spec(self, spec: Dict):
"""initialize the agent spec"""
if not isinstance(spec, dict):
spec = spec.dict()
if spec["type"] == "groupchat":
return GroupChatFlowSpec(**spec)
else:
return AgentFlowSpec(**spec)
def __post_init__(self):
if self.id is None:
self.id = str(uuid.uuid4())
self.sender = self.init_spec(self.sender)
self.receiver = self.init_spec(self.receiver)
if self.user_id is None:
self.user_id = "default"
if self.timestamp is None:
self.timestamp = datetime.now().isoformat()
def dict(self):
result = asdict(self)
result["sender"] = self.sender.dict()
result["receiver"] = self.receiver.dict()
return result
@dataclass
class Session(object):
"""Data model for AutoGen Chat Session"""
user_id: str
id: Optional[str] = None
timestamp: Optional[str] = None
flow_config: AgentWorkFlowConfig = None
name: Optional[str] = None
description: Optional[str] = None
def __post_init__(self):
if self.timestamp is None:
self.timestamp = datetime.now().isoformat()
if self.id is None:
self.id = str(uuid.uuid4())
def dict(self):
result = asdict(self)
result["flow_config"] = self.flow_config.dict()
return result
@dataclass
class Gallery(object):
"""Data model for Gallery Item"""
session: Session
messages: List[Message]
tags: List[str]
id: Optional[str] = None
timestamp: Optional[str] = None
def __post_init__(self):
if self.timestamp is None:
self.timestamp = datetime.now().isoformat()
if self.id is None:
self.id = str(uuid.uuid4())
def dict(self):
result = asdict(self)
return result
@dataclass
class ChatWebRequestModel(object):
"""Data model for Chat Web Request for Web End"""
message: Message
flow_config: AgentWorkFlowConfig
@dataclass
class DeleteMessageWebRequestModel(object):
user_id: str
msg_id: str
session_id: Optional[str] = None
@dataclass
class DBWebRequestModel(object):
user_id: str
msg_id: Optional[str] = None
session: Optional[Session] = None
skill: Optional[Skill] = None
tags: Optional[List[str]] = None
agent: Optional[AgentFlowSpec] = None
workflow: Optional[AgentWorkFlowConfig] = None
model: Optional[Model] = None
message: Optional[Message] = None
connection_id: Optional[str] = None
@dataclass
class SocketMessage(object):
connection_id: str
data: Dict[str, Any]
type: str
def dict(self):
result = asdict(self)
return result