80 lines
3.0 KiB
Python
80 lines
3.0 KiB
Python
import asyncio
|
|
|
|
import httpx
|
|
|
|
from .constants import DEFAULT_OLLAMA_TIMEOUT, OLLAMA_CA_FILE, OLLAMA_TIMEOUT, OLLAMA_TOKEN, OLLAMA_URL
|
|
|
|
VERSION_PATH = "/api/version"
|
|
TAGS_PATH = "/api/tags"
|
|
CHAT_PATH = "/api/chat"
|
|
DEFAULT_TIMEOUT = 60.0
|
|
DEFAULT_RETRIES = 2
|
|
DEFAULT_BACKOFF = 1.0
|
|
SERVER_ERROR_MIN = 500
|
|
CLIENT_ERROR_MIN = 400
|
|
|
|
|
|
class OllamaError(Exception):
|
|
pass
|
|
|
|
|
|
class OllamaClient:
|
|
def __init__(self, base_url: str, token: str | None = None, timeout: float = DEFAULT_TIMEOUT,
|
|
retries: int = DEFAULT_RETRIES, backoff: float = DEFAULT_BACKOFF,
|
|
verify: bool | str = True,
|
|
transport: httpx.AsyncBaseTransport | None = None):
|
|
headers = {"Authorization": f"Bearer {token}"} if token else {}
|
|
self._client = httpx.AsyncClient(
|
|
base_url=base_url.rstrip("/"), headers=headers, timeout=timeout,
|
|
verify=verify, transport=transport,
|
|
)
|
|
self._retries = retries
|
|
self._backoff = backoff
|
|
|
|
@classmethod
|
|
def from_settings(cls, store) -> "OllamaClient | None":
|
|
url = store.get(OLLAMA_URL)
|
|
if not url:
|
|
return None
|
|
timeout = float(store.get(OLLAMA_TIMEOUT, DEFAULT_OLLAMA_TIMEOUT))
|
|
verify: bool | str = store.get(OLLAMA_CA_FILE) or True
|
|
return cls(url, store.get(OLLAMA_TOKEN), timeout=timeout, verify=verify)
|
|
|
|
async def _request(self, method: str, path: str, **kwargs) -> dict:
|
|
last: Exception | None = None
|
|
for attempt in range(self._retries + 1):
|
|
if attempt:
|
|
await asyncio.sleep(self._backoff * attempt)
|
|
try:
|
|
response = await self._client.request(method, path, **kwargs)
|
|
except httpx.HTTPError as exc:
|
|
last = OllamaError(f"{type(exc).__name__}: {exc}")
|
|
continue
|
|
if response.status_code >= SERVER_ERROR_MIN:
|
|
last = OllamaError(f"HTTP {response.status_code} on {path}")
|
|
continue
|
|
if response.status_code >= CLIENT_ERROR_MIN:
|
|
raise OllamaError(f"HTTP {response.status_code} on {path}: {response.text[:200]}")
|
|
return response.json()
|
|
raise last or OllamaError("request failed")
|
|
|
|
async def version(self) -> str:
|
|
return (await self._request("GET", VERSION_PATH))["version"]
|
|
|
|
async def list_models(self) -> list[str]:
|
|
body = await self._request("GET", TAGS_PATH)
|
|
return [model["name"] for model in body.get("models", [])]
|
|
|
|
async def chat(self, model: str, messages: list[dict], tools: list[dict] | None = None,
|
|
options: dict | None = None) -> dict:
|
|
payload: dict = {"model": model, "messages": messages, "stream": False}
|
|
if tools:
|
|
payload["tools"] = tools
|
|
if options:
|
|
payload["options"] = options
|
|
body = await self._request("POST", CHAT_PATH, json=payload)
|
|
return body["message"]
|
|
|
|
async def aclose(self) -> None:
|
|
await self._client.aclose()
|