102 lines
2.4 KiB
Python
102 lines
2.4 KiB
Python
"""
|
|
Ollama endpoint resolver for non-critical workload placement.
|
|
|
|
ADR-110 gives AWOOOI three Ollama endpoints. This resolver is intentionally
|
|
small: it chooses the preferred endpoint by workload class, while health-aware
|
|
failover remains owned by ollama_failover_manager.py.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from typing import Literal, Protocol
|
|
|
|
from src.core.config import settings
|
|
|
|
OllamaWorkloadType = Literal[
|
|
"interactive",
|
|
"healthcheck",
|
|
"batch",
|
|
"embedding",
|
|
"rag",
|
|
"code_review",
|
|
"shadow",
|
|
"canary",
|
|
"local_required",
|
|
"privacy_sensitive",
|
|
"dr",
|
|
]
|
|
|
|
_GCP_B_PREFERRED_WORKLOADS = {
|
|
"batch",
|
|
"embedding",
|
|
"rag",
|
|
"code_review",
|
|
"shadow",
|
|
"canary",
|
|
}
|
|
|
|
_LOCAL_PREFERRED_WORKLOADS = {
|
|
"local_required",
|
|
"privacy_sensitive",
|
|
"dr",
|
|
}
|
|
|
|
|
|
class _OllamaSettings(Protocol):
|
|
OLLAMA_URL: str
|
|
OLLAMA_SECONDARY_URL: str
|
|
OLLAMA_FALLBACK_URL: str
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class OllamaEndpointSelection:
|
|
url: str
|
|
provider_name: str
|
|
workload_type: OllamaWorkloadType
|
|
reason: str
|
|
|
|
|
|
def resolve_ollama_selection(
|
|
workload_type: OllamaWorkloadType = "interactive",
|
|
*,
|
|
config: _OllamaSettings | None = None,
|
|
) -> OllamaEndpointSelection:
|
|
"""Return the preferred Ollama endpoint for a workload class."""
|
|
cfg = config or settings
|
|
primary = cfg.OLLAMA_URL
|
|
secondary = cfg.OLLAMA_SECONDARY_URL
|
|
fallback = cfg.OLLAMA_FALLBACK_URL
|
|
|
|
if workload_type in _GCP_B_PREFERRED_WORKLOADS and secondary:
|
|
return OllamaEndpointSelection(
|
|
url=secondary,
|
|
provider_name="ollama_gcp_b",
|
|
workload_type=workload_type,
|
|
reason="gcp_b_batch_lane",
|
|
)
|
|
|
|
if workload_type in _LOCAL_PREFERRED_WORKLOADS and fallback:
|
|
return OllamaEndpointSelection(
|
|
url=fallback,
|
|
provider_name="ollama_local",
|
|
workload_type=workload_type,
|
|
reason="local_privacy_or_dr_lane",
|
|
)
|
|
|
|
return OllamaEndpointSelection(
|
|
url=primary,
|
|
provider_name="ollama_gcp_a",
|
|
workload_type=workload_type,
|
|
reason="primary_interactive_lane",
|
|
)
|
|
|
|
|
|
def resolve_ollama_endpoint(
|
|
workload_type: OllamaWorkloadType = "interactive",
|
|
*,
|
|
config: _OllamaSettings | None = None,
|
|
) -> str:
|
|
"""Return only the preferred Ollama base URL."""
|
|
return resolve_ollama_selection(workload_type, config=config).url
|