model_client.py

"""Text inference client (HF Inference API, chat_completion) with graceful fallbacks."""
import json
import re
from typing import Optional, Dict, Any

try:
    from huggingface_hub import InferenceClient
except ImportError:
    InferenceClient = None

import config


class TextClient:
    """Wrapper for text inference via the Hugging Face Inference API."""

    def __init__(self, hf_token: Optional[str] = None):
        self.hf_token = hf_token or config.HF_TOKEN
        self.model = config.MODEL_NAME_TEXT
        self.client = None
        self.available = False

        if not InferenceClient:
            print("[warn]huggingface_hub not installed; text generation disabled")
            return
        if not self.hf_token:
            print("[warn]HF_TOKEN not set; text generation disabled (using fallbacks)")
            return

        try:
            self.client = InferenceClient(token=self.hf_token, timeout=config.TIMEOUT_INFERENCE)
            self.available = True
        except Exception as e:
            print(f"[warn]Text client initialization failed: {e}")
            self.available = False

    def infer(
        self,
        system_prompt: str,
        user_prompt: str,
        temperature: float = 0.65,
        max_tokens: int = 512,
    ) -> Optional[str]:
        """Generate text via chat completion."""
        if not self.available:
            return None

        try:
            response = self.client.chat_completion(
                model=self.model,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt},
                ],
                temperature=max(temperature, 0.01),
                max_tokens=max_tokens,
                top_p=0.9,
            )
            content = response.choices[0].message.content
            return content.strip() if content else None
        except Exception as e:
            print(f"[warn]Text inference error: {e}")
            return None

    def infer_json(
        self,
        system_prompt: str,
        user_prompt: str,
        temperature: float = 0.1,
        max_tokens: int = 1024,
    ) -> Optional[Dict[str, Any]]:
        """Generate and parse a JSON response."""
        raw = self.infer(system_prompt, user_prompt, temperature, max_tokens)
        if not raw:
            return None

        # Try direct JSON parse
        try:
            return json.loads(raw)
        except json.JSONDecodeError:
            pass

        # Try stripping markdown fences
        try:
            clean = re.sub(r"```(?:json)?\n?(.*?)\n?```", r"\1", raw, flags=re.DOTALL)
            return json.loads(clean)
        except json.JSONDecodeError:
            pass

        # Try extracting the first JSON object
        try:
            match = re.search(r"\{.*\}", raw, flags=re.DOTALL)
            if match:
                return json.loads(match.group(0))
        except json.JSONDecodeError:
            pass

        return None


# Global client instance
_client = None


def get_client() -> TextClient:
    """Get or create the global text client."""
    global _client
    if _client is None:
        _client = TextClient()
    return _client


def infer_text(
    system_prompt: str,
    user_prompt: str,
    temperature: float = 0.65,
    max_tokens: int = 512,
) -> Optional[str]:
    """Inference wrapper function."""
    return get_client().infer(system_prompt, user_prompt, temperature, max_tokens)


def infer_json(
    system_prompt: str,
    user_prompt: str,
    temperature: float = 0.1,
    max_tokens: int = 1024,
) -> Optional[Dict[str, Any]]:
    """JSON inference wrapper function."""
    return get_client().infer_json(system_prompt, user_prompt, temperature, max_tokens)