Add ability for agents to do webpage retrieval

Author: tmichaeldb

mindsdb_sdk/agents.py

@@ -1,5 +1,6 @@
 from requests.exceptions import HTTPError
 from typing import List, Union
+from urllib.parse import urlparse
 from uuid import uuid4
 import datetime
 import pandas as pd
@@ -86,6 +87,14 @@ def add_file(self, file_path: str, description: str, knowledge_base: str = None)
         """
         self.collection.add_file(self.name, file_path, description, knowledge_base)
 
+    def add_webpage(self, url: str, description: str, knowledge_base: str = None):
+        """
+        Add a crawled URL to the agent for retrieval.
+
+        :param url: URL of the page to be crawled and added.
+        """
+        self.collection.add_webpage(self.name, url, description, knowledge_base)
+
     def __repr__(self):
         return f'{self.__class__.__name__}(name: {self.name})'
 
@@ -205,6 +214,44 @@ def add_file(self, name: str, file_path: str, description: str, knowledge_base:
         agent.skills.append(file_retrieval_skill)
         self.update(agent.name, agent)
 
+    def add_webpage(self, name: str, url: str, description: str, knowledge_base: str = None):
+        """
+        Add a webpage to the agent for retrieval.
+
+        :param name: Name of the agent
+        :param file_path: URL of the webpage to be added, or name of existing webpage.
+        :param description: Description of the webpage. Used by agent to know when to do retrieval.
+        :param knowledge_base: Name of an existing knowledge base to be used. Will create a default knowledge base if not given.
+        """
+        parsed_url = urlparse(url)
+        domain = parsed_url.netloc.replace('.', '_')
+        path = parsed_url.path.replace('/', '_')
+        if knowledge_base is not None:
+            kb = self.knowledge_bases.get(knowledge_base)
+        else:
+            kb_name = f'{name}_{domain}{path}_kb'
+            try:
+                kb = self.knowledge_bases.get(kb_name)
+            except AttributeError:
+                # Create KB if it doesn't exist.
+                kb = self.knowledge_bases.create(kb_name)
+                # Wait for underlying embedding model to finish training.
+                kb.model.wait_complete()
+
+        # Insert crawled webpage.
+        kb.insert_webpages([url])
+
+        # Make sure skill name is unique.
+        skill_name = f'{domain}{path}_retrieval_skill_{uuid4()}'
+        retrieval_params = {
+            'source': kb.name,
+            'description': description,
+        }
+        webpage_retrieval_skill = self.skills.create(skill_name, 'retrieval', retrieval_params)
+        agent = self.get(name)
+        agent.skills.append(webpage_retrieval_skill)
+        self.update(agent.name, agent)
+
     def create(
             self,
             name: str,

mindsdb_sdk/connectors/rest_api.py

@@ -338,3 +338,17 @@ def insert_files_into_knowledge_base(self, project: str, knowledge_base_name: st
         _raise_for_status(r)
 
         return r.json()
+
+    @_try_relogin
+    def insert_webpages_into_knowledge_base(self, project: str, knowledge_base_name: str, urls: List[str]):
+        r = self.session.put(
+            self.url + f'/api/projects/{project}/knowledge_bases/{knowledge_base_name}',
+            json={
+                'knowledge_base': {
+                    'urls': urls
+                }
+            }
+        )
+        _raise_for_status(r)
+
+        return r.json()

mindsdb_sdk/knowledge_bases.py

@@ -124,6 +124,12 @@ def insert_files(self, file_paths: List[str]):
         """
         self.api.insert_files_into_knowledge_base(self.project.name, self.name, file_paths)
 
