136 lines
4.9 KiB
Python
136 lines
4.9 KiB
Python
import aiohttp
|
|
import base64
|
|
from database import Database
|
|
|
|
API_KEY = "eh"
|
|
MODEL = "p620"
|
|
DEFAULT_NAME = "Aoi"
|
|
NAME_PROMPT = "reply with your name, nothing else, no punctuation"
|
|
|
|
|
|
async def get_name(client, model, prompt):
|
|
"""Generates an assistant name for the given prompt."""
|
|
name_response = await client.chat.completions.create(
|
|
model=model,
|
|
messages=[
|
|
{"role": "system", "content": prompt},
|
|
{"role": "user", "content": NAME_PROMPT}
|
|
],
|
|
)
|
|
return name_response.choices[0].message.content.split('\n')[0]
|
|
|
|
|
|
class ConversationManager:
|
|
"""Creates and retrieves Conversations."""
|
|
def __init__(self, openai_client, model, db, default_prompt):
|
|
self.model = model
|
|
self.client = openai_client
|
|
self.db = db
|
|
self.default_prompt = default_prompt
|
|
|
|
async def get(self, key, create_if_missing=True):
|
|
"""Gets a conversation based on |key|, optionally create when not found."""
|
|
convo_data = self.db.get_conversation(key)
|
|
if convo_data:
|
|
history, bot_name, last_messages = convo_data
|
|
return Conversation(
|
|
key, bot_name, history, last_messages,
|
|
self.client, self.model, self.db,
|
|
)
|
|
if create_if_missing:
|
|
return await self.new_conversation(key, self.default_prompt)
|
|
return None
|
|
|
|
async def new_conversation(self, key, prompt = None):
|
|
"""Creates a new Conversation with key based on given prompt."""
|
|
prompt = prompt or self.default_prompt
|
|
name = await get_name(self.client, self.model, prompt)
|
|
history = [{"role": "system", "content": prompt}]
|
|
last_messages = []
|
|
convo = Conversation(
|
|
key, name, history, last_messages, self.client, self.model, self.db,
|
|
)
|
|
await convo.save()
|
|
return convo
|
|
|
|
|
|
class Conversation:
|
|
"""Holds data about a conversation thread."""
|
|
def __init__(
|
|
self, convo_id, name, history, last_messages, api_client, model, db,
|
|
):
|
|
self.id = convo_id
|
|
self.bot_name = name
|
|
self.history = history
|
|
self.last_messages = last_messages
|
|
self.client = api_client
|
|
self.model = model
|
|
self.db = db
|
|
|
|
async def save(self):
|
|
"""Saves the conversation to the DB."""
|
|
self.db.save(self.id, self.history, self.bot_name, self.last_messages)
|
|
|
|
def add_message_pair(self, user, assistant):
|
|
"""Adds a user/assistant convesation turn pair."""
|
|
self.history.extend([
|
|
{"role": "user", "content": user},
|
|
{"role": "assistant", "content": assistant},
|
|
])
|
|
|
|
async def pop(self):
|
|
"""Removes one user/assistant converation turn pair."""
|
|
if len(self.history) >= 3:
|
|
self.history = self.history[:-2]
|
|
await self.save()
|
|
|
|
async def update_prompt(self, prompt):
|
|
"""Changes current prompt to a new one, keeping the rest of history."""
|
|
self.history[0] = {"role": "system", "content": prompt}
|
|
self.bot_name = await get_name(self.client, self.model, prompt)
|
|
await self.save()
|
|
|
|
async def generate(self, text, media=tuple()):
|
|
"""Generates next assistant conversation turn."""
|
|
# prepare text part
|
|
if text:
|
|
openai_content = [{"type": "text", "text": text}]
|
|
else:
|
|
openai_content = [{"type": "text", "text": "."}]
|
|
|
|
# prepare images part
|
|
async with aiohttp.ClientSession() as session:
|
|
for (content_type, url) in media:
|
|
if "image" not in content_type:
|
|
continue
|
|
try:
|
|
async with session.get(url) as resp:
|
|
resp.raise_for_status()
|
|
image_data = await resp.read()
|
|
b64_image = base64.b64encode(image_data).decode('utf-8')
|
|
b64_url = f"data:{content_type};base64,{b64_image}"
|
|
openai_content.append({
|
|
"type": "image_url",
|
|
"image_url": {"url": b64_url}
|
|
})
|
|
except Exception as e:
|
|
print(f"Error downloading or processing attachment: {e}")
|
|
|
|
# send request to openai api and return response
|
|
request = self.history + [{"role": "user", "content": openai_content}]
|
|
llm_response = await self.client.chat.completions.create(
|
|
model=MODEL, messages=request,
|
|
)
|
|
response = llm_response.choices[0].message.content
|
|
self.add_message_pair(openai_content, response)
|
|
return response
|
|
|
|
async def regenerate(self):
|
|
"""Regenerates the last assistant turn."""
|
|
llm_response = await self.client.chat.completions.create(
|
|
model=MODEL, messages=self.history[:-1]
|
|
)
|
|
response = llm_response.choices[0].message.content
|
|
self.history[-1] = {"role": "assistant", "content": response}
|
|
return response
|