diff --git a/src/codegate/cli.py b/src/codegate/cli.py index ec32a89d..49b798c1 100644 --- a/src/codegate/cli.py +++ b/src/codegate/cli.py @@ -9,6 +9,7 @@ from codegate.codegate_logging import LogFormat, LogLevel, setup_logging from codegate.config import Config, ConfigurationError +from codegate.db.connection import init_db_sync from codegate.server import init_app @@ -164,6 +165,7 @@ def serve( }, ) + init_db_sync() app = init_app() import uvicorn diff --git a/src/codegate/db/connection.py b/src/codegate/db/connection.py index 34cb67e5..098c3007 100644 --- a/src/codegate/db/connection.py +++ b/src/codegate/db/connection.py @@ -35,7 +35,6 @@ def __init__(self, sqlite_path: Optional[str] = None): } self._async_db_engine = create_async_engine(**engine_dict) self._db_engine = create_engine(**engine_dict) - if not self.does_db_exist(): logger.info(f"Database does not exist at {self._db_path}. Creating..") asyncio.run(self.init_db()) @@ -153,7 +152,6 @@ async def record_output_stream( if output_chunks: # Record the output chunks output_str = json.dumps(output_chunks) - logger.info(f"Recorded chunks: {output_chunks}. Str: {output_str}") await self._record_output(prompt, output_str) async def record_output_non_stream( @@ -181,7 +179,12 @@ async def record_output_non_stream( def init_db_sync(): """DB will be initialized in the constructor in case it doesn't exist.""" - DbRecorder() + db = DbRecorder() + # Remove the DB file if exists for the moment to not cause issues at schema change. + # We can replace this in the future with migrations or something similar. + if db.does_db_exist(): + db._db_path.unlink() + asyncio.run(db.init_db()) if __name__ == "__main__":