Merge pull request #17 from mindsdb/docstring

Docstrings and changed prompt template param

Author: ZoranPandovski

examples/base_usage.py

@@ -35,14 +35,17 @@
 
 # or separately
 datasource = client.datasources.create(postgres_config)
-mind2 = client.minds.create(name='mind_name', datasources=[datasource] )
+mind = client.minds.create(name='mind_name', datasources=[datasource] )
+
+# with prompt template
+mind = client.minds.create(name='mind_name', prompt_template='You are codding assistant')
 
 # or add to existed mind
-mind3 = client.minds.create(name='mind_name')
+mind = client.minds.create(name='mind_name')
 # by config
-mind2.add_datasource(postgres_config)
+mind.add_datasource(postgres_config)
 # or by datasource
-mind2.add_datasource(datasource)
+mind.add_datasource(datasource)
 
 
 # --- managing minds ---

minds/datasources/datasources.py

@@ -21,6 +21,18 @@ def __init__(self, client):
         self.api = client.api
 
     def create(self, ds_config: DatabaseConfig, replace=False):
+        """
+        Create new datasource and return it
+
+        :param ds_config: datasource configuration, properties:
+           - name: str, name of datatasource
+           - engine: str, type of database handler, for example 'postgres', 'mysql', ...
+           - description: str, description of the database. Used by mind to know what data can be got from it.
+           - connection_data: dict, optional, credentials to connect to database
+           - tables: list of str, optional, list of allowed tables
+        :return: datasource object
+        """
+
         name = ds_config.name
 
         if replace:
@@ -34,6 +46,12 @@ def create(self, ds_config: DatabaseConfig, replace=False):
         return self.get(name)
 
     def list(self) -> List[Datasource]:
+        """
+        Returns list of datasources
+
+        :return: iterable datasources
+        """
+
         data = self.api.get('/datasources').json()
         ds_list = []
         for item in data:
@@ -44,6 +62,13 @@ def list(self) -> List[Datasource]:
         return ds_list
 
     def get(self, name: str) -> Datasource:
+        """
+        Get datasource by name
+
+        :param name: name of datasource
+        :return: datasource object
+        """
+
         data = self.api.get(f'/datasources/{name}').json()
 
         # TODO skip not sql skills
@@ -52,4 +77,10 @@ def get(self, name: str) -> Datasource:
         return Datasource(**data)
 
     def drop(self, name: str):
+        """
+        Drop datasource by name
+
+        :param name: name of datasource
+        """
+
         self.api.delete(f'/datasources/{name}')

minds/minds.py

