diff --git a/docker-compose.yml b/docker-compose.yml index 7a11aa8f..6ecafbba 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -2,7 +2,7 @@ version: "3" services: postgres: - image: ankane/pgvector + image: postgres:15 ports: - 5432:5432 environment: diff --git a/src/database.py b/src/database.py index 16187a52..1729d255 100644 --- a/src/database.py +++ b/src/database.py @@ -4,7 +4,6 @@ import asyncpg import logger import migration -from pgvector.asyncpg import register_vector log = logger.get("DATABASE") @@ -31,16 +30,11 @@ async def get_pool(): _connect_string, min_size=MIN_CONNECTION_POOL_SIZE, max_size=MAX_CONNECTION_POOL_SIZE, - init=init_connection, ) setattr(_pool_holder, "pool", pool) return pool -async def init_connection(conn): - await register_vector(conn) - - async def close_pool(): global _pool_holder pool = getattr(_pool_holder, "pool", None) diff --git a/src/migration/00055-vector-for-openai-embeddings.py b/src/migration/00055-vector-for-openai-embeddings.py index c42eef99..a9744b3b 100644 --- a/src/migration/00055-vector-for-openai-embeddings.py +++ b/src/migration/00055-vector-for-openai-embeddings.py @@ -1,4 +1,9 @@ async def exec(log, tx): + # These fail if run without the extension installed which + # means these migrations can't be run on a fresh database + # without the now unnecessary pgvector extension + return + await tx.execute("CREATE EXTENSION vector") await tx.execute(""" CREATE TABLE openaiembedding ( diff --git a/src/migration/00056-remove-pgvector-used-for-embeddings.py b/src/migration/00056-remove-pgvector-used-for-embeddings.py new file mode 100644 index 00000000..f1dee28f --- /dev/null +++ b/src/migration/00056-remove-pgvector-used-for-embeddings.py @@ -0,0 +1,3 @@ +async def exec(log, tx): + await tx.execute("DROP TABLE IF EXISTS openaiembedding") + await tx.execute("DROP EXTENSION IF EXISTS vector") diff --git a/src/openai.py b/src/openai.py index ab0391a3..32f611cf 100644 --- a/src/openai.py +++ b/src/openai.py @@ -11,7 +11,6 @@ import retry import database as db import util -import perf log = logger.get("OPENAI") @@ -29,66 +28,9 @@ def register(client): return { 'setprompt': cmd_setprompt, 'genimage': cmd_genimage, - 'search': cmd_search, } -async def main(): - while True: - try: - await generate_embeddings() - except Exception: - await util.log_exception(log) - await asyncio.sleep(120) - - -async def cmd_search(client, message, arg): - messages = await search_embedding(message.guild.id, arg) - contents = list(map(lambda r: f"Result with distance {r['distance']}:\n{r['content']}", messages)) - response = "\n".join(contents) - reply_target = message - for msg in util.split_message_for_sending(response.split("\n")): - reply_target = await reply_target.reply(msg) - - -async def search_embedding(guild_id, query): - response = await embeddings([query]) - embedding = response["data"][0]["embedding"] - return await db.fetch(""" - SELECT message_id, content, embedding <#> $2 AS distance - FROM message JOIN openaiembedding USING (message_id) - WHERE guild_id = $1 -- AND (embedding <#> $2) < -0.85 - -- OpenAI embeddings are normalized to length 1 so this is best performance for exact search - -- https://platform.openai.com/docs/guides/embeddings/which-distance-function-should-i-use - ORDER BY embedding <#> $2 - LIMIT 5 - """, str(guild_id), embedding) - - -@perf.time_async("generate_embeddings") -async def generate_embeddings(): - rows = await db.fetch(""" - SELECT message_id, content - FROM message - WHERE NOT bot and content != '' AND content NOT LIKE '!%' - AND NOT EXISTS (SELECT 1 FROM openaiembedding WHERE openaiembedding.message_id = message.message_id) - LIMIT 1000 - """) - if len(rows) == 0: - return - - contents = list(map(lambda r: r["content"], rows)) - response = await embeddings(contents) - - async with db.transaction() as tx: - def mk_row(output): return rows[output["index"]]["message_id"], output["embedding"] - queryparams = list(mk_row(output) for output in response["data"]) - await tx.executemany( - "INSERT INTO openaiembedding(message_id, embedding) VALUES ($1, $2)", - queryparams - ) - - async def handle_message(client, message): bot_mentioned = any(user for user in message.mentions if user.id == client.user.id) if not bot_mentioned: @@ -194,22 +136,13 @@ async def get_response_for_messages(messages): }) -async def embeddings(input, model="text-embedding-ada-002"): - status, response = await _call_api("/v1/embeddings", json_body={ - "model": model, - "input": input, - }, skip_log=True) - return response - - @retry.on_any_exception(max_attempts = 1, init_delay = 1, max_delay = 30) -async def _call_api(path, json_body=None, query=None, skip_log=False): +async def _call_api(path, json_body=None, query=None): url = "https://api.openai.com{0}{1}".format(path, http_util.make_query_string(query)) async with aiohttp.ClientSession() as session: for ratelimit_delay in retry.jitter(retry.exponential(1, 128)): response = await session.post(url, headers=AUTH_HEADER, json=json_body) - log_fn = log.debug if skip_log and response.status == 200 else log.info - log_fn({ + log.info({ "requestMethod": response.method, "requestUrl": str(response.url), "responseStatus": response.status, diff --git a/src/run_lemon_bot.py b/src/run_lemon_bot.py index 14322bf9..31be386f 100755 --- a/src/run_lemon_bot.py +++ b/src/run_lemon_bot.py @@ -91,7 +91,6 @@ async def main(): # Database schema has to be initialized before running the bot await db.initialize_schema() asyncio.create_task(archiver.main()) - asyncio.create_task(openai.main()) await trophies.main() for module in [casino, sqlcommands, osu, feed, reminder, youtube, lan, steam, anssicommands, trophies, laiva,