"""FallbackOllamaBackend — walks a route's chain until one entry answers.
A route is an ordered chain of ``model on target`` entries (see ``config.py``).
The first entry is the primary; the rest are backups. The engine:
* starts every call from the entry that is currently active for the route
(``begin_turn`` pins it at the start of a user turn — see below),
* skips entries whose target or model is known to be unhealthy, unless the
recorded retry time has passed,
* on a typed failure records *why* it failed (network / quota / auth / 5xx /
missing model) and moves to the next entry,
* retries an entry only when it is the last one left — that is what stops a
dead cloud model from costing ``retries × first_chunk_timeout`` on every
call while a healthy local model sits one step down the chain,
* reports the model *and the target* that served, so the UI can say which
side of the pair the user is on.
Switch-back is deliberately not part of a call: ``health.RouteHealth`` probes
degraded targets in the background every minute, and ``begin_turn`` starts the
next user turn on the primary entry once it is healthy again. A running turn
is never moved between hosts mid-stream.
"""
from __future__ import annotations
import asyncio
from typing import AsyncGenerator
import structlog
from navi.config import settings
from navi.exceptions import (
LLMAuthError,
LLMBackendError,
LLMConnectionError,
LLMModelNotFoundError,
LLMQuotaExceededError,
LLMServerUnavailableError,
)
from .base import LLMBackend, LLMChunk, LLMResponse, Message, ToolSchema
from .config import LLMConfig, RouteEntry
from .health import RouteHealth
from .ollama import OllamaBackend
log = structlog.get_logger()
# Errors that are a statement about the *moment*: the same entry may work on a
# retry. Quota, auth and a missing model are not.
_RETRYABLE = (LLMConnectionError, LLMServerUnavailableError)
_ALL_FAILURES = (
LLMConnectionError,
LLMServerUnavailableError,
LLMQuotaExceededError,
LLMAuthError,
LLMModelNotFoundError,
)
def _route_from_context() -> str | None:
"""The route of the run this call belongs to, if any.
The agent sets it once per run next to ``current_model``; tools, the
compressor and ``AIHelper`` inherit it without every signature growing a
``route`` parameter. A call may still pass ``route=`` explicitly.
"""
try:
from navi.tools._internal.base import current_route
return current_route.get()
except Exception: # pragma: no cover - defensive, import cycle guard
return None
def _model_list(model: "list[str] | str | None") -> list[str]:
if isinstance(model, list):
return model if model else [""]
return [model] if model else [""]
class FallbackOllamaBackend(LLMBackend):
"""Ollama backend with route-based failover and background recovery."""
def __init__(self, config: LLMConfig, health: RouteHealth | None = None) -> None:
self._config = config
self._health = health or RouteHealth(config)
self._clients: dict[str, OllamaBackend] = {}
# Route id → index of the entry currently serving it.
self._active: dict[str, int] = {}
# ── lifecycle ────────────────────────────────────────────────────────
@property
def config(self) -> LLMConfig:
return self._config
@property
def health(self) -> RouteHealth:
return self._health
def start(self) -> None:
self._health.start()
async def stop(self) -> None:
await self._health.stop()
# ── route state ──────────────────────────────────────────────────────
def begin_turn(self, route_id: str | None) -> None:
"""Pin the entry this turn starts on.
Called once per user turn. If the primary is healthy again (the
background probe noticed), the turn starts there; otherwise it starts
on the best entry we know of.
"""
route = self._config.route(route_id)
if route is None:
return
chosen = 0
for i, entry in enumerate(route.entries):
if self._health.entry_usable(entry):
chosen = i
break
previous = self._active.get(route.id)
self._active[route.id] = chosen
if previous != chosen:
log.info(
"llm.route_active",
route=route.id,
entry=f"{route.entries[chosen].model}@{route.entries[chosen].target.name}",
previous=(
f"{route.entries[previous].model}@{route.entries[previous].target.name}"
if previous is not None and previous < len(route.entries)
else None
),
)
def active_entry(self, route_id: str | None) -> RouteEntry | None:
route = self._config.route(route_id)
if route is None:
return None
return route.entries[min(self._active.get(route.id, 0), len(route.entries) - 1)]
def reset_health(self) -> None:
"""Forget recorded failures and start every route at its primary."""
self._health.reset()
self._active.clear()
def serving_primary(
self, route_id: str | None, model: str | None, target: str | None = None
) -> bool:
"""Is this the route's primary entry (i.e. not a fallback)?
A route may serve one model from several hosts (the same weights on two
machines), so the target decides too when it is known: comparing model
names alone would call the second host's answer "primary".
"""
route = self._config.route(route_id)
if route is None or model is None:
return True
if route.primary.model != model:
return False
return target is None or route.primary.target.name == target
def active_window(self, route_id: str | None) -> int | None:
"""Context window of the entry that will serve the next call."""
entry = self.active_entry(route_id)
return entry.context_window if entry is not None else None
def active_report(self) -> dict:
"""What the engine is currently doing, for /health and logs."""
out = {}
for route_id, route in self._config.routes.items():
entry = self.active_entry(route_id)
if entry is None:
continue
out[route_id] = {
"model": entry.model,
"target": entry.target.name,
"primary": entry is route.primary,
"primary_model": route.primary.model,
}
return out
def _remember_active(self, entry: RouteEntry) -> None:
for route_id, route in self._config.routes.items():
for i, candidate in enumerate(route.entries):
if candidate == entry:
self._active[route_id] = i
return
# ── chain resolution ─────────────────────────────────────────────────
def _entries(
self, route: str | None, model: "list[str] | str | None"
) -> tuple[str | None, list[RouteEntry]]:
route_id = route or _route_from_context()
route_obj = self._config.route(route_id)
if route_obj is not None:
entries = list(route_obj.entries)
start = self._active.get(route_obj.id, 0)
if 0 < start < len(entries):
entries = entries[start:] + entries[:start]
return route_obj.id, entries
# No route: a bare model list, one name per candidate. Resolve each
# name to the target the config says serves it, so a mixed list
# (cloud + local names) no longer tries every name on every host.
entries = [
self._config.entry_for_model(name) or self._config.fallback_entry(name)
for name in _model_list(model)
]
return None, entries
def _order_usable(self, entries: list[RouteEntry]) -> list[RouteEntry]:
"""Healthy entries first, in chain order; unhealthy ones last.
Unhealthy entries are kept as a last resort: a stale health record is
better than refusing to call the LLM at all.
"""
healthy = [e for e in entries if self._health.entry_usable(e)]
degraded = [e for e in entries if not self._health.entry_usable(e)]
return healthy + degraded
def _client(self, entry: RouteEntry) -> OllamaBackend:
target = entry.target
if target.name not in self._clients:
ollama_http_timeout = max(
settings.ollama_request_timeout,
settings.llm_complete_timeout,
settings.llm_stream_first_chunk_timeout,
)
self._clients[target.name] = OllamaBackend(
model="", host=target.host, api_key=target.api_key,
timeout=ollama_http_timeout,
)
return self._clients[target.name]
@staticmethod
def _num_ctx(caller: int | None, entry: RouteEntry) -> int | None:
"""Never ask for more context than the entry actually serves."""
if entry.context_window is None:
return caller
if caller is None:
return entry.context_window
return min(caller, entry.context_window)
# ── complete ─────────────────────────────────────────────────────────
async def complete(
self,
messages: list[Message],
tools: list[ToolSchema] | None = None,
temperature: float = 0.7,
model: "list[str] | str | None" = None,
think: bool | None = None,
max_tokens: int | None = None,
top_k: int | None = None,
top_p: float | None = None,
num_thread: int | None = None,
num_ctx: int | None = None,
route: str | None = None,
) -> LLMResponse:
_, entries = self._entries(route, model)
if not entries:
raise LLMBackendError("No LLM entries configured")
last_err: Exception = LLMBackendError("No LLM entries configured")
ordered = self._order_usable(entries)
for i, entry in enumerate(ordered):
attempts = (
max(1, settings.llm_complete_retries) if i == len(ordered) - 1 else 1
)
for attempt in range(attempts):
try:
resp = await self._client(entry).complete(
messages, tools=tools, temperature=temperature,
model=entry.model, think=think, max_tokens=max_tokens,
top_k=top_k, top_p=top_p, num_thread=num_thread,
num_ctx=self._num_ctx(num_ctx, entry),
extra_options=entry.options or None,
)
except _ALL_FAILURES as e:
last_err = e
self._health.report_failure(entry, e)
if isinstance(e, _RETRYABLE) and attempt + 1 < attempts:
log.warning(
"llm.retry_same_entry", target=entry.target.name,
model=entry.model, attempt=attempt + 1, error=str(e),
)
await asyncio.sleep(settings.llm_retry_backoff_sec * (attempt + 1))
continue
break
else:
self._health.report_success(entry)
self._remember_active(entry)
if resp.model is None:
resp.model = entry.model
resp.target = entry.target.name
return resp
raise LLMBackendError(f"All LLM entries exhausted: {last_err}") from last_err
# ── embed ────────────────────────────────────────────────────────────
async def embed(
self,
texts: list[str],
model: "list[str] | str | None" = None,
route: str | None = None,
) -> list[list[float]]:
_, entries = self._entries(route, model)
if not entries:
raise LLMBackendError("No LLM entries configured")
last_err: Exception = LLMBackendError("No LLM entries configured")
for entry in self._order_usable(entries):
try:
vectors = await self._client(entry).embed(texts, model=entry.model)
except _ALL_FAILURES as e:
last_err = e
self._health.report_failure(entry, e)
continue
self._health.report_success(entry)
return vectors
raise LLMBackendError(f"All LLM entries exhausted: {last_err}") from last_err
# ── stream ───────────────────────────────────────────────────────────
async def stream_complete(
self,
messages: list[Message],
tools: list[ToolSchema] | None = None,
temperature: float = 0.7,
model: "list[str] | str | None" = None,
think: bool | None = None,
top_k: int | None = None,
top_p: float | None = None,
num_thread: int | None = None,
num_ctx: int | None = None,
route: str | None = None,
) -> AsyncGenerator[LLMChunk, None]:
route_id, entries = self._entries(route, model)
if not entries:
raise LLMBackendError("No LLM entries configured")
last_err: Exception = LLMBackendError("No LLM entries configured")
ordered = self._order_usable(entries)
for i, entry in enumerate(ordered):
attempts = (
max(1, settings.llm_complete_retries) if i == len(ordered) - 1 else 1
)
for attempt in range(attempts):
gen = None
first: LLMChunk | None = None
err: Exception | None = None
try:
gen = self._client(entry).stream_complete(
messages, tools=tools, temperature=temperature,
model=entry.model, think=think, top_k=top_k, top_p=top_p,
num_thread=num_thread,
num_ctx=self._num_ctx(num_ctx, entry),
extra_options=entry.options or None,
)
first = await asyncio.wait_for(
gen.__anext__(), timeout=settings.llm_stream_first_chunk_timeout
)
except StopAsyncIteration:
err = LLMModelNotFoundError(f"Empty stream from {entry.model}")
except asyncio.TimeoutError:
err = LLMConnectionError(
"First-chunk timeout after "
f"{settings.llm_stream_first_chunk_timeout}s"
)
except asyncio.CancelledError:
raise
except _ALL_FAILURES as e:
err = e
except LLMBackendError:
raise
if err is None and first is not None:
# Success: the first chunk is in hand, so nothing below can
# move the answer to another host without losing text.
self._health.report_success(entry)
self._remember_active(entry)
if first.model is None or first.target is None:
first = first.model_copy(
update={
"model": first.model or entry.model,
"target": first.target or entry.target.name,
"is_fallback": not self.serving_primary(
route_id, entry.model, entry.target.name
),
}
)
yield first
async for chunk in gen:
yield chunk
return
# The attempt failed before producing anything — release the
# underlying connection before trying the next entry.
if gen is not None:
try:
await gen.aclose()
except Exception:
pass
err = err or LLMBackendError(f"{entry.model} produced no chunk")
last_err = err
self._health.report_failure(entry, err)
if isinstance(err, _RETRYABLE) and attempt + 1 < attempts:
log.warning(
"llm.retry_same_entry", target=entry.target.name,
model=entry.model, attempt=attempt + 1, error=str(err),
)
await asyncio.sleep(settings.llm_retry_backoff_sec * (attempt + 1))
continue
log.warning(
"llm.entry_failed", target=entry.target.name,
model=entry.model, error=str(err),
)
break
raise LLMBackendError(f"All LLM entries exhausted: {last_err}") from last_err
