From f36c51e8586895033c06f2f4677e8b13a4249254 Mon Sep 17 00:00:00 2001
From: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com>
Date: Tue, 27 May 2025 18:18:58 -0400
Subject: [PATCH] chore: Code Rabbit Nitpicks

---
 src/agent.py         | 10 +++++++++-
 src/config.py        | 20 +++++++++++---------
 src/deriver/queue.py |  2 +-
 src/main.py          |  2 +-
 tests/conftest.py    |  2 +-
 5 files changed, 23 insertions(+), 13 deletions(-)

diff --git a/src/agent.py b/src/agent.py
index d82fc7f9..7b00c25e 100644
--- a/src/agent.py
+++ b/src/agent.py
@@ -365,9 +365,17 @@ async def generate_semantic_queries(query: str) -> list[str]:
     logger.debug("Calling LLM for query generation")
     llm_start = asyncio.get_event_loop().time()
 
+    try:
+        provider = ModelProvider(settings.LLM.QUERY_GENERATION_PROVIDER)
+    except ValueError as e:
+        logger.error(
+            f"Invalid query-generation provider '{settings.LLM.QUERY_GENERATION_PROVIDER}': {e}"
+        )
+        raise
+
     # Create a new model client
     client = ModelClient(
-        provider=ModelProvider(settings.LLM.QUERY_GENERATION_PROVIDER),
+        provider=provider,
         model=settings.LLM.QUERY_GENERATION_MODEL,
     )
 
diff --git a/src/config.py b/src/config.py
index 07482177..451d1845 100644
--- a/src/config.py
+++ b/src/config.py
@@ -1,6 +1,6 @@
 import logging
 from pathlib import Path
-from typing import Annotated, Any, Optional
+from typing import Annotated, Any, ClassVar, Optional
 
 import tomllib
 from dotenv import load_dotenv
@@ -15,6 +15,8 @@ from pydantic_settings import (
 # Make sure this is called before AppSettings is instantiated if you rely on .env for AppSettings construction.
 load_dotenv()
 
+logger = logging.getLogger(__name__)
+
 
 def load_toml_config(config_path: str = "config.toml") -> dict[str, Any]:
     """Load configuration from TOML file if it exists."""
@@ -24,7 +26,7 @@ def load_toml_config(config_path: str = "config.toml") -> dict[str, Any]:
             with open(config_file, "rb") as f:
                 return tomllib.load(f)
         except (tomllib.TOMLDecodeError, OSError) as exc:
-            logging.warning("Failed to load %s: %s", config_path, exc)
+            logger.warning("Failed to load %s: %s", config_path, exc)
             return {}
     return {}
 
@@ -36,10 +38,10 @@ TOML_CONFIG = load_toml_config()
 class TomlConfigSettingsSource(PydanticBaseSettingsSource):
     """Custom settings source for loading from TOML file."""
 
-    def __init__(self, settings_cls: type[BaseSettings]):
+    def __init__(self, settings_cls: type[BaseSettings]) -> None:
         super().__init__(settings_cls)
 
-    SECTION_MAP: dict[str, str] = {
+    SECTION_MAP: ClassVar[dict[str, str]] = {
         "DB": "db",
         "AUTH": "auth",
         "SENTRY": "sentry",
@@ -90,11 +92,11 @@ class TomlSettings(BaseSettings):
     @classmethod
     def settings_customise_sources(
         cls,
-        settings_cls,
-        init_settings,
-        env_settings,
-        dotenv_settings,
-        file_secret_settings,
+        settings_cls: type[BaseSettings],
+        init_settings: PydanticBaseSettingsSource,
+        env_settings: PydanticBaseSettingsSource,
+        dotenv_settings: PydanticBaseSettingsSource,
+        file_secret_settings: PydanticBaseSettingsSource,
     ):
         # Return sources in priority order (first is lowest priority)
         return (
diff --git a/src/deriver/queue.py b/src/deriver/queue.py
index b133aa85..8e1a9f00 100644
--- a/src/deriver/queue.py
+++ b/src/deriver/queue.py
@@ -195,7 +195,7 @@ class QueueManager:
                                 settings.DERIVER.POLLING_SLEEP_INTERVAL_SECONDS
                             )
                     except Exception as e:
-                        logger.exception("Error in polling loop: %s", str(e))
+                        logger.exception("Error in polling loop")
                         if settings.SENTRY.ENABLED:
                             sentry_sdk.capture_exception(e)
                         # Note: rollback is handled by tracked_db dependency
diff --git a/src/main.py b/src/main.py
index 44c9cf0d..74914b94 100644
--- a/src/main.py
+++ b/src/main.py
@@ -27,7 +27,7 @@ from src.routers import (
 from src.security import create_admin_jwt
 
 
-def get_log_level():
+def get_log_level() -> int:
     """
     Convert log level string from settings to logging module constant.
 
diff --git a/tests/conftest.py b/tests/conftest.py
index eeb58956..78a2b8b8 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -56,7 +56,7 @@ DEFAULT_DB_URL = str(CONNECTION_URI.set(database="postgres"))
 
 # Test API authorization
 USE_AUTH = settings.AUTH.USE_AUTH
-AUTH_JWT_SECRET = settings.AUTH.JWT_SECRET
+AUTH_JWT_SECRET = settings.AUTH.JWT_SECRET or "test-secret"
 
 
 def create_test_database(db_url):