Update src/backends.py

src/backends.py

@@ -1,176 +1,58 @@
 import os
-import base64
-import mimetypes
 from dataclasses import dataclass
-from typing import Optional, Dict, Any, Protocol, Tuple
+from typing import Optional, Dict, Any, Union
+from huggingface_hub import InferenceClient
 
-import requests
-
-# OpenAI
-from openai import OpenAI
-
-# Gemini
-from google import genai
-from google.genai import types
-
-
-class LLMBackend(Protocol):
-    def generate(
-        self,
-        prompt: str,
-        *,
-        system: Optional[str],
-        params: Dict[str, Any],
-        image_path: Optional[str] = None,
-    ) -> str:
-        ...
-
-
-def _file_to_data_url(path: str) -> Tuple[str, str]:
-    mime, _ = mimetypes.guess_type(path)
-    mime = mime or "image/png"
-    with open(path, "rb") as f:
-        b64 = base64.b64encode(f.read()).decode("utf-8")
-    return f"data:{mime};base64,{b64}", mime
-
-
-@dataclass
-class OpenAIBackend:
-    model_id: str
-    api_key: Optional[str] = None
-
-    def __post_init__(self):
-        self.api_key = self.api_key or os.getenv("OPENAI_API_KEY")
-        if not self.api_key:
-            raise RuntimeError("OPENAI_API_KEY is not set (Spaces → Settings → Secrets).")
-        self.client = OpenAI(api_key=self.api_key)
-
-    def generate(self, prompt: str, *, system: Optional[str], params: Dict[str, Any], image_path: Optional[str] = None) -> str:
-        temperature = float(params.get("temperature", 0.2))
-        max_new_tokens = int(params.get("max_new_tokens", 800))
-        top_p = float(params.get("top_p", 0.95))
-
-        user_content = [{"type": "input_text", "text": prompt}]
-        if image_path:
-            data_url, _ = _file_to_data_url(image_path)
-            user_content.append({"type": "input_image", "image_url": data_url})
-
-        # Responses API: supports image inputs via input_image items. :contentReference[oaicite:4]{index=4}
-        input_items = []
-        if system:
-            input_items.append({
-                "role": "developer",
-                "content": [{"type": "input_text", "text": system}]
-            })
-        input_items.append({"role": "user", "content": user_content})
-
-        resp = self.client.responses.create(
-            model=self.model_id,
-            input=input_items,
-            temperature=temperature,
-            top_p=top_p,
-            max_output_tokens=max_new_tokens,
-        )
-        return resp.output_text
+try:
+    from PIL import Image
+except Exception:
+    Image = None
 
 
 @dataclass
-class GeminiBackend:
+class HFInferenceAPIBackend:
     model_id: str
-    api_key: Optional[str] = None
+    token: Optional[str] = None
+    timeout_s: int = 180
 
     def __post_init__(self):
-        self.api_key = self.api_key or os.getenv("GEMINI_API_KEY")
-        if not self.api_key:
-            raise RuntimeError("GEMINI_API_KEY is not set (Spaces → Settings → Secrets).")
-        self.client = genai.Client(api_key=self.api_key)
+        self.token = self.token or os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACEHUB_API_TOKEN")
+        self.client = InferenceClient(model=self.model_id, token=self.token, timeout=self.timeout_s)
 
-    def generate(self, prompt: str, *, system: Optional[str], params: Dict[str, Any], image_path: Optional[str] = None) -> str:
+    def generate(self, prompt: str, *, system: Optional[str], params: Dict[str, Any]) -> str:
         temperature = float(params.get("temperature", 0.2))
-        max_new_tokens = int(params.get("max_new_tokens", 800))
+        max_new_tokens = int(params.get("max_new_tokens", 600))
         top_p = float(params.get("top_p", 0.95))
 
