client.py

Go to the documentation of this file.

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# providers/ollama/client.py
 
"""
Ollama HTTP client — native /api/chat streaming.
 
Purpose
-------
:class:`OllamaClient` implements :class:`BaseClient` against Ollama's
native ``/api/chat`` endpoint (not the OpenAI-compat shim). Streaming is
parsed via :func:`providers.ollama.ndjson.parse_ndjson_stream`.
 
Responsibilities
----------------
- :class:`OllamaClient` — :class:`BaseClient` for the Ollama provider.
  ``stream_chat``, ``one_shot``, ``list_models``, ``show_model``.
- :func:`_connect` — build an ``HTTPConnection`` / ``HTTPSConnection``
  from the base URL (stdlib only, zero extra deps).
- Friendly ``RuntimeError`` when the Ollama daemon is not running.
 
Diagnostics
-----------
Domain: OC_CHAT.CLIENT
Levels:
    L2 — lifecycle events (connect, response status)
    L3 — request details (model, endpoint, chunk counts)
    L4 — deep tracing
 
Contracts
---------
- ``stream_chat`` yields ``("content", str)``, ``("tool_call", dict)``,
  and ``("usage", dict)`` tuples. ``"reasoning"`` is never yielded (local
  models don't separate reasoning from content in the Ollama protocol).
- ``tool_call`` dict: ``{"id": str, "name": str, "arguments": str}``
  where *arguments* is a JSON string (matches the rest of OC Chat).
- ``one_shot`` returns a plain ``str``; uses ``stream: false``.
- ``list_models`` hits ``GET /api/tags``; ``show_model`` hits
  ``POST /api/show``.
- Connection errors raise ``RuntimeError`` with a user-friendly message
  suggesting ``ollama serve``.
 
Dependencies
------------
- oc_diagnostics (_dbg)
- providers.ollama.ndjson (parse_ndjson_stream)
- http.client (HTTP and HTTPS connections to the Ollama daemon)
- json (encoding request bodies and tool-call arguments, decoding responses)
- urllib.parse (splitting the base URL into scheme, host and port)
- oc_chat.core.egress_guard (lazy; refusing a forbidden internal address as the Ollama host)
- option_c_agent_contracts (lazy; SupportsCancellation type of the stream_chat cancellation token)
"""
from __future__ import annotations
 
import http.client
import json
import urllib.parse
from collections.abc import Iterator
from typing import Any
 
from typing import TYPE_CHECKING
 
from oc_diagnostics import _dbg
 
from .ndjson import parse_ndjson_stream
 
if TYPE_CHECKING:
    from option_c_agent_contracts import SupportsCancellation
 
 
def _connect(base_url: str, timeout: int = 300) -> http.client.HTTPConnection:
    """Return an ``HTTPConnection`` or ``HTTPSConnection`` for *base_url*."""
    parsed = urllib.parse.urlparse(base_url)
    host = parsed.hostname or "localhost"
    port = parsed.port or (443 if parsed.scheme == "https" else 11434)
    # M8.0 WU-4c (F72): refuse a forbidden internal address (IMDS / link-local /
    # multicast / reserved). Ollama on loopback/LAN stays legal. Defense-in-depth:
    # ClientInterface already refuses such a base_url at bind time.
    from oc_chat.core.egress_guard import is_forbidden_provider_host
    if is_forbidden_provider_host(host):
        raise PermissionError(f"Ollama base_url is a forbidden internal address: {base_url}")
    if parsed.scheme == "https":
        return http.client.HTTPSConnection(host, port, timeout=timeout)
    return http.client.HTTPConnection(host, port, timeout=timeout)
 
 
#: Maximum bytes of error-body to include in raised exception messages.
#: 5xx error pages from a reverse-proxy in front of Ollama can be huge HTML;
#: 2 KB is enough for the daemon's "what went wrong" hint (which is always
#: at the head of the body) without blowing out logs or tracebacks.
_ERROR_BODY_TRUNCATE = 2048
 
 
def _raise_for_status(
    resp: http.client.HTTPResponse, base_url: str, method: str, path: str
) -> None:
    """Raise :class:`RuntimeError` when *resp* is a 4xx/5xx HTTP response.
 
    Plan 077: the Ollama daemon returns clear, useful error bodies on 4xx/5xx
    (e.g. ``{"error":"model 'X' not found"}``). The previous implementation
    handed the response straight to the NDJSON parser, which silently yielded
    zero chunks for an error body, masquerading as "no response content
    received". This helper surfaces the daemon's actual message instead.
 
    Reads the body eagerly (truncated to :data:`_ERROR_BODY_TRUNCATE`) and
    raises so callers never see a partially-consumed response on the error
    path. Callers consuming success responses (200-299) get the response
    untouched.
    """
    func = "_raise_for_status"
    if resp.status < 400:
        return
    body = resp.read(_ERROR_BODY_TRUNCATE + 1)
    if len(body) > _ERROR_BODY_TRUNCATE:
        body_text = body[:_ERROR_BODY_TRUNCATE].decode("utf-8", errors="replace") + "...[truncated]"
    else:
        body_text = body.decode("utf-8", errors="replace")
    _dbg(
        "OC_CHAT.CLIENT",
        func,
        f"HTTP {resp.status} on {method} {path}: {body_text!r}",
        1,
        override=True,
    )
    raise RuntimeError(
        f"Ollama API error {resp.status} on {method} {base_url}{path}: {body_text}"
    )
 
 
