src/generation/llm_backend.py

"""
src/generation/llm_backend.py
Backends LLM disponibles :
  - LiteRTLMBackend     : Gemma via mediapipe (local, CPU)
  - TransformersBackend : Qwen3-0.6B via HuggingFace (local)
  - ClaudeHaikuBackend  : Claude via Anthropic API
  - GroqBackend         : Llama-3.1 via Groq API (rapide, recommandé HF)
"""

from pathlib import Path


# =============================================================================
# BASE CLASS
# =============================================================================
class LLMBackend:
    name: str = "abstract"

    def generate(self, prompt: str, max_new_tokens: int = 512) -> str:
        raise NotImplementedError


# =============================================================================
# 1. LiteRT-LM (Gemma local)
# =============================================================================
class LiteRTLMBackend(LLMBackend):
    name = "liteRT-LM (gemma4-E2B-it)"

    def __init__(self, model_path: str, max_tokens: int = 2048,
                 temperature: float = 0.0, top_k: int = 40):

        from mediapipe.tasks.python.genai import llm_inference

        p = Path(model_path)
        if not p.exists():
            raise FileNotFoundError(f"Model not found: {model_path}")

        options = llm_inference.LlmInferenceOptions(
            model_path=str(p),
            max_tokens=max_tokens,
            top_k=top_k,
            temperature=temperature,
            random_seed=42,
        )

        self.llm = llm_inference.LlmInference.create_from_options(options)

    def generate(self, prompt: str, max_new_tokens: int = 512) -> str:
        return self.llm.generate_response(prompt)


# =============================================================================
# 2. Transformers backend (Qwen local)
# =============================================================================
class TransformersBackend(LLMBackend):
    name = "transformers (Qwen3-0.6B)"

    def __init__(self, model_name: str = "Qwen/Qwen3-0.6B"):
        import torch
        from transformers import AutoTokenizer, AutoModelForCausalLM

        self.device = "cuda" if torch.cuda.is_available() else "cpu"

        self.tok = AutoTokenizer.from_pretrained(model_name)

        dtype = torch.bfloat16 if self.device == "cuda" else torch.float32

        self.model = AutoModelForCausalLM.from_pretrained(
            model_name,
            torch_dtype=dtype,
            device_map=self.device,
        )

        self.model.eval()

    def generate(self, prompt: str, max_new_tokens: int = 512) -> str:
        import torch

        enc = self.tok(prompt, return_tensors="pt").to(self.device)

        with torch.no_grad():
            out = self.model.generate(
                **enc,
                max_new_tokens=max_new_tokens,
                do_sample=False,
                repetition_penalty=1.3,
                no_repeat_ngram_size=4,
            )

        return self.tok.decode(
            out[0][enc["input_ids"].shape[-1]:],
            skip_special_tokens=True
        )


# =============================================================================
# 3. Claude Haiku backend (Anthropic API)
# =============================================================================
class ClaudeHaikuBackend(LLMBackend):
    name = "claude-haiku (api)"

    def __init__(self, model: str = "claude-haiku-4-5"):
        import os

        # 🔥 IMPORTANT: NO .env LOADING ON HF
        self.api_key = os.getenv("ANTHROPIC_API_KEY", "")

        if not self.api_key:
            raise ValueError(
                "ANTHROPIC_API_KEY non trouvée (mettre dans Hugging Face Secrets)"
            )

        self.model = model
        print(f"✅ Claude backend prêt ({model})")

    def generate(self, prompt: str, max_new_tokens: int = 512) -> str:
        import urllib.request, json, urllib.error

        if len(prompt) > 6000:
            prompt = prompt[:6000] + "\n\n[contexte tronqué]"

        payload = json.dumps({
            "model": self.model,
            "max_tokens": min(max_new_tokens, 1024),
            "messages": [
                {"role": "user", "content": prompt}
            ],
        }).encode()

        req = urllib.request.Request(
            "https://api.anthropic.com/v1/messages",
            data=payload,
            headers={
                "Content-Type": "application/json",
                "x-api-key": self.api_key,
                "anthropic-version": "2023-06-01",
            },
        )

        try:
            with urllib.request.urlopen(req, timeout=60) as resp:
                data = json.loads(resp.read())
            return data["content"][0]["text"]

        except urllib.error.HTTPError as e:
            raise RuntimeError(f"Anthropic API error {e.code}: {e.read().decode()}")


# =============================================================================
# 4. GROQ backend (RECOMMANDÉ POUR HF)
# =============================================================================
class GroqBackend(LLMBackend):
    name = "groq (llama-3.1-8b)"

    def __init__(self, model: str = "llama-3.1-8b-instant"):
        import os

        # 🔥 SAFE HF WAY
        self.api_key = os.getenv("GROQ_API_KEY", "")

        if not self.api_key:
            raise ValueError(
                "GROQ_API_KEY non trouvée (Hugging Face Secrets requis)"
            )

        self.model = model
        print(f"✅ Groq backend prêt ({model})")

    def generate(self, prompt: str, max_new_tokens: int = 512) -> str:
        from groq import Groq

        if len(prompt) > 6000:
            prompt = prompt[:6000] + "\n\n[contexte tronqué]"

        client = Groq(api_key=self.api_key)

        response = client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "user", "content": prompt}
            ],
            max_tokens=min(max_new_tokens, 1024),
            temperature=0.1,
        )

        return response.choices[0].message.content


# =============================================================================
# 5. FACTORY (choix automatique backend)
# =============================================================================
def make_llm(
    model_path: str = "data/gemma4-e2b-it.task",
    fallback_hf: str = "Qwen/Qwen3-0.6B",
    use_claude: bool = False
) -> LLMBackend:

    # Claude (optionnel)
    if use_claude:
        try:
            return ClaudeHaikuBackend()
        except Exception as e:
            print(f"Claude failed: {e}")

    # LiteRT local model
    if Path(model_path).exists():
        try:
            return LiteRTLMBackend(model_path)
        except Exception as e:
            print(f"LiteRT failed: {e}")

    # default fallback
    return TransformersBackend(fallback_hf)