api.py

import httpx
import time
import json
from typing import Optional, List, Dict, Generator, Any

MAX_RETRIES = 3
RETRY_DELAY = 1.0  # seconds

class OllamaClient:
    def __init__(self, host: str = "http://localhost:11434"):
        self.host = host.rstrip('/')

    def health_check(self) -> bool:
        """Check if the Ollama server is reachable."""
        try:
            response = httpx.get(self.host, timeout=5.0)
            return response.status_code == 200
        except httpx.RequestError:
            return False

    def list_models(self) -> List[Dict[str, Any]]:
        try:
            response = httpx.get(f"{self.host}/api/tags")
            self._handle_error(response)
            return response.json().get("models", [])
        except httpx.RequestError as exc:
            raise ConnectionError(self._connection_hint()) from exc

    def _connection_hint(self) -> str:
        return (
            f"Unable to connect to Ollama at {self.host}.\n"
            "  - Is Ollama running? Start it with: ollama serve\n"
            "  - Wrong host? Set OLLAMA_HOST environment variable."
        )

    def _handle_error(self, response):
        try:
            response.raise_for_status()
        except httpx.HTTPStatusError as e:
            try:
                error_msg = e.response.json().get("error", e.response.text)
            except Exception:
                error_msg = str(e)
            raise ValueError(f"Ollama Error: {error_msg}") from e

    def pull_model(self, name: str) -> Generator[Dict[str, Any], None, None]:
        url = f"{self.host}/api/pull"
        data = {"name": name, "stream": True}
        try:
            with httpx.stream("POST", url, json=data, timeout=None) as response:
                self._handle_error(response)
                for line in response.iter_lines():
                    if line:
                        try:
                            yield json.loads(line)
                        except json.JSONDecodeError:
                            continue
        except httpx.RequestError as exc:
            raise ConnectionError(self._connection_hint()) from exc

    def delete_model(self, name: str) -> bool:
        url = f"{self.host}/api/delete"
        data = {"name": name}
        try:
            response = httpx.delete(url, json=data, timeout=10.0)
            self._handle_error(response)
            return True
        except ValueError:
            return False
        except httpx.RequestError:
            return False

    def show_model_info(self, name: str) -> Dict[str, Any]:
        url = f"{self.host}/api/show"
        data = {"name": name}
        try:
            response = httpx.post(url, json=data, timeout=10.0)
            self._handle_error(response)
            return response.json()
        except httpx.RequestError as exc:
            raise ConnectionError(self._connection_hint()) from exc

    def generate(self, model: str, prompt: str, system: Optional[str] = None, 
                 options: Optional[Dict[str, Any]] = None, 
                 images: Optional[List[str]] = None) -> Generator[str, None, None]:
        url = f"{self.host}/api/generate"
        data = {
            "model": model,
            "prompt": prompt,
            "stream": True,
            "system": system,
            "options": options,
            "images": images
        }
        # filter None values
        data = {k: v for k, v in data.items() if v is not None}

        last_exc = None
        for attempt in range(MAX_RETRIES):
            try:
                with httpx.stream("POST", url, json=data, timeout=None) as response:
                    self._handle_error(response)
                    for line in response.iter_lines():
                        if line:
                            try:
                                json_response = json.loads(line)
                                yield json_response.get("response", "")
                                if json_response.get("done", False):
                                    return
                            except json.JSONDecodeError:
                                continue
                return  # Stream finished successfully
            except httpx.RequestError as exc:
                last_exc = exc
                if attempt < MAX_RETRIES - 1:
                    time.sleep(RETRY_DELAY * (attempt + 1))
                    continue
        raise ConnectionError(self._connection_hint()) from last_exc

    def chat(self, model: str, messages: List[Dict[str, str]], 
             options: Optional[Dict[str, Any]] = None) -> Generator[str, None, None]:
        url = f"{self.host}/api/chat"
        data = {
            "model": model,
            "messages": messages,
            "stream": True,
            "options": options
        }
        data = {k: v for k, v in data.items() if v is not None}

        last_exc = None
        for attempt in range(MAX_RETRIES):
            try:
                with httpx.stream("POST", url, json=data, timeout=None) as response:
                    self._handle_error(response)
                    for line in response.iter_lines():
                        if line:
                            try:
                                json_response = json.loads(line)
                                message = json_response.get("message", {})
                                yield message.get("content", "")
                                if json_response.get("done", False):
                                    return
                            except json.JSONDecodeError:
                                continue
                return  # Stream finished successfully
            except httpx.RequestError as exc:
                last_exc = exc
                if attempt < MAX_RETRIES - 1:
                    time.sleep(RETRY_DELAY * (attempt + 1))
                    continue
        raise ConnectionError(self._connection_hint()) from last_exc