from __future__ import annotations import json import base64 import urllib.error import urllib.request from pathlib import Path from typing import Any, Callable class OllamaError(RuntimeError): pass class OllamaClient: def __init__( self, base_url: str, model: str, timeout: float = 2.5, chat_max_tokens: int | None = None, chat_response_length: str = "automatic", ) -> None: self.base_url = base_url.rstrip("/") self.model = model self.timeout = timeout self.chat_max_tokens = chat_max_tokens self.chat_response_length = str(chat_response_length or "automatic").casefold() def _chat_system(self, system: str) -> str: """Return the application system prompt unchanged.""" return system def _num_predict(self, default: int, *, chat: bool = True, prompt: str = "", image_count: int = 0) -> int: """Qwen3's reasoning commonly needs more than a short-chat token budget.""" if chat and self.chat_max_tokens is not None: limit = max(64, int(self.chat_max_tokens)) choices = {"short": 256, "balanced": 512, "detailed": limit} selected = choices.get(self.chat_response_length) if selected is not None: return min(limit, selected) text = prompt.casefold() if image_count: budget = 384 elif any(word in text for word in ("compare", "explain", "research", "plan", "review", "why", "how")): budget = 768 elif any(word in text for word in ("summarize", "details", "ideas", "examples")): budget = 512 else: budget = 256 return min(limit, max(128, budget)) return 1024 if self.model.casefold().startswith("qwen3") else default @staticmethod def prepare_image(path: str | Path, *, maximum_side: int = 1536) -> str: """Return a compact PNG attachment without changing the user's source image.""" source = Path(path).expanduser() if not source.is_file(): raise OllamaError("The attached image no longer exists.") try: from PIL import Image, ImageOps with Image.open(source) as opened: image = ImageOps.exif_transpose(opened).convert("RGB") image.thumbnail((maximum_side, maximum_side), Image.Resampling.LANCZOS) from io import BytesIO buffer = BytesIO() image.save(buffer, format="PNG", optimize=True) except (OSError, ValueError) as exc: raise OllamaError(f"ADAM could not read {source.name} as an image.") from exc return base64.b64encode(buffer.getvalue()).decode("ascii") def is_available(self, timeout: float = 0.35) -> bool: request = urllib.request.Request(f"{self.base_url}/api/tags", method="GET") try: with urllib.request.urlopen(request, timeout=timeout) as response: return response.status == 200 except (OSError, urllib.error.URLError, TimeoutError): return False def list_models(self, timeout: float = 2.0) -> list[str]: request = urllib.request.Request(f"{self.base_url}/api/tags", method="GET") try: with urllib.request.urlopen(request, timeout=timeout) as response: payload = json.loads(response.read().decode("utf-8")) except (OSError, urllib.error.URLError, TimeoutError, json.JSONDecodeError): return [] return [ str(model.get("name")) for model in payload.get("models", []) if isinstance(model, dict) and model.get("name") ] def generate_json(self, system: str, prompt: str) -> dict[str, Any]: raw_response = self._generate(system, prompt, json_format=True) try: parsed = json.loads(raw_response) except (TypeError, json.JSONDecodeError) as exc: raise OllamaError("Ollama did not return a valid JSON plan.") from exc if not isinstance(parsed, dict): raise OllamaError("Ollama returned an unsupported plan shape.") return parsed def generate_text( self, system: str, prompt: str, *, image_paths: list[str | Path] | None = None, ) -> str: response = self._generate(system, prompt, json_format=False, image_paths=image_paths).strip() if not response: raise OllamaError("Ollama returned an empty response.") return response def generate_text_stream( self, system: str, prompt: str, on_chunk: Callable[[str], None], image_paths: list[str | Path] | None = None, ) -> str: images = [self.prepare_image(path) for path in (image_paths or [])] payload: dict[str, Any] = { "model": self.model, "system": self._chat_system(system), "prompt": prompt, # Free GPU memory for image generation and other local workloads as # soon as this one-shot text request finishes. "keep_alive": 0, "stream": True, "options": {"temperature": 0.35, "num_predict": self._num_predict(180, prompt=prompt, image_count=len(images))}, } if images: payload["images"] = images request = urllib.request.Request( f"{self.base_url}/api/generate", data=json.dumps(payload).encode("utf-8"), headers={"Content-Type": "application/json"}, method="POST", ) pieces: list[str] = [] try: with urllib.request.urlopen(request, timeout=self.timeout) as response: for raw_line in response: if not raw_line.strip(): continue event = json.loads(raw_line.decode("utf-8")) chunk = str(event.get("response", "")) if chunk: pieces.append(chunk) on_chunk(chunk) if event.get("done"): break except (OSError, urllib.error.URLError, json.JSONDecodeError) as exc: raise OllamaError(f"Ollama streaming request failed: {exc}") from exc result = "".join(pieces).strip() if not result: raise OllamaError("Ollama returned an empty response.") return result def _generate(self, system: str, prompt: str, *, json_format: bool, image_paths: list[str | Path] | None = None) -> str: images = [self.prepare_image(path) for path in (image_paths or [])] payload: dict[str, Any] = { "model": self.model, "system": self._chat_system(system), "prompt": prompt, "keep_alive": 0, "stream": False, "options": { "temperature": 0.1 if json_format else 0.35, "num_predict": self._num_predict(300 if json_format else 180, chat=not json_format, prompt=prompt, image_count=len(images)), }, } if json_format: payload["format"] = "json" if images: payload["images"] = images body = json.dumps( payload ).encode("utf-8") request = urllib.request.Request( f"{self.base_url}/api/generate", data=body, headers={"Content-Type": "application/json"}, method="POST", ) try: with urllib.request.urlopen(request, timeout=self.timeout) as response: payload = json.loads(response.read().decode("utf-8")) except urllib.error.HTTPError as exc: detail = "" try: error_payload = json.loads(exc.read().decode("utf-8")) detail = str(error_payload.get("error", "")) except Exception: pass raise OllamaError(detail or f"Ollama returned HTTP {exc.code}.") from exc except (OSError, urllib.error.URLError, json.JSONDecodeError) as exc: raise OllamaError(f"Ollama request failed: {exc}") from exc return str(payload.get("response", ""))