From 95bc6ce041808967745e6062677ed738652066a8 Mon Sep 17 00:00:00 2001 From: Oleksandr Kozachuk Date: Tue, 22 Aug 2023 11:05:17 +0200 Subject: [PATCH] Add possibility to set `retry-model`, if AI request failes because of limits. Improve error handling in leanardo AI drawing. --- fjerkroa_bot/ai_responder.py | 68 ++++++++++++++++++++++++------------ 1 file changed, 45 insertions(+), 23 deletions(-) diff --git a/fjerkroa_bot/ai_responder.py b/fjerkroa_bot/ai_responder.py index 49ad722..f9db137 100644 --- a/fjerkroa_bot/ai_responder.py +++ b/fjerkroa_bot/ai_responder.py @@ -33,7 +33,7 @@ def parse_json(content: str) -> Dict: raise err -def exponential_backoff(base=2, max_delay=60, factor=1, jitter=0.1): +def exponential_backoff(base=2, max_delay=60, factor=1, jitter=0.1, max_attempts=None): """Generate sleep intervals for exponential backoff with jitter. Args: @@ -52,6 +52,8 @@ def exponential_backoff(base=2, max_delay=60, factor=1, jitter=0.1): sleep += random.uniform(-jitter_amount, jitter_amount) yield sleep attempt += 1 + if max_attempts is not None and attempt > max_attempts: + raise RuntimeError("Max attempts reached in exponential backoff.") def parse_maybe_json(json_string): @@ -158,36 +160,51 @@ class AIResponder(object): raise RuntimeError(f"Failed to generate image {repr(description)} after multiple retries") async def _draw_leonardo(self, description: str) -> BytesIO: - for _ in range(3): + error_backoff = exponential_backoff(max_attempts=12) + generation_id = None + image_url = None + image_bytes = None + while True: + error_sleep = next(error_backoff) try: async with aiohttp.ClientSession() as session: - async with session.post("https://cloud.leonardo.ai/api/rest/v1/generations", - json={"prompt": description, - "modelId": "6bef9f1b-29cb-40c7-b9df-32b51c1f67d3", - "num_images": 1, - "sd_version": "v2", - "promptMagic": True, - "width": 512, - "height": 512}, - headers={"Authorization": f"Bearer {self.config['leonardo-token']}", - "Accept": "application/json", - "Content-Type": "application/json"}, - ) as response: - response = await response.json() - generation_id = response["sdGenerationJob"]["generationId"] - while True: + if generation_id is None: + async with session.post("https://cloud.leonardo.ai/api/rest/v1/generations", + json={"prompt": description, + "modelId": "6bef9f1b-29cb-40c7-b9df-32b51c1f67d3", + "num_images": 1, + "sd_version": "v2", + "promptMagic": True, + "unzoomAmount": 1, + "width": 512, + "height": 512}, + headers={"Authorization": f"Bearer {self.config['leonardo-token']}", + "Accept": "application/json", + "Content-Type": "application/json"}, + ) as response: + response = await response.json() + if "sdGenerationJob" not in response: + logging.warning(f"No 'sdGenerationJob' found in response: {repr(response)}") + await asyncio.sleep(error_sleep) + continue + generation_id = response["sdGenerationJob"]["generationId"] + if image_url is None: async with session.get(f"https://cloud.leonardo.ai/api/rest/v1/generations/{generation_id}", headers={"Authorization": f"Bearer {self.config['leonardo-token']}", "Accept": "application/json"}, ) as response: response = await response.json() + if "generations_by_pk" not in response: + logging.warning(f"Unexpected response: {repr(response)}") + await asyncio.sleep(error_sleep) + continue if len(response["generations_by_pk"]["generated_images"]) == 0: - await asyncio.sleep(0.1) + await asyncio.sleep(error_sleep) continue image_url = response["generations_by_pk"]["generated_images"][0]["url"] - break - async with session.get(image_url) as response: - image_bytes = BytesIO(await response.read()) + if image_bytes is None: + async with session.get(image_url) as response: + image_bytes = BytesIO(await response.read()) async with session.delete(f"https://cloud.leonardo.ai/api/rest/v1/generations/{generation_id}", headers={"Authorization": f"Bearer {self.config['leonardo-token']}"}, ) as response: @@ -195,7 +212,9 @@ class AIResponder(object): logging.info(f'Drawed a picture with leonardo AI on this description: {repr(description)}') return image_bytes except Exception as err: - logging.warning(f"Failed to generate image: {repr(err)}") + logging.warning(f"Failed to generate image: {repr(description)}\n{repr(err)}") + else: + logging.warning(f"Failed to generate image: {repr(description)}") raise RuntimeError(f"Failed to generate image {repr(description)}") async def post_process(self, message: AIMessage, response: Dict[str, Any]) -> AIResponse: @@ -245,8 +264,9 @@ class AIResponder(object): return False async def _acreate(self, messages: List[Dict[str, Any]], limit: int) -> Tuple[Optional[Dict[str, Any]], int]: + model = self.config["model"] try: - result = await openai.ChatCompletion.acreate(model=self.config["model"], + result = await openai.ChatCompletion.acreate(model=model, messages=messages, temperature=self.config["temperature"], max_tokens=self.config["max-tokens"], @@ -267,6 +287,8 @@ class AIResponder(object): raise err except openai.error.RateLimitError as err: rate_limit_sleep = next(self.rate_limit_backoff) + if "retry-model" in self.config: + model = self.config["retry-model"] logging.warning(f"got an rate limit error, sleep for {rate_limit_sleep} seconds: {str(err)}") await asyncio.sleep(rate_limit_sleep) except Exception as err: