Fixed issues related formatting and linting

Author: bharath03-a

langextract/annotation.py

@@ -216,41 +216,43 @@ def _process_batch_with_retry(
       **kwargs,
   ) -> Iterator[list[core_types.ScoredOutput]]:
     """Process a batch of prompts with individual chunk retry capability.
-    
+
     This method processes each chunk individually and retries failed chunks
     due to transient errors (like 503 "model overloaded") while preserving
     successful chunks from the same batch.
-    
+
     Args:
       batch_prompts: List of prompts for the batch
       batch: List of TextChunk objects corresponding to the prompts
       **kwargs: Additional arguments passed to the language model
-      
+
     Yields:
       Lists of ScoredOutputs, with retries for failed chunks
     """
     try:
-      batch_results = list(self._language_model.infer(
-          batch_prompts=batch_prompts,
-          **kwargs,
-      ))
-      
+      batch_results = list(
+          self._language_model.infer(
+              batch_prompts=batch_prompts,
+              **kwargs,
+          )
+      )
+
       for result in batch_results:
         yield result
       return
-      
+
     except Exception as e:
       if not retry_utils.is_transient_error(e):
         raise
-      
+
       logging.warning(
           "Batch processing failed with transient error: %s. "
           "Falling back to individual chunk processing with retry.",
-          str(e)
+          str(e),
       )
-    
+
     individual_results = []
-    
+
     for i, (prompt, chunk) in enumerate(zip(batch_prompts, batch)):
       try:
         chunk_result = self._process_single_chunk_with_retry(
@@ -264,16 +266,19 @@ def _process_batch_with_retry(
             **kwargs,
         )
         individual_results.append(chunk_result)
-        
+
       except Exception as e:
         logging.error(
             "Failed to process chunk %d after retries: %s. "
             "Chunk info: document_id=%s, text_length=%d. "
             "Stopping document processing.",
-            i, str(e), chunk.document_id, len(chunk.chunk_text)
+            i,
+            str(e),
+            chunk.document_id,
+            len(chunk.chunk_text),
         )
         raise
-    
+
     for result in individual_results:
       yield result
 
@@ -289,7 +294,7 @@ def _process_single_chunk_with_retry(
       **kwargs,
   ) -> list[core_types.ScoredOutput]:
     """Process a single chunk with retry logic.
-    
+
     Args:
       prompt: The prompt for this chunk
       chunk: The TextChunk object
@@ -299,59 +304,69 @@ def _process_single_chunk_with_retry(
       retry_backoff_factor: Backoff multiplier for retries
       retry_max_delay: Maximum delay between retries
       **kwargs: Additional arguments for the language model
-      
+
     Returns:
       List containing a single ScoredOutput for this chunk
     """
     last_exception = None
     delay = retry_initial_delay
-    
+
     for attempt in range(max_retries + 1):
       try:
-        batch_results = list(self._language_model.infer(
-            batch_prompts=[prompt],
-            **kwargs,
-        ))
-        
+        batch_results = list(
+            self._language_model.infer(
+                batch_prompts=[prompt],
+                **kwargs,
+            )
+        )
+
         if not batch_results:
           raise exceptions.InferenceOutputError(
               f"No results returned for chunk in document {chunk.document_id}"
           )
-        
+
         return batch_results[0]
-        
+
       except Exception as e:
         last_exception = e
-        
+
         if not retry_transient_errors or not retry_utils.is_transient_error(e):
           logging.debug(
-              "Not retrying chunk processing: retry_disabled=%s, is_transient=%s, error=%s",
-              not retry_transient_errors, retry_utils.is_transient_error(e), str(e)
+              "Not retrying chunk processing: retry_disabled=%s,"
+              " is_transient=%s, error=%s",
+              not retry_transient_errors,
+              retry_utils.is_transient_error(e),
+              str(e),
           )
           raise
-        
+
         if attempt >= max_retries:
           logging.error(
               "Chunk processing failed after %d retries: %s",
-              max_retries, str(e)
+              max_retries,
+              str(e),
           )
           raise
-        
+
         current_delay = min(delay, retry_max_delay)
-        
+
         import random
+
         jitter_amount = current_delay * 0.1 * random.random()
         current_delay += jitter_amount
-        
+
         logging.warning(
-            "Chunk processing failed on attempt %d/%d due to transient error: %s. "
-            "Retrying in %.2f seconds...",
-            attempt + 1, max_retries + 1, str(e), current_delay
+            "Chunk processing failed on attempt %d/%d due to transient error:"
+            " %s. Retrying in %.2f seconds...",
+            attempt + 1,
+            max_retries + 1,
+            str(e),
+            current_delay,
         )
-        
+
         time.sleep(current_delay)
         delay = min(delay * retry_backoff_factor, retry_max_delay)
-    
+
     if last_exception:
       raise last_exception
     raise RuntimeError("Chunk retry logic failed unexpectedly")

langextract/extraction.py

@@ -333,11 +333,11 @@ def extract(
 
   # Add retry parameters to alignment kwargs
   retry_kwargs = {
-      'retry_transient_errors': retry_transient_errors,
-      'max_retries': max_retries,
-      'retry_initial_delay': retry_initial_delay,
-      'retry_backoff_factor': retry_backoff_factor,
-      'retry_max_delay': retry_max_delay,
+      "retry_transient_errors": retry_transient_errors,
+      "max_retries": max_retries,
+      "retry_initial_delay": retry_initial_delay,
+      "retry_backoff_factor": retry_backoff_factor,
+      "retry_max_delay": retry_max_delay,
   }
   alignment_kwargs.update(retry_kwargs)
 

langextract/providers/gemini.py

@@ -23,6 +23,7 @@
 
 from absl import logging
 
+from langextract import retry_utils
 from langextract.core import base_model
 from langextract.core import data
 from langextract.core import exceptions
@@ -31,7 +32,6 @@
 from langextract.providers import patterns
 from langextract.providers import router
 from langextract.providers import schemas
-from langextract import retry_utils
 
 _API_CONFIG_KEYS: Final[set[str]] = {
     'response_mime_type',