tech-envision commited on
Commit
7765906
·
1 Parent(s): 5cdf20a

Add Discord bot with file uploads

Browse files
Files changed (4) hide show
  1. README.md +17 -0
  2. bot/__init__.py +5 -0
  3. bot/discord_bot.py +95 -0
  4. src/__init__.py +1 -2
README.md CHANGED
@@ -36,3 +36,20 @@ async with ChatSession() as chat:
36
  When using the Discord bot, attach one or more text files to a message to
37
  upload them automatically. The bot responds with the location of each document
38
  inside the VM so they can be referenced in subsequent prompts.
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
36
  When using the Discord bot, attach one or more text files to a message to
37
  upload them automatically. The bot responds with the location of each document
38
  inside the VM so they can be referenced in subsequent prompts.
39
+
40
+ ## Discord Bot
41
+
42
+ Create a `.env` file with your Discord token:
43
+
44
+ ```bash
45
+ DISCORD_TOKEN="your-token"
46
+ ```
47
+
48
+ Then start the bot:
49
+
50
+ ```bash
51
+ python -m bot.discord_bot
52
+ ```
53
+
54
+ Any attachments sent to the bot are uploaded to the VM and the bot replies with
55
+ their paths so they can be used in later messages.
bot/__init__.py ADDED
@@ -0,0 +1,5 @@
 
 
 
 
 
 
1
+ """Discord bot package."""
2
+
3
+ from .discord_bot import main
4
+
5
+ __all__ = ["main"]
bot/discord_bot.py ADDED
@@ -0,0 +1,95 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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.chat import ChatSession
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.default()
20
+ intents.messages = True
21
+ intents.message_content = True
22
+ return commands.Bot(command_prefix="!", intents=intents)
23
+
24
+
25
+ bot = _create_bot()
26
+
27
+
28
+ @bot.event
29
+ async def on_ready() -> None:
30
+ _LOG.info("Logged in as %s", bot.user)
31
+
32
+
33
+ @bot.command(name="reset")
34
+ async def reset(ctx: commands.Context) -> None:
35
+ deleted = reset_history(str(ctx.author.id), str(ctx.channel.id))
36
+ await ctx.reply(f"Chat history cleared ({deleted} messages deleted).")
37
+
38
+
39
+ async def _handle_attachments(chat: ChatSession, message: discord.Message) -> list[tuple[str, str]]:
40
+ if not message.attachments:
41
+ return []
42
+
43
+ uploaded: list[tuple[str, str]] = []
44
+ tmpdir = Path(tempfile.mkdtemp(prefix="discord_upload_"))
45
+ try:
46
+ for attachment in message.attachments:
47
+ dest = tmpdir / attachment.filename
48
+ await attachment.save(dest)
49
+ vm_path = chat.upload_document(str(dest))
50
+ uploaded.append((attachment.filename, vm_path))
51
+ finally:
52
+ shutil.rmtree(tmpdir, ignore_errors=True)
53
+
54
+ return uploaded
55
+
56
+
57
+ @bot.event
58
+ async def on_message(message: discord.Message) -> None:
59
+ if message.author.bot:
60
+ return
61
+
62
+ await bot.process_commands(message)
63
+ if message.content.startswith("!"):
64
+ return
65
+
66
+ async with ChatSession(user=str(message.author.id), session=str(message.channel.id)) as chat:
67
+ docs = await _handle_attachments(chat, message)
68
+ if docs:
69
+ info = "\n".join(f"{name} -> {path}" for name, path in docs)
70
+ await message.channel.send(f"Uploaded:\n{info}")
71
+
72
+ if message.content.strip():
73
+ try:
74
+ reply = await chat.chat(message.content)
75
+ except Exception as exc: # pragma: no cover - runtime errors
76
+ _LOG.error("Failed to process message: %s", exc)
77
+ await message.channel.send(f"Error: {exc}")
78
+ else:
79
+ await message.channel.send(reply)
80
+
81
+
82
+ def main() -> None:
83
+ load_dotenv()
84
+ token = os.getenv("DISCORD_TOKEN")
85
+ if not token:
86
+ raise RuntimeError("DISCORD_TOKEN environment variable not set")
87
+
88
+ bot.run(token)
89
+
90
+
91
+ if __name__ == "__main__":
92
+ try:
93
+ main()
94
+ except KeyboardInterrupt: # pragma: no cover - manual exit
95
+ pass
src/__init__.py CHANGED
@@ -1,6 +1,5 @@
1
  from .chat import ChatSession
2
  from .tools import execute_terminal, set_vm
3
  from .vm import LinuxVM
4
- from .api import create_app
5
 
6
- __all__ = ["ChatSession", "execute_terminal", "set_vm", "LinuxVM", "create_app"]
 
1
  from .chat import ChatSession
2
  from .tools import execute_terminal, set_vm
3
  from .vm import LinuxVM
 
4
 
5
+ __all__ = ["ChatSession", "execute_terminal", "set_vm", "LinuxVM"]