def _post(base_url: str, path: str, body: dict, timeout: int = 300) -> http.client.HTTPResponse:
    """POST *body* as JSON to *base_url + path*. Returns the response.
 
    Raises :class:`RuntimeError` on connection failure (unreachable daemon)
    and on HTTP 4xx/5xx responses — the latter via :func:`_raise_for_status`,
    which embeds the daemon's response body in the exception message
    (plan 077; previously the error body was silently swallowed).
    """
    func = "_post"
    _dbg("OC_CHAT.CLIENT", func, f"POST {base_url}{path}", 3)
    try:
        conn = _connect(base_url, timeout=timeout)
        data = json.dumps(body).encode()
        conn.request(
            "POST",
            path,
            body=data,
            headers={"Content-Type": "application/json"},
        )
        resp = conn.getresponse()
        _dbg("OC_CHAT.CLIENT", func, f"status={resp.status}", 3)
    except (ConnectionRefusedError, OSError, TimeoutError) as exc:
        raise RuntimeError(
            f"Ollama unreachable at {base_url!r} — is 'ollama serve' running? ({exc})"
        ) from exc
    _raise_for_status(resp, base_url, "POST", path)
    return resp
 
 
def _get(base_url: str, path: str, timeout: int = 30) -> http.client.HTTPResponse:
    """GET *base_url + path*. Returns the response.
 
    Raises identically to :func:`_post` — see its docstring for the HTTP-error
    semantics (plan 077). The same swallow-on-error bug applied here for
    ``/api/tags`` / ``/api/show`` callers.
    """
    func = "_get"
    _dbg("OC_CHAT.CLIENT", func, f"GET {base_url}{path}", 3)
    try:
        conn = _connect(base_url, timeout=timeout)
        conn.request("GET", path)
        resp = conn.getresponse()
        _dbg("OC_CHAT.CLIENT", func, f"status={resp.status}", 3)
    except (ConnectionRefusedError, OSError, TimeoutError) as exc:
        raise RuntimeError(
            f"Ollama unreachable at {base_url!r} — is 'ollama serve' running? ({exc})"
        ) from exc
    _raise_for_status(resp, base_url, "GET", path)
    return resp
 
 
