| 1 | import json |
| 2 | import os |
| 3 | import tempfile |
| 4 | import uuid |
| 5 | from collections import OrderedDict |
| 6 | from datetime import datetime |
| 7 | from typing import Any |
| 8 | |
| 9 | from agent import Agent, AgentConfig, AgentContext, AgentContextType |
| 10 | from helpers import files, history |
| 11 | from helpers.litellm_transport import delete_stored_response_ids |
| 12 | from helpers.localization import Localization |
| 13 | from initialize import initialize_agent |
| 14 | |
| 15 | from helpers.log import Log, LogItem |
| 16 | |
| 17 | CHATS_FOLDER = "usr/chats" |
| 18 | LOG_SIZE = 1000 |
| 19 | CHAT_FILE_NAME = "chat.json" |
| 20 | SAVED_CHAT_CONTEXT_DATA_KEY = "_persist_chat_saved" |
| 21 | |
| 22 | |
| 23 | def _fallback_datetime_iso() -> str: |
| 24 | return datetime.fromtimestamp(0, tz=Localization.get().get_tzinfo()).isoformat() |
| 25 | |
| 26 | |
| 27 | def _parse_persisted_datetime(value: str | None) -> datetime: |
| 28 | raw_value = value or _fallback_datetime_iso() |
| 29 | dt = datetime.fromisoformat(raw_value) |
| 30 | if dt.tzinfo is None: |
| 31 | dt = Localization.get().localize_naive_datetime(dt) |
| 32 | return dt |
| 33 | |
| 34 | |
| 35 | def get_chat_folder_path(ctxid: str): |
| 36 | """ |
| 37 | Get the folder path for any context (chat or task). |
| 38 | |
| 39 | Args: |
| 40 | ctxid: The context ID |
| 41 | |
| 42 | Returns: |
| 43 | The absolute path to the context folder |
| 44 | """ |
| 45 | return files.get_abs_path(CHATS_FOLDER, ctxid) |
| 46 | |
| 47 | def get_chat_msg_files_folder(ctxid: str): |
| 48 | return files.get_abs_path(get_chat_folder_path(ctxid), "messages") |
| 49 | |
| 50 | def save_tmp_chat(context: AgentContext): |
| 51 | """Save context to the chats folder""" |
| 52 | # Skip saving BACKGROUND contexts as they should be ephemeral |
| 53 | if context.type == AgentContextType.BACKGROUND: |
| 54 | return |
| 55 | |
| 56 | path = _get_chat_file_path(context.id) |
| 57 | data = _serialize_context(context) |
| 58 | js = _safe_json_serialize(data, ensure_ascii=False) |
| 59 | _write_atomic(path, js) |
| 60 | mark_chat_saved(context) |
| 61 | |
| 62 | |
| 63 | def save_tmp_chats(): |
| 64 | """Save all contexts to the chats folder""" |
| 65 | for context in AgentContext.all(): |
| 66 | # Skip BACKGROUND contexts as they should be ephemeral |
| 67 | if context.type == AgentContextType.BACKGROUND: |
| 68 | continue |
| 69 | save_tmp_chat(context) |
| 70 | |
| 71 | |
| 72 | def load_tmp_chats(): |
| 73 | """Load all contexts from the chats folder""" |
| 74 | _convert_v080_chats() |
| 75 | folders = files.list_files(CHATS_FOLDER, "*") |
| 76 | json_files = [] |
| 77 | for folder_name in folders: |
| 78 | chat_file = _get_chat_file_path(folder_name) |
| 79 | if files.exists(chat_file): |
| 80 | json_files.append(chat_file) |
| 81 | |
| 82 | ctxids = [] |
| 83 | for file in json_files: |
| 84 | try: |
| 85 | js = files.read_file(file) |
| 86 | data = json.loads(js) |
| 87 | ctx = _deserialize_context(data) |
| 88 | mark_chat_saved(ctx) |
| 89 | ctxids.append(ctx.id) |
| 90 | except Exception as e: |
| 91 | print(f"Error loading chat {file}: {e}") |
| 92 | return ctxids |
| 93 | |
| 94 | |
| 95 | def _get_chat_file_path(ctxid: str): |
| 96 | return files.get_abs_path(CHATS_FOLDER, ctxid, CHAT_FILE_NAME) |
| 97 | |
| 98 | |
| 99 | def _write_atomic(path: str, content: str) -> None: |
| 100 | directory = os.path.dirname(path) |
| 101 | os.makedirs(directory, exist_ok=True) |
| 102 | fd, tmp_path = tempfile.mkstemp( |
| 103 | prefix=f".{os.path.basename(path)}.", suffix=".tmp", dir=directory |
| 104 | ) |
| 105 | try: |
| 106 | with os.fdopen(fd, "w", encoding="utf-8") as handle: |
| 107 | handle.write(content.encode("utf-8", "replace").decode("utf-8")) |
| 108 | handle.flush() |
| 109 | os.fsync(handle.fileno()) |
| 110 | os.replace(tmp_path, path) |
| 111 | directory_fd = os.open(directory, os.O_RDONLY) |
| 112 | try: |
| 113 | os.fsync(directory_fd) |
| 114 | finally: |
| 115 | os.close(directory_fd) |
| 116 | finally: |
| 117 | if os.path.exists(tmp_path): |
| 118 | os.unlink(tmp_path) |
| 119 | |
| 120 | |
| 121 | def mark_chat_saved(context: AgentContext) -> None: |
| 122 | context.data[SAVED_CHAT_CONTEXT_DATA_KEY] = True |
| 123 | |
| 124 | |
| 125 | def saved_chat_ids() -> set[str]: |
| 126 | return { |
| 127 | files.basename(files.dirname(path)) |
| 128 | for path in files.find_existing_paths_by_pattern( |
| 129 | files.get_abs_path(CHATS_FOLDER, "*", CHAT_FILE_NAME) |
| 130 | ) |
| 131 | } |
| 132 | |
| 133 | |
| 134 | def _convert_v080_chats(): |
| 135 | json_files = files.list_files(CHATS_FOLDER, "*.json") |
| 136 | for file in json_files: |
| 137 | path = files.get_abs_path(CHATS_FOLDER, file) |
| 138 | name = file.rstrip(".json") |
| 139 | new = _get_chat_file_path(name) |
| 140 | files.move_file(path, new) |
| 141 | |
| 142 | |
| 143 | def load_json_chats(jsons: list[str]): |
| 144 | """Load contexts from JSON strings""" |
| 145 | ctxids = [] |
| 146 | for js in jsons: |
| 147 | data = json.loads(js) |
| 148 | if "id" in data: |
| 149 | del data["id"] # remove id to get new |
| 150 | ctx = _deserialize_context(data) |
| 151 | ctxids.append(ctx.id) |
| 152 | return ctxids |
| 153 | |
| 154 | |
| 155 | def export_json_chat(context: AgentContext): |
| 156 | """Export context as JSON string""" |
| 157 | data = _serialize_context(context) |
| 158 | js = _safe_json_serialize(data, ensure_ascii=False) |
| 159 | return js |
| 160 | |
| 161 | |
| 162 | def remove_chat(ctxid): |
| 163 | """Remove a chat or task context""" |
| 164 | if not isinstance(ctxid, str) or not ctxid.strip(): |
| 165 | raise ValueError("remove_chat: context id must not be empty") |
| 166 | _delete_provider_responses_for_chat(ctxid) |
| 167 | path = get_chat_folder_path(ctxid) |
| 168 | files.delete_dir(path) |
| 169 | |
| 170 | |
| 171 | def remove_msg_files(ctxid): |
| 172 | """Remove all message files for a chat or task context""" |
| 173 | if not isinstance(ctxid, str) or not ctxid.strip(): |
| 174 | raise ValueError("remove_msg_files: context id must not be empty") |
| 175 | path = get_chat_msg_files_folder(ctxid) |
| 176 | files.delete_dir(path) |
| 177 | |
| 178 | |
| 179 | def _serialize_context(context: AgentContext): |
| 180 | profile = str( |
| 181 | getattr(context.agent0.config, "profile", None) |
| 182 | or getattr(context.config, "profile", None) |
| 183 | or "" |
| 184 | ) |
| 185 | |
| 186 | # serialize agents |
| 187 | agents = [] |
| 188 | agent = context.agent0 |
| 189 | while agent: |
| 190 | agents.append(_serialize_agent(agent)) |
| 191 | agent = agent.data.get(Agent.DATA_NAME_SUBORDINATE, None) |
| 192 | |
| 193 | |
| 194 | data = {k: v for k, v in context.data.items() if not k.startswith("_")} |
| 195 | output_data = {k: v for k, v in context.output_data.items() if not k.startswith("_")} |
| 196 | |
| 197 | return { |
| 198 | "id": context.id, |
| 199 | "name": context.name, |
| 200 | "created_at": ( |
| 201 | Localization.get().serialize_datetime(context.created_at) |
| 202 | if context.created_at |
| 203 | else _fallback_datetime_iso() |
| 204 | ), |
| 205 | "type": context.type.value, |
| 206 | "last_message": ( |
| 207 | Localization.get().serialize_datetime(context.last_message) |
| 208 | if context.last_message |
| 209 | else _fallback_datetime_iso() |
| 210 | ), |
| 211 | "agents": agents, |
| 212 | "streaming_agent": ( |
| 213 | context.streaming_agent.number if context.streaming_agent else 0 |
| 214 | ), |
| 215 | "agent_profile": profile, |
| 216 | "log": _serialize_log(context.log), |
| 217 | "data": data, |
| 218 | "output_data": output_data, |
| 219 | } |
| 220 | |
| 221 | |
| 222 | def _serialize_agent(agent: Agent): |
| 223 | data = {k: v for k, v in agent.data.items() if not k.startswith("_")} |
| 224 | |
| 225 | history = agent.history.serialize() |
| 226 | |
| 227 | return { |
| 228 | "number": agent.number, |
| 229 | "agent_profile": str(getattr(agent.config, "profile", "") or ""), |
| 230 | "data": data, |
| 231 | "history": history, |
| 232 | } |
| 233 | |
| 234 | |
| 235 | def _serialize_log(log: Log): |
| 236 | # Guard against concurrent log mutations while serializing. |
| 237 | with log._lock: |
| 238 | logs = [item.output() for item in log.logs[-LOG_SIZE:]] # serialize LogItem objects |
| 239 | guid = log.guid |
| 240 | progress = log.progress |
| 241 | progress_no = log.progress_no |
| 242 | return { |
| 243 | "guid": guid, |
| 244 | "logs": logs, |
| 245 | "progress": progress, |
| 246 | "progress_no": progress_no, |
| 247 | } |
| 248 | |
| 249 | |
| 250 | def _deserialize_context(data): |
| 251 | profile = data.get("agent_profile") |
| 252 | override_settings = {"agent_profile": profile} if profile else None |
| 253 | config = initialize_agent(override_settings=override_settings) |
| 254 | log = _deserialize_log(data.get("log", None)) |
| 255 | |
| 256 | context = AgentContext( |
| 257 | config=config, |
| 258 | id=data.get("id", None), # get new id |
| 259 | name=data.get("name", None), |
| 260 | created_at=( |
| 261 | _parse_persisted_datetime(data.get("created_at")) |
| 262 | ), |
| 263 | type=AgentContextType(data.get("type", AgentContextType.USER.value)), |
| 264 | last_message=( |
| 265 | _parse_persisted_datetime(data.get("last_message")) |
| 266 | ), |
| 267 | log=log, |
| 268 | paused=False, |
| 269 | data=data.get("data", {}), |
| 270 | output_data=data.get("output_data", {}), |
| 271 | # agent0=agent0, |
| 272 | # streaming_agent=straming_agent, |
| 273 | ) |
| 274 | |
| 275 | agents = data.get("agents", []) |
| 276 | agent0 = _deserialize_agents(agents, config, context) |
| 277 | streaming_agent = agent0 |
| 278 | while streaming_agent and streaming_agent.number != data.get("streaming_agent", 0): |
| 279 | streaming_agent = streaming_agent.data.get(Agent.DATA_NAME_SUBORDINATE, None) |
| 280 | |
| 281 | context.agent0 = agent0 |
| 282 | context.config = agent0.config |
| 283 | context.streaming_agent = streaming_agent |
| 284 | |
| 285 | return context |
| 286 | |
| 287 | |
| 288 | def _deserialize_agent_config( |
| 289 | agent_data: dict[str, Any], fallback_config: AgentConfig |
| 290 | ) -> AgentConfig: |
| 291 | fallback_profile = str(getattr(fallback_config, "profile", "") or "") |
| 292 | profile = str(agent_data.get("agent_profile") or fallback_profile).strip() |
| 293 | if profile == fallback_profile: |
| 294 | return fallback_config |
| 295 | override_settings = {"agent_profile": profile} if profile else None |
| 296 | return initialize_agent(override_settings=override_settings) |
| 297 | |
| 298 | |
| 299 | def _deserialize_agents( |
| 300 | agents: list[dict[str, Any]], config: AgentConfig, context: AgentContext |
| 301 | ) -> Agent: |
| 302 | prev: Agent | None = None |
| 303 | zero: Agent | None = None |
| 304 | |
| 305 | for ag in agents: |
| 306 | current = Agent( |
| 307 | number=ag["number"], |
| 308 | config=_deserialize_agent_config(ag, config), |
| 309 | context=context, |
| 310 | ) |
| 311 | current.data = ag.get("data", {}) |
| 312 | current.history = history.deserialize_history( |
| 313 | ag.get("history", ""), agent=current |
| 314 | ) |
| 315 | if not zero: |
| 316 | zero = current |
| 317 | |
| 318 | if prev: |
| 319 | prev.set_data(Agent.DATA_NAME_SUBORDINATE, current) |
| 320 | current.set_data(Agent.DATA_NAME_SUPERIOR, prev) |
| 321 | prev = current |
| 322 | |
| 323 | return zero or Agent(0, config, context) |
| 324 | |
| 325 | |
| 326 | # def _deserialize_history(history: list[dict[str, Any]]): |
| 327 | # result = [] |
| 328 | # for hist in history: |
| 329 | # content = hist.get("content", "") |
| 330 | # msg = ( |
| 331 | # HumanMessage(content=content) |
| 332 | # if hist.get("type") == "human" |
| 333 | # else AIMessage(content=content) |
| 334 | # ) |
| 335 | # result.append(msg) |
| 336 | # return result |
| 337 | |
| 338 | |
| 339 | def _deserialize_log(data: dict[str, Any]) -> "Log": |
| 340 | log = Log() |
| 341 | log.guid = data.get("guid", str(uuid.uuid4())) |
| 342 | log.set_initial_progress() |
| 343 | |
| 344 | # Deserialize the list of LogItem objects |
| 345 | i = 0 |
| 346 | for item_data in data.get("logs", []): |
| 347 | agentno = item_data.get("agentno") |
| 348 | if agentno is None: |
| 349 | agentno = item_data.get("agent_number", 0) |
| 350 | log.logs.append( |
| 351 | LogItem( |
| 352 | log=log, # restore the log reference |
| 353 | no=i, # item_data["no"], |
| 354 | type=item_data["type"], |
| 355 | heading=item_data.get("heading", ""), |
| 356 | content=item_data.get("content", ""), |
| 357 | kvps=OrderedDict(item_data["kvps"]) if item_data["kvps"] else None, |
| 358 | timestamp=item_data.get("timestamp", 0.0), |
| 359 | agentno=agentno, |
| 360 | id=item_data.get("id"), |
| 361 | ) |
| 362 | ) |
| 363 | log.updates.append(i) |
| 364 | i += 1 |
| 365 | |
| 366 | return log |
| 367 | |
| 368 | |
| 369 | def _safe_json_serialize(obj, **kwargs): |
| 370 | def serializer(o): |
| 371 | if isinstance(o, dict): |
| 372 | return {k: v for k, v in o.items() if is_json_serializable(v)} |
| 373 | elif isinstance(o, (list, tuple)): |
| 374 | return [item for item in o if is_json_serializable(item)] |
| 375 | elif is_json_serializable(o): |
| 376 | return o |
| 377 | else: |
| 378 | return None # Skip this property |
| 379 | |
| 380 | def is_json_serializable(item): |
| 381 | try: |
| 382 | json.dumps(item) |
| 383 | return True |
| 384 | except (TypeError, OverflowError): |
| 385 | return False |
| 386 | |
| 387 | return json.dumps(obj, default=serializer, **kwargs) |
| 388 | |
| 389 | |
| 390 | def _delete_provider_responses_for_chat(ctxid: str) -> None: |
| 391 | try: |
| 392 | data = json.loads(files.read_file(_get_chat_file_path(ctxid))) |
| 393 | except Exception: |
| 394 | return |
| 395 | if _responses_delete_disabled(data): |
| 396 | return |
| 397 | response_ids = _collect_response_ids(data) |
| 398 | if not response_ids: |
| 399 | return |
| 400 | delete_stored_response_ids(response_ids) |
| 401 | |
| 402 | |
| 403 | def _responses_delete_disabled(data: dict[str, Any]) -> bool: |
| 404 | if data.get("responses_delete_on_chat_delete") is False: |
| 405 | return True |
| 406 | context_data = data.get("data") |
| 407 | if isinstance(context_data, dict) and context_data.get("responses_delete_on_chat_delete") is False: |
| 408 | return True |
| 409 | for agent_data in data.get("agents", []) or []: |
| 410 | if not isinstance(agent_data, dict): |
| 411 | continue |
| 412 | state = agent_data.get("data") |
| 413 | if isinstance(state, dict) and state.get("responses_delete_on_chat_delete") is False: |
| 414 | return True |
| 415 | return False |
| 416 | |
| 417 | |
| 418 | def _collect_response_ids(data: Any) -> list[str]: |
| 419 | found: list[str] = [] |
| 420 | seen: set[str] = set() |
| 421 | |
| 422 | def add(value: Any) -> None: |
| 423 | response_id = str(value or "").strip() |
| 424 | if response_id and response_id not in seen: |
| 425 | seen.add(response_id) |
| 426 | found.append(response_id) |
| 427 | |
| 428 | def walk(obj: Any) -> None: |
| 429 | if isinstance(obj, dict): |
| 430 | state = obj.get(Agent.DATA_NAME_RESPONSES_STATE) |
| 431 | if isinstance(state, dict): |
| 432 | add(state.get("response_id")) |
| 433 | for response_id in state.get("response_ids") or []: |
| 434 | add(response_id) |
| 435 | |
| 436 | metadata = obj.get("metadata") |
| 437 | if isinstance(metadata, dict): |
| 438 | responses = metadata.get("responses") |
| 439 | if isinstance(responses, dict): |
| 440 | add(responses.get("response_id")) |
| 441 | |
| 442 | for value in obj.values(): |
| 443 | walk(value) |
| 444 | elif isinstance(obj, list): |
| 445 | for value in obj: |
| 446 | walk(value) |
| 447 | elif isinstance(obj, str) and '"response_id"' in obj: |
| 448 | try: |
| 449 | walk(json.loads(obj)) |
| 450 | except Exception: |
| 451 | return |
| 452 | |
| 453 | walk(data) |
| 454 | return found |