From 4980c8f2a7e99b5c04ac08b5911f0353ce1652c2 Mon Sep 17 00:00:00 2001 From: bless-stanford <107686196+bless-stanford@users.noreply.github.com> Date: Mon, 27 Jan 2025 11:47:04 -0800 Subject: [PATCH 01/19] Remove example environment file --- .example.env | 2 -- 1 file changed, 2 deletions(-) delete mode 100644 .example.env diff --git a/.example.env b/.example.env deleted file mode 100644 index 535e7b45..00000000 --- a/.example.env +++ /dev/null @@ -1,2 +0,0 @@ -DISCORD_TOKEN= -MISTRAL_API_KEY= \ No newline at end of file From c64348308609d1ab2943c82d4fd0223b0d22a0e9 Mon Sep 17 00:00:00 2001 From: bless-stanford <107686196+bless-stanford@users.noreply.github.com> Date: Mon, 27 Jan 2025 13:37:30 -0800 Subject: [PATCH 02/19] Refactor agent to use Semantic Kernel and improve chat message handling --- agent.py | 36 +++++++++++------------- bot.py | 12 ++++---- kernel/kernel_builder.py | 60 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 81 insertions(+), 27 deletions(-) create mode 100644 kernel/kernel_builder.py diff --git a/agent.py b/agent.py index fb886381..3231ad88 100644 --- a/agent.py +++ b/agent.py @@ -1,29 +1,25 @@ -import os -from mistralai import Mistral import discord +from semantic_kernel.contents import ChatHistory +from kernel.kernel_builder import KernelBuilder MISTRAL_MODEL = "mistral-large-latest" -SYSTEM_PROMPT = "You are a helpful assistant." - +SYSTEM_PROMPT = "You are a helpful assistant. Your name is Dodobot" class MistralAgent: def __init__(self): - MISTRAL_API_KEY = os.getenv("MISTRAL_API_KEY") - - self.client = Mistral(api_key=MISTRAL_API_KEY) + self.kernel = KernelBuilder.create_kernel(model_id=MISTRAL_MODEL) + self.settings = KernelBuilder.get_default_settings() + self.chat_service = self.kernel.get_service() async def run(self, message: discord.Message): - # The simplest form of an agent - # Send the message's content to Mistral's API and return Mistral's response - - messages = [ - {"role": "system", "content": SYSTEM_PROMPT}, - {"role": "user", "content": message.content}, - ] - - response = await self.client.chat.complete_async( - model=MISTRAL_MODEL, - messages=messages, + # Create chat history and add messages + chat_history = ChatHistory() + chat_history.add_system_message(SYSTEM_PROMPT) + chat_history.add_user_message(message.content) + + response = await self.chat_service.get_chat_message_content( + chat_history=chat_history, + settings=self.settings ) - - return response.choices[0].message.content + + return response.content \ No newline at end of file diff --git a/bot.py b/bot.py index d146885b..756fde73 100644 --- a/bot.py +++ b/bot.py @@ -42,8 +42,6 @@ async def on_ready(): async def on_message(message: discord.Message): """ Called when a message is sent in any channel the bot can see. - - https://discordpy.readthedocs.io/en/latest/api.html#discord.on_message """ # Don't delete this line! It's necessary for the bot to process commands. await bot.process_commands(message) @@ -53,12 +51,12 @@ async def on_message(message: discord.Message): return # Process the message with the agent you wrote - # Open up the agent.py file to customize the agent logger.info(f"Processing message from {message.author}: {message.content}") - response = await agent.run(message) - - # Send the response back to the channel - await message.reply(response) + + # Show typing indicator while processing + async with message.channel.typing(): + response = await agent.run(message) + await message.reply(response) # Commands diff --git a/kernel/kernel_builder.py b/kernel/kernel_builder.py new file mode 100644 index 00000000..ea688424 --- /dev/null +++ b/kernel/kernel_builder.py @@ -0,0 +1,60 @@ +import os +from dotenv import load_dotenv +from semantic_kernel import Kernel +from semantic_kernel.connectors.ai.mistral_ai import ( + MistralAIChatCompletion, + MistralAIChatPromptExecutionSettings +) + +class KernelBuilder: + @staticmethod + def create_kernel( + model_id: str = "mistral-large-latest", # or "mistral-small", "mistral-large" + load_env: bool = True + ) -> Kernel: + """ + Creates and configures a Semantic Kernel instance with Mistral AI. + + Args: + model_id (str): The Mistral AI model to use + service_id (str): Service identifier for the chat completion service + load_env (bool): Whether to load environment variables from .env file + + Returns: + Kernel: Configured Semantic Kernel instance + + Environment Variables Required: + MISTRAL_API_KEY: Your Mistral AI API key + """ + # Load environment variables if requested + if load_env: + load_dotenv() + + # Get API key from environment variables + api_key = os.getenv("MISTRAL_API_KEY") + if not api_key: + raise ValueError("MISTRAL_API_KEY environment variable is not set") + + # Create kernel instance + kernel = Kernel() + + # Create chat completion service + chat_completion_service = MistralAIChatCompletion( + ai_model_id=model_id, + api_key=api_key, + ) + + # Add the chat service to the kernel + kernel.add_service(chat_completion_service) + + return kernel + + @staticmethod + def get_default_settings() -> MistralAIChatPromptExecutionSettings: + """ + Creates default execution settings for Mistral AI chat completion. + + Returns: + MistralAIChatPromptExecutionSettings: Default settings for chat completion + """ + return MistralAIChatPromptExecutionSettings() \ No newline at end of file From 0adf14bf5ef3a9de1f73b6b0cbc53490e3f735b5 Mon Sep 17 00:00:00 2001 From: bless-stanford <107686196+bless-stanford@users.noreply.github.com> Date: Mon, 27 Jan 2025 14:10:36 -0800 Subject: [PATCH 03/19] Refactor chat history management to persist across messages --- agent.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/agent.py b/agent.py index 3231ad88..ef32e56a 100644 --- a/agent.py +++ b/agent.py @@ -10,16 +10,20 @@ def __init__(self): self.kernel = KernelBuilder.create_kernel(model_id=MISTRAL_MODEL) self.settings = KernelBuilder.get_default_settings() self.chat_service = self.kernel.get_service() + # Initialize chat history once + self.chat_history = ChatHistory() + self.chat_history.add_system_message(SYSTEM_PROMPT) async def run(self, message: discord.Message): - # Create chat history and add messages - chat_history = ChatHistory() - chat_history.add_system_message(SYSTEM_PROMPT) - chat_history.add_user_message(message.content) + # Add new message to existing history + self.chat_history.add_user_message(message.content) response = await self.chat_service.get_chat_message_content( - chat_history=chat_history, + chat_history=self.chat_history, settings=self.settings ) + # Add assistant's response to history + self.chat_history.add_assistant_message(response.content) + return response.content \ No newline at end of file From ad4aa30315a864e61fc35f9d738b561e3f962267 Mon Sep 17 00:00:00 2001 From: bless-stanford <107686196+bless-stanford@users.noreply.github.com> Date: Mon, 27 Jan 2025 14:48:24 -0800 Subject: [PATCH 04/19] Enhance message handling by implementing response splitting and markdown formatting for Discord compatibility --- agent.py | 124 ++++++++++++++++++++++++++++++++++++++++++++++++++++--- bot.py | 64 +++++++++++++++++----------- 2 files changed, 159 insertions(+), 29 deletions(-) diff --git a/agent.py b/agent.py index ef32e56a..daac3f11 100644 --- a/agent.py +++ b/agent.py @@ -1,21 +1,31 @@ import discord +import re from semantic_kernel.contents import ChatHistory from kernel.kernel_builder import KernelBuilder MISTRAL_MODEL = "mistral-large-latest" -SYSTEM_PROMPT = "You are a helpful assistant. Your name is Dodobot" +SYSTEM_PROMPT = """You are a helpful assistant. Your name is Dodobot. Format your responses using Discord-compatible Markdown: +- Use **bold** for emphasis +- Use *italics* for secondary emphasis +- Use `code` for technical terms, commands, or short code snippets +- Use ```language + code block + ``` for multi-line code (where 'language' is python, javascript, etc) +- Use > for quotes +- Use ||text|| for spoilers +Do not use # for headers or * - for bullet points as these don't render in Discord. +Keep responses concise when possible, as Discord has a 2000-character limit per message.""" class MistralAgent: def __init__(self): self.kernel = KernelBuilder.create_kernel(model_id=MISTRAL_MODEL) self.settings = KernelBuilder.get_default_settings() self.chat_service = self.kernel.get_service() - # Initialize chat history once self.chat_history = ChatHistory() self.chat_history.add_system_message(SYSTEM_PROMPT) + self.MAX_LENGTH = 1900 # Leave room for extra characters async def run(self, message: discord.Message): - # Add new message to existing history self.chat_history.add_user_message(message.content) response = await self.chat_service.get_chat_message_content( @@ -23,7 +33,111 @@ async def run(self, message: discord.Message): settings=self.settings ) - # Add assistant's response to history self.chat_history.add_assistant_message(response.content) - return response.content \ No newline at end of file + # Always return a list of chunks, even for short messages + if len(response.content) > self.MAX_LENGTH: + return self.split_response(response.content) + return [response.content] + + def split_response(self, content: str) -> list[str]: + chunks = [] + + # Use regex to split content into code blocks and regular text + code_block_pattern = r'(```(?:\w+\n)?[\s\S]*?```)' + parts = re.split(code_block_pattern, content) + + current_chunk = "" + + for part in parts: + if part.strip() == "": + continue + + is_code_block = part.startswith('```') and part.endswith('```') + + if is_code_block: + # If current chunk plus code block would exceed limit + if len(current_chunk) + len(part) > self.MAX_LENGTH: + if current_chunk: + chunks.append(current_chunk.strip()) + # If code block itself exceeds limit, split it + if len(part) > self.MAX_LENGTH: + code_chunks = self._split_code_block(part) + chunks.extend(code_chunks) + else: + current_chunk = part + else: + current_chunk += ('\n' if current_chunk else '') + part + else: + # Split non-code text by sentences + sentences = re.split(r'([.!?]\s+)', part) + + for i in range(0, len(sentences), 2): + sentence = sentences[i] + punctuation = sentences[i + 1] if i + 1 < len(sentences) else '' + full_sentence = sentence + punctuation + + # If adding this sentence would exceed limit + if len(current_chunk) + len(full_sentence) > self.MAX_LENGTH: + if current_chunk: + chunks.append(current_chunk.strip()) + current_chunk = full_sentence + else: + current_chunk += full_sentence + + if current_chunk: + chunks.append(current_chunk.strip()) + + # Post-process chunks to ensure proper markdown closing + return self._ensure_markdown_consistency(chunks) + + def _split_code_block(self, code_block: str) -> list[str]: + """Split a large code block into smaller chunks while preserving syntax.""" + # Extract language if present + first_line_end = code_block.find('\n') + language = code_block[3:first_line_end].strip() if first_line_end > 3 else '' + + # Remove original backticks and language + code = code_block[3 + len(language):].strip('`').strip() + + chunks = [] + current_chunk = '' + + for line in code.split('\n'): + if len(current_chunk) + len(line) + 8 > self.MAX_LENGTH: # 8 accounts for backticks and newline + if current_chunk: + chunks.append(f"```{language}\n{current_chunk.strip()}```") + current_chunk = line + else: + current_chunk += ('\n' if current_chunk else '') + line + + if current_chunk: + chunks.append(f"```{language}\n{current_chunk.strip()}```") + + return chunks + + def _ensure_markdown_consistency(self, chunks: list[str]) -> list[str]: + """Ensure that markdown formatting is properly closed in each chunk.""" + processed_chunks = [] + + for i, chunk in enumerate(chunks): + # Track open formatting + bold_count = chunk.count('**') % 2 + italic_count = chunk.count('*') % 2 + spoiler_count = chunk.count('||') % 2 + + # Close any open formatting + if bold_count: + chunk += '**' + if italic_count: + chunk += '*' + if spoiler_count: + chunk += '||' + + # If this is not the last chunk and ends with a partial code block + if i < len(chunks) - 1 and chunk.count('```') % 2: + chunk += '\n```' + + processed_chunks.append(chunk) + + return processed_chunks \ No newline at end of file diff --git a/bot.py b/bot.py index 756fde73..dadbc1b7 100644 --- a/bot.py +++ b/bot.py @@ -1,7 +1,6 @@ import os import discord import logging - from discord.ext import commands from dotenv import load_dotenv from agent import MistralAgent @@ -15,56 +14,74 @@ load_dotenv() # Create the bot with all intents -# The message content and members intent must be enabled in the Discord Developer Portal for the bot to work. intents = discord.Intents.all() bot = commands.Bot(command_prefix=PREFIX, intents=intents) # Import the Mistral agent from the agent.py file agent = MistralAgent() - # Get the token from the environment variables token = os.getenv("DISCORD_TOKEN") +async def send_split_message(message: discord.Message, response: str | list[str]): + """ + Sends a message that might be longer than Discord's character limit. + Handles both string and list responses from the agent. + """ + try: + if isinstance(response, str): + if len(response) <= 2000: + await message.reply(response) + else: + # Send first chunk as reply + chunks = [response[i:i+1900] for i in range(0, len(response), 1900)] + await message.reply(chunks[0]) + # Send remaining chunks as regular messages + for chunk in chunks[1:]: + await message.channel.send(chunk) + elif isinstance(response, list): + # Send first chunk as reply + if response: + await message.reply(response[0]) + # Send remaining chunks as regular messages + for chunk in response[1:]: + await message.channel.send(chunk) + except discord.errors.HTTPException as e: + error_msg = f"Error sending message: {str(e)}" + logger.error(error_msg) + await message.channel.send(error_msg[:1900]) @bot.event async def on_ready(): """ Called when the client is done preparing the data received from Discord. - Prints message on terminal when bot successfully connects to discord. - - https://discordpy.readthedocs.io/en/latest/api.html#discord.on_ready """ logger.info(f"{bot.user} has connected to Discord!") - @bot.event async def on_message(message: discord.Message): """ Called when a message is sent in any channel the bot can see. """ - # Don't delete this line! It's necessary for the bot to process commands. + # Process commands first await bot.process_commands(message) - # Ignore messages from self or other bots to prevent infinite loops. + # Ignore messages from self or other bots to prevent infinite loops if message.author.bot or message.content.startswith("!"): return - # Process the message with the agent you wrote + # Log the incoming message logger.info(f"Processing message from {message.author}: {message.content}") - # Show typing indicator while processing - async with message.channel.typing(): - response = await agent.run(message) - await message.reply(response) - - -# Commands + try: + async with message.channel.typing(): + response = await agent.run(message) + await send_split_message(message, response) + except Exception as e: + error_msg = f"An error occurred while processing the message: {str(e)}" + logger.error(error_msg) + await message.channel.send(error_msg[:1900]) - -# This example command is here to show you how to add commands to the bot. -# Run !ping with any number of arguments to see the command in action. -# Feel free to delete this if your project will not need commands. @bot.command(name="ping", help="Pings the bot.") async def ping(ctx, *, arg=None): if arg is None: @@ -72,6 +89,5 @@ async def ping(ctx, *, arg=None): else: await ctx.send(f"Pong! Your argument was {arg}") - -# Start the bot, connecting it to the gateway -bot.run(token) +# Start the bot +bot.run(token) \ No newline at end of file From bf1f5d38bfb8a45c9f47af45d2d881c74c9f1b06 Mon Sep 17 00:00:00 2001 From: bless-stanford <107686196+bless-stanford@users.noreply.github.com> Date: Mon, 27 Jan 2025 19:55:09 -0800 Subject: [PATCH 05/19] Remove unnecessary import comment for the Mistral agent in bot.py --- bot.py | 1 - 1 file changed, 1 deletion(-) diff --git a/bot.py b/bot.py index dadbc1b7..c138026c 100644 --- a/bot.py +++ b/bot.py @@ -17,7 +17,6 @@ intents = discord.Intents.all() bot = commands.Bot(command_prefix=PREFIX, intents=intents) -# Import the Mistral agent from the agent.py file agent = MistralAgent() # Get the token from the environment variables From 687a41cf3831e6fedd0fe6a404a618b65eafdbb2 Mon Sep 17 00:00:00 2001 From: bless-stanford <107686196+bless-stanford@users.noreply.github.com> Date: Mon, 27 Jan 2025 20:13:25 -0800 Subject: [PATCH 06/19] Implement chat history trimming to maintain context within a specified limit --- agent.py | 31 ++++++++++++++++++++++++++++++- 1 file changed, 30 insertions(+), 1 deletion(-) diff --git a/agent.py b/agent.py index daac3f11..325da95c 100644 --- a/agent.py +++ b/agent.py @@ -17,16 +17,45 @@ Keep responses concise when possible, as Discord has a 2000-character limit per message.""" class MistralAgent: - def __init__(self): + def __init__(self, max_context_messages=10): self.kernel = KernelBuilder.create_kernel(model_id=MISTRAL_MODEL) self.settings = KernelBuilder.get_default_settings() self.chat_service = self.kernel.get_service() self.chat_history = ChatHistory() self.chat_history.add_system_message(SYSTEM_PROMPT) self.MAX_LENGTH = 1900 # Leave room for extra characters + self.max_context_messages = max_context_messages + + def _trim_chat_history(self): + """Keep only the most recent messages within the context window.""" + # Count number of non-system messages + messages = [msg for msg in self.chat_history.messages if msg.role != "system"] + + # If we have more messages than our limit, remove oldest ones + if len(messages) > self.max_context_messages: + # Get the system message(s) + system_messages = [msg for msg in self.chat_history.messages if msg.role == "system"] + + # Keep only the most recent messages + recent_messages = messages[-self.max_context_messages:] + + # Reset chat history with system messages and recent context + self.chat_history = ChatHistory() + for msg in system_messages: + self.chat_history.add_system_message(msg.content) + + # Add back recent messages in order + for msg in recent_messages: + if msg.role == "user": + self.chat_history.add_user_message(msg.content) + elif msg.role == "assistant": + self.chat_history.add_assistant_message(msg.content) async def run(self, message: discord.Message): self.chat_history.add_user_message(message.content) + + # Trim history before getting response + self._trim_chat_history() response = await self.chat_service.get_chat_message_content( chat_history=self.chat_history, From c6659e21cafc0e2ab055930cecd41a35e7c45a6b Mon Sep 17 00:00:00 2001 From: bless-stanford <107686196+bless-stanford@users.noreply.github.com> Date: Sun, 9 Mar 2025 20:55:02 -0700 Subject: [PATCH 07/19] Add Box integration with commands for authorization, file upload, folder creation, and search --- .gitignore | 5 +- bot.py | 106 +++++++ plugins/__init__.py | 0 plugins/box_plugin.py | 0 pyproject.toml | 3 + server.py | 129 +++++++++ services/__init__.py | 0 services/box_models.py | 108 ++++++++ services/box_service.py | 592 ++++++++++++++++++++++++++++++++++++++++ 9 files changed, 942 insertions(+), 1 deletion(-) create mode 100644 plugins/__init__.py create mode 100644 plugins/box_plugin.py create mode 100644 server.py create mode 100644 services/__init__.py create mode 100644 services/box_models.py create mode 100644 services/box_service.py diff --git a/.gitignore b/.gitignore index cddb34f9..97edb728 100644 --- a/.gitignore +++ b/.gitignore @@ -10,4 +10,7 @@ wheels/ .venv # Secrets -.env \ No newline at end of file +.env + +user_tokens.json +temp/ \ No newline at end of file diff --git a/bot.py b/bot.py index c138026c..cdb55eea 100644 --- a/bot.py +++ b/bot.py @@ -4,6 +4,8 @@ from discord.ext import commands from dotenv import load_dotenv from agent import MistralAgent +from services.box_service import BoxService +from server import start_server # Import the server starter PREFIX = "!" @@ -88,5 +90,109 @@ async def ping(ctx, *, arg=None): else: await ctx.send(f"Pong! Your argument was {arg}") +@bot.command(name="authorize-box", help="Authorize the bot to access your Box account") +async def authorize_box(ctx): + """ + Sends a Box authorization link to the user via DM. + """ + try: + # Create Box service + box_service = BoxService() + + # Get authorization URL for the user + auth_url = await box_service.get_authorization_url(str(ctx.author.id)) + + # Send the URL as a DM to the user + await ctx.author.send(f"Please authorize access to your Box account by clicking this link: {auth_url}") + await ctx.send("I've sent you a DM with the authorization link!") + except Exception as e: + error_msg = f"Error generating authorization link: {str(e)}" + logger.error(error_msg) + await ctx.send(error_msg[:1900]) + +@bot.command(name="box-upload", help="Upload a file to Box") +async def box_upload(ctx): + """ + Uploads an attached file to Box. + """ + if not ctx.message.attachments: + await ctx.send("Please attach a file to upload.") + return + + attachment = ctx.message.attachments[0] + + # Create temp directory if it doesn't exist + if not os.path.exists("temp"): + os.makedirs("temp") + + # Download the attachment + file_path = f"temp/{attachment.filename}" + await attachment.save(file_path) + + try: + # Upload to Box + box_service = BoxService() + file_info = await box_service.upload_file(str(ctx.author.id), file_path, attachment.filename) + + # Send confirmation + await ctx.send(f"File uploaded to Box! File ID: {file_info['id']}") + + # Clean up temp file + if os.path.exists(file_path): + os.remove(file_path) + except Exception as e: + error_msg = f"Error uploading file: {str(e)}" + logger.error(error_msg) + await ctx.send(error_msg[:1900]) + + # Clean up temp file on error too + if os.path.exists(file_path): + os.remove(file_path) + +@bot.command(name="box-create-folder", help="Create a folder in Box") +async def box_create_folder(ctx, folder_name: str): + """ + Creates a folder in Box. + """ + try: + box_service = BoxService() + folder = await box_service.create_folder(str(ctx.author.id), folder_name) + await ctx.send(f"Folder created successfully! Folder ID: {folder['id']}") + except Exception as e: + error_msg = f"Error creating folder: {str(e)}" + logger.error(error_msg) + await ctx.send(error_msg[:1900]) + +@bot.command(name="box-search", help="Search for files in Box") +async def box_search(ctx, *, query: str): + """ + Searches for files in Box. + """ + try: + box_service = BoxService() + results = await box_service.search_for_file(str(ctx.author.id), query) + + if results and results.get('entries') and len(results['entries']) > 0: + files = results['entries'] + + # Create a nice formatted list of the files + response = "**Search Results:**\n\n" + for i, file in enumerate(files[:10], 1): # Limit to 10 files + response += f"{i}. **{file['name']}** (ID: {file['id']})\n" + + if len(files) > 10: + response += f"\n...and {len(files) - 10} more results." + + await ctx.send(response) + else: + await ctx.send("No files found matching your search query.") + except Exception as e: + error_msg = f"Error searching for files: {str(e)}" + logger.error(error_msg) + await ctx.send(error_msg[:1900]) + +# Start the web server in the background +server_thread = start_server(bot) + # Start the bot bot.run(token) \ No newline at end of file diff --git a/plugins/__init__.py b/plugins/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/plugins/box_plugin.py b/plugins/box_plugin.py new file mode 100644 index 00000000..e69de29b diff --git a/pyproject.toml b/pyproject.toml index 1295abc7..cf699405 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,4 +9,7 @@ dependencies = [ "discord-py>=2.4.0", "mistralai>=1.4.0", "python-dotenv>=1.0.1", + "fastapi>=0.95.0", + "uvicorn>=0.21.0", + "cryptography>=40.0.0", ] diff --git a/server.py b/server.py new file mode 100644 index 00000000..c1a4cb39 --- /dev/null +++ b/server.py @@ -0,0 +1,129 @@ +from fastapi import FastAPI, Request +from fastapi.responses import HTMLResponse +import uvicorn +from services.box_service import BoxService, TokenEncryptionHelper +import asyncio +import logging +import threading + +# Setup logging +logger = logging.getLogger("box_server") + +app = FastAPI() +box_service = BoxService() + +# This will be set from bot.py +bot = None + +@app.get("/") +async def root(): + return {"message": "Box OAuth Callback Server"} + +@app.get("/box/callback") +async def box_callback(code: str, state: str): + """ + Handle the OAuth callback from Box. + + This endpoint receives the authorization code from Box after a user + authorizes the application. It exchanges the code for access and + refresh tokens, stores them securely, and notifies the user. + """ + try: + # Get user ID from state + user_id = TokenEncryptionHelper.decrypt_token(state, box_service.encryption_key) + logger.info(f"Received callback for user {user_id}") + + # Handle the callback - this stores the tokens + await box_service.handle_auth_callback(state, code) + + # Notify the user through Discord + if bot: + # Schedule the notification in the bot's event loop + asyncio.run_coroutine_threadsafe(notify_user(user_id), bot.loop) + + # Return a nice HTML page with correct Content-Type + html_content = """ + + + + Authorization Successful + + + +
+
✅ Authorization Successful!
+
Your Box account has been connected to the Discord bot.
+
You can close this window and return to Discord.
+
You will also receive a confirmation message in Discord.
+
+ + + """ + + return HTMLResponse(content=html_content) + except Exception as e: + logger.error(f"Error in callback: {str(e)}") + return {"error": str(e)} + +async def notify_user(user_id): + """Send a Discord message to notify the user that authorization was successful.""" + try: + user = await bot.fetch_user(int(user_id)) + if user: + await user.send("✅ Your Box account has been successfully connected! You can now use Box commands.") + except Exception as e: + logger.error(f"Error notifying user: {str(e)}") + +def start_server(bot_instance=None): + """ + Start the FastAPI server in a background thread. + + Args: + bot_instance: The Discord bot instance, used for user notifications + + Returns: + thread: The thread running the server + """ + global bot + bot = bot_instance + + # Start the server in a separate thread + def run_server(): + logger.info("Starting Box OAuth callback server on port 8000") + uvicorn.run(app, host="0.0.0.0", port=8000, log_level="info") + + server_thread = threading.Thread(target=run_server, daemon=True) + server_thread.start() + logger.info("Server thread started") + + return server_thread \ No newline at end of file diff --git a/services/__init__.py b/services/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/services/box_models.py b/services/box_models.py new file mode 100644 index 00000000..ac40ea3a --- /dev/null +++ b/services/box_models.py @@ -0,0 +1,108 @@ +class BoxFolder: + """ + Model representing a Box folder. + """ + def __init__(self, + id=None, + name=None, + parent=None, + created_at=None, + modified_at=None, + item_collection=None, + **kwargs): + self.id = id + self.name = name + self.parent = parent + self.created_at = created_at + self.modified_at = modified_at + self.item_collection = item_collection + + # Store additional properties + for key, value in kwargs.items(): + setattr(self, key, value) + + @classmethod + def from_dict(cls, data): + """Create a BoxFolder instance from a dictionary.""" + return cls(**data) + + +class BoxFile: + """ + Model representing a Box file. + """ + def __init__(self, + id=None, + name=None, + parent=None, + created_at=None, + modified_at=None, + size=None, + extension=None, + shared_link=None, + **kwargs): + self.id = id + self.name = name + self.parent = parent + self.created_at = created_at + self.modified_at = modified_at + self.size = size + self.extension = extension + self.shared_link = shared_link + + # Store additional properties + for key, value in kwargs.items(): + setattr(self, key, value) + + @classmethod + def from_dict(cls, data): + """Create a BoxFile instance from a dictionary.""" + return cls(**data) + + +class BoxTokenData: + """ + Model representing Box token data for storage. + """ + def __init__(self, access_token=None, refresh_token=None, expires_at=None): + self.access_token = access_token + self.refresh_token = refresh_token + self.expires_at = expires_at + + def to_dict(self): + """Convert to dictionary for serialization.""" + return { + "access_token": self.access_token, + "refresh_token": self.refresh_token, + "expires_at": self.expires_at + } + + @classmethod + def from_dict(cls, data): + """Create a BoxTokenData instance from a dictionary.""" + return cls( + access_token=data.get("access_token"), + refresh_token=data.get("refresh_token"), + expires_at=data.get("expires_at") + ) + + +class BoxTokenResponse: + """ + Model representing the response from Box token endpoint. + """ + def __init__(self, access_token=None, refresh_token=None, expires_in=None, token_type=None): + self.access_token = access_token + self.refresh_token = refresh_token + self.expires_in = expires_in + self.token_type = token_type + + @classmethod + def from_dict(cls, data): + """Create a BoxTokenResponse instance from a dictionary.""" + return cls( + access_token=data.get("access_token"), + refresh_token=data.get("refresh_token"), + expires_in=data.get("expires_in"), + token_type=data.get("token_type") + ) \ No newline at end of file diff --git a/services/box_service.py b/services/box_service.py new file mode 100644 index 00000000..0e6988dc --- /dev/null +++ b/services/box_service.py @@ -0,0 +1,592 @@ +import os +import json +import requests +from datetime import datetime, timedelta +from cryptography.fernet import Fernet +from dotenv import load_dotenv +from urllib.parse import urlencode +import logging + +# Setup logging +logger = logging.getLogger("box_service") + +# Constants for platform and service +PLATFORM = "Box" +SERVICE = "BoxService" + +# API URLs +BOX_API_BASE_URL = "https://api.box.com/2.0/" +BOX_AUTH_BASE_URL = "https://account.box.com/api/oauth2/" +BOX_UPLOAD_API_BASE_URL = "https://upload.box.com/api/2.0/" + +class TokenEncryptionHelper: + @staticmethod + def encrypt_token(token_str, encryption_key): + """Encrypts a token string using Fernet symmetric encryption.""" + f = Fernet(encryption_key) + return f.encrypt(token_str.encode()).decode() + + @staticmethod + def decrypt_token(encrypted_token, encryption_key): + """Decrypts an encrypted token string using Fernet symmetric encryption.""" + f = Fernet(encryption_key) + return f.decrypt(encrypted_token.encode()).decode() + +class TokenStorageManager: + """A simple file-based token storage system.""" + + def __init__(self, storage_file="user_tokens.json"): + self.storage_file = storage_file + # Initialize the storage file if it doesn't exist + if not os.path.exists(storage_file): + with open(storage_file, 'w') as f: + json.dump({}, f) + + def get_token(self, user_id, platform, service): + """Retrieve a token from storage.""" + try: + with open(self.storage_file, 'r') as f: + tokens = json.load(f) + + key = f"{user_id}_{platform}_{service}" + return tokens.get(key) + except Exception as e: + logger.error(f"Error retrieving token: {str(e)}") + return None + + def store_token(self, user_id, platform, service, token_data): + """Store a token in storage.""" + try: + with open(self.storage_file, 'r') as f: + tokens = json.load(f) + + key = f"{user_id}_{platform}_{service}" + tokens[key] = token_data + + with open(self.storage_file, 'w') as f: + json.dump(tokens, f) + + return True + except Exception as e: + logger.error(f"Error storing token: {str(e)}") + return False + + def delete_token(self, user_id, platform, service): + """Delete a token from storage.""" + try: + with open(self.storage_file, 'r') as f: + tokens = json.load(f) + + key = f"{user_id}_{platform}_{service}" + if key in tokens: + del tokens[key] + + with open(self.storage_file, 'w') as f: + json.dump(tokens, f) + + return True + except Exception as e: + logger.error(f"Error deleting token: {str(e)}") + return False + +class BoxService: + def __init__(self, config=None): + """ + Initialize the Box service with configuration. + + Args: + config: Configuration dictionary or None to load from .env + """ + if config is None: + load_dotenv() + self.client_id = os.getenv("BOX_CLIENT_ID") + self.client_secret = os.getenv("BOX_CLIENT_SECRET") + self.redirect_uri = os.getenv("BOX_REDIRECT_URI") + + # Get or generate encryption key + encryption_key = os.getenv("ENCRYPTION_KEY") + if not encryption_key: + # Generate a new key if none exists + encryption_key = Fernet.generate_key().decode() + # Add this to your .env file manually or update it + logger.warning("No encryption key found. Generated new key. Add to .env: " + f"ENCRYPTION_KEY={encryption_key}") + + self.encryption_key = encryption_key.encode() if isinstance(encryption_key, str) else encryption_key + else: + self.client_id = config.get("client_id") + self.client_secret = config.get("client_secret") + self.redirect_uri = config.get("redirect_uri") + self.encryption_key = config.get("encryption_key", Fernet.generate_key()) + + # Initialize token storage + self.token_storage = TokenStorageManager() + + async def get_authorization_url(self, user_id): + """ + Get the authorization URL for Box OAuth flow. + + Args: + user_id: The user's ID + + Returns: + str: The authorization URL + """ + if not self.client_id: + raise ValueError("Box Client ID is not set in configuration.") + if not self.redirect_uri: + raise ValueError("Box Redirect URI is not set in configuration.") + + # Encrypt user_id as state parameter + state = TokenEncryptionHelper.encrypt_token(user_id, self.encryption_key) + + query = { + "client_id": self.client_id, + "response_type": "code", + "redirect_uri": self.redirect_uri, + "state": state + } + + query_string = urlencode(query) + return f"{BOX_AUTH_BASE_URL}authorize?{query_string}" + + async def handle_auth_callback(self, state, code): + """ + Handle the authorization callback from Box. + + Args: + state: The state parameter from the callback + code: The authorization code from the callback + """ + if not self.client_id: + raise ValueError("Box Client ID is not set in configuration.") + if not self.client_secret: + raise ValueError("Box Client Secret is not set in configuration.") + if not self.redirect_uri: + raise ValueError("Box Redirect URI is not set in configuration.") + + # Decrypt the user_id from state + user_id = TokenEncryptionHelper.decrypt_token(state, self.encryption_key) + + payload = { + "grant_type": "authorization_code", + "code": code, + "client_id": self.client_id, + "client_secret": self.client_secret, + "redirect_uri": self.redirect_uri + } + + response = requests.post(f"{BOX_AUTH_BASE_URL}token", data=payload) + response_data = response.json() + + if response.status_code == 200 and "access_token" in response_data: + await self._store_token( + user_id, + response_data["access_token"], + response_data["refresh_token"], + response_data["expires_in"] + ) + else: + error_msg = response_data.get("error_description", "Unknown error") + raise Exception(f"Failed to obtain user access token: {error_msg}") + + async def revoke_access(self, user_id): + """ + Revoke the Box access for a user. + + Args: + user_id: The user's ID + """ + token = await self._load_token(user_id) + if not token: + raise ValueError("No valid token found for user") + + if not self.client_id: + raise ValueError("Box Client ID is not set in configuration.") + if not self.client_secret: + raise ValueError("Box Client Secret is not set in configuration.") + + payload = { + "client_id": self.client_id, + "client_secret": self.client_secret, + "token": token + } + + response = requests.post(f"{BOX_AUTH_BASE_URL}revoke", data=payload) + + if response.status_code == 200: + # Delete the token from storage + self.token_storage.delete_token(user_id, PLATFORM, SERVICE) + else: + raise Exception(f"Failed to revoke token: {response.status_code}") + + async def create_folder(self, user_id, folder_name, parent_folder_id="0"): + """ + Create a folder in Box. + + Args: + user_id: The user's ID + folder_name: Name of the folder to create + parent_folder_id: ID of the parent folder (default: "0" for root) + + Returns: + dict: The created folder information + """ + token = await self._load_token(user_id) + if not token: + raise Exception("Failed to load a valid token") + + headers = { + "Authorization": f"Bearer {token}", + "Content-Type": "application/json" + } + + payload = { + "name": folder_name, + "parent": {"id": parent_folder_id} + } + + response = requests.post( + f"{BOX_API_BASE_URL}folders", + headers=headers, + json=payload + ) + + if response.status_code in (200, 201): + return response.json() + else: + error_msg = response.json().get("message", "Unknown error") + raise Exception(f"Box API request failed: {error_msg}") + + async def search_for_file(self, user_id, query, limit=100): + """ + Search for files in Box. + + Args: + user_id: The user's ID + query: Search query + limit: Maximum number of results to return + + Returns: + dict: Search results + """ + token = await self._load_token(user_id) + if not token: + raise Exception("Failed to load a valid token") + + headers = { + "Authorization": f"Bearer {token}" + } + + params = { + "query": query, + "limit": limit, + "type": "file" + } + + response = requests.get( + f"{BOX_API_BASE_URL}search", + headers=headers, + params=params + ) + + if response.status_code == 200: + return response.json() + else: + error_msg = response.json().get("message", "Unknown error") + raise Exception(f"Box API request failed: {error_msg}") + + async def delete_file(self, user_id, file_id): + """ + Delete a file from Box. + + Args: + user_id: The user's ID + file_id: ID of the file to delete + """ + token = await self._load_token(user_id) + if not token: + raise Exception("Failed to load a valid token") + + headers = { + "Authorization": f"Bearer {token}" + } + + response = requests.delete( + f"{BOX_API_BASE_URL}files/{file_id}", + headers=headers + ) + + if response.status_code != 204: # 204 No Content is success + error_msg = "Unknown error" + try: + error_msg = response.json().get("message", "Unknown error") + except: + pass + raise Exception(f"Box API request failed: {error_msg}") + + async def upload_file(self, user_id, file_path, original_file_name, folder_id="0"): + """ + Upload a file to Box. + + Args: + user_id: The user's ID + file_path: Path to the local file + original_file_name: Original name of the file + folder_id: ID of the folder to upload to (default: "0" for root) + + Returns: + dict: The uploaded file information + """ + token = await self._load_token(user_id) + if not token: + raise Exception("Failed to load a valid token") + + headers = { + "Authorization": f"Bearer {token}" + } + + attributes = json.dumps({ + "name": original_file_name, + "parent": {"id": folder_id} + }) + + with open(file_path, 'rb') as file: + files = { + 'attributes': (None, attributes, 'application/json'), + 'file': (original_file_name, file, 'application/octet-stream') + } + + response = requests.post( + f"{BOX_UPLOAD_API_BASE_URL}files/content", + headers=headers, + files=files + ) + + if response.status_code in (200, 201): + return response.json()['entries'][0] # Box returns an entries array + else: + error_msg = "Unknown error" + try: + error_msg = response.json().get("message", "Unknown error") + except: + pass + raise Exception(f"Box API request failed: {error_msg}") + + async def get_file_download_link(self, user_id, file_id): + """ + Get a download link for a file. + + Args: + user_id: The user's ID + file_id: ID of the file + + Returns: + str: Download URL + """ + token = await self._load_token(user_id) + if not token: + raise Exception("Failed to load a valid token") + + headers = { + "Authorization": f"Bearer {token}" + } + + # Box API redirects to the actual download URL, so we need to disable redirects + response = requests.get( + f"{BOX_API_BASE_URL}files/{file_id}/content", + headers=headers, + allow_redirects=False + ) + + if response.status_code == 302: # Redirect status code + return response.headers.get('Location') + else: + error_msg = "Unknown error" + try: + error_msg = response.json().get("message", "Unknown error") + except: + pass + raise Exception(f"Box API request failed: {error_msg}") + + async def get_file_view_link(self, user_id, file_id): + """ + Get a shared view link for a file. + + Args: + user_id: The user's ID + file_id: ID of the file + + Returns: + str: Shared URL + """ + token = await self._load_token(user_id) + if not token: + raise Exception("Failed to load a valid token") + + headers = { + "Authorization": f"Bearer {token}", + "Content-Type": "application/json" + } + + payload = { + "shared_link": {"access": "open"} + } + + response = requests.put( + f"{BOX_API_BASE_URL}files/{file_id}", + headers=headers, + json=payload + ) + + if response.status_code == 200: + data = response.json() + if "shared_link" in data and "url" in data["shared_link"]: + return data["shared_link"]["url"] + else: + raise Exception("Shared link URL not found in response") + else: + error_msg = response.json().get("message", "Unknown error") + raise Exception(f"Box API request failed: {error_msg}") + + async def share_file(self, user_id, file_id, email, role): + """ + Share a file with another user. + + Args: + user_id: The user's ID + file_id: ID of the file to share + email: Email of the user to share with + role: Role to assign (editor, viewer, etc.) + """ + token = await self._load_token(user_id) + if not token: + raise Exception("Failed to load a valid token") + + headers = { + "Authorization": f"Bearer {token}", + "Content-Type": "application/json" + } + + payload = { + "item": {"id": file_id, "type": "file"}, + "accessible_by": {"type": "user", "login": email}, + "role": role + } + + response = requests.post( + f"{BOX_API_BASE_URL}collaborations", + headers=headers, + json=payload + ) + + if response.status_code not in (200, 201): + error_msg = response.json().get("message", "Unknown error") + raise Exception(f"Box API request failed: {error_msg}") + + async def _store_token(self, user_id, access_token, refresh_token, expires_in): + """ + Store a token in the token storage. + + Args: + user_id: The user's ID + access_token: The access token + refresh_token: The refresh token + expires_in: Expiration time in seconds + """ + token_data = { + "access_token": access_token, + "refresh_token": refresh_token, + "expires_at": (datetime.utcnow() + timedelta(seconds=expires_in)).timestamp() + } + + # Serialize and encrypt the token data + serialized_token = json.dumps(token_data) + encrypted_token = TokenEncryptionHelper.encrypt_token(serialized_token, self.encryption_key) + + # Store in the token storage + token_record = { + "encrypted_token": encrypted_token, + "is_active": True, + "is_revoked": False, + "created_at": datetime.utcnow().timestamp() + } + + self.token_storage.store_token(user_id, PLATFORM, SERVICE, token_record) + + async def _load_token(self, user_id): + """ + Load a token from the token storage. + + Args: + user_id: The user's ID + + Returns: + str: The access token, or None if not found or expired + """ + token_record = self.token_storage.get_token(user_id, PLATFORM, SERVICE) + + if not token_record or not token_record.get("is_active") or token_record.get("is_revoked"): + logger.info("No valid token found in the storage") + return None + + try: + encrypted_token = token_record.get("encrypted_token") + if not encrypted_token: + return None + + decrypted_token = TokenEncryptionHelper.decrypt_token(encrypted_token, self.encryption_key) + token_data = json.loads(decrypted_token) + + if not token_data: + logger.error("Failed to deserialize token data") + return None + + # Check if token is expired + expires_at = token_data.get("expires_at") + if expires_at and expires_at <= datetime.utcnow().timestamp(): + logger.info("Token expired, attempting to refresh") + return await self._refresh_token(user_id, token_data.get("refresh_token")) + + return token_data.get("access_token") + except Exception as e: + logger.error(f"Error loading token: {str(e)}") + return None + + async def _refresh_token(self, user_id, refresh_token): + """ + Refresh an expired token. + + Args: + user_id: The user's ID + refresh_token: The refresh token + + Returns: + str: The new access token + """ + if not self.client_id: + raise ValueError("Box Client ID is not set in configuration.") + if not self.client_secret: + raise ValueError("Box Client Secret is not set in configuration.") + + payload = { + "grant_type": "refresh_token", + "refresh_token": refresh_token, + "client_id": self.client_id, + "client_secret": self.client_secret + } + + response = requests.post(f"{BOX_AUTH_BASE_URL}token", data=payload) + response_data = response.json() + + if response.status_code == 200 and "access_token" in response_data: + await self._store_token( + user_id, + response_data["access_token"], + response_data["refresh_token"], + response_data["expires_in"] + ) + return response_data["access_token"] + else: + error_msg = response_data.get("error_description", "Unknown error") + raise Exception(f"Failed to refresh token: {error_msg}") \ No newline at end of file From cd4cdc3fe1fa51aee3e7d2f9e624a10bc3c16c6a Mon Sep 17 00:00:00 2001 From: bless-stanford <107686196+bless-stanford@users.noreply.github.com> Date: Sun, 9 Mar 2025 23:49:57 -0700 Subject: [PATCH 08/19] Enhance Mistral agent to process natural language requests for Box integration and improve logging --- agent.py | 104 ++++++++-- bot.py | 52 +---- plugins/box_plugin.py | 440 ++++++++++++++++++++++++++++++++++++++++ services/box_service.py | 129 ++++++++---- 4 files changed, 631 insertions(+), 94 deletions(-) diff --git a/agent.py b/agent.py index 325da95c..5fbea6c2 100644 --- a/agent.py +++ b/agent.py @@ -2,9 +2,26 @@ import re from semantic_kernel.contents import ChatHistory from kernel.kernel_builder import KernelBuilder +from services.box_service import BoxService +from plugins.box_plugin import BoxPlugins +from semantic_kernel.functions import KernelArguments +from semantic_kernel.connectors.ai.function_choice_behavior import FunctionChoiceBehavior +import logging +import json + +logger = logging.getLogger("agent") MISTRAL_MODEL = "mistral-large-latest" -SYSTEM_PROMPT = """You are a helpful assistant. Your name is Dodobot. Format your responses using Discord-compatible Markdown: +SYSTEM_PROMPT = """You are a helpful assistant named Dodobot that can access and manage a user's Box cloud storage. +You can search for files, create folders, get download links, view links, share files, and delete files. +You can process natural language requests about Box files. + +When a user asks about Box files, folders, or cloud storage, use the appropriate Box function to handle their request. +Never ask the user for their user ID, as it is automatically provided by the system. + +If a user's Box account needs to be authorized or reauthorized, tell them to use the !authorize-box command. + +Format your responses using Discord-compatible Markdown: - Use **bold** for emphasis - Use *italics* for secondary emphasis - Use `code` for technical terms, commands, or short code snippets @@ -13,6 +30,7 @@ ``` for multi-line code (where 'language' is python, javascript, etc) - Use > for quotes - Use ||text|| for spoilers + Do not use # for headers or * - for bullet points as these don't render in Discord. Keep responses concise when possible, as Discord has a 2000-character limit per message.""" @@ -20,11 +38,20 @@ class MistralAgent: def __init__(self, max_context_messages=10): self.kernel = KernelBuilder.create_kernel(model_id=MISTRAL_MODEL) self.settings = KernelBuilder.get_default_settings() + + # Enable function calling in the settings + self.settings.function_choice_behavior = FunctionChoiceBehavior.Auto() + self.chat_service = self.kernel.get_service() self.chat_history = ChatHistory() self.chat_history.add_system_message(SYSTEM_PROMPT) self.MAX_LENGTH = 1900 # Leave room for extra characters self.max_context_messages = max_context_messages + + # Register Box plugins with the kernel + self.box_service = BoxService() + self.box_plugins = BoxPlugins(self.box_service) + self.kernel.add_plugin(self.box_plugins, "Box") def _trim_chat_history(self): """Keep only the most recent messages within the context window.""" @@ -52,22 +79,75 @@ def _trim_chat_history(self): self.chat_history.add_assistant_message(msg.content) async def run(self, message: discord.Message): - self.chat_history.add_user_message(message.content) + original_content = message.content + user_id = str(message.author.id) + + augmented_content = f"{original_content}\n\n[system: user_id={user_id}]" + + # Add the user's message to the chat history + self.chat_history.add_user_message(augmented_content) # Trim history before getting response self._trim_chat_history() - - response = await self.chat_service.get_chat_message_content( - chat_history=self.chat_history, - settings=self.settings - ) - self.chat_history.add_assistant_message(response.content) + # Add the user ID to the kernel arguments for Box plugin access + kernel_arguments = KernelArguments() + # Store the user ID with the exact key name expected by the plugin functions + user_id = str(message.author.id) + kernel_arguments["user_id"] = user_id + + # Log the user ID to verify it's correct + logger.info(f"Setting user_id in kernel arguments to: {user_id}") - # Always return a list of chunks, even for short messages - if len(response.content) > self.MAX_LENGTH: - return self.split_response(response.content) - return [response.content] + try: + # Log the request for debugging + logger.info(f"Processing request from user {message.author.id}: {message.content}") + + # Get response with function calling enabled + response = await self.chat_service.get_chat_message_content( + chat_history=self.chat_history, + settings=self.settings, + kernel=self.kernel, + arguments=kernel_arguments + ) + + # Check if response content contains Box auth error + if "Your Box authorization has expired" in response.content or \ + "Please use the `!authorize-box` command" in response.content: + error_message = ( + "It looks like I need access to your Box account to perform this task. " + "Please use the `!authorize-box` command to connect your Box account." + ) + self.chat_history.add_assistant_message(error_message) + return [error_message] + + # Add the assistant's response to the chat history + self.chat_history.add_assistant_message(response.content) + + # Log the response for debugging (exclude sensitive data) + logger.info(f"Generated response for user {message.author.id} (length: {len(response.content)})") + + # Always return a list of chunks, even for short messages + if len(response.content) > self.MAX_LENGTH: + return self.split_response(response.content) + return [response.content] + + except Exception as e: + logger.error(f"Error processing request: {str(e)}", exc_info=True) + + # Check if this is an authentication error + if "Box authorization has expired" in str(e) or "!authorize-box" in str(e): + error_message = ( + "I need to connect to your Box account to perform this task. " + "Please use the `!authorize-box` command to authorize access." + ) + self.chat_history.add_assistant_message(error_message) + return [error_message] + + # For other errors + error_message = f"Sorry, I encountered an error while processing your request. Please try again later." + self.chat_history.add_assistant_message(error_message) + return [error_message] def split_response(self, content: str) -> list[str]: chunks = [] diff --git a/bot.py b/bot.py index cdb55eea..cc983bc7 100644 --- a/bot.py +++ b/bot.py @@ -5,12 +5,15 @@ from dotenv import load_dotenv from agent import MistralAgent from services.box_service import BoxService -from server import start_server # Import the server starter +from server import start_server PREFIX = "!" # Setup logging logger = logging.getLogger("discord") +logging.basicConfig(level=logging.INFO, + format='[%(asctime)s] %(levelname)s - %(name)s: %(message)s', + datefmt='%Y-%m-%d %H:%M:%S') # Load the environment variables load_dotenv() @@ -19,6 +22,7 @@ intents = discord.Intents.all() bot = commands.Bot(command_prefix=PREFIX, intents=intents) +# Initialize agent with Semantic Kernel and Box plugins agent = MistralAgent() # Get the token from the environment variables @@ -68,7 +72,7 @@ async def on_message(message: discord.Message): await bot.process_commands(message) # Ignore messages from self or other bots to prevent infinite loops - if message.author.bot or message.content.startswith("!"): + if message.author.bot or message.content.startswith(PREFIX): return # Log the incoming message @@ -76,6 +80,8 @@ async def on_message(message: discord.Message): try: async with message.channel.typing(): + # The agent will now use Semantic Kernel to process natural language + # requests related to Box, without requiring specific commands response = await agent.run(message) await send_split_message(message, response) except Exception as e: @@ -149,48 +155,6 @@ async def box_upload(ctx): if os.path.exists(file_path): os.remove(file_path) -@bot.command(name="box-create-folder", help="Create a folder in Box") -async def box_create_folder(ctx, folder_name: str): - """ - Creates a folder in Box. - """ - try: - box_service = BoxService() - folder = await box_service.create_folder(str(ctx.author.id), folder_name) - await ctx.send(f"Folder created successfully! Folder ID: {folder['id']}") - except Exception as e: - error_msg = f"Error creating folder: {str(e)}" - logger.error(error_msg) - await ctx.send(error_msg[:1900]) - -@bot.command(name="box-search", help="Search for files in Box") -async def box_search(ctx, *, query: str): - """ - Searches for files in Box. - """ - try: - box_service = BoxService() - results = await box_service.search_for_file(str(ctx.author.id), query) - - if results and results.get('entries') and len(results['entries']) > 0: - files = results['entries'] - - # Create a nice formatted list of the files - response = "**Search Results:**\n\n" - for i, file in enumerate(files[:10], 1): # Limit to 10 files - response += f"{i}. **{file['name']}** (ID: {file['id']})\n" - - if len(files) > 10: - response += f"\n...and {len(files) - 10} more results." - - await ctx.send(response) - else: - await ctx.send("No files found matching your search query.") - except Exception as e: - error_msg = f"Error searching for files: {str(e)}" - logger.error(error_msg) - await ctx.send(error_msg[:1900]) - # Start the web server in the background server_thread = start_server(bot) diff --git a/plugins/box_plugin.py b/plugins/box_plugin.py index e69de29b..6e379370 100644 --- a/plugins/box_plugin.py +++ b/plugins/box_plugin.py @@ -0,0 +1,440 @@ +from semantic_kernel.functions import kernel_function +from semantic_kernel.functions.kernel_function_from_prompt import KernelFunctionFromPrompt +from services.box_service import BoxService +import logging + +logger = logging.getLogger("box_plugins") + +class BoxPlugins: + """ + Plugins for interacting with Box cloud storage. + """ + + def __init__(self, box_service=None): + """ + Initialize the Box plugins with a BoxService. + If no service is provided, a new one will be created. + """ + self.box_service = box_service or BoxService() + + @kernel_function( + name="create_folder", + description="Creates a new folder in the user's Box account" + ) + async def create_folder( + self, + folder_name: str, + parent_folder_id: str = "0", + user_id: str = None, + kernel = None + ) -> str: + """ + Creates a new folder in the user's Box account. + + Args: + folder_name: Name of the folder to create + parent_folder_id: ID of the parent folder (default: "0" for root) + user_id: The user's ID (automatically provided) + + Returns: + str: Success message with folder details or error message + """ + try: + if not user_id: + return "Error: User ID not available. Please try again later." + + folder = await self.box_service.create_folder(user_id, folder_name, parent_folder_id) + + if folder: + return f"Folder '{folder_name}' created successfully with ID: {folder['id']}" + else: + return f"Failed to create folder '{folder_name}'." + + except Exception as e: + logger.error(f"Error creating folder: {str(e)}") + return f"An error occurred while creating the folder: {str(e)}" + + @kernel_function( + name="search_file", + description="Searches for files in the user's Box account and returns details in a user friendly way" + ) + async def search_file( + self, + query: str, + user_id: str = None, + kernel = None + ) -> str: + """ + Searches for files in the user's Box account. + + Args: + query: Search query or file name + user_id: The user's ID (automatically provided) + + Returns: + str: File details or search results summary + """ + try: + # Get user_id from kernel.data instead of function parameter + if not user_id and kernel and hasattr(kernel, 'arguments'): + user_id = kernel.arguments.get("user_id") + + if not user_id: + return "Error: User ID not available. Please try again later." + + search_results = await self.box_service.search_for_file(user_id, query) + + if not search_results or not search_results.get('entries') or len(search_results['entries']) == 0: + return f"No files found matching '{query}'." + + files = search_results['entries'] + + if len(files) == 1: + return self._create_file_detail(files[0]) + + # If multiple files, return a summary + return self._create_search_results_summary(files) + + except Exception as e: + logger.error(f"Error searching files: {str(e)}") + return f"An error occurred while searching for files: {str(e)}" + + @kernel_function( + name="delete_file", + description="Searches and deletes a file from the user's Box account" + ) + async def delete_file( + self, + query: str, + user_id: str = None, + kernel = None + ) -> str: + """ + Searches and deletes a file from the user's Box account. + + Args: + query: Search query or file name + user_id: The user's ID (automatically provided) + kernel: Semantic Kernel instance for finding most relevant file + + Returns: + str: Success message or error message + """ + try: + if not user_id: + return "Error: User ID not available. Please try again later." + + search_results = await self.box_service.search_for_file(user_id, query) + + if not search_results or not search_results.get('entries') or len(search_results['entries']) == 0: + return f"No files found matching '{query}'." + + files = search_results['entries'] + + if len(files) == 1: + file = files[0] + await self.box_service.delete_file(user_id, file['id']) + return f"File '{file['name']}' has been successfully deleted." + + # If multiple files and kernel is provided, find most relevant + if kernel and len(files) > 1: + most_relevant_file = await self._find_most_relevant_file(kernel, files, query) + + if most_relevant_file: + await self.box_service.delete_file(user_id, most_relevant_file['id']) + return f"File '{most_relevant_file['name']}' has been successfully deleted." + + # If multiple files and no most relevant found, return summary + return f"Multiple files found matching '{query}'. Please be more specific:\n" + \ + "\n".join([f"- {file['name']}" for file in files[:5]]) + + except Exception as e: + logger.error(f"Error deleting file: {str(e)}") + return f"An error occurred while deleting the file: {str(e)}" + + @kernel_function( + name="get_file_download_link", + description="Gets a download link for a file in the user's Box account" + ) + async def get_file_download_link( + self, + query: str, + user_id: str = None, + kernel = None + ) -> str: + """ + Gets a download link for a file in the user's Box account. + + Args: + query: Search query or file name + user_id: The user's ID (automatically provided) + kernel: Semantic Kernel instance for finding most relevant file + + Returns: + str: Download link or error message + """ + try: + if not user_id: + return "Error: User ID not available. Please try again later." + + search_results = await self.box_service.search_for_file(user_id, query) + + if not search_results or not search_results.get('entries') or len(search_results['entries']) == 0: + return f"No files found matching '{query}'." + + files = search_results['entries'] + + if len(files) == 1: + file = files[0] + download_link = await self.box_service.get_file_download_link(user_id, file['id']) + return f"Download link for file '{file['name']}':\n{download_link}" + + # If multiple files and kernel is provided, find most relevant + if kernel and len(files) > 1: + most_relevant_file = await self._find_most_relevant_file(kernel, files, query) + + if most_relevant_file: + download_link = await self.box_service.get_file_download_link(user_id, most_relevant_file['id']) + return f"Download link for file '{most_relevant_file['name']}':\n{download_link}" + + # If multiple files and no most relevant found, return summary + return f"Multiple files found matching '{query}'. Please be more specific:\n" + \ + "\n".join([f"- {file['name']}" for file in files[:5]]) + + except Exception as e: + logger.error(f"Error getting download link: {str(e)}") + return f"An error occurred while getting the download link: {str(e)}" + + @kernel_function( + name="get_file_view_link", + description="Gets a shareable view link for a file in the user's Box account" + ) + async def get_file_view_link( + self, + query: str, + user_id: str = None, + kernel = None + ) -> str: + """ + Gets a shareable view link for a file in the user's Box account. + + Args: + query: Search query or file name + user_id: The user's ID (automatically provided) + kernel: Semantic Kernel instance for finding most relevant file + + Returns: + str: View link or error message + """ + try: + if not user_id: + return "Error: User ID not available. Please try again later." + + search_results = await self.box_service.search_for_file(user_id, query) + + if not search_results or not search_results.get('entries') or len(search_results['entries']) == 0: + return f"No files found matching '{query}'." + + files = search_results['entries'] + + if len(files) == 1: + file = files[0] + view_link = await self.box_service.get_file_view_link(user_id, file['id']) + return f"View link for file '{file['name']}':\n{view_link}" + + # If multiple files and kernel is provided, find most relevant + if kernel and len(files) > 1: + most_relevant_file = await self._find_most_relevant_file(kernel, files, query) + + if most_relevant_file: + view_link = await self.box_service.get_file_view_link(user_id, most_relevant_file['id']) + return f"View link for file '{most_relevant_file['name']}':\n{view_link}" + + # If multiple files and no most relevant found, return summary + return f"Multiple files found matching '{query}'. Please be more specific:\n" + \ + "\n".join([f"- {file['name']}" for file in files[:5]]) + + except Exception as e: + logger.error(f"Error getting view link: {str(e)}") + return f"An error occurred while getting the view link: {str(e)}" + + @kernel_function( + name="share_file", + description="Shares a file with another user" + ) + async def share_file( + self, + query: str, + email: str, + role: str = "viewer", + user_id: str = None, + kernel = None + ) -> str: + """ + Shares a file with another user. + + Args: + query: Search query or file name + email: Email of the user to share with + role: Role to assign (viewer, editor, etc.) + user_id: The user's ID (automatically provided) + kernel: Semantic Kernel instance for finding most relevant file + + Returns: + str: Success message or error message + """ + try: + if not user_id: + return "Error: User ID not available. Please try again later." + + search_results = await self.box_service.search_for_file(user_id, query) + + if not search_results or not search_results.get('entries') or len(search_results['entries']) == 0: + return f"No files found matching '{query}'." + + files = search_results['entries'] + + if len(files) == 1: + file = files[0] + await self.box_service.share_file(user_id, file['id'], email, role) + view_link = await self.box_service.get_file_view_link(user_id, file['id']) + return f"File '{file['name']}' has been shared with {email} as a {role}. They can access the file at: {view_link}" + + # If multiple files and kernel is provided, find most relevant + if kernel and len(files) > 1: + most_relevant_file = await self._find_most_relevant_file(kernel, files, query) + + if most_relevant_file: + await self.box_service.share_file(user_id, most_relevant_file['id'], email, role) + view_link = await self.box_service.get_file_view_link(user_id, most_relevant_file['id']) + return f"File '{most_relevant_file['name']}' has been shared with {email} as a {role}. They can access the file at: {view_link}" + + # If multiple files and no most relevant found, return summary + return f"Multiple files found matching '{query}'. Please be more specific:\n" + \ + "\n".join([f"- {file['name']}" for file in files[:5]]) + + except Exception as e: + logger.error(f"Error sharing file: {str(e)}") + return f"An error occurred while sharing the file: {str(e)}" + + async def _find_most_relevant_file(self, kernel, files, user_query): + """ + Find the most relevant file from a list based on user query. + + Args: + kernel: Semantic Kernel instance + files: List of files + user_query: The user's query + + Returns: + dict: The most relevant file or None + """ + try: + # Create a function from prompt + rank_files_function = KernelFunctionFromPrompt( + function_name="RankFilesByRelevance", + plugin_name=None, + prompt="Given the user query: '{{$userQuery}}' and a list of file names, " + "rank them by relevance and return the index of the most relevant file. " + "Do not add any comments or explanation to the response.\n" + "File list: {{$fileList}}", + template_format="semantic-kernel" + ) + + # Create file list string + file_list = "\n".join([f"{i}: Name: {file['name']}" for i, file in enumerate(files)]) + + # Create kernel arguments + kernel_arguments = { + "userQuery": user_query, + "fileList": file_list + } + + # Invoke the function + result = await kernel.invoke(rank_files_function, **kernel_arguments) + + # Get the value from the result - might be a list or a string + result_value = result.value + + # Handle different result types + if isinstance(result_value, list) and len(result_value) > 0: + result_text = str(result_value[0]).strip() + elif isinstance(result_value, str): + result_text = result_value.strip() + else: + # Fallback + result_text = str(result_value).strip() + + try: + most_relevant_index = int(result_text) + if 0 <= most_relevant_index < len(files): + return files[most_relevant_index] + except ValueError: + logger.warning(f"Could not parse the relevance index from AI result: {result_text}") + + return None + + except Exception as e: + logger.error(f"Error finding most relevant file: {str(e)}") + return None + + def _create_file_detail(self, file): + """Create a detailed text representation of a file.""" + detail = "**File Details:**\n" + detail += f"**Name:** {file.get('name', 'Unknown')}\n" + detail += f"**ID:** {file.get('id', 'Unknown')}\n" + + # Format file size if available + if 'size' in file: + size_bytes = file['size'] + size_str = self._format_file_size(size_bytes) + detail += f"**Size:** {size_str}\n" + + # Add dates if available + if 'created_at' in file: + detail += f"**Created At:** {file['created_at']}\n" + if 'modified_at' in file: + detail += f"**Modified At:** {file['modified_at']}\n" + + # Add shared link if available + if 'shared_link' in file and file['shared_link']: + shared_link = file['shared_link'] + if 'url' in shared_link: + detail += f"**Shared Link:** {shared_link['url']}\n" + if 'download_url' in shared_link: + detail += f"**Download URL:** {shared_link['download_url']}\n" + + return detail + + def _create_search_results_summary(self, files): + """Create a summary of multiple search results.""" + summary = "**Multiple files found. Here are the details:**\n\n" + + for i, file in enumerate(files[:5], 1): # Limit to 5 files and number them + summary += f"**{i}. {file.get('name', 'Unknown')}**\n" + summary += f" ID: {file.get('id', 'Unknown')}\n" + + # Add size if available + if 'size' in file: + size_str = self._format_file_size(file['size']) + summary += f" Size: {size_str}\n" + + summary += "\n" + + if len(files) > 5: + summary += f"\n...and {len(files) - 5} more files.\n" + + summary += "\nPlease provide a more specific query to find the exact file you want." + + return summary + + def _format_file_size(self, bytes): + """Format file size in human-readable form.""" + sizes = ["B", "KB", "MB", "GB", "TB"] + order = 0 + size = float(bytes) + while size >= 1024 and order < len(sizes) - 1: + order += 1 + size /= 1024 + + return f"{size:.2f} {sizes[order]}" \ No newline at end of file diff --git a/services/box_service.py b/services/box_service.py index 0e6988dc..d44938cf 100644 --- a/services/box_service.py +++ b/services/box_service.py @@ -66,6 +66,7 @@ def store_token(self, user_id, platform, service, token_data): with open(self.storage_file, 'w') as f: json.dump(tokens, f) + logger.info(f"Token stored successfully for user {user_id}") return True except Exception as e: logger.error(f"Error storing token: {str(e)}") @@ -84,6 +85,7 @@ def delete_token(self, user_id, platform, service): with open(self.storage_file, 'w') as f: json.dump(tokens, f) + logger.info(f"Token deleted successfully for user {user_id}") return True except Exception as e: logger.error(f"Error deleting token: {str(e)}") @@ -148,7 +150,9 @@ async def get_authorization_url(self, user_id): } query_string = urlencode(query) - return f"{BOX_AUTH_BASE_URL}authorize?{query_string}" + auth_url = f"{BOX_AUTH_BASE_URL}authorize?{query_string}" + logger.info(f"Generated authorization URL for user {user_id}") + return auth_url async def handle_auth_callback(self, state, code): """ @@ -167,6 +171,7 @@ async def handle_auth_callback(self, state, code): # Decrypt the user_id from state user_id = TokenEncryptionHelper.decrypt_token(state, self.encryption_key) + logger.info(f"Processing authorization callback for user {user_id}") payload = { "grant_type": "authorization_code", @@ -186,8 +191,10 @@ async def handle_auth_callback(self, state, code): response_data["refresh_token"], response_data["expires_in"] ) + logger.info(f"Successfully obtained and stored access token for user {user_id}") else: error_msg = response_data.get("error_description", "Unknown error") + logger.error(f"Failed to obtain access token: {error_msg}") raise Exception(f"Failed to obtain user access token: {error_msg}") async def revoke_access(self, user_id): @@ -217,7 +224,9 @@ async def revoke_access(self, user_id): if response.status_code == 200: # Delete the token from storage self.token_storage.delete_token(user_id, PLATFORM, SERVICE) + logger.info(f"Successfully revoked access for user {user_id}") else: + logger.error(f"Failed to revoke token: {response.status_code}") raise Exception(f"Failed to revoke token: {response.status_code}") async def create_folder(self, user_id, folder_name, parent_folder_id="0"): @@ -234,7 +243,7 @@ async def create_folder(self, user_id, folder_name, parent_folder_id="0"): """ token = await self._load_token(user_id) if not token: - raise Exception("Failed to load a valid token") + raise self._create_auth_exception(user_id) headers = { "Authorization": f"Bearer {token}", @@ -255,8 +264,7 @@ async def create_folder(self, user_id, folder_name, parent_folder_id="0"): if response.status_code in (200, 201): return response.json() else: - error_msg = response.json().get("message", "Unknown error") - raise Exception(f"Box API request failed: {error_msg}") + self._handle_api_error(response, user_id) async def search_for_file(self, user_id, query, limit=100): """ @@ -272,7 +280,7 @@ async def search_for_file(self, user_id, query, limit=100): """ token = await self._load_token(user_id) if not token: - raise Exception("Failed to load a valid token") + raise self._create_auth_exception(user_id) headers = { "Authorization": f"Bearer {token}" @@ -293,8 +301,7 @@ async def search_for_file(self, user_id, query, limit=100): if response.status_code == 200: return response.json() else: - error_msg = response.json().get("message", "Unknown error") - raise Exception(f"Box API request failed: {error_msg}") + self._handle_api_error(response, user_id) async def delete_file(self, user_id, file_id): """ @@ -306,7 +313,7 @@ async def delete_file(self, user_id, file_id): """ token = await self._load_token(user_id) if not token: - raise Exception("Failed to load a valid token") + raise self._create_auth_exception(user_id) headers = { "Authorization": f"Bearer {token}" @@ -318,12 +325,7 @@ async def delete_file(self, user_id, file_id): ) if response.status_code != 204: # 204 No Content is success - error_msg = "Unknown error" - try: - error_msg = response.json().get("message", "Unknown error") - except: - pass - raise Exception(f"Box API request failed: {error_msg}") + self._handle_api_error(response, user_id) async def upload_file(self, user_id, file_path, original_file_name, folder_id="0"): """ @@ -340,7 +342,7 @@ async def upload_file(self, user_id, file_path, original_file_name, folder_id="0 """ token = await self._load_token(user_id) if not token: - raise Exception("Failed to load a valid token") + raise self._create_auth_exception(user_id) headers = { "Authorization": f"Bearer {token}" @@ -366,12 +368,7 @@ async def upload_file(self, user_id, file_path, original_file_name, folder_id="0 if response.status_code in (200, 201): return response.json()['entries'][0] # Box returns an entries array else: - error_msg = "Unknown error" - try: - error_msg = response.json().get("message", "Unknown error") - except: - pass - raise Exception(f"Box API request failed: {error_msg}") + self._handle_api_error(response, user_id) async def get_file_download_link(self, user_id, file_id): """ @@ -386,7 +383,7 @@ async def get_file_download_link(self, user_id, file_id): """ token = await self._load_token(user_id) if not token: - raise Exception("Failed to load a valid token") + raise self._create_auth_exception(user_id) headers = { "Authorization": f"Bearer {token}" @@ -402,12 +399,7 @@ async def get_file_download_link(self, user_id, file_id): if response.status_code == 302: # Redirect status code return response.headers.get('Location') else: - error_msg = "Unknown error" - try: - error_msg = response.json().get("message", "Unknown error") - except: - pass - raise Exception(f"Box API request failed: {error_msg}") + self._handle_api_error(response, user_id) async def get_file_view_link(self, user_id, file_id): """ @@ -422,7 +414,7 @@ async def get_file_view_link(self, user_id, file_id): """ token = await self._load_token(user_id) if not token: - raise Exception("Failed to load a valid token") + raise self._create_auth_exception(user_id) headers = { "Authorization": f"Bearer {token}", @@ -446,8 +438,7 @@ async def get_file_view_link(self, user_id, file_id): else: raise Exception("Shared link URL not found in response") else: - error_msg = response.json().get("message", "Unknown error") - raise Exception(f"Box API request failed: {error_msg}") + self._handle_api_error(response, user_id) async def share_file(self, user_id, file_id, email, role): """ @@ -461,7 +452,7 @@ async def share_file(self, user_id, file_id, email, role): """ token = await self._load_token(user_id) if not token: - raise Exception("Failed to load a valid token") + raise self._create_auth_exception(user_id) headers = { "Authorization": f"Bearer {token}", @@ -481,8 +472,7 @@ async def share_file(self, user_id, file_id, email, role): ) if response.status_code not in (200, 201): - error_msg = response.json().get("message", "Unknown error") - raise Exception(f"Box API request failed: {error_msg}") + self._handle_api_error(response, user_id) async def _store_token(self, user_id, access_token, refresh_token, expires_in): """ @@ -527,7 +517,7 @@ async def _load_token(self, user_id): token_record = self.token_storage.get_token(user_id, PLATFORM, SERVICE) if not token_record or not token_record.get("is_active") or token_record.get("is_revoked"): - logger.info("No valid token found in the storage") + logger.info(f"No valid token found in the storage for user {user_id}") return None try: @@ -545,8 +535,15 @@ async def _load_token(self, user_id): # Check if token is expired expires_at = token_data.get("expires_at") if expires_at and expires_at <= datetime.utcnow().timestamp(): - logger.info("Token expired, attempting to refresh") - return await self._refresh_token(user_id, token_data.get("refresh_token")) + logger.info(f"Token expired for user {user_id}, attempting to refresh") + refresh_token = token_data.get("refresh_token") + if refresh_token: + try: + return await self._refresh_token(user_id, refresh_token) + except Exception as e: + logger.error(f"Error refreshing token: {str(e)}") + return None + return None return token_data.get("access_token") except Exception as e: @@ -576,6 +573,7 @@ async def _refresh_token(self, user_id, refresh_token): "client_secret": self.client_secret } + logger.info(f"Attempting to refresh token for user {user_id}") response = requests.post(f"{BOX_AUTH_BASE_URL}token", data=payload) response_data = response.json() @@ -586,7 +584,62 @@ async def _refresh_token(self, user_id, refresh_token): response_data["refresh_token"], response_data["expires_in"] ) + logger.info(f"Successfully refreshed token for user {user_id}") return response_data["access_token"] else: error_msg = response_data.get("error_description", "Unknown error") - raise Exception(f"Failed to refresh token: {error_msg}") \ No newline at end of file + logger.error(f"Failed to refresh token: {error_msg}") + # If refresh fails, mark the token as revoked so we don't keep trying + token_record = self.token_storage.get_token(user_id, PLATFORM, SERVICE) + if token_record: + token_record["is_revoked"] = True + self.token_storage.store_token(user_id, PLATFORM, SERVICE, token_record) + raise Exception(f"Failed to refresh token: {error_msg}") + + def _handle_api_error(self, response, user_id): + """ + Handle API errors and check for authentication issues. + + Args: + response: The response object + user_id: The user's ID + + Raises: + Exception: With appropriate error message + """ + try: + error_data = response.json() + error_msg = error_data.get("message", "Unknown error") + + # Check if this is an authentication error + if response.status_code in (401, 403): + # Mark token as revoked + token_record = self.token_storage.get_token(user_id, PLATFORM, SERVICE) + if token_record: + token_record["is_revoked"] = True + self.token_storage.store_token(user_id, PLATFORM, SERVICE, token_record) + + # Raise authentication exception + raise self._create_auth_exception(user_id) + + # For other errors + raise Exception(f"Box API request failed: {error_msg}") + except ValueError: + # Response couldn't be parsed as JSON + raise Exception(f"Box API request failed with status code: {response.status_code}") + + def _create_auth_exception(self, user_id): + """ + Create an authentication exception with reauthorization instructions. + + Args: + user_id: The user's ID + + Returns: + Exception: With reauthorization instructions + """ + # Don't try to generate an auth URL here, just return the instruction + return Exception( + "Your Box authorization has expired or is invalid. " + "Please use the `!authorize-box` command to reconnect your Box account." + ) \ No newline at end of file From 18507ec7f9cb29774bb8ddf00026bb03a61fb4e1 Mon Sep 17 00:00:00 2001 From: bless-stanford <107686196+bless-stanford@users.noreply.github.com> Date: Mon, 10 Mar 2025 00:15:11 -0700 Subject: [PATCH 09/19] Refactor Mistral agent to support multiple cloud services, enhance response formatting, and improve authorization handling --- agent.py | 134 +++++++++++++++++++++++++++++++++++++++++++++---------- 1 file changed, 110 insertions(+), 24 deletions(-) diff --git a/agent.py b/agent.py index 5fbea6c2..d3706b07 100644 --- a/agent.py +++ b/agent.py @@ -12,14 +12,18 @@ logger = logging.getLogger("agent") MISTRAL_MODEL = "mistral-large-latest" -SYSTEM_PROMPT = """You are a helpful assistant named Dodobot that can access and manage a user's Box cloud storage. -You can search for files, create folders, get download links, view links, share files, and delete files. -You can process natural language requests about Box files. +SYSTEM_PROMPT = """You are a helpful assistant named Dodobot that can access and manage various cloud services. +You can interact with services like Box, Dropbox, Gmail, and others to search for files, create folders, get download links, etc. -When a user asks about Box files, folders, or cloud storage, use the appropriate Box function to handle their request. +When a user asks about files, folders, or cloud storage, use the appropriate function to handle their request. Never ask the user for their user ID, as it is automatically provided by the system. -If a user's Box account needs to be authorized or reauthorized, tell them to use the !authorize-box command. +For download and view links, always format your response consistently like this: +1. Start with a brief confirmation message (e.g., "Here is the download link for [filename]:") +2. Then provide the actual link on a separate line +3. Do not include raw function call data in your responses + +If a service needs authorization, tell the user to use the !authorize-[service] command (e.g., !authorize-box, !authorize-dropbox). Format your responses using Discord-compatible Markdown: - Use **bold** for emphasis @@ -90,10 +94,8 @@ async def run(self, message: discord.Message): # Trim history before getting response self._trim_chat_history() - # Add the user ID to the kernel arguments for Box plugin access + # Add the user ID to the kernel arguments for plugin access kernel_arguments = KernelArguments() - # Store the user ID with the exact key name expected by the plugin functions - user_id = str(message.author.id) kernel_arguments["user_id"] = user_id # Log the user ID to verify it's correct @@ -111,35 +113,119 @@ async def run(self, message: discord.Message): arguments=kernel_arguments ) - # Check if response content contains Box auth error - if "Your Box authorization has expired" in response.content or \ - "Please use the `!authorize-box` command" in response.content: + # Handle raw function call responses + if response.content.startswith('[{"name":') and '"arguments":' in response.content: + try: + # Try to parse it and format it nicely + function_data = json.loads(response.content) + if isinstance(function_data, list) and len(function_data) > 0: + func_call = function_data[0] + func_name = func_call.get("name", "") + args = func_call.get("arguments", {}) + query = args.get("query", "") + + # Generic handling for various services and functions + service_name = func_name.split('-')[0] if '-' in func_name else "" + action_type = func_name.split('-')[1] if '-' in func_name else func_name + + if "get_file_download_link" in action_type or "download" in action_type: + formatted_response = f"I'll retrieve the download link for '{query}' from {service_name}..." + elif "search" in action_type: + formatted_response = f"I'm searching for '{query}' in your {service_name} account..." + elif "share" in action_type: + formatted_response = f"I'll prepare to share '{query}' from your {service_name} account..." + elif "create" in action_type: + formatted_response = f"I'll create '{query}' in your {service_name} account..." + elif "delete" in action_type: + formatted_response = f"I'll prepare to delete '{query}' from your {service_name} account..." + else: + formatted_response = f"I'm processing your {service_name} request..." + + response.content = formatted_response + except Exception as parse_error: + logger.error(f"Error parsing function call: {str(parse_error)}") + response.content = "I'm processing your request..." + + # Check for authorization errors across different services + auth_error_phrases = [ + "authorization has expired", + "needs to be authorized", + "Please use the `!authorize", + "not authorized", + "authorization required" + ] + + if any(phrase in response.content for phrase in auth_error_phrases): + # Extract the service name if available + service_match = re.search(r'!authorize-(\w+)', response.content) + service_name = service_match.group(1) if service_match else "service" + error_message = ( - "It looks like I need access to your Box account to perform this task. " - "Please use the `!authorize-box` command to connect your Box account." + f"I need access to your {service_name} account to perform this task. " + f"Please use the `!authorize-{service_name}` command to connect your account." ) self.chat_history.add_assistant_message(error_message) return [error_message] - + # Add the assistant's response to the chat history self.chat_history.add_assistant_message(response.content) - # Log the response for debugging (exclude sensitive data) - logger.info(f"Generated response for user {message.author.id} (length: {len(response.content)})") + # Format links consistently for better UI + formatted_content = response.content - # Always return a list of chunks, even for short messages - if len(response.content) > self.MAX_LENGTH: - return self.split_response(response.content) - return [response.content] + # Generic link pattern that should match most download links + link_pattern = r'(https?://\S+)' + + if re.search(link_pattern, formatted_content): + file_name = None + + # Try to extract filename from the context + filename_patterns = [ + r"file ['\"]([^'\"]+)['\"]", + r"file (\S+\.\w+)", + r"download link for ['\"]?([^'\"]+)['\"]?", + r"link for ['\"]?([^'\"]+)['\"]?" + ] + + for pattern in filename_patterns: + match = re.search(pattern, formatted_content, re.IGNORECASE) + if match: + file_name = match.group(1) + break + + # If we found a filename and it looks like a proper filename with extension + if file_name and '.' in file_name: + links = re.findall(link_pattern, formatted_content) + for link in links: + # Skip links that appear to be part of instructions or formatting + if "!authorize" in link or "example" in link.lower(): + continue + + # Format a download button + if not formatted_content.startswith("Download"): + # Only insert a clear Download line if we don't already have one + if "download" not in formatted_content.lower()[:50]: + formatted_content = f"Download {file_name}\n\n{formatted_content}" + + # Log the response for debugging + logger.info(f"Generated response for user {message.author.id} (length: {len(formatted_content)})") + + # Always return a list of chunks + if len(formatted_content) > self.MAX_LENGTH: + return self.split_response(formatted_content) + return [formatted_content] except Exception as e: logger.error(f"Error processing request: {str(e)}", exc_info=True) - # Check if this is an authentication error - if "Box authorization has expired" in str(e) or "!authorize-box" in str(e): + # Check for authentication errors in the exception + if any(phrase in str(e) for phrase in ["authorization", "authorize", "authenticate"]): + service_match = re.search(r'!authorize-(\w+)', str(e)) + service_name = service_match.group(1) if service_match else "service" + error_message = ( - "I need to connect to your Box account to perform this task. " - "Please use the `!authorize-box` command to authorize access." + f"I need to connect to your {service_name} account to perform this task. " + f"Please use the `!authorize-{service_name}` command to authorize access." ) self.chat_history.add_assistant_message(error_message) return [error_message] From 953ffb9352cde2d779ed6678fb25ae135dfca74e Mon Sep 17 00:00:00 2001 From: bless-stanford <107686196+bless-stanford@users.noreply.github.com> Date: Mon, 10 Mar 2025 13:54:59 -0700 Subject: [PATCH 10/19] Remove Box model classes and implement CloudPluginManager for Box and Dropbox integration --- agent.py | 67 +++- bot.py | 94 ++++- helpers/token_helpers.py | 191 +++++++++ plugins/cloud_plugin_manager.py | 105 +++++ plugins/dropbox_plugin.py | 552 ++++++++++++++++++++++++++ server.py | 172 +++++--- services/box_models.py | 108 ----- services/box_service.py | 102 +---- services/dropbox_service.py | 676 ++++++++++++++++++++++++++++++++ 9 files changed, 1790 insertions(+), 277 deletions(-) create mode 100644 helpers/token_helpers.py create mode 100644 plugins/cloud_plugin_manager.py create mode 100644 plugins/dropbox_plugin.py delete mode 100644 services/box_models.py create mode 100644 services/dropbox_service.py diff --git a/agent.py b/agent.py index d3706b07..15850274 100644 --- a/agent.py +++ b/agent.py @@ -2,13 +2,16 @@ import re from semantic_kernel.contents import ChatHistory from kernel.kernel_builder import KernelBuilder -from services.box_service import BoxService -from plugins.box_plugin import BoxPlugins from semantic_kernel.functions import KernelArguments from semantic_kernel.connectors.ai.function_choice_behavior import FunctionChoiceBehavior import logging import json +from services.box_service import BoxService +from services.dropbox_service import DropboxService + +from plugins.cloud_plugin_manager import CloudPluginManager + logger = logging.getLogger("agent") MISTRAL_MODEL = "mistral-large-latest" @@ -17,6 +20,7 @@ When a user asks about files, folders, or cloud storage, use the appropriate function to handle their request. Never ask the user for their user ID, as it is automatically provided by the system. +Do not expose implementation, internal values or functions to the user For download and view links, always format your response consistently like this: 1. Start with a brief confirmation message (e.g., "Here is the download link for [filename]:") @@ -35,6 +39,20 @@ - Use > for quotes - Use ||text|| for spoilers +When deciding which cloud service to use (Box or Dropbox), base your decision on: +1. If the user specifically mentions a service by name, use that service +2. If the user doesn't specify, use the service that appears to be more appropriate for their needs or the one they've used most recently + +For Box: +- Use Box for enterprise-focused needs +- Box uses folder IDs and file IDs for operations +- File operations focus on sharing with specific permissions + +For Dropbox: +- Use Dropbox for personal storage needs +- Dropbox uses file paths for operations +- File operations focus on temporary links and direct access + Do not use # for headers or * - for bullet points as these don't render in Discord. Keep responses concise when possible, as Discord has a 2000-character limit per message.""" @@ -52,10 +70,22 @@ def __init__(self, max_context_messages=10): self.MAX_LENGTH = 1900 # Leave room for extra characters self.max_context_messages = max_context_messages - # Register Box plugins with the kernel + # Initialize cloud services self.box_service = BoxService() - self.box_plugins = BoxPlugins(self.box_service) - self.kernel.add_plugin(self.box_plugins, "Box") + self.dropbox_service = DropboxService() + + # Initialize plugin manager and register plugins + self.cloud_plugin_manager = CloudPluginManager( + box_service=self.box_service, + dropbox_service=self.dropbox_service + ) + + # Register all cloud plugins with the kernel + self.cloud_plugin_manager.register_plugins(self.kernel) + + # Add plugin descriptions to chat history to help the model understand available functions + plugin_descriptions = self.cloud_plugin_manager.get_plugin_descriptions() + self.chat_history.add_system_message(plugin_descriptions) def _trim_chat_history(self): """Keep only the most recent messages within the context window.""" @@ -98,6 +128,9 @@ async def run(self, message: discord.Message): kernel_arguments = KernelArguments() kernel_arguments["user_id"] = user_id + # Update user context in cloud plugin manager + self.cloud_plugin_manager.update_user_context(self.kernel, user_id) + # Log the user ID to verify it's correct logger.info(f"Setting user_id in kernel arguments to: {user_id}") @@ -124,20 +157,28 @@ async def run(self, message: discord.Message): args = func_call.get("arguments", {}) query = args.get("query", "") - # Generic handling for various services and functions - service_name = func_name.split('-')[0] if '-' in func_name else "" - action_type = func_name.split('-')[1] if '-' in func_name else func_name + # Determine which service is being used + if "box" in func_name.lower(): + service_name = "Box" + elif "dropbox" in func_name.lower(): + service_name = "Dropbox" + else: + service_name = "cloud storage" - if "get_file_download_link" in action_type or "download" in action_type: + # Determine action type + if "get_file_download_link" in func_name or "download" in func_name: formatted_response = f"I'll retrieve the download link for '{query}' from {service_name}..." - elif "search" in action_type: + elif "search" in func_name: formatted_response = f"I'm searching for '{query}' in your {service_name} account..." - elif "share" in action_type: + elif "share" in func_name: formatted_response = f"I'll prepare to share '{query}' from your {service_name} account..." - elif "create" in action_type: + elif "create" in func_name: formatted_response = f"I'll create '{query}' in your {service_name} account..." - elif "delete" in action_type: + elif "delete" in func_name: formatted_response = f"I'll prepare to delete '{query}' from your {service_name} account..." + elif "list" in func_name: + path = args.get("path", "root folder") + formatted_response = f"I'll list the contents of '{path}' in your {service_name} account..." else: formatted_response = f"I'm processing your {service_name} request..." diff --git a/bot.py b/bot.py index cc983bc7..5b08055d 100644 --- a/bot.py +++ b/bot.py @@ -5,6 +5,7 @@ from dotenv import load_dotenv from agent import MistralAgent from services.box_service import BoxService +from services.dropbox_service import DropboxService from server import start_server PREFIX = "!" @@ -22,12 +23,16 @@ intents = discord.Intents.all() bot = commands.Bot(command_prefix=PREFIX, intents=intents) -# Initialize agent with Semantic Kernel and Box plugins +# Initialize agent with Semantic Kernel and cloud storage plugins agent = MistralAgent() # Get the token from the environment variables token = os.getenv("DISCORD_TOKEN") +# Initialize cloud service instances +box_service = BoxService() +dropbox_service = DropboxService() + async def send_split_message(message: discord.Message, response: str | list[str]): """ Sends a message that might be longer than Discord's character limit. @@ -81,7 +86,7 @@ async def on_message(message: discord.Message): try: async with message.channel.typing(): # The agent will now use Semantic Kernel to process natural language - # requests related to Box, without requiring specific commands + # requests related to cloud storage, without requiring specific commands response = await agent.run(message) await send_split_message(message, response) except Exception as e: @@ -102,9 +107,6 @@ async def authorize_box(ctx): Sends a Box authorization link to the user via DM. """ try: - # Create Box service - box_service = BoxService() - # Get authorization URL for the user auth_url = await box_service.get_authorization_url(str(ctx.author.id)) @@ -112,7 +114,7 @@ async def authorize_box(ctx): await ctx.author.send(f"Please authorize access to your Box account by clicking this link: {auth_url}") await ctx.send("I've sent you a DM with the authorization link!") except Exception as e: - error_msg = f"Error generating authorization link: {str(e)}" + error_msg = f"Error generating Box authorization link: {str(e)}" logger.error(error_msg) await ctx.send(error_msg[:1900]) @@ -155,6 +157,86 @@ async def box_upload(ctx): if os.path.exists(file_path): os.remove(file_path) +@bot.command(name="authorize-dropbox", help="Authorize the bot to access your Dropbox account") +async def authorize_dropbox(ctx): + """ + Sends a Dropbox authorization link to the user via DM. + """ + try: + # Get authorization URL for the user + auth_url = await dropbox_service.get_authorization_url(str(ctx.author.id)) + + # Send the URL as a DM to the user + await ctx.author.send(f"Please authorize access to your Dropbox account by clicking this link: {auth_url}") + await ctx.send("I've sent you a DM with the authorization link!") + except Exception as e: + error_msg = f"Error generating Dropbox authorization link: {str(e)}" + logger.error(error_msg) + await ctx.send(error_msg[:1900]) + +@bot.command(name="cloud-status", help="Check your cloud service connections") +async def cloud_status(ctx): + """ + Checks and reports the connection status for configured cloud services. + """ + embed = discord.Embed( + title="Cloud Services Status", + description="Current status of your connected cloud services", + color=discord.Color.blue() + ) + + embed.set_author(name=ctx.author.display_name, icon_url=ctx.author.avatar.url if ctx.author.avatar else None) + embed.set_footer(text="Use !authorize-box or !authorize-dropbox to connect services") + embed.timestamp = discord.utils.utcnow() + + # Check Box connection + try: + # Try to load the token to see if the user is authenticated + box_token = await box_service._load_token(str(ctx.author.id)) + if box_token: + embed.add_field( + name="Box Status", + value="✅ Connected", + inline=False + ) + else: + embed.add_field( + name="Box Status", + value="❌ Not connected\n*Use !authorize-box to connect*", + inline=False + ) + except Exception as e: + embed.add_field( + name="Box Status", + value=f"⚠️ Error checking connection\n```{str(e)}```", + inline=False + ) + + # Check Dropbox connection + try: + # Try to load the token to see if the user is authenticated + dropbox_token = await dropbox_service._load_token(str(ctx.author.id)) + if dropbox_token: + embed.add_field( + name="Dropbox Status", + value="✅ Connected", + inline=False + ) + else: + embed.add_field( + name="Dropbox Status", + value="❌ Not connected\n*Use !authorize-dropbox to connect*", + inline=False + ) + except Exception as e: + embed.add_field( + name="Dropbox Status", + value=f"⚠️ Error checking connection\n```{str(e)}```", + inline=False + ) + + await ctx.send(embed=embed) + # Start the web server in the background server_thread = start_server(bot) diff --git a/helpers/token_helpers.py b/helpers/token_helpers.py new file mode 100644 index 00000000..bb7d6b5c --- /dev/null +++ b/helpers/token_helpers.py @@ -0,0 +1,191 @@ +import os +import json +import logging +from datetime import datetime +from cryptography.fernet import Fernet + +# Setup logging +logger = logging.getLogger("token_helpers") + +class TokenEncryptionHelper: + """Helper class for encrypting and decrypting tokens.""" + + @staticmethod + def encrypt_token(token_str, encryption_key): + """ + Encrypts a token string using Fernet symmetric encryption. + + Args: + token_str (str): The token string to encrypt + encryption_key (bytes): The encryption key + + Returns: + str: The encrypted token as a string + """ + f = Fernet(encryption_key) + return f.encrypt(token_str.encode()).decode() + + @staticmethod + def decrypt_token(encrypted_token, encryption_key): + """ + Decrypts an encrypted token string using Fernet symmetric encryption. + + Args: + encrypted_token (str): The encrypted token string + encryption_key (bytes): The encryption key + + Returns: + str: The decrypted token string + """ + f = Fernet(encryption_key) + return f.decrypt(encrypted_token.encode()).decode() + + @staticmethod + def generate_key(): + """ + Generates a new Fernet encryption key. + + Returns: + bytes: A new encryption key + """ + return Fernet.generate_key() + + +class TokenStorageManager: + """A file-based token storage system for managing OAuth tokens.""" + + def __init__(self, storage_file="user_tokens.json"): + """ + Initialize the token storage. + + Args: + storage_file (str): Path to the token storage file + """ + self.storage_file = storage_file + # Initialize the storage file if it doesn't exist + if not os.path.exists(storage_file): + with open(storage_file, 'w') as f: + json.dump({}, f) + + def get_token(self, user_id, platform, service): + """ + Retrieve a token from storage. + + Args: + user_id (str): The user's ID + platform (str): The platform name (e.g., "Box", "Dropbox") + service (str): The service name (e.g., "BoxService") + + Returns: + dict: The token record or None if not found + """ + try: + with open(self.storage_file, 'r') as f: + tokens = json.load(f) + + key = f"{user_id}_{platform}_{service}" + return tokens.get(key) + except Exception as e: + logger.error(f"Error retrieving token: {str(e)}") + return None + + def store_token(self, user_id, platform, service, token_data): + """ + Store a token in storage. + + Args: + user_id (str): The user's ID + platform (str): The platform name (e.g., "Box", "Dropbox") + service (str): The service name (e.g., "BoxService") + token_data (dict): The token data to store + + Returns: + bool: True if successful, False otherwise + """ + try: + with open(self.storage_file, 'r') as f: + tokens = json.load(f) + + key = f"{user_id}_{platform}_{service}" + tokens[key] = token_data + + with open(self.storage_file, 'w') as f: + json.dump(tokens, f) + + logger.info(f"Token stored successfully for user {user_id}") + return True + except Exception as e: + logger.error(f"Error storing token: {str(e)}") + return False + + def delete_token(self, user_id, platform, service): + """ + Delete a token from storage. + + Args: + user_id (str): The user's ID + platform (str): The platform name (e.g., "Box", "Dropbox") + service (str): The service name (e.g., "BoxService") + + Returns: + bool: True if successful, False otherwise + """ + try: + with open(self.storage_file, 'r') as f: + tokens = json.load(f) + + key = f"{user_id}_{platform}_{service}" + if key in tokens: + del tokens[key] + + with open(self.storage_file, 'w') as f: + json.dump(tokens, f) + + logger.info(f"Token deleted successfully for user {user_id}") + return True + except Exception as e: + logger.error(f"Error deleting token: {str(e)}") + return False + + +def create_token_record(encrypted_token): + """ + Create a standard token record structure. + + Args: + encrypted_token (str): The encrypted token string + + Returns: + dict: A standardized token record + """ + return { + "encrypted_token": encrypted_token, + "is_active": True, + "is_revoked": False, + "created_at": datetime.utcnow().timestamp() + } + + +def load_or_generate_encryption_key(env_key_name="ENCRYPTION_KEY"): + """ + Load an encryption key from environment variable or generate a new one. + + Args: + env_key_name (str): Name of the environment variable + + Returns: + bytes: The encryption key + """ + import os + from dotenv import load_dotenv + + load_dotenv() + encryption_key = os.getenv(env_key_name) + + if not encryption_key: + # Generate a new key if none exists + encryption_key = Fernet.generate_key().decode() + # Log a warning since we should save this key + logger.warning(f"No encryption key found. Generated new key. Add to .env: {env_key_name}={encryption_key}") + + return encryption_key.encode() if isinstance(encryption_key, str) else encryption_key \ No newline at end of file diff --git a/plugins/cloud_plugin_manager.py b/plugins/cloud_plugin_manager.py new file mode 100644 index 00000000..8ed6e50f --- /dev/null +++ b/plugins/cloud_plugin_manager.py @@ -0,0 +1,105 @@ +from semantic_kernel.kernel import Kernel +from services.box_service import BoxService +from services.dropbox_service import DropboxService +from plugins.box_plugin import BoxPlugins +from plugins.dropbox_plugin import DropboxPlugins +import logging + +logger = logging.getLogger("cloud_plugin_manager") + +class CloudPluginManager: + """ + Manager for cloud storage plugins to use with Semantic Kernel. + Consolidates Box and Dropbox plugins into a single interface. + """ + + def __init__(self, box_service=None, dropbox_service=None): + """ + Initialize the cloud plugin manager with service instances. + If no services are provided, new ones will be created. + + Args: + box_service: BoxService instance or None + dropbox_service: DropboxService instance or None + """ + self.box_service = box_service or BoxService() + self.dropbox_service = dropbox_service or DropboxService() + + # Initialize plugin instances + self.box_plugins = BoxPlugins(self.box_service) + self.dropbox_plugins = DropboxPlugins(self.dropbox_service) + + def register_plugins(self, kernel: Kernel) -> Kernel: + """ + Register all cloud storage plugins with the given kernel. + + Args: + kernel: The Semantic Kernel instance + + Returns: + Kernel: The same kernel with plugins registered + """ + try: + # Register Box plugins + kernel.add_plugin(self.box_plugins, "box") + logger.info("Box plugins registered with kernel") + + # Register Dropbox plugins + kernel.add_plugin(self.dropbox_plugins, "dropbox") + logger.info("Dropbox plugins registered with kernel") + + return kernel + except Exception as e: + logger.error(f"Error registering cloud plugins: {str(e)}") + raise + + def get_plugin_descriptions(self) -> str: + """ + Get a user-friendly description of all available cloud plugins. + + Returns: + str: Formatted text with plugin descriptions + """ + descriptions = "# Available Cloud Storage Plugins\n\n" + + # Box plugins + descriptions += "## Box Plugins\n" + descriptions += "Use these to interact with your Box account:\n" + descriptions += "- `box.create_folder`: Create a new folder in Box\n" + descriptions += "- `box.search_file`: Search for files in Box\n" + descriptions += "- `box.delete_file`: Delete a file from Box\n" + descriptions += "- `box.get_file_download_link`: Get a download link for a Box file\n" + descriptions += "- `box.get_file_view_link`: Get a shareable view link for a Box file\n" + descriptions += "- `box.share_file`: Share a Box file with another user\n\n" + + # Dropbox plugins + descriptions += "## Dropbox Plugins\n" + descriptions += "Use these to interact with your Dropbox account:\n" + descriptions += "- `dropbox.create_folder`: Create a new folder in Dropbox\n" + descriptions += "- `dropbox.search_file`: Search for files in Dropbox\n" + descriptions += "- `dropbox.list_folder`: List files and folders in a Dropbox path\n" + descriptions += "- `dropbox.delete_file`: Delete a file from Dropbox\n" + descriptions += "- `dropbox.get_file_download_link`: Get a temporary download link for a Dropbox file\n" + descriptions += "- `dropbox.share_file`: Create a shared link for a Dropbox file\n" + + return descriptions + + def update_user_context(self, kernel: Kernel, user_id: str) -> None: + """ + Update the kernel context with user ID for cloud storage operations. + + Args: + kernel: The Semantic Kernel instance + user_id: The user's ID for cloud storage authentication + """ + try: + # Set user_id in variables + if hasattr(kernel, 'data'): + kernel.data["user_id"] = user_id + elif hasattr(kernel, 'variables'): + kernel.variables["user_id"] = user_id + + logger.info(f"Kernel context updated with user ID: {user_id}") + except Exception as e: + logger.error(f"Error updating kernel context: {str(e)}") + raise \ No newline at end of file diff --git a/plugins/dropbox_plugin.py b/plugins/dropbox_plugin.py new file mode 100644 index 00000000..b3e9fdaf --- /dev/null +++ b/plugins/dropbox_plugin.py @@ -0,0 +1,552 @@ +from semantic_kernel.functions import kernel_function +from semantic_kernel.functions.kernel_function_from_prompt import KernelFunctionFromPrompt +from services.dropbox_service import DropboxService +import logging + +logger = logging.getLogger("dropbox_plugins") + +class DropboxPlugins: + """ + Plugins for interacting with Dropbox cloud storage. + """ + + def __init__(self, dropbox_service=None): + """ + Initialize the Dropbox plugins with a DropboxService. + If no service is provided, a new one will be created. + """ + self.dropbox_service = dropbox_service or DropboxService() + + @kernel_function( + name="create_folder", + description="Creates a new folder in the user's Dropbox account" + ) + async def create_folder( + self, + folder_path: str, + user_id: str = None, + kernel = None + ) -> str: + """ + Creates a new folder in the user's Dropbox account. + + Args: + folder_path: Full path of the folder to create (e.g., "/Documents/Projects") + user_id: The user's ID (automatically provided) + + Returns: + str: Success message with folder details or error message + """ + try: + if not user_id: + return "Error: User ID not available. Please try again later." + + # Ensure path starts with / + if not folder_path.startswith('/'): + folder_path = '/' + folder_path + + folder = await self.dropbox_service.create_folder(user_id, folder_path) + + if folder and 'metadata' in folder: + return f"Folder '{folder_path}' created successfully!" + else: + return f"Failed to create folder '{folder_path}'." + + except Exception as e: + logger.error(f"Error creating folder: {str(e)}") + return f"An error occurred while creating the folder: {str(e)}" + + @kernel_function( + name="search_file", + description="Searches for files in the user's Dropbox account and returns details in a user friendly way" + ) + async def search_file( + self, + query: str, + path: str = "", + user_id: str = None, + kernel = None + ) -> str: + """ + Searches for files in the user's Dropbox account. + + Args: + query: Search query or file name + path: Path to search in (default: root folder) + user_id: The user's ID (automatically provided) + + Returns: + str: File details or search results summary + """ + try: + # Get user_id from kernel.data instead of function parameter + if not user_id and kernel and hasattr(kernel, 'arguments'): + user_id = kernel.arguments.get("user_id") + + if not user_id: + return "Error: User ID not available. Please try again later." + + search_results = await self.dropbox_service.search_files(user_id, query, path) + + if not search_results or not search_results.get('matches') or len(search_results['matches']) == 0: + return f"No files found matching '{query}'." + + # In Dropbox API, the metadata is nested under the match object + files = [match['metadata'] for match in search_results['matches']] + + if len(files) == 1: + return self._create_file_detail(files[0]) + + # If multiple files, return a summary + return self._create_search_results_summary(files) + + except Exception as e: + logger.error(f"Error searching files: {str(e)}") + return f"An error occurred while searching for files: {str(e)}" + + @kernel_function( + name="list_folder", + description="Lists files and folders in a specific path in the user's Dropbox" + ) + async def list_folder( + self, + path: str = "", + user_id: str = None, + kernel = None + ) -> str: + """ + Lists files and folders in a specific path in the user's Dropbox. + + Args: + path: Path to list (default: root folder) + user_id: The user's ID (automatically provided) + + Returns: + str: List of files and folders + """ + try: + if not user_id: + return "Error: User ID not available. Please try again later." + + folder_contents = await self.dropbox_service.list_folder(user_id, path) + + if not folder_contents or not folder_contents.get('entries') or len(folder_contents['entries']) == 0: + return f"No files or folders found in path '{path or 'root'}'." + + entries = folder_contents['entries'] + + # Return formatted list of entries + return self._create_folder_listing(entries, path) + + except Exception as e: + logger.error(f"Error listing folder: {str(e)}") + return f"An error occurred while listing the folder contents: {str(e)}" + + @kernel_function( + name="delete_file", + description="Searches and deletes a file from the user's Dropbox account" + ) + async def delete_file( + self, + query: str, + path: str = "", + user_id: str = None, + kernel = None + ) -> str: + """ + Searches and deletes a file from the user's Dropbox account. + + Args: + query: Search query or file name + path: Path to search in (default: root folder) + user_id: The user's ID (automatically provided) + kernel: Semantic Kernel instance for finding most relevant file + + Returns: + str: Success message or error message + """ + try: + if not user_id: + return "Error: User ID not available. Please try again later." + + search_results = await self.dropbox_service.search_files(user_id, query, path) + + if not search_results or not search_results.get('matches') or len(search_results['matches']) == 0: + return f"No files found matching '{query}'." + + # In Dropbox API, the metadata is nested under the match object + files = [match['metadata'] for match in search_results['matches']] + + if len(files) == 1: + file = files[0] + file_path = file.get('path_display', file.get('path_lower', 'unknown_path')) + await self.dropbox_service.delete_file(user_id, file_path) + return f"File '{file_path}' has been successfully deleted." + + # If multiple files and kernel is provided, find most relevant + if kernel and len(files) > 1: + most_relevant_file = await self._find_most_relevant_file(kernel, files, query) + + if most_relevant_file: + file_path = most_relevant_file.get('path_display', most_relevant_file.get('path_lower', 'unknown_path')) + await self.dropbox_service.delete_file(user_id, file_path) + return f"File '{file_path}' has been successfully deleted." + + # If multiple files and no most relevant found, return summary + return f"Multiple files found matching '{query}'. Please be more specific:\n" + \ + "\n".join([f"- {file.get('path_display', file.get('name', 'Unnamed'))}" for file in files[:5]]) + + except Exception as e: + logger.error(f"Error deleting file: {str(e)}") + return f"An error occurred while deleting the file: {str(e)}" + + @kernel_function( + name="get_file_download_link", + description="Gets a temporary download link for a file in the user's Dropbox account" + ) + async def get_file_download_link( + self, + query: str, + path: str = "", + user_id: str = None, + kernel = None + ) -> str: + """ + Gets a temporary download link for a file in the user's Dropbox account. + + Args: + query: Search query or file name + path: Path to search in (default: root folder) + user_id: The user's ID (automatically provided) + kernel: Semantic Kernel instance for finding most relevant file + + Returns: + str: Download link or error message + """ + try: + if not user_id: + return "Error: User ID not available. Please try again later." + + search_results = await self.dropbox_service.search_files(user_id, query, path) + + if not search_results or not search_results.get('matches') or len(search_results['matches']) == 0: + return f"No files found matching '{query}'." + + # In Dropbox API, the metadata is nested under the match object + files = [] + for match in search_results['matches']: + metadata = match['metadata'] + # Check if we have a double-nested metadata structure + if metadata.get('.tag') == 'metadata' and 'metadata' in metadata: + files.append(metadata['metadata']) + else: + files.append(metadata) + + if len(files) == 1: + file = files[0] + file_path = file.get('path_display', file.get('path_lower', 'unknown_path')) + + # Try using file ID if available, otherwise use path + file_id = file.get('id') + if file_id: + download_link = await self.dropbox_service.get_temporary_link(user_id, file_id) + else: + download_link = await self.dropbox_service.get_temporary_link(user_id, file_path) + + return f"Download link for file '{file_path}':\n{download_link}" + + # If multiple files and kernel is provided, find most relevant + if kernel and len(files) > 1: + most_relevant_file = await self._find_most_relevant_file(kernel, files, query) + + if most_relevant_file: + file_path = most_relevant_file.get('path_display', most_relevant_file.get('path_lower', 'unknown_path')) + + # Try using file ID if available, otherwise use path + file_id = most_relevant_file.get('id') + if file_id: + download_link = await self.dropbox_service.get_temporary_link(user_id, file_id) + else: + download_link = await self.dropbox_service.get_temporary_link(user_id, file_path) + + return f"Download link for file '{file_path}':\n{download_link}" + + # If multiple files and no most relevant found, return summary + return f"Multiple files found matching '{query}'. Please be more specific:\n" + \ + "\n".join([f"- {file.get('path_display', file.get('name', 'Unnamed'))}" for file in files[:5]]) + + except Exception as e: + logger.error(f"Error getting download link: {str(e)}") + return f"An error occurred while getting the download link: {str(e)}" + + @kernel_function( + name="share_file", + description="Creates a shared link for a file in the user's Dropbox account" + ) + async def share_file( + self, + query: str, + path: str = "", + user_id: str = None, + kernel = None + ) -> str: + """ + Creates a shared link for a file in the user's Dropbox account. + + Args: + query: Search query or file name + path: Path to search in (default: root folder) + user_id: The user's ID (automatically provided) + kernel: Semantic Kernel instance for finding most relevant file + + Returns: + str: Shared link or error message + """ + try: + if not user_id: + return "Error: User ID not available. Please try again later." + + search_results = await self.dropbox_service.search_files(user_id, query, path) + + if not search_results or not search_results.get('matches') or len(search_results['matches']) == 0: + return f"No files found matching '{query}'." + + # In Dropbox API, the metadata is nested under the match object + files = [match['metadata'] for match in search_results['matches']] + + if len(files) == 1: + file = files[0] + file_path = file.get('path_display', file.get('path_lower', 'unknown_path')) + sharing_info = await self.dropbox_service.share_file(user_id, file_path) + + # Extract the URL from sharing info + shared_url = self._extract_shared_url(sharing_info) + if shared_url: + return f"Shared link for file '{file_path}':\n{shared_url}" + else: + return f"File was shared but couldn't retrieve the URL." + + # If multiple files and kernel is provided, find most relevant + if kernel and len(files) > 1: + most_relevant_file = await self._find_most_relevant_file(kernel, files, query) + + if most_relevant_file: + file_path = most_relevant_file.get('path_display', most_relevant_file.get('path_lower', 'unknown_path')) + sharing_info = await self.dropbox_service.share_file(user_id, file_path) + + # Extract the URL from sharing info + shared_url = self._extract_shared_url(sharing_info) + if shared_url: + return f"Shared link for file '{file_path}':\n{shared_url}" + else: + return f"File was shared but couldn't retrieve the URL." + + # If multiple files and no most relevant found, return summary + return f"Multiple files found matching '{query}'. Please be more specific:\n" + \ + "\n".join([f"- {file.get('path_display', file.get('name', 'Unnamed'))}" for file in files[:5]]) + + except Exception as e: + logger.error(f"Error sharing file: {str(e)}") + return f"An error occurred while sharing the file: {str(e)}" + + def _extract_shared_url(self, sharing_info): + """Extract shared URL from Dropbox sharing info response.""" + if not sharing_info: + return None + + # Handle different response formats + # First, try the direct response from create_shared_link_with_settings + if 'url' in sharing_info: + return sharing_info['url'] + + # Then, try the list_shared_links response format + if 'links' in sharing_info and sharing_info['links']: + for link in sharing_info['links']: + if 'url' in link: + return link['url'] + + return None + + async def _find_most_relevant_file(self, kernel, files, user_query): + """ + Find the most relevant file from a list based on user query. + + Args: + kernel: Semantic Kernel instance + files: List of files + user_query: The user's query + + Returns: + dict: The most relevant file or None + """ + try: + # Create a function from prompt + rank_files_function = KernelFunctionFromPrompt( + function_name="RankFilesByRelevance", + plugin_name=None, + prompt="Given the user query: '{{$userQuery}}' and a list of file paths, " + "rank them by relevance and return the index of the most relevant file. " + "Do not add any comments or explanation to the response.\n" + "File list: {{$fileList}}", + template_format="semantic-kernel" + ) + + # Create file list string with paths that are more relevant for Dropbox + file_list = "\n".join([ + f"{i}: Path: {file.get('path_display', file.get('path_lower', file.get('name', 'Unnamed')))}" + for i, file in enumerate(files) + ]) + + # Create kernel arguments + kernel_arguments = { + "userQuery": user_query, + "fileList": file_list + } + + # Invoke the function + result = await kernel.invoke(rank_files_function, **kernel_arguments) + + # Get the value from the result - might be a list or a string + result_value = result.value + + # Handle different result types + if isinstance(result_value, list) and len(result_value) > 0: + result_text = str(result_value[0]).strip() + elif isinstance(result_value, str): + result_text = result_value.strip() + else: + # Fallback + result_text = str(result_value).strip() + + try: + most_relevant_index = int(result_text) + if 0 <= most_relevant_index < len(files): + return files[most_relevant_index] + except ValueError: + logger.warning(f"Could not parse the relevance index from AI result: {result_text}") + + return None + + except Exception as e: + logger.error(f"Error finding most relevant file: {str(e)}") + return None + + def _create_file_detail(self, file): + """Create a detailed text representation of a file.""" + detail = "**File Details:**\n" + + # Use path_display as primary identifier + path = file.get('path_display', file.get('path_lower', 'Unknown path')) + name = file.get('name', path.split('/')[-1] if path != 'Unknown path' else 'Unknown') + + detail += f"**Name:** {name}\n" + detail += f"**Path:** {path}\n" + + # Add ID if available + if 'id' in file: + detail += f"**ID:** {file['id']}\n" + + # Format file size if available + if 'size' in file: + size_bytes = file['size'] + size_str = self._format_file_size(size_bytes) + detail += f"**Size:** {size_str}\n" + + # Add dates if available + if 'server_modified' in file: + detail += f"**Modified At:** {file['server_modified']}\n" + if 'client_modified' in file: + detail += f"**Client Modified At:** {file['client_modified']}\n" + + # Add file type if available + if '.tag' in file: + detail += f"**Type:** {file['.tag']}\n" + + # Add content hash if available (for version tracking) + if 'content_hash' in file: + detail += f"**Content Hash:** {file['content_hash'][:10]}...\n" + + return detail + + def _create_search_results_summary(self, files): + """Create a summary of multiple search results.""" + summary = "**Multiple files found. Here are the details:**\n\n" + + for i, file in enumerate(files[:5], 1): # Limit to 5 files and number them + path = file.get('path_display', file.get('path_lower', 'Unknown path')) + name = file.get('name', path.split('/')[-1] if path != 'Unknown path' else 'Unknown') + + summary += f"**{i}. {name}**\n" + summary += f" Path: {path}\n" + + # Add size if available + if 'size' in file: + size_str = self._format_file_size(file['size']) + summary += f" Size: {size_str}\n" + + # Add type if available + if '.tag' in file: + summary += f" Type: {file['.tag']}\n" + + summary += "\n" + + if len(files) > 5: + summary += f"\n...and {len(files) - 5} more files.\n" + + summary += "\nPlease provide a more specific query to find the exact file you want." + + return summary + + def _create_folder_listing(self, entries, path): + """Create a formatted listing of folder contents.""" + path_display = path or "root folder" + listing = f"**Contents of {path_display}:**\n\n" + + # Separate folders and files + folders = [entry for entry in entries if entry.get('.tag') == 'folder'] + files = [entry for entry in entries if entry.get('.tag') == 'file'] + + # Sort by name + folders.sort(key=lambda x: x.get('name', '').lower()) + files.sort(key=lambda x: x.get('name', '').lower()) + + # Add folders first + if folders: + listing += "**Folders:**\n" + for folder in folders: + name = folder.get('name', 'Unnamed folder') + path = folder.get('path_display', folder.get('path_lower', 'Unknown path')) + listing += f"📁 {name} (Path: {path})\n" + listing += "\n" + + # Then add files + if files: + listing += "**Files:**\n" + for file in files: + name = file.get('name', 'Unnamed file') + path = file.get('path_display', file.get('path_lower', 'Unknown path')) + + # Add size if available + size_info = "" + if 'size' in file: + size_str = self._format_file_size(file['size']) + size_info = f", Size: {size_str}" + + listing += f"📄 {name} (Path: {path}{size_info})\n" + listing += "\n" + + if not folders and not files: + listing += "This folder is empty." + + return listing + + def _format_file_size(self, bytes): + """Format file size in human-readable form.""" + sizes = ["B", "KB", "MB", "GB", "TB"] + order = 0 + size = float(bytes) + while size >= 1024 and order < len(sizes) - 1: + order += 1 + size /= 1024 + + return f"{size:.2f} {sizes[order]}" \ No newline at end of file diff --git a/server.py b/server.py index c1a4cb39..87fd62ea 100644 --- a/server.py +++ b/server.py @@ -1,23 +1,85 @@ -from fastapi import FastAPI, Request +from fastapi import FastAPI from fastapi.responses import HTMLResponse import uvicorn -from services.box_service import BoxService, TokenEncryptionHelper +from services.box_service import BoxService +from services.dropbox_service import DropboxService +from helpers.token_helpers import TokenEncryptionHelper import asyncio import logging import threading # Setup logging -logger = logging.getLogger("box_server") +logger = logging.getLogger("oauth_server") app = FastAPI() box_service = BoxService() +dropbox_service = DropboxService() # This will be set from bot.py bot = None +# Reusable HTML template function +def get_success_html(service_name): + """ + Generate HTML for successful authorization. + + Args: + service_name: The name of the service (Box, Dropbox, etc.) + + Returns: + str: HTML content for the success page + """ + return f""" + + + + Authorization Successful + + + +
+
✅ Authorization Successful!
+
Your {service_name} account has been connected to the Discord bot.
+
You can close this window and return to Discord.
+
You will also receive a confirmation message in Discord.
+
+ + + """ + @app.get("/") async def root(): - return {"message": "Box OAuth Callback Server"} + return {"message": "OAuth Callback Server for Box and Dropbox"} @app.get("/box/callback") async def box_callback(code: str, state: str): @@ -31,7 +93,7 @@ async def box_callback(code: str, state: str): try: # Get user ID from state user_id = TokenEncryptionHelper.decrypt_token(state, box_service.encryption_key) - logger.info(f"Received callback for user {user_id}") + logger.info(f"Received Box callback for user {user_id}") # Handle the callback - this stores the tokens await box_service.handle_auth_callback(state, code) @@ -39,70 +101,60 @@ async def box_callback(code: str, state: str): # Notify the user through Discord if bot: # Schedule the notification in the bot's event loop - asyncio.run_coroutine_threadsafe(notify_user(user_id), bot.loop) + asyncio.run_coroutine_threadsafe(notify_user(user_id, "Box"), bot.loop) - # Return a nice HTML page with correct Content-Type - html_content = """ - - - - Authorization Successful - - - -
-
✅ Authorization Successful!
-
Your Box account has been connected to the Discord bot.
-
You can close this window and return to Discord.
-
You will also receive a confirmation message in Discord.
-
- - - """ + # Use the reusable HTML template + html_content = get_success_html("Box") return HTMLResponse(content=html_content) except Exception as e: - logger.error(f"Error in callback: {str(e)}") + logger.error(f"Error in Box callback: {str(e)}") return {"error": str(e)} -async def notify_user(user_id): - """Send a Discord message to notify the user that authorization was successful.""" +@app.get("/dropbox/callback") +async def dropbox_callback(code: str, state: str): + """ + Handle the OAuth callback from Dropbox. + + This endpoint receives the authorization code from Dropbox after a user + authorizes the application. It exchanges the code for access and + refresh tokens, stores them securely, and notifies the user. + """ + try: + # Get user ID from state + user_id = TokenEncryptionHelper.decrypt_token(state, dropbox_service.encryption_key) + logger.info(f"Received Dropbox callback for user {user_id}") + + # Handle the callback - this stores the tokens + await dropbox_service.handle_auth_callback(state, code) + + # Notify the user through Discord + if bot: + # Schedule the notification in the bot's event loop + asyncio.run_coroutine_threadsafe(notify_user(user_id, "Dropbox"), bot.loop) + + # Use the reusable HTML template + html_content = get_success_html("Dropbox") + + return HTMLResponse(content=html_content) + except Exception as e: + logger.error(f"Error in Dropbox callback: {str(e)}") + return {"error": str(e)} + +async def notify_user(user_id, service_name): + """ + Send a Discord message to notify the user that authorization was successful. + + Args: + user_id: The Discord user ID + service_name: The name of the service (Box, Dropbox, etc.) + """ try: user = await bot.fetch_user(int(user_id)) if user: - await user.send("✅ Your Box account has been successfully connected! You can now use Box commands.") + await user.send(f"✅ Your {service_name} account has been successfully connected! You can now use {service_name} commands.") except Exception as e: - logger.error(f"Error notifying user: {str(e)}") + logger.error(f"Error notifying user about {service_name}: {str(e)}") def start_server(bot_instance=None): """ @@ -119,7 +171,7 @@ def start_server(bot_instance=None): # Start the server in a separate thread def run_server(): - logger.info("Starting Box OAuth callback server on port 8000") + logger.info("Starting OAuth callback server on port 8000") uvicorn.run(app, host="0.0.0.0", port=8000, log_level="info") server_thread = threading.Thread(target=run_server, daemon=True) diff --git a/services/box_models.py b/services/box_models.py deleted file mode 100644 index ac40ea3a..00000000 --- a/services/box_models.py +++ /dev/null @@ -1,108 +0,0 @@ -class BoxFolder: - """ - Model representing a Box folder. - """ - def __init__(self, - id=None, - name=None, - parent=None, - created_at=None, - modified_at=None, - item_collection=None, - **kwargs): - self.id = id - self.name = name - self.parent = parent - self.created_at = created_at - self.modified_at = modified_at - self.item_collection = item_collection - - # Store additional properties - for key, value in kwargs.items(): - setattr(self, key, value) - - @classmethod - def from_dict(cls, data): - """Create a BoxFolder instance from a dictionary.""" - return cls(**data) - - -class BoxFile: - """ - Model representing a Box file. - """ - def __init__(self, - id=None, - name=None, - parent=None, - created_at=None, - modified_at=None, - size=None, - extension=None, - shared_link=None, - **kwargs): - self.id = id - self.name = name - self.parent = parent - self.created_at = created_at - self.modified_at = modified_at - self.size = size - self.extension = extension - self.shared_link = shared_link - - # Store additional properties - for key, value in kwargs.items(): - setattr(self, key, value) - - @classmethod - def from_dict(cls, data): - """Create a BoxFile instance from a dictionary.""" - return cls(**data) - - -class BoxTokenData: - """ - Model representing Box token data for storage. - """ - def __init__(self, access_token=None, refresh_token=None, expires_at=None): - self.access_token = access_token - self.refresh_token = refresh_token - self.expires_at = expires_at - - def to_dict(self): - """Convert to dictionary for serialization.""" - return { - "access_token": self.access_token, - "refresh_token": self.refresh_token, - "expires_at": self.expires_at - } - - @classmethod - def from_dict(cls, data): - """Create a BoxTokenData instance from a dictionary.""" - return cls( - access_token=data.get("access_token"), - refresh_token=data.get("refresh_token"), - expires_at=data.get("expires_at") - ) - - -class BoxTokenResponse: - """ - Model representing the response from Box token endpoint. - """ - def __init__(self, access_token=None, refresh_token=None, expires_in=None, token_type=None): - self.access_token = access_token - self.refresh_token = refresh_token - self.expires_in = expires_in - self.token_type = token_type - - @classmethod - def from_dict(cls, data): - """Create a BoxTokenResponse instance from a dictionary.""" - return cls( - access_token=data.get("access_token"), - refresh_token=data.get("refresh_token"), - expires_in=data.get("expires_in"), - token_type=data.get("token_type") - ) \ No newline at end of file diff --git a/services/box_service.py b/services/box_service.py index d44938cf..1f183c0d 100644 --- a/services/box_service.py +++ b/services/box_service.py @@ -2,11 +2,17 @@ import json import requests from datetime import datetime, timedelta -from cryptography.fernet import Fernet from dotenv import load_dotenv from urllib.parse import urlencode import logging +from helpers.token_helpers import ( + TokenEncryptionHelper, + TokenStorageManager, + create_token_record, + load_or_generate_encryption_key +) + # Setup logging logger = logging.getLogger("box_service") @@ -19,77 +25,6 @@ BOX_AUTH_BASE_URL = "https://account.box.com/api/oauth2/" BOX_UPLOAD_API_BASE_URL = "https://upload.box.com/api/2.0/" -class TokenEncryptionHelper: - @staticmethod - def encrypt_token(token_str, encryption_key): - """Encrypts a token string using Fernet symmetric encryption.""" - f = Fernet(encryption_key) - return f.encrypt(token_str.encode()).decode() - - @staticmethod - def decrypt_token(encrypted_token, encryption_key): - """Decrypts an encrypted token string using Fernet symmetric encryption.""" - f = Fernet(encryption_key) - return f.decrypt(encrypted_token.encode()).decode() - -class TokenStorageManager: - """A simple file-based token storage system.""" - - def __init__(self, storage_file="user_tokens.json"): - self.storage_file = storage_file - # Initialize the storage file if it doesn't exist - if not os.path.exists(storage_file): - with open(storage_file, 'w') as f: - json.dump({}, f) - - def get_token(self, user_id, platform, service): - """Retrieve a token from storage.""" - try: - with open(self.storage_file, 'r') as f: - tokens = json.load(f) - - key = f"{user_id}_{platform}_{service}" - return tokens.get(key) - except Exception as e: - logger.error(f"Error retrieving token: {str(e)}") - return None - - def store_token(self, user_id, platform, service, token_data): - """Store a token in storage.""" - try: - with open(self.storage_file, 'r') as f: - tokens = json.load(f) - - key = f"{user_id}_{platform}_{service}" - tokens[key] = token_data - - with open(self.storage_file, 'w') as f: - json.dump(tokens, f) - - logger.info(f"Token stored successfully for user {user_id}") - return True - except Exception as e: - logger.error(f"Error storing token: {str(e)}") - return False - - def delete_token(self, user_id, platform, service): - """Delete a token from storage.""" - try: - with open(self.storage_file, 'r') as f: - tokens = json.load(f) - - key = f"{user_id}_{platform}_{service}" - if key in tokens: - del tokens[key] - - with open(self.storage_file, 'w') as f: - json.dump(tokens, f) - - logger.info(f"Token deleted successfully for user {user_id}") - return True - except Exception as e: - logger.error(f"Error deleting token: {str(e)}") - return False class BoxService: def __init__(self, config=None): @@ -105,21 +40,13 @@ def __init__(self, config=None): self.client_secret = os.getenv("BOX_CLIENT_SECRET") self.redirect_uri = os.getenv("BOX_REDIRECT_URI") - # Get or generate encryption key - encryption_key = os.getenv("ENCRYPTION_KEY") - if not encryption_key: - # Generate a new key if none exists - encryption_key = Fernet.generate_key().decode() - # Add this to your .env file manually or update it - logger.warning("No encryption key found. Generated new key. Add to .env: " - f"ENCRYPTION_KEY={encryption_key}") - - self.encryption_key = encryption_key.encode() if isinstance(encryption_key, str) else encryption_key + # Get or generate encryption key using our helper + self.encryption_key = load_or_generate_encryption_key() else: self.client_id = config.get("client_id") self.client_secret = config.get("client_secret") self.redirect_uri = config.get("redirect_uri") - self.encryption_key = config.get("encryption_key", Fernet.generate_key()) + self.encryption_key = config.get("encryption_key") # Initialize token storage self.token_storage = TokenStorageManager() @@ -494,13 +421,8 @@ async def _store_token(self, user_id, access_token, refresh_token, expires_in): serialized_token = json.dumps(token_data) encrypted_token = TokenEncryptionHelper.encrypt_token(serialized_token, self.encryption_key) - # Store in the token storage - token_record = { - "encrypted_token": encrypted_token, - "is_active": True, - "is_revoked": False, - "created_at": datetime.utcnow().timestamp() - } + # Store in the token storage using the helper function + token_record = create_token_record(encrypted_token) self.token_storage.store_token(user_id, PLATFORM, SERVICE, token_record) diff --git a/services/dropbox_service.py b/services/dropbox_service.py new file mode 100644 index 00000000..08ee1a86 --- /dev/null +++ b/services/dropbox_service.py @@ -0,0 +1,676 @@ +import os +import json +import requests +import base64 +from datetime import datetime, timedelta +from dotenv import load_dotenv +from urllib.parse import urlencode +import logging + +from helpers.token_helpers import ( + TokenEncryptionHelper, + TokenStorageManager, + create_token_record, + load_or_generate_encryption_key +) + +# Setup logging +logger = logging.getLogger("dropbox_service") + +# Constants for platform and service +PLATFORM = "Dropbox" +SERVICE = "DropboxService" + +# API URLs +DROPBOX_API_BASE_URL = "https://api.dropboxapi.com/2/" +DROPBOX_CONTENT_API_BASE_URL = "https://content.dropboxapi.com/2/" +DROPBOX_AUTH_BASE_URL = "https://www.dropbox.com/oauth2/" + + +class DropboxService: + def __init__(self, config=None): + """ + Initialize the Dropbox service with configuration. + + Args: + config: Configuration dictionary or None to load from .env + """ + if config is None: + load_dotenv() + self.client_id = os.getenv("DROPBOX_CLIENT_ID") + self.client_secret = os.getenv("DROPBOX_CLIENT_SECRET") + self.redirect_uri = os.getenv("DROPBOX_REDIRECT_URI") + self.app_name = os.getenv("DROPBOX_APP_NAME", "DropboxApp") + + # Get or generate encryption key using our helper + self.encryption_key = load_or_generate_encryption_key() + else: + self.client_id = config.get("client_id") + self.client_secret = config.get("client_secret") + self.redirect_uri = config.get("redirect_uri") + self.app_name = config.get("app_name", "DropboxApp") + self.encryption_key = config.get("encryption_key") + + # Initialize token storage + self.token_storage = TokenStorageManager() + + async def get_authorization_url(self, user_id): + """ + Get the authorization URL for Dropbox OAuth flow. + + Args: + user_id: The user's ID + + Returns: + str: The authorization URL + """ + if not self.client_id: + raise ValueError("Dropbox Client ID is not set in configuration.") + if not self.redirect_uri: + raise ValueError("Dropbox Redirect URI is not set in configuration.") + + # Encrypt user_id as state parameter + state = TokenEncryptionHelper.encrypt_token(user_id, self.encryption_key) + + query = { + "client_id": self.client_id, + "response_type": "code", + "redirect_uri": self.redirect_uri, + "state": state, + "token_access_type": "offline" # Request a refresh token + } + + query_string = urlencode(query) + auth_url = f"{DROPBOX_AUTH_BASE_URL}authorize?{query_string}" + logger.info(f"Generated authorization URL for user {user_id}") + return auth_url + + async def handle_auth_callback(self, state, code): + """ + Handle the authorization callback from Dropbox. + + Args: + state: The state parameter from the callback + code: The authorization code from the callback + """ + if not self.client_id: + raise ValueError("Dropbox Client ID is not set in configuration.") + if not self.client_secret: + raise ValueError("Dropbox Client Secret is not set in configuration.") + if not self.redirect_uri: + raise ValueError("Dropbox Redirect URI is not set in configuration.") + + # Decrypt the user_id from state + user_id = TokenEncryptionHelper.decrypt_token(state, self.encryption_key) + logger.info(f"Processing authorization callback for user {user_id}") + + payload = { + "grant_type": "authorization_code", + "code": code, + "client_id": self.client_id, + "client_secret": self.client_secret, + "redirect_uri": self.redirect_uri + } + + response = requests.post(f"{DROPBOX_AUTH_BASE_URL}token", data=payload) + response_data = response.json() + + if response.status_code == 200 and "access_token" in response_data: + # Calculate expiry time (Dropbox tokens usually last 4 hours by default) + expires_in = response_data.get("expires_in", 14400) # 4 hours in seconds + + await self._store_token( + user_id, + response_data["access_token"], + response_data.get("refresh_token"), # Might be None if scope doesn't include offline access + expires_in + ) + logger.info(f"Successfully obtained and stored access token for user {user_id}") + else: + error_msg = response_data.get("error_description", "Unknown error") + logger.error(f"Failed to obtain access token: {error_msg}") + raise Exception(f"Failed to obtain user access token: {error_msg}") + + async def revoke_access(self, user_id): + """ + Revoke the Dropbox access for a user. + + Args: + user_id: The user's ID + """ + token = await self._load_token(user_id) + if not token: + raise ValueError("No valid token found for user") + + headers = { + "Content-Type": "application/json" + } + + # Dropbox requires token to be in Authorization header + auth_value = f"{self.client_id}:{self.client_secret}" + auth_bytes = auth_value.encode('ascii') + base64_auth = base64.b64encode(auth_bytes).decode('ascii') + headers["Authorization"] = f"Basic {base64_auth}" + + payload = { + "token": token + } + + response = requests.post( + f"{DROPBOX_API_BASE_URL}auth/token/revoke", + headers=headers, + json=payload + ) + + if response.status_code == 200: + # Delete the token from storage + self.token_storage.delete_token(user_id, PLATFORM, SERVICE) + logger.info(f"Successfully revoked access for user {user_id}") + else: + logger.error(f"Failed to revoke token: {response.status_code}") + raise Exception(f"Failed to revoke token: {response.status_code}") + + async def list_folder(self, user_id, path=""): + """ + List contents of a folder in Dropbox. + + Args: + user_id: The user's ID + path: Path to the folder (default: "" for root) + + Returns: + dict: The folder contents + """ + token = await self._load_token(user_id) + if not token: + raise self._create_auth_exception(user_id) + + headers = { + "Authorization": f"Bearer {token}", + "Content-Type": "application/json" + } + + payload = { + "path": path, + "recursive": False, + "include_media_info": False, + "include_deleted": False, + "include_has_explicit_shared_members": False + } + + response = requests.post( + f"{DROPBOX_API_BASE_URL}files/list_folder", + headers=headers, + json=payload + ) + + if response.status_code == 200: + return response.json() + else: + self._handle_api_error(response, user_id) + + async def search_files(self, user_id, query, path="", max_results=10): + """ + Search for files in Dropbox. + + Args: + user_id: The user's ID + query: Search query + path: Path to search in (default: "" for root) + max_results: Maximum number of results to return + + Returns: + dict: Search results + """ + token = await self._load_token(user_id) + if not token: + raise self._create_auth_exception(user_id) + + headers = { + "Authorization": f"Bearer {token}", + "Content-Type": "application/json" + } + + payload = { + "query": query, + "path": path if path else "", + "max_results": max_results, + "mode": { + ".tag": "filename_and_content" + } + } + + response = requests.post( + f"{DROPBOX_API_BASE_URL}files/search_v2", + headers=headers, + json=payload + ) + + if response.status_code == 200: + return response.json() + else: + self._handle_api_error(response, user_id) + + async def create_folder(self, user_id, path): + """ + Create a folder in Dropbox. + + Args: + user_id: The user's ID + path: Path of the folder to create + + Returns: + dict: The created folder information + """ + token = await self._load_token(user_id) + if not token: + raise self._create_auth_exception(user_id) + + headers = { + "Authorization": f"Bearer {token}", + "Content-Type": "application/json" + } + + payload = { + "path": path, + "autorename": False + } + + response = requests.post( + f"{DROPBOX_API_BASE_URL}files/create_folder_v2", + headers=headers, + json=payload + ) + + if response.status_code == 200: + return response.json() + elif response.status_code == 409: + # Folder already exists + logger.info(f"Folder already exists at path: {path}") + return {"metadata": {"path": path, "name": path.split('/')[-1]}} + else: + self._handle_api_error(response, user_id) + + async def upload_file(self, user_id, local_file_path, dropbox_path): + """ + Upload a file to Dropbox. + + Args: + user_id: The user's ID + local_file_path: Path to the local file + dropbox_path: Path where to store the file in Dropbox + + Returns: + dict: The uploaded file information + """ + token = await self._load_token(user_id) + if not token: + raise self._create_auth_exception(user_id) + + headers = { + "Authorization": f"Bearer {token}", + "Dropbox-API-Arg": json.dumps({ + "path": dropbox_path, + "mode": "overwrite", + "autorename": True, + "mute": False + }), + "Content-Type": "application/octet-stream" + } + + with open(local_file_path, "rb") as f: + file_data = f.read() + + response = requests.post( + f"{DROPBOX_CONTENT_API_BASE_URL}files/upload", + headers=headers, + data=file_data + ) + + if response.status_code in (200, 201): + return response.json() + else: + self._handle_api_error(response, user_id) + + async def delete_file(self, user_id, path): + """ + Delete a file from Dropbox. + + Args: + user_id: The user's ID + path: Path of the file to delete + """ + token = await self._load_token(user_id) + if not token: + raise self._create_auth_exception(user_id) + + headers = { + "Authorization": f"Bearer {token}", + "Content-Type": "application/json" + } + + payload = { + "path": path + } + + response = requests.post( + f"{DROPBOX_API_BASE_URL}files/delete_v2", + headers=headers, + json=payload + ) + + if response.status_code != 200: + self._handle_api_error(response, user_id) + + async def get_temporary_link(self, user_id, path): + """ + Get a temporary download link for a file. + + Args: + user_id: The user's ID + path: Path of the file + + Returns: + str: Temporary download URL + """ + token = await self._load_token(user_id) + if not token: + raise self._create_auth_exception(user_id) + + # Check if path starts with an ID + if path.startswith('id:'): + logger.info(f"Using ID format for path: {path}") + else: + # Ensure path starts with a slash + if not path.startswith('/'): + path = '/' + path + logger.info(f"Path was reformatted to include leading slash: {path}") + + headers = { + "Authorization": f"Bearer {token}", + "Content-Type": "application/json" + } + + payload = { + "path": path + } + + response = requests.post( + f"{DROPBOX_API_BASE_URL}files/get_temporary_link", + headers=headers, + json=payload + ) + + if response.status_code == 200: + data = response.json() + if "link" in data: + return data["link"] + else: + logger.error(f"Link not found in response: {json.dumps(data)}") + raise Exception("Link not found in response") + else: + try: + error_data = response.json() + logger.error(f"Error response: {json.dumps(error_data)}") + + # Check for specific error information + if "error" in error_data: + if isinstance(error_data["error"], dict) and ".tag" in error_data["error"]: + error_type = error_data["error"][".tag"] + logger.error(f"Error type: {error_type}") + + # Special handling for common error types + if error_type == "path": + # Extract more details about path errors + path_error = error_data["error"].get("path", {}) + path_error_tag = path_error.get(".tag") if isinstance(path_error, dict) else None + logger.error(f"Path error type: {path_error_tag}") + + except ValueError: + logger.error(f"Response was not valid JSON: {response.text[:200]}") + + self._handle_api_error(response, user_id) + + async def share_file(self, user_id, path, settings=None): + """ + Create a shared link for a file. + + Args: + user_id: The user's ID + path: Path of the file to share + settings: Optional sharing settings + + Returns: + dict: The sharing metadata + """ + token = await self._load_token(user_id) + if not token: + raise self._create_auth_exception(user_id) + + headers = { + "Authorization": f"Bearer {token}", + "Content-Type": "application/json" + } + + payload = { + "path": path, + "settings": settings or {} + } + + response = requests.post( + f"{DROPBOX_API_BASE_URL}sharing/create_shared_link_with_settings", + headers=headers, + json=payload + ) + + if response.status_code == 200: + return response.json() + elif response.status_code == 409 and "shared_link_already_exists" in response.text: + # Link already exists, get existing links + return await self.get_shared_links(user_id, path) + else: + self._handle_api_error(response, user_id) + + async def get_shared_links(self, user_id, path): + """ + Get existing shared links for a file. + + Args: + user_id: The user's ID + path: Path of the file + + Returns: + dict: Existing shared links + """ + token = await self._load_token(user_id) + if not token: + raise self._create_auth_exception(user_id) + + headers = { + "Authorization": f"Bearer {token}", + "Content-Type": "application/json" + } + + payload = { + "path": path + } + + response = requests.post( + f"{DROPBOX_API_BASE_URL}sharing/list_shared_links", + headers=headers, + json=payload + ) + + if response.status_code == 200: + return response.json() + else: + self._handle_api_error(response, user_id) + + async def _store_token(self, user_id, access_token, refresh_token, expires_in): + """ + Store a token in the token storage. + + Args: + user_id: The user's ID + access_token: The access token + refresh_token: The refresh token + expires_in: Expiration time in seconds + """ + token_data = { + "access_token": access_token, + "refresh_token": refresh_token, + "expires_at": (datetime.utcnow() + timedelta(seconds=expires_in)).timestamp() + } + + # Serialize and encrypt the token data + serialized_token = json.dumps(token_data) + encrypted_token = TokenEncryptionHelper.encrypt_token(serialized_token, self.encryption_key) + + # Store in the token storage using the helper function + token_record = create_token_record(encrypted_token) + + self.token_storage.store_token(user_id, PLATFORM, SERVICE, token_record) + + async def _load_token(self, user_id): + """ + Load a token from the token storage. + + Args: + user_id: The user's ID + + Returns: + str: The access token, or None if not found or expired + """ + token_record = self.token_storage.get_token(user_id, PLATFORM, SERVICE) + + if not token_record or not token_record.get("is_active") or token_record.get("is_revoked"): + logger.info(f"No valid token found in the storage for user {user_id}") + return None + + try: + encrypted_token = token_record.get("encrypted_token") + if not encrypted_token: + return None + + decrypted_token = TokenEncryptionHelper.decrypt_token(encrypted_token, self.encryption_key) + token_data = json.loads(decrypted_token) + + if not token_data: + logger.error("Failed to deserialize token data") + return None + + # Check if token is expired + expires_at = token_data.get("expires_at") + if expires_at and expires_at <= datetime.utcnow().timestamp(): + logger.info(f"Token expired for user {user_id}, attempting to refresh") + refresh_token = token_data.get("refresh_token") + if refresh_token: + try: + return await self._refresh_token(user_id, refresh_token) + except Exception as e: + logger.error(f"Error refreshing token: {str(e)}") + return None + return None + + return token_data.get("access_token") + except Exception as e: + logger.error(f"Error loading token: {str(e)}") + return None + + async def _refresh_token(self, user_id, refresh_token): + """ + Refresh an expired token. + + Args: + user_id: The user's ID + refresh_token: The refresh token + + Returns: + str: The new access token + """ + if not self.client_id: + raise ValueError("Dropbox Client ID is not set in configuration.") + if not self.client_secret: + raise ValueError("Dropbox Client Secret is not set in configuration.") + + payload = { + "grant_type": "refresh_token", + "refresh_token": refresh_token, + "client_id": self.client_id, + "client_secret": self.client_secret + } + + logger.info(f"Attempting to refresh token for user {user_id}") + response = requests.post(f"{DROPBOX_AUTH_BASE_URL}token", data=payload) + response_data = response.json() + + if response.status_code == 200 and "access_token" in response_data: + # Note: Dropbox might not return a new refresh token, so keep the old one if none returned + new_refresh_token = response_data.get("refresh_token", refresh_token) + expires_in = response_data.get("expires_in", 14400) # 4 hours in seconds + + await self._store_token( + user_id, + response_data["access_token"], + new_refresh_token, + expires_in + ) + logger.info(f"Successfully refreshed token for user {user_id}") + return response_data["access_token"] + else: + error_msg = response_data.get("error_description", "Unknown error") + logger.error(f"Failed to refresh token: {error_msg}") + # If refresh fails, mark the token as revoked so we don't keep trying + token_record = self.token_storage.get_token(user_id, PLATFORM, SERVICE) + if token_record: + token_record["is_revoked"] = True + self.token_storage.store_token(user_id, PLATFORM, SERVICE, token_record) + raise Exception(f"Failed to refresh token: {error_msg}") + + def _handle_api_error(self, response, user_id): + """ + Handle API errors and check for authentication issues. + + Args: + response: The response object + user_id: The user's ID + + Raises: + Exception: With appropriate error message + """ + try: + error_data = response.json() + error_summary = error_data.get("error_summary", "Unknown error") + + # Check if this is an authentication error + if response.status_code in (401, 403): + # Mark token as revoked + token_record = self.token_storage.get_token(user_id, PLATFORM, SERVICE) + if token_record: + token_record["is_revoked"] = True + self.token_storage.store_token(user_id, PLATFORM, SERVICE, token_record) + + # Raise authentication exception + raise self._create_auth_exception(user_id) + + # For other errors + raise Exception(f"Dropbox API request failed: {error_summary}") + except ValueError: + # Response couldn't be parsed as JSON + raise Exception(f"Dropbox API request failed with status code: {response.status_code}") + + def _create_auth_exception(self, user_id): + """ + Create an authentication exception with reauthorization instructions. + + Args: + user_id: The user's ID + + Returns: + Exception: With reauthorization instructions + """ + # Don't try to generate an auth URL here, just return the instruction + return Exception( + "Your Dropbox authorization has expired or is invalid. " + "Please use the `!authorize-dropbox` command to reconnect your Dropbox account." + ) \ No newline at end of file From 951119d575ecc95a27a5b61e4ec99cbc99f7a4eb Mon Sep 17 00:00:00 2001 From: bless-stanford <107686196+bless-stanford@users.noreply.github.com> Date: Mon, 10 Mar 2025 14:05:22 -0700 Subject: [PATCH 11/19] Add file attachment handling and upload functionality to Mistral agent and Box plugin --- agent.py | 50 ++++++++++++++++++++++++++++++++- plugins/box_plugin.py | 65 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 114 insertions(+), 1 deletion(-) diff --git a/agent.py b/agent.py index 15850274..a6addff9 100644 --- a/agent.py +++ b/agent.py @@ -6,6 +6,7 @@ from semantic_kernel.connectors.ai.function_choice_behavior import FunctionChoiceBehavior import logging import json +import os from services.box_service import BoxService from services.dropbox_service import DropboxService @@ -53,6 +54,9 @@ - Dropbox uses file paths for operations - File operations focus on temporary links and direct access +When a user attaches a file and asks to upload it, use the upload_file function from the Box plugins. +You can find the attached file path in the file_paths parameter that will be provided to you. + Do not use # for headers or * - for bullet points as these don't render in Discord. Keep responses concise when possible, as Discord has a 2000-character limit per message.""" @@ -116,7 +120,26 @@ async def run(self, message: discord.Message): original_content = message.content user_id = str(message.author.id) - augmented_content = f"{original_content}\n\n[system: user_id={user_id}]" + # Handle file attachments + attachment_info = "" + file_paths = [] + if message.attachments: + # Create temp directory if it doesn't exist + if not os.path.exists("temp"): + os.makedirs("temp") + + # Download all attachments + for i, attachment in enumerate(message.attachments): + file_path = f"temp/{attachment.filename}" + await attachment.save(file_path) + file_paths.append(file_path) + attachment_info += f"\n[Attachment {i+1}: {attachment.filename}, path: {file_path}]" + + # Add attachment info to the message + if attachment_info: + augmented_content = f"{original_content}\n\n[system: user_id={user_id}, attached_files=true]{attachment_info}" + else: + augmented_content = f"{original_content}\n\n[system: user_id={user_id}]" # Add the user's message to the chat history self.chat_history.add_user_message(augmented_content) @@ -128,6 +151,12 @@ async def run(self, message: discord.Message): kernel_arguments = KernelArguments() kernel_arguments["user_id"] = user_id + # Add file paths to the kernel arguments if there are any + if file_paths: + kernel_arguments["file_paths"] = file_paths + if len(file_paths) == 1: + kernel_arguments["file_path"] = file_paths[0] + # Update user context in cloud plugin manager self.cloud_plugin_manager.update_user_context(self.kernel, user_id) @@ -179,6 +208,9 @@ async def run(self, message: discord.Message): elif "list" in func_name: path = args.get("path", "root folder") formatted_response = f"I'll list the contents of '{path}' in your {service_name} account..." + elif "upload" in func_name: + file_name = args.get("file_name", "your file") + formatted_response = f"I'm uploading '{file_name}' to your {service_name} account..." else: formatted_response = f"I'm processing your {service_name} request..." @@ -206,6 +238,12 @@ async def run(self, message: discord.Message): f"Please use the `!authorize-{service_name}` command to connect your account." ) self.chat_history.add_assistant_message(error_message) + + # Clean up temporary files + for path in file_paths: + if os.path.exists(path): + os.remove(path) + return [error_message] # Add the assistant's response to the chat history @@ -251,6 +289,11 @@ async def run(self, message: discord.Message): # Log the response for debugging logger.info(f"Generated response for user {message.author.id} (length: {len(formatted_content)})") + # Clean up temporary files + for path in file_paths: + if os.path.exists(path): + os.remove(path) + # Always return a list of chunks if len(formatted_content) > self.MAX_LENGTH: return self.split_response(formatted_content) @@ -259,6 +302,11 @@ async def run(self, message: discord.Message): except Exception as e: logger.error(f"Error processing request: {str(e)}", exc_info=True) + # Clean up temporary files on error + for path in file_paths: + if os.path.exists(path): + os.remove(path) + # Check for authentication errors in the exception if any(phrase in str(e) for phrase in ["authorization", "authorize", "authenticate"]): service_match = re.search(r'!authorize-(\w+)', str(e)) diff --git a/plugins/box_plugin.py b/plugins/box_plugin.py index 6e379370..e18b903e 100644 --- a/plugins/box_plugin.py +++ b/plugins/box_plugin.py @@ -1,3 +1,4 @@ +import os from semantic_kernel.functions import kernel_function from semantic_kernel.functions.kernel_function_from_prompt import KernelFunctionFromPrompt from services.box_service import BoxService @@ -151,6 +152,70 @@ async def delete_file( except Exception as e: logger.error(f"Error deleting file: {str(e)}") return f"An error occurred while deleting the file: {str(e)}" + @kernel_function( + name="upload_file", + description="Uploads an attached file to the user's Box account" + ) + async def upload_file( + self, + file_url: str, + file_name: str = None, + user_id: str = None, + kernel = None + ) -> str: + """ + Uploads an attached file to the user's Box account. + + Args: + file_url: URL or path to the local file + file_name: Optional name to use when storing the file (if different from source) + user_id: The user's ID (automatically provided) + + Returns: + str: Success message with file details or error message + """ + try: + if not user_id: + return "Error: User ID not available. Please try again later." + + # Create temp directory if it doesn't exist + if not os.path.exists("temp"): + os.makedirs("temp") + + # If no file name provided, use the original name from URL + if not file_name and file_url: + file_name = os.path.basename(file_url) + + # Handle the case where file is already downloaded + if os.path.exists(file_url): + local_file_path = file_url + else: + # Could add code here to download from a URL if needed + return "Error: File not found. Please attach a file directly to your message." + + # Upload to Box + file_info = await self.box_service.upload_file(user_id, local_file_path, file_name) + + # Create a response with file details + if file_info and 'id' in file_info: + response = f"✅ File '{file_name}' uploaded successfully to Box!\n" + response += f"**File ID:** {file_info['id']}\n" + + # Get a view link if possible + try: + view_link = await self.box_service.get_file_view_link(user_id, file_info['id']) + response += f"**View Link:** {view_link}" + except: + # If getting a view link fails, that's okay + pass + + return response + else: + return f"File upload completed, but no file information was returned." + + except Exception as e: + logger.error(f"Error uploading file: {str(e)}") + return f"An error occurred while uploading the file: {str(e)}" @kernel_function( name="get_file_download_link", From fdc8dce14518925a87cfc80b45a9c73c49cb8b8e Mon Sep 17 00:00:00 2001 From: bless-stanford <107686196+bless-stanford@users.noreply.github.com> Date: Mon, 10 Mar 2025 14:10:50 -0700 Subject: [PATCH 12/19] Add Dropbox upload functionality and enhance file handling in Dropbox plugin --- bot.py | 40 +++++++++++++++ plugins/dropbox_plugin.py | 101 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 141 insertions(+) diff --git a/bot.py b/bot.py index 5b08055d..c6e35e47 100644 --- a/bot.py +++ b/bot.py @@ -157,6 +157,46 @@ async def box_upload(ctx): if os.path.exists(file_path): os.remove(file_path) +@bot.command(name="dropbox-upload", help="Upload a file to Dropbox") +async def dropbox_upload(ctx): + """ + Uploads an attached file to Dropbox. + """ + if not ctx.message.attachments: + await ctx.send("Please attach a file to upload.") + return + + attachment = ctx.message.attachments[0] + + # Create temp directory if it doesn't exist + if not os.path.exists("temp"): + os.makedirs("temp") + + # Download the attachment + file_path = f"temp/{attachment.filename}" + await attachment.save(file_path) + + try: + # Upload to Dropbox + dropbox_service = DropboxService() + dropbox_path = f"/{attachment.filename}" # Will be stored in root folder + file_info = await dropbox_service.upload_file(str(ctx.author.id), file_path, dropbox_path) + + # Send confirmation + await ctx.send(f"File uploaded to Dropbox! Path: {file_info['path_display']}") + + # Clean up temp file + if os.path.exists(file_path): + os.remove(file_path) + except Exception as e: + error_msg = f"Error uploading file: {str(e)}" + logger.error(error_msg) + await ctx.send(error_msg[:1900]) + + # Clean up temp file on error too + if os.path.exists(file_path): + os.remove(file_path) + @bot.command(name="authorize-dropbox", help="Authorize the bot to access your Dropbox account") async def authorize_dropbox(ctx): """ diff --git a/plugins/dropbox_plugin.py b/plugins/dropbox_plugin.py index b3e9fdaf..0f90c129 100644 --- a/plugins/dropbox_plugin.py +++ b/plugins/dropbox_plugin.py @@ -2,6 +2,7 @@ from semantic_kernel.functions.kernel_function_from_prompt import KernelFunctionFromPrompt from services.dropbox_service import DropboxService import logging +import os logger = logging.getLogger("dropbox_plugins") @@ -199,6 +200,106 @@ async def delete_file( except Exception as e: logger.error(f"Error deleting file: {str(e)}") return f"An error occurred while deleting the file: {str(e)}" + + @kernel_function( + name="upload_file", + description="Uploads an attached file to the user's Dropbox account" + ) + async def upload_file( + self, + file_url: str, + file_name: str = None, + dropbox_path: str = None, + user_id: str = None, + kernel = None + ) -> str: + """ + Uploads an attached file to the user's Dropbox account. + + Args: + file_url: URL or path to the local file + file_name: Optional name to use when storing the file (if different from source) + dropbox_path: Path where to store the file in Dropbox (default: root folder) + user_id: The user's ID (automatically provided) + + Returns: + str: Success message with file details or error message + """ + try: + if not user_id: + return "Error: User ID not available. Please try again later." + + # Create temp directory if it doesn't exist + if not os.path.exists("temp"): + os.makedirs("temp") + + # If no file name provided, use the original name from URL + if not file_name and file_url: + file_name = os.path.basename(file_url) + + # Handle the case where file is already downloaded + if os.path.exists(file_url): + local_file_path = file_url + else: + # Could add code here to download from a URL if needed + return "Error: File not found. Please attach a file directly to your message." + + # Determine the dropbox path + if not dropbox_path: + # Store in root folder if no path specified + dropbox_path = f"/{file_name}" + else: + # Ensure path starts with a slash + if not dropbox_path.startswith('/'): + dropbox_path = f"/{dropbox_path}" + + # If path doesn't end with the filename, append it + if not dropbox_path.endswith(file_name): + # Check if path ends with a slash + if dropbox_path.endswith('/'): + dropbox_path = f"{dropbox_path}{file_name}" + else: + dropbox_path = f"{dropbox_path}/{file_name}" + + # Upload to Dropbox + file_info = await self.dropbox_service.upload_file(user_id, local_file_path, dropbox_path) + + # Create a response with file details + if file_info: + response = f"✅ File '{file_name}' uploaded successfully to Dropbox!\n" + + # Add path information + path_display = file_info.get('path_display', dropbox_path) + response += f"**Path:** {path_display}\n" + + # Add size information if available + if 'size' in file_info: + size_str = self._format_file_size(file_info['size']) + response += f"**Size:** {size_str}\n" + + # Try to create a sharing link + try: + sharing_info = await self.dropbox_service.share_file(user_id, path_display) + shared_url = self._extract_shared_url(sharing_info) + if shared_url: + response += f"**Shared Link:** {shared_url}\n" + except Exception as share_err: + logger.warning(f"Could not create sharing link: {str(share_err)}") + # Try to get a temporary link instead + try: + temp_link = await self.dropbox_service.get_temporary_link(user_id, path_display) + if temp_link: + response += f"**Temporary Link:** {temp_link}\n" + except Exception as temp_err: + logger.warning(f"Could not create temporary link: {str(temp_err)}") + + return response + else: + return f"File upload completed, but no file information was returned." + + except Exception as e: + logger.error(f"Error uploading file: {str(e)}") + return f"An error occurred while uploading the file: {str(e)}" @kernel_function( name="get_file_download_link", From f79f2528fe4a35f64b90432f2f6e64042ddd638f Mon Sep 17 00:00:00 2001 From: bless-stanford <107686196+bless-stanford@users.noreply.github.com> Date: Mon, 10 Mar 2025 15:08:32 -0700 Subject: [PATCH 13/19] Add Google Drive integration to Mistral agent and update cloud plugin manager --- agent.py | 15 +- bot.py | 85 +++ plugins/cloud_plugin_manager.py | 27 +- plugins/google_drive_plugin.py | 609 +++++++++++++++++++ server.py | 34 +- services/google_drive_service.py | 971 +++++++++++++++++++++++++++++++ 6 files changed, 1735 insertions(+), 6 deletions(-) create mode 100644 plugins/google_drive_plugin.py create mode 100644 services/google_drive_service.py diff --git a/agent.py b/agent.py index a6addff9..bab06aa7 100644 --- a/agent.py +++ b/agent.py @@ -10,6 +10,7 @@ from services.box_service import BoxService from services.dropbox_service import DropboxService +from services.google_drive_service import GoogleDriveService from plugins.cloud_plugin_manager import CloudPluginManager @@ -54,6 +55,11 @@ - Dropbox uses file paths for operations - File operations focus on temporary links and direct access +For Google Drive: +- Use Google Drive for collaborative work and integration with Google Workspace +- Google Drive uses file IDs and folder IDs for operations +- File operations include sharing, viewing, and downloading + When a user attaches a file and asks to upload it, use the upload_file function from the Box plugins. You can find the attached file path in the file_paths parameter that will be provided to you. @@ -77,11 +83,13 @@ def __init__(self, max_context_messages=10): # Initialize cloud services self.box_service = BoxService() self.dropbox_service = DropboxService() + self.google_drive_service = GoogleDriveService() # Initialize plugin manager and register plugins self.cloud_plugin_manager = CloudPluginManager( box_service=self.box_service, - dropbox_service=self.dropbox_service + dropbox_service=self.dropbox_service, + google_drive_service=self.google_drive_service ) # Register all cloud plugins with the kernel @@ -191,6 +199,8 @@ async def run(self, message: discord.Message): service_name = "Box" elif "dropbox" in func_name.lower(): service_name = "Dropbox" + elif "gdrive" in func_name.lower() or "google" in func_name.lower(): + service_name = "Google Drive" else: service_name = "cloud storage" @@ -225,7 +235,8 @@ async def run(self, message: discord.Message): "needs to be authorized", "Please use the `!authorize", "not authorized", - "authorization required" + "authorization required", + "Google Drive authorization has expired" ] if any(phrase in response.content for phrase in auth_error_phrases): diff --git a/bot.py b/bot.py index c6e35e47..56b2a5bc 100644 --- a/bot.py +++ b/bot.py @@ -6,6 +6,7 @@ from agent import MistralAgent from services.box_service import BoxService from services.dropbox_service import DropboxService +from services.google_drive_service import GoogleDriveService from server import start_server PREFIX = "!" @@ -32,6 +33,7 @@ # Initialize cloud service instances box_service = BoxService() dropbox_service = DropboxService() +google_drive_service = GoogleDriveService() async def send_split_message(message: discord.Message, response: str | list[str]): """ @@ -214,6 +216,66 @@ async def authorize_dropbox(ctx): logger.error(error_msg) await ctx.send(error_msg[:1900]) +@bot.command(name="authorize-gdrive", help="Authorize the bot to access your Google Drive account") +async def authorize_gdrive(ctx): + """ + Sends a Google Drive authorization link to the user via DM. + """ + try: + # Get authorization URL for the user + auth_url = await google_drive_service.get_authorization_url(str(ctx.author.id)) + + # Send the URL as a DM to the user + await ctx.author.send(f"Please authorize access to your Google Drive account by clicking this link: {auth_url}") + await ctx.send("I've sent you a DM with the authorization link!") + except Exception as e: + error_msg = f"Error generating Google Drive authorization link: {str(e)}" + logger.error(error_msg) + await ctx.send(error_msg[:1900]) + +@bot.command(name="gdrive-upload", help="Upload a file to Google Drive") +async def gdrive_upload(ctx): + """ + Uploads an attached file to Google Drive. + """ + if not ctx.message.attachments: + await ctx.send("Please attach a file to upload.") + return + + attachment = ctx.message.attachments[0] + + # Create temp directory if it doesn't exist + if not os.path.exists("temp"): + os.makedirs("temp") + + # Download the attachment + file_path = f"temp/{attachment.filename}" + await attachment.save(file_path) + + try: + # Upload to Google Drive + file_info = await google_drive_service.upload_file( + str(ctx.author.id), + file_path, + attachment.filename + ) + + # Send confirmation with view link + view_link = file_info.get('webViewLink', 'No view link available') + await ctx.send(f"File uploaded to Google Drive! File ID: {file_info['id']}\nView link: {view_link}") + + # Clean up temp file + if os.path.exists(file_path): + os.remove(file_path) + except Exception as e: + error_msg = f"Error uploading file: {str(e)}" + logger.error(error_msg) + await ctx.send(error_msg[:1900]) + + # Clean up temp file on error too + if os.path.exists(file_path): + os.remove(file_path) + @bot.command(name="cloud-status", help="Check your cloud service connections") async def cloud_status(ctx): """ @@ -274,6 +336,29 @@ async def cloud_status(ctx): value=f"⚠️ Error checking connection\n```{str(e)}```", inline=False ) + + # Check Google Drive connection + try: + # Try to load the token to see if the user is authenticated + gdrive_token = await google_drive_service._load_token(str(ctx.author.id)) + if gdrive_token: + embed.add_field( + name="Google Drive Status", + value="✅ Connected", + inline=False + ) + else: + embed.add_field( + name="Google Drive Status", + value="❌ Not connected\n*Use !authorize-gdrive to connect*", + inline=False + ) + except Exception as e: + embed.add_field( + name="Google Drive Status", + value=f"⚠️ Error checking connection\n```{str(e)}```", + inline=False + ) await ctx.send(embed=embed) diff --git a/plugins/cloud_plugin_manager.py b/plugins/cloud_plugin_manager.py index 8ed6e50f..c3ca1937 100644 --- a/plugins/cloud_plugin_manager.py +++ b/plugins/cloud_plugin_manager.py @@ -1,8 +1,10 @@ from semantic_kernel.kernel import Kernel from services.box_service import BoxService from services.dropbox_service import DropboxService +from services.google_drive_service import GoogleDriveService from plugins.box_plugin import BoxPlugins from plugins.dropbox_plugin import DropboxPlugins +from plugins.google_drive_plugin import GoogleDrivePlugins import logging logger = logging.getLogger("cloud_plugin_manager") @@ -10,10 +12,10 @@ class CloudPluginManager: """ Manager for cloud storage plugins to use with Semantic Kernel. - Consolidates Box and Dropbox plugins into a single interface. + Consolidates Box, Dropbox, and Google Drive plugins into a single interface. """ - def __init__(self, box_service=None, dropbox_service=None): + def __init__(self, box_service=None, dropbox_service=None, google_drive_service=None): """ Initialize the cloud plugin manager with service instances. If no services are provided, new ones will be created. @@ -21,13 +23,16 @@ def __init__(self, box_service=None, dropbox_service=None): Args: box_service: BoxService instance or None dropbox_service: DropboxService instance or None + google_drive_service: GoogleDriveService instance or None """ self.box_service = box_service or BoxService() self.dropbox_service = dropbox_service or DropboxService() + self.google_drive_service = google_drive_service or GoogleDriveService() # Initialize plugin instances self.box_plugins = BoxPlugins(self.box_service) self.dropbox_plugins = DropboxPlugins(self.dropbox_service) + self.google_drive_plugins = GoogleDrivePlugins(self.google_drive_service) def register_plugins(self, kernel: Kernel) -> Kernel: """ @@ -48,6 +53,10 @@ def register_plugins(self, kernel: Kernel) -> Kernel: kernel.add_plugin(self.dropbox_plugins, "dropbox") logger.info("Dropbox plugins registered with kernel") + # Register Google Drive plugins + kernel.add_plugin(self.google_drive_plugins, "gdrive") + logger.info("Google Drive plugins registered with kernel") + return kernel except Exception as e: logger.error(f"Error registering cloud plugins: {str(e)}") @@ -80,7 +89,19 @@ def get_plugin_descriptions(self) -> str: descriptions += "- `dropbox.list_folder`: List files and folders in a Dropbox path\n" descriptions += "- `dropbox.delete_file`: Delete a file from Dropbox\n" descriptions += "- `dropbox.get_file_download_link`: Get a temporary download link for a Dropbox file\n" - descriptions += "- `dropbox.share_file`: Create a shared link for a Dropbox file\n" + descriptions += "- `dropbox.share_file`: Create a shared link for a Dropbox file\n\n" + + # Google Drive plugins + descriptions += "## Google Drive Plugins\n" + descriptions += "Use these to interact with your Google Drive account:\n" + descriptions += "- `gdrive.create_folder`: Create a new folder in Google Drive\n" + descriptions += "- `gdrive.search_file`: Search for files in Google Drive\n" + descriptions += "- `gdrive.delete_file`: Delete a file from Google Drive\n" + descriptions += "- `gdrive.upload_file`: Upload a file to Google Drive\n" + descriptions += "- `gdrive.get_file_download_link`: Get a download link for a Google Drive file\n" + descriptions += "- `gdrive.get_file_view_link`: Get a view link for a Google Drive file\n" + descriptions += "- `gdrive.share_file`: Share a Google Drive file with another user\n" + descriptions += "- `gdrive.move_file`: Move a file to a different folder in Google Drive\n" return descriptions diff --git a/plugins/google_drive_plugin.py b/plugins/google_drive_plugin.py new file mode 100644 index 00000000..3626dcac --- /dev/null +++ b/plugins/google_drive_plugin.py @@ -0,0 +1,609 @@ +import json +import os +from semantic_kernel.functions import kernel_function +from semantic_kernel.functions.kernel_function_from_prompt import KernelFunctionFromPrompt +from services.google_drive_service import GoogleDriveService +import logging + +logger = logging.getLogger("google_drive_plugins") + +class GoogleDrivePlugins: + """ + Plugins for interacting with Google Drive cloud storage. + """ + + def __init__(self, drive_service=None): + """ + Initialize the Google Drive plugins with a GoogleDriveService. + If no service is provided, a new one will be created. + """ + self.drive_service = drive_service or GoogleDriveService() + + @kernel_function( + name="create_folder", + description="Creates a new folder in the user's Google Drive account" + ) + async def create_folder( + self, + folder_name: str, + parent_folder_id: str = "root", + user_id: str = None, + kernel = None + ) -> str: + """ + Creates a new folder in the user's Google Drive account. + + Args: + folder_name: Name of the folder to create + parent_folder_id: ID of the parent folder (default: "root" for root) + user_id: The user's ID (automatically provided) + + Returns: + str: Success message with folder details or error message + """ + try: + if not user_id: + return "Error: User ID not available. Please try again later." + + folder = await self.drive_service.create_folder(user_id, folder_name, parent_folder_id) + + if folder: + return f"Folder '{folder_name}' created successfully with ID: {folder['id']}" + else: + return f"Failed to create folder '{folder_name}'." + + except Exception as e: + logger.error(f"Error creating folder: {str(e)}") + return f"An error occurred while creating the folder: {str(e)}" + + @kernel_function( + name="search_file", + description="Searches for files in the user's Google Drive account and returns details in a user friendly way" + ) + async def search_file( + self, + query: str, + user_id: str = None, + kernel = None + ) -> str: + """ + Searches for files in the user's Google Drive account. + + Args: + query: Search query or file name + user_id: The user's ID (automatically provided) + + Returns: + str: File details or search results summary + """ + try: + # Get user_id from kernel.data instead of function parameter + if not user_id and kernel and hasattr(kernel, 'arguments'): + user_id = kernel.arguments.get("user_id") + + if not user_id: + return "Error: User ID not available. Please try again later." + + files = await self.drive_service.search_files(user_id, query) + + if not files or len(files) == 0: + return f"No files found matching '{query}'." + + if len(files) == 1: + return self._create_file_detail(files[0]) + + # If multiple files, return a summary + return self._create_search_results_summary(files) + + except Exception as e: + logger.error(f"Error searching files: {str(e)}") + return f"An error occurred while searching for files: {str(e)}" + + @kernel_function( + name="delete_file", + description="Searches and deletes a file from the user's Google Drive account" + ) + async def delete_file( + self, + query: str, + user_id: str = None, + kernel = None + ) -> str: + """ + Searches and deletes a file from the user's Google Drive account. + + Args: + query: Search query or file name + user_id: The user's ID (automatically provided) + kernel: Semantic Kernel instance for finding most relevant file + + Returns: + str: Success message or error message + """ + try: + if not user_id: + return "Error: User ID not available. Please try again later." + + files = await self.drive_service.search_files(user_id, query) + + if not files or len(files) == 0: + return f"No files found matching '{query}'." + + if len(files) == 1: + file = files[0] + await self.drive_service.delete_file(user_id, file['id']) + return f"File '{file['name']}' has been successfully deleted." + + # If multiple files and kernel is provided, find most relevant + if kernel and len(files) > 1: + most_relevant_file = await self._find_most_relevant_file(kernel, files, query) + + if most_relevant_file: + await self.drive_service.delete_file(user_id, most_relevant_file['id']) + return f"File '{most_relevant_file['name']}' has been successfully deleted." + + # If multiple files and no most relevant found, return summary + return f"Multiple files found matching '{query}'. Please be more specific:\n" + \ + "\n".join([f"- {file['name']}" for file in files[:5]]) + + except Exception as e: + logger.error(f"Error deleting file: {str(e)}") + return f"An error occurred while deleting the file: {str(e)}" + + @kernel_function( + name="upload_file", + description="Uploads an attached file to the user's Google Drive account" + ) + async def upload_file( + self, + file_url: str, + file_name: str = None, + parent_folder_id: str = "root", + user_id: str = None, + kernel = None + ) -> str: + """ + Uploads an attached file to the user's Google Drive account. + + Args: + file_url: URL or path to the local file + file_name: Optional name to use when storing the file (if different from source) + parent_folder_id: ID of the parent folder (default: "root") + user_id: The user's ID (automatically provided) + + Returns: + str: Success message with file details or error message + """ + try: + if not user_id: + return "Error: User ID not available. Please try again later." + + # Create temp directory if it doesn't exist + if not os.path.exists("temp"): + os.makedirs("temp") + + # If no file name provided, use the original name from URL + if not file_name and file_url: + file_name = os.path.basename(file_url) + + # Handle the case where file is already downloaded + if os.path.exists(file_url): + local_file_path = file_url + else: + # Could add code here to download from a URL if needed + return "Error: File not found. Please attach a file directly to your message." + + # Upload to Google Drive + file_info = await self.drive_service.upload_file( + user_id, + local_file_path, + file_name, + parent_folder_id + ) + + # Create a response with file details + if file_info and 'id' in file_info: + response = f"✅ File '{file_name}' uploaded successfully to Google Drive!\n" + response += f"**File ID:** {file_info['id']}\n" + + # Add webViewLink if available + if 'webViewLink' in file_info: + response += f"**View Link:** {file_info['webViewLink']}" + + return response + else: + return f"File upload completed, but no file information was returned." + + except Exception as e: + logger.error(f"Error uploading file: {str(e)}") + return f"An error occurred while uploading the file: {str(e)}" + + @kernel_function( + name="get_file_download_link", + description="Gets a download link for a file in the user's Google Drive account" + ) + async def get_file_download_link( + self, + query: str, + user_id: str = None, + kernel = None + ) -> str: + """ + Gets a download link for a file in the user's Google Drive account. + + Args: + query: Search query or file name + user_id: The user's ID (automatically provided) + kernel: Semantic Kernel instance for finding most relevant file + + Returns: + str: Download link or error message + """ + try: + logger.info(f"Getting download link for query '{query}' from user {user_id}") + + if not user_id: + logger.error("User ID not available") + return "Error: User ID not available. Please try again later." + + logger.info(f"Searching for files matching '{query}'") + files = await self.drive_service.search_files(user_id, query) + + logger.info(f"Found {len(files) if files else 0} files matching '{query}'") + + if not files or len(files) == 0: + return f"No files found matching '{query}'." + + if len(files) == 1: + file = files[0] + logger.info(f"Found single file: {file.get('name', 'Unknown')}, ID: {file.get('id', 'Unknown')}") + + logger.info(f"Getting file details for ID: {file['id']}") + file_info = await self.drive_service.get_file(user_id, file['id']) + + logger.info(f"File details: {json.dumps(file_info, indent=2)}") + + if 'webContentLink' in file_info and file_info['webContentLink']: + logger.info(f"Found download link: {file_info['webContentLink']}") + return f"Download link for file '{file['name']}':\n{file_info['webContentLink']}" + elif 'webViewLink' in file_info and file_info['webViewLink']: + logger.info(f"No download link found, using view link: {file_info['webViewLink']}") + return f"No direct download link available for '{file['name']}'. You can access it via this view link instead:\n{file_info['webViewLink']}" + else: + logger.warning(f"No links found for file: {file['name']}") + return f"No direct links available for '{file['name']}'. You may need to access it via the Google Drive web interface." + + # If multiple files and kernel is provided, find most relevant + if kernel and len(files) > 1: + logger.info(f"Multiple files found, attempting to find most relevant") + most_relevant_file = await self._find_most_relevant_file(kernel, files, query) + + if most_relevant_file: + logger.info(f"Most relevant file: {most_relevant_file.get('name', 'Unknown')}") + + file_info = await self.drive_service.get_file(user_id, most_relevant_file['id']) + logger.info(f"File details for most relevant: {json.dumps(file_info, indent=2)}") + + if 'webContentLink' in file_info and file_info['webContentLink']: + return f"Download link for file '{most_relevant_file['name']}':\n{file_info['webContentLink']}" + elif 'webViewLink' in file_info and file_info['webViewLink']: + return f"No direct download link available for '{most_relevant_file['name']}'. You can access it via this view link instead:\n{file_info['webViewLink']}" + else: + return f"No direct links available for '{most_relevant_file['name']}'." + + # If multiple files and no most relevant found, return summary + file_list = "\n".join([f"- {file['name']}" for file in files[:5]]) + logger.info(f"Returning list of multiple files: {file_list}") + return f"Multiple files found matching '{query}'. Please be more specific:\n{file_list}" + + except Exception as e: + logger.error(f"Error getting download link: {str(e)}", exc_info=True) + return f"An error occurred while getting the download link: {str(e)}" + + @kernel_function( + name="get_file_view_link", + description="Gets a shareable view link for a file in the user's Google Drive account" + ) + async def get_file_view_link( + self, + query: str, + user_id: str = None, + kernel = None + ) -> str: + """ + Gets a shareable view link for a file in the user's Google Drive account. + + Args: + query: Search query or file name + user_id: The user's ID (automatically provided) + kernel: Semantic Kernel instance for finding most relevant file + + Returns: + str: View link or error message + """ + try: + if not user_id: + return "Error: User ID not available. Please try again later." + + files = await self.drive_service.search_files(user_id, query) + + if not files or len(files) == 0: + return f"No files found matching '{query}'." + + if len(files) == 1: + file = files[0] + file_info = await self.drive_service.get_file(user_id, file['id']) + + if 'webViewLink' in file_info and file_info['webViewLink']: + return f"View link for file '{file['name']}':\n{file_info['webViewLink']}" + else: + return f"No view link available for '{file['name']}'." + + # If multiple files and kernel is provided, find most relevant + if kernel and len(files) > 1: + most_relevant_file = await self._find_most_relevant_file(kernel, files, query) + + if most_relevant_file: + file_info = await self.drive_service.get_file(user_id, most_relevant_file['id']) + + if 'webViewLink' in file_info and file_info['webViewLink']: + return f"View link for file '{most_relevant_file['name']}':\n{file_info['webViewLink']}" + else: + return f"No view link available for '{most_relevant_file['name']}'." + + # If multiple files and no most relevant found, return summary + return f"Multiple files found matching '{query}'. Please be more specific:\n" + \ + "\n".join([f"- {file['name']}" for file in files[:5]]) + + except Exception as e: + logger.error(f"Error getting view link: {str(e)}") + return f"An error occurred while getting the view link: {str(e)}" + + @kernel_function( + name="share_file", + description="Shares a file with another user" + ) + async def share_file( + self, + query: str, + email: str, + role: str = "reader", + user_id: str = None, + kernel = None + ) -> str: + """ + Shares a file with another user. + + Args: + query: Search query or file name + email: Email of the user to share with + role: Role to assign (reader, writer, commenter) + user_id: The user's ID (automatically provided) + kernel: Semantic Kernel instance for finding most relevant file + + Returns: + str: Success message or error message + """ + try: + if not user_id: + return "Error: User ID not available. Please try again later." + + # Validate role (Google Drive uses reader, writer, commenter) + valid_roles = ["reader", "writer", "commenter"] + if role not in valid_roles: + return f"Invalid role '{role}'. Valid roles are: {', '.join(valid_roles)}" + + files = await self.drive_service.search_files(user_id, query) + + if not files or len(files) == 0: + return f"No files found matching '{query}'." + + if len(files) == 1: + file = files[0] + await self.drive_service.share_file(user_id, file['id'], email, role) + + file_info = await self.drive_service.get_file(user_id, file['id']) + view_link = file_info.get('webViewLink', 'No view link available') + + return f"File '{file['name']}' has been shared with {email} as a {role}. They can access the file at: {view_link}" + + # If multiple files and kernel is provided, find most relevant + if kernel and len(files) > 1: + most_relevant_file = await self._find_most_relevant_file(kernel, files, query) + + if most_relevant_file: + await self.drive_service.share_file(user_id, most_relevant_file['id'], email, role) + + file_info = await self.drive_service.get_file(user_id, most_relevant_file['id']) + view_link = file_info.get('webViewLink', 'No view link available') + + return f"File '{most_relevant_file['name']}' has been shared with {email} as a {role}. They can access the file at: {view_link}" + + # If multiple files and no most relevant found, return summary + return f"Multiple files found matching '{query}'. Please be more specific:\n" + \ + "\n".join([f"- {file['name']}" for file in files[:5]]) + + except Exception as e: + logger.error(f"Error sharing file: {str(e)}") + return f"An error occurred while sharing the file: {str(e)}" + + @kernel_function( + name="move_file", + description="Moves a file to a different folder in the user's Google Drive account" + ) + async def move_file( + self, + query: str, + destination_folder_id: str, + user_id: str = None, + kernel = None + ) -> str: + """ + Moves a file to a different folder. + + Args: + query: Search query or file name + destination_folder_id: ID of the destination folder + user_id: The user's ID (automatically provided) + kernel: Semantic Kernel instance for finding most relevant file + + Returns: + str: Success message or error message + """ + try: + if not user_id: + return "Error: User ID not available. Please try again later." + + files = await self.drive_service.search_files(user_id, query) + + if not files or len(files) == 0: + return f"No files found matching '{query}'." + + if len(files) == 1: + file = files[0] + await self.drive_service.move_file(user_id, file['id'], destination_folder_id) + return f"File '{file['name']}' has been successfully moved to the destination folder." + + # If multiple files and kernel is provided, find most relevant + if kernel and len(files) > 1: + most_relevant_file = await self._find_most_relevant_file(kernel, files, query) + + if most_relevant_file: + await self.drive_service.move_file(user_id, most_relevant_file['id'], destination_folder_id) + return f"File '{most_relevant_file['name']}' has been successfully moved to the destination folder." + + # If multiple files and no most relevant found, return summary + return f"Multiple files found matching '{query}'. Please be more specific:\n" + \ + "\n".join([f"- {file['name']}" for file in files[:5]]) + + except Exception as e: + logger.error(f"Error moving file: {str(e)}") + return f"An error occurred while moving the file: {str(e)}" + + async def _find_most_relevant_file(self, kernel, files, user_query): + """ + Find the most relevant file from a list based on user query. + + Args: + kernel: Semantic Kernel instance + files: List of files + user_query: The user's query + + Returns: + dict: The most relevant file or None + """ + try: + # Create a function from prompt + rank_files_function = KernelFunctionFromPrompt( + function_name="RankFilesByRelevance", + plugin_name=None, + prompt="Given the user query: '{{$userQuery}}' and a list of file names, " + "rank them by relevance and return the index of the most relevant file. " + "Do not add any comments or explanation to the response.\n" + "File list: {{$fileList}}", + template_format="semantic-kernel" + ) + + # Create file list string + file_list = "\n".join([f"{i}: Name: {file['name']}" for i, file in enumerate(files)]) + + # Create kernel arguments + kernel_arguments = { + "userQuery": user_query, + "fileList": file_list + } + + # Invoke the function + result = await kernel.invoke(rank_files_function, **kernel_arguments) + + # Get the value from the result - might be a list or a string + result_value = result.value + + # Handle different result types + if isinstance(result_value, list) and len(result_value) > 0: + result_text = str(result_value[0]).strip() + elif isinstance(result_value, str): + result_text = result_value.strip() + else: + # Fallback + result_text = str(result_value).strip() + + try: + most_relevant_index = int(result_text) + if 0 <= most_relevant_index < len(files): + return files[most_relevant_index] + except ValueError: + logger.warning(f"Could not parse the relevance index from AI result: {result_text}") + + return None + + except Exception as e: + logger.error(f"Error finding most relevant file: {str(e)}") + return None + + def _create_file_detail(self, file): + """Create a detailed text representation of a file.""" + detail = "**File Details:**\n" + detail += f"**Name:** {file.get('name', 'Unknown')}\n" + detail += f"**ID:** {file.get('id', 'Unknown')}\n" + + # Format file size if available + if 'size' in file: + size_bytes = int(file['size']) + size_str = self._format_file_size(size_bytes) + detail += f"**Size:** {size_str}\n" + + # Add mime type if available + if 'mimeType' in file: + detail += f"**Type:** {file['mimeType']}\n" + + # Add dates if available + if 'modifiedTime' in file: + detail += f"**Modified At:** {file['modifiedTime']}\n" + + # Add view link if available + if 'webViewLink' in file: + detail += f"**View Link:** {file['webViewLink']}\n" + + # Add download link if available + if 'webContentLink' in file: + detail += f"**Download Link:** {file['webContentLink']}\n" + + return detail + + def _create_search_results_summary(self, files): + """Create a summary of multiple search results.""" + summary = "**Multiple files found. Here are the details:**\n\n" + + for i, file in enumerate(files[:5], 1): # Limit to 5 files and number them + summary += f"**{i}. {file.get('name', 'Unknown')}**\n" + summary += f" ID: {file.get('id', 'Unknown')}\n" + + # Add size if available + if 'size' in file: + size_str = self._format_file_size(int(file['size'])) + summary += f" Size: {size_str}\n" + + # Add type if available + if 'mimeType' in file: + summary += f" Type: {file['mimeType']}\n" + + summary += "\n" + + if len(files) > 5: + summary += f"\n...and {len(files) - 5} more files.\n" + + summary += "\nPlease provide a more specific query to find the exact file you want." + + return summary + + def _format_file_size(self, bytes): + """Format file size in human-readable form.""" + sizes = ["B", "KB", "MB", "GB", "TB"] + order = 0 + size = float(bytes) + while size >= 1024 and order < len(sizes) - 1: + order += 1 + size /= 1024 + + return f"{size:.2f} {sizes[order]}" \ No newline at end of file diff --git a/server.py b/server.py index 87fd62ea..164576b4 100644 --- a/server.py +++ b/server.py @@ -3,6 +3,7 @@ import uvicorn from services.box_service import BoxService from services.dropbox_service import DropboxService +from services.google_drive_service import GoogleDriveService from helpers.token_helpers import TokenEncryptionHelper import asyncio import logging @@ -14,6 +15,7 @@ app = FastAPI() box_service = BoxService() dropbox_service = DropboxService() +google_drive_service = GoogleDriveService() # This will be set from bot.py bot = None @@ -79,7 +81,7 @@ def get_success_html(service_name): @app.get("/") async def root(): - return {"message": "OAuth Callback Server for Box and Dropbox"} + return {"message": "OAuth Callback Server for Box, Dropbox, and Google Drive"} @app.get("/box/callback") async def box_callback(code: str, state: str): @@ -141,6 +143,36 @@ async def dropbox_callback(code: str, state: str): logger.error(f"Error in Dropbox callback: {str(e)}") return {"error": str(e)} +@app.get("/gdrive/callback") +async def gdrive_callback(code: str, state: str): + """ + Handle the OAuth callback from Google Drive. + + This endpoint receives the authorization code from Google Drive after a user + authorizes the application. It exchanges the code for access and + refresh tokens, stores them securely, and notifies the user. + """ + try: + # Get user ID from state + user_id = TokenEncryptionHelper.decrypt_token(state, google_drive_service.encryption_key) + logger.info(f"Received Google Drive callback for user {user_id}") + + # Handle the callback - this stores the tokens + await google_drive_service.handle_auth_callback(state, code) + + # Notify the user through Discord + if bot: + # Schedule the notification in the bot's event loop + asyncio.run_coroutine_threadsafe(notify_user(user_id, "Google Drive"), bot.loop) + + # Use the reusable HTML template + html_content = get_success_html("Google Drive") + + return HTMLResponse(content=html_content) + except Exception as e: + logger.error(f"Error in Google Drive callback: {str(e)}") + return {"error": str(e)} + async def notify_user(user_id, service_name): """ Send a Discord message to notify the user that authorization was successful. diff --git a/services/google_drive_service.py b/services/google_drive_service.py new file mode 100644 index 00000000..d1de3941 --- /dev/null +++ b/services/google_drive_service.py @@ -0,0 +1,971 @@ +import os +import json +import logging +from datetime import datetime, timedelta +from urllib.parse import urlencode +from dotenv import load_dotenv + +import requests +from google.oauth2.credentials import Credentials +from google_auth_oauthlib.flow import Flow +from googleapiclient.discovery import build +from googleapiclient.http import MediaFileUpload, MediaIoBaseDownload + +from helpers.token_helpers import ( + TokenEncryptionHelper, + TokenStorageManager, + create_token_record, + load_or_generate_encryption_key +) + +# Setup logging +logger = logging.getLogger("google_drive_service") + +# Constants for platform and service +PLATFORM = "Google" +SERVICE = "GoogleDriveService" + +# API URLs +GOOGLE_AUTH_BASE_URL = "https://accounts.google.com/o/oauth2/" +GOOGLE_TOKEN_URL = "https://oauth2.googleapis.com/token" +GOOGLE_REVOKE_URL = "https://oauth2.googleapis.com/revoke" + +# Scopes for Google Drive API +SCOPES = ['https://www.googleapis.com/auth/drive'] + + +class GoogleDriveService: + def __init__(self, config=None): + """ + Initialize the Google Drive service with configuration. + + Args: + config: Configuration dictionary or None to load from .env + """ + if config is None: + load_dotenv() + self.client_id = os.getenv("GOOGLE_CLIENT_ID") + self.client_secret = os.getenv("GOOGLE_CLIENT_SECRET") + self.redirect_uri = os.getenv("GOOGLE_DRIVE_REDIRECT_URI") + self.app_name = os.getenv("GOOGLE_APP_NAME", "GoogleDriveApp") + + # Get or generate encryption key using our helper + self.encryption_key = load_or_generate_encryption_key() + else: + self.client_id = config.get("client_id") + self.client_secret = config.get("client_secret") + self.redirect_uri = config.get("redirect_uri") + self.app_name = config.get("app_name", "GoogleDriveApp") + self.encryption_key = config.get("encryption_key") + + # Initialize token storage + self.token_storage = TokenStorageManager() + + async def get_authorization_url(self, user_id): + """ + Get the authorization URL for Google OAuth flow. + + Args: + user_id: The user's ID + + Returns: + str: The authorization URL + """ + if not self.client_id: + raise ValueError("Google Client ID is not set in configuration.") + if not self.redirect_uri: + raise ValueError("Google Redirect URI is not set in configuration.") + + # Encrypt user_id as state parameter + state = TokenEncryptionHelper.encrypt_token(user_id, self.encryption_key) + + # Create a Flow instance + flow = Flow.from_client_config( + { + "web": { + "client_id": self.client_id, + "client_secret": self.client_secret, + "auth_uri": f"{GOOGLE_AUTH_BASE_URL}auth", + "token_uri": GOOGLE_TOKEN_URL, + "redirect_uris": [self.redirect_uri] + } + }, + scopes=SCOPES, + redirect_uri=self.redirect_uri + ) + + # Generate authorization URL + auth_url, _ = flow.authorization_url( + access_type='offline', + include_granted_scopes='true', + state=state, + prompt='consent' # Always show consent screen to get refresh token + ) + + logger.info(f"Generated authorization URL for user {user_id}") + return auth_url + + async def handle_auth_callback(self, state, code): + """ + Handle the authorization callback from Google. + + Args: + state: The state parameter from the callback + code: The authorization code from the callback + """ + if not self.client_id: + raise ValueError("Google Client ID is not set in configuration.") + if not self.client_secret: + raise ValueError("Google Client Secret is not set in configuration.") + if not self.redirect_uri: + raise ValueError("Google Redirect URI is not set in configuration.") + + # Decrypt the user_id from state + user_id = TokenEncryptionHelper.decrypt_token(state, self.encryption_key) + logger.info(f"Processing authorization callback for user {user_id}") + + # Create a Flow instance - but don't specify scopes this time + # This lets the flow accept whatever scopes Google returns + flow = Flow.from_client_config( + { + "web": { + "client_id": self.client_id, + "client_secret": self.client_secret, + "auth_uri": f"{GOOGLE_AUTH_BASE_URL}auth", + "token_uri": GOOGLE_TOKEN_URL, + "redirect_uris": [self.redirect_uri] + } + }, + scopes=None, # Allow any scope to be returned + redirect_uri=self.redirect_uri + ) + + # Exchange code for token + try: + flow.fetch_token(code=code) + credentials = flow.credentials + + # Store token + await self._store_token( + user_id, + credentials.token, + credentials.refresh_token, + credentials.expiry.timestamp() - datetime.now().timestamp() + ) + logger.info(f"Successfully obtained and stored access token for user {user_id}") + except Exception as e: + logger.error(f"Failed to obtain access token: {str(e)}") + raise Exception(f"Failed to obtain user access token: {str(e)}") + + async def revoke_access(self, user_id): + """ + Revoke the Google Drive access for a user. + + Args: + user_id: The user's ID + """ + token_data = await self._get_token_data(user_id) + if not token_data: + raise ValueError("No valid token found for user") + + token = token_data.get("access_token") + if not token: + raise ValueError("No valid access token found for user") + + # Revoke the token + params = {'token': token} + response = requests.post(GOOGLE_REVOKE_URL, params=params) + + if response.status_code in (200, 204): + # Delete the token from storage + self.token_storage.delete_token(user_id, PLATFORM, SERVICE) + logger.info(f"Successfully revoked access for user {user_id}") + else: + logger.error(f"Failed to revoke token: {response.status_code}") + raise Exception(f"Failed to revoke token: {response.status_code}") + + async def create_folder(self, user_id, folder_name, parent_folder_id="root"): + """ + Create a folder in Google Drive. + + Args: + user_id: The user's ID + folder_name: Name of the folder to create + parent_folder_id: ID of the parent folder (default: "root") + + Returns: + dict: The created folder information + """ + service = await self._get_drive_service(user_id) + + folder_metadata = { + 'name': folder_name, + 'mimeType': 'application/vnd.google-apps.folder', + 'parents': [parent_folder_id] + } + + try: + folder = service.files().create( + body=folder_metadata, + fields='id, name, mimeType, webViewLink' + ).execute() + + return folder + except Exception as e: + logger.error(f"Failed to create folder: {str(e)}") + raise Exception(f"Failed to create folder: {str(e)}") + + async def upload_file(self, user_id, file_path, file_name=None, parent_folder_id="root", mime_type=None, description=None): + """ + Upload a file to Google Drive. + + Args: + user_id: The user's ID + file_path: Path to the local file + file_name: Name to use for the file in Google Drive (defaults to local filename) + parent_folder_id: ID of the parent folder (default: "root") + mime_type: MIME type of the file (if None, will be guessed) + description: Optional description for the file + + Returns: + dict: The uploaded file information with file ID and webViewLink + """ + service = await self._get_drive_service(user_id) + + # Use the original filename if none provided + if not file_name: + file_name = os.path.basename(file_path) + + # Create file metadata + file_metadata = { + 'name': file_name, + 'parents': [parent_folder_id] + } + + if description: + file_metadata['description'] = description + + # Create a media upload instance + media = MediaFileUpload(file_path, mimetype=mime_type, resumable=True) + + # Upload the file + try: + file = service.files().create( + body=file_metadata, + media_body=media, + fields='id, name, mimeType, webViewLink' + ).execute() + + return file + except Exception as e: + logger.error(f"Failed to upload file: {str(e)}") + raise Exception(f"Failed to upload file: {str(e)}") + + async def delete_file(self, user_id, file_id): + """ + Delete a file from Google Drive. + + Args: + user_id: The user's ID + file_id: ID of the file to delete + """ + service = await self._get_drive_service(user_id) + + try: + service.files().delete(fileId=file_id).execute() + logger.info(f"Successfully deleted file {file_id}") + except Exception as e: + logger.error(f"Failed to delete file: {str(e)}") + raise Exception(f"Failed to delete file: {str(e)}") + + async def list_files(self, user_id, folder_id="root", page_size=100, query=None): + """ + List files in a folder in Google Drive. + + Args: + user_id: The user's ID + folder_id: ID of the folder (default: "root") + page_size: Maximum number of files to return + query: Optional query string to filter results + + Returns: + list: The files in the folder + """ + service = await self._get_drive_service(user_id) + + # Build the query string + q = f"'{folder_id}' in parents and trashed = false" + if query: + q += f" and {query}" + + # List files in the folder + results = [] + page_token = None + + while True: + try: + response = service.files().list( + q=q, + pageSize=page_size, + spaces='drive', + fields='nextPageToken, files(id, name, mimeType, size, modifiedTime, webViewLink)', + pageToken=page_token + ).execute() + + results.extend(response.get('files', [])) + page_token = response.get('nextPageToken') + + if not page_token: + break + except Exception as e: + logger.error(f"Failed to list files: {str(e)}") + raise Exception(f"Failed to list files: {str(e)}") + + return results + + async def get_file(self, user_id, file_id): + """ + Get a file's metadata from Google Drive. + + Args: + user_id: The user's ID + file_id: ID of the file + + Returns: + dict: The file metadata + """ + service = await self._get_drive_service(user_id) + + try: + file = service.files().get( + fileId=file_id, + fields='id, name, mimeType, size, modifiedTime, webViewLink, webContentLink' + ).execute() + + return file + except Exception as e: + logger.error(f"Failed to get file: {str(e)}") + raise Exception(f"Failed to get file: {str(e)}") + + async def download_file(self, user_id, file_id, local_path=None): + """ + Download a file from Google Drive. + + Args: + user_id: The user's ID + file_id: ID of the file to download + local_path: Path where to save the file locally (if None, returns file content as bytes) + + Returns: + bytes or None: File content as bytes if local_path is None, otherwise None + """ + service = await self._get_drive_service(user_id) + + try: + # Get file metadata to get file name if not provided + file_metadata = service.files().get(fileId=file_id, fields='name').execute() + + # Create request to download file + request = service.files().get_media(fileId=file_id) + + if local_path: + # If path is a directory, append the file name + if os.path.isdir(local_path): + local_path = os.path.join(local_path, file_metadata['name']) + + # Download to file + with open(local_path, 'wb') as f: + downloader = MediaIoBaseDownload(f, request) + done = False + while not done: + status, done = downloader.next_chunk() + logger.info(f"Download progress: {int(status.progress() * 100)}%") + return None + else: + # Download to memory + from io import BytesIO + fh = BytesIO() + downloader = MediaIoBaseDownload(fh, request) + done = False + while not done: + status, done = downloader.next_chunk() + logger.info(f"Download progress: {int(status.progress() * 100)}%") + fh.seek(0) + return fh.read() + except Exception as e: + logger.error(f"Failed to download file: {str(e)}") + raise Exception(f"Failed to download file: {str(e)}") + + async def move_file(self, user_id, file_id, new_parent_folder_id): + """ + Move a file to a different folder in Google Drive. + + Args: + user_id: The user's ID + file_id: ID of the file to move + new_parent_folder_id: ID of the destination folder + + Returns: + dict: The updated file metadata + """ + service = await self._get_drive_service(user_id) + + try: + # Get current parents + file = service.files().get(fileId=file_id, fields='parents').execute() + previous_parents = ",".join(file.get('parents', [])) + + # Move the file to the new folder + file = service.files().update( + fileId=file_id, + addParents=new_parent_folder_id, + removeParents=previous_parents, + fields='id, parents' + ).execute() + + return file + except Exception as e: + logger.error(f"Failed to move file: {str(e)}") + raise Exception(f"Failed to move file: {str(e)}") + + async def search_files(self, user_id, query, max_results=10): + """ + Search for files in Google Drive. + + Args: + user_id: The user's ID + query: Search query + max_results: Maximum number of results to return + + Returns: + list: Search results + """ + logger.info(f"Searching for files with query '{query}' for user {user_id}") + + service = await self._get_drive_service(user_id) + + # Sanitize the query to prevent injection + query = query.replace("'", "\\'") + + try: + results = [] + page_token = None + + while True: + logger.info(f"Making API request to search files with query: name contains '{query}'") + response = service.files().list( + q=f"name contains '{query}' and trashed = false", + spaces='drive', + fields='nextPageToken, files(id, name, mimeType, size, modifiedTime, webViewLink)', + pageSize=max_results, + pageToken=page_token + ).execute() + + files_found = response.get('files', []) + logger.info(f"API returned {len(files_found)} files for this page") + + results.extend(files_found) + page_token = response.get('nextPageToken') + + if not page_token or len(results) >= max_results: + break + + logger.info(f"Total files found: {len(results)}") + for i, file in enumerate(results[:5]): + logger.info(f"File {i+1}: {file.get('name', 'Unknown')} (ID: {file.get('id', 'Unknown')})") + + return results[:max_results] + except Exception as e: + logger.error(f"Failed to search files: {str(e)}", exc_info=True) + raise Exception(f"Failed to search files: {str(e)}") + + async def search_files_content(self, user_id, query, max_results=10, mime_type=None): + """ + Search for files with both titles and content that match the query. + + Args: + user_id: The user's ID + query: Search query + max_results: Maximum number of results to return + mime_type: Optional MIME type filter (e.g., 'application/vnd.google-apps.document') + + Returns: + list: Search results + """ + service = await self._get_drive_service(user_id) + + # Sanitize the query to prevent injection + query = query.replace("'", "\\'") + + # Build the query string + q = f"fullText contains '{query}' and trashed = false" + if mime_type: + q += f" and mimeType='{mime_type}'" + + try: + results = [] + page_token = None + + while True: + response = service.files().list( + q=q, + spaces='drive', + fields='nextPageToken, files(id, name, mimeType, size, modifiedTime, webViewLink)', + pageSize=max_results, + pageToken=page_token + ).execute() + + results.extend(response.get('files', [])) + page_token = response.get('nextPageToken') + + if not page_token or len(results) >= max_results: + break + + return results[:max_results] + except Exception as e: + logger.error(f"Failed to search files content: {str(e)}") + raise Exception(f"Failed to search files content: {str(e)}") + + async def search_google_docs(self, user_id, query, max_results=10): + """ + Search specifically for Google Docs. + + Args: + user_id: The user's ID + query: Search query + max_results: Maximum number of results to return + + Returns: + list: Search results for Google Docs + """ + return await self.search_files_content( + user_id, + query, + max_results, + mime_type='application/vnd.google-apps.document' + ) + + async def search_google_forms(self, user_id, query, max_results=10): + """ + Search specifically for Google Forms. + + Args: + user_id: The user's ID + query: Search query + max_results: Maximum number of results to return + + Returns: + list: Search results for Google Forms + """ + return await self.search_files_content( + user_id, + query, + max_results, + mime_type='application/vnd.google-apps.form' + ) + + async def search_google_sheets(self, user_id, query, max_results=10): + """ + Search specifically for Google Sheets. + + Args: + user_id: The user's ID + query: Search query + max_results: Maximum number of results to return + + Returns: + list: Search results for Google Sheets + """ + return await self.search_files_content( + user_id, + query, + max_results, + mime_type='application/vnd.google-apps.spreadsheet' + ) + + async def share_file(self, user_id, file_id, email, role): + """ + Share a file with another user. + + Args: + user_id: The user's ID + file_id: ID of the file to share + email: Email of the user to share with + role: Role to assign (reader, writer, commenter) + + Returns: + dict: The created permission + """ + service = await self._get_drive_service(user_id) + + # Define the permission + permission = { + 'type': 'user', + 'role': role, + 'emailAddress': email + } + + try: + # Create the permission + result = service.permissions().create( + fileId=file_id, + body=permission, + fields='id' + ).execute() + + return result + except Exception as e: + logger.error(f"Failed to share file: {str(e)}") + raise Exception(f"Failed to share file: {str(e)}") + + async def get_document_comments(self, user_id, document_id): + """ + Get comments for a Google Doc. + + Args: + user_id: The user's ID + document_id: ID of the document + + Returns: + list: Comments on the document + """ + service = await self._get_drive_service(user_id) + + try: + # Get comments for the document + result = service.comments().list( + fileId=document_id, + fields='comments(id, content, anchor, htmlContent, quotedFileContent)' + ).execute() + + # Format the comments similar to the C# implementation + formatted_comments = [] + for comment in result.get('comments', []): + formatted_comment = { + 'comment_id': comment.get('id'), + 'content': comment.get('content'), + } + + # Add quoted file content if available + quoted_content = comment.get('quotedFileContent') + if quoted_content: + formatted_comment['quoted_file_content'] = { + 'mime_type': quoted_content.get('mimeType'), + 'value': quoted_content.get('value') + } + + formatted_comments.append(formatted_comment) + + return formatted_comments + except Exception as e: + logger.error(f"Failed to get document comments: {str(e)}") + raise Exception(f"Failed to get document comments: {str(e)}") + + async def copy_document(self, user_id, source_file_id, new_title): + """ + Create a copy of a document. + + Args: + user_id: The user's ID + source_file_id: ID of the document to copy + new_title: Title for the new document + + Returns: + str: The ID of the new document + """ + service = await self._get_drive_service(user_id) + + try: + # Copy the document + body = {'name': new_title} + file = service.files().copy( + fileId=source_file_id, + body=body, + fields='id, webViewLink' + ).execute() + + return file['id'] + except Exception as e: + logger.error(f"Failed to copy document: {str(e)}") + raise Exception(f"Failed to copy document: {str(e)}") + + async def get_folders(self, user_id, parent_folder_id=None): + """ + Get folders in Google Drive. + + Args: + user_id: The user's ID + parent_folder_id: Optional parent folder ID to list folders within + + Returns: + list: Folders + """ + service = await self._get_drive_service(user_id) + + # Build the query string + q = "mimeType='application/vnd.google-apps.folder' and trashed = false" + if parent_folder_id: + q += f" and '{parent_folder_id}' in parents" + + try: + results = [] + page_token = None + + while True: + response = service.files().list( + q=q, + spaces='drive', + fields='nextPageToken, files(id, name, createdTime)', + pageToken=page_token + ).execute() + + results.extend(response.get('files', [])) + page_token = response.get('nextPageToken') + + if not page_token: + break + + return results + except Exception as e: + logger.error(f"Failed to get folders: {str(e)}") + raise Exception(f"Failed to get folders: {str(e)}") + + async def _store_token(self, user_id, access_token, refresh_token, expires_in): + """ + Store a token in the token storage. + + Args: + user_id: The user's ID + access_token: The access token + refresh_token: The refresh token + expires_in: Expiration time in seconds + """ + token_data = { + "access_token": access_token, + "refresh_token": refresh_token, + "expires_at": (datetime.utcnow() + timedelta(seconds=expires_in)).timestamp() + } + + # Serialize and encrypt the token data + serialized_token = json.dumps(token_data) + encrypted_token = TokenEncryptionHelper.encrypt_token(serialized_token, self.encryption_key) + + # Store in the token storage using the helper function + token_record = create_token_record(encrypted_token) + + self.token_storage.store_token(user_id, PLATFORM, SERVICE, token_record) + + async def _get_token_data(self, user_id): + """ + Get token data from storage. + + Args: + user_id: The user's ID + + Returns: + dict: The token data or None if not found + """ + token_record = self.token_storage.get_token(user_id, PLATFORM, SERVICE) + + if not token_record or not token_record.get("is_active") or token_record.get("is_revoked"): + logger.info(f"No valid token found in the storage for user {user_id}") + return None + + try: + encrypted_token = token_record.get("encrypted_token") + if not encrypted_token: + return None + + decrypted_token = TokenEncryptionHelper.decrypt_token(encrypted_token, self.encryption_key) + token_data = json.loads(decrypted_token) + + return token_data + except Exception as e: + logger.error(f"Error getting token data: {str(e)}") + return None + + async def _load_token(self, user_id): + """ + Load a token from the token storage. + + Args: + user_id: The user's ID + + Returns: + str: The access token, or None if not found or expired + """ + token_data = await self._get_token_data(user_id) + + if not token_data: + return None + + # Check if token is expired + expires_at = token_data.get("expires_at") + if expires_at and expires_at <= datetime.utcnow().timestamp(): + logger.info(f"Token expired for user {user_id}, attempting to refresh") + refresh_token = token_data.get("refresh_token") + if refresh_token: + try: + return await self._refresh_token(user_id, refresh_token) + except Exception as e: + logger.error(f"Error refreshing token: {str(e)}") + return None + return None + + return token_data.get("access_token") + + async def _refresh_token(self, user_id, refresh_token): + """ + Refresh an expired token. + + Args: + user_id: The user's ID + refresh_token: The refresh token + + Returns: + str: The new access token + """ + if not self.client_id: + raise ValueError("Google Client ID is not set in configuration.") + if not self.client_secret: + raise ValueError("Google Client Secret is not set in configuration.") + + payload = { + "grant_type": "refresh_token", + "refresh_token": refresh_token, + "client_id": self.client_id, + "client_secret": self.client_secret + } + + logger.info(f"Attempting to refresh token for user {user_id}") + response = requests.post(GOOGLE_TOKEN_URL, data=payload) + response_data = response.json() + + if response.status_code == 200 and "access_token" in response_data: + # Store the new token + expires_in = response_data.get("expires_in", 3600) # Default to 1 hour + await self._store_token( + user_id, + response_data["access_token"], + refresh_token, # Keep the existing refresh token if not provided + expires_in + ) + logger.info(f"Successfully refreshed token for user {user_id}") + return response_data["access_token"] + else: + error_msg = response_data.get("error_description", "Unknown error") + logger.error(f"Failed to refresh token: {error_msg}") + # If refresh fails, mark the token as revoked so we don't keep trying + token_record = self.token_storage.get_token(user_id, PLATFORM, SERVICE) + if token_record: + token_record["is_revoked"] = True + self.token_storage.store_token(user_id, PLATFORM, SERVICE, token_record) + raise Exception(f"Failed to refresh token: {error_msg}") + + async def _get_drive_service(self, user_id): + """ + Get an authenticated Google Drive service instance. + + Args: + user_id: The user's ID + + Returns: + Resource: The Drive service instance + """ + token_data = await self._get_token_data(user_id) + + if not token_data: + raise self._create_auth_exception(user_id) + + # Check if token is expired + expires_at = token_data.get("expires_at") + if expires_at and expires_at <= datetime.utcnow().timestamp(): + # Refresh the token + refresh_token = token_data.get("refresh_token") + if not refresh_token: + raise self._create_auth_exception(user_id) + + try: + token_data["access_token"] = await self._refresh_token(user_id, refresh_token) + except Exception: + raise self._create_auth_exception(user_id) + + # Create credentials from token data + expiry = datetime.fromtimestamp(token_data.get("expires_at", 0)) + credentials = Credentials( + token=token_data.get("access_token"), + refresh_token=token_data.get("refresh_token"), + token_uri=GOOGLE_TOKEN_URL, + client_id=self.client_id, + client_secret=self.client_secret, + expiry=expiry + ) + + # Build the Drive service + try: + service = build('drive', 'v3', credentials=credentials) + return service + except Exception as e: + logger.error(f"Failed to build Drive service: {str(e)}") + raise self._create_auth_exception(user_id) + + def _create_auth_exception(self, user_id): + """ + Create an authentication exception with reauthorization instructions. + + Args: + user_id: The user's ID + + Returns: + Exception: With reauthorization instructions + """ + # Don't try to generate an auth URL here, just return the instruction + return Exception( + "Your Google Drive authorization has expired or is invalid. " + "Please use the `!authorize-gdrive` command to reconnect your Google Drive account." + ) + + async def add_comment_to_document(self, user_id, file_id, content, target_text=None, anchor=None): + """ + Add a comment to a Google Doc. + + Args: + user_id: The user's ID + file_id: ID of the document + content: Comment content + target_text: Optional text to anchor the comment to + anchor: Optional anchor object (used instead of target_text if provided) + + Returns: + dict: The created comment + """ + service = await self._get_drive_service(user_id) + + try: + comment = { + 'content': content + } + + # If anchor is provided, use it + if anchor: + comment['anchor'] = anchor + # If target_text is provided, create an anchor for it + elif target_text: + # Get the document content to find the target text + # This requires using the Google Docs API, not just Drive + # We'd need to add the Documents API scope and build a docs service + # For simplicity, we're just using a placeholder here + comment['quotedFileContent'] = { + 'value': target_text + } + + result = service.comments().create( + fileId=file_id, + body=comment, + fields='id, content, anchor' + ).execute() + + return result + except Exception as e: + logger.error(f"Failed to add comment to document: {str(e)}") + raise Exception(f"Failed to add comment to document: {str(e)}") \ No newline at end of file From ea4f7510c6078e47f1a04fd5c6da04a7dd1790b1 Mon Sep 17 00:00:00 2001 From: bless-stanford <107686196+bless-stanford@users.noreply.github.com> Date: Mon, 10 Mar 2025 15:39:48 -0700 Subject: [PATCH 14/19] Update dependencies and project description for cloud workspace integration --- local_env.yml | 7 +++++++ pyproject.toml | 8 ++++++-- 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/local_env.yml b/local_env.yml index 9e619b19..90f09353 100644 --- a/local_env.yml +++ b/local_env.yml @@ -9,3 +9,10 @@ dependencies: - discord-py>=2.4.0 - mistralai>=1.4.0 - python-dotenv>=1.0.1 + - fastapi>=0.95.0 + - uvicorn>=0.21.0 + - cryptography>=40.0.0 + - requests>=2.28.0 + - semantic-kernel>=0.4.0 + - google-auth-oauthlib>=1.0.0 + - google-api-python-client>=2.70.0 \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index cf699405..0e1a1181 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "ai-agent" version = "0.1.0" -description = "Add your description here" +description = "Discord bot for cloud workspace integration" readme = "README.md" requires-python = ">=3.13" dependencies = [ @@ -12,4 +12,8 @@ dependencies = [ "fastapi>=0.95.0", "uvicorn>=0.21.0", "cryptography>=40.0.0", -] + "requests>=2.28.0", + "semantic-kernel>=0.4.0", + "google-auth-oauthlib>=1.0.0", + "google-api-python-client>=2.70.0", +] \ No newline at end of file From f0c565445ad89e333d71b352fddf6a464748669f Mon Sep 17 00:00:00 2001 From: bless-stanford <107686196+bless-stanford@users.noreply.github.com> Date: Mon, 10 Mar 2025 15:42:08 -0700 Subject: [PATCH 15/19] Update dependencies to specific versions and add new packages for enhanced functionality --- local_env.yml | 30 +++++++++++++++++++----------- pyproject.toml | 26 +++++++++++++++----------- 2 files changed, 34 insertions(+), 22 deletions(-) diff --git a/local_env.yml b/local_env.yml index 90f09353..0521c8e9 100644 --- a/local_env.yml +++ b/local_env.yml @@ -1,18 +1,26 @@ name: discord_bot channels: - defaults + - conda-forge dependencies: - python>=3.13 - pip - pip: - - audioop-lts>=0.2.1 - - discord-py>=2.4.0 - - mistralai>=1.4.0 - - python-dotenv>=1.0.1 - - fastapi>=0.95.0 - - uvicorn>=0.21.0 - - cryptography>=40.0.0 - - requests>=2.28.0 - - semantic-kernel>=0.4.0 - - google-auth-oauthlib>=1.0.0 - - google-api-python-client>=2.70.0 \ No newline at end of file + - audioop-lts==0.2.1 + - discord-py==2.4.0 + - mistralai==1.4.0 + - python-dotenv==1.0.0 + - fastapi==0.111.0 + - uvicorn==0.27.1 + - cryptography==41.0.8 + - requests==2.32.3 + - semantic-kernel==1.23.1 + - google-auth-oauthlib==1.2.1 + - google-api-python-client==2.163.0 + - aiohttp==3.11.11 + - httpx==0.27.0 + - pydantic==2.6.3 + - pydantic-settings==2.7.1 + - protobuf==5.29.3 + - pyOpenSSL==25.0.0 + - azure-identity==1.19.0 \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 0e1a1181..2419e78a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,15 +5,19 @@ description = "Discord bot for cloud workspace integration" readme = "README.md" requires-python = ">=3.13" dependencies = [ - "audioop-lts>=0.2.1", - "discord-py>=2.4.0", - "mistralai>=1.4.0", - "python-dotenv>=1.0.1", - "fastapi>=0.95.0", - "uvicorn>=0.21.0", - "cryptography>=40.0.0", - "requests>=2.28.0", - "semantic-kernel>=0.4.0", - "google-auth-oauthlib>=1.0.0", - "google-api-python-client>=2.70.0", + "audioop-lts==0.2.1", + "discord-py==2.4.0", + "mistralai==1.4.0", + "python-dotenv==1.0.0", + "fastapi==0.111.0", + "uvicorn==0.27.1", + "cryptography==41.0.8", + "requests==2.32.3", + "semantic-kernel==1.23.1", + "google-auth-oauthlib==1.2.1", + "google-api-python-client==2.163.0", + "aiohttp==3.11.11", + "httpx==0.27.0", + "pydantic==2.6.3", + "pydantic-settings==2.7.1", ] \ No newline at end of file From b6933f371bde0d868d089321764e30fd87bd1bd6 Mon Sep 17 00:00:00 2001 From: bless-stanford <107686196+bless-stanford@users.noreply.github.com> Date: Mon, 10 Mar 2025 20:50:38 -0700 Subject: [PATCH 16/19] Add requirements.txt with specified package versions for project dependencies --- requirements.txt | 15 +++++++++++++++ 1 file changed, 15 insertions(+) create mode 100644 requirements.txt diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 00000000..74c56b85 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,15 @@ +audioop-lts==0.2.1 +discord.py==2.4.0 +mistralai==1.4.0 +python-dotenv==1.0.0 +fastapi==0.111.0 +uvicorn==0.27.1 +cryptography==41.0.8 +requests==2.32.3 +semantic-kernel==1.23.1 +google-auth-oauthlib==1.2.1 +google-api-python-client==2.163.0 +aiohttp==3.11.11 +httpx==0.27.0 +pydantic==2.6.3 +pydantic-settings==2.7.1 \ No newline at end of file From a58d34c6d68c9ae3d8a323bfb53a8ba8dcd5608f Mon Sep 17 00:00:00 2001 From: christinaduong Date: Tue, 11 Mar 2025 00:24:57 -0700 Subject: [PATCH 17/19] initialize gcal support --- plugins/google_calendar_plugin.py | 0 requirements.txt | 10 ++++++---- services/google_calendar_service.py | 0 3 files changed, 6 insertions(+), 4 deletions(-) create mode 100644 plugins/google_calendar_plugin.py create mode 100644 services/google_calendar_service.py diff --git a/plugins/google_calendar_plugin.py b/plugins/google_calendar_plugin.py new file mode 100644 index 00000000..e69de29b diff --git a/requirements.txt b/requirements.txt index 74c56b85..5221b9c1 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,15 +1,17 @@ -audioop-lts==0.2.1 +# audioop-lts==0.2.1 discord.py==2.4.0 mistralai==1.4.0 python-dotenv==1.0.0 fastapi==0.111.0 uvicorn==0.27.1 -cryptography==41.0.8 +# cryptography==41.0.8 requests==2.32.3 semantic-kernel==1.23.1 google-auth-oauthlib==1.2.1 google-api-python-client==2.163.0 aiohttp==3.11.11 httpx==0.27.0 -pydantic==2.6.3 -pydantic-settings==2.7.1 \ No newline at end of file +# pydantic==2.6.3 +pydantic-settings==2.7.1 +cryptography>=41.0.8 +pydantic>=2.9.0 diff --git a/services/google_calendar_service.py b/services/google_calendar_service.py new file mode 100644 index 00000000..e69de29b From 570b45da73bff6e903740a3c193ab757f70e7203 Mon Sep 17 00:00:00 2001 From: christinaduong Date: Tue, 11 Mar 2025 01:06:02 -0700 Subject: [PATCH 18/19] add gcal --- bot.py | 23 ++ local_env.yml | 1 + plugins/cloud_plugin_manager.py | 23 +- plugins/google_calendar_plugin.py | 602 +++++++++++++++++++++++++++ services/google_calendar_service.py | 621 ++++++++++++++++++++++++++++ 5 files changed, 1269 insertions(+), 1 deletion(-) diff --git a/bot.py b/bot.py index 56b2a5bc..33809a45 100644 --- a/bot.py +++ b/bot.py @@ -7,7 +7,9 @@ from services.box_service import BoxService from services.dropbox_service import DropboxService from services.google_drive_service import GoogleDriveService +from services.google_calendar_service import GoogleCalendarService from server import start_server +from datetime import datetime, timedelta PREFIX = "!" @@ -34,6 +36,7 @@ box_service = BoxService() dropbox_service = DropboxService() google_drive_service = GoogleDriveService() +google_calendar_service = GoogleCalendarService() async def send_split_message(message: discord.Message, response: str | list[str]): """ @@ -276,6 +279,26 @@ async def gdrive_upload(ctx): if os.path.exists(file_path): os.remove(file_path) +@bot.command(name="gcalendar-create", help="Create a new calendar") +async def gcalendar_create(ctx, *, calendar_name=None): + """ + Creates a new calendar in the user's Google Calendar account. + + Usage: !gcalendar-create My New Calendar + """ + if not calendar_name: + await ctx.send("Please provide a name for the calendar.\nUsage: `!gcalendar-create My New Calendar`") + return + + try: + calendar_id = await google_calendar_service.create_calendar(str(ctx.author.id), calendar_name) + + await ctx.send(f"Calendar created successfully!\nName: {calendar_name}\nID: {calendar_id}") + except Exception as e: + error_msg = f"Error creating calendar: {str(e)}" + logger.error(error_msg) + await ctx.send(error_msg[:1900]) + @bot.command(name="cloud-status", help="Check your cloud service connections") async def cloud_status(ctx): """ diff --git a/local_env.yml b/local_env.yml index 0521c8e9..e0ae5c60 100644 --- a/local_env.yml +++ b/local_env.yml @@ -1,3 +1,4 @@ + name: discord_bot channels: - defaults diff --git a/plugins/cloud_plugin_manager.py b/plugins/cloud_plugin_manager.py index c3ca1937..02fc1c2a 100644 --- a/plugins/cloud_plugin_manager.py +++ b/plugins/cloud_plugin_manager.py @@ -2,9 +2,11 @@ from services.box_service import BoxService from services.dropbox_service import DropboxService from services.google_drive_service import GoogleDriveService +from services.google_calendar_service import GoogleCalendarService from plugins.box_plugin import BoxPlugins from plugins.dropbox_plugin import DropboxPlugins from plugins.google_drive_plugin import GoogleDrivePlugins +from plugins.google_calendar_plugin import GoogleCalendarPlugins import logging logger = logging.getLogger("cloud_plugin_manager") @@ -15,7 +17,7 @@ class CloudPluginManager: Consolidates Box, Dropbox, and Google Drive plugins into a single interface. """ - def __init__(self, box_service=None, dropbox_service=None, google_drive_service=None): + def __init__(self, box_service=None, dropbox_service=None, google_drive_service=None, google_calendar_service=None): """ Initialize the cloud plugin manager with service instances. If no services are provided, new ones will be created. @@ -24,15 +26,19 @@ def __init__(self, box_service=None, dropbox_service=None, google_drive_service= box_service: BoxService instance or None dropbox_service: DropboxService instance or None google_drive_service: GoogleDriveService instance or None + google_calendar_service: GoogleCalendarService instance or None """ self.box_service = box_service or BoxService() self.dropbox_service = dropbox_service or DropboxService() self.google_drive_service = google_drive_service or GoogleDriveService() + self.google_calendar_service = google_calendar_service or GoogleCalendarService() + # Initialize plugin instances self.box_plugins = BoxPlugins(self.box_service) self.dropbox_plugins = DropboxPlugins(self.dropbox_service) self.google_drive_plugins = GoogleDrivePlugins(self.google_drive_service) + self.google_calendar_plugins = GoogleCalendarPlugins(self.google_calendar_service) def register_plugins(self, kernel: Kernel) -> Kernel: """ @@ -56,6 +62,10 @@ def register_plugins(self, kernel: Kernel) -> Kernel: # Register Google Drive plugins kernel.add_plugin(self.google_drive_plugins, "gdrive") logger.info("Google Drive plugins registered with kernel") + + # Register Google Calendar plugins + kernel.add_plugin(self.google_calendar_plugins, "gcalendar") + logger.info("Google Calendar plugins registered with kernel") return kernel except Exception as e: @@ -103,6 +113,17 @@ def get_plugin_descriptions(self) -> str: descriptions += "- `gdrive.share_file`: Share a Google Drive file with another user\n" descriptions += "- `gdrive.move_file`: Move a file to a different folder in Google Drive\n" + # Google Calendar plugins + descriptions += "## Google Calendar Plugins\n" + descriptions += "Use these to interact with your Google Calendar account:\n" + descriptions += "- `gcalendar.create_calendar`: Create a new calendar in Google Calendar\n" + descriptions += "- `gcalendar.add_event`: Add a new event to your calendar\n" + descriptions += "- `gcalendar.delete_event`: Delete an event from your calendar\n" + descriptions += "- `gcalendar.get_event`: Get details of a specific event\n" + descriptions += "- `gcalendar.get_events`: Get events within a date range\n" + descriptions += "- `gcalendar.update_event`: Update an existing event\n" + descriptions += "- `gcalendar.share_event`: Share an event with another user\n" + return descriptions def update_user_context(self, kernel: Kernel, user_id: str) -> None: diff --git a/plugins/google_calendar_plugin.py b/plugins/google_calendar_plugin.py index e69de29b..2852b8e0 100644 --- a/plugins/google_calendar_plugin.py +++ b/plugins/google_calendar_plugin.py @@ -0,0 +1,602 @@ +import json +import os +from datetime import datetime, timedelta +from semantic_kernel.functions import kernel_function +from semantic_kernel.functions.kernel_function_from_prompt import KernelFunctionFromPrompt +from services.google_calendar_service import GoogleCalendarService +import logging + +logger = logging.getLogger("google_calendar_plugins") + +class GoogleCalendarPlugins: + """ + Plugins for interacting with Google Calendar. + """ + + def __init__(self, calendar_service=None): + """ + Initialize the Google Calendar plugins with a GoogleCalendarService. + If no service is provided, a new one will be created. + """ + self.calendar_service = calendar_service or GoogleCalendarService() + + @kernel_function( + name="create_calendar", + description="Creates a new calendar in the user's Google Calendar account" + ) + async def create_calendar( + self, + calendar_name: str, + user_id: str = None, + kernel = None + ) -> str: + """ + Creates a new calendar in the user's Google Calendar account. + + Args: + calendar_name: Name of the calendar to create + user_id: The user's ID (automatically provided) + + Returns: + str: Success message with calendar details or error message + """ + try: + if not user_id: + return "Error: User ID not available. Please try again later." + + calendar_id = await self.calendar_service.create_calendar(user_id, calendar_name) + + if calendar_id: + return f"Calendar '{calendar_name}' created successfully with ID: {calendar_id}." + else: + return f"Failed to create calendar '{calendar_name}'." + + except Exception as e: + logger.error(f"Error creating calendar: {str(e)}") + return f"An error occurred while creating the calendar: {str(e)}" + + @kernel_function( + name="add_event", + description="Adds an event to the user's Google Calendar" + ) + async def add_event( + self, + summary: str, + description: str, + start_date_time: str, + end_date_time: str, + location: str = "", + is_all_day: str = "false", + attendee_emails: str = "", + user_id: str = None, + kernel = None + ) -> str: + """ + Adds an event to the user's Google Calendar. + + Args: + summary: Event summary/title + description: Event description + start_date_time: Start time (ISO 8601 - YYYY-MM-DDTHH:MM:SS±hh:mm) + end_date_time: End time (ISO 8601 - YYYY-MM-DDTHH:MM:SS±hh:mm) + location: Event location (optional) + is_all_day: Whether this is an all-day event (default: "false") + attendee_emails: Comma-separated list of attendee emails + user_id: The user's ID (automatically provided) + + Returns: + str: Success message with event details or error message + """ + try: + if not user_id: + return "Error: User ID not available. Please try again later." + + # Parse datetime strings and is_all_day + try: + start_dt = datetime.fromisoformat(start_date_time.replace('Z', '+00:00')) + end_dt = datetime.fromisoformat(end_date_time.replace('Z', '+00:00')) + is_all_day_event = is_all_day.lower() == "true" + except ValueError as e: + return f"Error parsing date: {str(e)}" + + # Parse attendee emails + attendees = [] + if attendee_emails: + attendees = [{"email": email.strip()} for email in attendee_emails.split(",") if email.strip()] + + # Fetch user's timezone + try: + timezone_info = await self.calendar_service.get_user_timezone(user_id) + user_timezone = timezone_info.get("timezone", "UTC") + except Exception: + user_timezone = "UTC" + + # Create event dictionary + event = { + "summary": summary, + "description": description, + "location": location, + "start": {}, + "end": {}, + "attendees": attendees + } + + # Set start and end dates/times + if is_all_day_event: + event["start"] = {"date": start_dt.strftime("%Y-%m-%d")} + event["end"] = {"date": end_dt.strftime("%Y-%m-%d")} + else: + event["start"] = { + "dateTime": start_dt.isoformat(), + "timeZone": user_timezone + } + event["end"] = { + "dateTime": end_dt.isoformat(), + "timeZone": user_timezone + } + + # Add the event + added_event = await self.calendar_service.add_event(user_id, event) + + if added_event: + event_link = "https://calendar.google.com/calendar/u/0/r" + return f"Event added: {added_event.get('summary')}\nEvent link: {event_link}" + else: + return "Failed to add the event." + + except Exception as e: + logger.error(f"Error adding event: {str(e)}") + return f"An error occurred while adding the event: {str(e)}" + + @kernel_function( + name="delete_event", + description="Deletes an event from the user's Google Calendar" + ) + async def delete_event( + self, + event_id_or_query: str, + user_id: str = None, + kernel = None + ) -> str: + """ + Deletes an event from the user's Google Calendar. + + Args: + event_id_or_query: Event ID or search query + user_id: The user's ID (automatically provided) + + Returns: + str: Success message or error message + """ + try: + if not user_id: + return "Error: User ID not available. Please try again later." + + # First try to delete assuming event_id_or_query is an event ID + try: + await self.calendar_service.delete_event(user_id, event_id_or_query) + return f"Event deleted successfully. ID: {event_id_or_query}" + except Exception: + # If deletion fails, assume it's not an ID and proceed with search + pass + + # Search for events matching the query + search_results = await self.calendar_service.search_events(user_id, event_id_or_query) + events = search_results.get("items", []) + + if not events: + return f"No events found matching '{event_id_or_query}'." + + if len(events) == 1: + event = events[0] + await self.calendar_service.delete_event(user_id, event["id"]) + return f"Event deleted: {event.get('summary')} (ID: {event.get('id')})" + + # If multiple events and kernel is provided, find most relevant + if kernel and len(events) > 1: + most_relevant_event = await self._find_most_relevant_event(kernel, events, event_id_or_query) + + if most_relevant_event: + await self.calendar_service.delete_event(user_id, most_relevant_event["id"]) + return f"Event deleted: {most_relevant_event.get('summary')} (ID: {most_relevant_event.get('id')})" + + # If multiple events and no most relevant found, return summary + return "Multiple events found. Please be more specific:\n" + self._create_search_results_summary(events) + + except Exception as e: + logger.error(f"Error deleting event: {str(e)}") + return f"An error occurred while deleting the event: {str(e)}" + + @kernel_function( + name="get_event", + description="Gets details of a specific event from the user's Google Calendar" + ) + async def get_event( + self, + search_query: str, + user_id: str = None, + kernel = None + ) -> str: + """ + Gets details of a specific event from the user's Google Calendar. + + Args: + search_query: Search query or event ID + user_id: The user's ID (automatically provided) + + Returns: + str: Event details or error message + """ + try: + if not user_id: + return "Error: User ID not available. Please try again later." + + # First try to get the event directly (in case search_query is an event ID) + try: + event = await self.calendar_service.get_event(user_id, search_query) + return json.dumps(event, indent=2) + except Exception: + # If getting the event fails, assume it's not an ID and proceed with search + pass + + # Search for events matching the query + search_results = await self.calendar_service.search_events(user_id, search_query) + events = search_results.get("items", []) + + if not events: + return f"No events found matching '{search_query}'." + + if len(events) == 1: + return json.dumps(events[0], indent=2) + + # If multiple events and kernel is provided, find most relevant + if kernel and len(events) > 1: + most_relevant_event = await self._find_most_relevant_event(kernel, events, search_query) + + if most_relevant_event: + return json.dumps(most_relevant_event, indent=2) + + # If multiple events and no most relevant found, return summary + return "Multiple events found. Here's a summary:\n" + self._create_search_results_summary(events) + + except Exception as e: + logger.error(f"Error getting event: {str(e)}") + return f"An error occurred while getting the event: {str(e)}" + + @kernel_function( + name="get_events", + description="Gets events from the user's Google Calendar within a date range" + ) + async def get_events( + self, + start_date: str, + end_date: str, + max_results: int = 10, + user_id: str = None, + kernel = None + ) -> str: + """ + Gets events from the user's Google Calendar within a date range. + + Args: + start_date: Start date (YYYY-MM-DD) + end_date: End date (YYYY-MM-DD) + max_results: Maximum number of events to return + user_id: The user's ID (automatically provided) + + Returns: + str: List of events or error message + """ + try: + if not user_id: + return "Error: User ID not available. Please try again later." + + try: + start_dt = datetime.fromisoformat(start_date) + end_dt = datetime.fromisoformat(end_date) + except ValueError: + # Try to parse as just date + try: + start_dt = datetime.strptime(start_date, "%Y-%m-%d") + end_dt = datetime.strptime(end_date, "%Y-%m-%d") + except ValueError as e: + return f"Error parsing date: {str(e)}" + + events = await self.calendar_service.get_events(user_id, start_dt, end_dt, max_results) + + if not events: + return f"No events found in the date range {start_date} to {end_date}." + + return self._format_events(events) + + except Exception as e: + logger.error(f"Error getting events: {str(e)}") + return f"An error occurred while getting events: {str(e)}" + + @kernel_function( + name="update_event", + description="Updates an existing event in the user's Google Calendar" + ) + async def update_event( + self, + event_id_or_query: str, + summary: str, + description: str, + start_date_time: str, + end_date_time: str, + user_id: str = None, + kernel = None + ) -> str: + """ + Updates an existing event in the user's Google Calendar. + + Args: + event_id_or_query: Event ID or search query + summary: New event summary/title + description: New event description + start_date_time: New start time (ISO 8601) + end_date_time: New end time (ISO 8601) + user_id: The user's ID (automatically provided) + + Returns: + str: Success message or error message + """ + try: + if not user_id: + return "Error: User ID not available. Please try again later." + + # Parse datetime strings + try: + start_dt = datetime.fromisoformat(start_date_time.replace('Z', '+00:00')) + end_dt = datetime.fromisoformat(end_date_time.replace('Z', '+00:00')) + except ValueError as e: + return f"Error parsing date: {str(e)}" + + # Create updated event + updated_event = { + "summary": summary, + "description": description, + "start": { + "dateTime": start_dt.isoformat(), + "timeZone": "UTC" # Default to UTC + }, + "end": { + "dateTime": end_dt.isoformat(), + "timeZone": "UTC" # Default to UTC + } + } + + # First try to update assuming event_id_or_query is an event ID + try: + result = await self.calendar_service.update_event(user_id, event_id_or_query, updated_event) + event_link = f"https://www.google.com/calendar/event?eid={result.get('id')}" + return f"Event updated: {result.get('summary')} (ID: {result.get('id')})\nEvent link: {event_link}" + except Exception: + # If update fails, assume it's not an ID and proceed with search + pass + + # Search for events matching the query + search_results = await self.calendar_service.search_events(user_id, event_id_or_query) + events = search_results.get("items", []) + + if not events: + return f"No events found matching '{event_id_or_query}'." + + if len(events) == 1: + event = events[0] + result = await self.calendar_service.update_event(user_id, event["id"], updated_event) + event_link = f"https://www.google.com/calendar/event?eid={result.get('id')}" + return f"Event updated: {result.get('summary')} (ID: {result.get('id')})\nEvent link: {event_link}" + + # If multiple events and kernel is provided, find most relevant + if kernel and len(events) > 1: + most_relevant_event = await self._find_most_relevant_event(kernel, events, event_id_or_query) + + if most_relevant_event: + result = await self.calendar_service.update_event(user_id, most_relevant_event["id"], updated_event) + event_link = f"https://www.google.com/calendar/event?eid={result.get('id')}" + return f"Event updated: {result.get('summary')} (ID: {result.get('id')})\nEvent link: {event_link}" + + # If multiple events and no most relevant found, return summary + return "Multiple events found. Please be more specific:\n" + self._create_search_results_summary(events) + + except Exception as e: + logger.error(f"Error updating event: {str(e)}") + return f"An error occurred while updating the event: {str(e)}" + + @kernel_function( + name="share_event", + description="Shares an event with another user by adding them as an attendee" + ) + async def share_event( + self, + event_id_or_query: str, + shared_email: str, + user_id: str = None, + kernel = None + ) -> str: + """ + Shares an event with another user by adding them as an attendee. + + Args: + event_id_or_query: Event ID or search query + shared_email: Email of the user to share with + user_id: The user's ID (automatically provided) + + Returns: + str: Success message or error message + """ + try: + if not user_id: + return "Error: User ID not available. Please try again later." + + # First try to share assuming event_id_or_query is an event ID + try: + await self.calendar_service.share_event(user_id, event_id_or_query, shared_email) + event_link = f"https://www.google.com/calendar/event?eid={event_id_or_query}" + return f"Event (ID: {event_id_or_query}) successfully shared with {shared_email}.\nEvent link: {event_link}" + except Exception: + # If sharing fails, assume it's not an ID and proceed with search + pass + + # Search for events matching the query + search_results = await self.calendar_service.search_events(user_id, event_id_or_query) + events = search_results.get("items", []) + + if not events: + return f"No events found matching '{event_id_or_query}'." + + if len(events) == 1: + event = events[0] + await self.calendar_service.share_event(user_id, event["id"], shared_email) + event_link = f"https://www.google.com/calendar/event?eid={event.get('id')}" + return f"Event shared: {event.get('summary')} (ID: {event.get('id')}) with {shared_email}\nEvent link: {event_link}" + + # If multiple events and kernel is provided, find most relevant + if kernel and len(events) > 1: + most_relevant_event = await self._find_most_relevant_event(kernel, events, event_id_or_query) + + if most_relevant_event: + await self.calendar_service.share_event(user_id, most_relevant_event["id"], shared_email) + event_link = f"https://www.google.com/calendar/event?eid={most_relevant_event.get('id')}" + return f"Event shared: {most_relevant_event.get('summary')} (ID: {most_relevant_event.get('id')}) with {shared_email}\nEvent link: {event_link}" + + # If multiple events and no most relevant found, return summary + return "Multiple events found. Please be more specific:\n" + self._create_search_results_summary(events) + + except Exception as e: + logger.error(f"Error sharing event: {str(e)}") + return f"An error occurred while sharing the event: {str(e)}" + + async def _find_most_relevant_event(self, kernel, events, user_query): + """ + Find the most relevant event from a list based on user query. + + Args: + kernel: Semantic Kernel instance + events: List of events + user_query: The user's query + + Returns: + dict: The most relevant event or None + """ + try: + # Create a function from prompt + rank_events_function = KernelFunctionFromPrompt( + function_name="RankEventsByRelevance", + plugin_name=None, + prompt="Given the user query: '{{$userQuery}}' and a list of event summaries and descriptions, " + "rank them by relevance and return the index of the most relevant event. " + "Do not add any comments or explanation to the response.\n" + "Event list: {{$eventList}}", + template_format="semantic-kernel" + ) + + # Create event list string + event_list = "\n".join([f"{i}: Summary: {event.get('summary', 'No title')}, Description: {event.get('description', 'No description')}" + for i, event in enumerate(events)]) + + # Create kernel arguments + kernel_arguments = { + "userQuery": user_query, + "eventList": event_list + } + + # Invoke the function + result = await kernel.invoke(rank_events_function, **kernel_arguments) + + # Get the value from the result - might be a list or a string + result_value = result.value + + # Handle different result types + if isinstance(result_value, list) and len(result_value) > 0: + result_text = str(result_value[0]).strip() + elif isinstance(result_value, str): + result_text = result_value.strip() + else: + # Fallback + result_text = str(result_value).strip() + + try: + most_relevant_index = int(result_text) + if 0 <= most_relevant_index < len(events): + return events[most_relevant_index] + except ValueError: + logger.warning(f"Could not parse the relevance index from AI result: {result_text}") + + return None + + except Exception as e: + logger.error(f"Error finding most relevant event: {str(e)}") + return None + + def _format_events(self, events): + """Format a list of events into a human-readable string.""" + if not events: + return "No events found." + + formatted_events = [] + for event in events: + formatted_event = f"Event: {event.get('summary', 'No title')}\n" + + # Format start time + start = event.get('start', {}) + if 'dateTime' in start: + start_time = datetime.fromisoformat(start['dateTime'].replace('Z', '+00:00')) + formatted_event += f"Start: {start_time.strftime('%Y-%m-%d %H:%M:%S')}\n" + elif 'date' in start: + formatted_event += f"Start: {start['date']} (All day)\n" + + # Format end time + end = event.get('end', {}) + if 'dateTime' in end: + end_time = datetime.fromisoformat(end['dateTime'].replace('Z', '+00:00')) + formatted_event += f"End: {end_time.strftime('%Y-%m-%d %H:%M:%S')}\n" + elif 'date' in end: + formatted_event += f"End: {end['date']} (All day)\n" + + # Add location if available + if 'location' in event and event['location']: + formatted_event += f"Location: {event['location']}\n" + + # Add description if available + if 'description' in event and event['description']: + formatted_event += f"Description: {event['description']}\n" + + formatted_events.append(formatted_event) + + return "\n".join(formatted_events) + + def _create_search_results_summary(self, events): + """Create a summary of multiple search results.""" + if not events: + return "No events found." + + summary = [] + for event in events: + event_summary = f"ID: {event.get('id')}\n" + event_summary += f"Summary: {event.get('summary', 'No title')}\n" + + # Format start time + start = event.get('start', {}) + if 'dateTime' in start: + start_time = datetime.fromisoformat(start['dateTime'].replace('Z', '+00:00')) + event_summary += f"Start: {start_time.strftime('%Y-%m-%d %H:%M:%S')}\n" + elif 'date' in start: + event_summary += f"Start: {start['date']} (All day)\n" + + # Format end time + end = event.get('end', {}) + if 'dateTime' in end: + end_time = datetime.fromisoformat(end['dateTime'].replace('Z', '+00:00')) + event_summary += f"End: {end_time.strftime('%Y-%m-%d %H:%M:%S')}\n" + elif 'date' in end: + event_summary += f"End: {end['date']} (All day)\n" + + # Add event link + event_link = f"https://www.google.com/calendar/event?eid={event.get('id')}" + event_summary += f"Link: {event_link}\n" + + summary.append(event_summary) + + return "\n".join(summary) \ No newline at end of file diff --git a/services/google_calendar_service.py b/services/google_calendar_service.py index e69de29b..d750b9ac 100644 --- a/services/google_calendar_service.py +++ b/services/google_calendar_service.py @@ -0,0 +1,621 @@ +import os +import json +import logging +from datetime import datetime, timedelta +from urllib.parse import urlencode +from dotenv import load_dotenv + +import requests +from google.oauth2.credentials import Credentials +from google_auth_oauthlib.flow import Flow +from googleapiclient.discovery import build + +from helpers.token_helpers import ( + TokenEncryptionHelper, + TokenStorageManager, + create_token_record, + load_or_generate_encryption_key +) + +# Setup logging +logger = logging.getLogger("google_calendar_service") + +# Constants for platform and service +PLATFORM = "Google" +SERVICE = "GoogleCalendarService" + +# API URLs +GOOGLE_AUTH_BASE_URL = "https://accounts.google.com/o/oauth2/" +GOOGLE_TOKEN_URL = "https://oauth2.googleapis.com/token" +GOOGLE_REVOKE_URL = "https://oauth2.googleapis.com/revoke" + +# Scopes for Google Calendar API +SCOPES = ['https://www.googleapis.com/auth/calendar'] + + +class GoogleCalendarService: + def __init__(self, config=None): + """ + Initialize the Google Calendar service with configuration. + + Args: + config: Configuration dictionary or None to load from .env + """ + if config is None: + load_dotenv() + self.client_id = os.getenv("GOOGLE_CLIENT_ID") + self.client_secret = os.getenv("GOOGLE_CLIENT_SECRET") + self.redirect_uri = os.getenv("GOOGLE_CALENDAR_REDIRECT_URI") + self.app_name = os.getenv("GOOGLE_APP_NAME", "GoogleCalendarApp") + + # Get or generate encryption key using our helper + self.encryption_key = load_or_generate_encryption_key() + else: + self.client_id = config.get("client_id") + self.client_secret = config.get("client_secret") + self.redirect_uri = config.get("redirect_uri") + self.app_name = config.get("app_name", "GoogleCalendarApp") + self.encryption_key = config.get("encryption_key") + + # Initialize token storage + self.token_storage = TokenStorageManager() + + async def get_authorization_url(self, user_id): + """ + Get the authorization URL for Google OAuth flow. + + Args: + user_id: The user's ID + + Returns: + str: The authorization URL + """ + if not self.client_id: + raise ValueError("Google Client ID is not set in configuration.") + if not self.redirect_uri: + raise ValueError("Google Calendar Redirect URI is not set in configuration.") + + # Encrypt user_id as state parameter + state = TokenEncryptionHelper.encrypt_token(user_id, self.encryption_key) + + # Create a Flow instance + flow = Flow.from_client_config( + { + "web": { + "client_id": self.client_id, + "client_secret": self.client_secret, + "auth_uri": f"{GOOGLE_AUTH_BASE_URL}auth", + "token_uri": GOOGLE_TOKEN_URL, + "redirect_uris": [self.redirect_uri] + } + }, + scopes=SCOPES, + redirect_uri=self.redirect_uri + ) + + # Generate authorization URL + auth_url, _ = flow.authorization_url( + access_type='offline', + include_granted_scopes='true', + state=state, + prompt='consent' # Always show consent screen to get refresh token + ) + + logger.info(f"Generated authorization URL for user {user_id}") + return auth_url + + async def handle_auth_callback(self, state, code): + """ + Handle the authorization callback from Google. + + Args: + state: The state parameter from the callback + code: The authorization code from the callback + """ + if not self.client_id: + raise ValueError("Google Client ID is not set in configuration.") + if not self.client_secret: + raise ValueError("Google Client Secret is not set in configuration.") + if not self.redirect_uri: + raise ValueError("Google Calendar Redirect URI is not set in configuration.") + + # Decrypt the user_id from state + user_id = TokenEncryptionHelper.decrypt_token(state, self.encryption_key) + logger.info(f"Processing authorization callback for user {user_id}") + + # Create a Flow instance - but don't specify scopes this time + # This lets the flow accept whatever scopes Google returns + flow = Flow.from_client_config( + { + "web": { + "client_id": self.client_id, + "client_secret": self.client_secret, + "auth_uri": f"{GOOGLE_AUTH_BASE_URL}auth", + "token_uri": GOOGLE_TOKEN_URL, + "redirect_uris": [self.redirect_uri] + } + }, + scopes=None, # Allow any scope to be returned + redirect_uri=self.redirect_uri + ) + + # Exchange code for token + try: + flow.fetch_token(code=code) + credentials = flow.credentials + + # Store token + await self._store_token( + user_id, + credentials.token, + credentials.refresh_token, + credentials.expiry.timestamp() - datetime.now().timestamp() + ) + logger.info(f"Successfully obtained and stored access token for user {user_id}") + except Exception as e: + logger.error(f"Failed to obtain access token: {str(e)}") + raise Exception(f"Failed to obtain user access token: {str(e)}") + + async def revoke_access(self, user_id): + """ + Revoke the Google Calendar access for a user. + + Args: + user_id: The user's ID + """ + token_data = await self._get_token_data(user_id) + if not token_data: + raise ValueError("No valid token found for user") + + token = token_data.get("access_token") + if not token: + raise ValueError("No valid access token found for user") + + # Revoke the token + params = {'token': token} + response = requests.post(GOOGLE_REVOKE_URL, params=params) + + if response.status_code in (200, 204): + # Delete the token from storage + self.token_storage.delete_token(user_id, PLATFORM, SERVICE) + logger.info(f"Successfully revoked access for user {user_id}") + else: + logger.error(f"Failed to revoke token: {response.status_code}") + raise Exception(f"Failed to revoke token: {response.status_code}") + + async def get_user_timezone(self, user_id): + """ + Get the user's calendar timezone. + + Args: + user_id: The user's ID + + Returns: + dict: Object containing the user's timezone + """ + service = await self._get_calendar_service(user_id) + + try: + calendar = service.calendars().get(calendarId='primary').execute() + return {"timezone": calendar['timeZone']} + except Exception as e: + logger.error(f"Failed to get user timezone: {str(e)}") + raise Exception(f"Failed to get user timezone: {str(e)}") + + async def get_events(self, user_id, start_date, end_date, max_results=10): + """ + Get events from the user's primary calendar. + + Args: + user_id: The user's ID + start_date: Start date for events (datetime) + end_date: End date for events (datetime) + max_results: Maximum number of events to return + + Returns: + list: The calendar events + """ + service = await self._get_calendar_service(user_id) + + try: + # Format dates to RFC3339 timestamp + start_date_rfc = start_date.isoformat() + 'Z' + end_date_rfc = end_date.isoformat() + 'Z' + + events_result = service.events().list( + calendarId='primary', + timeMin=start_date_rfc, + timeMax=end_date_rfc, + maxResults=max_results, + singleEvents=True, + orderBy='startTime' + ).execute() + + events = events_result.get('items', []) + return events + except Exception as e: + logger.error(f"Failed to get events: {str(e)}") + raise Exception(f"Failed to get events: {str(e)}") + + async def add_event(self, user_id, event_details): + """ + Add an event to the user's primary calendar. + + Args: + user_id: The user's ID + event_details: Dictionary with event details + + Returns: + dict: The created event + """ + service = await self._get_calendar_service(user_id) + + try: + event = service.events().insert( + calendarId='primary', + body=event_details + ).execute() + + return event + except Exception as e: + logger.error(f"Failed to add event: {str(e)}") + raise Exception(f"Failed to add event: {str(e)}") + + async def update_event(self, user_id, event_id, updated_event): + """ + Update an event in the user's primary calendar. + + Args: + user_id: The user's ID + event_id: ID of the event to update + updated_event: Dictionary with updated event details + + Returns: + dict: The updated event + """ + service = await self._get_calendar_service(user_id) + + try: + event = service.events().update( + calendarId='primary', + eventId=event_id, + body=updated_event + ).execute() + + return event + except Exception as e: + logger.error(f"Failed to update event: {str(e)}") + raise Exception(f"Failed to update event: {str(e)}") + + async def delete_event(self, user_id, event_id): + """ + Delete an event from the user's primary calendar. + + Args: + user_id: The user's ID + event_id: ID of the event to delete + """ + service = await self._get_calendar_service(user_id) + + try: + service.events().delete( + calendarId='primary', + eventId=event_id + ).execute() + + logger.info(f"Successfully deleted event {event_id}") + except Exception as e: + logger.error(f"Failed to delete event: {str(e)}") + raise Exception(f"Failed to delete event: {str(e)}") + + async def get_event(self, user_id, event_id): + """ + Get a specific event from the user's primary calendar. + + Args: + user_id: The user's ID + event_id: ID of the event to retrieve + + Returns: + dict: The event details + """ + service = await self._get_calendar_service(user_id) + + try: + event = service.events().get( + calendarId='primary', + eventId=event_id + ).execute() + + return event + except Exception as e: + logger.error(f"Failed to get event: {str(e)}") + raise Exception(f"Failed to get event: {str(e)}") + + async def search_events(self, user_id, query, max_results=10): + """ + Search for events in the user's primary calendar. + + Args: + user_id: The user's ID + query: Search query string + max_results: Maximum number of results to return + + Returns: + dict: Search results containing events + """ + service = await self._get_calendar_service(user_id) + + try: + events_result = service.events().list( + calendarId='primary', + q=query, + maxResults=max_results, + singleEvents=True, + orderBy='startTime' + ).execute() + + return events_result + except Exception as e: + logger.error(f"Failed to search events: {str(e)}") + raise Exception(f"Failed to search events: {str(e)}") + + async def share_event(self, user_id, event_id, shared_email): + """ + Share an event with another user by adding them as an attendee. + + Args: + user_id: The user's ID + event_id: ID of the event to share + shared_email: Email address of the user to share with + """ + service = await self._get_calendar_service(user_id) + + try: + # Get the event first + event = service.events().get( + calendarId='primary', + eventId=event_id + ).execute() + + # Initialize attendees list if it doesn't exist + if 'attendees' not in event: + event['attendees'] = [] + + # Add the new attendee + event['attendees'].append({'email': shared_email}) + + # Update the event + updated_event = service.events().update( + calendarId='primary', + eventId=event_id, + body=event + ).execute() + + logger.info(f"Successfully shared event {event_id} with {shared_email}") + return updated_event + except Exception as e: + logger.error(f"Failed to share event: {str(e)}") + raise Exception(f"Failed to share event: {str(e)}") + + async def create_calendar(self, user_id, calendar_name): + """ + Create a new calendar for the user. + + Args: + user_id: The user's ID + calendar_name: Name for the new calendar + + Returns: + str: ID of the created calendar + """ + service = await self._get_calendar_service(user_id) + + try: + calendar = { + 'summary': calendar_name, + 'timeZone': 'UTC' + } + + created_calendar = service.calendars().insert(body=calendar).execute() + + logger.info(f"Successfully created calendar: {calendar_name}") + return created_calendar['id'] + except Exception as e: + logger.error(f"Failed to create calendar: {str(e)}") + raise Exception(f"Failed to create calendar: {str(e)}") + + async def _store_token(self, user_id, access_token, refresh_token, expires_in): + """ + Store a token in the token storage. + + Args: + user_id: The user's ID + access_token: The access token + refresh_token: The refresh token + expires_in: Expiration time in seconds + """ + token_data = { + "access_token": access_token, + "refresh_token": refresh_token, + "expires_at": (datetime.utcnow() + timedelta(seconds=expires_in)).timestamp() + } + + # Serialize and encrypt the token data + serialized_token = json.dumps(token_data) + encrypted_token = TokenEncryptionHelper.encrypt_token(serialized_token, self.encryption_key) + + # Store in the token storage using the helper function + token_record = create_token_record(encrypted_token) + + self.token_storage.store_token(user_id, PLATFORM, SERVICE, token_record) + + async def _get_token_data(self, user_id): + """ + Get token data from storage. + + Args: + user_id: The user's ID + + Returns: + dict: The token data or None if not found + """ + token_record = self.token_storage.get_token(user_id, PLATFORM, SERVICE) + + if not token_record or not token_record.get("is_active") or token_record.get("is_revoked"): + logger.info(f"No valid token found in the storage for user {user_id}") + return None + + try: + encrypted_token = token_record.get("encrypted_token") + if not encrypted_token: + return None + + decrypted_token = TokenEncryptionHelper.decrypt_token(encrypted_token, self.encryption_key) + token_data = json.loads(decrypted_token) + + return token_data + except Exception as e: + logger.error(f"Error getting token data: {str(e)}") + return None + + async def _load_token(self, user_id): + """ + Load a token from the token storage. + + Args: + user_id: The user's ID + + Returns: + str: The access token, or None if not found or expired + """ + token_data = await self._get_token_data(user_id) + + if not token_data: + return None + + # Check if token is expired + expires_at = token_data.get("expires_at") + if expires_at and expires_at <= datetime.utcnow().timestamp(): + logger.info(f"Token expired for user {user_id}, attempting to refresh") + refresh_token = token_data.get("refresh_token") + if refresh_token: + try: + return await self._refresh_token(user_id, refresh_token) + except Exception as e: + logger.error(f"Error refreshing token: {str(e)}") + return None + return None + + return token_data.get("access_token") + + async def _refresh_token(self, user_id, refresh_token): + """ + Refresh an expired token. + + Args: + user_id: The user's ID + refresh_token: The refresh token + + Returns: + str: The new access token + """ + if not self.client_id: + raise ValueError("Google Client ID is not set in configuration.") + if not self.client_secret: + raise ValueError("Google Client Secret is not set in configuration.") + + payload = { + "grant_type": "refresh_token", + "refresh_token": refresh_token, + "client_id": self.client_id, + "client_secret": self.client_secret + } + + logger.info(f"Attempting to refresh token for user {user_id}") + response = requests.post(GOOGLE_TOKEN_URL, data=payload) + response_data = response.json() + + if response.status_code == 200 and "access_token" in response_data: + # Store the new token + expires_in = response_data.get("expires_in", 3600) # Default to 1 hour + await self._store_token( + user_id, + response_data["access_token"], + refresh_token, # Keep the existing refresh token if not provided + expires_in + ) + logger.info(f"Successfully refreshed token for user {user_id}") + return response_data["access_token"] + else: + error_msg = response_data.get("error_description", "Unknown error") + logger.error(f"Failed to refresh token: {error_msg}") + # If refresh fails, mark the token as revoked so we don't keep trying + token_record = self.token_storage.get_token(user_id, PLATFORM, SERVICE) + if token_record: + token_record["is_revoked"] = True + self.token_storage.store_token(user_id, PLATFORM, SERVICE, token_record) + raise Exception(f"Failed to refresh token: {error_msg}") + + async def _get_calendar_service(self, user_id): + """ + Get an authenticated Google Calendar service instance. + + Args: + user_id: The user's ID + + Returns: + Resource: The Calendar service instance + """ + token_data = await self._get_token_data(user_id) + + if not token_data: + raise self._create_auth_exception(user_id) + + # Check if token is expired + expires_at = token_data.get("expires_at") + if expires_at and expires_at <= datetime.utcnow().timestamp(): + # Refresh the token + refresh_token = token_data.get("refresh_token") + if not refresh_token: + raise self._create_auth_exception(user_id) + + try: + token_data["access_token"] = await self._refresh_token(user_id, refresh_token) + except Exception: + raise self._create_auth_exception(user_id) + + # Create credentials from token data + expiry = datetime.fromtimestamp(token_data.get("expires_at", 0)) + credentials = Credentials( + token=token_data.get("access_token"), + refresh_token=token_data.get("refresh_token"), + token_uri=GOOGLE_TOKEN_URL, + client_id=self.client_id, + client_secret=self.client_secret, + expiry=expiry + ) + + # Build the Calendar service + try: + service = build('calendar', 'v3', credentials=credentials) + return service + except Exception as e: + logger.error(f"Failed to build Calendar service: {str(e)}") + raise self._create_auth_exception(user_id) + + def _create_auth_exception(self, user_id): + """ + Create an authentication exception with reauthorization instructions. + + Args: + user_id: The user's ID + + Returns: + Exception: With reauthorization instructions + """ + # Don't try to generate an auth URL here, just return the instruction + return Exception( + "Your Google Calendar authorization has expired or is invalid. " + "Please use the `!authorize-gcalendar` command to reconnect your Google Calendar account." + ) \ No newline at end of file From aff4e710a253f44a4e072165d9c2f7829c891fbb Mon Sep 17 00:00:00 2001 From: christinaduong Date: Tue, 11 Mar 2025 14:55:23 -0700 Subject: [PATCH 19/19] update necessary files for gcal functionality --- agent.py | 35 ++++++-- bot.py | 260 +++++++++++++++++++++++++++++++++++++++++++++++++++++- server.py | 34 ++++++- 3 files changed, 318 insertions(+), 11 deletions(-) diff --git a/agent.py b/agent.py index bab06aa7..bd662b0e 100644 --- a/agent.py +++ b/agent.py @@ -11,6 +11,7 @@ from services.box_service import BoxService from services.dropbox_service import DropboxService from services.google_drive_service import GoogleDriveService +from services.google_calendar_service import GoogleCalendarService from plugins.cloud_plugin_manager import CloudPluginManager @@ -18,9 +19,9 @@ MISTRAL_MODEL = "mistral-large-latest" SYSTEM_PROMPT = """You are a helpful assistant named Dodobot that can access and manage various cloud services. -You can interact with services like Box, Dropbox, Gmail, and others to search for files, create folders, get download links, etc. +You can interact with services like Box, Dropbox, Gmail, Google Drive, Google Calendar and others to search for files, create folders, get download links, manage calendars, etc. -When a user asks about files, folders, or cloud storage, use the appropriate function to handle their request. +When a user asks about files, folders, cloud storage, or calendar events, use the appropriate function to handle their request. Never ask the user for their user ID, as it is automatically provided by the system. Do not expose implementation, internal values or functions to the user @@ -29,7 +30,7 @@ 2. Then provide the actual link on a separate line 3. Do not include raw function call data in your responses -If a service needs authorization, tell the user to use the !authorize-[service] command (e.g., !authorize-box, !authorize-dropbox). +If a service needs authorization, tell the user to use the !authorize-[service] command (e.g., !authorize-box, !authorize-dropbox, !authorize-gcalendar). Format your responses using Discord-compatible Markdown: - Use **bold** for emphasis @@ -60,6 +61,11 @@ - Google Drive uses file IDs and folder IDs for operations - File operations include sharing, viewing, and downloading +For Google Calendar: +- Use Google Calendar for managing events, meetings, and appointments +- You can create calendars, add events, view upcoming events, and share events with others +- Google Calendar uses event IDs and calendar IDs for operations + When a user attaches a file and asks to upload it, use the upload_file function from the Box plugins. You can find the attached file path in the file_paths parameter that will be provided to you. @@ -84,12 +90,14 @@ def __init__(self, max_context_messages=10): self.box_service = BoxService() self.dropbox_service = DropboxService() self.google_drive_service = GoogleDriveService() + self.google_calendar_service = GoogleCalendarService() # Initialize plugin manager and register plugins self.cloud_plugin_manager = CloudPluginManager( box_service=self.box_service, dropbox_service=self.dropbox_service, - google_drive_service=self.google_drive_service + google_drive_service=self.google_drive_service, + google_calendar_service=self.google_calendar_service ) # Register all cloud plugins with the kernel @@ -199,10 +207,12 @@ async def run(self, message: discord.Message): service_name = "Box" elif "dropbox" in func_name.lower(): service_name = "Dropbox" - elif "gdrive" in func_name.lower() or "google" in func_name.lower(): + elif "gdrive" in func_name.lower() or "google_drive" in func_name.lower(): service_name = "Google Drive" + elif "gcalendar" in func_name.lower() or "google_calendar" in func_name.lower(): + service_name = "Google Calendar" else: - service_name = "cloud storage" + service_name = "cloud service" # Determine action type if "get_file_download_link" in func_name or "download" in func_name: @@ -211,13 +221,19 @@ async def run(self, message: discord.Message): formatted_response = f"I'm searching for '{query}' in your {service_name} account..." elif "share" in func_name: formatted_response = f"I'll prepare to share '{query}' from your {service_name} account..." + elif "create_calendar" in func_name: + calendar_name = args.get("calendar_name", "new calendar") + formatted_response = f"I'm creating a new calendar '{calendar_name}' in your Google Calendar account..." + elif "add_event" in func_name or "create_event" in func_name: + summary = args.get("summary", "event") + formatted_response = f"I'm adding the event '{summary}' to your Google Calendar..." elif "create" in func_name: formatted_response = f"I'll create '{query}' in your {service_name} account..." elif "delete" in func_name: formatted_response = f"I'll prepare to delete '{query}' from your {service_name} account..." - elif "list" in func_name: + elif "list" in func_name or "get_events" in func_name: path = args.get("path", "root folder") - formatted_response = f"I'll list the contents of '{path}' in your {service_name} account..." + formatted_response = f"I'll list the contents in your {service_name} account..." elif "upload" in func_name: file_name = args.get("file_name", "your file") formatted_response = f"I'm uploading '{file_name}' to your {service_name} account..." @@ -236,7 +252,8 @@ async def run(self, message: discord.Message): "Please use the `!authorize", "not authorized", "authorization required", - "Google Drive authorization has expired" + "Google Drive authorization has expired", + "Google Calendar authorization has expired" ] if any(phrase in response.content for phrase in auth_error_phrases): diff --git a/bot.py b/bot.py index 33809a45..4a553fb3 100644 --- a/bot.py +++ b/bot.py @@ -236,6 +236,23 @@ async def authorize_gdrive(ctx): logger.error(error_msg) await ctx.send(error_msg[:1900]) +@bot.command(name="authorize-gcalendar", help="Authorize the bot to access your Google Calendar account") +async def authorize_gcalendar(ctx): + """ + Sends a Google Calendar authorization link to the user via DM. + """ + try: + # Get authorization URL for the user + auth_url = await google_calendar_service.get_authorization_url(str(ctx.author.id)) + + # Send the URL as a DM to the user + await ctx.author.send(f"Please authorize access to your Google Calendar account by clicking this link: {auth_url}") + await ctx.send("I've sent you a DM with the authorization link!") + except Exception as e: + error_msg = f"Error generating Google Calendar authorization link: {str(e)}" + logger.error(error_msg) + await ctx.send(error_msg[:1900]) + @bot.command(name="gdrive-upload", help="Upload a file to Google Drive") async def gdrive_upload(ctx): """ @@ -299,6 +316,224 @@ async def gcalendar_create(ctx, *, calendar_name=None): logger.error(error_msg) await ctx.send(error_msg[:1900]) +@bot.command(name="gcalendar-add-event", help="Add an event to your Google Calendar") +async def gcalendar_add_event(ctx, *, event_data=None): + """ + Adds an event to the user's Google Calendar. + + Usage: !gcalendar-add-event title | description | start_time | end_time | location + Example: !gcalendar-add-event Team Meeting | Weekly sync | 2024-03-15T14:00:00 | 2024-03-15T15:00:00 | Conference Room + """ + if not event_data: + await ctx.send("Please provide event details in this format:\n`!gcalendar-add-event title | description | start_time | end_time | location`\n\nExample: `!gcalendar-add-event Team Meeting | Weekly sync | 2024-03-15T14:00:00 | 2024-03-15T15:00:00 | Conference Room`") + return + + try: + # Parse event data + parts = event_data.split('|') + if len(parts) < 4: + await ctx.send("Please provide at least title, description, start time, and end time, separated by '|'") + return + + title = parts[0].strip() + description = parts[1].strip() + start_time = parts[2].strip() + end_time = parts[3].strip() + location = parts[4].strip() if len(parts) > 4 else "" + + # Create event object + event = { + "summary": title, + "description": description, + "location": location, + "start": { + "dateTime": start_time, + "timeZone": "UTC" + }, + "end": { + "dateTime": end_time, + "timeZone": "UTC" + } + } + + # Add the event + result = await google_calendar_service.add_event(str(ctx.author.id), event) + + # Create embed for nice display + embed = discord.Embed( + title="Event Added to Google Calendar", + description=description, + color=discord.Color.green() + ) + + embed.add_field(name="Title", value=title, inline=False) + embed.add_field(name="Start", value=start_time, inline=True) + embed.add_field(name="End", value=end_time, inline=True) + if location: + embed.add_field(name="Location", value=location, inline=False) + + event_link = f"https://calendar.google.com/calendar/event?eid={result.get('id', 'unknown')}" + embed.add_field(name="Calendar Link", value=f"[Open in Google Calendar]({event_link})", inline=False) + + await ctx.send(embed=embed) + except Exception as e: + error_msg = f"Error adding event to Google Calendar: {str(e)}" + logger.error(error_msg) + await ctx.send(error_msg[:1900]) + +@bot.command(name="gcalendar-events", help="Get your upcoming events from Google Calendar") +async def gcalendar_events(ctx, days: int = 7): + """ + Gets the user's upcoming events from Google Calendar. + + Args: + days: Number of days to look ahead (default: 7) + + Usage: !gcalendar-events + Alternative: !gcalendar-events 14 + """ + try: + # Calculate date range + start_date = datetime.utcnow() + end_date = start_date + timedelta(days=days) + + # Get events + events = await google_calendar_service.get_events( + str(ctx.author.id), + start_date, + end_date, + max_results=10 + ) + + if not events: + await ctx.send(f"No events found in the next {days} days.") + return + + # Create an embed for nice formatting + embed = discord.Embed( + title=f"Your Calendar: Next {days} Days", + description=f"Showing your upcoming events from {start_date.strftime('%Y-%m-%d')} to {end_date.strftime('%Y-%m-%d')}", + color=discord.Color.blue() + ) + + # Add events to the embed + for event in events: + # Get event details + title = event.get('summary', 'No title') + + # Format start time + start = event.get('start', {}) + if 'dateTime' in start: + start_time = datetime.fromisoformat(start['dateTime'].replace('Z', '+00:00')) + start_str = start_time.strftime('%Y-%m-%d %H:%M') + elif 'date' in start: + start_str = f"{start['date']} (All day)" + else: + start_str = "Unknown time" + + # Format location if available + location = event.get('location', '') + location_str = f"\nLocation: {location}" if location else "" + + # Add to embed + embed.add_field( + name=title, + value=f"When: {start_str}{location_str}", + inline=False + ) + + await ctx.send(embed=embed) + except Exception as e: + error_msg = f"Error retrieving events from Google Calendar: {str(e)}" + logger.error(error_msg) + await ctx.send(error_msg[:1900]) + +@bot.command(name="gcalendar-delete", help="Delete an event from your Google Calendar") +async def gcalendar_delete(ctx, *, event_query=None): + """ + Deletes an event from the user's Google Calendar by searching for it. + + Usage: !gcalendar-delete Team Meeting + Alternative: !gcalendar-delete [event-id] + """ + if not event_query: + await ctx.send("Please provide an event title or event ID to delete.\nUsage: `!gcalendar-delete Team Meeting`") + return + + try: + result = await google_calendar_service.delete_event(str(ctx.author.id), event_query) + await ctx.send(result) + except Exception as e: + error_msg = f"Error deleting event: {str(e)}" + logger.error(error_msg) + await ctx.send(error_msg[:1900]) + +@bot.command(name="gcalendar-update", help="Update an existing event in your Google Calendar") +async def gcalendar_update(ctx, event_id=None, *, event_data=None): + """ + Updates an existing event in the user's Google Calendar. + + Usage: !gcalendar-update event_id title | description | start_time | end_time + Example: !gcalendar-update abc123 Updated Meeting | New description | 2024-03-15T15:00:00 | 2024-03-15T16:00:00 + """ + if not event_id or not event_data: + await ctx.send("Please provide both an event ID and updated event details.\nUsage: `!gcalendar-update event_id title | description | start_time | end_time`") + return + + try: + # Parse event data + parts = event_data.split('|') + if len(parts) < 4: + await ctx.send("Please provide at least title, description, start time, and end time, separated by '|'") + return + + title = parts[0].strip() + description = parts[1].strip() + start_time = parts[2].strip() + end_time = parts[3].strip() + + result = await google_calendar_service.update_event( + str(ctx.author.id), + event_id, + { + "summary": title, + "description": description, + "start": { + "dateTime": start_time, + "timeZone": "UTC" + }, + "end": { + "dateTime": end_time, + "timeZone": "UTC" + } + } + ) + + await ctx.send(f"Event updated successfully!\nTitle: {result.get('summary')}\nID: {result.get('id')}") + except Exception as e: + error_msg = f"Error updating event: {str(e)}" + logger.error(error_msg) + await ctx.send(error_msg[:1900]) + +@bot.command(name="gcalendar-share", help="Share a calendar event with another user") +async def gcalendar_share(ctx, event_id=None, *, email=None): + """ + Shares a calendar event with another user by adding them as an attendee. + + Usage: !gcalendar-share event_id email@example.com + """ + if not event_id or not email: + await ctx.send("Please provide both an event ID and an email to share with.\nUsage: `!gcalendar-share event_id email@example.com`") + return + + try: + result = await google_calendar_service.share_event(str(ctx.author.id), event_id, email) + await ctx.send(f"Event shared successfully with {email}!") + except Exception as e: + error_msg = f"Error sharing event: {str(e)}" + logger.error(error_msg) + await ctx.send(error_msg[:1900]) + @bot.command(name="cloud-status", help="Check your cloud service connections") async def cloud_status(ctx): """ @@ -311,7 +546,7 @@ async def cloud_status(ctx): ) embed.set_author(name=ctx.author.display_name, icon_url=ctx.author.avatar.url if ctx.author.avatar else None) - embed.set_footer(text="Use !authorize-box or !authorize-dropbox to connect services") + embed.set_footer(text="Use !authorize-* commands to connect services") embed.timestamp = discord.utils.utcnow() # Check Box connection @@ -382,6 +617,29 @@ async def cloud_status(ctx): value=f"⚠️ Error checking connection\n```{str(e)}```", inline=False ) + + # Check Google Calendar connection + try: + # Try to load the token to see if the user is authenticated + gcalendar_token = await google_calendar_service._load_token(str(ctx.author.id)) + if gcalendar_token: + embed.add_field( + name="Google Calendar Status", + value="✅ Connected", + inline=False + ) + else: + embed.add_field( + name="Google Calendar Status", + value="❌ Not connected\n*Use !authorize-gcalendar to connect*", + inline=False + ) + except Exception as e: + embed.add_field( + name="Google Calendar Status", + value=f"⚠️ Error checking connection\n```{str(e)}```", + inline=False + ) await ctx.send(embed=embed) diff --git a/server.py b/server.py index 164576b4..c7f519f7 100644 --- a/server.py +++ b/server.py @@ -4,6 +4,7 @@ from services.box_service import BoxService from services.dropbox_service import DropboxService from services.google_drive_service import GoogleDriveService +from services.google_calendar_service import GoogleCalendarService from helpers.token_helpers import TokenEncryptionHelper import asyncio import logging @@ -16,6 +17,7 @@ box_service = BoxService() dropbox_service = DropboxService() google_drive_service = GoogleDriveService() +google_calendar_service = GoogleCalendarService() # This will be set from bot.py bot = None @@ -81,7 +83,7 @@ def get_success_html(service_name): @app.get("/") async def root(): - return {"message": "OAuth Callback Server for Box, Dropbox, and Google Drive"} + return {"message": "OAuth Callback Server for Box, Dropbox, Google Drive, and Google Calendar"} @app.get("/box/callback") async def box_callback(code: str, state: str): @@ -173,6 +175,36 @@ async def gdrive_callback(code: str, state: str): logger.error(f"Error in Google Drive callback: {str(e)}") return {"error": str(e)} +@app.get("/gcalendar/callback") +async def gcalendar_callback(code: str, state: str): + """ + Handle the OAuth callback from Google Calendar. + + This endpoint receives the authorization code from Google Calendar after a user + authorizes the application. It exchanges the code for access and + refresh tokens, stores them securely, and notifies the user. + """ + try: + # Get user ID from state + user_id = TokenEncryptionHelper.decrypt_token(state, google_calendar_service.encryption_key) + logger.info(f"Received Google Calendar callback for user {user_id}") + + # Handle the callback - this stores the tokens + await google_calendar_service.handle_auth_callback(state, code) + + # Notify the user through Discord + if bot: + # Schedule the notification in the bot's event loop + asyncio.run_coroutine_threadsafe(notify_user(user_id, "Google Calendar"), bot.loop) + + # Use the reusable HTML template + html_content = get_success_html("Google Calendar") + + return HTMLResponse(content=html_content) + except Exception as e: + logger.error(f"Error in Google Calendar callback: {str(e)}") + return {"error": str(e)} + async def notify_user(user_id, service_name): """ Send a Discord message to notify the user that authorization was successful.