src/explanation/model_loader.py

"""
model_loader.py
---------------
NVIDIA Build API Client for RazorShield SLM Explainer.
Operates 100% via API (Zero GPU / Zero Local Model Weight Overhead).
"""

from __future__ import annotations

import os
import logging
from typing import Any

HAS_SPACES = False
LOGGER = logging.getLogger("slm-model-loader")


class SLMModelLoader:
    """NVIDIA Build API model loader for zero-shot explanation generation."""

    def __init__(
        self,
        model_name: str | None = None,
        device: str | None = None,
        max_new_tokens: int | None = None,
        temperature: float | None = None,
    ):
        if os.getenv("NVIDIA_API_KEY") and os.getenv("NVIDIA_API_KEY").strip():
            if not model_name or "Qwen" in model_name:
                model_name = os.getenv("NVIDIA_MODEL", "openai/gpt-oss-20b")
        self.model_name = model_name or os.getenv("NVIDIA_MODEL", "openai/gpt-oss-20b")
        self.max_tokens = max_new_tokens or int(os.getenv("SLM_MAX_TOKENS", "4096"))
        self.temperature = temperature or float(os.getenv("SLM_TEMPERATURE", "1.0"))
        self._is_loaded = False

    @property
    def is_loaded(self) -> bool:
        """Returns True if NVIDIA_API_KEY environment variable is configured."""
        key = os.getenv("NVIDIA_API_KEY", "").strip()
        return bool(key) or self._is_loaded

    @is_loaded.setter
    def is_loaded(self, value: bool):
        self._is_loaded = value

    def load_model(self) -> bool:
        """Validates NVIDIA Build API configuration."""
        key = os.getenv("NVIDIA_API_KEY", "").strip()
        if key:
            LOGGER.info("NVIDIA_API_KEY detected. Active SLM: NVIDIA Build API (%s).", self.model_name)
            self._is_loaded = True
            return True

        LOGGER.warning("NVIDIA_API_KEY is not set. SLM will fall back to Deterministic Explainer (0 GPU).")
        self._is_loaded = False
        return False

    def generate(self, prompt: str) -> str:
        """Generates response using NVIDIA Build API (openai/gpt-oss-20b)."""
        nvidia_key = os.getenv("NVIDIA_API_KEY", "").strip()
        if not nvidia_key:
            raise RuntimeError("Model is not loaded. Please set NVIDIA_API_KEY environment variable.")

        try:
            from openai import OpenAI
            client = OpenAI(
                base_url="https://integrate.api.nvidia.com/v1",
                api_key=nvidia_key,
            )
            LOGGER.info("Executing SLM generation via NVIDIA Build API (%s)...", self.model_name)
            completion = client.chat.completions.create(
                model=self.model_name,
                messages=[{"role": "user", "content": prompt}],
                temperature=self.temperature,
                top_p=1,
                max_tokens=self.max_tokens,
                stream=False,
            )
            reasoning = getattr(completion.choices[0].message, "reasoning_content", None)
            if reasoning:
                LOGGER.info("NVIDIA Build API reasoning trace: %s", reasoning[:200])

            content = completion.choices[0].message.content
            if content and content.strip():
                return content.strip()
            raise RuntimeError("NVIDIA Build API returned empty response content.")
        except Exception as exc:
            LOGGER.error("NVIDIA Build API call failed: %s", exc)
            raise exc