config.py

from __future__ import annotations

import os
from dataclasses import dataclass, field

import i18n

APP_LANG = os.environ.get("APP_LANG", "es")
PACK = i18n.get(APP_LANG)

ROOT = os.path.dirname(os.path.abspath(__file__))
MODELS_DIR = os.path.join(ROOT, "models")
DATA_DIR = os.path.join(ROOT, "data")


def _physical_cores() -> int:
    """Count physical CPU cores.

    llama.cpp is fastest at the physical-core count; logical cores (SMT) hurt
    (on a 12-core/24-thread Ryzen, 24 threads doubled generation time). Falls
    back to the logical count where /proc/cpuinfo lacks topology (e.g. phones).
    """
    try:
        cores, phys, core = set(), None, None
        with open("/proc/cpuinfo") as fh:
            for line in fh:
                if line.startswith("physical id"):
                    phys = line.split(":")[1].strip()
                elif line.startswith("core id"):
                    core = line.split(":")[1].strip()
                elif not line.strip() and phys is not None and core is not None:
                    cores.add((phys, core))
                    phys = core = None
        if cores:
            return len(cores)
    except OSError:
        pass
    return os.cpu_count() or 8


@dataclass
class LLMBackend:
    name: str
    gguf_path: str
    # None uses the GGUF's embedded chat template, which Qwen3.5 needs: its
    # template honors "/no_think" while the generic "chatml" handler does not,
    # so the model thinks and a turn balloons from ~7s to ~80s.
    chat_format: str | None = None
    n_ctx: int = 8192
    supports_tools: bool = True
    lora_path: str | None = None
    # Hub repos the startup downloader pulls from when a file is missing, so a
    # Space (which ships code only) fetches its models at boot.
    hf_repo: str | None = None
    lora_repo: str | None = None


LLM_BACKENDS: dict[str, LLMBackend] = {
    # Strong Spanish, reliable tool use, big context.
    "qwen3.5-4b": LLMBackend(
        name="Qwen3.5-4B",
        gguf_path="Qwen3.5-4B-Q4_K_M.gguf",
        hf_repo="unsloth/Qwen3.5-4B-GGUF",
    ),
    # Smaller and faster.
    "qwen3.5-2b": LLMBackend(
        name="Qwen3.5-2B",
        gguf_path="Qwen3.5-2B-Q4_K_M.gguf",
        hf_repo="unsloth/Qwen3.5-2B-GGUF",
    ),
    # A/B challenger: Apache 2.0, weaker tool use.
    "gemma4-e4b": LLMBackend(
        name="Gemma-4-E4B-it",
        gguf_path="gemma-4-E4B-it-Q4_K_M.gguf",
        hf_repo="unsloth/gemma-4-E4B-it-GGUF",
    ),
    # Base 2B plus the amigo persona LoRA; the default on a Space.
    "amigo-2b": LLMBackend(
        name="Qwen3.5-2B-amigo (LoRA)",
        gguf_path="Qwen3.5-2B-Q4_K_M.gguf",
        hf_repo="unsloth/Qwen3.5-2B-GGUF",
        lora_path="amigo-lora-Q8_0.gguf",
        lora_repo="pebeto/amigo-lora",
    ),
}

# On a Space (SPACE_ID set) default to the small amigo-2b, the only model fast
# enough on a CPU Space; locally default to the 4B for the best profile use.
_DEFAULT_MODEL = "amigo-2b" if os.environ.get("SPACE_ID") else "qwen3.5-4b"
MODEL_KEY = os.environ.get("MODEL_KEY", _DEFAULT_MODEL)


@dataclass
class Config:
    lang: str = APP_LANG
    pack: dict = field(default_factory=lambda: PACK)

    llm: LLMBackend = field(default_factory=lambda: LLM_BACKENDS[MODEL_KEY])
    n_threads: int = int(os.environ.get("N_THREADS", _physical_cores()))
    max_tokens: int = 220
    temperature: float = 0.6
    # Lower temperature on web answers, so the model copies names and figures
    # from the results instead of confabulating from memory.
    search_temperature: float = float(os.environ.get("SEARCH_TEMPERATURE", "0.3"))

    # medium is markedly more accurate on Peruvian Spanish than small, at ~1.5GB
    # and slower on CPU. Drop to WHISPER_SIZE=small/base if a Space feels slow.
    whisper_size: str = os.environ.get("WHISPER_SIZE", "medium")
    whisper_compute: str = "int8"
    language: str = PACK["whisper"]

    piper_voice: str = os.environ.get("PIPER_VOICE", PACK["voice"])
    # Piper's stock values read rushed and flat. A slower pace (length_scale > 1)
    # and a touch more duration variation (noise_w) sound warmer and suit an
    # older listener, at negligible latency cost.
    tts_length_scale: float = float(os.environ.get("TTS_LENGTH_SCALE", "1.2"))
    tts_noise_scale: float = float(os.environ.get("TTS_NOISE_SCALE", "0.667"))
    tts_noise_w: float = float(os.environ.get("TTS_NOISE_W", "0.9"))

    # Granite R2 multilingual: Apache 2.0, 384-dim (compatible with the old
    # e5-small store) and stronger multilingual retrieval at the same ~97M size.
    embed_model: str = "ibm-granite/granite-embedding-97m-multilingual-r2"
    chroma_dir: str = os.path.join(DATA_DIR, "chroma")
    profile_path: str = os.path.join(DATA_DIR, "profile.yaml")
    rag_top_k: int = 3

    search_enabled: bool = os.environ.get("SEARCH", "1") == "1"
    # More than a few results pulls in conflicting stories.
    search_max_results: int = 3

    def model_path(self) -> str:
        """Absolute path to the active model's GGUF."""
        return os.path.join(MODELS_DIR, self.llm.gguf_path)

    def lora_path(self) -> str | None:
        """Absolute path to the active LoRA adapter, or None."""
        if not self.llm.lora_path:
            return None
        return os.path.join(MODELS_DIR, self.llm.lora_path)

    def piper_path(self) -> str:
        """Absolute path to the active Piper voice."""
        return os.path.join(MODELS_DIR, self.piper_voice)


CONFIG = Config()