Praxis CI
147 lines
5.2 KiB
Python
147 lines
5.2 KiB
Python
"""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"] |