Merge pull request #13 from PerfectThymeTech/marvinbuss/handle_files

Handle Files in Framework

code/backend/bots/assistant_bot.py

@@ -1,26 +1,45 @@
+import json
+import os
+import urllib
 from typing import List
 
 from botbuilder.core import ActivityHandler, MessageFactory, TurnContext
-from botbuilder.schema import ActionTypes, CardAction, ChannelAccount, SuggestedActions
+from botbuilder.schema import (
+    ActionTypes,
+    Attachment,
+    CardAction,
+    ChannelAccount,
+    SuggestedActions,
+)
+from core.config import settings
 from llm.assisstant import assistant_handler
+from models.assistant_bot_models import FileInfo
+from utils import get_logger
+
+logger = get_logger(__name__)
 
 
 class AssistantBot(ActivityHandler):
     thread_id = None
+    vector_store_id = []
 
     async def on_members_added_activity(
         self, members_added: List[ChannelAccount], turn_context: TurnContext
-    ):
+    ) -> None:
         """Onboards new members to the assistant by creating a new thread and adding a initial welcome message.
 
         members_added (List[ChannelAccount]): The list of channel accounts.
         turn_context (TurnContext): The turn context.
-        RETURNS (str): The welcome message actvity is being returned.
+        RETURNS (None): No return value.
         """
         for member in members_added:
             if member.id != turn_context.activity.recipient.id:
                 # Initialize thread in assistant
                 self.thread_id = assistant_handler.create_thread()
