Spaces:
Runtime error
Runtime error
| """ | |
| Model caller with retry and fallback support. | |
| """ | |
| import time | |
| import re | |
| import json | |
| from openai import OpenAI | |
| from src.config import config | |
| from src.utils.logging import get_logger | |
| logger = get_logger(__name__) | |
| # openai client pointing at github models | |
| client = OpenAI( | |
| base_url="https://models.github.ai/inference", | |
| api_key=config.GITHUB_TOKEN, | |
| ) | |
| def _is_quota_error(e: Exception) -> bool: | |
| """Check if this is a rate-limit/quota error.""" | |
| error_str = str(e).lower() | |
| return "429" in error_str or "quota" in error_str or "rate" in error_str or "resource exhausted" in error_str | |
| def call_with_fallback( | |
| messages: list, | |
| response_format: type = None, | |
| primary_model_name: str = None, | |
| fallback_model_names: list[str] = None, | |
| max_retries: int = None, | |
| base_delay: float = None, | |
| client_override: OpenAI = None, | |
| model_override: str = None, | |
| ) -> str: | |
| """ | |
| Try the primary model, retry on 429s, then fall back to the next model. | |
| Returns raw text (or JSON string if response_format is set). | |
| If client_override and model_override are provided, they take priority | |
| over the default client and model chain (used for user-supplied API keys). | |
| """ | |
| if max_retries is None: | |
| max_retries = config.MAX_RETRIES | |
| if base_delay is None: | |
| base_delay = config.RETRY_BASE_DELAY | |
| active_client = client_override or client | |
| primary_model = model_override or primary_model_name or config.ROUTER_MODEL_NAME | |
| # Skip fallbacks when using a user-supplied override | |
| fallbacks = [] if (client_override or model_override) else (fallback_model_names or config.ROUTER_MODEL_FALLBACKS) | |
| model_chain = [primary_model] + list(fallbacks) | |
| last_exception = None | |
| for model_index, model_name in enumerate(model_chain): | |
| for attempt in range(max_retries): | |
| try: | |
| logger.info( | |
| f"Calling model '{model_name}' (attempt {attempt + 1}/{max_retries})", | |
| extra={"custom_fields": {"model": model_name, "attempt": attempt + 1}} | |
| ) | |
| if response_format: | |
| response = active_client.beta.chat.completions.parse( | |
| model=model_name, | |
| messages=messages, | |
| response_format=response_format | |
| ) | |
| result_text = response.choices[0].message.content | |
| else: | |
| response = active_client.chat.completions.create( | |
| model=model_name, | |
| messages=messages, | |
| ) | |
| result_text = response.choices[0].message.content | |
| logger.info( | |
| f"Model '{model_name}' succeeded on attempt {attempt + 1}", | |
| extra={"custom_fields": {"model": model_name}} | |
| ) | |
| return result_text | |
| except Exception as e: | |
| last_exception = e | |
| if not _is_quota_error(e): | |
| logger.error( | |
| f"Model '{model_name}' failed with non-quota error: {e}", | |
| extra={"custom_fields": {"model": model_name, "error": str(e)}} | |
| ) | |
| raise | |
| wait_time = 5.0 | |
| if attempt < max_retries - 1: | |
| logger.warning( | |
| f"Model '{model_name}' quota hit (attempt {attempt + 1}/{max_retries}). Retrying in {wait_time:.1f}s...", | |
| extra={"custom_fields": {"model": model_name, "wait": wait_time}} | |
| ) | |
| time.sleep(wait_time) | |
| else: | |
| if model_index < len(model_chain) - 1: | |
| next_model = model_chain[model_index + 1] | |
| logger.warning( | |
| f"Model '{model_name}' exhausted all {max_retries} retries. Falling back to '{next_model}'...", | |
| extra={"custom_fields": {"model": model_name, "fallback": next_model}} | |
| ) | |
| else: | |
| logger.error("All models and retries exhausted.") | |
| raise last_exception | |