config: async driver

Author: LongBaoCoder2

backend/qllm/core/config.py

@@ -2,8 +2,7 @@
 from enum import Enum
 from typing import Annotated, Any, List, Literal, Optional
 
-from pydantic import AnyUrl, BeforeValidator, HttpUrl, PostgresDsn
-from pydantic_core import MultiHostUrl
+from pydantic import AnyUrl, BeforeValidator, HttpUrl
 from pydantic_settings import BaseSettings
 
 
@@ -18,6 +17,15 @@ def parse_cors(v: Any) -> List[str] | str:
     raise ValueError(v)
 
 
+def replace_postgresql_asyncpg(url: str) -> str:
+    """
+    Replaces `postgresql://` with `postgresql+asyncpg://` in the URL.
+    """
+    if "postgresql://" in url:
+        return url.replace("postgresql://", "postgresql+asyncpg://")
+    return url
+
+
 class AppEnvironment(str, Enum):
     """
     Enum for app environments.
@@ -36,6 +44,7 @@ class ApplicationSetting(BaseSettings):
     PROJECT_NAME: str = "QLLM"
     API_PREFIX: str = "/api/v1"
     SECRET_KEY: str = secrets.token_urlsafe(32)
+    PORT: int = 10000
     ACCESS_TOKEN_EXPIRE_MINUTES: int = 11520
     UVICORN_WORKER_COUNT: Optional[int] = None
     BACKEND_CORS_ORIGINS: Annotated[
@@ -53,7 +62,7 @@ class ApplicationSetting(BaseSettings):
     ALGORITHM: str = "HS256"
 
     # Database config
-    DATABASE_URL: str
+    DATABASE_URL: Annotated[str, BeforeValidator(replace_postgresql_asyncpg)]
 
     # LLM Config
     LLM_TEMPERATURE: float = 0.5

backend/qllm/core/database.py

@@ -4,8 +4,6 @@
 
 DATABASE_URL = settings.DATABASE_URL
 
-if "postgresql://" in DATABASE_URL:
-    DATABASE_URL = DATABASE_URL.replace("postgresql://", "postgresql+asyncpg://")
 
 print("DATABASE_URL: ", DATABASE_URL)
 engine = create_async_engine(DATABASE_URL, echo=True)

backend/qllm/main.py

@@ -45,9 +45,6 @@ def __setup_logging(log_level: str):
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    import os
-    from pathlib import Path
-
     # first wait for DB to be connectable
     await check_database_connection()
 
@@ -145,7 +142,7 @@ def start():
     uvicorn.run(
         "qllm.main:app",
         host="0.0.0.0",
-        port=8000,
+        port=settings.PORT,
         reload=live_reload,
         workers=settings.UVICORN_WORKER_COUNT,
     )