main
py 574 lines 18.1 KB
Raw
1 from __future__ import annotations
2
3 import json
4 from dataclasses import dataclass
5
6 from agent import AgentContext
7 from helpers import files, projects, subagents
8 from helpers import integration_commands
9 from helpers.persist_chat import save_tmp_chat
10 from helpers.state_monitor_integration import mark_dirty_for_context
11 from plugins._model_config.helpers import model_config
12 from plugins._telegram_integration.helpers import telegram_client as tc
13 from plugins._telegram_integration.helpers.constants import (
14 CTX_TG_STREAM_ENABLED,
15 CTX_TG_TOOLS_ENABLED,
16 CTX_TG_BOT,
17 CTX_TG_CHAT_ID,
18 CTX_TG_CHAT_TYPE,
19 CTX_TG_USER_ID,
20 CTX_TG_USERNAME,
21 STATE_FILE,
22 )
23
24 PAGE_SIZE = 8
25 SESSION_PAGE_SIZE = 4
26 CALLBACK_PREFIX = "tg"
27
28
29 @dataclass(frozen=True)
30 class PickerItem:
31 key: str
32 label: str
33
34
35 @dataclass(frozen=True)
36 class SessionItem:
37 context: AgentContext
38 label: str
39 last_message: str
40 running: bool
41
42
43 async def handle_command(
44 context: AgentContext,
45 token: str,
46 chat_id: int,
47 reply_to_message_id: int | None,
48 text: str,
49 ) -> bool:
50 parsed = integration_commands.parse_command(text or "", integration="telegram")
51 if not parsed:
52 return False
53 command, args = parsed
54 if command in {"/model", "/config", "/preset"} and not args:
55 await send_model_picker(context, token, chat_id, reply_to_message_id, 0)
56 return True
57 if command == "/project" and not args:
58 await send_project_picker(context, token, chat_id, reply_to_message_id, 0)
59 return True
60 if command in {"/agent", "/profile"} and not args:
61 await send_agent_picker(context, token, chat_id, reply_to_message_id, 0)
62 return True
63 if command in {"/sessions", "/session"} and not args:
64 await send_session_picker(context, token, chat_id, reply_to_message_id, 0)
65 return True
66 if command == "/stream":
67 await send_toggle_picker(
68 context,
69 token,
70 chat_id,
71 reply_to_message_id,
72 CTX_TG_STREAM_ENABLED,
73 "Response streaming",
74 args,
75 )
76 return True
77 if command == "/tools":
78 await send_toggle_picker(
79 context,
80 token,
81 chat_id,
82 reply_to_message_id,
83 CTX_TG_TOOLS_ENABLED,
84 "Tool progress",
85 args,
86 )
87 return True
88 return False
89
90
91 async def handle_callback(
92 context: AgentContext,
93 token: str,
94 chat_id: int,
95 message_id: int,
96 data: str,
97 ) -> bool:
98 parts = (data or "").split(":")
99 if len(parts) < 3 or parts[0] != CALLBACK_PREFIX:
100 return False
101 kind, action = parts[1], parts[2]
102 value = parts[3] if len(parts) > 3 else ""
103 if action == "noop":
104 return True
105 if action == "page":
106 page = _safe_int(value)
107 if kind == "model":
108 await edit_model_picker(context, token, chat_id, message_id, page)
109 elif kind == "project":
110 await edit_project_picker(context, token, chat_id, message_id, page)
111 elif kind == "agent":
112 await edit_agent_picker(context, token, chat_id, message_id, page)
113 elif kind == "session":
114 await edit_session_picker(context, token, chat_id, message_id, page)
115 return True
116 if kind == "model" and action in {"set", "clear"}:
117 await _select_model(context, _safe_int(value), clear=(action == "clear"))
118 await edit_model_picker(context, token, chat_id, message_id, 0, selected=True)
119 return True
120 if kind == "project" and action in {"set", "clear"}:
121 await _select_project(context, _safe_int(value), clear=(action == "clear"))
122 await edit_project_picker(context, token, chat_id, message_id, 0, selected=True)
123 return True
124 if kind == "agent" and action == "set":
125 await _select_agent(context, _safe_int(value))
126 await edit_agent_picker(context, token, chat_id, message_id, 0, selected=True)
127 return True
128 if kind == "session" and action == "set":
129 selected_context = await _select_session(context, _safe_int(value))
130 await edit_session_picker(selected_context or context, token, chat_id, message_id, 0, selected=bool(selected_context))
131 return True
132 if kind in {"stream", "tools"} and action in {"on", "off"}:
133 key = CTX_TG_STREAM_ENABLED if kind == "stream" else CTX_TG_TOOLS_ENABLED
134 label = "Response streaming" if kind == "stream" else "Tool progress"
135 context.set_data(key, action == "on")
136 save_tmp_chat(context)
137 mark_dirty_for_context(context.id, reason=f"telegram.{kind}_toggle")
138 await edit_toggle_picker(context, token, chat_id, message_id, key, label)
139 return True
140 return True
141
142
143 async def send_model_picker(
144 context: AgentContext,
145 token: str,
146 chat_id: int,
147 reply_to_message_id: int | None,
148 page: int,
149 ) -> None:
150 text, markup = _model_view(context, page)
151 await tc.raw_send_text(token, chat_id, text, reply_to_message_id, "HTML", markup)
152
153
154 async def edit_model_picker(
155 context: AgentContext,
156 token: str,
157 chat_id: int,
158 message_id: int,
159 page: int,
160 *,
161 selected: bool = False,
162 ) -> None:
163 text, markup = _model_view(context, page, selected=selected)
164 await tc.raw_edit_text(token, chat_id, message_id, text, "HTML", markup)
165
166
167 async def send_project_picker(
168 context: AgentContext,
169 token: str,
170 chat_id: int,
171 reply_to_message_id: int | None,
172 page: int,
173 ) -> None:
174 text, markup = _project_view(context, page)
175 await tc.raw_send_text(token, chat_id, text, reply_to_message_id, "HTML", markup)
176
177
178 async def edit_project_picker(
179 context: AgentContext,
180 token: str,
181 chat_id: int,
182 message_id: int,
183 page: int,
184 *,
185 selected: bool = False,
186 ) -> None:
187 text, markup = _project_view(context, page, selected=selected)
188 await tc.raw_edit_text(token, chat_id, message_id, text, "HTML", markup)
189
190
191 async def send_agent_picker(
192 context: AgentContext,
193 token: str,
194 chat_id: int,
195 reply_to_message_id: int | None,
196 page: int,
197 ) -> None:
198 text, markup = _agent_view(context, page)
199 await tc.raw_send_text(token, chat_id, text, reply_to_message_id, "HTML", markup)
200
201
202 async def edit_agent_picker(
203 context: AgentContext,
204 token: str,
205 chat_id: int,
206 message_id: int,
207 page: int,
208 *,
209 selected: bool = False,
210 ) -> None:
211 text, markup = _agent_view(context, page, selected=selected)
212 await tc.raw_edit_text(token, chat_id, message_id, text, "HTML", markup)
213
214
215 async def send_toggle_picker(
216 context: AgentContext,
217 token: str,
218 chat_id: int,
219 reply_to_message_id: int | None,
220 key: str,
221 label: str,
222 args: str = "",
223 ) -> None:
224 desired = _parse_toggle(args)
225 if desired is not None:
226 context.set_data(key, desired)
227 save_tmp_chat(context)
228 mark_dirty_for_context(context.id, reason=f"telegram.{key}")
229 text, markup = _toggle_view(context, key, label)
230 await tc.raw_send_text(token, chat_id, text, reply_to_message_id, "HTML", markup)
231
232
233 async def edit_toggle_picker(
234 context: AgentContext,
235 token: str,
236 chat_id: int,
237 message_id: int,
238 key: str,
239 label: str,
240 ) -> None:
241 text, markup = _toggle_view(context, key, label)
242 await tc.raw_edit_text(token, chat_id, message_id, text, "HTML", markup)
243
244
245 async def send_session_picker(
246 context: AgentContext,
247 token: str,
248 chat_id: int,
249 reply_to_message_id: int | None,
250 page: int,
251 ) -> None:
252 text, markup = _session_view(context, page)
253 await tc.raw_send_text(token, chat_id, text, reply_to_message_id, "HTML", markup)
254
255
256 async def edit_session_picker(
257 context: AgentContext,
258 token: str,
259 chat_id: int,
260 message_id: int,
261 page: int,
262 *,
263 selected: bool = False,
264 ) -> None:
265 text, markup = _session_view(context, page, selected=selected)
266 await tc.raw_edit_text(token, chat_id, message_id, text, "HTML", markup)
267
268
269 def _model_view(
270 context: AgentContext,
271 page: int,
272 *,
273 selected: bool = False,
274 ) -> tuple[str, dict | None]:
275 presets = [
276 PickerItem(str(preset.get("name", "")), str(preset.get("name", "")))
277 for preset in model_config.get_presets()
278 if isinstance(preset, dict) and preset.get("name")
279 ]
280 current = context.get_data("chat_model_override")
281 current_name = current.get("preset_name") if isinstance(current, dict) else ""
282 effective_name = model_config.get_effective_preset_name(context.agent0)
283 status = f"Current model: <b>{_html(effective_name)}</b>"
284 if not model_config.is_chat_override_allowed(context.agent0):
285 return status + "\nPer-chat model switching is disabled.", None
286 if selected:
287 status = "Model updated.\n" + status
288 rows = _paged_buttons("model", presets, page, effective_name)
289 rows.append([{"text": "Use scoped preset", "callback_data": "tg:model:clear"}])
290 return status, {"inline_keyboard": rows}
291
292
293 def _project_view(
294 context: AgentContext,
295 page: int,
296 *,
297 selected: bool = False,
298 ) -> tuple[str, dict | None]:
299 items = [
300 PickerItem(str(item.get("name", "")), str(item.get("title") or item.get("name") or ""))
301 for item in projects.get_active_projects_list() or []
302 if item.get("name")
303 ]
304 current = context.get_data("project") or ""
305 status = f"Current project: <b>{_html(_label_for(items, current) or 'none')}</b>"
306 if selected:
307 status = "Project updated.\n" + status
308 rows = _paged_buttons("project", items, page, current)
309 rows.append([{"text": "No project", "callback_data": "tg:project:clear"}])
310 return status, {"inline_keyboard": rows}
311
312
313 def _agent_view(
314 context: AgentContext,
315 page: int,
316 *,
317 selected: bool = False,
318 ) -> tuple[str, dict | None]:
319 items = [
320 PickerItem(str(item.get("key", "")), str(item.get("label") or item.get("key") or ""))
321 for item in subagents.get_all_agents_list()
322 if item.get("key")
323 ]
324 current = getattr(context.agent0.config, "profile", "") or "agent0"
325 status = f"Current agent: <b>{_html(_label_for(items, current) or current)}</b>"
326 if context.is_running():
327 status += "\nAgent profile can be changed after the current run finishes."
328 elif selected:
329 status = "Agent updated.\n" + status
330 rows = _paged_buttons("agent", items, page, current, disabled=context.is_running())
331 if not rows:
332 return status + "\nNo agent profiles were found.", None
333 return status, {"inline_keyboard": rows}
334
335
336 def _toggle_view(context: AgentContext, key: str, label: str) -> tuple[str, dict]:
337 enabled = _toggle_enabled(context, key)
338 kind = "stream" if key == CTX_TG_STREAM_ENABLED else "tools"
339 state = "enabled" if enabled else "disabled"
340 text = f"{_html(label)}: <b>{state}</b>"
341 rows = [[
342 {"text": ("On" if enabled else "Turn on"), "callback_data": f"tg:{kind}:on"},
343 {"text": ("Off" if not enabled else "Turn off"), "callback_data": f"tg:{kind}:off"},
344 ]]
345 return text, {"inline_keyboard": rows}
346
347
348 def _session_view(
349 context: AgentContext,
350 page: int,
351 *,
352 selected: bool = False,
353 ) -> tuple[str, dict | None]:
354 items = _session_items()
355 current_label = _session_label(context)
356 status = f"Current session: <b>{_html(current_label)}</b>"
357 if context.is_running():
358 status += "\nSession switching is available after the current run finishes."
359 elif selected:
360 status = "Session switched.\n" + status
361 if not items:
362 return status + "\nNo sessions were found.", None
363
364 total = len(items)
365 page = _clamp_page(page, total, SESSION_PAGE_SIZE)
366 start = page * SESSION_PAGE_SIZE
367 end = min(start + SESSION_PAGE_SIZE, total)
368 rows: list[list[dict[str, str]]] = []
369 for index, item in enumerate(items[start:end], start=start):
370 marker = "" if item.context.id == context.id else ""
371 suffix = " (running)" if item.running else ""
372 action = "noop" if context.is_running() or item.running else "set"
373 rows.append([{
374 "text": f"{marker}{item.label}{suffix}"[:64],
375 "callback_data": f"tg:session:{action}:{index}",
376 }])
377
378 nav: list[dict[str, str]] = []
379 if page > 0:
380 nav.append({"text": "Prev", "callback_data": f"tg:session:page:{page - 1}"})
381 if end < total:
382 nav.append({"text": "Next", "callback_data": f"tg:session:page:{page + 1}"})
383 if nav:
384 rows.append(nav)
385
386 range_text = f"\nShowing {start + 1}-{end} of {total}."
387 return status + range_text, {"inline_keyboard": rows}
388
389
390 async def _select_model(context: AgentContext, index: int, *, clear: bool = False) -> None:
391 if not model_config.is_chat_override_allowed(context.agent0):
392 return
393 if clear:
394 context.set_data("chat_model_override", None)
395 else:
396 presets = [preset for preset in model_config.get_presets() if preset.get("name")]
397 if index < 0 or index >= len(presets):
398 return
399 context.set_data("chat_model_override", {"preset_name": presets[index]["name"]})
400 save_tmp_chat(context)
401 mark_dirty_for_context(context.id, reason="telegram.model_select")
402
403
404 async def _select_project(context: AgentContext, index: int, *, clear: bool = False) -> None:
405 if clear:
406 projects.deactivate_project(context.id)
407 return
408 items = [item for item in projects.get_active_projects_list() or [] if item.get("name")]
409 if index < 0 or index >= len(items):
410 return
411 projects.activate_project(context.id, str(items[index]["name"]))
412
413
414 async def _select_agent(context: AgentContext, index: int) -> None:
415 if context.is_running():
416 return
417 from initialize import initialize_agent
418
419 items = [item for item in subagents.get_all_agents_list() if item.get("key")]
420 if index < 0 or index >= len(items):
421 return
422 profile = str(items[index]["key"])
423 config = initialize_agent(override_settings={"agent_profile": profile})
424 context.config = config
425 context.agent0.config = config
426 save_tmp_chat(context)
427 mark_dirty_for_context(context.id, reason="telegram.agent_select")
428
429
430 async def _select_session(context: AgentContext, index: int) -> AgentContext | None:
431 if context.is_running():
432 return None
433 items = _session_items()
434 if index < 0 or index >= len(items):
435 return None
436 target = items[index].context
437 if target.is_running():
438 return None
439
440 _copy_telegram_binding(context, target)
441 _set_session_mapping(target)
442 save_tmp_chat(target)
443 mark_dirty_for_context(context.id, reason="telegram.session_unselect")
444 mark_dirty_for_context(target.id, reason="telegram.session_select")
445 return target
446
447
448 def _session_items() -> list[SessionItem]:
449 contexts = sorted(
450 AgentContext.all(),
451 key=lambda item: str(item.output().get("last_message") or ""),
452 reverse=True,
453 )
454 return [
455 SessionItem(
456 context=item,
457 label=_session_label(item),
458 last_message=str(item.output().get("last_message") or ""),
459 running=item.is_running(),
460 )
461 for item in contexts
462 ]
463
464
465 def _session_label(context: AgentContext) -> str:
466 return str(context.name or context.id or "Session")
467
468
469 def _copy_telegram_binding(source: AgentContext, target: AgentContext) -> None:
470 for key in (
471 CTX_TG_BOT,
472 CTX_TG_CHAT_ID,
473 CTX_TG_CHAT_TYPE,
474 CTX_TG_USER_ID,
475 CTX_TG_USERNAME,
476 ):
477 if key in source.data:
478 target.data[key] = source.data[key]
479
480
481 def _set_session_mapping(context: AgentContext) -> None:
482 bot_name = str(context.data.get(CTX_TG_BOT) or "")
483 user_id = context.data.get(CTX_TG_USER_ID)
484 chat_id = context.data.get(CTX_TG_CHAT_ID)
485 if not bot_name or user_id is None or chat_id is None:
486 return
487 key = f"{bot_name}:{int(user_id)}:{int(chat_id)}"
488 state = _load_telegram_state()
489 chats = state.setdefault("chats", {})
490 chats[key] = context.id
491 _save_telegram_state(state)
492
493
494 def _load_telegram_state() -> dict:
495 path = files.get_abs_path(STATE_FILE)
496 if not files.exists(path):
497 return {}
498 try:
499 return json.loads(files.read_file(path))
500 except Exception:
501 return {}
502
503
504 def _save_telegram_state(state: dict) -> None:
505 path = files.get_abs_path(STATE_FILE)
506 files.make_dirs(path)
507 files.write_file(path, json.dumps(state))
508
509
510 def _paged_buttons(
511 kind: str,
512 items: list[PickerItem],
513 page: int,
514 current: str,
515 *,
516 disabled: bool = False,
517 ) -> list[list[dict[str, str]]]:
518 page = max(0, page)
519 total_pages = max(1, (len(items) + PAGE_SIZE - 1) // PAGE_SIZE)
520 page = min(page, total_pages - 1)
521 start = page * PAGE_SIZE
522 rows: list[list[dict[str, str]]] = []
523 for offset, item in enumerate(items[start:start + PAGE_SIZE], start=start):
524 marker = "" if item.key == current else ""
525 action = "noop" if disabled else "set"
526 rows.append([{
527 "text": f"{marker}{item.label}"[:64],
528 "callback_data": f"tg:{kind}:{action}:{offset}",
529 }])
530 nav: list[dict[str, str]] = []
531 if page > 0:
532 nav.append({"text": "Prev", "callback_data": f"tg:{kind}:page:{page - 1}"})
533 if page < total_pages - 1:
534 nav.append({"text": "Next", "callback_data": f"tg:{kind}:page:{page + 1}"})
535 if nav:
536 rows.append(nav)
537 return rows
538
539
540 def _toggle_enabled(context: AgentContext, key: str) -> bool:
541 value = context.get_data(key)
542 return True if value is None else bool(value)
543
544
545 def _parse_toggle(args: str) -> bool | None:
546 value = (args or "").strip().lower()
547 if value in {"on", "enable", "enabled", "yes", "true", "1"}:
548 return True
549 if value in {"off", "disable", "disabled", "no", "false", "0"}:
550 return False
551 return None
552
553
554 def _safe_int(value: str) -> int:
555 try:
556 return int(value)
557 except (TypeError, ValueError):
558 return 0
559
560
561 def _clamp_page(page: int, total_items: int, page_size: int) -> int:
562 total_pages = max(1, (total_items + page_size - 1) // page_size)
563 return min(max(0, page), total_pages - 1)
564
565
566 def _label_for(items: list[PickerItem], key: str) -> str:
567 for item in items:
568 if item.key == key:
569 return item.label
570 return key
571
572
573 def _html(value: str) -> str:
574 return str(value).replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")