praxis/server/llm/ollama_cloud.py

"""Ollama Cloud LLM adapter behind the LLMProvider interface (D-020).

Direct API to https://ollama.com/api/chat with OLLAMA_API_KEY bearer,
stream=True. Two models:
  - gemma4:cloud (role-play fast path, 256K ctx)
  - deepseek-v4-flash:cloud (debrief + branch classifier, no-think mode)

R6 resolution: Pipecat's OLLamaLLMService accepts a custom base_url + bearer
(see docs/latency-report.md). This adapter is a thin wrapper over the raw
/api/chat streaming endpoint so the pipeline has a stable, testable contract
independent of Pipecat's OpenAI-compat shim. The Pipecat pipeline wires the
LLM via this adapter (TASK-02-04) so a swap (e.g. self-hosted gemma4:e4b
post-pilot) requires no pipeline change.
"""

from __future__ import annotations

import json
import os
from typing import Any, AsyncIterator

from server.services.base import LLMProvider, LLMStreamChunk

CHAT_URL_DEFAULT = "https://ollama.com/api/chat"


class OllamaCloudLLM(LLMProvider):
    """Ollama Cloud direct-API LLM adapter (D-020)."""

    name = "ollama-cloud"

    def __init__(
        self,
        *,
        api_key: str | None = None,
        chat_url: str | None = None,
        roleplay_model: str | None = None,
        debrief_model: str | None = None,
    ) -> None:
        self._api_key = (api_key or os.environ.get("OLLAMA_API_KEY", "")).strip()
        self._chat_url = (chat_url or os.environ.get("OLLAMA_CHAT_URL", CHAT_URL_DEFAULT)).strip()
        self._roleplay_model = (
            roleplay_model or os.environ.get("OLLAMA_ROLEPLAY_MODEL", "gemma4:cloud")
        ).strip()
        self._debrief_model = (
            debrief_model
            or os.environ.get("OLLAMA_DEBRIEF_MODEL", "deepseek-v4-flash:cloud")
        ).strip()

    @property
    def roleplay_model(self) -> str:
        return self._roleplay_model

    @property
    def debrief_model(self) -> str:
        return self._debrief_model

    def _missing(self) -> bool:
        return not self._api_key

    def _headers(self) -> dict[str, str]:
        return {"Authorization": f"Bearer {self._api_key}", "Content-Type": "application/json"}

    def _body(
        self,
        messages: list[dict[str, str]],
        model: str,
        stream: bool,
        no_think: bool,
    ) -> dict[str, Any]:
        body: dict[str, Any] = {
            "model": model,
            "messages": messages,
            "stream": stream,
        }
        if no_think:
            # deepseek-v4-flash:cloud no-think mode (D-020) — skips reasoning
            # tokens for latency on the debrief / branch-classifier path.
            body["think"] = False
        return body

    async def chat(
        self,
        messages: list[dict[str, str]],
        *,
        stream: bool = True,
        model: str | None = None,
        no_think: bool = False,
    ) -> AsyncIterator[LLMStreamChunk]:
        """Stream chat-completion chunks from Ollama Cloud /api/chat."""
        mdl = model or self._roleplay_model
        if self._missing():
            # Graceful: yield a single empty chunk so callers don't crash.
            return
        import httpx

        try:
            async with httpx.AsyncClient(timeout=60.0) as client:
                async with client.stream(
                    "POST", self._chat_url, headers=self._headers(),
                    json=self._body(messages, mdl, stream, no_think),
                ) as resp:
                    if resp.status_code != 200:
                        # Auth/error — degrade to no chunks (pipeline stays up).
                        return
                    is_first = True
                    async for line in resp.aiter_lines():
                        if not line:
                            continue
                        try:
                            chunk = json.loads(line)
                        except json.JSONDecodeError:
                            continue
                        content = chunk.get("message", {}).get("content", "")
                        if content:
                            yield LLMStreamChunk(
                                content=content,
                                is_first=is_first,
                                finish_reason=chunk.get("done") and "stop" or None,
                                extra={"eval_count": chunk.get("eval_count")},
                            )
                            is_first = False
        except Exception:
            # Network/auth errors degrade to no chunks; the pipeline stays up.
            return

    async def chat_full(
        self,
        messages: list[dict[str, str]],
        *,
        model: str | None = None,
        no_think: bool = False,
    ) -> tuple[str, dict[str, Any]]:
        """Return (full_text, usage) for non-streaming (debrief / classifier)."""
        mdl = model or self._debrief_model
        parts: list[str] = []
        usage: dict[str, Any] = {"input_tokens": 0, "output_tokens": 0, "model": mdl}
        async for chunk in self.chat(
            messages, stream=True, model=mdl, no_think=no_think
        ):
            parts.append(chunk.content)
            if chunk.extra.get("eval_count"):
                usage["output_tokens"] = chunk.extra["eval_count"]
        return "".join(parts), usage


__all__ = ["OllamaCloudLLM"]