Merge pull request #192 from mindsdb/feature/update_kb_apis

Updated the SDK Operations for Knowledge Bases

Author: MinuraPunchihewa

mindsdb_sdk/knowledge_bases.py

@@ -4,7 +4,6 @@
 
 import pandas as pd
 
-from mindsdb_sql_parser.ast.mindsdb import CreateKnowledgeBase, DropKnowledgeBase
 from mindsdb_sql_parser.ast import Identifier, Star, Select, BinaryOperation, Constant, Insert
 
 from mindsdb_sdk.utils.sql import dict_to_binary_op, query_to_native_query
@@ -57,9 +56,14 @@ def __init__(self, api, project, data: dict):
             table = Table(database, data['vector_database_table'])
             self.storage = table
 
-        self.model = None
-        if data['embedding_model'] is not None:
-            self.model = Model(self.project, {'name': data['embedding_model']})
+        # models
+        self.embedding_model = data.get('embedding_model', {})
+        self.reranking_model = data.get('reranking_model', {})
+
+        # columns
+        self.metadata_columns = data.get('metadata_columns', [])
+        self.content_columns = data.get('content_columns', [])
+        self.id_column = data.get('id_column', None)
 
         params = data.get('params', {})
         if isinstance(params, str):
@@ -68,11 +72,6 @@ def __init__(self, api, project, data: dict):
             except json.JSONDecodeError:
                 params = {}
 
-        # columns
-        self.metadata_columns = params.pop('metadata_columns', [])
-        self.content_columns = params.pop('content_columns', [])
-        self.id_column = params.pop('id_column', None)
-
         self.params = params
 
         # query behavior
