converting to pydantic-settings

2026-06-14 21:29:45 -04:00
parent a5d7c3be4f
commit 68b3a38b81
10 changed files with 111 additions and 110 deletions
@@ -2,88 +2,15 @@

 from __future__ import annotations

-from dataclasses import dataclass
 from os import getenv
+from typing import Annotated
+
+from pydantic import AliasChoices, Field, field_validator
+from pydantic_settings import BaseSettings, NoDecode, SettingsConfigDict


-def getenv_bool(name: str, *, default: bool) -> bool:
-    """Read a boolean environment variable with a default fallback."""
-    value = getenv(name)
-    if value is None:
-        return default
-    return value.strip().lower() in {"1", "true", "yes", "on"}
-
-
-def getenv_int(name: str, *, default: int) -> int:
-    """Read an integer environment variable with a default fallback."""
-    value = getenv(name)
-    if value is None or not value.strip():
-        return default
-    return int(value)
-
-
-@dataclass(frozen=True)
-class RerankConfig:
-    """vLLM reranker settings."""
-
-    enabled: bool = False
-    base_url: str = "http://192.168.90.25:8001"
-    model: str = "qwen3-reranker-06b"
-    candidates: int = 24
-    timeout_seconds: float = 30.0
-
-
-@dataclass(frozen=True)
-class EbookSearchConfig:
-    """Runtime settings for EPUB search."""
-
-    rerank: RerankConfig
-    top_k: int = 12
-    library_paths: tuple[str, ...] = ()
-    vllm_base_url: str = "https://ollama.com/v1"
-    vllm_api_key: str = "not-needed"
-    chat_model: str = "deepseek-v4-flash"
-    answer_enabled: bool = True
-    embedding_base_url: str = "http://192.168.90.25:8000/v1"
-    embedding_api_key: str = "not-needed"
-    embedding_model: str = "qwen3-embedding-0.6b"
-    embedding_batch_size: int = 32
-    bm25_index_dir: str = ".ebook_search_bm25"
-    bm25_refresh_delay_seconds: int = 60
-
-
-def load_rerank_config() -> RerankConfig:
-    """Load reranker config from environment variables."""
-    return RerankConfig(
-        enabled=getenv_bool("EBOOK_SEARCH_RERANK_ENABLED", default=False),
-        base_url=getenv("EBOOK_SEARCH_RERANK_BASE_URL", "http://192.168.90.25:8001"),
-        model=getenv("EBOOK_SEARCH_RERANK_MODEL", "qwen3-reranker-06b"),
-        candidates=getenv_int("EBOOK_SEARCH_RERANK_CANDIDATES", default=24),
-        timeout_seconds=float(getenv_int("EBOOK_SEARCH_RERANK_TIMEOUT_SECONDS", default=30)),
-    )
-
-
-def load_config() -> EbookSearchConfig:
-    """Load EPUB search config from environment variables."""
-    return EbookSearchConfig(
-        rerank=load_rerank_config(),
-        top_k=getenv_int("EBOOK_SEARCH_TOP_K", default=12),
-        library_paths=library_paths_from_env(),
-        vllm_base_url=getenv("EBOOK_SEARCH_VLLM_BASE_URL", "https://ollama.com/v1"),
-        vllm_api_key=getenv("EBOOK_SEARCH_VLLM_API_KEY") or getenv("OLLAMA_API_KEY") or "not-needed",
-        chat_model=getenv("EBOOK_SEARCH_CHAT_MODEL", "deepseek-v4-flash"),
-        answer_enabled=getenv_bool("EBOOK_SEARCH_ANSWER_ENABLED", default=True),
-        embedding_base_url=getenv("EBOOK_SEARCH_EMBEDDING_BASE_URL", "http://192.168.90.25:8000/v1"),
-        embedding_api_key=getenv("EBOOK_SEARCH_EMBEDDING_API_KEY", "not-needed"),
-        embedding_model=normalize_embedding_model(),
-        embedding_batch_size=getenv_int("EBOOK_SEARCH_EMBEDDING_BATCH_SIZE", default=32),
-        bm25_index_dir=getenv("EBOOK_SEARCH_BM25_INDEX_DIR", ".ebook_search_bm25"),
-        bm25_refresh_delay_seconds=getenv_int("EBOOK_SEARCH_BM25_REFRESH_DELAY_SECONDS", default=60),
-    )
-
-
-def normalize_embedding_model(default: str = "qwen3-embedding-0.6b") -> str:
-    """Normalize supported embedding aliases to provider model names."""
+def normalize_embedding_alias(model: str) -> str:
+    """Normalize a supported embedding alias to its provider model name."""
    aliases = {
        "Qwen3-Embedding-0.6B": "qwen3-embedding-0.6b",
        "Qwen3-Embedding-4B": "qwen3-embedding-4b",
@@ -98,20 +25,86 @@ def normalize_embedding_model(default: str = "qwen3-embedding-0.6b") -> str:
        "qwen3-embedding-4b": "qwen3-embedding-4b",
        "qwen3-embedding-8b": "qwen3-embedding-8b",
    }