-        parts = []
+        # Chat when possible
+        try:
+            messages = []
+            if system:
+                messages.append({"role": "system", "content": system})
+            messages.append({"role": "user", "content": prompt})
 
-        # Gemini: inline bytes via Part.from_bytes (официальный пример). :contentReference[oaicite:5]{index=5}
-        if image_path:
-            mime, _ = mimetypes.guess_type(image_path)
-            mime = mime or "image/png"
-            with open(image_path, "rb") as f:
-                img_bytes = f.read()
-            parts.append(types.Part.from_bytes(data=img_bytes, mime_type=mime))
-
-        text = prompt if not system else f"{system}\n\n{prompt}"
-        parts.append(text)
-
-        resp = self.client.models.generate_content(
-            model=self.model_id,
-            contents=parts,
-            config=types.GenerateContentConfig(
+            resp = self.client.chat.completions.create(
+                model=self.model_id,
+                messages=messages,
                 temperature=temperature,
+                max_tokens=max_new_tokens,
                 top_p=top_p,
-                max_output_tokens=max_new_tokens,
             )
-        )
-        return resp.text or ""
-
-
-@dataclass
-class DeepSeekBackend:
-    model_id: str
-    api_key: Optional[str] = None
-    base_url: str = "https://api.deepseek.com"
-
-    def __post_init__(self):
-        self.api_key = self.api_key or os.getenv("DEEPSEEK_API_KEY")
-        if not self.api_key:
-            raise RuntimeError("DEEPSEEK_API_KEY is not set (Spaces → Settings → Secrets).")
-
-    def generate(self, prompt: str, *, system: Optional[str], params: Dict[str, Any], image_path: Optional[str] = None) -> str:
-        # DeepSeek official docs show text chat completions. :contentReference[oaicite:6]{index=6}
-        temperature = float(params.get("temperature", 0.2))
-        max_tokens = int(params.get("max_new_tokens", 800))
-        top_p = float(params.get("top_p", 0.95))
-
-        if image_path:
-            prompt = (
-                "ВАЖНО: Пользователь приложил изображение (диаграмму), "
-                "но этот провайдер в текущей конфигурации работает только с текстом. "
-                "Попроси пользователя описать диаграмму текстом, либо продолжи только по тексту.\n\n"
-                + prompt
+            return resp.choices[0].message.content
+        except Exception:
+            out = self.client.text_generation(
+                prompt=(f"{system}\n\n{prompt}" if system else prompt),
+                temperature=temperature,
+                max_new_tokens=max_new_tokens,
+                top_p=top_p,
+                do_sample=True,
+                return_full_text=False,
             )
-
-        messages = []
-        if system:
-            messages.append({"role": "system", "content": system})
-        messages.append({"role": "user", "content": prompt})
-
-        r = requests.post(
-            f"{self.base_url}/chat/completions",
-            headers={
-                "Authorization": f"Bearer {self.api_key}",
-                "Content-Type": "application/json",
-            },
-            json={
-                "model": self.model_id,
-                "messages": messages,
-                "temperature": temperature,
-                "top_p": top_p,
-                "max_tokens": max_tokens,
-                "stream": False,
-            },
-            timeout=120,
-        )
-        r.raise_for_status()
-        data = r.json()
-        return data["choices"][0]["message"]["content"]
-
-
-def make_backend(provider: str, model_id: str) -> LLMBackend:
-    if provider == "openai":
-        return OpenAIBackend(model_id=model_id)
-    if provider == "gemini":
-        return GeminiBackend(model_id=model_id)
-    if provider == "deepseek":
-        return DeepSeekBackend(model_id=model_id)
-    raise ValueError(f"Unknown provider: {provider}")
+            return out
+
+    # --- NEW: image -> text (OCR / caption) ---
+    def image_to_text(self, image: "Image.Image") -> str:
+        """
+        Uses HF task 'image-to-text' for models like TrOCR or BLIP-caption.
+        """
+        return self.client.image_to_text(image).generated_text