diff --git a/application/prompt_client/litellm_router.py b/application/prompt_client/litellm_router.py index e9a8b5114..60db86dcd 100644 --- a/application/prompt_client/litellm_router.py +++ b/application/prompt_client/litellm_router.py @@ -21,12 +21,38 @@ LlmFn = Callable[[str, str], str] +DEFAULT_MAX_RETRIES = 2 +DEFAULT_RETRY_SLEEP_SECONDS = 15 + + def retry_policy() -> tuple[int, int]: - """Return (max_retries, sleep_seconds) from CRE_LLM_* env vars.""" - return ( - int(os.environ.get("CRE_LLM_MAX_RETRIES", "2")), - int(os.environ.get("CRE_LLM_RETRY_SLEEP_SECONDS", "15")), - ) + """Return (max_retries, sleep_seconds) from CRE_LLM_* env vars with defensive fallback.""" + raw_retries = os.environ.get("CRE_LLM_MAX_RETRIES") + raw_sleep = os.environ.get("CRE_LLM_RETRY_SLEEP_SECONDS") + + max_retries = DEFAULT_MAX_RETRIES + if raw_retries is not None: + try: + max_retries = max(0, int(raw_retries.strip())) + except (ValueError, AttributeError): + logger.warning( + "Invalid CRE_LLM_MAX_RETRIES=%r; falling back to default %d", + raw_retries, + DEFAULT_MAX_RETRIES, + ) + + sleep_seconds = DEFAULT_RETRY_SLEEP_SECONDS + if raw_sleep is not None: + try: + sleep_seconds = max(0, int(raw_sleep.strip())) + except (ValueError, AttributeError): + logger.warning( + "Invalid CRE_LLM_RETRY_SLEEP_SECONDS=%r; falling back to default %d", + raw_sleep, + DEFAULT_RETRY_SLEEP_SECONDS, + ) + + return (max_retries, sleep_seconds) def get_litellm() -> Any: diff --git a/application/tests/litellm_router_test.py b/application/tests/litellm_router_test.py index dcced3048..44ff96657 100644 --- a/application/tests/litellm_router_test.py +++ b/application/tests/litellm_router_test.py @@ -83,6 +83,45 @@ def test_extract_content_text_lenient_on_empty(self) -> None: "", ) + def test_retry_policy_defaults(self) -> None: + with patch.dict(os.environ, {}, clear=True): + retries, sleep = litellm_router.retry_policy() + self.assertEqual(retries, 2) + self.assertEqual(sleep, 15) + + def test_retry_policy_valid_env(self) -> None: + with patch.dict( + os.environ, + {"CRE_LLM_MAX_RETRIES": "5", "CRE_LLM_RETRY_SLEEP_SECONDS": "20"}, + clear=True, + ): + retries, sleep = litellm_router.retry_policy() + self.assertEqual(retries, 5) + self.assertEqual(sleep, 20) + + def test_retry_policy_invalid_strings_fallback(self) -> None: + with patch.dict( + os.environ, + {"CRE_LLM_MAX_RETRIES": "invalid", "CRE_LLM_RETRY_SLEEP_SECONDS": "not-a-number"}, + clear=True, + ): + with self.assertLogs("application.prompt_client.litellm_router", level="WARNING") as cm: + retries, sleep = litellm_router.retry_policy() + self.assertEqual(retries, 2) + self.assertEqual(sleep, 15) + self.assertTrue(any("Invalid CRE_LLM_MAX_RETRIES" in log for log in cm.output)) + self.assertTrue(any("Invalid CRE_LLM_RETRY_SLEEP_SECONDS" in log for log in cm.output)) + + def test_retry_policy_negative_values_clamped(self) -> None: + with patch.dict( + os.environ, + {"CRE_LLM_MAX_RETRIES": "-2", "CRE_LLM_RETRY_SLEEP_SECONDS": "-10"}, + clear=True, + ): + retries, sleep = litellm_router.retry_policy() + self.assertEqual(retries, 0) + self.assertEqual(sleep, 0) + if __name__ == "__main__": unittest.main()