Coverage for src/asketmc_bot/discord_bot.py: 42%
260 statements
« prev ^ index » next coverage.py v7.5.1, created at 2026-04-08 17:14 +0000
1#!/usr/bin/env python3.10
2from __future__ import annotations
4import asyncio
5import hashlib
6import logging
7import time
8from dataclasses import dataclass
9from pathlib import Path
10from typing import Any, Awaitable, Callable, Dict, Optional
12import aiohttp
13import discord
14from aiohttp import ClientSession, TCPConnector
15from discord.ext import commands
17from asketmc_bot import config as cfg
18from asketmc_bot.rag_filter import purge_filter_cache
20disc_log = logging.getLogger("asketmc.discord")
21disc_log.setLevel(logging.DEBUG if getattr(cfg, "DEBUG", False) else logging.INFO)
23RagGenFn = Callable[..., Awaitable[str]]
24QueryModelFn = Callable[..., Awaitable[tuple[str, bool]]]
25CallLocalFn = Callable[[str], Awaitable[str]]
26BuildIndexFn = Callable[[], Awaitable[Any]]
27IsBlockedFn = Callable[[], bool]
28ShutdownFn = Callable[[], Awaitable[None]]
31@dataclass
32class AppDeps:
33 index: Any
34 retriever: Any
35 generate_rag_answer: RagGenFn
36 query_model: QueryModelFn
37 call_local_llm: CallLocalFn
38 build_index: BuildIndexFn
39 is_openrouter_blocked: IsBlockedFn
40 on_core_shutdown: Optional[ShutdownFn] = None
43_STATE: Optional[AppDeps] = None
44_user_last: Dict[int, float] = {}
47class _AsyncSessionHolder:
48 def __init__(self) -> None:
49 self._lock = asyncio.Lock()
50 self._session: Optional[ClientSession] = None
52 async def get(self) -> ClientSession:
53 async with self._lock:
54 if self._session is None or self._session.closed:
55 self._session = ClientSession(
56 connector=TCPConnector(limit=cfg.HTTP_CONN_LIMIT)
57 )
58 return self._session
60 async def close(self) -> None:
61 async with self._lock:
62 if self._session and not self._session.closed: 62 ↛ 61line 62 didn't jump to line 61
63 await self._session.close()
64 self._session = None
67_SESSION_HOLDER = _AsyncSessionHolder()
70def _configure_logging(*, verbose: bool) -> None:
71 """
72 Keep semantic bot logs; suppress discord.py transport/debug noise.
74 This config is local to loggers' levels and does not require changes in main.py.
75 """
76 disc_log.setLevel(logging.DEBUG if verbose else logging.INFO)
78 # discord.py internal chatter
79 logging.getLogger("discord").setLevel(logging.WARNING)
80 logging.getLogger("discord.gateway").setLevel(logging.WARNING)
81 logging.getLogger("discord.client").setLevel(logging.WARNING)
82 logging.getLogger("discord.http").setLevel(logging.WARNING)
84 # optional: aiohttp can be noisy too
85 logging.getLogger("aiohttp").setLevel(logging.WARNING)
88def _require_state() -> AppDeps:
89 if _STATE is None:
90 raise RuntimeError(
91 "Discord module is not initialized. Call run_bot(...) or start_bot_async(...)."
92 )
93 return _STATE
96def _check_cooldown(user_id: int) -> bool:
97 now = time.monotonic()
98 last = _user_last.get(user_id, 0.0)
99 if now - last < cfg.USER_COOLDOWN:
100 return False
101 _user_last[user_id] = now
102 return True
105def _sanitize(text: str) -> str:
106 return (
107 text.replace("@", "@\u200b")
108 .replace("```", " ")
109 .replace("</sys>", " ")
110 .replace("<sys>", " ")
111 )
114def _split_for_discord(text: str, limit: int = 2000) -> list[str]:
115 parts: list[str] = []
116 while text: 116 ↛ 125line 116 didn't jump to line 125, because the condition on line 116 was always true
117 if len(text) <= limit:
118 parts.append(text)
119 break
120 cut = text.rfind("\n", 0, limit)
121 if cut == -1:
122 cut = limit
123 parts.append(text[:cut])
124 text = text[cut:].lstrip("\n")
125 return parts
128def _qhash(text: str) -> str:
129 return hashlib.sha1(text.encode("utf-8")).hexdigest()[:10]
132async def _send_long(ctx: commands.Context, text: str, limit: int = 1900) -> None:
133 parts = _split_for_discord(text, limit)
134 disc_log.debug(
135 "send parts=%s total_len=%s limit=%s ch_id=%s",
136 len(parts),
137 len(text),
138 limit,
139 getattr(ctx.channel, "id", "?"),
140 )
141 for i, chunk in enumerate(parts, start=1):
142 disc_log.debug("send_part i=%s len=%s", i, len(chunk))
143 await ctx.send(chunk)
146_bot: Optional[commands.Bot] = None
147_shutdown_lock = asyncio.Lock()
148_shutdown_flag = False
151def _build_bot() -> commands.Bot:
152 # Minimal intents for prefix commands; avoids presence/member noise.
153 intents = discord.Intents.none()
154 intents.guilds = True
155 intents.messages = True
156 intents.message_content = True
157 return commands.Bot(command_prefix="!", intents=intents)
160def _check_admin_only():
161 def predicate(ctx: commands.Context) -> bool:
162 if ctx.author.id not in cfg.ADMIN_IDS:
163 raise commands.CheckFailure("not_admin")
164 return True
166 return commands.check(predicate)
169def _check_channel_allowed():
170 def predicate(ctx: commands.Context) -> bool:
171 if ctx.channel.id not in cfg.ALLOWED_CHANNELS:
172 raise commands.CheckFailure("not_allowed_channel")
173 return True
175 return commands.check(predicate)
178def _read_prompt_file(prompt_path: Path) -> str:
179 return prompt_path.read_text(encoding="utf-8")
182async def _answer_rag(
183 ctx: commands.Context,
184 question_raw: str,
185 *,
186 prompt_path: Path,
187 force_local: bool | None = None,
188) -> None:
189 deps = _require_state()
191 q = _sanitize((question_raw or "").strip().replace("\n", " "))
192 if not q or len(q) > cfg.MAX_QUESTION_LEN or not cfg.ALLOWED_CHARS.match(q):
193 await ctx.send("❌ Invalid query format.")
194 return
196 if not _check_cooldown(ctx.author.id):
197 await ctx.send("⏳ Please wait before sending another query.")
198 return
200 is_blocked = deps.is_openrouter_blocked()
201 use_remote = (force_local is None and not is_blocked) or (force_local is False)
203 disc_log.debug(
204 "rq user_id=%s ch_id=%s cmd=%s qlen=%s qhash=%s mode=%s",
205 getattr(ctx.author, "id", "?"),
206 getattr(ctx.channel, "id", "?"),
207 getattr(getattr(ctx, "command", None), "qualified_name", "?"),
208 len(q),
209 _qhash(q),
210 "remote" if use_remote else "local",
211 )
213 await ctx.send("🧠 Thinking locally…" if not use_remote else "🔍 Thinking…")
215 start = time.monotonic()
216 try:
217 async with cfg.REQUEST_SEMAPHORE:
218 sys_prompt = _read_prompt_file(Path(prompt_path))
219 answer = await deps.generate_rag_answer(q, sys_prompt, use_remote=use_remote)
220 except FileNotFoundError:
221 await ctx.send("⚠️ System prompt file is missing.")
222 return
223 except UnicodeError:
224 await ctx.send("⚠️ System prompt file is not valid UTF-8.")
225 return
226 except asyncio.TimeoutError:
227 await ctx.send("⚠️ Request timed out. Try again later.")
228 return
229 except aiohttp.ClientError as exc:
230 disc_log.warning("[_answer_rag] network error: %s", exc)
231 await ctx.send("⚠️ Network error while contacting the model. Try again later.")
232 return
233 except ValueError as exc:
234 disc_log.info("[_answer_rag] bad request: %s", exc)
235 await ctx.send("⚠️ Request rejected by validation.")
236 return
237 except discord.HTTPException as exc:
238 disc_log.warning("[_answer_rag] discord HTTP error: %s", exc)
239 await ctx.send("⚠️ Discord API error while sending the response.")
240 return
241 except Exception as exc:
242 disc_log.exception("[_answer_rag] unexpected error: %r", exc)
243 await ctx.send("❌ Internal error. Try again later.")
244 return
246 out = answer or "❌ No answer."
247 disc_log.debug(
248 "rs user_id=%s ch_id=%s olen=%s dt_ms=%s",
249 getattr(ctx.author, "id", "?"),
250 getattr(ctx.channel, "id", "?"),
251 len(out),
252 int((time.monotonic() - start) * 1000),
253 )
254 await _send_long(ctx, out)
256 if force_local is None and is_blocked:
257 await ctx.send("⚠️ OpenRouter unavailable, local model used.")
260def _register_commands(bot: commands.Bot) -> None:
261 @bot.command(name="strict", help="Answer using strict prompt and RAG.")
262 @_check_channel_allowed()
263 async def cmd_strict(ctx: commands.Context, *, q: str) -> None:
264 await _answer_rag(ctx, q, prompt_path=Path(cfg.PROMPT_STRICT), force_local=None)
266 @bot.command(name="think", help="Answer using reasoning prompt and RAG.")
267 @_check_channel_allowed()
268 async def cmd_think(ctx: commands.Context, *, q: str) -> None:
269 await _answer_rag(ctx, q, prompt_path=Path(cfg.PROMPT_REASON), force_local=None)
271 @bot.command(name="local", help="Answer only with local LLM; cloud is skipped.")
272 @_check_channel_allowed()
273 async def cmd_local(ctx: commands.Context, *, q: str) -> None:
274 await _answer_rag(ctx, q, prompt_path=Path(cfg.PROMPT_STRICT), force_local=True)
276 @bot.command(name="status", help="Show bot/index/cache status.")
277 async def cmd_status(ctx: commands.Context) -> None:
278 deps = _require_state()
279 try:
280 docs = len(deps.index.docstore.docs) # type: ignore[attr-defined]
281 except Exception:
282 docs = -1
283 await ctx.send(
284 f"🧠 Documents: {docs}\n"
285 f"💾 Cache: {'yes' if cfg.CACHE_PATH.exists() else 'no'}\n"
286 f"🌐 OpenRouter: {'blocked' if deps.is_openrouter_blocked() else 'ok'}"
287 )
289 @bot.command(name="reload_index", help="Rebuild the index (admin only).")
290 @_check_admin_only()
291 async def cmd_reload_index(ctx: commands.Context) -> None:
292 deps = _require_state()
293 try:
294 purge_filter_cache()
295 new_index = await deps.build_index()
296 deps.index = new_index
297 deps.retriever = new_index.as_retriever(similarity_top_k=cfg.TOP_K)
298 await ctx.send("✅ Index reloaded.")
299 except Exception as exc:
300 disc_log.exception("[reload_index] error: %r", exc)
301 await ctx.send("❌ Internal error while rebuilding index.")
303 @bot.command(name="stop", help="Stop the bot (admin only).")
304 @_check_admin_only()
305 async def cmd_stop(ctx: commands.Context) -> None:
306 await ctx.send("🛑 Shutting down bot...")
307 await _shutdown_once(bot)
309 @bot.event
310 async def on_ready() -> None:
311 disc_log.info("Bot started as %s (ID: %s)", bot.user, getattr(bot.user, "id", "?"))
312 try:
313 ch = bot.get_channel(next(iter(cfg.ALLOWED_CHANNELS)))
314 if ch:
315 await ch.send("✅ Bot started and ready.")
316 except Exception as exc:
317 disc_log.error("[on_ready] startup notify failed: %s", exc)
319 @bot.event
320 async def on_command_error(ctx: commands.Context, error: Exception) -> None:
321 if isinstance(error, commands.CheckFailure):
322 msg = str(error)
323 if msg == "not_admin":
324 await ctx.send("❌ Access denied.")
325 elif msg == "not_allowed_channel":
326 await ctx.send("❌ This command is not allowed in this channel.")
327 else:
328 await ctx.send("❌ You are not allowed to use this command here.")
329 return
330 if isinstance(error, commands.MissingRequiredArgument):
331 await ctx.send("❌ Invalid command usage.")
332 return
333 disc_log.exception("[on_command_error] Unhandled: %r", error)
334 await ctx.send("❌ Internal error.")
337async def _shutdown_once(bot: commands.Bot) -> None:
338 global _shutdown_flag
339 async with _shutdown_lock:
340 if _shutdown_flag:
341 return
342 _shutdown_flag = True
344 try:
345 ch = bot.get_channel(next(iter(cfg.ALLOWED_CHANNELS))) # type: ignore[arg-type]
346 if ch:
347 await ch.send("🛑 Bot stopped.")
348 except Exception as exc:
349 disc_log.error("[shutdown] notify failed: %s", exc)
351 try:
352 await _SESSION_HOLDER.close()
353 finally:
354 deps = _require_state()
355 if deps.on_core_shutdown:
356 try:
357 await deps.on_core_shutdown()
358 except Exception as exc:
359 disc_log.exception("[shutdown] on_core_shutdown error: %s", exc)
360 await bot.close()
363def run_bot(
364 *,
365 token: str,
366 index: Any,
367 retriever: Any,
368 generate_rag_answer: RagGenFn,
369 query_model: QueryModelFn,
370 call_local_llm: CallLocalFn,
371 build_index: BuildIndexFn,
372 is_openrouter_blocked: IsBlockedFn,
373 on_core_shutdown: Optional[ShutdownFn] = None,
374) -> None:
375 global _STATE, _bot, _shutdown_flag
376 _shutdown_flag = False
378 _configure_logging(verbose=bool(getattr(cfg, "DEBUG", False)))
380 _STATE = AppDeps(
381 index=index,
382 retriever=retriever,
383 generate_rag_answer=generate_rag_answer,
384 query_model=query_model,
385 call_local_llm=call_local_llm,
386 build_index=build_index,
387 is_openrouter_blocked=is_openrouter_blocked,
388 on_core_shutdown=on_core_shutdown,
389 )
390 _bot = _build_bot()
391 _register_commands(_bot)
392 _bot.run(token, log_handler=None)
395async def start_bot_async(
396 *,
397 token: str,
398 index: Any,
399 retriever: Any,
400 generate_rag_answer: RagGenFn,
401 query_model: QueryModelFn,
402 call_local_llm: CallLocalFn,
403 build_index: BuildIndexFn,
404 is_openrouter_blocked: IsBlockedFn,
405 on_core_shutdown: Optional[ShutdownFn] = None,
406) -> None:
407 global _STATE, _bot, _shutdown_flag
408 _shutdown_flag = False
410 _configure_logging(verbose=bool(getattr(cfg, "DEBUG", False)))
412 _STATE = AppDeps(
413 index=index,
414 retriever=retriever,
415 generate_rag_answer=generate_rag_answer,
416 query_model=query_model,
417 call_local_llm=call_local_llm,
418 build_index=build_index,
419 is_openrouter_blocked=is_openrouter_blocked,
420 on_core_shutdown=on_core_shutdown,
421 )
422 _bot = _build_bot()
423 _register_commands(_bot)
424 try:
425 await _bot.start(token)
426 finally:
427 if _bot is not None:
428 await _shutdown_once(_bot)