Merge pull request unclecode#1590 from unclecode/fix/async-llm-extraction-arunMany

ntohidi · web-flow · commit 66175e132b6f · 2025-11-06T18:40:42.000+08:00
This commit resolves issue unclecode#1055 where LLM extraction was blocking async
diff --git a/crawl4ai/async_webcrawler.py b/crawl4ai/async_webcrawler.py
@@ -617,7 +617,17 @@ async def aprocess_html(
                 else config.chunking_strategy
             )
             sections = chunking.chunk(content)
-            extracted_content = config.extraction_strategy.run(url, sections)
+            # extracted_content = config.extraction_strategy.run(url, sections)
+
+            # Use async version if available for better parallelism
+            if hasattr(config.extraction_strategy, 'arun'):
+                extracted_content = await config.extraction_strategy.arun(url, sections)
+            else:
+                # Fallback to sync version run in thread pool to avoid blocking
+                extracted_content = await asyncio.to_thread(
+                    config.extraction_strategy.run, url, sections
+                )
+
             extracted_content = json.dumps(
                 extracted_content, indent=4, default=str, ensure_ascii=False
             )
diff --git a/crawl4ai/extraction_strategy.py b/crawl4ai/extraction_strategy.py
@@ -94,6 +94,20 @@ def run(self, url: str, sections: List[str], *q, **kwargs) -> List[Dict[str, Any
                 extracted_content.extend(future.result())
         return extracted_content
 
+    async def arun(self, url: str, sections: List[str], *q, **kwargs) -> List[Dict[str, Any]]:
+        """
+        Async version: Process sections of text in parallel using asyncio.
+
+        Default implementation runs the sync version in a thread pool.
+        Subclasses can override this for true async processing.
+
+        :param url: The URL of the webpage.
+        :param sections: List of sections (strings) to process.
+        :return: A list of processed JSON blocks.
+        """
+        import asyncio
+        return await asyncio.to_thread(self.run, url, sections, *q, **kwargs)
+
 
 class NoExtractionStrategy(ExtractionStrategy):
     """
@@ -780,6 +794,177 @@ def run(self, url: str, sections: List[str]) -> List[Dict[str, Any]]:
 
         return extracted_content
 
+    async def aextract(self, url: str, ix: int, html: str) -> List[Dict[str, Any]]:
+        """
+        Async version: Extract meaningful blocks or chunks from the given HTML using an LLM.
+
+        How it works:
+        1. Construct a prompt with variables.
+        2. Make an async request to the LLM using the prompt.
+        3. Parse the response and extract blocks or chunks.
+
+        Args:
+            url: The URL of the webpage.
+            ix: Index of the block.
+            html: The HTML content of the webpage.
+
+        Returns:
+            A list of extracted blocks or chunks.
+        """
+        from .utils import aperform_completion_with_backoff
+
+        if self.verbose:
+            print(f"[LOG] Call LLM for {url} - block index: {ix}")
+
+        variable_values = {
+            "URL": url,
+            "HTML": escape_json_string(sanitize_html(html)),
+        }
+
+        prompt_with_variables = PROMPT_EXTRACT_BLOCKS
+        if self.instruction:
+            variable_values["REQUEST"] = self.instruction
+            prompt_with_variables = PROMPT_EXTRACT_BLOCKS_WITH_INSTRUCTION
+
+        if self.extract_type == "schema" and self.schema:
+            variable_values["SCHEMA"] = json.dumps(self.schema, indent=2)
+            prompt_with_variables = PROMPT_EXTRACT_SCHEMA_WITH_INSTRUCTION
+
+        if self.extract_type == "schema" and not self.schema:
+            prompt_with_variables = PROMPT_EXTRACT_INFERRED_SCHEMA
+
+        for variable in variable_values:
+            prompt_with_variables = prompt_with_variables.replace(
+                "{" + variable + "}", variable_values[variable]
+            )
+
+        try:
+            response = await aperform_completion_with_backoff(
+                self.llm_config.provider,
+                prompt_with_variables,
+                self.llm_config.api_token,
+                base_url=self.llm_config.base_url,
+                json_response=self.force_json_response,
+                extra_args=self.extra_args,
+            )
+            # Track usage
+            usage = TokenUsage(
+                completion_tokens=response.usage.completion_tokens,
+                prompt_tokens=response.usage.prompt_tokens,
+                total_tokens=response.usage.total_tokens,
+                completion_tokens_details=response.usage.completion_tokens_details.__dict__
+                if response.usage.completion_tokens_details
+                else {},
+                prompt_tokens_details=response.usage.prompt_tokens_details.__dict__
+                if response.usage.prompt_tokens_details
+                else {},
+            )
+            self.usages.append(usage)
+
+            # Update totals
+            self.total_usage.completion_tokens += usage.completion_tokens
+            self.total_usage.prompt_tokens += usage.prompt_tokens
+            self.total_usage.total_tokens += usage.total_tokens
+
+            try:
+                content = response.choices[0].message.content
+                blocks = None
+
+                if self.force_json_response:
+                    blocks = json.loads(content)
+                    if isinstance(blocks, dict):
+                        if len(blocks) == 1 and isinstance(list(blocks.values())[0], list):
+                            blocks = list(blocks.values())[0]
+                        else:
+                            blocks = [blocks]
+                    elif isinstance(blocks, list):
+                        blocks = blocks
+                else:
+                    blocks = extract_xml_data(["blocks"], content)["blocks"]
+                    blocks = json.loads(blocks)
+
+                for block in blocks:
+                    block["error"] = False
+            except Exception:
+                parsed, unparsed = split_and_parse_json_objects(
+                    response.choices[0].message.content
+                )
+                blocks = parsed
+                if unparsed:
+                    blocks.append(
+                        {"index": 0, "error": True, "tags": ["error"], "content": unparsed}
+                    )
+
+            if self.verbose:
+                print(
+                    "[LOG] Extracted",
+                    len(blocks),
+                    "blocks from URL:",
+                    url,
+                    "block index:",
+                    ix,
+                )
+            return blocks
+        except Exception as e:
+            if self.verbose:
+                print(f"[LOG] Error in LLM extraction: {e}")
+            return [
+                {
+                    "index": ix,
+                    "error": True,
+                    "tags": ["error"],
+                    "content": str(e),
+                }
+            ]
+
+    async def arun(self, url: str, sections: List[str]) -> List[Dict[str, Any]]:
+        """
+        Async version: Process sections with true parallelism using asyncio.gather.
+
+        Args:
+            url: The URL of the webpage.
+            sections: List of sections (strings) to process.
+
+        Returns:
+            A list of extracted blocks or chunks.
+        """
+        import asyncio
+
+        merged_sections = self._merge(
+            sections,
+            self.chunk_token_threshold,
+            overlap=int(self.chunk_token_threshold * self.overlap_rate),
+        )
+
+        extracted_content = []
+
+        # Create tasks for all sections to run in parallel
+        tasks = [
+            self.aextract(url, ix, sanitize_input_encode(section))
+            for ix, section in enumerate(merged_sections)
+        ]
+
+        # Execute all tasks concurrently
+        results = await asyncio.gather(*tasks, return_exceptions=True)
+
+        # Process results
+        for result in results:
+            if isinstance(result, Exception):
+                if self.verbose:
+                    print(f"Error in async extraction: {result}")
+                extracted_content.append(
+                    {
+                        "index": 0,
+                        "error": True,
+                        "tags": ["error"],
+                        "content": str(result),
+                    }
+                )
+            else:
+                extracted_content.extend(result)
+
+        return extracted_content
+
     def show_usage(self) -> None:
         """Print a detailed token usage report showing total and per-request usage."""
         print("\n=== Token Usage Summary ===")
diff --git a/crawl4ai/utils.py b/crawl4ai/utils.py
@@ -1825,6 +1825,82 @@ def perform_completion_with_backoff(
             # ]
 
 
+async def aperform_completion_with_backoff(
+    provider,
+    prompt_with_variables,
+    api_token,
+    json_response=False,
+    base_url=None,
+    **kwargs,
+):
+    """
+    Async version: Perform an API completion request with exponential backoff.
+
+    How it works:
+    1. Sends an async completion request to the API.
+    2. Retries on rate-limit errors with exponential delays (async).
+    3. Returns the API response or an error after all retries.
+
+    Args:
+        provider (str): The name of the API provider.
+        prompt_with_variables (str): The input prompt for the completion request.
+        api_token (str): The API token for authentication.
+        json_response (bool): Whether to request a JSON response. Defaults to False.
+        base_url (Optional[str]): The base URL for the API. Defaults to None.
+        **kwargs: Additional arguments for the API request.
+
+    Returns:
+        dict: The API response or an error message after all retries.
+    """
+
+    from litellm import acompletion
+    from litellm.exceptions import RateLimitError
+    import asyncio
+
+    max_attempts = 3
+    base_delay = 2  # Base delay in seconds, you can adjust this based on your needs
+
+    extra_args = {"temperature": 0.01, "api_key": api_token, "base_url": base_url}
+    if json_response:
+        extra_args["response_format"] = {"type": "json_object"}
+
+    if kwargs.get("extra_args"):
+        extra_args.update(kwargs["extra_args"])
+
+    for attempt in range(max_attempts):
+        try:
+            response = await acompletion(
+                model=provider,
+                messages=[{"role": "user", "content": prompt_with_variables}],
+                **extra_args,
+            )
+            return response  # Return the successful response
+        except RateLimitError as e:
+            print("Rate limit error:", str(e))
+
+            if attempt == max_attempts - 1:
+                # Last attempt failed, raise the error.
+                raise
+
+            # Check if we have exhausted our max attempts
+            if attempt < max_attempts - 1:
+                # Calculate the delay and wait
+                delay = base_delay * (2**attempt)  # Exponential backoff formula
+                print(f"Waiting for {delay} seconds before retrying...")
+                await asyncio.sleep(delay)
+            else:
+                # Return an error response after exhausting all retries
+                return [
+                    {
+                        "index": 0,
+                        "tags": ["error"],
+                        "content": ["Rate limit error. Please try again later."],
+                    }
+                ]
+        except Exception as e:
+            raise e  # Raise any other exceptions immediately
+
+
 def extract_blocks(url, html, provider=DEFAULT_PROVIDER, api_token=None, base_url=None):
     """
     Extract content blocks from website HTML using an AI provider.
diff --git a/tests/test_llm_extraction_parallel_issue_1055.py b/tests/test_llm_extraction_parallel_issue_1055.py
