Changes.
This commit is contained in:
+21
-160
@@ -1,10 +1,7 @@
|
||||
import os
|
||||
import json
|
||||
import asyncio
|
||||
import random
|
||||
import multiline
|
||||
import openai
|
||||
import aiohttp
|
||||
import logging
|
||||
import time
|
||||
import re
|
||||
@@ -84,16 +81,6 @@ def async_cache_to_file(filename):
|
||||
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
|
||||
@@ -152,12 +139,17 @@ class AIResponse(AIMessageBase):
|
||||
self.hack = hack
|
||||
|
||||
|
||||
class AIResponder(object):
|
||||
class AIResponderBase(object):
|
||||
def __init__(self, config: Dict[str, Any], channel: Optional[str] = None) -> None:
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.history: List[Dict[str, Any]] = []
|
||||
self.channel = channel if channel is not None else 'system'
|
||||
self.client = openai.AsyncOpenAI(api_key=self.config['openai-token'])
|
||||
|
||||
|
||||
class AIResponder(AIResponderBase):
|
||||
def __init__(self, config: Dict[str, Any], channel: Optional[str] = None) -> None:
|
||||
super().__init__(config, channel)
|
||||
self.history: List[Dict[str, Any]] = []
|
||||
self.rate_limit_backoff = exponential_backoff()
|
||||
self.history_file: Optional[Path] = None
|
||||
if 'history-directory' in self.config:
|
||||
@@ -166,7 +158,7 @@ class AIResponder(object):
|
||||
with open(self.history_file, 'rb') as fd:
|
||||
self.history = pickle.load(fd)
|
||||
|
||||
def _message(self, message: AIMessage, limit: Optional[int] = None) -> List[Dict[str, Any]]:
|
||||
def message(self, message: AIMessage, limit: Optional[int] = None) -> List[Dict[str, Any]]:
|
||||
messages = []
|
||||
system = self.config.get(self.channel, self.config['system'])
|
||||
system = system.replace('{date}', time.strftime('%Y-%m-%d'))\
|
||||
@@ -187,79 +179,14 @@ class AIResponder(object):
|
||||
|
||||
async def draw(self, description: str) -> BytesIO:
|
||||
if self.config.get('leonardo-token') is not None:
|
||||
return await self._draw_leonardo(description)
|
||||
return await self._draw_openai(description)
|
||||
return await self.draw_leonardo(description)
|
||||
return await self.draw_openai(description)
|
||||
|
||||
async def _draw_openai(self, description: str) -> BytesIO:
|
||||
for _ in range(3):
|
||||
try:
|
||||
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)}')
|
||||
return BytesIO(await image.read())
|
||||
except Exception as err:
|
||||
logging.warning(f"Failed to generate image {repr(description)}: {repr(err)}")
|
||||
raise RuntimeError(f"Failed to generate image {repr(description)} after multiple retries")
|
||||
async def draw_leonardo(self, description: str) -> BytesIO:
|
||||
raise NotImplementedError()
|
||||
|
||||
async def _draw_leonardo(self, description: str) -> BytesIO:
|
||||
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:
|
||||
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, sleep for {error_sleep}s: {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, sleep for {error_sleep}s: {repr(response)}")
|
||||
await asyncio.sleep(error_sleep)
|
||||
continue
|
||||
if len(response["generations_by_pk"]["generated_images"]) == 0:
|
||||
await asyncio.sleep(error_sleep)
|
||||
continue
|
||||
image_url = response["generations_by_pk"]["generated_images"][0]["url"]
|
||||
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:
|
||||
await response.json()
|
||||
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, sleep for {error_sleep}s: {repr(description)}\n{repr(err)}")
|
||||
else:
|
||||
logging.warning(f"Failed to generate image, sleep for {error_sleep}s: {repr(description)}")
|
||||
await asyncio.sleep(error_sleep)
|
||||
raise RuntimeError(f"Failed to generate image {repr(description)}")
|
||||
async def draw_openai(self, description: str) -> BytesIO:
|
||||
raise NotImplementedError()
|
||||
|
||||
async def post_process(self, message: AIMessage, response: Dict[str, Any]) -> AIResponse:
|
||||
for fld in ('answer', 'channel', 'staff', 'picture', 'hack'):
|
||||
@@ -307,80 +234,14 @@ class AIResponder(object):
|
||||
return True
|
||||
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_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()
|
||||
logging.info(f"generated response {result.usage}: {repr(answer)}")
|
||||
return answer, limit
|
||||
except openai.BadRequestError as err:
|
||||
if 'maximum context length is' in str(err) and limit > 4:
|
||||
logging.warning(f"context length exceeded, reduce the limit {limit}: {str(err)}")
|
||||
limit -= 1
|
||||
return None, limit
|
||||
raise err
|
||||
except openai.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:
|
||||
logging.warning(f"failed to generate response: {repr(err)}")
|
||||
return None, limit
|
||||
async def chat(self, messages: List[Dict[str, Any]], limit: int) -> Tuple[Optional[Dict[str, Any]], int]:
|
||||
raise NotImplementedError()
|
||||
|
||||
async def fix(self, answer: str) -> str:
|
||||
if 'fix-model' not in self.config:
|
||||
return answer
|
||||
messages = [{"role": "system", "content": self.config["fix-description"]},
|
||||
{"role": "user", "content": answer}]
|
||||
try:
|
||||
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("}")
|
||||
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
|
||||
raise NotImplementedError()
|
||||
|
||||
async def translate(self, text: str, language: str = "english") -> str:
|
||||
if 'fix-model' not in self.config:
|
||||
return text
|
||||
message = [{"role": "system", "content": f"You are an professional translator to {language} language,"
|
||||
f" you translate everything you get directly to {language}"
|
||||
f" if it is not already in {language}, otherwise you just copy it."},
|
||||
{"role": "user", "content": text}]
|
||||
try:
|
||||
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
|
||||
except Exception as err:
|
||||
logging.warning(f"failed to translate the text: {repr(err)}")
|
||||
return text
|
||||
raise NotImplementedError()
|
||||
|
||||
def shrink_history_by_one(self, index: int = 0) -> None:
|
||||
if index >= len(self.history):
|
||||
@@ -420,11 +281,11 @@ class AIResponder(object):
|
||||
|
||||
while retries > 0:
|
||||
# Get the message queue
|
||||
messages = self._message(message, limit)
|
||||
messages = self.message(message, limit)
|
||||
logging.info(f"try to send this messages:\n{pp(messages)}")
|
||||
|
||||
# Attempt to send the message to the AI
|
||||
answer, limit = await self._acreate(messages, limit)
|
||||
answer, limit = await self.chat(messages, limit)
|
||||
|
||||
if answer is None:
|
||||
continue
|
||||
|
||||
Reference in New Issue
Block a user