api_utils.py

"""
Utilities for making API calls with retry logic and error handling.
"""

import time
import random
import traceback
from loguru import logger
from litellm import completion

def make_api_call_with_retry(model: str, prompt: str) -> str:
    """
    Makes an API call with a retry mechanism for error handling.

    Args:
        model: The model identifier to use.
        prompt: The prompt text to send to the model.

    Returns:
        The response from the model as a string.
    """
    max_attempts = 20
    base_delay = 10
    max_delay = 60
    attempt = 0
    last_exception = None
    
    while attempt < max_attempts:
        try:
            # Add a small random delay to prevent simultaneous requests
            jitter = random.uniform(0.1, 1.0)
            time.sleep(jitter)
            
            # If this is a retry attempt, add exponential backoff delay
            if attempt > 0:
                delay = min(base_delay * (2 ** (attempt - 1)), max_delay)
                time.sleep(delay)
            
            response = completion(
                model=model,
                messages=[{"role": "user", "content": prompt}],
                num_retries=2,  # Built-in retry mechanism of LiteLLM
            )

            # Try to extract the content from the response
            try:
                analysis_result = response.choices[0].message.content.strip()
                return analysis_result
            except (AttributeError, KeyError, IndexError):
                try:
                    analysis_result = response["choices"][0]["message"]["content"].strip()
                    return analysis_result
                except (AttributeError, KeyError, IndexError):
                    # If we couldn't extract the content, return an error
                    raise ValueError("Failed to extract content from response")
                    
        except (ConnectionError, TimeoutError) as e:
            last_exception = e
            logger.warning("API call failed (attempt {}/{}) for model {}: {}. Retrying...", attempt+1, max_attempts, model, str(e))
            attempt += 1
            continue
        except Exception as e:
            logger.error("Unexpected error: {}", str(e))
            logger.error(traceback.format_exc())
            raise  # For other exceptions, we don't retry
    
    # If all attempts failed, re-raise the last exception
    if last_exception:
        logger.error("All {} attempts failed. Last error: {}", max_attempts, str(last_exception))
        raise last_exception