+    def insert_webpages(self, urls: List[str]):
+        """
+        Insert data from crawled URLs to knowledge base
+        """
+        self.api.insert_webpages_into_knowledge_base(self.project.name, self.name, urls)
+
     def insert(self, data: Union[pd.DataFrame, Query, dict]):
         """
         Insert data to knowledge base

tests/test_sdk.py

@@ -38,6 +38,25 @@ def side_effect(*args, **kwargs):
     mock.side_effect = side_effect
 
 
+def responses_mock(mock, data):
+    side_effect_fns = []
+    for d in data:
+        if isinstance(d, pd.DataFrame):
+            # to sql/query format (mostly used)
+            pd_data = d.to_dict('split')
+            d = {
+                'type': 'table',
+                'column_names': pd_data['columns'],
+                'data': pd_data['data']
+            }
+        def side_effect(*args, **kwargs):
+            r_mock = Mock()
+            r_mock.status_code = 200
+            r_mock.json.return_value = d
+            return r_mock
+        side_effect_fns.append(side_effect())
+    mock.side_effect = side_effect_fns
+
 def check_sql_call(mock, sql, database=None, call_stack_num=None):
     if call_stack_num is not None:
         call_args = mock.mock_calls[call_stack_num]
@@ -1436,6 +1455,137 @@ def test_delete(self, mock_delete):
         # Check API call.
         assert mock_delete.call_args[0][0] == f'{DEFAULT_LOCAL_API_URL}/api/projects/mindsdb/agents/test_agent'
 
+    @patch('requests.Session.get')
+    @patch('requests.Session.put')
+    @patch('requests.Session.post')
+    def test_add_file(self, mock_post, mock_put, mock_get):
+        server = mindsdb_sdk.connect()
+        responses_mock(mock_get, [
+            # File metadata get.
+            [{'name': 'tokaido_rules'}],
+            # Existing agent get.
+            {
+                'name': 'test_agent',
+                'model_name': 'test_model',
+                'skills': [],
+                'params': {},
+                'created_at': None,
+                'updated_at': None
+            },
+            # Skills get in Agent update to check if it exists.
+            {'name': 'new_skill', 'type': 'retrieval', 'params': {'source': 'test_agent_tokaido_rules_kb'}},
+            # Existing agent get in Agent update.
+            {
+                'name': 'test_agent',
+                'model_name': 'test_model',
+                'skills': [],
+                'params': {},
+                'created_at': None,
+                'updated_at': None
+            },
+        ])
+        responses_mock(mock_post, [
+            # KB get (POST /sql).
+            pd.DataFrame([
+                {'name': 'test_agent_tokaido_rules_kb', 'storage': None, 'model': None},
+            ]),
+            # Skill creation.
+            {'name': 'new_skill', 'type': 'retrieval', 'params': {'source': 'test_agent_tokaido_rules_kb'}}
+        ])
+        responses_mock(mock_put, [
+            # KB update.
+            {'name': 'test_agent_tokaido_rules_kb'},
+            # Agent update with new skill.
+            {
+                'name': 'test_agent',
+                'model_name': 'test_model',
+                'skills': [{'name': 'new_skill', 'type': 'retrieval', 'params': {'source': 'test_agent_tokaido_rules_kb'}}],
+                'params': {},
+                'created_at': None,
+                'updated_at': None
+            },
+        ])
+        server.agents.add_file('test_agent', './tokaido_rules.pdf', 'Rules for the board game Tokaido')
+
+        # Check Agent was updated with a new skill.
+        agent_update_json = mock_put.call_args[-1]['json']
+        expected_agent_json = {
+            'agent': {
+                'name': 'test_agent',
+                'model_name': 'test_model',
+                # Skill name is a generated UUID.
+                'skills_to_add': [agent_update_json['agent']['skills_to_add'][0]],
+                'skills_to_remove': [],
+                'params': {},
+            }
+        }
+        assert agent_update_json == expected_agent_json
+
+
+    @patch('requests.Session.get')
+    @patch('requests.Session.put')
+    @patch('requests.Session.post')
+    def test_add_webpage(self, mock_post, mock_put, mock_get):
+        server = mindsdb_sdk.connect()
+        responses_mock(mock_get, [
+            # Existing agent get.
+            {
+                'name': 'test_agent',
+                'model_name': 'test_model',
+                'skills': [],
+                'params': {},
+                'created_at': None,
+                'updated_at': None
+            },
+            # Skills get in Agent update to check if it exists.
+            {'name': 'new_skill', 'type': 'retrieval', 'params': {'source': 'test_agent_tokaido_rules_kb'}},
+            # Existing agent get in Agent update.
+            {
+                'name': 'test_agent',
+                'model_name': 'test_model',
+                'skills': [],
+                'params': {},
+                'created_at': None,
+                'updated_at': None
+            },
+        ])
+        responses_mock(mock_post, [
+            # KB get (POST /sql).
+            pd.DataFrame([
+                {'name': 'test_agent_docs_mdb_ai_kb', 'storage': None, 'model': None},
+            ]),
+            # Skill creation.
+            {'name': 'new_skill', 'type': 'retrieval', 'params': {'source': 'test_agent_docs_mdb_ai_kb'}}
+        ])
+        responses_mock(mock_put, [
+            # KB update.
+            {'name': 'test_agent_docs_mdb_ai_kb'},
+            # Agent update with new skill.
+            {
+                'name': 'test_agent',
+                'model_name': 'test_model',
+                'skills': [{'name': 'new_skill', 'type': 'retrieval', 'params': {'source': 'test_agent_docs_mdb_ai_kb'}}],
+                'params': {},
+                'created_at': None,
+                'updated_at': None
+            },
+        ])
+        server.agents.add_webpage('test_agent', 'docs.mdb.ai', 'Documentation for MindsDB')
+
+        # Check Agent was updated with a new skill.
+        agent_update_json = mock_put.call_args[-1]['json']
+        expected_agent_json = {
+            'agent': {
+                'name': 'test_agent',
+                'model_name': 'test_model',
+                # Skill name is a generated UUID.
+                'skills_to_add': [agent_update_json['agent']['skills_to_add'][0]],
+                'skills_to_remove': [],
+                'params': {},
+            }
+        }
+        assert agent_update_json == expected_agent_json
+
 
 class TestSkills():
     @patch('requests.Session.get')