updated the Mind models and opreations

Author: MinuraPunchihewa

minds/datasources/datasources.py

@@ -54,11 +54,11 @@ def create(
         if description is not None:
             data['description'] = description
 
-        self.api.post(
+        response = self.api.post(
             '/datasources',
             data=data
         )
-        return self.get(name)
+        return Datasource(**response.json())
 
     def list(self) -> List[Datasource]:
         """

minds/minds.py

@@ -1,8 +1,9 @@
-from typing import List, Union, Iterable
+from openai import OpenAI
+from typing import Dict, List, Optional, Union, Iterable
+
 import minds.utils as utils
 import minds.exceptions as exc
-from minds.datasources import Datasource, DatabaseConfig, DatabaseTables, DatabaseConfigBase
-from minds.knowledge_bases import KnowledgeBase, KnowledgeBaseConfig
+# from minds.knowledge_bases import KnowledgeBase, KnowledgeBaseConfig
 
 
 class Mind:
@@ -44,132 +45,38 @@ def __repr__(self):
                 f'datasources={self.datasources}, '
                 f'status={self.status})')
 
-    def update(
-        self,
-        name: str = None,
-        model_name: str = None,
-        provider=None,
-        datasources=None,
-        # knowledge_bases=None,
-        parameters=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
-
-        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 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 = {}
-        
-        if name is not None:
-            utils.validate_mind_name(name)
-
-        if datasources is not None:
-            ds_list = []
-            for ds in datasources:
-                ds_list.append(self.client.minds._check_datasource(ds))
-            data['datasources'] = ds_list
-
-        # 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:
-            data['model_name'] = model_name
-        if provider is not None:
-            data['provider'] = provider
-        if parameters is not None:
-            data['parameters'] = parameters
-
-        self.api.put(
-            f'/minds/{self.name}',
-            data=data
-        )
-
-        if name is not None and name != self.name:
-            self.name = name
-
-        refreshed_mind = self.client.minds.get(self.name)
-        self.model_name = refreshed_mind.model_name
-        self.provider = refreshed_mind.provider
-        self.parameters = refreshed_mind.parameters
-        self.created_at = refreshed_mind.created_at
-        self.updated_at = refreshed_mind.updated_at
-        self.datasources = refreshed_mind.datasources
-        # self.knowledge_bases = refreshed_mind.knowledge_bases
-
-    def add_datasource(self, datasource: Datasource):
+    def add_datasource(self, datasource_name: str, tables: Optional[List[str]] = None) -> None:
         """
-        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
+        Add an existing Datasource to a Mind.
 
-        :param datasource: input datasource
+        :param datasource_name: name of the datasource to add.
+        :param tables: list of tables to use from the datasource, optional.
         """
-        ds_name = self.client.minds._check_datasource(datasource)['name']
-        existing_datasources = self.datasources or []
-
-        self.api.put(
+        response = self.api.put(
             f'/minds/{self.name}',
             data={
-                'datasources': existing_datasources + [{'name': ds_name}]
+                'datasources': self.datasources + [{'name': datasource_name, 'tables': tables}]
             }
         )
-        updated = self.client.minds.get(self.name)
+        updated_mind = response.json()
+        self.datasources = updated_mind['datasources']
+        self.status = updated_mind['status']
 
-        self.datasources = updated.datasources
-
-    def del_datasource(self, datasource: Union[Datasource, str]):
+    def remove_datasource(self, datasource_name: str) -> None:
         """
-        Delete datasource from mind
-
-        Datasource can be passed as
-         - name, str
-         - Datasource object (minds.datasources.Database)
+        Remove a datasource from a Mind.
 
-        :param datasource: datasource to delete
+        :param datasource_name: name of the datasource to remove.
         """
-        if isinstance(datasource, Datasource):
-            datasource = datasource.name
-        elif not isinstance(datasource, str):
-            raise ValueError(f'Unknown type of datasource: {datasource}')
-
-        updated_datasources = [ds for ds in (self.datasources or []) if ds['name'] != datasource]
-        if len(updated_datasources) == len(self.datasources or []):
-            raise exc.ObjectNotFound(f'Datasource {datasource} not found in mind {self.name}')
-        
-        self.api.put(
+        response = self.api.put(
             f'/minds/{self.name}',
             data={
-                'datasources': updated_datasources
+                'datasources': [ds for ds in (self.datasources or []) if ds['name'] != datasource_name]
             }
         )
-        updated = self.client.minds.get(self.name)
-
-        self.datasources = updated.datasources
+        updated_mind = response.json()
+        self.datasources = updated_mind['datasources']
+        self.status = updated_mind['status']
 
     # def add_knowledge_base(self, knowledge_base: Union[str, KnowledgeBase, KnowledgeBaseConfig]):
     #     """
@@ -224,7 +131,11 @@ def completion(self, message: str, stream: bool = False) -> Union[str, Iterable[
 
         :return: string if stream mode is off or iterator of strings if stream mode is on
         """
-        response = self.openai_client.chat.completions.create(
+        openai_client = OpenAI(
+            api_key=self.api.api_key,
+            base_url=self.api.base_url
+        )
+        response = openai_client.chat.completions.create(
             model=self.name,
             messages=[
                 {'role': 'user', 'content': message}
@@ -236,7 +147,7 @@ def completion(self, message: str, stream: bool = False) -> Union[str, Iterable[
         else:
             return response.choices[0].message.content
 
-    def _stream_response(self, response):
+    def _stream_response(self, response) -> Iterable[str]:
         for chunk in response:
             yield chunk.choices[0].delta.content
 
@@ -252,7 +163,6 @@ def list(self) -> List[Mind]:
 
         :return: iterable
         """
-
         data = self.api.get(f'/minds').json()
         minds_list = []
         for item in data:
@@ -266,32 +176,9 @@ def get(self, name: str) -> Mind:
         :param name: name of the mind
         :return: a mind object
         """
-        
         item = self.api.get(f'/minds/{name}').json()
         return Mind(self.client, **item)
 
-    def _check_datasource(self, ds) -> dict:
-        if isinstance(ds, DatabaseConfigBase):
-            res = {'name': ds.name}
-
-            if isinstance(ds, DatabaseTables):
-                if ds.tables:
-                    res['tables'] = ds.tables
-
-            if isinstance(ds, DatabaseConfig):
-                # if not exists - create
-                try:
-                    self.client.datasources.get(ds.name)
-                except exc.ObjectNotFound:
-                    self.client.datasources.create(ds)
-
-        elif isinstance(ds, str):
-            res = {'name': ds}
-        else:
-            raise ValueError(f'Unknown type of datasource: {ds}')
-
-        return res
-
     # def _check_knowledge_base(self, knowledge_base) -> str:
     #     if isinstance(knowledge_base, KnowledgeBase):
     #         knowledge_base = knowledge_base.name
@@ -309,22 +196,20 @@ def _check_datasource(self, ds) -> dict:
 
     def create(
         self,
-        name,
-        model_name=None,
-        provider=None,
-        datasources=None,
+        name: str,
+        model_name: Optional[str] = None,
+        provider: Optional[str] = None,
+        datasources: Optional[List[Dict[str, Union[str, List[str]]]]] = None,
         # knowledge_bases=None,
         parameters=None,
         replace=False,
-        update=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
+        Datasources should be a list of dicts with keys:
+        - name: str
+        - tables: Optional[List[str]]
 
         Knowledge base can be passed as
          - name, str
@@ -338,12 +223,9 @@ def create(
         :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
         :return: created mind
         """
-        
-        if name is not None:
-            utils.validate_mind_name(name)
+        utils.validate_mind_name(name)
 
         if replace:
             try:
@@ -352,48 +234,81 @@ def create(
             except exc.ObjectNotFound:
                 ...
 
-        ds_list = []
-        if datasources:
-            for ds in datasources:
-                ds_list.append(self._check_datasource(ds))
-
         # kb_names = []
         # if knowledge_bases:
         #     for kb in knowledge_bases:
         #         kb = self._check_knowledge_base(kb)
         #         kb_names.append(kb)
 
-        if update:
-            method = self.api.put
-            url = f'/minds/{name}'
-        else:
-            method = self.api.post
-            url = f'/minds'
-
         data = {
             'name': name,
             'model_name': model_name,
             'provider': provider,
-            'datasources': ds_list,
+            'datasources': datasources or [],
         }
         if parameters:
             data['parameters'] = parameters
         # if kb_names:
         #     data['knowledge_bases'] = kb_names
 
-        method(
-            url,
+        response = self.api.post(
+            '/minds',
+            data=data
+        )
+
+        return Mind(self.client, **response.json())
+    
+    def update(
+        self,
+        name: str,
+        new_name: Optional[str] = None,
+        model_name: Optional[str] = None,
+        provider: Optional[str] = None,
+        datasources: Optional[List[Dict[str, Union[str, List[str]]]]] = None,
+        parameters: Optional[Dict] = None,
+    ) -> Mind:
+        """
+        Update an existing Mind and return it
+
+        Datasources should be a list of dicts with keys:
+        - name: str
+        - tables: Optional[List[str]]
+
+        :param name: name of the mind to update
+        :param new_name: new name of the mind, optional
+        :param model_name: llm model name, optional
+        :param provider: llm provider, optional
+        :param datasources: list of datasources used by mind, optional
+        :param parameters, dict: other parameters of the mind, optional
+        :return: updated mind
+        """
+        utils.validate_mind_name(name)
+        if new_name:
+            utils.validate_mind_name(new_name)
+
+        data = {}
+        if new_name:
+            data['name'] = new_name
+        if model_name is not None:
+            data['model_name'] = model_name
+        if provider is not None:
+            data['provider'] = provider
+        if datasources is not None:
+            data['datasources'] = datasources
+        if parameters is not None:
+            data['parameters'] = parameters
+
+        response = self.api.put(
+            f'/minds/{name}',
             data=data
         )
-        mind = self.get(name)
 
-        return mind
+        return Mind(self.client, **response.json())
 
-    def drop(self, name: str):
+    def drop(self, name: str) -> None:
        """
        Drop mind by name
 
        :param name: name of the mind
        """
-
        self.api.delete(f'/minds/{name}')