Spaces:
Sleeping
Sleeping
File size: 3,212 Bytes
4ad7943 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | """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()
|