Repository: mindsdb/minds_python_sdk

Issue:
Minds Java SDK
## Short description of current behavior

The goal of this issue is to develop a new Java SDK for Minds.

## Steps 🕵️‍♂️ 🕵️‍♀️

- You can familiarize yourself with [Minds](https://mindsdb.com/).
- Review the existing [Python SDK](https://github.com/mindsdb/minds_python_sdk) as a reference.
- Create a new repository in your account and submit a Pull Request (PR) to this repository, adding the new SDK to the Other SDKs section of the README.
- Multiple contributors are welcome, but the SDK with the highest standards—such as clear documentation, robust tests, and overall code quality—will be accepted.

## 📢 Additional rewards 🏆

This issue brings 1️⃣ 2️⃣ 0️⃣ points for SWAG and entry into the draw for a 👇 :

3️⃣ ❌ 💻 [Razer Blade 16 Gaming Laptop](https://www.razer.com/gaming-laptops/Razer-Blade-16/RZ09-0510TEN3-R3U1)
👕 🐻 [MindsDB Swag](https://store.covver.io/mindsdb/collections/swag-collection)
ℹ️ For more info check out https://mindsdb.com/hacktoberfest 👈

Discussion:
@Better-Boy: Since this needs a separate repo, minds-java-sdk - https://github.com/Better-Boy/minds-java-sdk.

For testing, I'll need an account with minds.@ZoranPandovski: @Better-Boy Since you already made some progress, I will assign this to you. @chandrevdw31 Will send you an API key on slack@quest-bot: ### ![New Quest! image](https://quine-assets.s3.amazonaws.com/assets/quest-bot/quest-generic/opening28.svg) New Quest!

*A new Quest has been launched in @mindsdb’s repo.
Merge a PR that solves this issue to loot the Quest and earn your reward.*

---

Some loot has been stashed in this issue to reward the solver!

🗡 Comment `@quest-bot embark` to check-in for this Quest and start solving the issue. Other solvers will be notified!

⚔️ When you submit a PR, comment `@quest-bot loot #23` to link your PR to this Quest.

Questions? Check out [the docs](https://docs.quira.sh/for-developers/quests/solver-quests).@Better-Boy: @ZoranPandovski I'm done with the java sdk.  Please review once.

The following are done:
1. chat completion done
2. Unit test done
3. Integration test done
4. Documentation done
5. Readme update done

Open to any feedback@Better-Boy: @quest-bot embark@quest-bot: @Better-Boy has embarked on their Quest. 🗡

- @Better-Boy has been on GitHub since **2016**.
- They have merged **55** public PRs in that time.
- Their swords are blessed with `Python` and `HTML` magic ✨
- They have contributed to this repo before.

_This is not an assignment to the issue. Please check the repo’s contribution guidelines before submitting a PR._

Questions? Check out [the docs](https://docs.quira.sh/for-developers/quests/solver-quests).@quest-bot: 🧚 @Better-Boy has submitted PR https://github.com/mindsdb/minds_python_sdk/issues/35 and is claiming the loot.

Keep up the pace, or you'll be left in the shadows.

Questions? Check out [the docs](https://docs.quira.sh/for-developers/quests/solver-quests).

Solution patch:
diff --git a/README.md b/README.md
index 3cb4808..e953eb2 100644
--- a/README.md
+++ b/README.md
@@ -150,10 +150,10 @@ client.datasources.drop('my_datasource')
 >Note: The SDK currently does not support automatically removing a data source if it is no longer connected to any mind.
 
 ### Other SDKs
+- [Java-SDK](https://github.com/Better-Boy/minds-java-sdk)
 - [Ruby-SDK](https://github.com/tungnt1203/minds_ruby_sdk)
 - [Dart-SDK](https://github.com/ArnavK-09/mdb_dart)
 
 #### Command Line Tools
 - [Minds CLI](https://github.com/Better-Boy/minds-cli-sdk)
 
-
diff --git a/minds/__about__.py b/minds/__about__.py
index cf14e91..8309cc7 100644
--- a/minds/__about__.py
+++ b/minds/__about__.py
@@ -1,6 +1,6 @@
 __title__ = 'minds_sdk'
 __package_name__ = 'minds'
-__version__ = '1.0.9'
+__version__ = '1.2.0'
 __description__ = 'An AI-Data Mind is an LLM with the built-in power to answer data questions for Agents'
 __email__ = 'hello@mindsdb.com'
 __author__ = 'MindsDB Inc'
diff --git a/minds/client.py b/minds/client.py
index 738cf71..b7b7fcd 100644
--- a/minds/client.py
+++ b/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)
diff --git a/minds/knowledge_bases/__init__.py b/minds/knowledge_bases/__init__.py
new file mode 100644
index 0000000..7e08d9a
--- /dev/null
+++ b/minds/knowledge_bases/__init__.py
@@ -0,0 +1 @@
+from .knowledge_bases import *
diff --git a/minds/knowledge_bases/knowledge_bases.py b/minds/knowledge_bases/knowledge_bases.py
new file mode 100644
index 0000000..337a5b4
--- /dev/null
+++ b/minds/knowledge_bases/knowledge_bases.py
@@ -0,0 +1,175 @@
+from typing import Any, Dict, List, Optional, Union
+
+from pydantic import BaseModel
+
+from minds.knowledge_bases.preprocessing import PreprocessingConfig
+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
+    # Params to apply to retrieval pipeline.
+    params: Optional[Dict] = 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, preprocessing_config: PreprocessingConfig = None):
+        '''
+        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
+        }
+        if preprocessing_config is not None:
+            update_request['preprocessing'] = preprocessing_config.model_dump()
+        _ = self.api.put(f'/knowledge_bases/{self.name}', data=update_request)
+
+    def insert_documents(self, documents: List[KnowledgeBaseDocument], preprocessing_config: PreprocessingConfig = None):
+        '''
+        Inserts documents directly into this knowledge base
+
+        :param documents: The documents to insert
+        '''
+        update_request = {
+            'rows': [d.model_dump() for d in documents]
+        }
+        if preprocessing_config is not None:
+            update_request['preprocessing'] = preprocessing_config.model_dump()
+        _ = self.api.put(f'/knowledge_bases/{self.name}', data=update_request)
+
+    def insert_urls(self, urls: List[str], preprocessing_config: PreprocessingConfig = None):
+        '''
+        Crawls URLs & inserts the retrieved webpages into this knowledge base
+
+        :param urls: Valid URLs to crawl & insert
+        '''
+        update_request = {
+            'urls': urls
+        }
+        if preprocessing_config is not None:
+            update_request['preprocessing'] = preprocessing_config.model_dump()
+        _ = self.api.put(f'/knowledge_bases/{self.name}', data=update_request)
+
+    def insert_files(self, files: List[str], preprocessing_config: PreprocessingConfig = None):
+        '''
+        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
+        }
+        if preprocessing_config is not None:
+            update_request['preprocessing'] = preprocessing_config.model_dump()
+        _ = 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
+        if config.params is not None:
+            create_request['params'] = config.params
+
+        _ = 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)
diff --git a/minds/knowledge_bases/preprocessing.py b/minds/knowledge_bases/preprocessing.py
new file mode 100644
index 0000000..7ea4716
--- /dev/null
+++ b/minds/knowledge_bases/preprocessing.py
@@ -0,0 +1,78 @@
+from typing import Any, Dict, List, Literal, Optional
+
+from pydantic import BaseModel, Field, model_validator
+
+
+DEFAULT_LLM_MODEL = 'gpt-4o'
+DEFAULT_LLM_MODEL_PROVIDER = 'openai'
+
+
+class TextChunkingConfig(BaseModel):
+    '''Configuration for chunking text content before they are inserted into a knowledge base'''
+    separators: List[str] = Field(
+        default=['\n\n', '\n', ' ', ''],
+        description='List of separators to use for splitting text, in order of priority'
+    )
+    chunk_size: int = Field(
+        default=1000,
+        description='The target size of each text chunk',
+        gt=0
+    )
+    chunk_overlap: int = Field(
+        default=200,
+        description='The number of characters to overlap between chunks',
+        ge=0
+    )
+
+
+class LLMConfig(BaseModel):
+    model_name: str = Field(default=DEFAULT_LLM_MODEL, description='LLM model to use for context generation')
+    provider: str = Field(default=DEFAULT_LLM_MODEL_PROVIDER, description='LLM model provider to use for context generation')
+    params: Dict[str, Any] = Field(default={}, description='Additional parameters to pass in when initializing the LLM')
+
+
+class ContextualConfig(BaseModel):
+    '''Configuration specific to contextual preprocessing'''
+    llm_config: LLMConfig = Field(
+        default=LLMConfig(),
+        description='LLM configuration to use for context generation'
+    )
+    context_template: Optional[str] = Field(
+        default=None,
+        description='Custom template for context generation'
+    )
+    chunk_size: int = Field(
+        default=1000,
+        description='The target size of each text chunk',
+        gt=0
+    )
+    chunk_overlap: int = Field(
+        default=200,
+        description='The number of characters to overlap between chunks',
+        ge=0
+    )
+
+
+class PreprocessingConfig(BaseModel):
+    '''Complete preprocessing configuration'''
+    type: Literal['contextual', 'text_chunking'] = Field(
+        default='text_chunking',
+        description='Type of preprocessing to apply'
+    )
+    contextual_config: Optional[ContextualConfig] = Field(
+        default=None,
+        description='Configuration for contextual preprocessing'
+    )
+    text_chunking_config: Optional[TextChunkingConfig] = Field(
+        default=None,
+        description='Configuration for text chunking preprocessing'
+    )
+
+    @model_validator(mode='after')
+    def validate_config_presence(self) -> 'PreprocessingConfig':
+        '''Ensure the appropriate config is present for the chosen type'''
+        if self.type == 'contextual' and not self.contextual_config:
+            self.contextual_config = ContextualConfig()
+        if self.type == 'text_chunking' and not self.text_chunking_config:
+            self.text_chunking_config = TextChunkingConfig()
+        return self
diff --git a/minds/minds.py b/minds/minds.py
index 82508a3..43ee967 100644
--- a/minds/minds.py
+++ b/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)


Test patch:
diff --git a/tests/unit/test_unit.py b/tests/unit/test_unit.py
index a3125d1..daeac79 100644
--- a/tests/unit/test_unit.py
+++ b/tests/unit/test_unit.py
@@ -4,16 +4,20 @@
 
 
 from minds.datasources.examples import example_ds
+from minds.knowledge_bases import EmbeddingConfig, KnowledgeBaseConfig, VectorStoreConfig
+
 
 def get_client():
     from minds.client import Client
     return Client(API_KEY)
 
+
 from minds import rest_api
 
 # patch _raise_for_status
 rest_api._raise_for_status = Mock()
-#
+
+
 def response_mock(mock, data):
     def side_effect(*args, **kwargs):
         r_mock = Mock()
@@ -44,6 +48,7 @@ def test_create_datasources(self, mock_del, mock_post, mock_put, mock_get):
         response_mock(mock_get, example_ds.model_dump())
 
         ds = client.datasources.create(example_ds)
+
         def check_ds_created(ds, mock_post, url):
             self._compare_ds(ds, example_ds)
             args, kwargs = mock_post.call_args
@@ -76,7 +81,7 @@ def test_delete_datasource(self, mock_del):
         client.datasources.drop('ds_name')
 
         args, _ = mock_del.call_args
-        assert args[0].endswith(f'/api/datasources/ds_name')
+        assert args[0].endswith('/api/datasources/ds_name')
 
     @patch('requests.get')
     def test_list_datasources(self, mock_get):
@@ -89,7 +94,184 @@ def test_list_datasources(self, mock_get):
         self._compare_ds(ds, example_ds)
 
         args, _ = mock_get.call_args
-        assert args[0].endswith(f'/api/datasources')
+        assert args[0].endswith('/api/datasources')
+
+
+class TestKnowledgeBases:
+
+    def _compare_knowledge_base(self, knowledge_base, config):
+        assert knowledge_base.name == config.name
+
+    @patch('requests.get')
+    @patch('requests.post')
+    def test_create_knowledge_bases(self, mock_post, mock_get):
+        client = get_client()
+
+        test_embedding_config = EmbeddingConfig(
+            provider='openai',
+            model='gpt-4o',
+            params={
+                'k1': 'v1'
+            }
+        )
+        test_vector_store_connection_data = {
+            'user': 'test_user',
+            'password': 'test_password',
+            'host': 'boop.mindsdb.com',
+            'port': '5432',
+            'database': 'test',
+        }
+        test_vector_store_config = VectorStoreConfig(
+            engine='pgvector',
+            connection_data=test_vector_store_connection_data,
+            table='test_table'
+        )
+        test_knowledge_base_config = KnowledgeBaseConfig(
+            name='test_kb',
+            description='Test knowledge base',
+            vector_store_config=test_vector_store_config,
+            embedding_config=test_embedding_config,
+            params={
+                'k1': 'v1'
+            }
+        )
+        response_mock(mock_get, test_knowledge_base_config.model_dump())
+
+        created_knowledge_base = client.knowledge_bases.create(test_knowledge_base_config)
+        self._compare_knowledge_base(created_knowledge_base, test_knowledge_base_config)
+
+        args, kwargs = mock_post.call_args
+
+        assert kwargs['headers'] == {'Authorization': 'Bearer ' + API_KEY}
+
+        expected_create_request = {
+            'name': test_knowledge_base_config.name,
+            'description': test_knowledge_base_config.description,
+            'vector_store': {
+                'engine': test_vector_store_config.engine,
+                'connection_data': test_vector_store_config.connection_data
+            },
+            'embedding_model': {
+                'provider': test_embedding_config.provider,
+                'name': test_embedding_config.model,
+                'k1': 'v1'
+            },
+            'params': {
+                'k1': 'v1'
+            }
+        }
+
+        assert kwargs['json'] == expected_create_request
+        assert args[0] == 'https://mdb.ai/api/knowledge_bases'
+
+    @patch('requests.get')
+    def test_get_knowledge_base(self, mock_get):
+        client = get_client()
+
+        test_embedding_config = EmbeddingConfig(
+            provider='openai',
+            model='gpt-4o',
+            params={
+                'k1': 'v1'
+            }
+        )
+        test_vector_store_connection_data = {
+            "user": "test_user",
+            "password": "test_password",
+            "host": "boop.mindsdb.com",
+            "port": "5432",
+            "database": "test",
+        }
+        test_vector_store_config = VectorStoreConfig(
+            engine='pgvector',
+            connection_data=test_vector_store_connection_data,
+            table='test_table'
+        )
+        test_knowledge_base_config = KnowledgeBaseConfig(
+            name='test_kb',
+            description='Test knowledge base',
+            vector_store_config=test_vector_store_config,
+            embedding_config=test_embedding_config
+        )
+
+        # Expected response from MindsDB API server.
+        get_response = {
+            'created_at': '2024-11-15',
+            'embedding_model': 'test_kb_embeddings',
+            'id': 1,
+            'name': 'test_kb',
+            'params': {},
+            'project_id': 1,
+            'updated_at': '2024-11-15',
+            'vector_database': 'test_kb_vector_store',
+            'vector_database_table': 'test_table'
+        }
+        response_mock(mock_get, get_response)
+        get_knowledge_base = client.knowledge_bases.get(test_knowledge_base_config.name)
+        self._compare_knowledge_base(get_knowledge_base, test_knowledge_base_config)
+
+        args, _ = mock_get.call_args
+        assert args[0].endswith(f'/api/knowledge_bases/{test_knowledge_base_config.name}')
+
+    @patch('requests.delete')
+    def test_delete_knowledge_base(self, mock_del):
+        client = get_client()
+
+        client.knowledge_bases.drop('test_kb')
+
+        args, _ = mock_del.call_args
+        assert args[0].endswith('/api/knowledge_bases/test_kb')
+
+    @patch('requests.get')
+    def test_list_knowledge_bases(self, mock_get):
+        client = get_client()
+
+        test_embedding_config = EmbeddingConfig(
+            provider='openai',
+            model='gpt-4o',
+            params={
+                'k1': 'v1'
+            }
+        )
+        test_vector_store_connection_data = {
+            "user": "test_user",
+            "password": "test_password",
+            "host": "boop.mindsdb.com",
+            "port": "5432",
+            "database": "test",
+        }
+        test_vector_store_config = VectorStoreConfig(
+            engine='pgvector',
+            connection_data=test_vector_store_connection_data,
+            table='test_table'
+        )
+        test_knowledge_base_config = KnowledgeBaseConfig(
+            name='test_kb',
+            description='Test knowledge base',
+            vector_store_config=test_vector_store_config,
+            embedding_config=test_embedding_config
+        )
+
+        # Expected response from MindsDB API server.
+        get_response = {
+            'created_at': '2024-11-15',
+            'embedding_model': 'test_kb_embeddings',
+            'id': 1,
+            'name': 'test_kb',
+            'params': {},
+            'project_id': 1,
+            'updated_at': '2024-11-15',
+            'vector_database': 'test_kb_vector_store',
+            'vector_database_table': 'test_table'
+        }
+        response_mock(mock_get, [get_response])
+        knowledge_base_list = client.knowledge_bases.list()
+        assert len(knowledge_base_list) == 1
+        knowledge_base = knowledge_base_list[0]
+        self._compare_knowledge_base(knowledge_base, test_knowledge_base_config)
+
+        args, _ = mock_get.call_args
+        assert args[0].endswith('/api/knowledge_bases')
 
 
 class TestMinds:
@@ -98,6 +280,7 @@ class TestMinds:
         'model_name': 'gpt-4o',
         'name': 'test_mind',
         'datasources': ['example_ds'],
+        'knowledge_bases': ['example_kb'],
         'provider': 'openai',
         'parameters': {
             'prompt_template': "Answer the user's question"
@@ -122,13 +305,15 @@ def test_create(self, mock_del, mock_post, mock_put, mock_get):
         mind_name = 'test_mind'
         prompt_template = 'always agree'
         datasources = ['my_ds']
+        knowledge_bases = ['example_kb']
         provider = 'openai'
 
         response_mock(mock_get, self.mind_json)
         create_params = {
             'name': mind_name,
             'prompt_template': prompt_template,
-            'datasources': datasources
+            'datasources': datasources,
+            'knowledge_bases': knowledge_bases
         }
         mind = client.minds.create(**create_params)
 
@@ -136,7 +321,7 @@ def check_mind_created(mind, mock_post, create_params, url):
             args, kwargs = mock_post.call_args
             assert args[0].endswith(url)
             request = kwargs['json']
-            for key in ('name', 'datasources', 'provider', 'model_name'),:
+            for key in ('name', 'datasources', 'knowledge_bases', 'provider', 'model_name'),:
                 assert request.get(key) == create_params.get(key)
             assert create_params.get('prompt_template') == request.get('parameters', {}).get('prompt_template')
 
@@ -177,6 +362,7 @@ def test_update(self, mock_patch, mock_get):
         update_params = dict(
             name='mind_name2',
             datasources=['ds_name'],
+            knowledge_bases=['kb_name'],
             provider='ollama',
             model_name='llama',
             parameters={
@@ -200,7 +386,7 @@ def test_get(self, mock_get):
         self.compare_mind(mind, self.mind_json)
 
         args, _ = mock_get.call_args
-        assert args[0].endswith(f'/api/projects/mindsdb/minds/my_mind')
+        assert args[0].endswith('/api/projects/mindsdb/minds/my_mind')
 
     @patch('requests.get')
     def test_list(self, mock_get):
@@ -212,7 +398,7 @@ def test_list(self, mock_get):
         self.compare_mind(minds_list[0], self.mind_json)
 
         args, _ = mock_get.call_args
-        assert args[0].endswith(f'/api/projects/mindsdb/minds')
+        assert args[0].endswith('/api/projects/mindsdb/minds')
 
     @patch('requests.delete')
     def test_delete(self, mock_del):
@@ -220,7 +406,7 @@ def test_delete(self, mock_del):
         client.minds.drop('my_name')
 
         args, _ = mock_del.call_args
-        assert args[0].endswith(f'/api/projects/mindsdb/minds/my_name')
+        assert args[0].endswith('/api/projects/mindsdb/minds/my_name')
 
     @patch('requests.get')
     @patch('minds.minds.OpenAI')
@@ -230,8 +416,6 @@ def test_completion(self, mock_openai, mock_get):
         response_mock(mock_get, self.mind_json)
         mind = client.minds.get('mind_name')
 
-        openai_response = 'how can I assist you today?'
-
         def openai_completion_f(messages, *args, **kwargs):
             # echo question
             answer = messages[0]['content']
@@ -259,4 +443,3 @@ def openai_completion_f(messages, *args, **kwargs):
             if question == chunk.content.lower():
                 success = True
         assert success is True
-