tech-envision commited on
Commit
57fa15b
·
1 Parent(s): 4e9c2e5

Refactor Discord bot and add module entry point

Browse files
Files changed (3) hide show
  1. README.md +1 -1
  2. bot/__main__.py +4 -0
  3. bot/discord_bot.py +94 -68
README.md CHANGED
@@ -69,7 +69,7 @@ DISCORD_TOKEN="your-token"
69
  Then start the bot:
70
 
71
  ```bash
72
- python -m bot.discord_bot
73
  ```
74
 
75
  Any attachments sent to the bot are uploaded to the VM and the bot replies with
 
69
  Then start the bot:
70
 
71
  ```bash
72
+ python -m bot
73
  ```
74
 
75
  Any attachments sent to the bot are uploaded to the VM and the bot replies with
bot/__main__.py ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ from .discord_bot import main
2
+
3
+ if __name__ == "__main__":
4
+ main()
bot/discord_bot.py CHANGED
@@ -1,92 +1,118 @@
 
 
 
 
1
  import logging
2
  import os
3
  import shutil
4
  import tempfile
5
  from pathlib import Path
 
6
 
7
  import discord
8
  from discord.ext import commands
9
  from dotenv import load_dotenv
10
 
11
- from src.team import TeamChatSession
12
  from src.db import reset_history
13
  from src.log import get_logger
 
14
 
15
- _LOG = get_logger(__name__, level=logging.INFO)
16
-
17
-
18
- def _create_bot() -> commands.Bot:
19
- intents = discord.Intents.all()
20
- return commands.Bot(command_prefix="!", intents=intents)
21
-
22
-
23
- bot = _create_bot()
24
-
25
-
26
- @bot.event
27
- async def on_ready() -> None:
28
- _LOG.info("Logged in as %s", bot.user)
29
-
30
-
31
- @bot.command(name="reset")
32
- async def reset(ctx: commands.Context) -> None:
33
- deleted = reset_history(str(ctx.author.id), str(ctx.channel.id))
34
- await ctx.reply(f"Chat history cleared ({deleted} messages deleted).")
35
-
36
-
37
- async def _handle_attachments(chat: TeamChatSession, message: discord.Message) -> list[tuple[str, str]]:
38
- if not message.attachments:
39
- return []
40
-
41
- uploaded: list[tuple[str, str]] = []
42
- tmpdir = Path(tempfile.mkdtemp(prefix="discord_upload_"))
43
- try:
44
- for attachment in message.attachments:
45
- dest = tmpdir / attachment.filename
46
- await attachment.save(dest)
47
- vm_path = chat.upload_document(str(dest))
48
- uploaded.append((attachment.filename, vm_path))
49
- finally:
50
- shutil.rmtree(tmpdir, ignore_errors=True)
51
-
52
- return uploaded
53
-
54
-
55
- @bot.event
56
- async def on_message(message: discord.Message) -> None:
57
- if message.author.bot:
58
- return
59
-
60
- await bot.process_commands(message)
61
- if message.content.startswith("!"):
62
- return
63
-
64
- async with TeamChatSession(user=str(message.author.id), session=str(message.channel.id)) as chat:
65
- docs = await _handle_attachments(chat, message)
66
- if docs:
67
- info = "\n".join(f"{name} -> {path}" for name, path in docs)
68
- await message.reply(f"Uploaded:\n{info}", mention_author=False)
69
 
70
- if message.content.strip():
71
- try:
72
- async for part in chat.chat_stream(message.content):
73
- await message.reply(part, mention_author=False)
74
- except Exception as exc: # pragma: no cover - runtime errors
75
- _LOG.error("Failed to process message: %s", exc)
76
- await message.reply(f"Error: {exc}", mention_author=False)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
77
 
78
 
79
  def main() -> None:
 
 
80
  load_dotenv()
81
  token = os.getenv("DISCORD_TOKEN")
82
  if not token:
83
  raise RuntimeError("DISCORD_TOKEN environment variable not set")
84
 
85
- bot.run(token)
86
 
87
 
88
- if __name__ == "__main__":
89
- try:
90
- main()
91
- except KeyboardInterrupt: # pragma: no cover - manual exit
92
- pass
 
1
+ """Discord bot implementation."""
2
+
3
+ from __future__ import annotations
4
+
5
  import logging
