Skip to content

Commit 4875774

Browse files
updated the create() func for KBs
1 parent cfafb1d commit 4875774

File tree

1 file changed

+20
-23
lines changed

1 file changed

+20
-23
lines changed

mindsdb_sdk/knowledge_bases.py

Lines changed: 20 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -288,7 +288,8 @@ def get(self, name: str) -> KnowledgeBase:
288288
def create(
289289
self,
290290
name: str,
291-
model: Model = None,
291+
embedding_model: dict = None,
292+
reranking_model: dict = None,
292293
storage: Table = None,
293294
metadata_columns: list = None,
294295
content_columns: list = None,
@@ -301,7 +302,8 @@ def create(
301302
302303
>>> kb = server.knowledge_bases.create(
303304
... 'my_kb',
304-
... model=server.models.emb_model,
305+
... embedding_model={'provider': 'openai', 'model': 'text-embedding-ada-002', 'api_key': 'sk-...'},
306+
... reranking_model={'provider': 'openai', 'model': 'gpt-4', 'api_key': 'sk-...'},
305307
... storage=server.databases.pvec.tables.tbl1,
306308
... metadata_columns=['date', 'author'],
307309
... content_columns=['review', 'description'],
@@ -310,7 +312,8 @@ def create(
310312
...)
311313
312314
:param name: name of the knowledge base
313-
:param model: embedding model, optional. Default: 'sentence_transformers' will be used (defined in mindsdb server)
315+
:param embedding_model: embedding model, optional. Default: OpenAI will be the default provider
316+
:param reranking_model: reranking model, optional. Default: OpenAI will be the default provider
314317
:param storage: vector storage, optional. Default: chromadb database will be created
315318
:param metadata_columns: columns to use as metadata, optional. Default: all columns which are not content and id
316319
:param content_columns: columns to use as content, optional. Default: all columns except id column
@@ -319,30 +322,24 @@ def create(
319322
:return: created KnowledgeBase object
320323
"""
321324

322-
params_out = {}
323-
324-
if metadata_columns is not None:
325-
params_out['metadata_columns'] = metadata_columns
326-
327-
if content_columns is not None:
328-
params_out['content_columns'] = content_columns
329-
330-
if id_column is not None:
331-
params_out['id_column'] = id_column
332-
333-
if params is not None:
334-
params_out.update(params)
335-
336-
if model is not None:
337-
model = model.name
338-
339325
payload = {
340326
'name': name,
341-
'model': model,
342-
'params': params_out
343327
}
344328

345-
if storage is not None:
329+
if embedding_model:
330+
payload['embedding_model'] = embedding_model
331+
if reranking_model:
332+
payload['reranking_model'] = reranking_model
333+
if metadata_columns:
334+
payload['metadata_columns'] = metadata_columns
335+
if content_columns:
336+
payload['content_columns'] = content_columns
337+
if id_column:
338+
payload['id_column'] = id_column
339+
if params:
340+
payload['params'] = params
341+
342+
if storage:
346343
payload['storage'] = {
347344
'database': storage.db.name,
348345
'table': storage.name

0 commit comments

Comments
 (0)