bot.py 6.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170
  1. # bot.py
  2. import os
  3. import asyncio
  4. import discord
  5. from discord import app_commands
  6. import requests
  7. from dotenv import load_dotenv
  8. # --- Configuration ---
  9. load_dotenv()
  10. DISCORD_BOT_TOKEN = os.getenv("DISCORD_BOT_TOKEN")
  11. PERPLEXICA_API_URL = "http://localhost:3000/api/search"
  12. # Hardcoded model configurations
  13. CHAT_MODEL = { "provider": "ollama", "name": "ministral-3:8b" }
  14. EMBEDDING_MODEL = { "provider": "ollama", "name": "ryanshillington/Qwen3-Embedding-0.6B:latest" }
  15. # In-memory store for user-selected optimization mode
  16. user_modes = {}
  17. DEFAULT_OPTIMIZATION_MODE = "balanced"
  18. DEFAULT_FOCUS_MODE = "webSearch"
  19. HELP_TEXT = """
  20. Hello! I'm Perplexica Bot.
  21. To use me, please type `/search <your query>` to search for information.
  22. You can also change your search effort mode using `/mode`.
  23. I don't respond to direct messages or mentions outside of these commands.
  24. """
  25. # --- Discord Bot Setup ---
  26. intents = discord.Intents.default()
  27. intents.message_content = True
  28. client = discord.Client(intents=intents)
  29. tree = app_commands.CommandTree(client)
  30. @client.event
  31. async def on_ready():
  32. """Event handler for when the bot is ready."""
  33. print(f'Logged in as {client.user}')
  34. await tree.sync()
  35. print("Slash commands synced.")
  36. @client.event
  37. async def on_message(message: discord.Message):
  38. """Event handler for when a message is sent."""
  39. # Ignore messages from the bot itself
  40. if message.author == client.user:
  41. return
  42. # Respond to DMs or mentions
  43. if isinstance(message.channel, discord.DMChannel) or client.user.mentioned_in(message):
  44. await message.channel.send(HELP_TEXT)
  45. return
  46. await client.process_commands(message) # Important! Allows other commands to be processed.
  47. # --- Slash Commands ---
  48. @tree.command(name="search", description="Search with Perplexica")
  49. async def search(interaction: discord.Interaction, *, query: str):
  50. """Performs a search using the Perplexica API."""
  51. await interaction.response.send_message(f"<@{interaction.user.id}> 🧠 Thinking about your query: \"{query}\"...")
  52. # Get user's selected mode, default to balanced
  53. optimization_mode = user_modes.get(interaction.user.id, DEFAULT_OPTIMIZATION_MODE)
  54. # Construct the payload
  55. payload = {
  56. "query": query,
  57. "focusMode": DEFAULT_FOCUS_MODE,
  58. "optimizationMode": optimization_mode,
  59. "chatModel": CHAT_MODEL,
  60. "embeddingModel": EMBEDDING_MODEL,
  61. "stream": False
  62. }
  63. try:
  64. # Send request to Perplexica API
  65. response = requests.post(PERPLEXICA_API_URL, json=payload, timeout=300)
  66. response.raise_for_status() # Raise an exception for bad status codes
  67. data = response.json()
  68. message = data.get("message", "No answer found.")
  69. sources = data.get("sources", [])
  70. # Format the response
  71. embed = discord.Embed(
  72. title=f"🔎 {query}",
  73. description=message,
  74. color=discord.Color.blue()
  75. )
  76. if sources:
  77. source_links = []
  78. for i, source in enumerate(sources, 1):
  79. title = source.get("title")
  80. url = source.get("url")
  81. if title and url:
  82. source_links.append(f"[{i}] {title} ({url})")
  83. # Join sources and add to a field, handling character limits
  84. source_text = "\n".join(source_links)
  85. if len(source_text) > 1024:
  86. source_text = source_text[:1021] + "..."
  87. embed.add_field(name="🔗 Sources", value=source_text, inline=False)
  88. # Handle overall message length
  89. if len(embed) > 2000:
  90. original_description = embed.description
  91. # Calculate available space for description in the first message
  92. non_desc_len = len(embed) - len(original_description)
  93. # Max length for description, keeping total embed length under 2000.
  94. # With a small buffer for "..."
  95. available_len = 2000 - non_desc_len - 10
  96. # Ensure available_len is not negative
  97. if available_len < 0:
  98. available_len = 0
  99. first_part = original_description[:available_len]
  100. rest_of_message = original_description[available_len:]
  101. embed.description = first_part
  102. if rest_of_message:
  103. embed.description += "..."
  104. await interaction.edit_original_response(content=f"<@{interaction.user.id}>", embed=embed)
  105. if rest_of_message:
  106. # Split the rest of the message into chunks of 2000 characters
  107. chunks = [rest_of_message[i:i + 2000] for i in range(0, len(rest_of_message), 2000)]
  108. for i, chunk in enumerate(chunks):
  109. await asyncio.sleep(1) # Add a small delay to prevent rate limiting issues
  110. continuation_embed = discord.Embed(
  111. title=f"🔎 {query} (continued {i + 1}/{len(chunks)})",
  112. description=chunk,
  113. color=discord.Color.blue()
  114. )
  115. await interaction.followup.send(content=f"<@{interaction.user.id}>", embed=continuation_embed)
  116. else:
  117. await interaction.edit_original_response(content=f"<@{interaction.user.id}>", embed=embed)
  118. except requests.exceptions.RequestException as e:
  119. await interaction.edit_original_response(content=f"An error occurred while contacting the Perplexica API: {e}")
  120. except Exception as e:
  121. await interaction.edit_original_response(content=f"An unexpected error occurred: {e}")
  122. @tree.command(name="mode", description="Set the search optimization mode")
  123. @app_commands.choices(mode=[
  124. app_commands.Choice(name="Balanced", value="balanced"),
  125. app_commands.Choice(name="Fast", value="speed"),
  126. ])
  127. async def mode(interaction: discord.Interaction, mode: app_commands.Choice[str]):
  128. """Sets the optimization mode for the user."""
  129. user_modes[interaction.user.id] = mode.value
  130. await interaction.response.send_message(f"✅ Your search mode has been set to **{mode.name}**.", ephemeral=True)
  131. # --- Run the Bot ---
  132. if __name__ == "__main__":
  133. if DISCORD_BOT_TOKEN == "YOUR_DISCORD_BOT_TOKEN_HERE" or not DISCORD_BOT_TOKEN:
  134. print("ERROR: Please set your DISCORD_BOT_TOKEN in the .env file.")
  135. else:
  136. client.run(DISCORD_BOT_TOKEN)