Use new KB endpoint for inserting files

Author: tmichaeldb

mindsdb_sdk/agents.py

@@ -192,7 +192,7 @@ def add_file(self, name: str, file_path: str, description: str, knowledge_base:
                 kb.model.wait_complete()
 
         # Insert the entire file.
-        kb.insert(self.databases.files.tables.get(filename_no_extension))
+        kb.insert_files([filename_no_extension])
 
         # Make sure skill name is unique.
         skill_name = f'{filename_no_extension}_retrieval_skill_{uuid4()}'

mindsdb_sdk/connectors/rest_api.py

@@ -319,3 +319,18 @@ def delete_skill(self, project: str, name: str):
         url = self.url + f'/api/projects/{project}/skills/{name}'
         r = self.session.delete(url)
         _raise_for_status(r)
+
+    # Knowledge Base operations.
+    @_try_relogin
+    def insert_files_into_knowledge_base(self, project: str, knowledge_base_name: str, file_names: List[str]):
+        r = self.session.put(
+            self.url + f'/api/projects/{project}/knowledge_bases/{knowledge_base_name}',
+            json={
+                'knowledge_base': {
+                    'files': file_names
+                }
+            }
+        )
+        _raise_for_status(r)
+
+        return r.json()

mindsdb_sdk/knowledge_bases.py

@@ -31,8 +31,8 @@ class KnowledgeBase(Query):
 
     """
 
-    def __init__(self, project, data: dict):
-
+    def __init__(self, api, project, data: dict):
+        self.api = api
         self.project = project
         self.name = data['name']
 
@@ -118,6 +118,12 @@ def _update_query(self):
             ast_query.limit = Constant(self._limit)
         self.sql = ast_query.to_string()
 
+    def insert_files(self, file_paths: List[str]):
+        """
+        Insert data from file to knowledge base
+        """
+        self.api.insert_files_into_knowledge_base(self.project.name, self.name, file_paths)
+
     def insert(self, data: Union[pd.DataFrame, Query, dict]):
         """
         Insert data to knowledge base
@@ -202,7 +208,7 @@ def _list(self, name: str = None) -> List[KnowledgeBase]:
         df = df.rename(columns=cols_map)
 
         return [
-            KnowledgeBase(self.project, item)
+            KnowledgeBase(self.api, self.project, item)
             for item in df.to_dict('records')
         ]
 

mindsdb_sdk/skills.py

@@ -18,7 +18,7 @@ class Skill():
 
     Create a new SQL skill:
 
-    >>> text_to_sql_skill = skills.create('text_to_sql', 'sql', { 'tables': ['my_table'], 'database': 'my_database' })
+    >>> text_to_sql_skill = skills.create('text_to_sql', 'sql', { 'tables': ['my_table'], 'database': 'my_database', 'description': 'my_description'})
 
     Update a skill:
 
@@ -50,11 +50,14 @@ def __repr__(self):
 
     @classmethod
     def from_json(cls, json: dict):
+        name = json['name']
+        type = json['type']
+        params = json['params']
         if json['type'] == 'sql':
-            return SQLSkill(json['name'], json['params']['tables'], json['params']['database'], json['params']['description'])
+            return SQLSkill(name, params['tables'], params['database'], params.get('description', ''))
         if json['type'] == 'retrieval':
-            return RetrievalSkill(json['name'], json['params']['source'], json['params']['description'])
-        return Skill(json['name'], json['type'], json['params'])
+            return RetrievalSkill(name, params['source'], params.get('description', ''))
+        return Skill(name, type, params)
 
 
 class SQLSkill(Skill):

tests/test_sdk.py

@@ -1264,7 +1264,7 @@ def test_create(self, mock_post):
                 'name': 'test_skill',
                 'project_id': 1,
                 'type': 'sql',
-                'params': {'tables': ['test_table'], 'database': 'test_database'},
+                'params': {'tables': ['test_table'], 'database': 'test_database', 'description': 'test_description'},
             }],
             'params': {'k1': 'v1'},
             'created_at': created_at,
@@ -1277,7 +1277,7 @@ def test_create(self, mock_post):
         new_agent = server.agents.create(
             name='test_agent',
             model=Model(None, {'name':'m1'}),
-            skills=[SQLSkill('test_skill', ['test_table'], 'test_database'), 'test_description'],
+            skills=[SQLSkill('test_skill', ['test_table'], 'test_database', 'test_description')],
             params={'k1': 'v1'}
         )
         # Check API call.
