dotfiles/python/ebook_search/config.py

"""Configuration for the EPUB search app."""

from __future__ import annotations

from os import getenv
from typing import Annotated, Self

from pydantic import AliasChoices, Field, field_validator, model_validator
from pydantic_settings import BaseSettings, NoDecode, SettingsConfigDict


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",
        "Qwen3-Embedding-8B": "qwen3-embedding-8b",
        "Qwen/Qwen3-Embedding-0.6B": "qwen3-embedding-0.6b",
        "Qwen/Qwen3-Embedding-4B": "qwen3-embedding-4b",
        "Qwen/Qwen3-Embedding-8B": "qwen3-embedding-8b",
        "qwen3-embedding:0.6b": "qwen3-embedding-0.6b",
        "qwen3-embedding:4b": "qwen3-embedding-4b",
        "qwen3-embedding:8b": "qwen3-embedding-8b",
        "qwen3-embedding-0.6b": "qwen3-embedding-0.6b",
        "qwen3-embedding-4b": "qwen3-embedding-4b",
        "qwen3-embedding-8b": "qwen3-embedding-8b",
    }
    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 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 = True
    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
    min_retrieval_confidence: float = 0.0
    validate_citations_enabled: bool = True
    bm25_index_dir: str = ".ebook_search_bm25"
    bm25_refresh_delay_seconds: int = 60
    protected_phrase_max_candidates_per_book: int = 5000
    protected_phrase_llm_candidates_per_book: int = 500
    protected_phrase_extraction_workers: int = 16
    phrase_judge_book_workers: int = 20
    phrase_judge_phrase_workers: int = 100
    protected_phrase_confidence_threshold: float = 0.80
    phrase_matching_enabled: bool = True
    phrase_hit_boost: float = 0.25
    phrase_min_tokens: int = 2
    phrase_max_tokens: int = 5
    phrase_max_entity_tokens: int = 8
    phrase_raw_ngram_min_count: int = 2
    phrase_raw_count_score_threshold: int = 3
    phrase_raw_count_high_score_threshold: int = 10
    phrase_chapter_count_score_threshold: int = 2
    phrase_chapter_count_high_score_threshold: int = 5
    phrase_target_protected_per_book: int = 100
    phrase_default_allow_nested: bool = False
    phrase_default_suppress_children: bool = True

    @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)

    @model_validator(mode="after")
    def validate_runtime_consistency(self) -> Self:
        """Reject configurations that cannot serve the features they enable."""
        if not self.embedding_base_url.strip():
            msg = "embedding_base_url must be set"
            raise ValueError(msg)
        if self.answer_enabled and (not self.vllm_base_url.strip() or not self.chat_model.strip()):
            msg = "answer_enabled requires vllm_base_url and chat_model to be set"
            raise ValueError(msg)
        if self.rerank.enabled and not self.rerank.base_url.strip():
            msg = "rerank.enabled requires rerank.base_url to be set"
            raise ValueError(msg)
        return self


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()