diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..2b8c860 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,50 @@ +# Git +.git +.gitignore + +# Environment +.env +.env.local +.env.*.local + +# Python +__pycache__ +*.py[cod] +*$py.class +*.so +.Python +.venv +venv +*.egg-info +dist +build + +# Testing +.pytest_cache +.coverage +htmlcov +.tox +.nox + +# Type checking +.mypy_cache +.dmypy.json + +# IDE +.idea +.vscode +*.swp +*.swo + +# Documentation +*.md +!README.md + +# Docker +docker-compose*.yml +Dockerfile* + +# Other +.pre-commit-config.yaml +tests +searxng diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..e34de47 --- /dev/null +++ b/.env.example @@ -0,0 +1,33 @@ +# Discord Configuration +DISCORD_TOKEN=your_discord_bot_token_here +DISCORD_GUILD_ID= # Optional: for faster command sync during development + +# AI Provider Configuration +# Available providers: "openai", "openrouter", "anthropic" +AI_PROVIDER=openai +AI_MODEL=gpt-4o + +# Provider API Keys (set the ones you use) +OPENAI_API_KEY=sk-xxx +OPENROUTER_API_KEY=sk-or-xxx +ANTHROPIC_API_KEY=sk-ant-xxx + +# AI Settings +AI_MAX_TOKENS=1024 +AI_TEMPERATURE=0.7 + +# SearXNG Configuration +SEARXNG_BASE_URL=http://localhost:8080 +SEARXNG_TIMEOUT=10 + +# Rate Limiting +RATE_LIMIT_MESSAGES=10 # Messages per user per minute +RATE_LIMIT_SEARCHES=5 # Searches per user per minute + +# Logging +LOG_LEVEL=INFO + +# Bot Behavior +BOT_NAME=Daemon Boyfriend +BOT_PERSONALITY=helpful, witty, and slightly mischievous +MAX_CONVERSATION_HISTORY=20 diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..5b7d64d --- /dev/null +++ b/.gitignore @@ -0,0 +1,59 @@ +# Python +__pycache__/ +*.py[cod] +*$py.class +*.so +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +*.egg-info/ +.installed.cfg +*.egg + +# Virtual environments +.venv/ +venv/ +ENV/ + +# IDE +.idea/ +.vscode/ +*.swp +*.swo +*~ + +# Environment +.env +.env.local +.env.*.local + +# Testing +.pytest_cache/ +.coverage +htmlcov/ +.tox/ +.nox/ + +# Type checking +.mypy_cache/ +.dmypy.json + +# Docker +*.log + +# OS +.DS_Store +Thumbs.db + +# SearXNG config (may contain secrets) +searxng/settings.yml diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..7a07665 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,38 @@ +# Stage 1: Build dependencies +FROM python:3.11-slim as builder + +WORKDIR /app + +# Install build dependencies +RUN apt-get update && apt-get install -y --no-install-recommends \ + gcc \ + && rm -rf /var/lib/apt/lists/* + +# Install Python dependencies +COPY requirements.txt . +RUN pip wheel --no-cache-dir --no-deps --wheel-dir /app/wheels -r requirements.txt + +# Stage 2: Runtime +FROM python:3.11-slim + +WORKDIR /app + +# Create non-root user for security +RUN useradd --create-home --shell /bin/bash botuser + +# Copy wheels from builder and install +COPY --from=builder /app/wheels /wheels +RUN pip install --no-cache /wheels/* && rm -rf /wheels + +# Copy application code +COPY src/ ./src/ + +# Set Python path +ENV PYTHONPATH=/app/src +ENV PYTHONUNBUFFERED=1 + +# Switch to non-root user +USER botuser + +# Run the bot +CMD ["python", "-m", "daemon_boyfriend"] diff --git a/README.md b/README.md index 92c1a03..f7284e3 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,145 @@ -# MSC-Servo-Skull +# Daemon Boyfriend +A Discord bot for the MSC group with multi-provider AI chat and SearXNG web search. + +## Features + +- **Multi-Provider AI Chat**: Supports OpenAI, OpenRouter, and Anthropic (Claude) +- **Mention-based Chat**: Just @mention the bot to chat +- **Web Search**: Privacy-respecting search via SearXNG +- **Conversation Memory**: Remembers context per user +- **Rate Limiting**: Prevents abuse + +## Commands + +| Command | Description | +|---------|-------------| +| `@Daemon Boyfriend ` | Chat with the bot | +| `/chat ` | Chat via slash command | +| `/search ` | Search the web | +| `/image ` | Search for images | +| `/clear` | Clear your conversation history | +| `/status` | Show bot status | +| `/help` | Show all commands | + +## Quick Start + +### Prerequisites + +- Python 3.11+ +- Discord bot token ([Discord Developer Portal](https://discord.com/developers/applications)) +- AI provider API key (OpenAI, OpenRouter, or Anthropic) +- SearXNG instance (optional, for search) + +### Installation + +1. Clone the repository: +```bash +git clone https://github.com/your-username/Daemon-Boyfriend.git +cd Daemon-Boyfriend +``` + +2. Create a virtual environment: +```bash +python -m venv .venv +source .venv/bin/activate # Linux/Mac +# or: .venv\Scripts\activate # Windows +``` + +3. Install dependencies: +```bash +pip install -r requirements.txt +``` + +4. Configure the bot: +```bash +cp .env.example .env +# Edit .env with your tokens +``` + +5. Run the bot: +```bash +python -m daemon_boyfriend +``` + +## Configuration + +Edit `.env` to configure the bot: + +```env +# Required +DISCORD_TOKEN=your_bot_token + +# AI Provider (choose one) +AI_PROVIDER=openai # or: openrouter, anthropic +AI_MODEL=gpt-4o + +# API Keys (set the one you use) +OPENAI_API_KEY=sk-xxx +OPENROUTER_API_KEY=sk-or-xxx +ANTHROPIC_API_KEY=sk-ant-xxx + +# SearXNG (optional) +SEARXNG_BASE_URL=http://localhost:8080 +``` + +### AI Providers + +| Provider | Models | Notes | +|----------|--------|-------| +| OpenAI | gpt-4o, gpt-4-turbo, gpt-3.5-turbo | Official OpenAI API | +| OpenRouter | 100+ models | Access to Llama, Mistral, Claude, etc. | +| Anthropic | claude-3-5-sonnet, claude-3-opus | Direct Claude API | + +## Docker Deployment + +### With Docker Compose (includes SearXNG) + +```bash +cp .env.example .env +# Edit .env with your tokens + +docker-compose up -d +``` + +### Bot Only (external SearXNG) + +```bash +docker build -t daemon-boyfriend . +docker run -d --env-file .env daemon-boyfriend +``` + +## Development + +### Project Structure + +``` +src/daemon_boyfriend/ +├── bot.py # Main bot class +├── config.py # Configuration +├── cogs/ # Command modules +│ ├── ai_chat.py # Chat commands +│ ├── search.py # Search commands +│ └── admin.py # Admin commands +├── services/ # External integrations +│ ├── ai_service.py # AI provider factory +│ ├── providers/ # AI providers +│ ├── searxng.py # SearXNG client +│ └── conversation.py # Chat history +└── utils/ # Utilities + ├── logging.py + ├── rate_limiter.py + └── error_handler.py +``` + +### Running in Development + +Set `DISCORD_GUILD_ID` in `.env` for faster command sync during development: + +```env +DISCORD_GUILD_ID=123456789012345678 +``` + +## License + +MIT diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 0000000..677a762 --- /dev/null +++ b/docker-compose.yml @@ -0,0 +1,47 @@ +version: "3.8" + +services: + daemon-boyfriend: + build: . + container_name: daemon-boyfriend + restart: unless-stopped + env_file: + - .env + environment: + - PYTHONUNBUFFERED=1 + depends_on: + searxng: + condition: service_healthy + networks: + - bot-network + + # Optional: SearXNG instance + # Comment out if you're using an external SearXNG instance + searxng: + image: searxng/searxng:latest + container_name: searxng + restart: unless-stopped + ports: + - "8080:8080" + volumes: + - ./searxng:/etc/searxng:rw + environment: + - SEARXNG_BASE_URL=http://localhost:8080/ + cap_drop: + - ALL + cap_add: + - CHOWN + - SETGID + - SETUID + healthcheck: + test: ["CMD", "wget", "--spider", "-q", "http://localhost:8080/healthz"] + interval: 30s + timeout: 10s + retries: 3 + start_period: 10s + networks: + - bot-network + +networks: + bot-network: + driver: bridge diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..00a494b --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,51 @@ +[project] +name = "daemon-boyfriend" +version = "1.0.0" +description = "AI-powered Discord bot for the MSC group with multi-provider support and SearXNG integration" +readme = "README.md" +requires-python = ">=3.11" +dependencies = [ + "discord.py>=2.3.0", + "anthropic>=0.18.0", + "openai>=1.12.0", + "aiohttp>=3.9.0", + "pydantic>=2.6.0", + "pydantic-settings>=2.2.0", + "python-dotenv>=1.0.0", +] + +[project.optional-dependencies] +dev = [ + "pytest>=8.0.0", + "pytest-asyncio>=0.23.0", + "pytest-cov>=4.1.0", + "black>=24.0.0", + "ruff>=0.2.0", + "mypy>=1.8.0", +] + +[project.scripts] +daemon-boyfriend = "daemon_boyfriend.__main__:main" + +[build-system] +requires = ["setuptools>=61.0"] +build-backend = "setuptools.build_meta" + +[tool.setuptools.packages.find] +where = ["src"] + +[tool.black] +line-length = 100 +target-version = ["py311", "py312"] + +[tool.ruff] +line-length = 100 +select = ["E", "F", "I", "N", "W", "UP"] + +[tool.mypy] +python_version = "3.11" +strict = true + +[tool.pytest.ini_options] +asyncio_mode = "auto" +testpaths = ["tests"] diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..abbfc3a --- /dev/null +++ b/requirements.txt @@ -0,0 +1,14 @@ +# Discord +discord.py>=2.3.0 + +# AI Providers +anthropic>=0.18.0 +openai>=1.12.0 + +# HTTP Client +aiohttp>=3.9.0 + +# Configuration +pydantic>=2.6.0 +pydantic-settings>=2.2.0 +python-dotenv>=1.0.0 diff --git a/src/daemon_boyfriend/__init__.py b/src/daemon_boyfriend/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/daemon_boyfriend/__main__.py b/src/daemon_boyfriend/__main__.py new file mode 100644 index 0000000..589d912 --- /dev/null +++ b/src/daemon_boyfriend/__main__.py @@ -0,0 +1,16 @@ +"""Entry point for the Daemon Boyfriend bot.""" + +from daemon_boyfriend.bot import DaemonBoyfriend +from daemon_boyfriend.config import settings +from daemon_boyfriend.utils.logging import setup_logging + + +def main() -> None: + """Start the bot.""" + setup_logging() + bot = DaemonBoyfriend() + bot.run(settings.discord_token) + + +if __name__ == "__main__": + main() diff --git a/src/daemon_boyfriend/bot.py b/src/daemon_boyfriend/bot.py new file mode 100644 index 0000000..e68034e --- /dev/null +++ b/src/daemon_boyfriend/bot.py @@ -0,0 +1,89 @@ +"""Main Discord bot class.""" + +import logging +from pathlib import Path + +import discord +from discord.ext import commands + +from daemon_boyfriend.config import settings + +logger = logging.getLogger(__name__) + + +class DaemonBoyfriend(commands.Bot): + """The main bot class for Daemon Boyfriend.""" + + def __init__(self) -> None: + intents = discord.Intents.default() + intents.message_content = True + intents.guilds = True + intents.members = True + + super().__init__( + command_prefix=settings.command_prefix, + intents=intents, + help_command=None, # We use slash commands instead + ) + + async def setup_hook(self) -> None: + """Load cogs and sync commands on startup.""" + # Load all cogs + cogs_path = Path(__file__).parent / "cogs" + for cog_file in cogs_path.glob("*.py"): + if cog_file.name.startswith("_"): + continue + cog_name = f"daemon_boyfriend.cogs.{cog_file.stem}" + try: + await self.load_extension(cog_name) + logger.info(f"Loaded cog: {cog_name}") + except Exception as e: + logger.error(f"Failed to load cog {cog_name}: {e}") + + # Sync slash commands + if settings.discord_guild_id: + # Sync to specific guild for faster testing (instant) + guild = discord.Object(id=settings.discord_guild_id) + self.tree.copy_global_to(guild=guild) + await self.tree.sync(guild=guild) + logger.info(f"Synced commands to guild {settings.discord_guild_id}") + else: + # Global sync (can take up to 1 hour to propagate) + await self.tree.sync() + logger.info("Synced commands globally") + + async def on_ready(self) -> None: + """Called when the bot is ready.""" + if self.user is None: + return + + logger.info(f"Logged in as {self.user} (ID: {self.user.id})") + logger.info(f"Connected to {len(self.guilds)} guild(s)") + + # Set activity status + activity = discord.Activity( + type=discord.ActivityType.watching, + name="over the MSC group", + ) + await self.change_presence(activity=activity) + + async def on_command_error( + self, + ctx: commands.Context, + error: commands.CommandError, # type: ignore[type-arg] + ) -> None: + """Global error handler for prefix commands.""" + if isinstance(error, commands.CommandNotFound): + return # Ignore unknown commands + + if isinstance(error, commands.MissingPermissions): + await ctx.send("You don't have permission to use this command.") + return + + if isinstance(error, commands.CommandOnCooldown): + await ctx.send(f"Command on cooldown. Try again in {error.retry_after:.1f}s") + return + + # Log unexpected errors + logger.error(f"Command error: {error}", exc_info=error) + await ctx.send("An unexpected error occurred.") diff --git a/src/daemon_boyfriend/cogs/__init__.py b/src/daemon_boyfriend/cogs/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/daemon_boyfriend/cogs/admin.py b/src/daemon_boyfriend/cogs/admin.py new file mode 100644 index 0000000..ac1c108 --- /dev/null +++ b/src/daemon_boyfriend/cogs/admin.py @@ -0,0 +1,137 @@ +"""Admin cog - administrative commands for the bot.""" + +import logging +import platform +import sys +from datetime import datetime + +import discord +from discord import app_commands +from discord.ext import commands + +from daemon_boyfriend.config import settings +from daemon_boyfriend.services import SearXNGService + +logger = logging.getLogger(__name__) + + +class AdminCog(commands.Cog): + """Administrative commands for bot management.""" + + def __init__(self, bot: commands.Bot) -> None: + self.bot = bot + self.start_time = datetime.now() + + @app_commands.command(name="ping", description="Check bot latency") + async def ping(self, interaction: discord.Interaction) -> None: + """Check the bot's latency.""" + latency = round(self.bot.latency * 1000) + await interaction.response.send_message(f"Pong! Latency: {latency}ms") + + @app_commands.command(name="status", description="Show bot status and info") + async def status(self, interaction: discord.Interaction) -> None: + """Show bot status and information.""" + # Calculate uptime + uptime = datetime.now() - self.start_time + hours, remainder = divmod(int(uptime.total_seconds()), 3600) + minutes, seconds = divmod(remainder, 60) + uptime_str = f"{hours}h {minutes}m {seconds}s" + + # Check SearXNG status + searxng = SearXNGService() + searxng_status = "Online" if await searxng.health_check() else "Offline" + + embed = discord.Embed( + title=f"{settings.bot_name} Status", + color=discord.Color.green(), + ) + + embed.add_field(name="Uptime", value=uptime_str, inline=True) + embed.add_field(name="Latency", value=f"{round(self.bot.latency * 1000)}ms", inline=True) + embed.add_field(name="Guilds", value=str(len(self.bot.guilds)), inline=True) + + embed.add_field(name="AI Provider", value=settings.ai_provider, inline=True) + embed.add_field(name="AI Model", value=settings.ai_model, inline=True) + embed.add_field(name="SearXNG", value=searxng_status, inline=True) + + embed.add_field( + name="Python", value=f"{sys.version_info.major}.{sys.version_info.minor}", inline=True + ) + embed.add_field(name="Discord.py", value=discord.__version__, inline=True) + embed.add_field(name="Platform", value=platform.system(), inline=True) + + await interaction.response.send_message(embed=embed) + + @app_commands.command(name="provider", description="Show or change AI provider") + @app_commands.describe(provider="The AI provider to switch to (admin only)") + @app_commands.choices( + provider=[ + app_commands.Choice(name="OpenAI", value="openai"), + app_commands.Choice(name="OpenRouter", value="openrouter"), + app_commands.Choice(name="Anthropic (Claude)", value="anthropic"), + ] + ) + async def provider( + self, + interaction: discord.Interaction, + provider: str | None = None, + ) -> None: + """Show current AI provider or change it (info only, actual change requires restart).""" + if provider is None: + # Just show current provider + await interaction.response.send_message( + f"Current AI provider: **{settings.ai_provider}**\nModel: **{settings.ai_model}**", + ephemeral=True, + ) + else: + # Inform that changing requires restart + await interaction.response.send_message( + f"To change the AI provider to **{provider}**, update the `AI_PROVIDER` " + f"and corresponding API key in your `.env` file, then restart the bot.", + ephemeral=True, + ) + + @app_commands.command(name="help", description="Show available commands") + async def help_command(self, interaction: discord.Interaction) -> None: + """Show help information.""" + embed = discord.Embed( + title=f"{settings.bot_name} - Help", + description="Here are the available commands:", + color=discord.Color.blue(), + ) + + embed.add_field( + name="Chat", + value=( + f"**@{settings.bot_name} ** - Chat with me by mentioning me\n" + "**/chat ** - Chat using slash command\n" + "**/clear** - Clear your conversation history" + ), + inline=False, + ) + + embed.add_field( + name="Search", + value=("**/search ** - Search the web\n**/image ** - Search for images"), + inline=False, + ) + + embed.add_field( + name="Info", + value=( + "**/ping** - Check bot latency\n" + "**/status** - Show bot status\n" + "**/provider** - Show current AI provider\n" + "**/help** - Show this help message" + ), + inline=False, + ) + + embed.set_footer(text=f"Made with love for the MSC group") + + await interaction.response.send_message(embed=embed) + + +async def setup(bot: commands.Bot) -> None: + """Load the Admin cog.""" + await bot.add_cog(AdminCog(bot)) diff --git a/src/daemon_boyfriend/cogs/ai_chat.py b/src/daemon_boyfriend/cogs/ai_chat.py new file mode 100644 index 0000000..161b4e3 --- /dev/null +++ b/src/daemon_boyfriend/cogs/ai_chat.py @@ -0,0 +1,190 @@ +"""AI Chat cog - handles chat commands and mention responses.""" + +import logging +import re + +import discord +from discord import app_commands +from discord.ext import commands + +from daemon_boyfriend.config import settings +from daemon_boyfriend.services import AIService, ConversationManager, Message + +logger = logging.getLogger(__name__) + +# Discord message character limit +MAX_MESSAGE_LENGTH = 2000 + + +def split_message(content: str, max_length: int = MAX_MESSAGE_LENGTH) -> list[str]: + """Split a long message into chunks that fit Discord's limit. + + Tries to split on paragraph breaks, then sentence breaks, then word breaks. + """ + if len(content) <= max_length: + return [content] + + chunks: list[str] = [] + remaining = content + + while remaining: + if len(remaining) <= max_length: + chunks.append(remaining) + break + + # Find a good split point + split_point = max_length + + # Try to split on paragraph break + para_break = remaining.rfind("\n\n", 0, max_length) + if para_break > max_length // 2: + split_point = para_break + 2 + else: + # Try to split on line break + line_break = remaining.rfind("\n", 0, max_length) + if line_break > max_length // 2: + split_point = line_break + 1 + else: + # Try to split on sentence + sentence_end = max( + remaining.rfind(". ", 0, max_length), + remaining.rfind("! ", 0, max_length), + remaining.rfind("? ", 0, max_length), + ) + if sentence_end > max_length // 2: + split_point = sentence_end + 2 + else: + # Fall back to word break + word_break = remaining.rfind(" ", 0, max_length) + if word_break > 0: + split_point = word_break + 1 + + chunks.append(remaining[:split_point].rstrip()) + remaining = remaining[split_point:].lstrip() + + return chunks + + +class AIChatCog(commands.Cog): + """AI conversation commands and mention handling.""" + + def __init__(self, bot: commands.Bot) -> None: + self.bot = bot + self.ai_service = AIService() + self.conversations = ConversationManager() + + @app_commands.command(name="chat", description="Chat with Daemon Boyfriend") + @app_commands.describe(message="Your message to the bot") + async def chat(self, interaction: discord.Interaction, message: str) -> None: + """Slash command to chat with the AI.""" + await interaction.response.defer(thinking=True) + + try: + response_text = await self._generate_response(interaction.user.id, message) + + # Split long responses + chunks = split_message(response_text) + await interaction.followup.send(chunks[0]) + + # Send additional chunks as follow-up messages + for chunk in chunks[1:]: + await interaction.followup.send(chunk) + + except Exception as e: + logger.error(f"Chat error: {e}", exc_info=True) + await interaction.followup.send("Sorry, I encountered an error. Please try again.") + + @app_commands.command(name="clear", description="Clear your conversation history") + async def clear_history(self, interaction: discord.Interaction) -> None: + """Clear the user's conversation history.""" + self.conversations.clear_history(interaction.user.id) + await interaction.response.send_message( + "Your conversation history has been cleared!", ephemeral=True + ) + + @commands.Cog.listener() + async def on_message(self, message: discord.Message) -> None: + """Respond when the bot is mentioned.""" + # Ignore messages from bots + if message.author.bot: + return + + # Check if bot is mentioned + if self.bot.user is None or self.bot.user not in message.mentions: + return + + # Extract message content without the mention + content = self._extract_message_content(message) + + if not content: + # Just a mention with no message + await message.reply(f"Hey {message.author.display_name}! How can I help you?") + return + + # Show typing indicator while generating response + async with message.channel.typing(): + try: + response_text = await self._generate_response(message.author.id, content) + + # Split and send response + chunks = split_message(response_text) + await message.reply(chunks[0]) + + for chunk in chunks[1:]: + await message.channel.send(chunk) + + except Exception as e: + logger.error(f"Mention response error: {e}", exc_info=True) + await message.reply("Sorry, I encountered an error. Please try again.") + + def _extract_message_content(self, message: discord.Message) -> str: + """Extract the actual message content, removing bot mentions.""" + content = message.content + + # Remove all mentions of the bot + if self.bot.user: + # Remove <@BOT_ID> and <@!BOT_ID> patterns + content = re.sub( + rf"<@!?{self.bot.user.id}>", + "", + content, + ) + + return content.strip() + + async def _generate_response(self, user_id: int, user_message: str) -> str: + """Generate an AI response for a user message. + + Args: + user_id: Discord user ID + user_message: The user's message + + Returns: + The AI's response text + """ + # Get conversation history + history = self.conversations.get_history(user_id) + + # Add current message to history for the API call + messages = history + [Message(role="user", content=user_message)] + + # Generate response + response = await self.ai_service.chat( + messages=messages, + system_prompt=self.ai_service.get_system_prompt(), + ) + + # Save the exchange to history + self.conversations.add_exchange(user_id, user_message, response.content) + + logger.debug( + f"Generated response for user {user_id}: " + f"{len(response.content)} chars, {response.usage}" + ) + + return response.content + + +async def setup(bot: commands.Bot) -> None: + """Load the AI Chat cog.""" + await bot.add_cog(AIChatCog(bot)) diff --git a/src/daemon_boyfriend/cogs/search.py b/src/daemon_boyfriend/cogs/search.py new file mode 100644 index 0000000..1b892bd --- /dev/null +++ b/src/daemon_boyfriend/cogs/search.py @@ -0,0 +1,144 @@ +"""Search cog - web search using SearXNG.""" + +import logging + +import discord +from discord import app_commands +from discord.ext import commands + +from daemon_boyfriend.services import SearchResponse, SearXNGService + +logger = logging.getLogger(__name__) + + +def create_search_embed(response: SearchResponse) -> discord.Embed: + """Create a Discord embed for search results. + + Args: + response: The search response + + Returns: + A formatted Discord embed + """ + embed = discord.Embed( + title=f"Search: {response.query}", + color=discord.Color.blue(), + ) + + if not response.results: + embed.description = "No results found." + return embed + + # Add results as fields + for i, result in enumerate(response.results, 1): + # Truncate content if too long + content = result.content + if len(content) > 200: + content = content[:197] + "..." + + embed.add_field( + name=f"{i}. {result.title[:100]}", + value=f"{content}\n[Link]({result.url})", + inline=False, + ) + + # Add footer with result count + embed.set_footer(text=f"Found {response.number_of_results} results") + + return embed + + +class SearchCog(commands.Cog): + """Web search commands using SearXNG.""" + + def __init__(self, bot: commands.Bot) -> None: + self.bot = bot + self.search_service = SearXNGService() + + @app_commands.command(name="search", description="Search the web") + @app_commands.describe( + query="What to search for", + category="Search category", + ) + @app_commands.choices( + category=[ + app_commands.Choice(name="General", value="general"), + app_commands.Choice(name="Images", value="images"), + app_commands.Choice(name="News", value="news"), + app_commands.Choice(name="Science", value="science"), + app_commands.Choice(name="IT", value="it"), + app_commands.Choice(name="Videos", value="videos"), + ] + ) + async def search( + self, + interaction: discord.Interaction, + query: str, + category: str = "general", + ) -> None: + """Search the web using SearXNG.""" + await interaction.response.defer(thinking=True) + + try: + response = await self.search_service.search( + query=query, + categories=[category], + num_results=5, + ) + + embed = create_search_embed(response) + await interaction.followup.send(embed=embed) + + except Exception as e: + logger.error(f"Search error: {e}", exc_info=True) + await interaction.followup.send( + f"Search failed. Make sure SearXNG is running and accessible." + ) + + @app_commands.command(name="image", description="Search for images") + @app_commands.describe(query="What to search for") + async def image_search( + self, + interaction: discord.Interaction, + query: str, + ) -> None: + """Search for images using SearXNG.""" + await interaction.response.defer(thinking=True) + + try: + response = await self.search_service.search( + query=query, + categories=["images"], + num_results=5, + ) + + if not response.results: + await interaction.followup.send(f"No images found for: {query}") + return + + # Create embed with first image + embed = discord.Embed( + title=f"Image search: {query}", + color=discord.Color.purple(), + ) + + # Add image results + for i, result in enumerate(response.results[:5], 1): + embed.add_field( + name=f"{i}. {result.title[:50]}", + value=f"[View]({result.url})", + inline=True, + ) + + await interaction.followup.send(embed=embed) + + except Exception as e: + logger.error(f"Image search error: {e}", exc_info=True) + await interaction.followup.send( + "Image search failed. Make sure SearXNG is running and accessible." + ) + + +async def setup(bot: commands.Bot) -> None: + """Load the Search cog.""" + await bot.add_cog(SearchCog(bot)) diff --git a/src/daemon_boyfriend/config.py b/src/daemon_boyfriend/config.py new file mode 100644 index 0000000..da40ae8 --- /dev/null +++ b/src/daemon_boyfriend/config.py @@ -0,0 +1,78 @@ +"""Configuration management using pydantic-settings.""" + +from typing import Literal + +from pydantic import Field +from pydantic_settings import BaseSettings, SettingsConfigDict + + +class Settings(BaseSettings): + """Application settings loaded from environment variables.""" + + model_config = SettingsConfigDict( + env_file=".env", + env_file_encoding="utf-8", + case_sensitive=False, + ) + + # Discord Configuration + discord_token: str = Field(..., description="Discord bot token") + discord_guild_id: int | None = Field( + None, description="Test guild ID for faster command sync during development" + ) + command_prefix: str = Field("!", description="Legacy command prefix") + + # AI Provider Configuration + ai_provider: Literal["openai", "openrouter", "anthropic"] = Field( + "openai", description="Which AI provider to use" + ) + ai_model: str = Field("gpt-4o", description="AI model to use") + ai_max_tokens: int = Field(1024, ge=100, le=4096, description="Max tokens for AI response") + ai_temperature: float = Field(0.7, ge=0.0, le=2.0, description="AI temperature") + + # Provider API Keys + openai_api_key: str | None = Field(None, description="OpenAI API key") + openrouter_api_key: str | None = Field(None, description="OpenRouter API key") + anthropic_api_key: str | None = Field(None, description="Anthropic API key") + + # SearXNG Configuration + searxng_base_url: str = Field( + "http://localhost:8080", description="SearXNG instance URL" + ) + searxng_timeout: int = Field(10, description="Search timeout in seconds") + + # Rate Limiting + rate_limit_messages: int = Field(10, description="Messages per user per minute") + rate_limit_searches: int = Field(5, description="Searches per user per minute") + + # Logging + log_level: str = Field("INFO", description="Logging level") + + # Bot Behavior + bot_name: str = Field("Daemon Boyfriend", description="Bot display name") + bot_personality: str = Field( + "helpful, witty, and slightly mischievous", + description="Bot personality description for system prompt", + ) + max_conversation_history: int = Field( + 20, description="Max messages to keep in conversation memory per user" + ) + + def get_api_key(self) -> str: + """Get the API key for the configured provider.""" + key_map = { + "openai": self.openai_api_key, + "openrouter": self.openrouter_api_key, + "anthropic": self.anthropic_api_key, + } + key = key_map.get(self.ai_provider) + if not key: + raise ValueError( + f"No API key configured for provider '{self.ai_provider}'. " + f"Set {self.ai_provider.upper()}_API_KEY in your environment." + ) + return key + + +# Global settings instance +settings = Settings() diff --git a/src/daemon_boyfriend/models/__init__.py b/src/daemon_boyfriend/models/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/daemon_boyfriend/services/__init__.py b/src/daemon_boyfriend/services/__init__.py new file mode 100644 index 0000000..57231dd --- /dev/null +++ b/src/daemon_boyfriend/services/__init__.py @@ -0,0 +1,16 @@ +"""Services for external integrations.""" + +from .ai_service import AIService +from .conversation import ConversationManager +from .providers import AIResponse, Message +from .searxng import SearchResponse, SearchResult, SearXNGService + +__all__ = [ + "AIService", + "AIResponse", + "Message", + "ConversationManager", + "SearXNGService", + "SearchResponse", + "SearchResult", +] diff --git a/src/daemon_boyfriend/services/ai_service.py b/src/daemon_boyfriend/services/ai_service.py new file mode 100644 index 0000000..98b2270 --- /dev/null +++ b/src/daemon_boyfriend/services/ai_service.py @@ -0,0 +1,101 @@ +"""AI Service - Factory and facade for AI providers.""" + +import logging +from typing import Literal + +from daemon_boyfriend.config import Settings, settings + +from .providers import ( + AIProvider, + AIResponse, + AnthropicProvider, + Message, + OpenAIProvider, + OpenRouterProvider, +) + +logger = logging.getLogger(__name__) + +ProviderType = Literal["openai", "openrouter", "anthropic"] + + +class AIService: + """Factory and facade for AI providers. + + This class manages the creation and switching of AI providers, + and provides a unified interface for generating responses. + """ + + def __init__(self, config: Settings | None = None) -> None: + self._config = config or settings + self._provider: AIProvider | None = None + self._init_provider() + + def _init_provider(self) -> None: + """Initialize the AI provider based on configuration.""" + self._provider = self._create_provider( + self._config.ai_provider, + self._config.get_api_key(), + self._config.ai_model, + ) + + def _create_provider(self, provider_type: ProviderType, api_key: str, model: str) -> AIProvider: + """Create a provider instance.""" + providers: dict[ProviderType, type[AIProvider]] = { + "openai": OpenAIProvider, + "openrouter": OpenRouterProvider, + "anthropic": AnthropicProvider, + } + + provider_class = providers.get(provider_type) + if not provider_class: + raise ValueError(f"Unknown provider: {provider_type}") + + logger.info(f"Initializing {provider_type} provider with model {model}") + return provider_class(api_key=api_key, model=model) + + @property + def provider(self) -> AIProvider: + """Get the current provider.""" + if self._provider is None: + raise RuntimeError("AI provider not initialized") + return self._provider + + @property + def provider_name(self) -> str: + """Get the name of the current provider.""" + return self.provider.provider_name + + @property + def model(self) -> str: + """Get the current model name.""" + return self._config.ai_model + + async def chat( + self, + messages: list[Message], + system_prompt: str | None = None, + ) -> AIResponse: + """Generate a chat response. + + Args: + messages: List of conversation messages + system_prompt: Optional system prompt + + Returns: + AIResponse with the generated content + """ + return await self.provider.generate( + messages=messages, + system_prompt=system_prompt, + max_tokens=self._config.ai_max_tokens, + temperature=self._config.ai_temperature, + ) + + def get_system_prompt(self) -> str: + """Get the default system prompt for the bot.""" + return ( + f"You are {self._config.bot_name}, a {self._config.bot_personality} " + f"Discord bot for the MSC group. Keep your responses concise and engaging. " + f"You can use Discord markdown formatting in your responses." + ) diff --git a/src/daemon_boyfriend/services/conversation.py b/src/daemon_boyfriend/services/conversation.py new file mode 100644 index 0000000..e162fb4 --- /dev/null +++ b/src/daemon_boyfriend/services/conversation.py @@ -0,0 +1,81 @@ +"""Conversation history management.""" + +import logging +from collections import defaultdict + +from daemon_boyfriend.config import settings + +from .providers import Message + +logger = logging.getLogger(__name__) + + +class ConversationManager: + """Manages conversation history per user. + + Stores conversation messages in memory with a configurable + maximum history length per user. + """ + + def __init__(self, max_history: int | None = None) -> None: + self.max_history = max_history or settings.max_conversation_history + self._conversations: dict[int, list[Message]] = defaultdict(list) + + def get_history(self, user_id: int) -> list[Message]: + """Get the conversation history for a user. + + Args: + user_id: Discord user ID + + Returns: + List of messages in the conversation + """ + return list(self._conversations[user_id]) + + def add_message(self, user_id: int, message: Message) -> None: + """Add a message to a user's conversation history. + + Args: + user_id: Discord user ID + message: The message to add + """ + history = self._conversations[user_id] + history.append(message) + + # Trim history if it exceeds max length + if len(history) > self.max_history: + # Keep only the most recent messages + self._conversations[user_id] = history[-self.max_history :] + logger.debug(f"Trimmed conversation history for user {user_id}") + + def add_exchange(self, user_id: int, user_message: str, assistant_message: str) -> None: + """Add a user/assistant exchange to the conversation. + + Args: + user_id: Discord user ID + user_message: The user's message + assistant_message: The assistant's response + """ + self.add_message(user_id, Message(role="user", content=user_message)) + self.add_message(user_id, Message(role="assistant", content=assistant_message)) + + def clear_history(self, user_id: int) -> None: + """Clear the conversation history for a user. + + Args: + user_id: Discord user ID + """ + if user_id in self._conversations: + del self._conversations[user_id] + logger.debug(f"Cleared conversation history for user {user_id}") + + def get_history_length(self, user_id: int) -> int: + """Get the number of messages in a user's history. + + Args: + user_id: Discord user ID + + Returns: + Number of messages in history + """ + return len(self._conversations[user_id]) diff --git a/src/daemon_boyfriend/services/providers/__init__.py b/src/daemon_boyfriend/services/providers/__init__.py new file mode 100644 index 0000000..71e99b5 --- /dev/null +++ b/src/daemon_boyfriend/services/providers/__init__.py @@ -0,0 +1,15 @@ +"""AI Provider implementations.""" + +from .anthropic import AnthropicProvider +from .base import AIProvider, AIResponse, Message +from .openai import OpenAIProvider +from .openrouter import OpenRouterProvider + +__all__ = [ + "AIProvider", + "AIResponse", + "Message", + "OpenAIProvider", + "OpenRouterProvider", + "AnthropicProvider", +] diff --git a/src/daemon_boyfriend/services/providers/anthropic.py b/src/daemon_boyfriend/services/providers/anthropic.py new file mode 100644 index 0000000..c7ec7a4 --- /dev/null +++ b/src/daemon_boyfriend/services/providers/anthropic.py @@ -0,0 +1,59 @@ +"""Anthropic (Claude) provider implementation.""" + +import logging + +import anthropic + +from .base import AIProvider, AIResponse, Message + +logger = logging.getLogger(__name__) + + +class AnthropicProvider(AIProvider): + """Anthropic Claude API provider.""" + + def __init__(self, api_key: str, model: str = "claude-sonnet-4-20250514") -> None: + self.client = anthropic.AsyncAnthropic(api_key=api_key) + self.model = model + + @property + def provider_name(self) -> str: + return "anthropic" + + async def generate( + self, + messages: list[Message], + system_prompt: str | None = None, + max_tokens: int = 1024, + temperature: float = 0.7, + ) -> AIResponse: + """Generate a response using Claude.""" + # Build messages list (Anthropic format) + api_messages = [{"role": m.role, "content": m.content} for m in messages] + + logger.debug(f"Sending {len(api_messages)} messages to Anthropic") + + response = await self.client.messages.create( + model=self.model, + max_tokens=max_tokens, + temperature=temperature, + system=system_prompt or "", + messages=api_messages, # type: ignore[arg-type] + ) + + # Extract text from response + content = "" + for block in response.content: + if block.type == "text": + content += block.text + + usage = { + "input_tokens": response.usage.input_tokens, + "output_tokens": response.usage.output_tokens, + } + + return AIResponse( + content=content, + model=response.model, + usage=usage, + ) diff --git a/src/daemon_boyfriend/services/providers/base.py b/src/daemon_boyfriend/services/providers/base.py new file mode 100644 index 0000000..274e099 --- /dev/null +++ b/src/daemon_boyfriend/services/providers/base.py @@ -0,0 +1,52 @@ +"""Abstract base class for AI providers.""" + +from abc import ABC, abstractmethod +from dataclasses import dataclass + + +@dataclass +class Message: + """A chat message.""" + + role: str # "user", "assistant", "system" + content: str + + +@dataclass +class AIResponse: + """Response from an AI provider.""" + + content: str + model: str + usage: dict[str, int] # Token usage info + + +class AIProvider(ABC): + """Abstract base class for AI providers.""" + + @abstractmethod + async def generate( + self, + messages: list[Message], + system_prompt: str | None = None, + max_tokens: int = 1024, + temperature: float = 0.7, + ) -> AIResponse: + """Generate a response from the AI model. + + Args: + messages: List of conversation messages + system_prompt: Optional system prompt + max_tokens: Maximum tokens in response + temperature: Sampling temperature + + Returns: + AIResponse with the generated content + """ + pass + + @property + @abstractmethod + def provider_name(self) -> str: + """Return the name of this provider.""" + pass diff --git a/src/daemon_boyfriend/services/providers/openai.py b/src/daemon_boyfriend/services/providers/openai.py new file mode 100644 index 0000000..78713db --- /dev/null +++ b/src/daemon_boyfriend/services/providers/openai.py @@ -0,0 +1,61 @@ +"""OpenAI provider implementation.""" + +import logging + +from openai import AsyncOpenAI + +from .base import AIProvider, AIResponse, Message + +logger = logging.getLogger(__name__) + + +class OpenAIProvider(AIProvider): + """OpenAI API provider.""" + + def __init__(self, api_key: str, model: str = "gpt-4o") -> None: + self.client = AsyncOpenAI(api_key=api_key) + self.model = model + + @property + def provider_name(self) -> str: + return "openai" + + async def generate( + self, + messages: list[Message], + system_prompt: str | None = None, + max_tokens: int = 1024, + temperature: float = 0.7, + ) -> AIResponse: + """Generate a response using OpenAI.""" + # Build messages list + api_messages: list[dict[str, str]] = [] + + if system_prompt: + api_messages.append({"role": "system", "content": system_prompt}) + + api_messages.extend([{"role": m.role, "content": m.content} for m in messages]) + + logger.debug(f"Sending {len(api_messages)} messages to OpenAI") + + response = await self.client.chat.completions.create( + model=self.model, + max_tokens=max_tokens, + temperature=temperature, + messages=api_messages, # type: ignore[arg-type] + ) + + content = response.choices[0].message.content or "" + usage = {} + if response.usage: + usage = { + "prompt_tokens": response.usage.prompt_tokens, + "completion_tokens": response.usage.completion_tokens, + "total_tokens": response.usage.total_tokens, + } + + return AIResponse( + content=content, + model=response.model, + usage=usage, + ) diff --git a/src/daemon_boyfriend/services/providers/openrouter.py b/src/daemon_boyfriend/services/providers/openrouter.py new file mode 100644 index 0000000..82bf433 --- /dev/null +++ b/src/daemon_boyfriend/services/providers/openrouter.py @@ -0,0 +1,77 @@ +"""OpenRouter provider implementation. + +OpenRouter uses an OpenAI-compatible API, so we extend the OpenAI provider. +""" + +import logging + +from openai import AsyncOpenAI + +from .base import AIProvider, AIResponse, Message + +logger = logging.getLogger(__name__) + +# OpenRouter API base URL +OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1" + + +class OpenRouterProvider(AIProvider): + """OpenRouter API provider. + + OpenRouter provides access to 100+ models through an OpenAI-compatible API. + """ + + def __init__(self, api_key: str, model: str = "openai/gpt-4o") -> None: + self.client = AsyncOpenAI( + api_key=api_key, + base_url=OPENROUTER_BASE_URL, + ) + self.model = model + + @property + def provider_name(self) -> str: + return "openrouter" + + async def generate( + self, + messages: list[Message], + system_prompt: str | None = None, + max_tokens: int = 1024, + temperature: float = 0.7, + ) -> AIResponse: + """Generate a response using OpenRouter.""" + # Build messages list + api_messages: list[dict[str, str]] = [] + + if system_prompt: + api_messages.append({"role": "system", "content": system_prompt}) + + api_messages.extend([{"role": m.role, "content": m.content} for m in messages]) + + logger.debug(f"Sending {len(api_messages)} messages to OpenRouter ({self.model})") + + response = await self.client.chat.completions.create( + model=self.model, + max_tokens=max_tokens, + temperature=temperature, + messages=api_messages, # type: ignore[arg-type] + extra_headers={ + "HTTP-Referer": "https://github.com/daemon-boyfriend", + "X-Title": "Daemon Boyfriend", + }, + ) + + content = response.choices[0].message.content or "" + usage = {} + if response.usage: + usage = { + "prompt_tokens": response.usage.prompt_tokens, + "completion_tokens": response.usage.completion_tokens, + "total_tokens": response.usage.total_tokens, + } + + return AIResponse( + content=content, + model=response.model, + usage=usage, + ) diff --git a/src/daemon_boyfriend/services/searxng.py b/src/daemon_boyfriend/services/searxng.py new file mode 100644 index 0000000..d342444 --- /dev/null +++ b/src/daemon_boyfriend/services/searxng.py @@ -0,0 +1,119 @@ +"""SearXNG search service.""" + +import logging +from dataclasses import dataclass +from typing import Literal + +import aiohttp + +from daemon_boyfriend.config import settings + +logger = logging.getLogger(__name__) + + +@dataclass +class SearchResult: + """A single search result.""" + + title: str + url: str + content: str # snippet/description + engine: str + + +@dataclass +class SearchResponse: + """Response from a SearXNG search.""" + + query: str + results: list[SearchResult] + suggestions: list[str] + number_of_results: int + + +class SearXNGService: + """SearXNG search service client.""" + + def __init__( + self, + base_url: str | None = None, + timeout: int | None = None, + ) -> None: + self.base_url = (base_url or settings.searxng_base_url).rstrip("/") + self.timeout = aiohttp.ClientTimeout(total=timeout or settings.searxng_timeout) + + async def search( + self, + query: str, + categories: list[str] | None = None, + language: str = "en", + safesearch: Literal[0, 1, 2] = 1, # 0=off, 1=moderate, 2=strict + num_results: int = 5, + ) -> SearchResponse: + """Search using SearXNG. + + Args: + query: Search query + categories: Search categories (general, images, videos, news, it, science) + language: Language code (e.g., "en", "nl") + safesearch: Safe search level (0=off, 1=moderate, 2=strict) + num_results: Maximum number of results to return + + Returns: + SearchResponse with results + + Raises: + aiohttp.ClientError: On network errors + """ + params: dict[str, str | int] = { + "q": query, + "format": "json", + "language": language, + "safesearch": safesearch, + } + + if categories: + params["categories"] = ",".join(categories) + + logger.debug(f"Searching SearXNG: {query}") + + async with aiohttp.ClientSession(timeout=self.timeout) as session: + async with session.get(f"{self.base_url}/search", params=params) as response: + response.raise_for_status() + data = await response.json() + + # Parse results + results = [ + SearchResult( + title=r.get("title", "No title"), + url=r.get("url", ""), + content=r.get("content", "No description"), + engine=r.get("engine", "unknown"), + ) + for r in data.get("results", [])[:num_results] + ] + + return SearchResponse( + query=query, + results=results, + suggestions=data.get("suggestions", []), + number_of_results=data.get("number_of_results", len(results)), + ) + + async def health_check(self) -> bool: + """Check if the SearXNG instance is reachable. + + Returns: + True if healthy, False otherwise + """ + try: + async with aiohttp.ClientSession(timeout=self.timeout) as session: + # Try the search endpoint with a simple query + async with session.get( + f"{self.base_url}/search", + params={"q": "test", "format": "json"}, + ) as response: + return response.status == 200 + except Exception as e: + logger.error(f"SearXNG health check failed: {e}") + return False diff --git a/src/daemon_boyfriend/utils/__init__.py b/src/daemon_boyfriend/utils/__init__.py new file mode 100644 index 0000000..ce8fbd3 --- /dev/null +++ b/src/daemon_boyfriend/utils/__init__.py @@ -0,0 +1,13 @@ +"""Utility modules.""" + +from .logging import get_logger, setup_logging +from .rate_limiter import RateLimiter, chat_limiter, rate_limited, search_limiter + +__all__ = [ + "setup_logging", + "get_logger", + "RateLimiter", + "rate_limited", + "chat_limiter", + "search_limiter", +] diff --git a/src/daemon_boyfriend/utils/error_handler.py b/src/daemon_boyfriend/utils/error_handler.py new file mode 100644 index 0000000..ad97cb3 --- /dev/null +++ b/src/daemon_boyfriend/utils/error_handler.py @@ -0,0 +1,66 @@ +"""Global error handling for the bot.""" + +import logging +import traceback + +import discord +from discord import app_commands +from discord.ext import commands + +logger = logging.getLogger(__name__) + + +class ErrorHandler(commands.Cog): + """Global error handling cog.""" + + def __init__(self, bot: commands.Bot) -> None: + self.bot = bot + # Set up the app command error handler + self.bot.tree.on_error = self.on_app_command_error + + async def on_app_command_error( + self, + interaction: discord.Interaction, + error: app_commands.AppCommandError, + ) -> None: + """Handle slash command errors. + + Args: + interaction: The interaction that caused the error + error: The error that was raised + """ + # Determine the error message + if isinstance(error, app_commands.CommandOnCooldown): + message = f"This command is on cooldown. Try again in {error.retry_after:.1f} seconds." + elif isinstance(error, app_commands.MissingPermissions): + message = "You don't have permission to use this command." + elif isinstance(error, app_commands.BotMissingPermissions): + message = "I don't have the required permissions to do that." + elif isinstance(error, app_commands.CommandNotFound): + message = "That command doesn't exist." + elif isinstance(error, app_commands.TransformerError): + message = f"Invalid input: {error}" + elif isinstance(error, app_commands.CheckFailure): + message = "You can't use this command." + else: + # Log unexpected errors + logger.error( + f"Unhandled app command error: {error}", + exc_info=error, + ) + message = "An unexpected error occurred. Please try again later." + + # Send the error message + try: + if interaction.response.is_done(): + await interaction.followup.send(message, ephemeral=True) + else: + await interaction.response.send_message(message, ephemeral=True) + except discord.HTTPException: + # If we can't send a message, just log it + logger.error(f"Failed to send error message: {message}") + + +async def setup(bot: commands.Bot) -> None: + """Load the error handler cog.""" + await bot.add_cog(ErrorHandler(bot)) diff --git a/src/daemon_boyfriend/utils/logging.py b/src/daemon_boyfriend/utils/logging.py new file mode 100644 index 0000000..0c222df --- /dev/null +++ b/src/daemon_boyfriend/utils/logging.py @@ -0,0 +1,44 @@ +"""Logging configuration for the bot.""" + +import logging +import sys + +from daemon_boyfriend.config import settings + + +def setup_logging() -> None: + """Configure application-wide logging.""" + # Create formatter + formatter = logging.Formatter( + "%(asctime)s - %(name)s - %(levelname)s - %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", + ) + + # Configure root logger + root_logger = logging.getLogger() + root_logger.setLevel(settings.log_level) + + # Remove existing handlers + root_logger.handlers.clear() + + # Console handler + console_handler = logging.StreamHandler(sys.stdout) + console_handler.setFormatter(formatter) + root_logger.addHandler(console_handler) + + # Reduce noise from discord.py + logging.getLogger("discord").setLevel(logging.WARNING) + logging.getLogger("discord.http").setLevel(logging.WARNING) + logging.getLogger("discord.gateway").setLevel(logging.WARNING) + + # Reduce noise from aiohttp + logging.getLogger("aiohttp").setLevel(logging.WARNING) + + # Reduce noise from httpx (used by openai/anthropic) + logging.getLogger("httpx").setLevel(logging.WARNING) + logging.getLogger("httpcore").setLevel(logging.WARNING) + + +def get_logger(name: str) -> logging.Logger: + """Get a logger with the given name.""" + return logging.getLogger(name) diff --git a/src/daemon_boyfriend/utils/rate_limiter.py b/src/daemon_boyfriend/utils/rate_limiter.py new file mode 100644 index 0000000..158845f --- /dev/null +++ b/src/daemon_boyfriend/utils/rate_limiter.py @@ -0,0 +1,138 @@ +"""Rate limiting utilities.""" + +import asyncio +import logging +from collections import defaultdict +from datetime import datetime, timedelta +from functools import wraps +from typing import Any, Callable, Coroutine, TypeVar + +import discord +from discord import app_commands + +from daemon_boyfriend.config import settings + +logger = logging.getLogger(__name__) + +T = TypeVar("T") + + +class RateLimiter: + """Per-user rate limiting. + + Tracks the number of calls per user within a time window + and rejects calls that exceed the limit. + """ + + def __init__(self, max_calls: int, period_seconds: int) -> None: + """Initialize the rate limiter. + + Args: + max_calls: Maximum calls allowed per period + period_seconds: Length of the period in seconds + """ + self.max_calls = max_calls + self.period = timedelta(seconds=period_seconds) + self.calls: dict[int, list[datetime]] = defaultdict(list) + self._lock = asyncio.Lock() + + async def is_allowed(self, user_id: int) -> bool: + """Check if a user is within rate limit and record the call. + + Args: + user_id: Discord user ID + + Returns: + True if allowed, False if rate limited + """ + async with self._lock: + now = datetime.now() + cutoff = now - self.period + + # Clean old entries + self.calls[user_id] = [t for t in self.calls[user_id] if t > cutoff] + + if len(self.calls[user_id]) >= self.max_calls: + return False + + self.calls[user_id].append(now) + return True + + def remaining(self, user_id: int) -> int: + """Get remaining calls for a user in current period. + + Args: + user_id: Discord user ID + + Returns: + Number of remaining calls + """ + now = datetime.now() + cutoff = now - self.period + recent = [t for t in self.calls[user_id] if t > cutoff] + return max(0, self.max_calls - len(recent)) + + def reset_time(self, user_id: int) -> float: + """Get seconds until rate limit resets for a user. + + Args: + user_id: Discord user ID + + Returns: + Seconds until reset, or 0 if not rate limited + """ + if not self.calls[user_id]: + return 0 + + oldest = min(self.calls[user_id]) + reset_at = oldest + self.period + remaining = (reset_at - datetime.now()).total_seconds() + return max(0, remaining) + + +def rate_limited( + limiter: RateLimiter, +) -> Callable[ + [Callable[..., Coroutine[Any, Any, T]]], + Callable[..., Coroutine[Any, Any, T | None]], +]: + """Decorator for rate-limited slash commands. + + Args: + limiter: The RateLimiter instance to use + + Returns: + Decorated function that checks rate limit before execution + """ + + def decorator( + func: Callable[..., Coroutine[Any, Any, T]], + ) -> Callable[..., Coroutine[Any, Any, T | None]]: + @wraps(func) + async def wrapper( + self: Any, interaction: discord.Interaction, *args: Any, **kwargs: Any + ) -> T | None: + if not await limiter.is_allowed(interaction.user.id): + reset_time = limiter.reset_time(interaction.user.id) + await interaction.response.send_message( + f"You're sending messages too quickly. Please wait {reset_time:.0f} seconds.", + ephemeral=True, + ) + return None + return await func(self, interaction, *args, **kwargs) + + return wrapper + + return decorator + + +# Pre-configured rate limiters +chat_limiter = RateLimiter( + max_calls=settings.rate_limit_messages, + period_seconds=60, +) + +search_limiter = RateLimiter( + max_calls=settings.rate_limit_searches, + period_seconds=60, +)