Skip to content

Commit 2f5b51e

Browse files
committed
try fix llm with retry
1 parent 2d842dc commit 2f5b51e

File tree

1 file changed

+7
-2
lines changed

1 file changed

+7
-2
lines changed

src/agents/llm_factory.py

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -79,8 +79,13 @@ def __init__(self, temperature: int = 0):
7979
model=os.environ['GOOGLE_AI_MODEL'],
8080
temperature=temperature,
8181
google_api_key=os.environ['GOOGLE_AI_API_KEY'],
82-
max_retries=0,
82+
)
83+
84+
# This wrapper will catch errors and retry the call itself.
85+
# It does NOT pass 'max_retries' to the underlying model.
86+
self._google_llm_with_retries = self._google_llm.with_retry(
87+
stop_after_attempt=1, # Corresponds to max_retries
8388
)
8489

8590
def get_llm(self):
86-
return self._google_llm
91+
return self._google_llm_with_retries

0 commit comments

Comments
 (0)