model_client.py

"""Model client for Analog Town using Hugging Face Inference API."""

import io
import json
import logging
import os
import re
from dotenv import load_dotenv
from huggingface_hub import InferenceClient
from prompts import REPAIR_PROMPT

logger = logging.getLogger(__name__)

TTS_MODELS = [
    "microsoft/speecht5_tts",
    "facebook/mms-tts-eng",
]


def synthesize_speech(text: str, token: str | None = None) -> bytes | None:
    """Call HF Inference TTS and return raw WAV bytes, or None on failure."""
    _token = token or os.getenv("HF_TOKEN")
    if not _token:
        return None
    client = InferenceClient(token=_token)
    for model in TTS_MODELS:
        try:
            result = client.text_to_speech(text, model=model)
            if hasattr(result, "read"):
                return result.read()
            if isinstance(result, (bytes, bytearray)):
                return bytes(result)
        except Exception as exc:
            logger.warning("TTS model %s failed: %s", model, exc)
            continue
    return None

load_dotenv()

# Model fallback chain (ensuring all models are under 32B parameters for hackathon rules)
MODELS = [
    "Qwen/Qwen2.5-7B-Instruct",
    "Qwen/Qwen2.5-14B-Instruct",
]


class ModelClient:
    """Wrapper for HF Inference API with JSON validation and retry logic."""

    def __init__(
        self,
        model_id: str | None = None,
        token: str | None = None,
        temperature: float = 0.3,
        top_p: float = 0.8,
        max_tokens: int = 900,
    ):
        self.token = token or os.getenv("HF_TOKEN")
        if not self.token:
            raise ValueError("HF_TOKEN not found. Set it in .env or pass it directly.")
        self.model_id = model_id or MODELS[0]
        self.temperature = temperature
        self.top_p = top_p
        self.max_tokens = max_tokens
        self.client = InferenceClient(token=self.token)

    def generate(self, system_prompt: str, user_prompt: str) -> str:
        """Generate text from the model. Returns raw string response."""
        # Try each model in fallback chain
        models_to_try = [self.model_id] + [m for m in MODELS if m != self.model_id]
        last_error = None
        
        for model in models_to_try:
            try:
                response = self.client.chat_completion(
                    model=model,
                    messages=[
                        {"role": "system", "content": system_prompt},
                        {"role": "user", "content": user_prompt},
                    ],
                    max_tokens=self.max_tokens,
                    temperature=self.temperature,
                    top_p=self.top_p,
                )
                return response.choices[0].message.content
            except Exception as e:
                last_error = e
                continue
        
        raise RuntimeError(f"All models failed. Last error: {last_error}")

    def _extract_json(self, text: str) -> dict:
        """Extract JSON from model output, handling markdown code blocks."""
        # Try direct parse first
        text = text.strip()
        try:
            return json.loads(text)
        except json.JSONDecodeError:
            pass
        
        # Try extracting from markdown code block
        patterns = [
            r'```json\s*\n(.*?)\n\s*```',
            r'```\s*\n(.*?)\n\s*```',
            r'\{[\s\S]*\}',
        ]
        for pattern in patterns:
            match = re.search(pattern, text, re.DOTALL)
            if match:
                try:
                    candidate = match.group(1) if '```' in pattern else match.group(0)
                    return json.loads(candidate)
                except (json.JSONDecodeError, IndexError):
                    continue
        
        raise json.JSONDecodeError("No valid JSON found in output", text, 0)

    def generate_json(self, system_prompt: str, user_prompt: str) -> dict:
        """Generate and parse JSON from model. Includes retry with repair prompt."""
        # First attempt
        raw_output = self.generate(system_prompt, user_prompt)
        
        try:
            return self._extract_json(raw_output)
        except json.JSONDecodeError:
            pass
        
        # Repair attempt
        repair_prompt = REPAIR_PROMPT.format(
            bad_output=raw_output,
            schema="See the required JSON structure in the original prompt."
        )
        
        try:
            repaired_output = self.generate(system_prompt, repair_prompt)
            return self._extract_json(repaired_output)
        except (json.JSONDecodeError, Exception) as e:
            raise RuntimeError(
                f"JSON generation failed after repair attempt. "
                f"Original output: {raw_output[:200]}... "
                f"Error: {e}"
            )