list/get/drop/create via http api

Author: ea-rus

mindsdb_sdk/connectors/rest_api.py

@@ -409,3 +409,34 @@ def insert_into_knowledge_base(self, project: str, knowledge_base_name: str, dat
         _raise_for_status(r)
 
         return r.json()
+
+    @_try_relogin
+    def list_knowledge_bases(self, project: str):
+        r = self.session.get(self.url + f'/api/projects/{project}/knowledge_bases')
+        _raise_for_status(r)
+        return r.json()
+
+    @_try_relogin
+    def get_knowledge_base(self, project: str, knowledge_base_name):
+        r = self.session.get(self.url + f'/api/projects/{project}/knowledge_bases/{knowledge_base_name}')
+        _raise_for_status(r)
+        return r.json()
+
+    @_try_relogin
+    def delete_knowledge_base(self, project: str, knowledge_base_name):
+        r = self.session.delete(self.url + f'/api/projects/{project}/knowledge_bases/{knowledge_base_name}')
+        _raise_for_status(r)
+
+    @_try_relogin
+    def create_knowledge_base(self, project: str, data):
+        r = self.session.post(
+            self.url + f'/api/projects/{project}/knowledge_bases',
+            json={
+                'knowledge_base': data
+            }
+        )
+        _raise_for_status(r)
+
+        return r.json()
+
+

mindsdb_sdk/knowledge_bases.py

@@ -39,19 +39,14 @@ def __init__(self, api, project, data: dict):
         self.table_name = Identifier(parts=[self.project.name, self.name])
 
         self.storage = None
-        if data['storage'] is not None:
-            # if name contents '.' there could be errors
-
-            parts = data['storage'].split('.')
-            if len(parts) == 2:
-                database_name, table_name = parts
-                database = Database(project, database_name)
-                table = Table(database, table_name)
-                self.storage = table
+        if data.get('vector_database_table') is not None:
+            database = Database(project, data['vector_database'])
+            table = Table(database, data['vector_database_table'])
+            self.storage = table
 
         self.model = None
-        if data['model'] is not None:
-            self.model = Model(self.project, {'name': data['model']})
+        if data['embedding_model'] is not None:
+            self.model = Model(self.project, {'name': data['embedding_model']})
 
         params = data.get('params', {})
         if isinstance(params, str):
@@ -241,24 +236,6 @@ def __init__(self, project, api):
         self.project = project
         self.api = api
 
-    def _list(self, name: str = None) -> List[KnowledgeBase]:
-
-        # TODO add filter by project. for now 'project' is empty
-        ast_query = Select(targets=[Star()], from_table=Identifier(parts=['information_schema', 'knowledge_bases']))
-        if name is not None:
-            ast_query.where = dict_to_binary_op({'name': name})
-
-        df = self.api.sql_query(ast_query.to_string(), database=self.project.name)
-
-        # columns to lower case
-        cols_map = {i: i.lower() for i in df.columns}
-        df = df.rename(columns=cols_map)
-
-        return [
-            KnowledgeBase(self.api, self.project, item)
-            for item in df.to_dict('records')
-        ]
-
     def list(self) -> List[KnowledgeBase]:
         """
 
@@ -268,7 +245,11 @@ def list(self) -> List[KnowledgeBase]:
 
         :return: list of knowledge bases
         """
-        return self._list()
+
+        return [
+            KnowledgeBase(self.api, self.project, item)
+            for item in self.api.list_knowledge_bases(self.project.name)
+        ]
 
     def get(self, name: str) -> KnowledgeBase:
         """
@@ -277,13 +258,9 @@ def get(self, name: str) -> KnowledgeBase:
         :param name: name of the knowledge base
         :return: KnowledgeBase object
         """
-        item = self._list(name)
-        if len(item) == 1:
-            return item[0]
-        elif len(item) == 0:
-            raise AttributeError("KnowledgeBase doesn't exist")
-        else:
-            raise RuntimeError("Several knowledgeBases with the same name")
+
+        data = self.api.get_knowledge_base(self.project.name, name)
+        return KnowledgeBase(self.api, self.project, data)
 
     def create(
         self,
@@ -334,27 +311,21 @@ def create(
             params_out.update(params)
 
         if model is not None:
-            model_name = Identifier(parts=[model.project.name, model.name])
-        else:
-            model_name = None
+            model = model.name
 
-        if storage is not None:
-            storage_name = Identifier(parts=[storage.db.name, storage.name])
-        else:
-            storage_name = None
-
-        ast_query = CreateKnowledgeBase(
-            Identifier(parts=[self.project.name, name]),
-            model=model_name,
-            storage=storage_name,
-            params=params_out
-        )
-        sql = ast_query.to_string()
+        payload = {
+            'name': name,
+            'model': model,
+            'params': params_out
+        }
 
-        if is_saving():
-            return Query(self, sql)
+        if storage is not None:
+            payload['storage'] = {
+                'database': storage.db.name,
+                'table': storage.name
+            }
 
-        self.api.sql_query(sql)
+        self.api.create_knowledge_base(self.project.name, data=payload)
 
         return self.get(name)
 
@@ -365,10 +336,4 @@ def drop(self, name: str):
         :return:
         """
 
-        ast_query = DropKnowledgeBase(Identifier(parts=[self.project.name, name]))
-        sql = ast_query.to_string()
-
-        if is_saving():
-            return Query(self, sql)
-
-        self.api.sql_query(sql)
+        return self.api.delete_knowledge_base(self.project.name, name)

tests/test_sdk.py

@@ -1130,7 +1130,9 @@ def check_project_jobs(self, project, model, database, kb, mock_post):
 
     @patch('requests.Session.put')
     @patch('requests.Session.post')
-    def check_project_kb(self, project, model, database, mock_post, mock_put):
+    @patch('requests.Session.delete')
+    @patch('requests.Session.get')
+    def check_project_kb(self, project, model, database, mock_get, mock_del, mock_post, mock_put):
 
         response_mock(mock_post, pd.DataFrame([{
             'NAME': 'my_kb',
@@ -1140,10 +1142,24 @@ def check_project_kb(self, project, model, database, mock_post, mock_put):
             'PARAMS': {"id_column": "num"},
         }]))
 
+        example_kb = {
+            'id': 1,
+            'name': 'my_kb',
+            'project_id': 1,
+            'embedding_model': 'openai_emb',
+            'vector_database': 'pvec',
+            'vector_database_table': 'tbl1',
+            'updated_at': '2024-10-04 10:55:25.350799',
+            'created_at': '2024-10-04 10:55:25.350790',
+            'params': {}
+        }
+
+        mock_get().json.return_value = [example_kb]
+
         kbs = project.knowledge_bases.list()
 
-        # TODO add filter by project
-        check_sql_call(mock_post, "select * from information_schema.knowledge_bases")
+        args, kwargs = mock_get.call_args
+        assert args[0] == f'{DEFAULT_CLOUD_API_URL}/api/projects/{project.name}/knowledge_bases'
 
         kb = kbs[0]
 
@@ -1156,6 +1172,7 @@ def check_project_kb(self, project, model, database, mock_post, mock_put):
         assert kb.storage.name == 'tbl1'
         assert kb.storage.db.name == 'pvec'
 
+        mock_get().json.return_value = example_kb
         kb = project.knowledge_bases.my_kb
 
         str(kb)
@@ -1208,18 +1225,16 @@ def check_project_kb(self, project, model, database, mock_post, mock_put):
             metadata_columns=['date', 'author'],
             params={'k': 'v'}
         )
-
-        model_name = f'{model.project.name}.{model.name}'
-        check_sql_call(
-            mock_post,
-            f'''
-            CREATE KNOWLEDGE_BASE {project.name}.kb2
-              USING model={model_name}, 
-              metadata_columns=['date', 'author'],
-              k='v'
-            ''',
-            call_stack_num=-2
-        )
+        args, kwargs = mock_post.call_args
+        assert args[0] == f'{DEFAULT_CLOUD_API_URL}/api/projects/{project.name}/knowledge_bases'
+        assert kwargs == {'json': {'knowledge_base': {
+            'name': 'kb2',
+            'model': model.name,
+            'params': {
+                'k': 'v',
+                'metadata_columns': ['date', 'author']
+            }
+        }}}
 
         # create 2
         kb = project.knowledge_bases.create(
@@ -1229,25 +1244,26 @@ def check_project_kb(self, project, model, database, mock_post, mock_put):
             id_column='num'
         )
 
-        table_name = f'{database.name}.tbl1'
-        check_sql_call(
-            mock_post,
-            f'''
-            CREATE KNOWLEDGE_BASE {project.name}.kb2
-              USING storage={table_name}, 
-              content_columns=['review'],
-              id_column='num'
-            ''',
-            call_stack_num=-2
-        )
+        args, kwargs = mock_post.call_args
+        assert args[0] == f'{DEFAULT_CLOUD_API_URL}/api/projects/{project.name}/knowledge_bases'
+        assert kwargs == {'json': {'knowledge_base': {
+            'name': 'kb2',
+            'model': None,
+            'params': {
+                'content_columns': ['review'],
+                'id_column': 'num'
+            },
+            'storage': {
+                'database': database.name,
+                'table': 'tbl1'
+            },
+        }}}
 
         # drop
         project.knowledge_bases.drop('kb2')
 
-        check_sql_call(
-            mock_post,
-            f"DROP KNOWLEDGE_BASE {project.name}.kb2"
-        )
+        args, kwargs = mock_del.call_args
+        assert args[0] == f'{DEFAULT_CLOUD_API_URL}/api/projects/{project.name}/knowledge_bases/kb2'
 
         return kb
 
@@ -1501,6 +1517,18 @@ def test_add_file(self, mock_post, mock_put, mock_get):
                 'updated_at': None,
                 'provider': 'mindsdb'
             },
+            # get KB
+            {
+                'id': 1,
+                'name': 'my_kb',
+                'project_id': 1,
+                'embedding_model': 'openai_emb',
+                'vector_database': 'pvec',
+                'vector_database_table': 'tbl1',
+                'updated_at': '2024-10-04 10:55:25.350799',
+                'created_at': '2024-10-04 10:55:25.350790',
+                'params': {}
+            },
             # Skills get in Agent update to check if it exists.
             {'name': 'new_skill', 'type': 'retrieval', 'params': {'source': 'test_agent_tokaido_rules_kb'}},
             # Existing agent get in Agent update.
@@ -1515,10 +1543,6 @@ def test_add_file(self, mock_post, mock_put, mock_get):
             },
         ])
         responses_mock(mock_post, [
-            # KB get (POST /sql).
-            pd.DataFrame([
-                {'name': 'test_agent_tokaido_rules_kb', 'storage': None, 'model': None},
-            ]),
             # Skill creation.
             {'name': 'new_skill', 'type': 'retrieval', 'params': {'source': 'test_agent_tokaido_rules_kb'}}
         ])
@@ -1568,6 +1592,18 @@ def test_add_webpage(self, mock_post, mock_put, mock_get):
                 'updated_at':None,
                 'provider':'mindsdb'
             },
+            # get KB
+            {
+                'id': 1,
+                'name': 'my_kb',
+                'project_id': 1,
+                'embedding_model': 'openai_emb',
+                'vector_database': 'pvec',
+                'vector_database_table': 'tbl1',
+                'updated_at': '2024-10-04 10:55:25.350799',
+                'created_at': '2024-10-04 10:55:25.350790',
+                'params': {}
+            },
             # Skills get in Agent update to check if it exists.
             {'name':'new_skill', 'type':'retrieval', 'params':{'source':'test_agent_docs_mdb_ai_kb'}},
             # Existing agent get in Agent update.
@@ -1582,10 +1618,6 @@ def test_add_webpage(self, mock_post, mock_put, mock_get):
             },
         ])
         responses_mock(mock_post, [
-            # KB get (POST /sql).
-            pd.DataFrame([
-                {'name':'test_agent_docs_mdb_ai_kb', 'storage':None, 'model':None},
-            ]),
             # Skill creation.
             {'name':'new_skill', 'type':'retrieval', 'params':{'source':'test_agent_docs_mdb_ai_kb'}}
         ])