+
+                # # Initialize vector store in assistant
+                # self.vector_store_id = assistant_handler.create_vector_store(thread_id=self.thread_id)
+
                 # Respond with welcome message
                 welcome_message = (
                     "Hello and welcome! I am your personal joke assistant."
@@ -60,19 +79,83 @@ async def on_members_added_activity(
                 await turn_context.send_activity(suggested_topics)
 
                 # Add messages from assisstant to thread
-                assistant_handler.send_assisstant_message(welcome_message)
-                assistant_handler.send_assisstant_message(suggested_topics_message)
+                assistant_handler.send_assisstant_message(
+                    message=welcome_message, thread_id=self.thread_id
+                )
+                assistant_handler.send_assisstant_message(
+                    message=suggested_topics_message, thread_id=self.thread_id
+                )
+
+    async def on_message_activity(self, turn_context: TurnContext) -> None:
+        """Acts upon new messages or attachments added to a channel.
+
+        turn_context (TurnContext): The turn context.
+        RETURNS (None): No return value.
+        """
+        if (
+            turn_context.activity.attachments
+            and len(turn_context.activity.attachments) > 0
+        ):
+            # Download attachment and add it to thread
+            await self.__handle_incoming_attachment(turn_context)
+        else:
+            # Interact with assistant
+            message = assistant_handler.send_user_message(
+                message=turn_context.activity.text,
+                thread_id=self.thread_id,
+            )
+            if message:
+                await turn_context.send_activity(MessageFactory.text(message))
 
-    async def on_message_activity(self, turn_context: TurnContext):
-        """Acts upon new messages added to a channel.
+    async def __handle_incoming_attachment(self, turn_context: TurnContext) -> None:
+        """Handles all attachments uploaded by users.
 
         turn_context (TurnContext): The turn context.
-        RETURNS (str): The assistant message actvity is being returned.
+        RETURNS (None): No return value.
         """
-        # Interact with assistant
-        message = assistant_handler.send_user_message(
-            message=turn_context.activity.text,
-            thread_id=self.thread_id,
+        for attachment in turn_context.activity.attachments:
+            file_info = await self.__download_attachment_and_write(attachment)
+            if file_info:
+                self.vector_store_ids = assistant_handler.send_user_file(
+                    file_path=file_info.file_path, thread_id=self.thread_id
+                )
+        await turn_context.send_activity(
+            MessageFactory.text("The file was added to the context. How can I help?")
         )
-        if message:
-            return await turn_context.send_activity(MessageFactory.text(message))
+
+    async def __download_attachment_and_write(
+        self, attachment: Attachment
+    ) -> FileInfo | None:
+        """Retrieves the attachment via the attachment's contentUrl.
+
+        attachment (Attachment): Attachment sent by the user.
+        RETURNS (dict): Returns a dic containing the attachment details including the keys 'filename' and 'local_path'.
+        """
+        try:
+            response = urllib.request.urlopen(attachment.content_url)
+            headers = response.info()
+
+            # If user uploads JSON file, this prevents it from being written as
+            # "{"type":"Buffer","data":[123,13,10,32,32,34,108..."
+            if headers["content-type"] == "application/json":
+                data = bytes(json.load(response)["data"])
+            else:
+                data = response.read()
+
+            # Define directory
+            directory_path = os.path.join(settings.HOME_DIRECTORY, self.thread_id)
+            file_path = os.path.join(directory_path, attachment.name)
+            if not os.path.exists(directory_path):
+                os.makedirs(directory_path)
+
+            # Write file
+            with open(file_path, "wb") as file:
+                file.write(data)
+
+            return FileInfo(
+                file_name=attachment.name,
+                file_path=file_path,
+            )
+        except Exception as e:
+            logger.error(f"Failed to download file '{attachment.name}'", e)
+            return None

code/backend/core/config.py

@@ -14,6 +14,7 @@ class Settings(BaseSettings):
     # Web app settings
     WEBSITE_NAME: str = Field(default="test", alias="WEBSITE_SITE_NAME")
     WEBSITE_INSTANCE_ID: str = Field(default="0", alias="WEBSITE_INSTANCE_ID")
+    HOME_DIRECTORY: str = Field(default="", alias="HOME")
 
     # Logging settings
     LOGGING_LEVEL: int = logging.INFO

code/backend/llm/assisstant.py

@@ -1,5 +1,6 @@
 import json
 import time
+from typing import List
 
 from azure.identity import DefaultAzureCredential, get_bearer_token_provider
 from core.config import settings
@@ -39,7 +40,7 @@ def __init_assistant(self) -> str:
         assistant = self.client.beta.assistants.create(
             name=settings.PROJECT_NAME,
             instructions=settings.AZURE_OPENAI_SYSTEM_PROMPT,
-            tools=[],
+            tools=[{"type": "file_search"}],
             model=settings.AZURE_OPENAI_MODEL_NAME,
         )
         logger.info(f"Created new assisstant with assistant id: '{assistant.id}'")
@@ -50,10 +51,21 @@ def create_thread(self) -> str:
 
         RETURNS (str): Thread id of the newly created thread.
         """
-        thread = self.client.beta.threads.create()
+        thread = self.client.beta.threads.create(
+            # tool_resources={"file_search": {"vector_store_ids": []}},
+        )
         logger.debug(f"Created thread with thread id: '{thread.id}'")
         return thread.id
 
+    def create_vector_store(self, thread_id: str) -> str:
+        """Create a vector store in the assistant.
+
+        RETURNS (str): Vector store id of the newly created vector store.
+        """
+        vector_store = self.client.beta.vector_stores.create(name=thread_id)
+        logger.debug(f"Created vector store with id: '{vector_store.id}'")
+        return vector_store.id
+
     def send_user_message(self, message: str, thread_id: str) -> str | None:
         """Send a message to the thread and return the response from the assistant.
 
@@ -79,12 +91,12 @@ def send_user_message(self, message: str, thread_id: str) -> str | None:
         )
         run = self.__wait_for_run(run=run, thread_id=thread_id)
         run = self.__check_for_tools(run=run, thread_id=thread_id)
-        message = self.__get_assisstant_response(thread_id=thread_id)
+        response = self.__get_assisstant_response(thread_id=thread_id)
 
-        return message
+        return response
 
     def send_assisstant_message(self, message: str, thread_id: str) -> str | None:
-        """Send a message to the thread in teh context of the assisstant.
+        """Send a message to the thread in the context of the assisstant.
 
         message (str): The message to be sent to the thread in the contex of the assisstant.
         thread_id (str): The thread id to which the message should be sent to the assistant.
@@ -96,12 +108,45 @@ def send_assisstant_message(self, message: str, thread_id: str) -> str | None:
         if thread_id is None:
             return None
 
-        message = self.client.beta.threads.messages.create(
+        _ = self.client.beta.threads.messages.create(
             thread_id=thread_id,
             content=message,
             role="assistant",
         )
 
+    def send_user_file(self, file_path: str, thread_id: str) -> List[str]:
+        """Send a file to the thread in the context of the user and add it to the internal vector store.
+
+        file_path (str): The file path to the file that should be added to the fiel search.
+        thread_id (str): The thread id to which the message should be sent to the assistant.
+        RETURNS (List[str]): Returns the list of vector indexes.
+        """
+        # Upload file
+        logger.info(f"Uploading file '{file_path}' to assistant.")
+        file = self.client.files.create(
+            file=open(file_path, "rb"),
+            purpose="assistants",
+        )
+        # Attach file to thread
+        logger.info(f"Adding file '{file_path}' to thread '{thread_id}'")
+        _ = self.client.beta.threads.messages.create(
+            thread_id=thread_id,
+            content="File shared by the user.",
+            attachments=[{"file_id": file.id, "tools": [{"type": "file_search"}]}],
+            role="user",
+        )
+
+        # Return vector store id's
+        logger.info(f"Get thread '{thread_id}'")
+        thread = self.client.beta.threads.retrieve(
+            thread_id=thread_id,
+        )
+        vector_store_ids = thread.tool_resources.file_search.vector_store_ids
+        logger.info(
+            f"Vector indexes of thread '{thread_id}' are the following: '{vector_store_ids}'"
+        )
+        return vector_store_ids
+
     def __wait_for_run(self, run: Run, thread_id: str) -> Run:
         """Wait for the run to complete and return the run once completed.
 

code/backend/models/assistant_bot_models.py

@@ -0,0 +1,6 @@
+from pydantic import BaseModel
+
+
+class FileInfo(BaseModel):
+    file_name: str
+    file_path: str