mirror of
https://github.com/MODSetter/SurfSense.git
synced 2025-09-01 18:19:08 +00:00
fix: async thread on discord but
This commit is contained in:
parent
1d67a87b82
commit
a0f9efdb67
2 changed files with 46 additions and 9 deletions
|
@ -10,11 +10,12 @@ import logging
|
|||
import discord
|
||||
from discord.ext import commands
|
||||
import datetime
|
||||
import asyncio
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class DiscordConnector:
|
||||
class DiscordConnector(commands.Bot):
|
||||
"""Class for retrieving guild, channel, and message history from Discord."""
|
||||
|
||||
def __init__(self, token: str = None):
|
||||
|
@ -31,23 +32,28 @@ class DiscordConnector:
|
|||
intents.members = True # Required to fetch member information
|
||||
super().__init__(command_prefix="!", intents=intents) # command_prefix is required but not strictly used here
|
||||
self.token = token
|
||||
self._ready = False
|
||||
self._bot_task = None # Holds the async bot task
|
||||
|
||||
# Event to confirm bot is ready
|
||||
@self.event
|
||||
async def on_ready():
|
||||
logger.info(f"Logged in as {self.user} (ID: {self.user.id})")
|
||||
self._ready = True
|
||||
|
||||
async def start_bot(self):
|
||||
"""Starts the bot to connect to Discord."""
|
||||
logger.info("Starting Discord bot...")
|
||||
|
||||
if not self.token:
|
||||
raise ValueError("Discord bot token not set. Call set_token(token) first.")
|
||||
await self.start(self.token)
|
||||
logger.info("Discord bot started successfully.")
|
||||
|
||||
async def close_bot(self):
|
||||
"""Closes the bot's connection to Discord."""
|
||||
logger.info("Closing Discord bot connection...")
|
||||
|
||||
await self.close()
|
||||
logger.info("Discord bot connection closed.")
|
||||
|
||||
|
||||
def set_token(self, token: str) -> None:
|
||||
|
@ -57,14 +63,21 @@ class DiscordConnector:
|
|||
Args:
|
||||
token (str): The Discord bot token.
|
||||
"""
|
||||
logger.info("Setting Discord bot token.")
|
||||
self.token = token
|
||||
logger.info("Token set successfully. You can now start the bot with start_bot().")
|
||||
|
||||
async def _wait_until_ready(self):
|
||||
"""Helper to wait until the bot is connected and ready."""
|
||||
if not self._ready:
|
||||
logger.info("Bot not yet ready, waiting...")
|
||||
await self.wait_until_ready()
|
||||
logger.info("Bot is now ready.")
|
||||
logger.info("Waiting for the bot to be ready...")
|
||||
|
||||
# Give the event loop a chance to switch to the bot's startup task.
|
||||
# This allows self.start() to begin initializing the client.
|
||||
# Terrible solution, but necessary to avoid blocking the event loop.
|
||||
await asyncio.sleep(1) # Yield control to the event loop
|
||||
|
||||
await self.wait_until_ready()
|
||||
logger.info("Bot is ready.")
|
||||
|
||||
async def get_guilds(self) -> list[dict]:
|
||||
"""
|
||||
|
@ -78,6 +91,7 @@ class DiscordConnector:
|
|||
ValueError: If the token is not set.
|
||||
"""
|
||||
await self._wait_until_ready()
|
||||
logger.info("Fetching guilds...")
|
||||
|
||||
guilds_data = []
|
||||
for guild in self.guilds:
|
||||
|
@ -89,6 +103,8 @@ class DiscordConnector:
|
|||
"member_count": member_count,
|
||||
}
|
||||
)
|
||||
|
||||
logger.info(f"Fetched {len(guilds_data)} guilds.")
|
||||
return guilds_data
|
||||
|
||||
async def get_text_channels(self, guild_id: str) -> list[dict]:
|
||||
|
@ -106,6 +122,7 @@ class DiscordConnector:
|
|||
discord.NotFound: If the guild is not found.
|
||||
"""
|
||||
await self._wait_until_ready()
|
||||
logger.info(f"Fetching text channels for guild ID: {guild_id}")
|
||||
|
||||
guild = self.get_guild(int(guild_id))
|
||||
if not guild:
|
||||
|
@ -118,6 +135,8 @@ class DiscordConnector:
|
|||
channels_data.append(
|
||||
{"id": str(channel.id), "name": channel.name, "type": "text"}
|
||||
)
|
||||
|
||||
logger.info(f"Fetched {len(channels_data)} text channels from guild {guild_id}.")
|
||||
return channels_data
|
||||
|
||||
async def get_channel_history(
|
||||
|
@ -146,6 +165,7 @@ class DiscordConnector:
|
|||
discord.Forbidden: If the bot does not have permissions to read history in the channel.
|
||||
"""
|
||||
await self._wait_until_ready()
|
||||
logger.info(f"Fetching message history for channel ID: {channel_id}")
|
||||
|
||||
channel = self.get_channel(int(channel_id))
|
||||
if not channel:
|
||||
|
@ -190,7 +210,8 @@ class DiscordConnector:
|
|||
except discord.HTTPException as e:
|
||||
logger.error(f"Failed to fetch messages from channel {channel_id}: {e}")
|
||||
return []
|
||||
|
||||
|
||||
logger.info(f"Fetched {len(messages_data)} messages from channel {channel_id}.")
|
||||
return messages_data
|
||||
|
||||
async def get_user_info(self, guild_id: str, user_id: str) -> dict | None:
|
||||
|
@ -211,6 +232,7 @@ class DiscordConnector:
|
|||
permissions to view members.
|
||||
"""
|
||||
await self._wait_until_ready()
|
||||
logger.info(f"Fetching user info for user ID: {user_id} in guild ID: {guild_id}")
|
||||
|
||||
guild = self.get_guild(int(guild_id))
|
||||
if not guild:
|
||||
|
@ -221,12 +243,15 @@ class DiscordConnector:
|
|||
member = await guild.fetch_member(int(user_id))
|
||||
if member:
|
||||
roles = [role.name for role in member.roles if role.name != "@everyone"]
|
||||
logger.info(f"User {user_id} found in guild {guild_id}.")
|
||||
|
||||
return {
|
||||
"id": str(member.id),
|
||||
"name": member.name,
|
||||
"joined_at": member.joined_at.isoformat() if member.joined_at else None,
|
||||
"roles": roles,
|
||||
}
|
||||
logger.warning(f"User {user_id} not found in guild {guild_id}.")
|
||||
return None
|
||||
except discord.NotFound:
|
||||
logger.warning(f"User {user_id} not found in guild {guild_id}.")
|
||||
|
|
|
@ -15,6 +15,7 @@ from app.connectors.discord_connector import DiscordConnector
|
|||
from discord import DiscordException
|
||||
from slack_sdk.errors import SlackApiError
|
||||
import logging
|
||||
import asyncio
|
||||
|
||||
from app.utils.document_converters import generate_content_hash
|
||||
|
||||
|
@ -952,6 +953,9 @@ async def index_discord_messages(
|
|||
if not discord_token:
|
||||
return 0, "Discord token not found in connector config"
|
||||
|
||||
logger.info(f"Starting Discord indexing for connector {connector_id}")
|
||||
logger.info(f"discord token is: {discord_token[:4]}... (truncated for security)")
|
||||
|
||||
# Initialize Discord client
|
||||
discord_client = DiscordConnector(token=discord_token)
|
||||
|
||||
|
@ -975,13 +979,21 @@ async def index_discord_messages(
|
|||
skipped_guilds = []
|
||||
|
||||
try:
|
||||
await discord_client.start_bot()
|
||||
logger.info("Starting Discord bot to fetch guilds")
|
||||
discord_client._bot_task = asyncio.create_task(discord_client.start_bot())
|
||||
await discord_client._wait_until_ready()
|
||||
|
||||
logger.info("Fetching Discord guilds")
|
||||
guilds = await discord_client.get_guilds()
|
||||
logger.info(f"Found {len(guilds)} guilds")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get Discord guilds: {str(e)}", exc_info=True)
|
||||
|
||||
await discord_client.close_bot()
|
||||
return 0, f"Failed to get Discord guilds: {str(e)}"
|
||||
if not guilds:
|
||||
logger.info("No Discord guilds found to index")
|
||||
|
||||
await discord_client.close_bot()
|
||||
return 0, "No Discord guilds found"
|
||||
|
||||
|
|
Loading…
Add table
Reference in a new issue