@@ -28,6 +28,9 @@ def __init__(
         self.name = name
         self.model_name = model_name
         self.provider = provider
+        if parameters is None:
+            parameters = {}
+        self.prompt_template = parameters.pop('prompt_template', None)
         self.parameters = parameters
         self.created_at = created_at
         self.updated_at = updated_at
@@ -39,9 +42,27 @@ def update(
         name: str = None,
         model_name: str = None,
         provider=None,
+        prompt_template=None,
+        datasources=None,
         parameters=None,
-        datasources=None
     ):
+        """
+        Update mind
+
+        If parameter is set it will be applied to mind
+
+        Datasource can be passed as
+         - name, str
+         - Datasource object (minds.datasources.Database)
+         - datasource config (minds.datasources.DatabaseConfig), in this case datasource 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 parameters, dict: alter other parameters of the mind, optional
+        """
         data = {}
 
         if datasources is not None:
@@ -57,8 +78,13 @@ def update(
             data['model_name'] = model_name
         if provider is not None:
             data['provider'] = provider
-        if parameters is not None:
-            data['parameters'] = parameters
+        if parameters is None:
+            parameters = {}
+
+        data['parameters'] = parameters
+
+        if prompt_template is not None:
+            data['parameters']['prompt_template'] = prompt_template
 
         self.api.patch(
             f'/projects/{self.project}/minds/{self.name}',
@@ -68,6 +94,15 @@ def update(
             self.name = name
 
     def add_datasource(self, datasource: Datasource):
+        """
+        Add datasource to mind
+        Datasource can be passed as
+         - name, str
+         - Datasource object (minds.datasources.Database)
+         - datasource config (minds.datasources.DatabaseConfig), in this case datasource will be created
+
+        :param datasource: input datasource
+        """
 
         ds_name = self.client.minds._check_datasource(datasource)
 
@@ -82,6 +117,15 @@ def add_datasource(self, datasource: Datasource):
         self.datasources = updated.datasources
 
     def del_datasource(self, datasource: Union[Datasource, str]):
+        """
+        Delete datasource from mind
+
+        Datasource can be passed as
+         - name, str
+         - Datasource object (minds.datasources.Database)
+
+        :param datasource: datasource to delete
+        """
         if isinstance(datasource, Datasource):
             datasource = datasource.name
         elif not isinstance(datasource, str):
@@ -126,11 +170,11 @@ def completion(self, message: str, stream: bool = False) -> Union[str, Iterable[
             stream=stream
         )
         if stream:
-            return self.stream_response(response)
+            return self._stream_response(response)
         else:
             return response.choices[0].message.content
 
-    def stream_response(self, response):
+    def _stream_response(self, response):
         for chunk in response:
             yield chunk.choices[0].delta
 
@@ -143,13 +187,26 @@ def __init__(self, client):
         self.project = 'mindsdb'
 
     def list(self) -> List[Mind]:
+        """
+        Returns list of minds
+
+        :return: iterable
+        """
+
         data = self.api.get(f'/projects/{self.project}/minds').json()
         minds_list = []
         for item in data:
             minds_list.append(Mind(self.client, **item))
         return minds_list
 
     def get(self, name: str) -> Mind:
+        """
+        Get mind by name
+
+        :param name: name of the mind
+        :return: a mind object
+        """
+
         item = self.api.get(f'/projects/{self.project}/minds/{name}').json()
         return Mind(self.client, **item)
 
@@ -172,10 +229,28 @@ def create(
         self, name,
         model_name=None,
         provider=None,
-        parameters=None,
+        prompt_template=None,
         datasources=None,
+        parameters=None,
         replace=False,
     ) -> Mind:
+        """
+        Create a new mind and return it
+
+        Datasource can be passed as
+         - name, str
+         - Datasource object (minds.datasources.Database)
+         - datasource config (minds.datasources.DatabaseConfig), in this case datasource 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 parameters, dict: other parameters of the mind, optional
+        :param replace: if true - to remove existing mind, default is false
+        :return: created mind
+        """
 
         if replace:
             try:
@@ -193,6 +268,9 @@ def create(
 
         if parameters is None:
             parameters = {}
+
+        if prompt_template is not None:
+            parameters['prompt_template'] = prompt_template
         if 'prompt_template' not in parameters:
             parameters['prompt_template'] = DEFAULT_PROMPT_TEMPLATE
 
@@ -211,4 +289,10 @@ def create(
         return mind
 
     def drop(self, name: str):
+       """
+       Drop mind by name
+
+       :param name: name of the mind
+       """
+
        self.api.delete(f'/projects/{self.project}/minds/{name}')

tests/integration/test_base_flow.py

@@ -88,15 +88,13 @@ def test_minds():
         mind_name,
         replace=True,
         datasources=[ds.name, ds2_cfg],
-        parameters={
-            'prompt_template': prompt1
-        }
+        prompt_template=prompt1
     )
 
     # get
     mind = client.minds.get(mind_name)
     assert len(mind.datasources) == 2
-    assert mind.parameters['prompt_template'] == prompt1
+    assert mind.prompt_template == prompt1
 
     # list
     mind_list = client.minds.list()
@@ -106,17 +104,15 @@ def test_minds():
     mind.update(
         name=mind_name2,
         datasources=[ds.name],
-        parameters={
-            'prompt_template': prompt2
-        }
+        prompt_template=prompt2
     )
     with pytest.raises(ObjectNotFound):
         # this name not exists
         client.minds.get(mind_name)
 
     mind = client.minds.get(mind_name2)
     assert len(mind.datasources) == 1
-    assert mind.parameters['prompt_template'] == prompt2
+    assert mind.prompt_template == prompt2
 
     # add datasource
     mind.add_datasource(ds2_cfg)

tests/unit/test_unit.py

@@ -121,14 +121,14 @@ def test_create(self, mock_del, mock_post, mock_get):
         client = get_client()
 
         mind_name = 'test_mind'
-        parameters = {'prompt_template': 'always agree'}
+        prompt_template = 'always agree'
         datasources = ['my_ds']
         provider = 'openai'
 
         response_mock(mock_get, self.mind_json)
         create_params = {
             'name': mind_name,
-            'parameters': parameters,
+            'prompt_template': prompt_template,
             'datasources': datasources
         }
         mind = client.minds.create(**create_params)
@@ -137,8 +137,9 @@ def check_mind_created(mind, mock_post, create_params):
             args, kwargs = mock_post.call_args
             assert args[0].endswith('/api/projects/mindsdb/minds')
             request = kwargs['json']
-            for k, v in create_params.items():
-                assert request[k] == v
+            for key in ('name', 'datasources', 'provider', 'model_name'),:
+                assert request.get(key) == create_params.get(key)
+            assert create_params.get('prompt_template') == request.get('parameters', {}).get('prompt_template')
 
             self.compare_mind(mind, self.mind_json)
 
@@ -147,7 +148,7 @@ def check_mind_created(mind, mock_post, create_params):
         # with replace
         create_params = {
             'name': mind_name,
-            'parameters': parameters,
+            'prompt_template': prompt_template,
             'provider': provider,
         }
         mind = client.minds.create(replace=True, **create_params)