build-small-hackathon/Off-Grid-Field-Repair-Logbook

1from __future__ import annotations2 3import json4import os5import time6from dataclasses import dataclass7from functools import lru_cache8from pathlib import Path9from typing import Any, Sequence10 11REPO_ROOT = Path(__file__).resolve().parents[1]12MODEL_REGISTRY_PATH = REPO_ROOT / "configs" / "model_registry.yaml"13DEFAULT_COMPONENTS = ("reasoning_llm", "vision_llm", "manual_parser")14 15_COMPONENT_ENV_VARS: dict[str, tuple[str, ...]] = {16    "reasoning_llm": ("P3_REASONING_MODEL_PATH", "P3_MODEL_PATH"),17    "vision_llm": ("P3_VISION_MODEL_PATH", "P3_MODEL_PATH"),18    "manual_parser": ("P3_MANUAL_PARSER_MODEL_PATH", "P3_MODEL_PATH"),19}20 21 22def _json_safe(value: Any) -> Any:23    if isinstance(value, Path):24        return str(value)25    if isinstance(value, dict):26        return {str(key): _json_safe(item) for key, item in value.items()}27    if isinstance(value, (list, tuple)):28        return [_json_safe(item) for item in value]29    return value30 31 32@lru_cache(maxsize=1)33def load_model_registry() -> dict[str, Any]:34    try:35        import yaml36    except Exception as exc:  # pragma: no cover - dependency issue is environment-specific37        raise RuntimeError("PyYAML is required to read configs/model_registry.yaml") from exc38 39    if not MODEL_REGISTRY_PATH.exists():40        raise FileNotFoundError(f"Missing model registry: {MODEL_REGISTRY_PATH}")41    loaded = yaml.safe_load(MODEL_REGISTRY_PATH.read_text(encoding="utf-8"))42    if not isinstance(loaded, dict):43        raise ValueError(f"Invalid model registry format: {MODEL_REGISTRY_PATH}")44    return loaded45 46 47@dataclass(frozen=True)48class ComponentSpec:49    component: str50    expected_model_id: str51    backend: str | None52    runtime: str | None53 54 55@dataclass(frozen=True)56class ModelAvailability:57    component: str58    expected_model_id: str59    backend: str | None60    runtime: str | None61    available: bool62    resolved_path: Path | None63    checked_paths: tuple[str, ...]64    problem: str65 66    def to_blocker(self) -> dict[str, Any]:67        return {68            "component": self.component,69            "expected_model_id": self.expected_model_id,70            "backend": self.backend,71            "runtime": self.runtime,72            "available": self.available,73            "resolved_path": str(self.resolved_path) if self.resolved_path else None,74            "checked_paths": list(self.checked_paths),75            "problem": self.problem,76        }77 78 79class ModelUnavailableError(RuntimeError):80    def __init__(self, availability: ModelAvailability, *, reason: str | None = None):81        self.availability = availability82        self.reason = reason or availability.problem83        super().__init__(self.reason)84 85    def to_blocker(self) -> dict[str, Any]:86        payload = self.availability.to_blocker()87        payload["problem"] = self.reason88        return payload89 90 91@lru_cache(maxsize=None)92def get_component_spec(component: str) -> ComponentSpec:93    registry = load_model_registry()94    projects = registry.get("projects", {})95    if not isinstance(projects, dict):96        raise ValueError("model_registry.yaml is missing the projects mapping")97    p3 = projects.get("p3", {})98    if not isinstance(p3, dict):99        raise ValueError("model_registry.yaml is missing the projects.p3 mapping")100    spec = p3.get(component)101    if not isinstance(spec, dict):102        raise KeyError(f"Unknown P3 component: {component}")103    model_id = spec.get("model_id")104    if not isinstance(model_id, str) or not model_id.strip():105        raise ValueError(f"Invalid model_id for {component}")106    backend = spec.get("backend")107    runtime = spec.get("runtime")108    return ComponentSpec(109        component=component,110        expected_model_id=model_id.strip(),111        backend=backend.strip() if isinstance(backend, str) and backend.strip() else None,112        runtime=runtime.strip() if isinstance(runtime, str) and runtime.strip() else None,113    )114 115 116def _candidate_roots() -> list[Path]:117    roots: list[Path] = []118    for env_var in ("P3_MODEL_CACHE_DIR", "MODEL_CACHE_DIR"):119        value = os.environ.get(env_var)120        if value:121            roots.append(Path(value).expanduser())122    roots.extend([Path("/opt/data/workspace/model-cache"), REPO_ROOT / "models"])123    unique: list[Path] = []124    seen: set[str] = set()125    for root in roots:126        key = str(root)127        if key in seen:128            continue129        seen.add(key)130        unique.append(root)131    return unique132 133 134def _candidate_model_paths(component: str, model_id: str) -> list[Path]:135    candidates: list[Path] = []136    env_vars = _COMPONENT_ENV_VARS.get(component, ())137    for env_var in env_vars:138        value = os.environ.get(env_var)139        if value:140            candidates.append(Path(value).expanduser())141    model_name = model_id.strip()142    model_tail = model_name.split("/")[-1]143    for root in _candidate_roots():144        candidates.extend(145            [146                root / model_name,147                root / model_tail,148                root / model_tail.replace("-", "_"),149                root / model_name.replace("/", "-"),150            ]151        )152    unique: list[Path] = []153    seen: set[str] = set()154    for candidate in candidates:155        key = str(candidate)156        if key in seen:157            continue158        seen.add(key)159        unique.append(candidate)160    return unique161 162 163@lru_cache(maxsize=None)164def check_component_availability(component: str) -> ModelAvailability:165    spec = get_component_spec(component)166    candidates = _candidate_model_paths(component, spec.expected_model_id)167    for candidate in candidates:168        if candidate.exists():169            return ModelAvailability(170                component=spec.component,171                expected_model_id=spec.expected_model_id,172                backend=spec.backend,173                runtime=spec.runtime,174                available=True,175                resolved_path=candidate,176                checked_paths=tuple(str(path) for path in candidates),177                problem="available",178            )179    problem = f"required model not mounted: {spec.expected_model_id}"180    return ModelAvailability(181        component=spec.component,182        expected_model_id=spec.expected_model_id,183        backend=spec.backend,184        runtime=spec.runtime,185        available=False,186        resolved_path=None,187        checked_paths=tuple(str(path) for path in candidates),188        problem=problem,189    )190 191 192def check_required_components(components: Sequence[str] = DEFAULT_COMPONENTS) -> list[ModelAvailability]:193    return [check_component_availability(component) for component in components]194 195 196def summarize_blockers(availabilities: Sequence[ModelAvailability]) -> list[dict[str, Any]]:197    return [item.to_blocker() for item in availabilities if not item.available]198 199 200def format_blocker_markdown(availabilities: Sequence[ModelAvailability], *, title: str = "Model-backed diagnosis is blocked") -> str:201    blockers = [item for item in availabilities if not item.available]202    if not blockers:203        return ""204    lines = [f"⚠️ **{title}**", "", "The app will not emit a rule-based answer path while the sponsor model requirements are unmet.", "", "Missing components:"]205    for blocker in blockers:206        lines.append(f"- `{blocker.component}` → `{blocker.expected_model_id}`")207    lines.append("")208    lines.append("Checked local paths:")209    for blocker in blockers:210        checked = blocker.checked_paths[:3]211        suffix = " …" if len(blocker.checked_paths) > 3 else ""212        lines.append(f"- `{blocker.component}`: {', '.join(f'`{path}`' for path in checked)}{suffix}")213    return "\n".join(lines)214 215 216def _load_llama_cpp():217    try:218        import llama_cpp219    except Exception as exc:  # pragma: no cover - import error is environment-specific220        raise RuntimeError(221            "llama_cpp is required for GGUF-backed inference; install llama-cpp-python in the runtime environment.") from exc222    return llama_cpp223 224 225def _load_transformers():226    try:227        import torch228        from transformers import AutoModelForCausalLM, AutoTokenizer229    except Exception as exc:  # pragma: no cover - import error is environment-specific230        raise RuntimeError(231            "transformers/torch are required for Hugging Face model-backed inference; install them in the runtime environment.") from exc232    return torch, AutoModelForCausalLM, AutoTokenizer233 234 235def _extract_generation_stats(response: Any, *, prompt_tokens: int | None = None, completion_tokens: int | None = None, duration_s: float | None = None, load_s: float | None = None, adapter_name: str | None = None, model_path: Path | None = None) -> dict[str, Any]:236    usage = response.get("usage") if isinstance(response, dict) else None237    if isinstance(usage, dict):238        prompt_tokens = prompt_tokens if prompt_tokens is not None else usage.get("prompt_tokens")239        completion_tokens = completion_tokens if completion_tokens is not None else usage.get("completion_tokens")240        total_tokens = usage.get("total_tokens")241    else:242        total_tokens = None243    stats = {244        "prompt_tokens": prompt_tokens,245        "completion_tokens": completion_tokens,246        "total_tokens": total_tokens,247        "duration_s": duration_s,248        "load_s": load_s,249        "adapter_name": adapter_name,250        "model_path": str(model_path) if model_path else None,251    }252    return {key: value for key, value in stats.items() if value is not None}253 254 255def generate_text(component: str, prompt: str, *, max_tokens: int = 384, temperature: float = 0.2, top_p: float = 0.9, seed: int = 13) -> tuple[str, dict[str, Any]]:256    availability = check_component_availability(component)257    if not availability.available or availability.resolved_path is None:258        raise ModelUnavailableError(availability)259 260    model_path = availability.resolved_path261    if model_path.suffix.lower() == ".gguf":262        llama_cpp = _load_llama_cpp()263        started = time.perf_counter()264        llm = llama_cpp.Llama(265            model_path=str(model_path),266            n_ctx=4096,267            seed=seed,268            verbose=False,269        )270        loaded = time.perf_counter()271        kwargs = {272            "max_tokens": max_tokens,273            "temperature": temperature,274            "top_p": top_p,275            "seed": seed,276        }277        try:278            response = llm.create_chat_completion(279                messages=[{"role": "user", "content": prompt}],280                **kwargs,281            )282            choice = response["choices"][0]283            message = choice.get("message") or {}284            text = message.get("content") or ""285        except Exception:286            response = llm(f"{prompt}\n", echo=False, **kwargs)287            choice = response["choices"][0]288            text = choice.get("text") or ""289        if not text.strip():290            raise RuntimeError(f"{component} returned empty text from GGUF model {model_path}")291        stats = _extract_generation_stats(292            response,293            duration_s=round(time.perf_counter() - started, 3),294            load_s=round(loaded - started, 3),295            adapter_name="llama_cpp",296            model_path=model_path,297        )298        return text.strip(), {299            "component": component,300            "model_name": availability.expected_model_id,301            "model_id": availability.expected_model_id,302            "expected_model_id": availability.expected_model_id,303            "adapter_name": "llama_cpp",304            "backend": availability.backend or "llama_cpp",305            "resolved_model_path": str(model_path),306            "generation_stats": stats,307        }308 309    torch, AutoModelForCausalLM, AutoTokenizer = _load_transformers()310    started = time.perf_counter()311    tokenizer = AutoTokenizer.from_pretrained(str(model_path), trust_remote_code=True)312    model = AutoModelForCausalLM.from_pretrained(313        str(model_path),314        trust_remote_code=True,315        torch_dtype="auto",316        device_map="auto",317    )318    loaded = time.perf_counter()319    if hasattr(tokenizer, "apply_chat_template"):320        prompt_text = tokenizer.apply_chat_template(321            [{"role": "user", "content": prompt}],322            tokenize=False,323            add_generation_prompt=True,324        )325    else:326        prompt_text = prompt327    inputs = tokenizer(prompt_text, return_tensors="pt")328    try:329        inputs = {key: value.to(model.device) for key, value in inputs.items()}330    except Exception:331        pass332    with torch.no_grad():333        output_ids = model.generate(334            **inputs,335            max_new_tokens=max_tokens,336            temperature=temperature,337            top_p=top_p,338            do_sample=temperature > 0,339            pad_token_id=tokenizer.eos_token_id,340        )341    prompt_len = int(inputs["input_ids"].shape[-1])342    completion_ids = output_ids[0][prompt_len:]343    text = tokenizer.decode(completion_ids, skip_special_tokens=True).strip()344    if not text:345        raise RuntimeError(f"{component} returned empty text from transformers model {model_path}")346    stats = _extract_generation_stats(347        {},348        prompt_tokens=prompt_len,349        completion_tokens=int(completion_ids.shape[-1]),350        duration_s=round(time.perf_counter() - started, 3),351        load_s=round(loaded - started, 3),352        adapter_name="transformers",353        model_path=model_path,354    )355    return text, {356        "component": component,357        "model_name": availability.expected_model_id,358        "model_id": availability.expected_model_id,359        "expected_model_id": availability.expected_model_id,360        "adapter_name": "transformers",361        "backend": availability.backend or "transformers",362        "resolved_model_path": str(model_path),363        "generation_stats": stats,364    }365 366 367def blocker_payload(availability: Sequence[ModelAvailability], *, status: str = "blocked") -> dict[str, Any]:368    blockers = summarize_blockers(availability)369    return {370        "status": status,371        "blocked_by": blockers,372    }373 374 375def stringify_blocked_response(availability: Sequence[ModelAvailability], *, title: str = "Model-backed diagnosis is blocked") -> str:376    body = format_blocker_markdown(availability, title=title)377    if body:378        body += "\n\nNo deterministic fallback will be used until the required model assets are mounted."379    return body380