File size: 8,285 Bytes
e0265b9 f8c73f9 e0265b9 f8c73f9 e0265b9 f8c73f9 e0265b9 f8c73f9 e0265b9 f8c73f9 e0265b9 f8c73f9 e0265b9 f8c73f9 e0265b9 f8c73f9 e0265b9 f8c73f9 e0265b9 f8c73f9 e0265b9 f8c73f9 e0265b9 f8c73f9 e0265b9 f8c73f9 e0265b9 f8c73f9 e0265b9 f8c73f9 e0265b9 f8c73f9 e0265b9 f8c73f9 e0265b9 | 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 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 | 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", ""))
|