-
-    model = getenv("EBOOK_SEARCH_EMBEDDING_MODEL", default)
    standard_model = aliases.get(model)
-
    if standard_model is None:
        error = f"Embedding model {model} is not supported. Supported models are {aliases.keys()}"
        raise ValueError(error)
-
    return standard_model


-def library_paths_from_env() -> tuple[str, ...]:
-    """Read configured EPUB library paths from the environment."""
-    value = getenv("EBOOK_SEARCH_LIBRARY_PATHS")
-    if value is None:
-        return ()
-    return tuple(path for path in value.split(":") if path)
+def normalize_embedding_model(default: str = "qwen3-embedding-0.6b") -> str:
+    """Normalize the configured embedding alias to its provider model name."""
+    return normalize_embedding_alias(getenv("EBOOK_SEARCH_EMBEDDING_MODEL", default))
+
+
+class RerankConfig(BaseSettings):
+    """vLLM reranker settings."""
+
+    model_config = SettingsConfigDict(env_prefix="EBOOK_SEARCH_RERANK_", frozen=True, protected_namespaces=())
+
+    enabled: bool = False
+    base_url: str = "http://192.168.90.25:8001"
+    model: str = "qwen3-reranker-06b"
+    candidates: int = 24
+    timeout_seconds: float = 30.0
+    score_weight: float = 0.7
+    hybrid_weight: float = 0.3
+
+
+class EbookSearchConfig(BaseSettings):
+    """Runtime settings for EPUB search."""
+
+    model_config = SettingsConfigDict(
+        env_prefix="EBOOK_SEARCH_",
+        frozen=True,
+        populate_by_name=True,
+        protected_namespaces=(),
+    )
+
+    rerank: RerankConfig = Field(default_factory=RerankConfig)
+    top_k: int = 12
+    library_paths: Annotated[tuple[str, ...], NoDecode] = ()
+    chunk_tokens: int = 700
+    chunk_overlap: int = 100
+    vllm_base_url: str = "https://ollama.com/v1"
+    vllm_api_key: str = Field(
+        default="not-needed",
+        validation_alias=AliasChoices("EBOOK_SEARCH_VLLM_API_KEY", "OLLAMA_API_KEY"),
+    )
+    chat_model: str = "deepseek-v4-flash"
+    answer_enabled: bool = True
+    embedding_base_url: str = "http://192.168.90.25:8000/v1"
+    embedding_api_key: str = "not-needed"
+    embedding_model: str = "qwen3-embedding-0.6b"
+    embedding_batch_size: int = 32
+    embedding_timeout_seconds: float = 60.0
+    chat_timeout_seconds: float = 60.0
+    vector_candidate_multiplier: int = 4
+    bm25_candidate_limit: int = 120
+    rrf_rank_constant: int = 60
+    bm25_index_dir: str = ".ebook_search_bm25"
+    bm25_refresh_delay_seconds: int = 60
+
+    @field_validator("library_paths", mode="before")
+    @classmethod
+    def split_library_paths(cls, value: object) -> object:
+        """Split a colon-separated library path string into a tuple of paths."""
+        if isinstance(value, str):
+            return tuple(path for path in value.split(":") if path)
+        return value
+
+    @field_validator("embedding_model")
+    @classmethod
+    def normalize_embedding(cls, value: str) -> str:
+        """Normalize the configured embedding alias to its provider model name."""
+        return normalize_embedding_alias(value)
+
+
+def load_rerank_config() -> RerankConfig:
+    """Load reranker config from environment variables."""
+    return RerankConfig()
+
+
+def load_config() -> EbookSearchConfig:
+    """Load EPUB search config from environment variables."""
+    return EbookSearchConfig()