Fix and remove some tests, make openai calls cachable in a file.

fjerkroa_bot/ai_responder.py

@@ -1,3 +1,4 @@
+import sys
 import os
 import json
 import asyncio
@@ -12,7 +13,7 @@ import pickle
 from pathlib import Path
 from io import BytesIO
 from pprint import pformat
-from functools import lru_cache
+from functools import lru_cache, wraps
 from typing import Optional, List, Dict, Any, Tuple
 
 
@@ -57,6 +58,46 @@ def exponential_backoff(base=2, max_delay=60, factor=1, jitter=0.1, max_attempts
             raise RuntimeError("Max attempts reached in exponential backoff.")
 
 
+def async_cache_to_file(filename):
+    cache_file = Path(filename)
+    cache = None
+    if cache_file.exists():
+        try:
+            with cache_file.open('rb') as fd:
+                cache = pickle.load(fd)
+        except Exception:
+            cache = {}
+
+    def decorator(func):
+        @wraps(func)
+        async def wrapper(*args, **kwargs):
+            if cache is None:
+                sys.stderr.write(f'@@@ forward {func.__name__}({repr(args)}, {repr(kwargs)}')
+                return await func(*args, **kwargs)
+            key = json.dumps((func.__name__, list(args[1:]), kwargs), sort_keys=True)
+            if key in cache:
+                sys.stderr.write(f'@@@ cache {func.__name__}({repr(args)}, {repr(kwargs)} -> {cache[key]}')
+                return cache[key]
+            sys.stderr.write(f'@@@ execute {func.__name__}({repr(args)}, {repr(kwargs)}')
+            result = await func(*args, **kwargs)
+            cache[key] = result
+            with cache_file.open('wb') as fd:
+                pickle.dump(cache, fd)
+            return result
+        return wrapper
+    return decorator
+
+
+@async_cache_to_file('openai_chat.dat')
+async def openai_chat(client, *args, **kwargs):
+    return await client.chat.completions.create(*args, **kwargs)
+
+
+@async_cache_to_file('openai_chat.dat')
+async def openai_image(client, *args, **kwargs):
+    return await client.images.generate(*args, **kwargs)
+
+
 def parse_maybe_json(json_string):
     if json_string is None:
         return None
@@ -156,7 +197,7 @@ class AIResponder(object):
     async def _draw_openai(self, description: str) -> BytesIO:
         for _ in range(3):
             try:
-                response = await self.client.images.generate(prompt=description, n=1, size="1024x1024", model="dall-e-3")
+                response = await openai_image(self.client, prompt=description, n=1, size="1024x1024", model="dall-e-3")
                 async with aiohttp.ClientSession() as session:
                     async with session.get(response.data[0].url) as image:
                         logging.info(f'Drawed a picture with DALL-E on this description: {repr(description)}')
@@ -273,13 +314,14 @@ class AIResponder(object):
     async def _acreate(self, messages: List[Dict[str, Any]], limit: int) -> Tuple[Optional[Dict[str, Any]], int]:
         model = self.config["model"]
         try:
-            result = await self.client.chat.completions.create(model=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"])
+            result = await openai_chat(self.client,
+                                       model=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"])
             answer_obj = result.choices[0].message
             answer = {'content': answer_obj.content, 'role': answer_obj.role}
             self.rate_limit_backoff = exponential_backoff()
@@ -307,10 +349,11 @@ class AIResponder(object):
         messages = [{"role": "system", "content": self.config["fix-description"]},
                     {"role": "user", "content": answer}]
         try:
-            result = await self.client.chat.completions.create(model=self.config["fix-model"],
-                                                               messages=messages,
-                                                               temperature=0.2,
-                                                               max_tokens=2048)
+            result = await openai_chat(self.client,
+                                       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("}")
@@ -331,10 +374,11 @@ class AIResponder(object):
                                                  f" if it is not already in {language}, otherwise you just copy it."},
                    {"role": "user", "content": text}]
         try:
-            result = await self.client.chat.completions.create(model=self.config["fix-model"],
-                                                               messages=message,
-                                                               temperature=0.2,
-                                                               max_tokens=2048)
+            result = await openai_chat(self.client,
+                                       model=self.config["fix-model"],
+                                       messages=message,
+                                       temperature=0.2,
+                                       max_tokens=2048)
             response = result.choices[0].message.content
             logging.info(f"got this translated message:\n{pp(response)}")
             return response