From 8417442d0fabe0312c6e0bc70a7ca818428d49ad Mon Sep 17 00:00:00 2001 From: Justin Date: Fri, 19 May 2023 11:08:48 -0500 Subject: [PATCH 1/9] Fix migrations off by one --- migrations/__init__.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/migrations/__init__.py b/migrations/__init__.py index 1cf8d53..02a7168 100644 --- a/migrations/__init__.py +++ b/migrations/__init__.py @@ -36,7 +36,7 @@ def migrate(db: DuckDBPyConnection, from_version: Optional[int] = None, to_versi """ if from_version is None: - from_version = get_version(db) + from_version = get_version(db) + 1 if to_version is None: to_version = max(MIGRATIONS.keys()) @@ -44,8 +44,10 @@ def migrate(db: DuckDBPyConnection, from_version: Optional[int] = None, to_versi if from_version > to_version: raise ValueError("Cannot migrate from a higher version to a lower version.") + print(f"Migrating from version {from_version} to {to_version}...") for version in range(from_version, to_version): if version in MIGRATIONS: - MIGRATIONS[version + 1](db) + print(f"Running migration {version}...") + MIGRATIONS[version](db) - return from_version, to_version \ No newline at end of file + return from_version, to_version From e200393670f04395c6023a5b8a67ed35542cd7a8 Mon Sep 17 00:00:00 2001 From: Justin Date: Fri, 19 May 2023 11:12:31 -0500 Subject: [PATCH 2/9] Cleanup unused imports. Fixes a few variables --- classes/bot.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/classes/bot.py b/classes/bot.py index a8681e4..bbde9b2 100644 --- a/classes/bot.py +++ b/classes/bot.py @@ -1,7 +1,6 @@ import markdown2 import duckdb import tiktoken -import magic import asyncio from PIL import Image @@ -27,12 +26,11 @@ from nio import ( RoomLeaveError, RoomSendError, RoomVisibility, - RoomCreateResponse, RoomCreateError, ) from nio.crypto import Olm -from typing import Optional, List, Dict, Tuple +from typing import Optional, List from configparser import ConfigParser from datetime import datetime from io import BytesIO @@ -174,7 +172,7 @@ class GPTBot: async def _last_n_messages(self, room: str | MatrixRoom, n: Optional[int]): messages = [] - n = n or bot.max_messages + n = n or self.max_messages room_id = room.room_id if isinstance(room, MatrixRoom) else room self.logger.log( @@ -585,7 +583,7 @@ class GPTBot: self.logger.log( "No database connection set up, using in-memory database. Data will be lost on bot shutdown.") IN_MEMORY = True - self.database = DuckDBPyConnection(":memory:") + self.database = duckdb.DuckDBPyConnection(":memory:") self.logger.log("Running migrations...") before, after = migrate(self.database) @@ -803,7 +801,7 @@ class GPTBot: else: # Send a notice to the room if there was an error self.logger.log("Didn't get a response from GPT API", "error") - await send_message( + await self.send_message( room, "Something went wrong. Please try again.", True) await self.matrix_client.room_typing(room.room_id, False) From 9c2c4d5f6f463d41baddbeb2cf9863a476755639 Mon Sep 17 00:00:00 2001 From: Justin Date: Fri, 19 May 2023 11:26:07 -0500 Subject: [PATCH 3/9] Run OpenAI calls async --- classes/bot.py | 15 +++++++++++++-- commands/classify.py | 21 +++++++++++++++++++-- commands/imagine.py | 21 +++++++++++++++++++-- 3 files changed, 51 insertions(+), 6 deletions(-) diff --git a/classes/bot.py b/classes/bot.py index bbde9b2..149f930 100644 --- a/classes/bot.py +++ b/classes/bot.py @@ -2,6 +2,7 @@ import markdown2 import duckdb import tiktoken import asyncio +import functools from PIL import Image @@ -780,8 +781,18 @@ class GPTBot: chat_messages, self.max_tokens - 1, system_message=system_message) try: - response, tokens_used = self.chat_api.generate_chat_response( - chat_messages, user=room.room_id) + loop = asyncio.get_event_loop() + except Exception as e: + self.logger.log(f"Error getting event loop: {e}", "error") + await self.send_message( + room, "Something went wrong. Please try again.", True) + return + + try: + chat_partial = functools.partial(self.chat_api.generate_chat_response, truncated_messages, user=room.room_id) + response, tokens_used = await loop.run_in_executor(None, chat_partial) + # response, tokens_used = self.chat_api.generate_chat_response( + # chat_messages, user=room.room_id) except Exception as e: self.logger.log(f"Error generating response: {e}", "error") await self.send_message( diff --git a/commands/classify.py b/commands/classify.py index e5813ff..fba23a7 100644 --- a/commands/classify.py +++ b/commands/classify.py @@ -1,3 +1,6 @@ +import asyncio +import functools + from nio.events.room_events import RoomMessageText from nio.rooms import MatrixRoom @@ -8,7 +11,21 @@ async def command_classify(room: MatrixRoom, event: RoomMessageText, bot): if prompt: bot.logger.log("Classifying message...") - response, tokens_used = bot.classification_api.classify_message(prompt, user=room.room_id) + try: + loop = asyncio.get_event_loop() + except Exception as e: + bot.logger.log(f"Error getting event loop: {e}", "error") + await bot.send_message( + room, "Something went wrong. Please try again.", True) + return + + try: + classify_partial = functools.partial(bot.classification_api.classify_message, prompt, user=room.room_id) + response, tokens_used = await loop.run_in_executor(None, classify_partial) + except Exception as e: + bot.logger.log(f"Error classifying message: {e}", "error") + await bot.send_message(room, "Sorry, I couldn't classify the message. Please try again later.", True) + return message = f"The message you provided seems to be of type: {response['type']}." @@ -21,4 +38,4 @@ async def command_classify(room: MatrixRoom, event: RoomMessageText, bot): return - await bot.send_message(room, "You need to provide a prompt.", True) \ No newline at end of file + await bot.send_message(room, "You need to provide a prompt.", True) diff --git a/commands/imagine.py b/commands/imagine.py index 54e6c71..90b7930 100644 --- a/commands/imagine.py +++ b/commands/imagine.py @@ -1,3 +1,6 @@ +import asyncio +import functools + from nio.events.room_events import RoomMessageText from nio.rooms import MatrixRoom @@ -8,7 +11,21 @@ async def command_imagine(room: MatrixRoom, event: RoomMessageText, bot): if prompt: bot.logger.log("Generating image...") - images, tokens_used = bot.image_api.generate_image(prompt, user=room.room_id) + try: + loop = asyncio.get_event_loop() + except Exception as e: + bot.logger.log(f"Error getting event loop: {e}", "error") + await bot.send_message( + room, "Something went wrong. Please try again.", True) + return + + try: + image_partial = functools.partial(bot.image_api.generate_image, prompt, user=room.room_id) + images, tokens_used = await loop.run_in_executor(None, image_partial) + except Exception as e: + bot.logger.log(f"Error generating image: {e}", "error") + await bot.send_message(room, "Sorry, I couldn't generate an image. Please try again later.", True) + return for image in images: bot.logger.log(f"Sending image...") @@ -18,4 +35,4 @@ async def command_imagine(room: MatrixRoom, event: RoomMessageText, bot): return - await bot.send_message(room, "You need to provide a prompt.", True) \ No newline at end of file + await bot.send_message(room, "You need to provide a prompt.", True) From a546e969cc95c042d73ea221d11c96e6b9ec7f36 Mon Sep 17 00:00:00 2001 From: Justin Date: Fri, 19 May 2023 12:06:19 -0500 Subject: [PATCH 4/9] Fix latest migration not applying --- migrations/__init__.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/migrations/__init__.py b/migrations/__init__.py index 02a7168..9724944 100644 --- a/migrations/__init__.py +++ b/migrations/__init__.py @@ -24,7 +24,7 @@ def get_version(db: DuckDBPyConnection) -> int: try: return int(db.execute("SELECT MAX(id) FROM migrations").fetchone()[0]) except: - return 0 + return 1 def migrate(db: DuckDBPyConnection, from_version: Optional[int] = None, to_version: Optional[int] = None) -> None: """Migrate the database to a specific version. @@ -36,7 +36,7 @@ def migrate(db: DuckDBPyConnection, from_version: Optional[int] = None, to_versi """ if from_version is None: - from_version = get_version(db) + 1 + from_version = get_version(db) if to_version is None: to_version = max(MIGRATIONS.keys()) @@ -44,8 +44,12 @@ def migrate(db: DuckDBPyConnection, from_version: Optional[int] = None, to_versi if from_version > to_version: raise ValueError("Cannot migrate from a higher version to a lower version.") + # dont migrate if already at the target version + if from_version == to_version: + return from_version, to_version + print(f"Migrating from version {from_version} to {to_version}...") - for version in range(from_version, to_version): + for version in range(from_version, to_version + 1): if version in MIGRATIONS: print(f"Running migration {version}...") MIGRATIONS[version](db) From f118a23714ed23d8ff205f1b6ef0dcf25f931362 Mon Sep 17 00:00:00 2001 From: Justin Date: Fri, 19 May 2023 12:43:19 -0500 Subject: [PATCH 5/9] Makes OpenAI calls awaitable --- classes/bot.py | 24 +++++++--------- classes/openai.py | 68 ++++++++++++++++++++++++++++++-------------- commands/classify.py | 14 +-------- commands/imagine.py | 14 +-------- 4 files changed, 58 insertions(+), 62 deletions(-) diff --git a/classes/bot.py b/classes/bot.py index 149f930..c1679b5 100644 --- a/classes/bot.py +++ b/classes/bot.py @@ -746,8 +746,14 @@ class GPTBot: await self.matrix_client.room_read_markers(room.room_id, event.event_id) if (not from_chat_command) and self.room_uses_classification(room): - classification, tokens = self.classification_api.classify_message( - event.body, room.room_id) + try: + classification, tokens = await self.classification_api.classify_message( + event.body, room.room_id) + except Exception as e: + self.logger.log(f"Error classifying message: {e}", "error") + await self.send_message( + room, "Something went wrong. Please try again.", True) + return self.log_api_usage( event, room, f"{self.classification_api.api_code}-{self.classification_api.classification_api}", tokens) @@ -781,18 +787,8 @@ class GPTBot: chat_messages, self.max_tokens - 1, system_message=system_message) try: - loop = asyncio.get_event_loop() - except Exception as e: - self.logger.log(f"Error getting event loop: {e}", "error") - await self.send_message( - room, "Something went wrong. Please try again.", True) - return - - try: - chat_partial = functools.partial(self.chat_api.generate_chat_response, truncated_messages, user=room.room_id) - response, tokens_used = await loop.run_in_executor(None, chat_partial) - # response, tokens_used = self.chat_api.generate_chat_response( - # chat_messages, user=room.room_id) + response, tokens_used = await self.chat_api.generate_chat_response( + chat_messages, user=room.room_id) except Exception as e: self.logger.log(f"Error generating response: {e}", "error") await self.send_message( diff --git a/classes/openai.py b/classes/openai.py index 249701c..c43aba5 100644 --- a/classes/openai.py +++ b/classes/openai.py @@ -1,6 +1,8 @@ import openai import requests +import asyncio +import functools import json from .logging import Logger @@ -17,7 +19,7 @@ class OpenAI: @property def chat_api(self) -> str: return self.chat_model - + classification_api = chat_api image_api: str = "dalle" @@ -28,7 +30,7 @@ class OpenAI: self.chat_model = chat_model or self.chat_model self.logger = logger or Logger() - def generate_chat_response(self, messages: List[Dict[str, str]], user: Optional[str] = None) -> Tuple[str, int]: + async def generate_chat_response(self, messages: List[Dict[str, str]], user: Optional[str] = None) -> Tuple[str, int]: """Generate a response to a chat message. Args: @@ -37,22 +39,29 @@ class OpenAI: Returns: Tuple[str, int]: The response text and the number of tokens used. """ + try: + loop = asyncio.get_event_loop() + except Exception as e: + self.logger.log(f"Error getting event loop: {e}", "error") + return self.logger.log(f"Generating response to {len(messages)} messages using {self.chat_model}...") - response = openai.ChatCompletion.create( - model=self.chat_model, - messages=messages, - api_key=self.api_key, - user = user + chat_partial = functools.partial( + openai.ChatCompletion.create, + model=self.chat_model, + messages=messages, + api_key=self.api_key, + user = user ) + response = await loop.run_in_executor(None, chat_partial) result_text = response.choices[0].message['content'] tokens_used = response.usage["total_tokens"] self.logger.log(f"Generated response with {tokens_used} tokens.") return result_text, tokens_used - def classify_message(self, query: str, user: Optional[str] = None) -> Tuple[Dict[str, str], int]: + async def classify_message(self, query: str, user: Optional[str] = None) -> Tuple[Dict[str, str], int]: system_message = """You are a classifier for different types of messages. You decide whether an incoming message is meant to be a prompt for an AI chat model, or meant for a different API. You respond with a JSON object like this: { "type": event_type, "prompt": prompt } @@ -66,10 +75,15 @@ class OpenAI: - If for any reason you are unable to classify the message (for example, if it infringes on your terms of service), the event_type is "error", and the prompt is a message explaining why you are unable to process the message. Only the event_types mentioned above are allowed, you must not respond in any other way.""" + try: + loop = asyncio.get_event_loop() + except Exception as e: + self.logger.log(f"Error getting event loop: {e}", "error") + return messages = [ { - "role": "system", + "role": "system", "content": system_message }, { @@ -80,12 +94,14 @@ Only the event_types mentioned above are allowed, you must not respond in any ot self.logger.log(f"Classifying message '{query}'...") - response = openai.ChatCompletion.create( - model=self.chat_model, - messages=messages, - api_key=self.api_key, - user = user + chat_partial = functools.partial( + openai.ChatCompletion.create, + model=self.chat_model, + messages=messages, + api_key=self.api_key, + user=user ) + response = await loop.run_in_executor(None, chat_partial) try: result = json.loads(response.choices[0].message['content']) @@ -98,7 +114,7 @@ Only the event_types mentioned above are allowed, you must not respond in any ot return result, tokens_used - def generate_image(self, prompt: str, user: Optional[str] = None) -> Generator[bytes, None, None]: + async def generate_image(self, prompt: str, user: Optional[str] = None) -> Generator[bytes, None, None]: """Generate an image from a prompt. Args: @@ -107,16 +123,24 @@ Only the event_types mentioned above are allowed, you must not respond in any ot Yields: bytes: The image data. """ + try: + loop = asyncio.get_event_loop() + except Exception as e: + self.logger.log(f"Error getting event loop: {e}", "error") + return + self.logger.log(f"Generating image from prompt '{prompt}'...") - response = openai.Image.create( - prompt=prompt, - n=1, - api_key=self.api_key, - size="1024x1024", - user = user + image_partial = functools.partial( + openai.Image.create, + prompt=prompt, + n=1, + api_key=self.api_key, + size="1024x1024", + user = user ) + response = await loop.run_in_executor(None, image_partial) images = [] @@ -124,4 +148,4 @@ Only the event_types mentioned above are allowed, you must not respond in any ot image = requests.get(image.url).content images.append(image) - return images, len(images) \ No newline at end of file + return images, len(images) diff --git a/commands/classify.py b/commands/classify.py index fba23a7..c07bbfc 100644 --- a/commands/classify.py +++ b/commands/classify.py @@ -1,6 +1,3 @@ -import asyncio -import functools - from nio.events.room_events import RoomMessageText from nio.rooms import MatrixRoom @@ -12,16 +9,7 @@ async def command_classify(room: MatrixRoom, event: RoomMessageText, bot): bot.logger.log("Classifying message...") try: - loop = asyncio.get_event_loop() - except Exception as e: - bot.logger.log(f"Error getting event loop: {e}", "error") - await bot.send_message( - room, "Something went wrong. Please try again.", True) - return - - try: - classify_partial = functools.partial(bot.classification_api.classify_message, prompt, user=room.room_id) - response, tokens_used = await loop.run_in_executor(None, classify_partial) + response, tokens_used = await bot.classification_api.classify_message(prompt, user=room.room_id) except Exception as e: bot.logger.log(f"Error classifying message: {e}", "error") await bot.send_message(room, "Sorry, I couldn't classify the message. Please try again later.", True) diff --git a/commands/imagine.py b/commands/imagine.py index 90b7930..462a771 100644 --- a/commands/imagine.py +++ b/commands/imagine.py @@ -1,6 +1,3 @@ -import asyncio -import functools - from nio.events.room_events import RoomMessageText from nio.rooms import MatrixRoom @@ -12,16 +9,7 @@ async def command_imagine(room: MatrixRoom, event: RoomMessageText, bot): bot.logger.log("Generating image...") try: - loop = asyncio.get_event_loop() - except Exception as e: - bot.logger.log(f"Error getting event loop: {e}", "error") - await bot.send_message( - room, "Something went wrong. Please try again.", True) - return - - try: - image_partial = functools.partial(bot.image_api.generate_image, prompt, user=room.room_id) - images, tokens_used = await loop.run_in_executor(None, image_partial) + images, tokens_used = await bot.image_api.generate_image(prompt, user=room.room_id) except Exception as e: bot.logger.log(f"Error generating image: {e}", "error") await bot.send_message(room, "Sorry, I couldn't generate an image. Please try again later.", True) From cdca5fb124b6ac7b3667fada91d038f81d5d3326 Mon Sep 17 00:00:00 2001 From: Justin Date: Fri, 19 May 2023 14:52:20 -0500 Subject: [PATCH 6/9] using openai async methods --- classes/openai.py | 57 +++++++++++++---------------------------------- 1 file changed, 16 insertions(+), 41 deletions(-) diff --git a/classes/openai.py b/classes/openai.py index c43aba5..9c143a0 100644 --- a/classes/openai.py +++ b/classes/openai.py @@ -39,22 +39,14 @@ class OpenAI: Returns: Tuple[str, int]: The response text and the number of tokens used. """ - try: - loop = asyncio.get_event_loop() - except Exception as e: - self.logger.log(f"Error getting event loop: {e}", "error") - return - self.logger.log(f"Generating response to {len(messages)} messages using {self.chat_model}...") - chat_partial = functools.partial( - openai.ChatCompletion.create, - model=self.chat_model, - messages=messages, - api_key=self.api_key, - user = user + response = await openai.ChatCompletion.acreate( + model=self.chat_model, + messages=messages, + api_key=self.api_key, + user = user ) - response = await loop.run_in_executor(None, chat_partial) result_text = response.choices[0].message['content'] tokens_used = response.usage["total_tokens"] @@ -75,12 +67,6 @@ class OpenAI: - If for any reason you are unable to classify the message (for example, if it infringes on your terms of service), the event_type is "error", and the prompt is a message explaining why you are unable to process the message. Only the event_types mentioned above are allowed, you must not respond in any other way.""" - try: - loop = asyncio.get_event_loop() - except Exception as e: - self.logger.log(f"Error getting event loop: {e}", "error") - return - messages = [ { "role": "system", @@ -94,14 +80,12 @@ Only the event_types mentioned above are allowed, you must not respond in any ot self.logger.log(f"Classifying message '{query}'...") - chat_partial = functools.partial( - openai.ChatCompletion.create, - model=self.chat_model, - messages=messages, - api_key=self.api_key, - user=user + response = await openai.ChatCompletion.acreate( + model=self.chat_model, + messages=messages, + api_key=self.api_key, + user = user ) - response = await loop.run_in_executor(None, chat_partial) try: result = json.loads(response.choices[0].message['content']) @@ -123,24 +107,15 @@ Only the event_types mentioned above are allowed, you must not respond in any ot Yields: bytes: The image data. """ - try: - loop = asyncio.get_event_loop() - except Exception as e: - self.logger.log(f"Error getting event loop: {e}", "error") - return - - self.logger.log(f"Generating image from prompt '{prompt}'...") - image_partial = functools.partial( - openai.Image.create, - prompt=prompt, - n=1, - api_key=self.api_key, - size="1024x1024", - user = user + response = await openai.Image.acreate( + prompt=prompt, + n=1, + api_key=self.api_key, + size="1024x1024", + user = user ) - response = await loop.run_in_executor(None, image_partial) images = [] From 6c97c0f61ddac7c938318d7e34c66f978e8e4be1 Mon Sep 17 00:00:00 2001 From: Justin Date: Fri, 19 May 2023 14:54:51 -0500 Subject: [PATCH 7/9] Cleanup unused imports --- classes/openai.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/classes/openai.py b/classes/openai.py index 9c143a0..059dab9 100644 --- a/classes/openai.py +++ b/classes/openai.py @@ -1,8 +1,6 @@ import openai import requests -import asyncio -import functools import json from .logging import Logger From b41a9ecd14ab7f58c587b49454072bb5a479a28f Mon Sep 17 00:00:00 2001 From: Justin Date: Fri, 19 May 2023 15:37:04 -0500 Subject: [PATCH 8/9] Adds retry logic for failed openai requests --- classes/openai.py | 70 +++++++++++++++++++++++++++++++++++------------ 1 file changed, 53 insertions(+), 17 deletions(-) diff --git a/classes/openai.py b/classes/openai.py index 059dab9..3423ba2 100644 --- a/classes/openai.py +++ b/classes/openai.py @@ -1,11 +1,13 @@ import openai import requests +import asyncio import json +from functools import partial from .logging import Logger -from typing import Dict, List, Tuple, Generator, Optional +from typing import Dict, List, Tuple, Generator, AsyncGenerator, Optional, Any class OpenAI: api_key: str @@ -28,6 +30,32 @@ class OpenAI: self.chat_model = chat_model or self.chat_model self.logger = logger or Logger() + async def _request_with_retries(self, request: partial, attempts: int = 5, retry_interval: int = 2) -> AsyncGenerator[Any | list | Dict, None]: + """Retry a request a set number of times if it fails. + + Args: + request (partial): The request to make with retries. + attempts (int, optional): The number of attempts to make. Defaults to 5. + retry_interval (int, optional): The interval in seconds between attempts. Defaults to 2 seconds. + + Returns: + AsyncGenerator[Any | list | Dict, None]: The OpenAI response for the request. + """ + # call the request function and return the response if it succeeds, else retry + current_attempt = 1 + while current_attempt <= attempts: + try: + response = await request() + return response + except Exception as e: + self.logger.log(f"Request failed: {e}", "error") + self.logger.log(f"Retrying in {retry_interval} seconds...") + await asyncio.sleep(retry_interval) + current_attempt += 1 + + # if all attempts failed, raise an exception + raise Exception("Request failed after all attempts.") + async def generate_chat_response(self, messages: List[Dict[str, str]], user: Optional[str] = None) -> Tuple[str, int]: """Generate a response to a chat message. @@ -39,12 +67,16 @@ class OpenAI: """ self.logger.log(f"Generating response to {len(messages)} messages using {self.chat_model}...") - response = await openai.ChatCompletion.acreate( - model=self.chat_model, - messages=messages, - api_key=self.api_key, - user = user + + chat_partial = partial( + openai.ChatCompletion.acreate, + model=self.chat_model, + messages=messages, + api_key=self.api_key, + user=user ) + response = await self._request_with_retries(chat_partial) + result_text = response.choices[0].message['content'] tokens_used = response.usage["total_tokens"] @@ -78,12 +110,14 @@ Only the event_types mentioned above are allowed, you must not respond in any ot self.logger.log(f"Classifying message '{query}'...") - response = await openai.ChatCompletion.acreate( - model=self.chat_model, - messages=messages, - api_key=self.api_key, - user = user + chat_partial = partial( + openai.ChatCompletion.acreate, + model=self.chat_model, + messages=messages, + api_key=self.api_key, + user=user ) + response = await self._request_with_retries(chat_partial) try: result = json.loads(response.choices[0].message['content']) @@ -107,13 +141,15 @@ Only the event_types mentioned above are allowed, you must not respond in any ot """ self.logger.log(f"Generating image from prompt '{prompt}'...") - response = await openai.Image.acreate( - prompt=prompt, - n=1, - api_key=self.api_key, - size="1024x1024", - user = user + image_partial = partial( + openai.Image.acreate, + prompt=prompt, + n=1, + api_key=self.api_key, + size="1024x1024", + user=user ) + response = await self._request_with_retries(image_partial) images = [] From 861b82061f2313564186a9a420200fbdc69f72e2 Mon Sep 17 00:00:00 2001 From: Justin Date: Sat, 20 May 2023 11:20:54 -0500 Subject: [PATCH 9/9] Use suggested migration fix --- migrations/__init__.py | 14 ++++---------- 1 file changed, 4 insertions(+), 10 deletions(-) diff --git a/migrations/__init__.py b/migrations/__init__.py index 9724944..1085ee2 100644 --- a/migrations/__init__.py +++ b/migrations/__init__.py @@ -24,7 +24,7 @@ def get_version(db: DuckDBPyConnection) -> int: try: return int(db.execute("SELECT MAX(id) FROM migrations").fetchone()[0]) except: - return 1 + return 0 def migrate(db: DuckDBPyConnection, from_version: Optional[int] = None, to_version: Optional[int] = None) -> None: """Migrate the database to a specific version. @@ -44,14 +44,8 @@ def migrate(db: DuckDBPyConnection, from_version: Optional[int] = None, to_versi if from_version > to_version: raise ValueError("Cannot migrate from a higher version to a lower version.") - # dont migrate if already at the target version - if from_version == to_version: - return from_version, to_version - - print(f"Migrating from version {from_version} to {to_version}...") - for version in range(from_version, to_version + 1): - if version in MIGRATIONS: - print(f"Running migration {version}...") - MIGRATIONS[version](db) + for version in range(from_version, to_version): + if version + 1 in MIGRATIONS: + MIGRATIONS[version + 1](db) return from_version, to_version