helloparthshah commited on
Commit
7ce7683
·
1 Parent(s): 4f43fb0

Updating retry behavior

Browse files
Files changed (1) hide show
  1. src/manager/manager.py +7 -0
src/manager/manager.py CHANGED
@@ -16,6 +16,8 @@ from sentence_transformers import SentenceTransformer
16
  import torch
17
  from src.tools.default_tools.memory_manager import MemoryManager
18
  from pathlib import Path
 
 
19
 
20
  logger = logging.getLogger(__name__)
21
  handler = logging.StreamHandler(sys.stdout)
@@ -77,6 +79,10 @@ class GeminiManager:
77
  def check_mode(self, mode: Mode):
78
  return mode in self.modes
79
 
 
 
 
 
80
  def generate_response(self, messages):
81
  tools = self.toolsLoader.getTools()
82
  return self.client.models.generate_content(
@@ -286,6 +292,7 @@ class GeminiManager:
286
  "content": "No response from the model.",
287
  "metadata": {"title": "No response from the model."}
288
  })
 
289
  yield messages
290
  return messages
291
 
 
16
  import torch
17
  from src.tools.default_tools.memory_manager import MemoryManager
18
  from pathlib import Path
19
+ from google.genai.errors import APIError
20
+ import backoff
21
 
22
  logger = logging.getLogger(__name__)
23
  handler = logging.StreamHandler(sys.stdout)
 
79
  def check_mode(self, mode: Mode):
80
  return mode in self.modes
81
 
82
+ @backoff.on_exception(backoff.expo,
83
+ APIError,
84
+ max_tries=3,
85
+ jitter=None)
86
  def generate_response(self, messages):
87
  tools = self.toolsLoader.getTools()
88
  return self.client.models.generate_content(
 
292
  "content": "No response from the model.",
293
  "metadata": {"title": "No response from the model."}
294
  })
295
+ print(response)
296
  yield messages
297
  return messages
298