263 lines
10 KiB
Python
263 lines
10 KiB
Python
|
# retoor <retoor@molodetz.nl>
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
import html
|
||
|
|
import logging
|
||
|
|
import time
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
from devplacepy.database import get_table
|
||
|
|
from devplacepy.services.audit import record as audit
|
||
|
|
from devplacepy.services.manager import service_manager
|
||
|
|
from devplacepy.utils import is_admin, is_primary_admin
|
||
|
|
|
||
|
|
from . import store
|
||
|
|
from .format import markdown_to_telegram_html, split_for_telegram
|
||
|
|
|
||
|
|
logger = logging.getLogger("telegram.bridge")
|
||
|
|
|
||
|
|
TURN_TIMEOUT_SECONDS = 300.0
|
||
|
|
TYPING_INTERVAL_SECONDS = 4.0
|
||
|
|
PROGRESS_EDIT_INTERVAL_SECONDS = 1.5
|
||
|
|
ATTEMPT_LIMIT = 5
|
||
|
|
ATTEMPT_WINDOW_SECONDS = 600.0
|
||
|
|
THINKING_PLACEHOLDER = "<i>Devii is thinking...</i>"
|
||
|
|
|
||
|
|
WELCOME = (
|
||
|
|
"Welcome to Devii on Telegram. To connect your DevPlace account, open your profile, "
|
||
|
|
"request a Telegram pairing code, and send me the four digit code here."
|
||
|
|
)
|
||
|
|
PAIRED_TEMPLATE = "Paired. Hello {username}. Send me a message and I will help."
|
||
|
|
BAD_CODE = (
|
||
|
|
"That code is invalid or expired. Request a fresh code from your DevPlace profile and "
|
||
|
|
"send it here."
|
||
|
|
)
|
||
|
|
TOO_MANY = "Too many attempts. Wait a few minutes, request a new code, and try again."
|
||
|
|
|
||
|
|
|
||
|
|
class TelegramConnection:
|
||
|
|
def __init__(self, service: Any, session: Any, chat_id: int, message_id: int | None) -> None:
|
||
|
|
self._service = service
|
||
|
|
self._session = session
|
||
|
|
self._chat_id = chat_id
|
||
|
|
self._message_id = message_id
|
||
|
|
self.done = asyncio.Event()
|
||
|
|
self._finalized = False
|
||
|
|
self._last_status = ""
|
||
|
|
self._last_edit = 0.0
|
||
|
|
self._typing_task = asyncio.create_task(self._typing_loop())
|
||
|
|
|
||
|
|
async def send_json(self, payload: dict[str, Any]) -> None:
|
||
|
|
kind = payload.get("type")
|
||
|
|
if kind == "reply":
|
||
|
|
await self._finalize(markdown_to_telegram_html(payload.get("text", "")) or "(no response)")
|
||
|
|
elif kind == "error":
|
||
|
|
await self._finalize(html.escape("Error: " + str(payload.get("text", ""))))
|
||
|
|
elif kind in ("status", "trace"):
|
||
|
|
await self._progress(payload)
|
||
|
|
elif kind in ("avatar", "client"):
|
||
|
|
query_id = payload.get("id")
|
||
|
|
if query_id:
|
||
|
|
self._session.resolve_query(
|
||
|
|
query_id, {"error": "No browser is attached (Telegram session)."}
|
||
|
|
)
|
||
|
|
|
||
|
|
async def _typing_loop(self) -> None:
|
||
|
|
try:
|
||
|
|
while not self.done.is_set():
|
||
|
|
await self._service.chat_action(self._chat_id, "typing")
|
||
|
|
await asyncio.sleep(TYPING_INTERVAL_SECONDS)
|
||
|
|
except asyncio.CancelledError:
|
||
|
|
pass
|
||
|
|
except Exception: # noqa: BLE001 - typing is cosmetic, never fail a turn over it
|
||
|
|
pass
|
||
|
|
|
||
|
|
async def _progress(self, payload: dict[str, Any]) -> None:
|
||
|
|
if self._finalized or self._message_id is None:
|
||
|
|
return
|
||
|
|
if payload.get("type") == "trace":
|
||
|
|
if payload.get("event") != "call":
|
||
|
|
return
|
||
|
|
text = f"Working... ({payload.get('name', '')})"
|
||
|
|
else:
|
||
|
|
text = str(payload.get("text", ""))
|
||
|
|
if not text or text == self._last_status:
|
||
|
|
return
|
||
|
|
now = time.monotonic()
|
||
|
|
if now - self._last_edit < PROGRESS_EDIT_INTERVAL_SECONDS:
|
||
|
|
return
|
||
|
|
self._last_status = text
|
||
|
|
self._last_edit = now
|
||
|
|
await self._service.edit(
|
||
|
|
self._chat_id, self._message_id, f"<i>{html.escape(text)}</i>"
|
||
|
|
)
|
||
|
|
|
||
|
|
async def _finalize(self, html_text: str) -> None:
|
||
|
|
if self._finalized:
|
||
|
|
return
|
||
|
|
self._finalized = True
|
||
|
|
self._stop_typing()
|
||
|
|
chunks = split_for_telegram(html_text)
|
||
|
|
first = chunks[0] if chunks else "(no response)"
|
||
|
|
if self._message_id is not None:
|
||
|
|
ok = await self._service.edit(self._chat_id, self._message_id, first)
|
||
|
|
if not ok:
|
||
|
|
await self._service.send(self._chat_id, first)
|
||
|
|
else:
|
||
|
|
await self._service.send(self._chat_id, first)
|
||
|
|
for chunk in chunks[1:]:
|
||
|
|
await self._service.send(self._chat_id, chunk)
|
||
|
|
self.done.set()
|
||
|
|
|
||
|
|
def _stop_typing(self) -> None:
|
||
|
|
if self._typing_task is not None and not self._typing_task.done():
|
||
|
|
self._typing_task.cancel()
|
||
|
|
|
||
|
|
|
||
|
|
class TelegramBridge:
|
||
|
|
def __init__(self, service: Any, max_concurrent: int = 8) -> None:
|
||
|
|
self._service = service
|
||
|
|
self._semaphore = asyncio.Semaphore(max(1, max_concurrent))
|
||
|
|
self._chat_locks: dict[int, asyncio.Lock] = {}
|
||
|
|
self._attempts: dict[int, list[float]] = {}
|
||
|
|
|
||
|
|
def _chat_lock(self, chat_id: int) -> asyncio.Lock:
|
||
|
|
lock = self._chat_locks.get(chat_id)
|
||
|
|
if lock is None:
|
||
|
|
lock = asyncio.Lock()
|
||
|
|
self._chat_locks[chat_id] = lock
|
||
|
|
return lock
|
||
|
|
|
||
|
|
async def handle_inbound(self, event: dict[str, Any]) -> None:
|
||
|
|
chat_id = int(event["chat_id"])
|
||
|
|
from_id = int(event["from_id"])
|
||
|
|
text = str(event.get("text", "")).strip()
|
||
|
|
images = [img for img in event.get("images", []) if img]
|
||
|
|
link = store.user_for_chat(chat_id)
|
||
|
|
if link is None:
|
||
|
|
await self._handle_unpaired(chat_id, from_id, text)
|
||
|
|
return
|
||
|
|
if int(link.get("from_id", 0)) != from_id:
|
||
|
|
return
|
||
|
|
user = get_table("users").find_one(uid=link["user_uid"])
|
||
|
|
if not user:
|
||
|
|
await self._service.send(
|
||
|
|
chat_id, "Your DevPlace account was not found. Please re-pair from your profile."
|
||
|
|
)
|
||
|
|
return
|
||
|
|
async with self._chat_lock(chat_id):
|
||
|
|
async with self._semaphore:
|
||
|
|
await self._run_turn(chat_id, user, text, images)
|
||
|
|
|
||
|
|
async def _handle_unpaired(self, chat_id: int, from_id: int, text: str) -> None:
|
||
|
|
if text.isdigit() and len(text) == 4:
|
||
|
|
if self._too_many_attempts(chat_id):
|
||
|
|
await self._service.send(chat_id, TOO_MANY)
|
||
|
|
return
|
||
|
|
user = store.verify_code(text, chat_id, from_id)
|
||
|
|
if user:
|
||
|
|
self._attempts.pop(chat_id, None)
|
||
|
|
await self._service.send(
|
||
|
|
chat_id, PAIRED_TEMPLATE.format(username=user.get("username", "there"))
|
||
|
|
)
|
||
|
|
audit.record_system(
|
||
|
|
"telegram.pair.success",
|
||
|
|
actor_kind="user",
|
||
|
|
actor_uid=user["uid"],
|
||
|
|
actor_username=user.get("username", ""),
|
||
|
|
summary=f"Telegram paired for {user.get('username', user['uid'])}",
|
||
|
|
metadata={"chat_id": chat_id},
|
||
|
|
)
|
||
|
|
return
|
||
|
|
self._register_attempt(chat_id)
|
||
|
|
await self._service.send(chat_id, BAD_CODE)
|
||
|
|
audit.record_system(
|
||
|
|
"telegram.pair.failure",
|
||
|
|
actor_kind="guest",
|
||
|
|
result="failure",
|
||
|
|
summary="Telegram pairing code rejected",
|
||
|
|
metadata={"chat_id": chat_id},
|
||
|
|
)
|
||
|
|
return
|
||
|
|
await self._service.send(chat_id, WELCOME)
|
||
|
|
|
||
|
|
def _too_many_attempts(self, chat_id: int) -> bool:
|
||
|
|
now = time.monotonic()
|
||
|
|
recent = [t for t in self._attempts.get(chat_id, []) if now - t < ATTEMPT_WINDOW_SECONDS]
|
||
|
|
self._attempts[chat_id] = recent
|
||
|
|
return len(recent) >= ATTEMPT_LIMIT
|
||
|
|
|
||
|
|
def _register_attempt(self, chat_id: int) -> None:
|
||
|
|
self._attempts.setdefault(chat_id, []).append(time.monotonic())
|
||
|
|
|
||
|
|
async def _run_turn(
|
||
|
|
self, chat_id: int, user: dict[str, Any], text: str, images: list[str]
|
||
|
|
) -> None:
|
||
|
|
devii = service_manager.get_service("devii")
|
||
|
|
if devii is None or not devii.is_enabled():
|
||
|
|
await self._service.send(chat_id, "Devii is currently unavailable.")
|
||
|
|
return
|
||
|
|
owner_id = user["uid"]
|
||
|
|
owner_is_admin = is_admin(user)
|
||
|
|
owner_is_primary_admin = is_primary_admin(user)
|
||
|
|
if devii.quota_exceeded("user", owner_id, owner_is_admin):
|
||
|
|
limit = devii.daily_limit_for("user", owner_is_admin)
|
||
|
|
audit.record_system(
|
||
|
|
"ai.quota.exceeded",
|
||
|
|
actor_kind="user",
|
||
|
|
actor_uid=owner_id,
|
||
|
|
actor_username=user.get("username", ""),
|
||
|
|
actor_role="admin" if owner_is_admin else "member",
|
||
|
|
origin="telegram",
|
||
|
|
via_agent=1,
|
||
|
|
result="denied",
|
||
|
|
summary=f"Telegram AI request by {user.get('username', owner_id)} blocked - 24h quota reached",
|
||
|
|
metadata={"limit_usd": limit},
|
||
|
|
)
|
||
|
|
await self._service.send(
|
||
|
|
chat_id, "Your daily AI quota is reached (100%). Please try again later."
|
||
|
|
)
|
||
|
|
return
|
||
|
|
session = devii.hub().get_or_create(
|
||
|
|
"user",
|
||
|
|
owner_id,
|
||
|
|
user.get("username", ""),
|
||
|
|
user.get("api_key", ""),
|
||
|
|
devii.instance_base_url(),
|
||
|
|
is_admin=owner_is_admin,
|
||
|
|
is_primary_admin=owner_is_primary_admin,
|
||
|
|
channel="telegram",
|
||
|
|
)
|
||
|
|
session.set_timezone(user.get("timezone") or "")
|
||
|
|
placeholder_id = await self._service.send(chat_id, THINKING_PLACEHOLDER)
|
||
|
|
connection = TelegramConnection(self._service, session, chat_id, placeholder_id)
|
||
|
|
session.attach(connection)
|
||
|
|
content, audit_text = self._build_content(text, images)
|
||
|
|
started = time.monotonic()
|
||
|
|
try:
|
||
|
|
session.spawn_turn(content, audit_text=audit_text)
|
||
|
|
await asyncio.wait_for(connection.done.wait(), timeout=TURN_TIMEOUT_SECONDS)
|
||
|
|
except asyncio.TimeoutError:
|
||
|
|
await connection._finalize(
|
||
|
|
html.escape("Devii took too long to respond. Please try again.")
|
||
|
|
)
|
||
|
|
except Exception: # noqa: BLE001 - never let one turn crash the bridge
|
||
|
|
logger.exception("Telegram turn failed for %s", owner_id)
|
||
|
|
await connection._finalize(html.escape("Something went wrong handling your message."))
|
||
|
|
finally:
|
||
|
|
session.detach(connection)
|
||
|
|
self._service.record_latency(time.monotonic() - started)
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def _build_content(text: str, images: list[str]) -> tuple[Any, str]:
|
||
|
|
if not images:
|
||
|
|
return text, text
|
||
|
|
prompt = text or "Look at the attached image and respond."
|
||
|
|
content: list[dict[str, Any]] = [{"type": "text", "text": prompt}]
|
||
|
|
for uri in images:
|
||
|
|
content.append({"type": "image_url", "image_url": {"url": uri}})
|
||
|
|
audit_text = f"{text} [{len(images)} image(s)]".strip()
|
||
|
|
return content, audit_text
|