6
  import os
7
  import shutil
8
  import tempfile
9
  from pathlib import Path
10
+ from typing import Iterable
11
 
12
  import discord
13
  from discord.ext import commands
14
  from dotenv import load_dotenv
15
 
 
16
  from src.db import reset_history
17
  from src.log import get_logger
18
+ from src.team import TeamChatSession
19
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
20
 
21
+ class DiscordTeamBot(commands.Bot):
22
+ """Discord bot for interacting with :class:`TeamChatSession`."""
23
+
24
+ def __init__(self) -> None:
25
+ intents = discord.Intents.all()
26
+ super().__init__(command_prefix="!", intents=intents)
27
+ self._log = get_logger(__name__, level=logging.INFO)
28
+ self._register_commands()
29
+
30
+ # ------------------------------------------------------------------
31
+ # Lifecycle events
32
+ # ------------------------------------------------------------------
33
+ async def on_ready(self) -> None: # noqa: D401 - callback signature
34
+ """Log a message once the bot has connected."""
35
+
36
+ self._log.info("Logged in as %s", self.user)
37
+
38
+ async def on_message(self, message: discord.Message) -> None: # noqa: D401 - callback signature
39
+ """Process incoming messages and stream chat replies."""
40
+
41
+ if message.author.bot:
42
+ return
43
+
44
+ await self.process_commands(message)
45
+ if message.content.startswith("!"):
46
+ return
47
+
48
+ async with TeamChatSession(
49
+ user=str(message.author.id), session=str(message.channel.id)
50
+ ) as chat:
51
+ docs = await self._handle_attachments(chat, message.attachments)
52
+ if docs:
53
+ info = "\n".join(f"{name} -> {path}" for name, path in docs)
54
+ await message.reply(f"Uploaded:\n{info}", mention_author=False)
55
+
56
+ if message.content.strip():
57
+ try:
58
+ async for part in chat.chat_stream(message.content):
59
+ await message.reply(part, mention_author=False)
60
+ except Exception as exc: # pragma: no cover - runtime errors
61
+ self._log.error("Failed to process message: %s", exc)
62
+ await message.reply(f"Error: {exc}", mention_author=False)
63
+
64
+ # ------------------------------------------------------------------
65
+ # Commands
66
+ # ------------------------------------------------------------------
67
+ def _register_commands(self) -> None:
68
+ @self.command(name="reset")
69
+ async def reset(ctx: commands.Context) -> None:
70
+ deleted = reset_history(str(ctx.author.id), str(ctx.channel.id))
71
+ await ctx.reply(
72
+ f"Chat history cleared ({deleted} messages deleted).",
73
+ )
74
+
75
+ # ------------------------------------------------------------------
76
+ # Helpers
77
+ # ------------------------------------------------------------------
78
+ async def _handle_attachments(
79
+ self, chat: TeamChatSession, attachments: Iterable[discord.Attachment]
80
+ ) -> list[tuple[str, str]]:
81
+ """Download any attachments and return their VM paths."""
82
+
83
+ if not attachments:
84
+ return []
85
+
86
+ uploaded: list[tuple[str, str]] = []
87
+ tmpdir = Path(tempfile.mkdtemp(prefix="discord_upload_"))
88
+ try:
89
+ for attachment in attachments:
90
+ dest = tmpdir / attachment.filename
91
+ await attachment.save(dest)
92
+ vm_path = chat.upload_document(str(dest))
93
+ uploaded.append((attachment.filename, vm_path))
94
+ finally:
95
+ shutil.rmtree(tmpdir, ignore_errors=True)
96
+
97
+ return uploaded
98
+
99
+
100
+ def run_bot(token: str) -> None:
101
+ """Create and run the Discord bot."""
102
+
103
+ DiscordTeamBot().run(token)
104
 
105
 
106
  def main() -> None:
107
+ """Load environment and start the bot."""
108
+
109
  load_dotenv()
110
  token = os.getenv("DISCORD_TOKEN")
111
  if not token:
112
  raise RuntimeError("DISCORD_TOKEN environment variable not set")
113
 
114
+ run_bot(token)
115
 
116
 
117
+ if __name__ == "__main__": # pragma: no cover - manual execution
118
+ main()