| 1 | import os |
| 2 | import uuid |
| 3 | from typing import TYPE_CHECKING |
| 4 | from helpers import guids |
| 5 | |
| 6 | if TYPE_CHECKING: |
| 7 | from agent import AgentContext |
| 8 | |
| 9 | from helpers.print_style import PrintStyle |
| 10 | |
| 11 | QUEUE_KEY = "message_queue" |
| 12 | QUEUE_SEQ_KEY = "message_queue_seq" |
| 13 | UPLOAD_FOLDER = "/a0/usr/uploads" |
| 14 | |
| 15 | |
| 16 | def get_queue(context: "AgentContext") -> list: |
| 17 | """Get current queue from context.data.""" |
| 18 | return context.get_data(QUEUE_KEY) or [] |
| 19 | |
| 20 | |
| 21 | def _get_next_seq(context: "AgentContext") -> int: |
| 22 | """Get next sequence number.""" |
| 23 | seq = context.get_data(QUEUE_SEQ_KEY) or 0 |
| 24 | seq += 1 |
| 25 | context.set_data(QUEUE_SEQ_KEY, seq) |
| 26 | return seq |
| 27 | |
| 28 | |
| 29 | def _sync_output(context: "AgentContext"): |
| 30 | """Sync queue to output_data for frontend polling.""" |
| 31 | queue = get_queue(context) |
| 32 | # Truncate text for frontend display |
| 33 | truncated = [] |
| 34 | for item in queue: |
| 35 | truncated.append({ |
| 36 | "id": item["id"], |
| 37 | "seq": item.get("seq", 0), |
| 38 | "text": item["text"][:100] + "..." if len(item["text"]) > 100 else item["text"], |
| 39 | "attachments": [a.split("/")[-1] for a in item.get("attachments", [])], |
| 40 | "attachment_count": len(item.get("attachments", [])), |
| 41 | }) |
| 42 | context.set_output_data(QUEUE_KEY, truncated) |
| 43 | |
| 44 | |
| 45 | def add( |
| 46 | context: "AgentContext", |
| 47 | text: str, |
| 48 | attachments: list[str] | None = None, |
| 49 | item_id: str | None = None, |
| 50 | ) -> dict: |
| 51 | """Add message to queue. Attachments should be filenames, will be converted to full paths.""" |
| 52 | queue = get_queue(context) |
| 53 | |
| 54 | # Convert filenames to full paths |
| 55 | full_paths = [] |
| 56 | for att in (attachments or []): |
| 57 | if att.startswith("/"): |
| 58 | full_paths.append(att) |
| 59 | else: |
| 60 | full_paths.append(f"{UPLOAD_FOLDER}/{att}") |
| 61 | |
| 62 | item = { |
| 63 | "id": item_id or guids.generate_id(), |
| 64 | "seq": _get_next_seq(context), |
| 65 | "text": text, |
| 66 | "attachments": full_paths, |
| 67 | } |
| 68 | queue.append(item) |
| 69 | context.set_data(QUEUE_KEY, queue) |
| 70 | _sync_output(context) |
| 71 | return item |
| 72 | |
| 73 | |
| 74 | def remove(context: "AgentContext", item_id: str | None = None) -> int: |
| 75 | """Remove item(s). If item_id is None, clears all. Returns remaining count.""" |
| 76 | if not item_id: |
| 77 | context.set_data(QUEUE_KEY, []) |
| 78 | context.set_output_data(QUEUE_KEY, []) |
| 79 | return 0 |
| 80 | queue = [i for i in get_queue(context) if i["id"] != item_id] |
| 81 | context.set_data(QUEUE_KEY, queue) |
| 82 | _sync_output(context) |
| 83 | return len(queue) |
| 84 | |
| 85 | |
| 86 | def pop_first(context: "AgentContext") -> dict | None: |
| 87 | """Remove and return first item.""" |
| 88 | queue = get_queue(context) |
| 89 | if not queue: |
| 90 | return None |
| 91 | item = queue.pop(0) |
| 92 | context.set_data(QUEUE_KEY, queue) |
| 93 | _sync_output(context) |
| 94 | return item |
| 95 | |
| 96 | |
| 97 | def pop_item(context: "AgentContext", item_id: str) -> dict | None: |
| 98 | """Remove and return specific item.""" |
| 99 | queue = get_queue(context) |
| 100 | for i, item in enumerate(queue): |
| 101 | if item["id"] == item_id: |
| 102 | queue.pop(i) |
| 103 | context.set_data(QUEUE_KEY, queue) |
| 104 | _sync_output(context) |
| 105 | return item |
| 106 | return None |
| 107 | |
| 108 | |
| 109 | def has_queue(context: "AgentContext") -> bool: |
| 110 | """Check if queue has items.""" |
| 111 | return len(get_queue(context)) > 0 |
| 112 | |
| 113 | |
| 114 | def log_user_message( |
| 115 | context: "AgentContext", |
| 116 | message: str, |
| 117 | attachment_paths: list[str], |
| 118 | message_id: str | None = None, |
| 119 | source: str = "", |
| 120 | ): |
| 121 | """Log user message to console and UI. Used by message API and queue processing.""" |
| 122 | # Prepare attachment filenames for logging |
| 123 | attachment_filenames = ( |
| 124 | [os.path.basename(path) for path in attachment_paths] |
| 125 | if attachment_paths |
| 126 | else [] |
| 127 | ) |
| 128 | |
| 129 | # Print to console |
| 130 | label = f"User message{source}:" |
| 131 | PrintStyle( |
| 132 | background_color="#6C3483", font_color="white", bold=True, padding=True |
| 133 | ).print(label) |
| 134 | PrintStyle(font_color="white", padding=False).print(f"> {message}") |
| 135 | if attachment_filenames: |
| 136 | PrintStyle(font_color="white", padding=False).print("Attachments:") |
| 137 | for filename in attachment_filenames: |
| 138 | PrintStyle(font_color="white", padding=False).print(f"- {filename}") |
| 139 | |
| 140 | # Log to UI |
| 141 | context.log.log( |
| 142 | type="user", |
| 143 | heading="", |
| 144 | content=message, |
| 145 | kvps={"attachments": attachment_filenames}, |
| 146 | id=message_id, |
| 147 | ) |
| 148 | |
| 149 | |
| 150 | def send_message(context: "AgentContext", item: dict, source: str = " (from queue)"): |
| 151 | """Send a single queued message (log + communicate).""" |
| 152 | from agent import UserMessage # Import here to avoid circular import |
| 153 | |
| 154 | message = item.get("text", "") |
| 155 | attachments = item.get("attachments", []) |
| 156 | msg_id = str(uuid.uuid4()) |
| 157 | log_user_message(context, message, attachments, message_id=msg_id, source=source) |
| 158 | context.communicate(UserMessage(message, attachments, id=msg_id)) |
| 159 | |
| 160 | |
| 161 | def send_next(context: "AgentContext") -> bool: |
| 162 | """Send next queued message. Returns True if sent, False if queue empty.""" |
| 163 | if not has_queue(context): |
| 164 | return False |
| 165 | item = pop_first(context) |
| 166 | if item: |
| 167 | send_message(context, item) |
| 168 | return True |
| 169 | return False |
| 170 | |
| 171 | |
| 172 | def send_all_aggregated(context: "AgentContext") -> int: |
| 173 | """Aggregate and send all queued messages as one. Returns count of items sent.""" |
| 174 | from agent import UserMessage # Import here to avoid circular import |
| 175 | |
| 176 | if not has_queue(context): |
| 177 | return 0 |
| 178 | |
| 179 | items = [] |
| 180 | while has_queue(context): |
| 181 | items.append(pop_first(context)) |
| 182 | |
| 183 | # Combine texts with separator |
| 184 | text = "\n\n---\n\n".join(i["text"] for i in items if i["text"]) |
| 185 | attachments = [a for i in items for a in i.get("attachments", [])] |
| 186 | |
| 187 | msg_id = str(uuid.uuid4()) |
| 188 | log_user_message(context, text, attachments, message_id=msg_id, source=" (queued batch)") |
| 189 | context.communicate(UserMessage(text, attachments, id=msg_id)) |
| 190 | return len(items) |