@@ -1298,11 +1298,11 @@ def test_create(self, mock_post):
            'skill': {
                 'name': 'test_skill',
                 'type': 'sql',
-                'params': {'database': 'test_database', 'tables': ['test_table']}
+                'params': {'database': 'test_database', 'tables': ['test_table'], 'description': 'test_description'}
             }
         }
 
-        expected_skill = SQLSkill('test_skill', ['test_table'], 'test_database', 'test_descrition')
+        expected_skill = SQLSkill('test_skill', ['test_table'], 'test_database', 'test_description')
         expected_agent = Agent(
             'test_agent',
             'test_model',
@@ -1332,7 +1332,7 @@ def test_update(self, mock_get, mock_put, _):
                 'name': 'updated_skill',
                 'project_id': 1,
                 'type': 'sql',
-                'params': {'tables': ['updated_table'], 'database': 'updated_database'},
+                'params': {'tables': ['updated_table'], 'database': 'updated_database', 'description': 'test_description'},
             }],
             'params': {'k2': 'v2'},
             'created_at': created_at,
@@ -1422,7 +1422,7 @@ def test_list(self, mock_get):
                 'id': 1,
                 'name': 'test_skill',
                 'project_id': 1,
-                'params': {'tables': ['test_table'], 'database': 'test_database'},
+                'params': {'tables': ['test_table'], 'database': 'test_database', 'description': 'test_description' },
                 'type': 'sql'
             }
         ])
@@ -1440,7 +1440,7 @@ def test_get(self, mock_get):
                 'id': 1,
                 'name': 'test_skill',
                 'project_id': 1,
-                'params': {'tables': ['test_table'], 'database': 'test_database'},
+                'params': {'tables': ['test_table'], 'database': 'test_database', 'description': 'test_description'},
                 'type': 'sql'
             }
         )
@@ -1456,7 +1456,7 @@ def test_create(self, mock_post):
             'id': 1,
             'name': 'test_skill',
             'project_id': 1,
-            'params': {'k1': 'v1'},
+            'params': {'tables': ['test_table'], 'database': 'test_database', 'description': 'test_description'},
             'type': 'test'
         }
         response_mock(mock_post, data)
@@ -1466,18 +1466,18 @@ def test_create(self, mock_post):
         new_skill = server.skills.create(
             'test_skill',
             'sql',
-            params={'tables': ['test_table'], 'database': 'test_database'}
+            params={'tables': ['test_table'], 'database': 'test_database', 'description': 'test_description'}
         )
         # Check API call.
         assert mock_post.call_args[0][0] == f'{DEFAULT_LOCAL_API_URL}/api/projects/mindsdb/skills'
         assert mock_post.call_args[1]['json'] == {
            'skill': {
                 'name': 'test_skill',
                 'type': 'sql',
-                'params': {'database': 'test_database', 'tables': ['test_table']}
+                'params': {'database': 'test_database', 'tables': ['test_table'], 'description': 'test_description'}
             }
         }
-        expected_skill = SQLSkill('test_skill', ['test_table'], 'test_database')
+        expected_skill = SQLSkill('test_skill', ['test_table'], 'test_database', 'test_description')
 
         assert new_skill == expected_skill
 
@@ -1487,13 +1487,13 @@ def test_update(self, mock_put):
             'id': 1,
             'name': 'test_skill',
             'project_id': 1,
-            'params': {'tables': ['updated_table'], 'database': 'updated_database'},
+            'params': {'tables': ['updated_table'], 'database': 'updated_database', 'description': 'updated_description'},
             'type': 'sql'
         }
         response_mock(mock_put, data)
 
         server = mindsdb_sdk.connect()
-        expected_skill = SQLSkill('test_skill', ['updated_table'], 'updated_database', 'test_description')
+        expected_skill = SQLSkill('test_skill', ['updated_table'], 'updated_database', 'updated_description')
 
         updated_skill = server.skills.update('test_skill', expected_skill)
         # Check API call.
@@ -1502,7 +1502,7 @@ def test_update(self, mock_put):
            'skill': {
                 'name': 'test_skill',
                 'type': 'sql',
-                'params': {'tables': ['updated_table'], 'database': 'updated_database'}
+                'params': {'tables': ['updated_table'], 'database': 'updated_database', 'description': 'updated_description'}
             }
         }