Upload 2 files

src/agent.py

@@ -1,4 +1,6 @@
+import base64
 import json
+from os.path import join
 
 from langchain_core.messages import SystemMessage, HumanMessage
 from langchain_core.rate_limiters import InMemoryRateLimiter
@@ -20,6 +22,7 @@ class Agent:
         model: str,
         tools: list,
         system_prompt_path: str,
+        data_path: str,
         openai_api_key: str = None,
         langfuse_callback_handler: CallbackHandler = None
     ):
@@ -28,6 +31,7 @@ class Agent:
         :param model: The OpenAI model to use.
         :param tools: List of tools the agent can use.
         :param system_prompt_path: Path to the system prompt file.
+        :param data_path: Data to be used by the agent.
         :param openai_api_key: OpenAI API key for authentication.
         :param langfuse_callback_handler: Langfuse callback handler for
             tracking and logging interactions.
@@ -47,6 +51,7 @@ class Agent:
         )
         with open(system_prompt_path, "r") as file:
             self.system_prompt = file.read()
+        self.data_path = data_path
         self.tools = tools
         if langfuse_callback_handler is not None:
             self.chat_model.callbacks = [langfuse_callback_handler]
@@ -56,18 +61,27 @@ class Agent:
         )
         self.graph = self.__build_graph()
 
-    def __call__(self, question: str) -> tuple[str, str]:
+    def __call__(
+        self,
+        question: str,
+        question_file: str | None
+    ) -> tuple[str, str]:
         """
         Reply to a question using the agent and return the agents full reply
         with reasoning included.
         :param question: The question to ask the agent.
+        :param question_file: The file that comes with the question.
         :return: The agent's response.
         """
+        human_message = self.__format_human_message(
+            question=question,
+            question_file=question_file
+        )
         final_state = self.graph.invoke(
             input={
                 "messages": [
                     SystemMessage(content=self.system_prompt),
-                    HumanMessage(content=question)
+                    human_message
                 ]
             },
             config={
@@ -105,10 +119,77 @@ class Agent:
         response = self.chat_model_with_tools.invoke(state["messages"])
         return {"messages": [response]}
 
+    def __format_human_message(
+        self,
+        question: str,
+        question_file: str | None
+    ) -> HumanMessage:
+        """
+        Format the human message for the agent.
+        :param question: The question to ask the agent.
+        :param question_file: The file that comes with the question.
+        :return: Formatted HumanMessage.
+        """
+        if question_file is None:
+            human_message = HumanMessage(content=question)
+        else:
+            if '.png' in question_file:
+                with open(join(self.data_path, question_file), "rb") as file:
+                    file_content = base64.b64encode(file.read()).\
+                        decode("utf-8")
+                human_message = HumanMessage(
+                    content=[
+                        {
+                            'type': 'text',
+                            'text': question
+                        },
+                        {
+                            'type': 'image',
+                            'source_type': 'base64',
+                            'data': file_content,
+                            "mime_type": "image/png"
+                        }
+                    ]
+                )
+            elif '.mp3' in question_file:
+                with open(join(self.data_path, question_file), "rb") as file:
+                    file_content = base64.b64encode(file.read()).\
+                        decode("utf-8")
+                human_message = HumanMessage(
+                    content=[
+                        {
+                            'type': 'text',
+                            'text': question
+                        },
+                        {
+                            'type': 'audio',
+                            'source_type': 'base64',
+                            'data': file_content,
+                            'mime_type': 'audio/mp3'
+                        }
+                    ]
+                )
+            elif '.py' in question_file:
+                with open(join(self.data_path, question_file), "r") as file:
+                    file_content = file.read()
+                    human_message = HumanMessage(
+                        content=[
+                            {
+                                'type': 'text',
+                                'text': f'''{question}\n\nHere is the code:
+                                ```python\n{file_content}\n```'''
+                            },
+                        ]
+                    )
+            elif '.xlsx' in question_file:
+                human_message = HumanMessage(content=question)
+        return human_message
+
 
 if __name__ == "__main__":
 
     import os
+
     from langchain_community.tools import DuckDuckGoSearchResults
 
     from tools import multiply, add, subtract, divide, modulus

src/utils.py

@@ -0,0 +1,16 @@
+import requests
+
+
+def download_file(url: str, local_filepath: str) -> None:
+    """
+    Download a file from a URL and save it locally.
+
+    Args:
+        url (str): The URL of the file to download.
+        local_filename (str): The local path where the file will be saved.
+    """
+    with requests.get(url, stream=True) as response:
+        response.raise_for_status()
+        with open(local_filepath, 'wb') as file:
+            for chunk in response.iter_content(chunk_size=8192):
+                file.write(chunk)