captcha_solver/engines/ollama_engine.py

"""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()