refac

backend/open_webui/retrieval/web/utils.py

@@ -326,8 +326,9 @@ class SafeFireCrawlLoader(BaseLoader, RateLimitMixin, URLProcessingMixin):
         self.params = params or {}
 
     def lazy_load(self) -> Iterator[Document]:
-        try:
-            for url in self.web_paths:
+        for url in self.web_paths:
+            try:
+                self._sync_wait_for_rate_limit()
                 doc = scrape_firecrawl_url(
                     self.api_url,
                     self.api_key,
@@ -338,22 +339,32 @@ class SafeFireCrawlLoader(BaseLoader, RateLimitMixin, URLProcessingMixin):
                 )
                 if doc is not None:
                     yield doc
-        except Exception as e:
-            if self.continue_on_failure:
-                log.warning(f'Error extracting content from URLs with Firecrawl: {e}')
-            else:
-                raise e
+            except Exception as e:
+                if self.continue_on_failure:
+                    log.warning(f'Error extracting content from {url} with Firecrawl: {e}')
+                    continue
+                raise
 
     async def alazy_load(self):
-        try:
-            docs = await run_in_threadpool(lambda: list(self.lazy_load()))
-            for doc in docs:
-                yield doc
-        except Exception as e:
-            if self.continue_on_failure:
-                log.warning(f'Error extracting content from URLs with Firecrawl: {e}')
-            else:
-                raise e
+        for url in self.web_paths:
+            try:
+                await self._wait_for_rate_limit()
+                doc = await run_in_threadpool(
+                    scrape_firecrawl_url,
+                    self.api_url,
+                    self.api_key,
+                    url,
+                    verify_ssl=self.verify_ssl,
+                    timeout=self.timeout,
+                    params=self.params,
+                )
+                if doc is not None:
+                    yield doc
+            except Exception as e:
+                if self.continue_on_failure:
+                    log.warning(f'Error extracting content from {url} with Firecrawl: {e}')
+                    continue
+                raise
 
 
 class SafeTavilyLoader(BaseLoader, RateLimitMixin, URLProcessingMixin):