55 lines
1.5 KiB
Python
55 lines
1.5 KiB
Python
import os
|
|
import argparse
|
|
from pydantic_settings import BaseSettings, PydanticBaseSettingsSource, YamlConfigSettingsSource
|
|
|
|
|
|
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()
|
|
|
|
|
|
_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
|
|
bucket: str
|
|
index_name: str
|
|
deployed_index_id: str
|
|
endpoint_name: str
|
|
endpoint_domain: str
|
|
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,
|
|
)
|