Add Knowledge Bases to Core SDK (#59)

* Add Knowledge Bases to Core SDK

* EOF newline

Author: tmichaeldb

minds/client.py

@@ -2,6 +2,7 @@
 from minds.rest_api import RestAPI
 
 from minds.datasources import Datasources
+from minds.knowledge_bases import KnowledgeBases
 from minds.minds import Minds
 
 
@@ -12,5 +13,6 @@ def __init__(self, api_key, base_url=None):
         self.api = RestAPI(api_key, base_url)
 
         self.datasources = Datasources(self)
+        self.knowledge_bases = KnowledgeBases(self)
 
         self.minds = Minds(self)

minds/knowledge_bases/__init__.py

@@ -0,0 +1 @@
+from .knowledge_bases import *

minds/knowledge_bases/knowledge_bases.py

@@ -0,0 +1,162 @@
+from typing import Any, Dict, List, Optional, Union
+
+from pydantic import BaseModel
+
+from minds.rest_api import RestAPI
+
+
+class VectorStoreConfig(BaseModel):
+    '''Configuration for the underlying vector store for knowledge base embeddings'''
+    engine: str
+    connection_data: Dict[str, Any]
+    table: str = 'embeddings'
+
+
+class EmbeddingConfig(BaseModel):
+    '''Configuration for embeddings to use with underlying vector store for knowledge base'''
+    provider: str
+    model: str
+    params: Optional[Dict[str, Any]] = None
+
+
+class KnowledgeBaseConfig(BaseModel):
+    '''Configuration for a knowledge base'''
+    name: str
+    description: str
+    vector_store_config: Optional[VectorStoreConfig] = None
+    embedding_config: Optional[EmbeddingConfig] = None
+
+
+class KnowledgeBaseDocument(BaseModel):
+    '''Represents a document that can be inserted into a knowledge base'''
+    id: Union[int, str]
+    content: str
+    metadata: Optional[Dict[str, Any]] = {}
+
+
+class KnowledgeBase:
+    def __init__(self, name, api: RestAPI):
+        self.name = name
+        self.api = api
+
+    def insert_from_select(self, query: str):
+        '''
+        Inserts select content of a connected datasource into this knowledge base
+
+        :param query: The SQL SELECT query to use to retrieve content to be inserted
+        '''
+        update_request = {
+            'query': query
+        }
+        _ = self.api.put(f'/knowledge_bases/{self.name}', data=update_request)
+
+    def insert_documents(self, documents: List[KnowledgeBaseDocument]):
+        '''
+        Inserts documents directly into this knowledge base
+
+        :param documents: The documents to insert
+        '''
+        update_request = {
+            'rows': [d.model_dump() for d in documents]
+        }
+        _ = self.api.put(f'/knowledge_bases/{self.name}', data=update_request)
+
+    def insert_urls(self, urls: List[str]):
+        '''
+        Crawls URLs & inserts the retrieved webpages into this knowledge base
+
+        :param urls: Valid URLs to crawl & insert
+        '''
+        update_request = {
+            'urls': urls
+        }
+        _ = self.api.put(f'/knowledge_bases/{self.name}', data=update_request)
+
+    def insert_files(self, files: List[str]):
+        '''
+        Inserts files that have already been uploaded to MindsDB into this knowledge base
+
+        :param files: Names of preuploaded files to insert
+        '''
+        update_request = {
+            'files': files
+        }
+        _ = self.api.put(f'/knowledge_bases/{self.name}', data=update_request)
+
+
+class KnowledgeBases:
+    def __init__(self, client):
+        self.api = client.api
+
+    def create(self, config: KnowledgeBaseConfig) -> KnowledgeBase:
+        '''
+        Create new knowledge base and return it
+
+        :param config: knowledge base configuration, properties:
+           - name: str, name of knowledge base
+           - description: str, description of the knowledge base. Used by minds to know what data can be retrieved.
+           - vector_store_config: VectorStoreConfig, configuration for embeddings vector store.
+           - embedding_config: EmbeddingConfig, configuration for embeddings.
+        :return: knowledge base object
+        '''
+        create_request = {
+            'name': config.name,
+            'description': config.description
+        }
+        if config.vector_store_config is not None:
+            vector_store_data = {
+                'engine': config.vector_store_config.engine,
+                'connection_data': config.vector_store_config.connection_data
+            }
+            create_request['vector_store'] = vector_store_data
+        if config.embedding_config is not None:
+            embedding_data = {
+                'provider': config.embedding_config.provider,
+                'name': config.embedding_config.model
+            }
+            if config.embedding_config.params is not None:
+                embedding_data.update(config.embedding_config.params)
+            create_request['embedding_model'] = embedding_data
+
+        _ = self.api.post('/knowledge_bases', data=create_request)
+        return self.get(config.name)
+
+    def list(self) -> List[KnowledgeBase]:
+        '''
+        Returns list of knowledge bases
+
+        :return: iterable knowledge bases
+        '''
+
+        list_knowledge_bases_response = self.api.get('/knowledge_bases')
+        knowledge_bases = list_knowledge_bases_response.json()
+
+        all_knowledge_bases = []
+        for knowledge_base in knowledge_bases:
+            all_knowledge_bases.append(KnowledgeBase(knowledge_base['name'], self.api))
+        return all_knowledge_bases
+
+    def get(self, name: str) -> KnowledgeBase:
+        '''
+        Get knowledge base by name
+
+        :param name: name of knowledge base
+        :return: knowledge base object
+        '''
+
+        knowledge_base_response = self.api.get(f'/knowledge_bases/{name}')
+        knowledge_base = knowledge_base_response.json()
+        return KnowledgeBase(knowledge_base['name'], self.api)
+
+    def drop(self, name: str, force=False):
+        '''
+        Drop knowledge base by name
+
+        :param name: name of knowledge base
+        :param force: if True - remove from all minds, default: False
+        '''
+        data = None
+        if force:
+            data = {'cascade': True}
+
+        self.api.delete(f'/knowledge_bases/{name}', data=data)

minds/minds.py

@@ -3,6 +3,7 @@
 import minds.utils as utils
 import minds.exceptions as exc
 from minds.datasources import Datasource, DatabaseConfig
+from minds.knowledge_bases import KnowledgeBase, KnowledgeBaseConfig
 
 DEFAULT_PROMPT_TEMPLATE = 'Use your database tools to answer the user\'s question: {{question}}'
 
@@ -13,6 +14,7 @@ def __init__(
         provider=None,
         parameters=None,
         datasources=None,
+        knowledge_bases=None,
         created_at=None,
         updated_at=None,
         **kwargs
@@ -36,6 +38,7 @@ def __init__(
             base_url=base_url
         )
         self.datasources = datasources
+        self.knowledge_bases = knowledge_bases
 
     def __repr__(self):
         return (f'Mind(name={self.name}, '
@@ -44,6 +47,7 @@ def __repr__(self):
                 f'created_at="{self.created_at}", '
                 f'updated_at="{self.updated_at}", '
                 f'parameters={self.parameters}, '
+                f'knowledge_bases={self.knowledge_bases}, '
                 f'datasources={self.datasources})')
 
     def update(
@@ -53,6 +57,7 @@ def update(
         provider=None,
         prompt_template=None,
         datasources=None,
+        knowledge_bases=None,
         parameters=None,
     ):
         """
@@ -65,11 +70,17 @@ def update(
          - Datasource object (minds.datasources.Database)
          - datasource config (minds.datasources.DatabaseConfig), in this case datasource will be created
 
+        Knowledge base can be passed as
+         - name, str
+         - KnowledgeBase object (minds.knowledge_bases.KnowledgeBase)
+         - Knowledge base config (minds.knowledge_bases.KnowledgeBaseConfig), in this case knowledge base will be created
+
         :param name: new name of the mind, optional
         :param model_name: new llm model name, optional
         :param provider: new llm provider, optional
         :param prompt_template: new prompt template, optional
         :param datasources: alter list of datasources used by mind, optional
+        :param knowledge_bases: alter list of knowledge bases used by mind, optional
         :param parameters, dict: alter other parameters of the mind, optional
         """
         data = {}
@@ -84,6 +95,13 @@ def update(
                 ds_names.append(ds)
             data['datasources'] = ds_names
 
+        if knowledge_bases is not None:
+            kb_names = []
+            for kb in knowledge_bases:
+                kb = self.client.minds._check_knowledge_base(kb)
+                kb_names.append(kb)
+            data['knowledge_bases'] = kb_names
+
         if name is not None:
             data['name'] = name
         if model_name is not None:
@@ -149,6 +167,50 @@ def del_datasource(self, datasource: Union[Datasource, str]):
 
         self.datasources = updated.datasources
 
+    def add_knowledge_base(self, knowledge_base: Union[str, KnowledgeBase, KnowledgeBaseConfig]):
+        """
+        Add knowledge base to mind
+        Knowledge base can be passed as
+         - name, str
+         - Knowledge base object (minds.knowledge_bases.KnowledgeBase)
+         - Knowledge base config (minds.knowledge_bases.KnowledgeBaseConfig), in this case knowledge base will be created
+
+        :param knowledge_base: input knowledge base
+        """
+
+        kb_name = self.client.minds._check_knowledge_base(knowledge_base)
+
+        self.api.post(
+            f'/projects/{self.project}/minds/{self.name}/knowledge_bases',
+            data={
+                'name': kb_name,
+            }
+        )
+        updated = self.client.minds.get(self.name)
+
+        self.knowledge_bases = updated.knowledge_bases
+
+    def del_knowledge_base(self, knowledge_base: Union[KnowledgeBase, str]):
+        """
+        Delete knowledge base from mind
+
+        Knowledge base can be passed as
+         - name, str
+         - KnowledgeBase object (minds.knowledge_bases.KnowledgeBase)
+
+        :param knowledge_base: Knowledge base to delete
+        """
+        if isinstance(knowledge_base, KnowledgeBase):
+            knowledge_base = knowledge_base.name
+        elif not isinstance(knowledge_base, str):
+            raise ValueError(f'Unknown type of knowledge base: {knowledge_base}')
+        self.api.delete(
+            f'/projects/{self.project}/minds/{self.name}/knowledge_bases/{knowledge_base}',
+        )
+        updated = self.client.minds.get(self.name)
+
+        self.knowledge_bases = updated.knowledge_bases
+
     def completion(self, message: str, stream: bool = False) -> Union[str, Iterable[object]]:
         """
         Call mind completion
@@ -221,12 +283,28 @@ def _check_datasource(self, ds) -> str:
             raise ValueError(f'Unknown type of datasource: {ds}')
         return ds
 
+    def _check_knowledge_base(self, knowledge_base) -> str:
+        if isinstance(knowledge_base, KnowledgeBase):
+            knowledge_base = knowledge_base.name
+        elif isinstance(knowledge_base, KnowledgeBaseConfig):
+            # if not exists - create
+            try:
+                self.client.knowledge_bases.get(knowledge_base.name)
+            except exc.ObjectNotFound:
+                self.client.knowledge_bases.create(knowledge_base)
+
+            knowledge_base = knowledge_base.name
+        elif not isinstance(knowledge_base, str):
+            raise ValueError(f'Unknown type of knowledge base: {knowledge_base}')
+        return knowledge_base
+
     def create(
         self, name,
         model_name=None,
         provider=None,
         prompt_template=None,
         datasources=None,
+        knowledge_bases=None,
         parameters=None,
         replace=False,
         update=False,
@@ -239,11 +317,17 @@ def create(
          - Datasource object (minds.datasources.Database)
          - datasource config (minds.datasources.DatabaseConfig), in this case datasource will be created
 
+        Knowledge base can be passed as
+         - name, str
+         - KnowledgeBase object (minds.knowledge_bases.KnowledgeBase)
+         - Knowledge base config (minds.knowledge_bases.KnowledgeBaseConfig), in this case knowledge base will be created
+
         :param name: name of the mind
         :param model_name: llm model name, optional
         :param provider: llm provider, optional
         :param prompt_template: instructions to llm, optional
         :param datasources: list of datasources used by mind, optional
+        :param knowledge_bases: alter list of knowledge bases used by mind, optional
         :param parameters, dict: other parameters of the mind, optional
         :param replace: if true - to remove existing mind, default is false
         :param update: if true - to update mind if exists, default is false
@@ -267,6 +351,12 @@ def create(
 
                 ds_names.append(ds)
 
+        kb_names = []
+        if knowledge_bases:
+            for kb in knowledge_bases:
+                kb = self._check_knowledge_base(kb)
+                kb_names.append(kb)
+
         if parameters is None:
             parameters = {}
 
@@ -290,6 +380,7 @@ def create(
                 'provider': provider,
                 'parameters': parameters,
                 'datasources': ds_names,
+                'knowledge_bases': kb_names
             }
         )
         mind = self.get(name)