insert uses api endpoints instead of query

Author: ea-rus

mindsdb_sdk/knowledge_bases.py

@@ -150,35 +150,48 @@ def insert(self, data: Union[pd.DataFrame, Query, dict]):
         """
         Insert data to knowledge base
 
-        >>> # insert using query
-        >>> my_kb.insert(server.databases.example_db.tables.houses_sales.filter(type='house'))
         >>> # using dataframe
         >>> my_kb.insert(pd.read_csv('house_sales.csv'))
         >>> # using dict
         >>> my_kb.insert({'type': 'house', 'date': '2020-02-02'})
 
-        Data will be if id (defined by id_column param, see create knowledge base) is already exists in knowledge base
-        it will be replaced
+        If id is already exists in knowledge base:
+        - it will be replaced
+        - `id` column can be defined by id_column param, see create knowledge base
 
         :param data: Dataframe or Query object or dict.
         """
 
+        if isinstance(data, Query):
+            # for back compatibility
+            return self.insert_query(data)
+
         if isinstance(data, dict):
-            data = pd.DataFrame([data])
+            data = [data]
+        elif isinstance(data, pd.DataFrame):
+            data = data.to_dict('records')
+        else:
+             raise ValueError("Unknown data type, accepted types: DataFrame, Query, dict")
+
+        return self.api.insert_into_knowledge_base(
+            self.project.name,
+            self.name,
+            data={'rows': data}
+        )
 
-        if isinstance(data, pd.DataFrame):
-            # insert data
-            data_split = data.to_dict('split')
+    def insert_query(self, data: Query):
+        """
+        Insert data to knowledge base using query
 
-            ast_query = Insert(
-                table=self.table_name,
-                columns=data_split['columns'],
-                values=data_split['data']
-            )
-            sql = ast_query.to_string()
+        >>> my_kb.insert(server.databases.example_db.tables.houses_sales.filter(type='house'))
 
-        else:
-            # insert from select
+        Data will be if id (defined by id_column param, see create knowledge base) is already exists in knowledge base
+        it will be replaced
+
+        :param data: Dataframe or Query object or dict.
+        """
+        if is_saving():
+            # generate insert from select query
             if data.database is not None:
                 ast_query = Insert(
                     table=self.table_name,
@@ -188,11 +201,15 @@ def insert(self, data: Union[pd.DataFrame, Query, dict]):
             else:
                 sql = f'INSERT INTO {self.table_name.to_string()} ({data.sql})'
 
-        if is_saving():
             # don't execute it right now, return query object
             return Query(self, sql, self.database)
 
-        self.api.sql_query(sql, self.database)
+        # query have to be in context of mindsdb project
+        self.api.insert_into_knowledge_base(
+            self.project.name,
+            self.name,
+            data={'query': data.sql}
+        )
 
 
 class KnowledgeBases(CollectionBase):

tests/test_sdk.py

@@ -12,7 +12,7 @@
 import mindsdb_sdk
 
 from mindsdb_sdk.agents import Agent
-from mindsdb_sdk.connect import DEFAULT_LOCAL_API_URL
+from mindsdb_sdk.connect import DEFAULT_LOCAL_API_URL, DEFAULT_CLOUD_API_URL
 from mindsdb_sdk.skills import SQLSkill
 from mindsdb_sdk.connectors import rest_api
 
@@ -1128,8 +1128,9 @@ def check_project_jobs(self, project, model, database, kb, mock_post):
             call_stack_num=-2
         )
 
+    @patch('requests.Session.put')
     @patch('requests.Session.post')
-    def check_project_kb(self, project, model, database, mock_post):
+    def check_project_kb(self, project, model, database, mock_post, mock_put):
 
         response_mock(mock_post, pd.DataFrame([{
             'NAME': 'my_kb',
@@ -1162,34 +1163,35 @@ def check_project_kb(self, project, model, database, mock_post):
         assert kb.storage.db.name == 'pvec'
         assert kb.model.name == 'openai_emb'
 
-        # insert
+        # --- insert ---
 
+        # table
         kb.insert(
             database.tables.tbl2.filter(a=1)
         )
-        check_sql_call(
-            mock_post,
-            f''' insert into {project.name}.{kb.name} (
-               select * from {database.name}.tbl2 where a=1
-            )'''
-        )
+
+        args, kwargs = mock_put.call_args
+        assert args[0] == f'{DEFAULT_CLOUD_API_URL}/api/projects/{project.name}/knowledge_bases/my_kb'
+        assert kwargs == {'json': {'knowledge_base': {'query': 'SELECT * FROM pg1.tbl2 WHERE a = 1'}}}
+
+        # query
         kb.insert(
             database.query('select * from tbl2 limit 1')
         )
-        check_sql_call(
-            mock_post,
-            f''' insert into {project.name}.{kb.name} (
-               select * from {database.name} (select * from tbl2 limit 1)
-            )'''
-        )
+        args, kwargs = mock_put.call_args
+        assert args[0] == f'{DEFAULT_CLOUD_API_URL}/api/projects/{project.name}/knowledge_bases/my_kb'
+        assert kwargs == {'json': {'knowledge_base': {'query': 'select * from tbl2 limit 1'}}}
 
+        # dataframe
         kb.insert(
             pd.DataFrame([[1, 'Alice'], [2, 'Bob']], columns=['id', 'name'])
         )
-        check_sql_call(
-            mock_post,
-            f'''INSERT INTO {project.name}.{kb.name}(id, name) VALUES (1, 'Alice'), (2, 'Bob')'''
-        )
+
+        args, kwargs = mock_put.call_args
+        assert args[0] == f'{DEFAULT_CLOUD_API_URL}/api/projects/{project.name}/knowledge_bases/my_kb'
+        assert kwargs == {'json': {
+            'knowledge_base': {'rows': [{'id': 1, 'name': 'Alice'}, {'id': 2, 'name': 'Bob'}]}
+        }}
 
         # query
         df = kb.find(query='dog', limit=5).fetch()