Merge pull request #113 from mindsdb/fix-file-upload

Add functionality to upload files to MindsDB

Author: tmichaeldb

examples/data/tokaido-rulebook.pdf

@@ -0,0 +1,12 @@
+import mindsdb_sdk
+
+con = mindsdb_sdk.connect()
+
+# Now create an agent that will use the model we just created.
+agent = con.agents.get('agent_with_retrieval')
+agent.add_file('./data/tokaido-rulebook.pdf', 'rule book for the board game takaido')
+
+print('Ask a question: ')
+question = input()
+answer = agent.completion([{'question': question, 'answer': None}])
+print(answer.content)

examples/using_agents_with_retrieval.py

@@ -233,11 +233,8 @@ def add_files(self, name: str, file_paths: List[str], description: str, knowledg
             except HTTPError as e:
                 if e.response.status_code >= 400 and e.response.status_code != 404:
                     raise e
-                # Upload file if it doesn't exist.
-                with open(file_path, 'rb') as file:
-                    content = file.read()
-                    df = pd.DataFrame.from_records([{'content': content}])
-                    self.api.upload_file(filename_no_extension, df)
+                # upload file to mindsdb
+                self.api.upload_file(filename, file_path)
 
         # Insert uploaded files into new knowledge base.
         agent = self.get(name)

mindsdb_sdk/agents.py

@@ -1,5 +1,5 @@
 from functools import wraps
-from typing import List
+from typing import List, Union
 import io
 
 import requests
@@ -128,28 +128,72 @@ def objects_tree(self, item=''):
 
         return pd.DataFrame(r.json())
 
-    @_try_relogin
-    def upload_file(self, name: str, df: pd.DataFrame):
-
-        # convert to file
+    @staticmethod
+    def read_file_as_bytes(file_path: str):
+        """
+        Read and return content of a file in bytes, given its path.
+        :param file_path: Path of the file to read.
+        :return: File content in bytes.
+        """
+        try:
+            with open(file_path, 'rb+') as file:
+                return file.read()
+        except FileNotFoundError:
+            raise Exception(f'File {file_path} does not exist.')
+        except PermissionError:
+            raise Exception(f'Permission denied when reading file {file_path}.')
+        except Exception as e:
+            raise Exception(f'Unknown error occurred when reading file {file_path} - {str(e)}')
+    @staticmethod
+    def read_dataframe_as_csv(data: pd.DataFrame):
+        """
+        Read and return content of a DataFrame as CSV in bytes.
+        :param data: DataFrame to read.
+        :return: DataFrame content as CSV in bytes.
+        """
         fd = io.BytesIO()
-        df.to_csv(fd, index=False)
+        data.to_csv(fd, index=False)
         fd.seek(0)
+        return fd.read()
+
+    def upload_data(self, file_name: str, data: bytes):
+        """
+        Upload binary data to MindsDB.
+        :param file_name: Name of the file.
+        :param data: Binary data to upload.
+        """
+        # remove suffix from file if present
+        name = file_name.split('.')[0]
 
         url = self.url + f'/api/files/{name}'
         r = self.session.put(
             url,
             data={
-                'source': name,
-                'name': name,
-                'source_type': 'file',
+                'original_file_name':file_name,
+                'name':name,
+                'source_type':'file',
             },
             files={
-                'file': fd,
+                'file': (file_name, data)
+
             }
         )
         _raise_for_status(r)
 
+    @_try_relogin
+    def upload_file(self, name: str, data: Union[pd.DataFrame, str]):
+        """
+        Upload a file or a DataFrame to MindsDB.
+        :param name: Name of the file or DataFrame.
+        :param data: DataFrame data or file path.
+        """
+        if isinstance(data, pd.DataFrame):
+            data_in_bytes = self.read_dataframe_as_csv(data)
+        else:
+            data_in_bytes = self.read_file_as_bytes(data)
+
+        self.upload_data(name, data_in_bytes)
+
     @_try_relogin
     def get_file_metadata(self, name: str) -> dict:
         # No endpoint currently to get single file.