diff --git a/agentkit/config.py b/agentkit/config.py index 3fc702f..1121ada 100644 --- a/agentkit/config.py +++ b/agentkit/config.py @@ -71,6 +71,11 @@ class Settings(BaseSettings): gt=0, description="Maximum execution steps allowed per agent run.", ) + MAX_RETRIES: int = Field( + default=3, + ge=0, + description="Maximum retries on transient LLM errors.", + ) RUN_TIMEOUT_S: int = Field( default=60, gt=0, diff --git a/agentkit/llm/factory.py b/agentkit/llm/factory.py index e840cf0..56168d5 100644 --- a/agentkit/llm/factory.py +++ b/agentkit/llm/factory.py @@ -8,6 +8,7 @@ from agentkit.llm.fake import FakeLLMClient from agentkit.llm.gemini import GeminiClient from agentkit.llm.openai import OpenAIClient +from agentkit.llm.retry import RetryingLLMClient logger = logging.getLogger(__name__) @@ -59,7 +60,7 @@ def create_llm_client_from_settings(settings: Settings) -> LLMClient: settings: Application Settings instance. Returns: - Configured LLMClient instance. + Configured LLMClient instance (wrapped with RetryingLLMClient for live providers). """ provider = settings.LLM_PROVIDER.strip().lower() api_key: str | None = None @@ -69,8 +70,16 @@ def create_llm_client_from_settings(settings: Settings) -> LLMClient: elif provider == "openai": api_key = settings.OPENAI_API_KEY - return create_llm_client( + client = create_llm_client( provider=provider, api_key=api_key, model=settings.LLM_MODEL, ) + + if provider in ("gemini", "openai") and settings.MAX_RETRIES > 0: + return RetryingLLMClient( + client=client, + max_retries=settings.MAX_RETRIES, + ) + + return client diff --git a/tests/unit/test_llm_factory.py b/tests/unit/test_llm_factory.py index f0d6cc0..2f4afc9 100644 --- a/tests/unit/test_llm_factory.py +++ b/tests/unit/test_llm_factory.py @@ -5,6 +5,7 @@ from agentkit.llm.fake import FakeLLMClient from agentkit.llm.gemini import GeminiClient from agentkit.llm.openai import OpenAIClient +from agentkit.llm.retry import RetryingLLMClient def test_create_fake_client() -> None: @@ -47,10 +48,13 @@ def test_create_from_settings_gemini() -> None: LLM_PROVIDER="gemini", LLM_MODEL="gemini-2.0-flash", GEMINI_API_KEY="test-gemini-key", + MAX_RETRIES=3, ) client = create_llm_client_from_settings(settings) - assert isinstance(client, GeminiClient) + assert isinstance(client, RetryingLLMClient) + assert isinstance(client._client, GeminiClient) assert client.model == "gemini-2.0-flash" + assert client._max_retries == 3 def test_create_from_settings_openai() -> None: @@ -58,13 +62,29 @@ def test_create_from_settings_openai() -> None: LLM_PROVIDER="openai", LLM_MODEL="gpt-4o", OPENAI_API_KEY="test-openai-key", + MAX_RETRIES=2, ) client = create_llm_client_from_settings(settings) - assert isinstance(client, OpenAIClient) + assert isinstance(client, RetryingLLMClient) + assert isinstance(client._client, OpenAIClient) assert client.model == "gpt-4o" + assert client._max_retries == 2 + + +def test_create_from_settings_zero_retries() -> None: + settings = Settings( + LLM_PROVIDER="gemini", + LLM_MODEL="gemini-2.0-flash", + GEMINI_API_KEY="test-gemini-key", + MAX_RETRIES=0, + ) + client = create_llm_client_from_settings(settings) + assert isinstance(client, GeminiClient) + assert not isinstance(client, RetryingLLMClient) def test_create_from_settings_fake() -> None: settings = Settings(LLM_PROVIDER="fake") client = create_llm_client_from_settings(settings) assert isinstance(client, FakeLLMClient) + assert not isinstance(client, RetryingLLMClient)