class OllamaClient:
    """Native Ollama HTTP client implementing :class:`BaseClient`.
 
    Uses Ollama's ``/api/chat`` endpoint directly (not the OpenAI-compat
    shim at ``/v1/chat/completions``) to preserve Ollama-specific fields
    like ``eval_count``, ``eval_duration``, and ``tool_calls`` in their
    native shape.
    """
 
    # OCT-LINT: disable=no_hardcoded_secrets reason="Empty-string default for optional api_key param, not a credential"
    def __init__(self, api_key: str = "", base_url: str = "http://localhost:11434") -> None:
        func = "__init__"
        self.api_key = api_key  # unused by Ollama; kept for BaseClient compat
        self.base_url = base_url.rstrip("/")
        _dbg("OC_CHAT.CLIENT", func, f"OllamaClient base_url={self.base_url!r}", 3)
 
    # ── Streaming chat ──────────────────────────────────────────────────────
 
    def stream_chat(
        self,
        messages: list[dict],
        model: str,
        *,
        effort: str | None = None,
        extra_body: dict | None = None,
        tools: list[dict] | None = None,
        max_tokens: int | None = None,
        cancellation: "SupportsCancellation | None" = None,
    ) -> Iterator[tuple[str, Any]]:
        """Stream a chat completion via ``POST /api/chat``.
 
        Yields ``(chunk_type, data)`` tuples:
 
        * ``("content", str)`` — one or more tokens of assistant text
        * ``("tool_call", dict)`` — ``{id, name, arguments}`` where
          *arguments* is a JSON *string* (matching the OC Chat contract)
        * ``("usage", dict)`` — usage figures from the final chunk
        * ``("finish_reason", str)`` — the final chunk's ``done_reason``;
          ``"length"`` means generation stopped at ``num_predict``
 
        If *cancellation* is provided, the token is polled before yielding
        each chunk; a cancelled token raises :class:`CancellationError` and
        the NDJSON stream is abandoned (Phase O3c).
        """
        func = "stream_chat"
        payload: dict[str, Any] = {
            "model": model,
            "messages": messages,
            "stream": True,
        }
        if tools:
            payload["tools"] = tools
        # Merge extra_body first (provider generation options like temperature, top_p)
        if extra_body:
            for k, v in extra_body.items():
                if k == "options" and "options" in payload and isinstance(v, dict):
                    payload["options"] = {**payload["options"], **v}
                else:
                    payload[k] = v
        # max_tokens overrides num_predict specifically
        if max_tokens is not None:
            payload.setdefault("options", {})["num_predict"] = max_tokens
 
        _dbg("OC_CHAT.CLIENT", func,
             f"model={model} tools={'yes' if tools else 'no'} max_tokens={max_tokens}", 2)
 
        resp = _post(self.base_url, "/api/chat", payload)
        chunk_count = 0
 
        for chunk in parse_ndjson_stream(resp):
            if cancellation is not None:
                cancellation.raise_if_cancelled()
            chunk_count += 1
            msg = chunk.get("message") or {}
            content = msg.get("content") or ""
 
            if content:
                yield ("content", content)
 
            # Tool calls arrive in the final (done=True) chunk
            if chunk.get("done") and msg.get("tool_calls"):
                for i, tc in enumerate(msg["tool_calls"]):
                    fn = tc.get("function") or {}
                    args = fn.get("arguments", {})
                    # Ollama arguments is a dict; convert to JSON string for OC Chat
                    arguments_str = json.dumps(args) if isinstance(args, dict) else str(args)
                    yield ("tool_call", {
                        "id": f"ollama_tc_{i}",
                        "name": fn.get("name", ""),
                        "arguments": arguments_str,
                    })
 
            # Emit usage from the final chunk
            if chunk.get("done"):
                usage: dict[str, Any] = {
                    "prompt_tokens": chunk.get("prompt_eval_count", 0),
                    "completion_tokens": chunk.get("eval_count", 0),
                    "total_tokens": (
                        chunk.get("prompt_eval_count", 0)
                        + chunk.get("eval_count", 0)
                    ),
                    # Pass through raw Ollama perf fields for O5 parse_performance
                    "prompt_eval_count": chunk.get("prompt_eval_count", 0),
                    "prompt_eval_duration": chunk.get("prompt_eval_duration", 0),
                    "eval_count": chunk.get("eval_count", 0),
                    "eval_duration": chunk.get("eval_duration", 0),
                    "total_duration": chunk.get("total_duration", 0),
                }
                yield ("usage", usage)
                # Ollama's terminal-chunk ``done_reason`` uses the same
                # vocabulary as the OpenAI-compatible ``finish_reason``
                # ("stop" / "length" / "tool_calls"), so it crosses the
                # provider seam untranslated. "length" == num_predict hit.
                if chunk.get("done_reason"):
                    yield ("finish_reason", chunk["done_reason"])
 
        _dbg("OC_CHAT.CLIENT", func, f"stream done — {chunk_count} chunks", 3)
 
    # ── One-shot (non-streaming) ────────────────────────────────────────────
 
    def one_shot(
        self,
        messages: list[dict],
        model: str,
        *,
        effort: str | None = None,
        extra_body: dict | None = None,
        max_tokens: int | None = None,
    ) -> str:
        """Non-streaming completion via ``POST /api/chat`` with ``stream: false``."""
        func = "one_shot"
        payload: dict[str, Any] = {
            "model": model,
            "messages": messages,
            "stream": False,
        }
        if max_tokens is not None:
            payload["options"] = {"num_predict": max_tokens}
 
        _dbg("OC_CHAT.CLIENT", func, f"model={model}", 2)
        resp = _post(self.base_url, "/api/chat", payload)
        raw = resp.read()
        data = json.loads(raw)
        return (data.get("message") or {}).get("content") or ""
 
    # ── Model listing ───────────────────────────────────────────────────────
 
    def list_models(self) -> list[str]:
        """Return model names from ``GET /api/tags``."""
        func = "list_models"
        try:
            resp = _get(self.base_url, "/api/tags")
            raw = resp.read()
            data = json.loads(raw)
            names = [m["name"] for m in data.get("models", [])]
            _dbg("OC_CHAT.CLIENT", func, f"{len(names)} models", 3)
            return names
        except RuntimeError:
            raise
        except Exception as exc:
            _dbg("OC_CHAT.CLIENT", func, f"list_models failed: {exc}", 2)
            return []
 
    # ── Model info ──────────────────────────────────────────────────────────
 
    def show_model(self, name: str) -> dict:
        """Return raw metadata from ``POST /api/show``."""
        func = "show_model"
        _dbg("OC_CHAT.CLIENT", func, f"name={name!r}", 3)
        resp = _post(self.base_url, "/api/show", {"name": name}, timeout=30)
        raw = resp.read()
        return json.loads(raw)