@@ -311,7 +310,8 @@ def get(self, name: str) -> KnowledgeBase:
     def create(
         self,
         name: str,
-        model: Model = None,
+        embedding_model: dict = None,
+        reranking_model: dict = None,
         storage: Table = None,
         metadata_columns: list = None,
         content_columns: list = None,
@@ -324,7 +324,8 @@ def create(
 
         >>> kb = server.knowledge_bases.create(
         ...   'my_kb',
-        ...   model=server.models.emb_model,
+        ...   embedding_model={'provider': 'openai', 'model': 'text-embedding-ada-002', 'api_key': 'sk-...'},
+        ...   reranking_model={'provider': 'openai', 'model': 'gpt-4', 'api_key': 'sk-...'},
         ...   storage=server.databases.pvec.tables.tbl1,
         ...   metadata_columns=['date', 'author'],
         ...   content_columns=['review', 'description'],
@@ -333,7 +334,8 @@ def create(
         ...)
 
         :param name: name of the knowledge base
-        :param model: embedding model, optional. Default: 'sentence_transformers' will be used (defined in mindsdb server)
+        :param embedding_model: embedding model, optional. Default: OpenAI will be the default provider
+        :param reranking_model: reranking model, optional. Default: OpenAI will be the default provider
         :param storage: vector storage, optional. Default: chromadb database will be created
         :param metadata_columns: columns to use as metadata, optional. Default: all columns which are not content and id
         :param content_columns: columns to use as content, optional. Default: all columns except id column
@@ -342,30 +344,24 @@ def create(
         :return: created KnowledgeBase object
         """
 
-        params_out = {}
-
-        if metadata_columns is not None:
-            params_out['metadata_columns'] = metadata_columns
-
-        if content_columns is not None:
-            params_out['content_columns'] = content_columns
-
-        if id_column is not None:
-            params_out['id_column'] = id_column
-
-        if params is not None:
-            params_out.update(params)
-
-        if model is not None:
-            model = model.name
-
         payload = {
             'name': name,
-            'model': model,
-            'params': params_out
         }
 
-        if storage is not None:
+        if embedding_model:
+            payload['embedding_model'] = embedding_model
+        if reranking_model:
+            payload['reranking_model'] = reranking_model
+        if metadata_columns:
+            payload['metadata_columns'] = metadata_columns
+        if content_columns:
+            payload['content_columns'] = content_columns
+        if id_column:
+            payload['id_column'] = id_column
+        if params:
+            payload['params'] = params
+
+        if storage:
             payload['storage'] = {
                 'database': storage.db.name,
                 'table': storage.name

tests/test_sdk.py

@@ -775,7 +775,7 @@ def check_project(self, project, database):
 
         self.check_project_models_versions(project, database)
 
-        kb = self.check_project_kb(project, model, database)
+        kb = self.check_project_kb(project, database)
 
         self.check_project_jobs(project, model, database, kb)
 
@@ -1143,12 +1143,21 @@ def check_project_jobs(self, project, model, database, kb, mock_post):
     @patch('requests.Session.post')
     @patch('requests.Session.delete')
     @patch('requests.Session.get')
-    def check_project_kb(self, project, model, database, mock_get, mock_del, mock_post, mock_put):
+    def check_project_kb(self, project, database, mock_get, mock_del, mock_post, mock_put):
 
         response_mock(mock_post, pd.DataFrame([{
             'NAME': 'my_kb',
             'PROJECT': 'mindsdb',
-            'MODEL': 'openai_emb',
+            'EMBEDDING_MODEL': {
+                'PROVIDER': 'openai',
+                'MODEL_NAME': 'openai_emb',
+                'API_KEY': 'sk-...'
+            },
+            'RERANKING_MODEL': {
+                'PROVIDER': 'openai',
+                'MODEL_NAME': 'openai_rerank',
+                'API_KEY': 'sk-...'
+            },
             'STORAGE': 'pvec.tbl1',
             'PARAMS': {"id_column": "num"},
         }]))
@@ -1157,7 +1166,16 @@ def check_project_kb(self, project, model, database, mock_get, mock_del, mock_po
             'id': 1,
             'name': 'my_kb',
             'project_id': 1,
-            'embedding_model': 'openai_emb',
+            'embedding_model': {
+                'provider': 'openai',
+                'model_name': 'openai_emb',
+                'api_key': 'sk-...'
+            },
+            'reranking_model': {
+                'provider': 'openai',
+                'model_name': 'openai_rerank',
+                'api_key': 'sk-...'
+            },
             'vector_database': 'pvec',
             'vector_database_table': 'tbl1',
             'updated_at': '2024-10-04 10:55:25.350799',
@@ -1176,8 +1194,8 @@ def check_project_kb(self, project, model, database, mock_get, mock_del, mock_po
 
         assert kb.name == 'my_kb'
 
-        assert isinstance(kb.model, Model)
-        assert kb.model.name == 'openai_emb'
+        assert kb.embedding_model['model_name'] == 'openai_emb'
+        assert kb.reranking_model['model_name'] == 'openai_rerank'
 
         assert isinstance(kb.storage, Table)
         assert kb.storage.name == 'tbl1'
@@ -1189,7 +1207,8 @@ def check_project_kb(self, project, model, database, mock_get, mock_del, mock_po
         str(kb)
         assert kb.name == 'my_kb'
         assert kb.storage.db.name == 'pvec'
-        assert kb.model.name == 'openai_emb'
+        assert kb.embedding_model['model_name'] == 'openai_emb'
+        assert kb.reranking_model['model_name'] == 'openai_rerank'
 
         # --- insert ---
 
@@ -1232,18 +1251,36 @@ def check_project_kb(self, project, model, database, mock_get, mock_del, mock_po
         # create 1
         project.knowledge_bases.create(
             name='kb2',
-            model=model,
+            embedding_model={
+                'provider': 'openai',
+                'model_name': 'openai_emb',
+                'api_key': 'sk-...'
+            },
+            reranking_model={
+                'provider': 'openai',
+                'model_name': 'openai_rerank',
+                'api_key': 'sk-...'
+            },
             metadata_columns=['date', 'author'],
             params={'k': 'v'}
         )
         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,
+            'embedding_model': {
+                'provider': 'openai',
+                'model_name': 'openai_emb',
+                'api_key': 'sk-...'
+            },
+            'reranking_model': {
+                'provider': 'openai',
+                'model_name': 'openai_rerank',
+                'api_key': 'sk-...'
+            },
+            'metadata_columns': ['date', 'author'],
             'params': {
                 'k': 'v',
-                'metadata_columns': ['date', 'author']
             }
         }}}
 
@@ -1259,11 +1296,8 @@ def check_project_kb(self, project, model, database, mock_get, mock_del, mock_po
         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'
-            },
+            'content_columns': ['review'],
+            'id_column': 'num',
             'storage': {
                 'database': database.name,
                 'table': 'tbl1'