From 93c5758e029bf6a0efb95d09daca13db181f6f1f Mon Sep 17 00:00:00 2001
From: Oleksandr Kozachuk <ddeus.lp@mailnull.com>
Date: Wed, 29 Mar 2023 19:53:45 +0200
Subject: [PATCH] Some more fixes, reduce retries.

---
 fjerkroa_bot/ai_responder.py | 38 ++++++++++++++++++++----------------
 tests/test_main.py           |  2 +-
 2 files changed, 22 insertions(+), 18 deletions(-)

diff --git a/fjerkroa_bot/ai_responder.py b/fjerkroa_bot/ai_responder.py
index ec74fc8..ae3b4bd 100644
--- a/fjerkroa_bot/ai_responder.py
+++ b/fjerkroa_bot/ai_responder.py
@@ -83,7 +83,7 @@ class AIResponder(object):
         return messages
 
     async def draw(self, description: str) -> BytesIO:
-        for _ in range(7):
+        for _ in range(3):
             try:
                 response = await openai.Image.acreate(prompt=description, n=1, size="512x512")
                 async with aiohttp.ClientSession() as session:
@@ -133,6 +133,7 @@ class AIResponder(object):
             result = await openai.ChatCompletion.acreate(model=self.config["model"],
                                                          messages=messages,
                                                          temperature=self.config["temperature"],
+                                                         max_tokens=self.config["max-tokens"],
                                                          top_p=self.config["top-p"],
                                                          presence_penalty=self.config["presence-penalty"],
                                                          frequency_penalty=self.config["frequency-penalty"])
@@ -155,21 +156,21 @@ class AIResponder(object):
             return answer
         messages = [{"role": "system", "content": self.config["fix-description"]},
                     {"role": "user", "content": answer}]
-        for _ in range(3):
-            try:
-                result = await openai.ChatCompletion.acreate(model=self.config["fix-model"],
-                                                             messages=messages,
-                                                             temperature=0.2)
-                logging.info(f"got this message as fix:\n{pp(result['choices'][0]['message']['content'])}")
-                response = result['choices'][0]['message']['content']
-                start, end = response.find("```"), response.rfind("```")
-                if start == -1 or end == -1 or (start + 3) >= end:
-                    return answer
-                response = response[start + 3:end]
-                logging.info(f"fixed answer:\n{pp(response)}")
-                return response
-            except Exception as err:
-                logging.warning(f"failed to execute a fix for the answer: {repr(err)}")
+        try:
+            result = await openai.ChatCompletion.acreate(model=self.config["fix-model"],
+                                                         messages=messages,
+                                                         temperature=0.2,
+                                                         max_tokens=2048)
+            logging.info(f"got this message as fix:\n{pp(result['choices'][0]['message']['content'])}")
+            response = result['choices'][0]['message']['content']
+            start, end = response.find("{"), response.rfind("}")
+            if start == -1 or end == -1 or (start + 3) >= end:
+                return answer
+            response = response[start:end + 1]
+            logging.info(f"fixed answer:\n{pp(response)}")
+            return response
+        except Exception as err:
+            logging.warning(f"failed to execute a fix for the answer: {repr(err)}")
         return answer
 
     def update_history(self, question: Dict[str, Any], answer: Dict[str, Any], limit: int) -> None:
@@ -185,7 +186,8 @@ class AIResponder(object):
         limit = self.config["history-limit"]
         if self.short_path(message, limit):
             return AIResponse(None, False, None, None, False)
-        for _ in range(5):
+        retries = 3
+        while retries > 0:
             messages = self._message(message, limit)
             logging.info(f"try to send this messages:\n{pp(messages)}")
             answer, limit = await self._acreate(messages, limit)
@@ -200,8 +202,10 @@ class AIResponder(object):
                     response = parse_response(answer['content'])
                 except Exception as err:
                     logging.error(f"failed to parse the fixed answer: {pp(err)}\n{repr(answer['content'])}")
+                    retries -= 1
                     continue
             if 'hack' not in response or type(response.get('picture', None)) not in (type(None), str):
+                retries -= 1
                 continue
             logging.info(f"got this answer:\n{pp(response)}")
             self.update_history(messages[-1], answer, limit)
diff --git a/tests/test_main.py b/tests/test_main.py
index af878f8..1b34f17 100644
--- a/tests/test_main.py
+++ b/tests/test_main.py
@@ -21,7 +21,7 @@ class TestBotBase(unittest.IsolatedAsyncioTestCase):
         self.config_data = {
             "openai-token": os.environ['OPENAI_TOKEN'],
             "model": "gpt-4",
-            "max_tokens": 1024,
+            "max-tokens": 1024,
             "temperature": 0.9,
             "top-p": 1.0,
             "presence-penalty": 1.0,