From 466ec4ce90ab2dfbc584a3461c45340a82ddd54e Mon Sep 17 00:00:00 2001 From: Kangmoon Seo <100016044+KangmoonSeo@users.noreply.github.com> Date: Thu, 22 May 2025 15:04:16 +0900 Subject: [PATCH] fix: exclude RateLimitError from retry logic --- .../algo/ai_handlers/langchain_ai_handler.py | 22 +++++++++++------ .../algo/ai_handlers/litellm_ai_handler.py | 12 +++++----- .../algo/ai_handlers/openai_ai_handler.py | 24 ++++++++++--------- 3 files changed, 34 insertions(+), 24 deletions(-) diff --git a/pr_agent/algo/ai_handlers/langchain_ai_handler.py b/pr_agent/algo/ai_handlers/langchain_ai_handler.py index d4ea0aa5..4d708fcb 100644 --- a/pr_agent/algo/ai_handlers/langchain_ai_handler.py +++ b/pr_agent/algo/ai_handlers/langchain_ai_handler.py @@ -6,8 +6,8 @@ except: # we don't enforce langchain as a dependency, so if it's not installed, import functools -from openai import APIError, RateLimitError, Timeout -from retry import retry +import openai +from tenacity import retry, retry_if_exception_type, retry_if_not_exception_type, stop_after_attempt from pr_agent.algo.ai_handlers.base_ai_handler import BaseAiHandler from pr_agent.config_loader import get_settings @@ -36,8 +36,10 @@ class LangChainOpenAIHandler(BaseAiHandler): """ return get_settings().get("OPENAI.DEPLOYMENT_ID", None) - @retry(exceptions=(APIError, Timeout, AttributeError, RateLimitError), - tries=OPENAI_RETRIES, delay=2, backoff=2, jitter=(1, 3)) + @retry( + retry=retry_if_exception_type(openai.APIError) & retry_if_not_exception_type(openai.RateLimitError), + stop=stop_after_attempt(OPENAI_RETRIES), + ) async def chat_completion(self, model: str, system: str, user: str, temperature: float = 0.2): try: messages = [SystemMessage(content=system), HumanMessage(content=user)] @@ -47,9 +49,15 @@ class LangChainOpenAIHandler(BaseAiHandler): finish_reason = "completed" return resp.content, finish_reason - except (Exception) as e: - get_logger().error("Unknown error during OpenAI inference: ", e) - raise e + except openai.RateLimitError as e: + get_logger().error(f"Rate limit error during LLM inference: {e}") + raise + except openai.APIError as e: + get_logger().warning(f"Error during LLM inference: {e}") + raise + except Exception as e: + get_logger().warning(f"Unknown error during LLM inference: {e}") + raise openai.APIError from e def _create_chat(self, deployment_id=None): try: diff --git a/pr_agent/algo/ai_handlers/litellm_ai_handler.py b/pr_agent/algo/ai_handlers/litellm_ai_handler.py index 8d727b8b..f20b03f8 100644 --- a/pr_agent/algo/ai_handlers/litellm_ai_handler.py +++ b/pr_agent/algo/ai_handlers/litellm_ai_handler.py @@ -3,7 +3,7 @@ import litellm import openai import requests from litellm import acompletion -from tenacity import retry, retry_if_exception_type, stop_after_attempt +from tenacity import retry, retry_if_exception_type, retry_if_not_exception_type, stop_after_attempt from pr_agent.algo import CLAUDE_EXTENDED_THINKING_MODELS, NO_SUPPORT_TEMPERATURE_MODELS, SUPPORT_REASONING_EFFORT_MODELS, USER_MESSAGE_ONLY_MODELS from pr_agent.algo.ai_handlers.base_ai_handler import BaseAiHandler @@ -274,8 +274,8 @@ class LiteLLMAIHandler(BaseAiHandler): return get_settings().get("OPENAI.DEPLOYMENT_ID", None) @retry( - retry=retry_if_exception_type((openai.APIError, openai.APIConnectionError, openai.APITimeoutError)), # No retry on RateLimitError - stop=stop_after_attempt(OPENAI_RETRIES) + retry=retry_if_exception_type(openai.APIError) & retry_if_not_exception_type(openai.RateLimitError), + stop=stop_after_attempt(OPENAI_RETRIES), ) async def chat_completion(self, model: str, system: str, user: str, temperature: float = 0.2, img_path: str = None): try: @@ -371,13 +371,13 @@ class LiteLLMAIHandler(BaseAiHandler): get_logger().info(f"\nUser prompt:\n{user}") response = await acompletion(**kwargs) - except (openai.RateLimitError) as e: + except openai.RateLimitError as e: get_logger().error(f"Rate limit error during LLM inference: {e}") raise - except (openai.APIError, openai.APITimeoutError) as e: + except openai.APIError as e: get_logger().warning(f"Error during LLM inference: {e}") raise - except (Exception) as e: + except Exception as e: get_logger().warning(f"Unknown error during LLM inference: {e}") raise openai.APIError from e if response is None or len(response["choices"]) == 0: diff --git a/pr_agent/algo/ai_handlers/openai_ai_handler.py b/pr_agent/algo/ai_handlers/openai_ai_handler.py index f74444a1..253282b0 100644 --- a/pr_agent/algo/ai_handlers/openai_ai_handler.py +++ b/pr_agent/algo/ai_handlers/openai_ai_handler.py @@ -1,8 +1,8 @@ from os import environ from pr_agent.algo.ai_handlers.base_ai_handler import BaseAiHandler import openai -from openai import APIError, AsyncOpenAI, RateLimitError, Timeout -from retry import retry +from openai import AsyncOpenAI +from tenacity import retry, retry_if_exception_type, retry_if_not_exception_type, stop_after_attempt from pr_agent.algo.ai_handlers.base_ai_handler import BaseAiHandler from pr_agent.config_loader import get_settings @@ -38,8 +38,10 @@ class OpenAIHandler(BaseAiHandler): """ return get_settings().get("OPENAI.DEPLOYMENT_ID", None) - @retry(exceptions=(APIError, Timeout, AttributeError, RateLimitError), - tries=OPENAI_RETRIES, delay=2, backoff=2, jitter=(1, 3)) + @retry( + retry=retry_if_exception_type(openai.APIError) & retry_if_not_exception_type(openai.RateLimitError), + stop=stop_after_attempt(OPENAI_RETRIES), + ) async def chat_completion(self, model: str, system: str, user: str, temperature: float = 0.2): try: get_logger().info("System: ", system) @@ -57,12 +59,12 @@ class OpenAIHandler(BaseAiHandler): get_logger().info("AI response", response=resp, messages=messages, finish_reason=finish_reason, model=model, usage=usage) return resp, finish_reason - except (APIError, Timeout) as e: - get_logger().error("Error during OpenAI inference: ", e) + except openai.RateLimitError as e: + get_logger().error(f"Rate limit error during LLM inference: {e}") raise - except (RateLimitError) as e: - get_logger().error("Rate limit error during OpenAI inference: ", e) - raise - except (Exception) as e: - get_logger().error("Unknown error during OpenAI inference: ", e) + except openai.APIError as e: + get_logger().warning(f"Error during LLM inference: {e}") raise + except Exception as e: + get_logger().warning(f"Unknown error during LLM inference: {e}") + raise openai.APIError from e