Merge pull request #106 from mindsdb/retrieval-fix

Use Retrieval Skill for Adding Files to Agents

Author: tmichaeldb

mindsdb_sdk/agents.py

@@ -192,15 +192,15 @@ def add_file(self, name: str, file_path: str, description: str, knowledge_base:
                 kb.model.wait_complete()
 
         # Insert the entire file.
-        kb.insert(self.databases.files.tables.get(filename_no_extension))
+        kb.insert_files([filename_no_extension])
 
         # Make sure skill name is unique.
         skill_name = f'{filename_no_extension}_retrieval_skill_{uuid4()}'
         retrieval_params = {
             'source': kb.name,
             'description': description,
         }
-        file_retrieval_skill = self.skills.create(skill_name, 'knowledge_base', retrieval_params)
+        file_retrieval_skill = self.skills.create(skill_name, 'retrieval', retrieval_params)
         agent = self.get(name)
         agent.skills.append(file_retrieval_skill)
         self.update(agent.name, agent)

mindsdb_sdk/connectors/rest_api.py

@@ -323,3 +323,18 @@ def delete_skill(self, project: str, name: str):
         url = self.url + f'/api/projects/{project}/skills/{name}'
         r = self.session.delete(url)
         _raise_for_status(r)
+
+    # Knowledge Base operations.
+    @_try_relogin
+    def insert_files_into_knowledge_base(self, project: str, knowledge_base_name: str, file_names: List[str]):
+        r = self.session.put(
+            self.url + f'/api/projects/{project}/knowledge_bases/{knowledge_base_name}',
+            json={
+                'knowledge_base': {
+                    'files': file_names
+                }
+            }
+        )
+        _raise_for_status(r)
+
+        return r.json()

mindsdb_sdk/knowledge_bases.py

@@ -32,8 +32,8 @@ class KnowledgeBase(Query):
 
     """
 
-    def __init__(self, project, data: dict):
-
+    def __init__(self, api, project, data: dict):
+        self.api = api
         self.project = project
         self.name = data['name']
         self.table_name = Identifier(parts=[self.project.name, self.name])
@@ -118,6 +118,12 @@ def _update_query(self):
             ast_query.limit = Constant(self._limit)
         self.sql = ast_query.to_string()
 
+    def insert_files(self, file_paths: List[str]):
+        """
+        Insert data from file to knowledge base
+        """
+        self.api.insert_files_into_knowledge_base(self.project.name, self.name, file_paths)
+
     def insert(self, data: Union[pd.DataFrame, Query, dict]):
         """
         Insert data to knowledge base
@@ -210,7 +216,7 @@ def _list(self, name: str = None) -> List[KnowledgeBase]:
         df = df.rename(columns=cols_map)
 
         return [
-            KnowledgeBase(self.project, item)
+            KnowledgeBase(self.api, self.project, item)
             for item in df.to_dict('records')
         ]
 

mindsdb_sdk/skills.py

@@ -18,7 +18,7 @@ class Skill():
 
     Create a new SQL skill:
 
-    >>> text_to_sql_skill = skills.create('text_to_sql', 'sql', { 'tables': ['my_table'], 'database': 'my_database' })
+    >>> text_to_sql_skill = skills.create('text_to_sql', 'sql', { 'tables': ['my_table'], 'database': 'my_database', 'description': 'my_description'})
 
     Update a skill:
 
@@ -50,30 +50,34 @@ def __repr__(self):
 
     @classmethod
     def from_json(cls, json: dict):
+        name = json['name']
+        type = json['type']
+        params = json['params']
         if json['type'] == 'sql':
-            return SQLSkill(json['name'], json['params']['tables'], json['params']['database'])
+            return SQLSkill(name, params['tables'], params['database'], params.get('description', ''))
         if json['type'] == 'retrieval':
-            return RetrievalSkill(json['name'], json['params']['knowledge_base'], json['params']['description'])
-        return Skill(json['name'], json['type'], json['params'])
+            return RetrievalSkill(name, params['source'], params.get('description', ''))
+        return Skill(name, type, params)
 
 
 class SQLSkill(Skill):
     """Represents a MindsDB skill for agents to interact with MindsDB databases"""
-    def __init__(self, name: str, tables: List[str], database: str):
+    def __init__(self, name: str, tables: List[str], database: str, description: str):
         params = {
             'database': database,
             'tables': tables,
+            'description': description
         }
         super().__init__(name, 'sql', params)
 
 class RetrievalSkill(Skill):
     """Represents a MindsDB skill for agents to interact with MindsDB data sources"""
     def __init__(self, name: str, knowledge_base: str, description: str):
         params = {
-            'knowledge_base': knowledge_base,
+            'source': knowledge_base,
             'description': description
         }
-        super().__init__(name, 'knowledge_base', params)
+        super().__init__(name, 'retrieval', params)
 
 
 class Skills(CollectionBase):
@@ -114,7 +118,7 @@ def create(self, name: str, type: str, params: dict = None) -> Skill:
         """
         _ = self.api.create_skill(self.project, name, type, params)
         if type == 'sql':
-            return SQLSkill(name, params['tables'], params['database'])
+            return SQLSkill(name, params['tables'], params['database'], params['description'])
         return Skill(name, type, params)
 
     def update(self, name: str, updated_skill: Skill) -> Skill:

tests/test_sdk.py

@@ -1295,7 +1295,7 @@ def test_create(self, mock_post):
                 'name': 'test_skill',
                 'project_id': 1,
                 'type': 'sql',
-                'params': {'tables': ['test_table'], 'database': 'test_database'},
+                'params': {'tables': ['test_table'], 'database': 'test_database', 'description': 'test_description'},
             }],
             'params': {'k1': 'v1'},
             'created_at': created_at,
@@ -1308,7 +1308,7 @@ def test_create(self, mock_post):
         new_agent = server.agents.create(
             name='test_agent',
             model=Model(None, {'name':'m1'}),
-            skills=[SQLSkill('test_skill', ['test_table'], 'test_database')],
+            skills=[SQLSkill('test_skill', ['test_table'], 'test_database', 'test_description')],
             params={'k1': 'v1'}
         )
         # Check API call.
@@ -1329,11 +1329,11 @@ def test_create(self, mock_post):
            'skill': {
                 'name': 'test_skill',
                 'type': 'sql',
-                'params': {'database': 'test_database', 'tables': ['test_table']}
+                'params': {'database': 'test_database', 'tables': ['test_table'], 'description': 'test_description'}
             }
         }
 
-        expected_skill = SQLSkill('test_skill', ['test_table'], 'test_database')
+        expected_skill = SQLSkill('test_skill', ['test_table'], 'test_database', 'test_description')
         expected_agent = Agent(
             'test_agent',
             'test_model',
@@ -1363,7 +1363,7 @@ def test_update(self, mock_get, mock_put, _):
                 'name': 'updated_skill',
                 'project_id': 1,
                 'type': 'sql',
-                'params': {'tables': ['updated_table'], 'database': 'updated_database'},
+                'params': {'tables': ['updated_table'], 'database': 'updated_database', 'description': 'test_description'},
             }],
             'params': {'k2': 'v2'},
             'created_at': created_at,
@@ -1385,7 +1385,7 @@ def test_update(self, mock_get, mock_put, _):
         expected_agent = Agent(
             'test_agent',
             'updated_model',
-            [SQLSkill('updated_skill', ['updated_table'], 'updated_database')],
+            [SQLSkill('updated_skill', ['updated_table'], 'updated_database', 'test_description')],
             {'k2': 'v2'},
             created_at,
             updated_at
@@ -1453,14 +1453,14 @@ def test_list(self, mock_get):
                 'id': 1,
                 'name': 'test_skill',
                 'project_id': 1,
-                'params': {'tables': ['test_table'], 'database': 'test_database'},
+                'params': {'tables': ['test_table'], 'database': 'test_database', 'description': 'test_description' },
                 'type': 'sql'
             }
         ])
         all_skills = server.skills.list()
         assert len(all_skills) == 1
 
-        expected_skill = SQLSkill('test_skill', ['test_table'], 'test_database')
+        expected_skill = SQLSkill('test_skill', ['test_table'], 'test_database', 'test_description')
         assert all_skills[0] == expected_skill
 
     @patch('requests.Session.get')
@@ -1471,14 +1471,14 @@ def test_get(self, mock_get):
                 'id': 1,
                 'name': 'test_skill',
                 'project_id': 1,
-                'params': {'tables': ['test_table'], 'database': 'test_database'},
+                'params': {'tables': ['test_table'], 'database': 'test_database', 'description': 'test_description'},
                 'type': 'sql'
             }
         )
         skill = server.skills.get('test_skill')
         # Check API call.
         assert mock_get.call_args[0][0] == f'{DEFAULT_LOCAL_API_URL}/api/projects/mindsdb/skills/test_skill'
-        expected_skill = SQLSkill('test_skill', ['test_table'], 'test_database')
+        expected_skill = SQLSkill('test_skill', ['test_table'], 'test_database', 'test_description')
         assert skill == expected_skill
 
     @patch('requests.Session.post')
@@ -1487,7 +1487,7 @@ def test_create(self, mock_post):
             'id': 1,
             'name': 'test_skill',
             'project_id': 1,
-            'params': {'k1': 'v1'},
+            'params': {'tables': ['test_table'], 'database': 'test_database', 'description': 'test_description'},
             'type': 'test'
         }
         response_mock(mock_post, data)
@@ -1497,18 +1497,18 @@ def test_create(self, mock_post):
         new_skill = server.skills.create(
             'test_skill',
             'sql',
-            params={'tables': ['test_table'], 'database': 'test_database'}
+            params={'tables': ['test_table'], 'database': 'test_database', 'description': 'test_description'}
         )
         # Check API call.
         assert mock_post.call_args[0][0] == f'{DEFAULT_LOCAL_API_URL}/api/projects/mindsdb/skills'
         assert mock_post.call_args[1]['json'] == {
            'skill': {
                 'name': 'test_skill',
                 'type': 'sql',
-                'params': {'database': 'test_database', 'tables': ['test_table']}
+                'params': {'database': 'test_database', 'tables': ['test_table'], 'description': 'test_description'}
             }
         }
-        expected_skill = SQLSkill('test_skill', ['test_table'], 'test_database')
+        expected_skill = SQLSkill('test_skill', ['test_table'], 'test_database', 'test_description')
 
         assert new_skill == expected_skill
 
@@ -1518,13 +1518,13 @@ def test_update(self, mock_put):
             'id': 1,
             'name': 'test_skill',
             'project_id': 1,
-            'params': {'tables': ['updated_table'], 'database': 'updated_database'},
+            'params': {'tables': ['updated_table'], 'database': 'updated_database', 'description': 'updated_description'},
             'type': 'sql'
         }
         response_mock(mock_put, data)
 
         server = mindsdb_sdk.connect()
-        expected_skill = SQLSkill('test_skill', ['updated_table'], 'updated_database')
+        expected_skill = SQLSkill('test_skill', ['updated_table'], 'updated_database', 'updated_description')
 
         updated_skill = server.skills.update('test_skill', expected_skill)
         # Check API call.
@@ -1533,7 +1533,7 @@ def test_update(self, mock_put):
            'skill': {
                 'name': 'test_skill',
                 'type': 'sql',
-                'params': {'tables': ['updated_table'], 'database': 'updated_database'}
+                'params': {'tables': ['updated_table'], 'database': 'updated_database', 'description': 'updated_description'}
             }
         }