Add local config file support

main.py

@@ -5,6 +5,7 @@ import argparse
 import asyncio
 import io
 import logging
+import os
 from collections.abc import AsyncIterator, Sequence
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
@@ -16,7 +17,7 @@ from gcloud.aio.storage import Storage
 from google import genai
 from google.genai import types as genai_types
 from mcp.server.fastmcp import Context, FastMCP
-from pydantic_settings import BaseSettings
+from pydantic_settings import BaseSettings, PydanticBaseSettingsSource, YamlConfigSettingsSource
 
 logger = logging.getLogger(__name__)
 
@@ -282,10 +283,29 @@ class GoogleCloudVectorSearch:
 # ---------------------------------------------------------------------------
 
 
-class Settings(BaseSettings):
-    """Server configuration populated from environment variables."""
+def _parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--transport",
+        choices=["stdio", "sse"],
+        default="stdio",
+    )
+    parser.add_argument("--host", default="0.0.0.0")
+    parser.add_argument("--port", type=int, default=8080)
+    parser.add_argument(
+        "--config",
+        default=os.environ.get("CONFIG_FILE", "config.yaml"),
+    )
+    return parser.parse_args()
 
-    model_config = {"env_file": ".env"}
+
+_args = _parse_args()
+
+
+class Settings(BaseSettings):
+    """Server configuration populated from env vars and a YAML config file."""
+
+    model_config = {"env_file": ".env", "yaml_file": _args.config}
 
     project_id: str
     location: str
@@ -297,6 +317,23 @@ class Settings(BaseSettings):
     embedding_model: str = "gemini-embedding-001"
     search_limit: int = 10
 
+    @classmethod
+    def settings_customise_sources(
+        cls,
+        settings_cls: type[BaseSettings],
+        init_settings: PydanticBaseSettingsSource,
+        env_settings: PydanticBaseSettingsSource,
+        dotenv_settings: PydanticBaseSettingsSource,
+        file_secret_settings: PydanticBaseSettingsSource,
+    ) -> tuple[PydanticBaseSettingsSource, ...]:
+        return (
+            init_settings,
+            env_settings,
+            dotenv_settings,
+            YamlConfigSettingsSource(settings_cls),
+            file_secret_settings,
+        )
+
 
 @dataclass
 class AppContext:
@@ -334,19 +371,6 @@ async def lifespan(_server: FastMCP) -> AsyncIterator[AppContext]:
     )
 
 
-def _parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "--transport",
-        choices=["stdio", "sse"],
-        default="stdio",
-    )
-    parser.add_argument("--host", default="0.0.0.0")
-    parser.add_argument("--port", type=int, default=8080)
-    return parser.parse_args()
-
-
-_args = _parse_args()
 cfg = Settings.model_validate({})
 
 mcp = FastMCP(