127 lines
4.7 KiB
Python
127 lines
4.7 KiB
Python
"""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
|
|
|
|
@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()
|