Update src/backends.py

src/backends.py

@@ -1,130 +1,176 @@
 import os
-import time
-import json
-import requests
+import base64
+import mimetypes
 from dataclasses import dataclass
-from typing import Optional, Dict, Any, Protocol
+from typing import Optional, Dict, Any, Protocol, Tuple
+
+import requests
 
-from huggingface_hub import InferenceClient
+# OpenAI
+from openai import OpenAI
 
-# Local backend (optional)
-try:
-    from transformers import AutoTokenizer, AutoModelForCausalLM
-    import torch
-except Exception:
-    AutoTokenizer = None
-    AutoModelForCausalLM = None
-    torch = None
+# 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]) -> str:
+    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 HFInferenceAPIBackend:
-    """
-    Uses HF Inference API via huggingface_hub.InferenceClient.
-    Works well on Spaces for large models if you provide HF_TOKEN in Secrets.
-    """
+class OpenAIBackend:
     model_id: str
-    token: Optional[str] = None
-    timeout_s: int = 120
+    api_key: Optional[str] = None
 
     def __post_init__(self):
-        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)
+        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]) -> str:
-        # We use chat.completions when available (for chat-tuned models),
-        # otherwise fall back to text_generation.
-        # InferenceClient adapts per model capabilities.
+    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", 600))
+        max_new_tokens = int(params.get("max_new_tokens", 800))
         top_p = float(params.get("top_p", 0.95))
-        repetition_penalty = float(params.get("repetition_penalty", 1.05))
-
-        # Try chat first
-        try:
-            messages = []
-            if system:
-                messages.append({"role": "system", "content": system})
-            messages.append({"role": "user", "content": prompt})
-
-            resp = self.client.chat.completions.create(
-                model=self.model_id,
-                messages=messages,
-                temperature=temperature,
-                max_tokens=max_new_tokens,
-                top_p=top_p,
-            )
-            return resp.choices[0].message.content
-        except Exception:
-            # Fallback: text generation
-            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,
-                repetition_penalty=repetition_penalty,
-                do_sample=True,
-                return_full_text=False,
-            )
-            return out
+
+        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
 
 
 @dataclass
-class LocalTransformersBackend:
-    """
-    Loads model locally in the Space container.
-    Use only small models unless you have GPU Space and enough memory.
-    """
+class GeminiBackend:
     model_id: str
-    device: str = "cpu"
+    api_key: Optional[str] = None
 
     def __post_init__(self):
-        if AutoTokenizer is None or AutoModelForCausalLM is None:
-            raise RuntimeError("transformers/torch not available in this environment.")
+        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.tokenizer = AutoTokenizer.from_pretrained(self.model_id, use_fast=True)
-        self.model = AutoModelForCausalLM.from_pretrained(self.model_id)
-        if torch is not None:
-            self.model.to(self.device)
-
-    def generate(self, prompt: str, *, system: Optional[str], params: Dict[str, Any]) -> str:
+    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", 600))
+        max_new_tokens = int(params.get("max_new_tokens", 800))
         top_p = float(params.get("top_p", 0.95))
-        repetition_penalty = float(params.get("repetition_penalty", 1.05))
 
-        full_prompt = (f"{system}\n\n{prompt}" if system else prompt)
+        parts = []
+
+        # 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))
 
-        inputs = self.tokenizer(full_prompt, return_tensors="pt")
-        if torch is not None:
-            inputs = {k: v.to(self.device) for k, v in inputs.items()}
+        text = prompt if not system else f"{system}\n\n{prompt}"
+        parts.append(text)
 
-        with torch.no_grad():
-            output_ids = self.model.generate(
-                **inputs,
-                do_sample=True,
+        resp = self.client.models.generate_content(
+            model=self.model_id,
+            contents=parts,
+            config=types.GenerateContentConfig(
                 temperature=temperature,
                 top_p=top_p,
-                repetition_penalty=repetition_penalty,
-                max_new_tokens=max_new_tokens,
+                max_output_tokens=max_new_tokens,
             )
+        )
+        return resp.text or ""
+
 
-        text = self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
-        # Heuristic: remove the prompt prefix if present
-        if text.startswith(full_prompt):
-            text = text[len(full_prompt):].lstrip()
-        return text
+@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
+            )
 
-def make_backend(backend_type: str, model_id: str) -> LLMBackend:
-    if backend_type == "hf_inference_api":
-        return HFInferenceAPIBackend(model_id=model_id)
-    if backend_type == "local_transformers":
-        # auto-device for local; keep cpu by default
-        return LocalTransformersBackend(model_id=model_id, device="cpu")
-    raise ValueError(f"Unknown backend: {backend_type}")
+        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}")