bot.py 5.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149
  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": "granite4:micro-h" }
  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. # --- Discord Bot Setup ---
  20. intents = discord.Intents.default()
  21. intents.message_content = True
  22. client = discord.Client(intents=intents)
  23. tree = app_commands.CommandTree(client)
  24. @client.event
  25. async def on_ready():
  26. """Event handler for when the bot is ready."""
  27. print(f'Logged in as {client.user}')
  28. await tree.sync()
  29. print("Slash commands synced.")
  30. # --- Slash Commands ---
  31. @tree.command(name="search", description="Search with Perplexica")
  32. async def search(interaction: discord.Interaction, *, query: str):
  33. """Performs a search using the Perplexica API."""
  34. await interaction.response.send_message(f"<@{interaction.user.id}> 🧠 Thinking about your query: \"{query}\"...")
  35. # Get user's selected mode, default to balanced
  36. optimization_mode = user_modes.get(interaction.user.id, DEFAULT_OPTIMIZATION_MODE)
  37. # Construct the payload
  38. payload = {
  39. "query": query,
  40. "focusMode": DEFAULT_FOCUS_MODE,
  41. "optimizationMode": optimization_mode,
  42. "chatModel": CHAT_MODEL,
  43. "embeddingModel": EMBEDDING_MODEL,
  44. "stream": False
  45. }
  46. try:
  47. # Send request to Perplexica API
  48. response = requests.post(PERPLEXICA_API_URL, json=payload, timeout=300)
  49. response.raise_for_status() # Raise an exception for bad status codes
  50. data = response.json()
  51. message = data.get("message", "No answer found.")
  52. sources = data.get("sources", [])
  53. # Format the response
  54. embed = discord.Embed(
  55. title=f"🔎 {query}",
  56. description=message,
  57. color=discord.Color.blue()
  58. )
  59. if sources:
  60. source_links = []
  61. for i, source in enumerate(sources, 1):
  62. title = source.get("title")
  63. url = source.get("url")
  64. if title and url:
  65. source_links.append(f"[{i}] {title} ({url})")
  66. # Join sources and add to a field, handling character limits
  67. source_text = "\n".join(source_links)
  68. if len(source_text) > 1024:
  69. source_text = source_text[:1021] + "..."
  70. embed.add_field(name="🔗 Sources", value=source_text, inline=False)
  71. # Handle overall message length
  72. if len(embed) > 2000:
  73. original_description = embed.description
  74. # Calculate available space for description in the first message
  75. non_desc_len = len(embed) - len(original_description)
  76. # Max length for description, keeping total embed length under 2000.
  77. # With a small buffer for "..."
  78. available_len = 2000 - non_desc_len - 10
  79. # Ensure available_len is not negative
  80. if available_len < 0:
  81. available_len = 0
  82. first_part = original_description[:available_len]
  83. rest_of_message = original_description[available_len:]
  84. embed.description = first_part
  85. if rest_of_message:
  86. embed.description += "..."
  87. await interaction.edit_original_response(content=f"<@{interaction.user.id}>", embed=embed)
  88. if rest_of_message:
  89. # Split the rest of the message into chunks of 2000 characters
  90. chunks = [rest_of_message[i:i + 2000] for i in range(0, len(rest_of_message), 2000)]
  91. for i, chunk in enumerate(chunks):
  92. await asyncio.sleep(1) # Add a small delay to prevent rate limiting issues
  93. continuation_embed = discord.Embed(
  94. title=f"🔎 {query} (continued {i + 1}/{len(chunks)})",
  95. description=chunk,
  96. color=discord.Color.blue()
  97. )
  98. await interaction.followup.send(content=f"<@{interaction.user.id}>", embed=continuation_embed)
  99. else:
  100. await interaction.edit_original_response(content=f"<@{interaction.user.id}>", embed=embed)
  101. except requests.exceptions.RequestException as e:
  102. await interaction.edit_original_response(content=f"An error occurred while contacting the Perplexica API: {e}")
  103. except Exception as e:
  104. await interaction.edit_original_response(content=f"An unexpected error occurred: {e}")
  105. @tree.command(name="mode", description="Set the search optimization mode")
  106. @app_commands.choices(mode=[
  107. app_commands.Choice(name="Balanced", value="balanced"),
  108. app_commands.Choice(name="Fast", value="speed"),
  109. ])
  110. async def mode(interaction: discord.Interaction, mode: app_commands.Choice[str]):
  111. """Sets the optimization mode for the user."""
  112. user_modes[interaction.user.id] = mode.value
  113. await interaction.response.send_message(f"✅ Your search mode has been set to **{mode.name}**.", ephemeral=True)
  114. # --- Run the Bot ---
  115. if __name__ == "__main__":
  116. if DISCORD_BOT_TOKEN == "YOUR_DISCORD_BOT_TOKEN_HERE" or not DISCORD_BOT_TOKEN:
  117. print("ERROR: Please set your DISCORD_BOT_TOKEN in the .env file.")
  118. else:
  119. client.run(DISCORD_BOT_TOKEN)