Upload captcha_solver/engines/ollama_engine.py with huggingface_hub

captcha_solver/engines/ollama_engine.py

@@ -0,0 +1,96 @@
+"""Ollama HTTP client (optional upgrade).
+
+If ollama is enabled and running, this engine provides better quality
+text + vision inference by routing to local ollama models. The HF
+engines remain the default for fully-offline operation.
+"""
+
+from __future__ import annotations
+
+import base64
+import json
+import urllib.error
+import urllib.request
+from typing import Any, Optional
+
+from captcha_solver.engines.base import BaseEngine
+from captcha_solver.config import get_settings
+
+
+class OllamaEngine(BaseEngine):
+    name = "ollama"
+
+    def __init__(self) -> None:
+        super().__init__()
+        self._enabled = False
+        self._host = ""
+
+    def _do_load(self) -> None:
+        s = get_settings()
+        self._enabled = s.ollama_enabled
+        self._host = s.ollama_host.rstrip("/")
+        if not self._enabled:
+            raise RuntimeError("ollama disabled (ollama_enabled=false in config)")
+        try:
+            with urllib.request.urlopen(f"{self._host}/api/tags", timeout=3) as r:
+                if r.status != 200:
+                    raise RuntimeError(f"ollama returned {r.status}")
+        except Exception as exc:
+            self._enabled = False
+            raise RuntimeError(f"ollama not reachable at {self._host}: {exc}") from exc
+
+    def _do_unload(self) -> None:
+        self._enabled = False
+
+    @property
+    def enabled(self) -> bool:
+        return self._enabled and self._loaded
+
+    def _post(self, path: str, payload: dict, timeout: int = 30) -> dict:
+        data = json.dumps(payload).encode("utf-8")
+        req = urllib.request.Request(
+            f"{self._host}{path}",
+            data=data,
+            headers={"Content-Type": "application/json"},
+            method="POST",
+        )
+        with urllib.request.urlopen(req, timeout=timeout) as r:
+            return json.loads(r.read().decode("utf-8"))
+
+    def generate_text(
+        self,
+        prompt: str,
+        system: Optional[str] = None,
+        model: Optional[str] = None,
+        max_tokens: int = 64,
+    ) -> str:
+        if not self.enabled:
+            raise RuntimeError("ollama not enabled")
+        s = get_settings()
+        payload: dict[str, Any] = {
+            "model": model or s.ollama_text_model,
+            "prompt": prompt,
+            "stream": False,
+            "options": {"num_predict": max_tokens, "temperature": 0.0},
+        }
+        if system:
+            payload["system"] = system
+        return self._post("/api/generate", payload, s.ollama_timeout).get("response", "").strip()
+
+    def describe_image(self, pil_image, prompt: str, model: Optional[str] = None) -> str:
+        if not self.enabled:
+            raise RuntimeError("ollama not enabled")
+        s = get_settings()
+        import io
+
+        buf = io.BytesIO()
+        pil_image.save(buf, format="PNG")
+        b64 = base64.b64encode(buf.getvalue()).decode("ascii")
+        payload = {
+            "model": model or s.ollama_vision_model,
+            "prompt": prompt,
+            "images": [b64],
+            "stream": False,
+            "options": {"num_predict": 200, "temperature": 0.0},
+        }
+        return self._post("/api/generate", payload, s.ollama_timeout).get("response", "").strip()