Merge pull request #187 from mindsdb/fix-kb-web

Add `limit` parameter for using web.crawler from SDK

Author: ea-rus

mindsdb_sdk/agents.py

@@ -131,6 +131,7 @@ def add_webpages(
             description: str,
             knowledge_base: str = None,
             crawl_depth: int = 1,
+            limit: int = None,
             filters: List[str] = None):
         """
         Add a crawled URL to the agent for retrieval.
@@ -139,16 +140,19 @@ def add_webpages(
         :param description: Description of the webpages. 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.
         :param crawl_depth: How deep to crawl from each base URL. 0 = scrape given URLs only, -1 = default max
+        :param limit: max count of pages to crawl
         :param filters: Include only URLs that match these regex patterns
         """
-        self.collection.add_webpages(self.name, urls, description, knowledge_base=knowledge_base, crawl_depth=crawl_depth, filters=filters)
+        self.collection.add_webpages(self.name, urls, description, knowledge_base=knowledge_base,
+                                     crawl_depth=crawl_depth, limit=limit, filters=filters)
 
     def add_webpage(
             self,
             url: str,
             description: str,
             knowledge_base: str = None,
             crawl_depth: int = 1,
+            limit: int = None,
             filters: List[str] = None):
         """
         Add a crawled URL to the agent for retrieval.
@@ -157,9 +161,11 @@ def add_webpage(
         :param description: Description of the webpages. 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.
         :param crawl_depth: How deep to crawl from each base URL. 0 = scrape given URLs only, -1 = default max
+        :param limit: max count of pages to crawl
         :param filters: Include only URLs that match these regex patterns
         """
-        self.collection.add_webpage(self.name, url, description, knowledge_base=knowledge_base, crawl_depth=crawl_depth, filters=filters)
+        self.collection.add_webpage(self.name, url, description, knowledge_base=knowledge_base,
+                                    crawl_depth=crawl_depth, limit=limit, filters=filters)
 
     def add_database(self, database: str, tables: List[str], description: str):
         """
@@ -368,6 +374,7 @@ def add_webpages(
             description: str,
             knowledge_base: str = None,
             crawl_depth: int = 1,
+            limit: int = None,
             filters: List[str] = None
             ):
         """
@@ -378,6 +385,7 @@ def add_webpages(
         :param description: Description of the webpages. 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.
         :param crawl_depth: How deep to crawl from each base URL. 0 = scrape given URLs only
+        :param limit: max count of pages to crawl
         :param filters: Include only URLs that match these regex patterns
         """
         if not urls:
@@ -393,7 +401,7 @@ def add_webpages(
             kb = self._create_default_knowledge_base(agent, kb_name)
 
         # Insert crawled webpage.
-        kb.insert_webpages(urls, crawl_depth=crawl_depth, filters=filters)
+        kb.insert_webpages(urls, crawl_depth=crawl_depth, filters=filters, limit=limit)
 
         # Make sure skill name is unique.
         skill_name = f'web_retrieval_skill_{uuid4().hex}'
@@ -412,6 +420,7 @@ def add_webpage(
             description: str,
             knowledge_base: str = None,
             crawl_depth: int = 1,
+            limit: int = None,
             filters: List[str] = None):
         """
         Add a webpage to the agent for retrieval.
@@ -421,9 +430,11 @@ def add_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.
         :param crawl_depth: How deep to crawl from each base URL. 0 = scrape given URLs only
+        :param limit: max count of pages to crawl
         :param filters: Include only URLs that match these regex patterns
         """
-        self.add_webpages(name, [url], description, knowledge_base=knowledge_base, crawl_depth=crawl_depth, filters=filters)
+        self.add_webpages(name, [url], description, knowledge_base=knowledge_base,
+                          crawl_depth=crawl_depth, limit=limit, filters=filters)
 
     def add_database(self, name: str, database: str, tables: List[str], description: str):
         """

mindsdb_sdk/knowledge_bases.py

@@ -127,18 +127,21 @@ def insert_files(self, file_paths: List[str], params: dict = None):
             data=data
         )
 
-    def insert_webpages(self, urls: List[str], crawl_depth: int = 1, filters: List[str] = None, params: dict = None):
+    def insert_webpages(self, urls: List[str], crawl_depth: int = 1,
+                        filters: List[str] = None, limit=None, params: dict = None):
         """
         Insert data from crawled URLs to knowledge base.
 
         :param urls: URLs to be crawled and inserted.
         :param crawl_depth: How deep to crawl from each base URL. 0 = scrape given URLs only
         :param filters: Include only URLs that match these regex patterns
+        :param limit: max count of pages to crawl
         :param params: Runtime parameters for KB
         """
         data={
             'urls': urls,
             'crawl_depth': crawl_depth,
+            'limit': limit,
             'filters': [] if filters is None else filters,
         }
         if params: