235 lines
7.5 KiB
Python
235 lines
7.5 KiB
Python
"""Digest-locked Ollama candidates for the decision agent runtime."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import time
|
|
from dataclasses import asdict, dataclass
|
|
from typing import Any, Callable, Mapping
|
|
|
|
import requests
|
|
|
|
from services.ollama_service import (
|
|
OLLAMA_HOST_FALLBACK,
|
|
OLLAMA_HOST_PRIMARY,
|
|
OLLAMA_HOST_SECONDARY,
|
|
)
|
|
|
|
|
|
PRIMARY_MODEL = os.getenv("NEMOTRON_OLLAMA_MODEL", "qwen3:14b").strip()
|
|
PRIMARY_EXPECTED_DIGEST = os.getenv(
|
|
"NEMOTRON_OLLAMA_EXPECTED_DIGEST",
|
|
"bdbd181c33f2ed1b31c972991882db3cf4d192569092138a7d29e973cd9debe8",
|
|
).strip().lower()
|
|
FALLBACK_MODEL = os.getenv(
|
|
"NEMOTRON_OLLAMA_FALLBACK_MODEL",
|
|
"qwen3:8b",
|
|
).strip()
|
|
FALLBACK_EXPECTED_DIGEST = os.getenv(
|
|
"NEMOTRON_OLLAMA_FALLBACK_EXPECTED_DIGEST",
|
|
"500a1f067a9f782620b40bee6f7b0c89e17ae61f686b92c24933e4ca4b2b8b41",
|
|
).strip().lower()
|
|
NEMOTRON_FALLBACK_NUM_CTX = 4096
|
|
NEMOTRON_FALLBACK_NUM_PREDICT = 512
|
|
|
|
|
|
def _bounded_int_env(name: str, default: int, minimum: int, maximum: int) -> int:
|
|
try:
|
|
value = int(os.getenv(name, str(default)))
|
|
except (TypeError, ValueError):
|
|
value = default
|
|
return max(minimum, min(value, maximum))
|
|
|
|
|
|
IDENTITY_TIMEOUT_SEC = _bounded_int_env(
|
|
"NEMOTRON_MODEL_IDENTITY_TIMEOUT_SEC", 10, 2, 30
|
|
)
|
|
GCP_ATTEMPT_TIMEOUT_SEC = _bounded_int_env(
|
|
"NEMOTRON_GCP_ATTEMPT_TIMEOUT_SEC", 60, 30, 180
|
|
)
|
|
FALLBACK_ATTEMPT_TIMEOUT_SEC = _bounded_int_env(
|
|
"NEMOTRON_111_ATTEMPT_TIMEOUT_SEC", 45, 20, 120
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class NemotronRuntimeCandidate:
|
|
label: str
|
|
tier: str
|
|
host: str
|
|
model: str
|
|
expected_digest: str
|
|
request_timeout_sec: int
|
|
num_predict: int
|
|
num_ctx: int | None
|
|
|
|
@property
|
|
def is_fallback(self) -> bool:
|
|
return self.tier == "fallback"
|
|
|
|
def as_public_dict(self) -> dict[str, Any]:
|
|
payload = asdict(self)
|
|
payload["is_fallback"] = self.is_fallback
|
|
return payload
|
|
|
|
|
|
def build_nemotron_runtime_candidates(
|
|
*,
|
|
primary_model: str | None = None,
|
|
primary_expected_digest: str | None = None,
|
|
fallback_model: str | None = None,
|
|
fallback_expected_digest: str | None = None,
|
|
) -> tuple[NemotronRuntimeCandidate, ...]:
|
|
"""Return the approved GCP-first, 111-final candidate chain."""
|
|
resolved_primary_model = str(primary_model or PRIMARY_MODEL).strip()
|
|
resolved_primary_digest = str(
|
|
primary_expected_digest or PRIMARY_EXPECTED_DIGEST
|
|
).strip().lower()
|
|
resolved_fallback_model = str(fallback_model or FALLBACK_MODEL).strip()
|
|
resolved_fallback_digest = str(
|
|
fallback_expected_digest or FALLBACK_EXPECTED_DIGEST
|
|
).strip().lower()
|
|
return (
|
|
NemotronRuntimeCandidate(
|
|
label="gcp_primary",
|
|
tier="primary",
|
|
host=OLLAMA_HOST_PRIMARY.rstrip("/"),
|
|
model=resolved_primary_model,
|
|
expected_digest=resolved_primary_digest,
|
|
request_timeout_sec=GCP_ATTEMPT_TIMEOUT_SEC,
|
|
num_predict=2048,
|
|
num_ctx=None,
|
|
),
|
|
NemotronRuntimeCandidate(
|
|
label="gcp_secondary",
|
|
tier="secondary",
|
|
host=OLLAMA_HOST_SECONDARY.rstrip("/"),
|
|
model=resolved_primary_model,
|
|
expected_digest=resolved_primary_digest,
|
|
request_timeout_sec=GCP_ATTEMPT_TIMEOUT_SEC,
|
|
num_predict=2048,
|
|
num_ctx=None,
|
|
),
|
|
NemotronRuntimeCandidate(
|
|
label="ollama_111_fallback",
|
|
tier="fallback",
|
|
host=OLLAMA_HOST_FALLBACK.rstrip("/"),
|
|
model=resolved_fallback_model,
|
|
expected_digest=resolved_fallback_digest,
|
|
request_timeout_sec=FALLBACK_ATTEMPT_TIMEOUT_SEC,
|
|
num_predict=NEMOTRON_FALLBACK_NUM_PREDICT,
|
|
num_ctx=NEMOTRON_FALLBACK_NUM_CTX,
|
|
),
|
|
)
|
|
|
|
|
|
def inspect_nemotron_model_identity(
|
|
candidate: NemotronRuntimeCandidate,
|
|
*,
|
|
request_get: Callable[..., Any] | None = None,
|
|
timeout_sec: int | None = None,
|
|
) -> dict[str, Any]:
|
|
"""Verify an exact model digest without loading the model."""
|
|
started = time.monotonic()
|
|
getter = request_get or requests.get
|
|
timeout = max(0.1, min(float(timeout_sec or IDENTITY_TIMEOUT_SEC), 30.0))
|
|
try:
|
|
response = getter(f"{candidate.host}/api/tags", timeout=timeout)
|
|
response.raise_for_status()
|
|
payload = response.json()
|
|
if not isinstance(payload, Mapping):
|
|
raise ValueError("malformed_tags_payload")
|
|
models = payload.get("models", [])
|
|
if models is None:
|
|
models = []
|
|
if not isinstance(models, list):
|
|
raise ValueError("malformed_models_schema")
|
|
except Exception as exc:
|
|
return {
|
|
**candidate.as_public_dict(),
|
|
"ok": False,
|
|
"digest": None,
|
|
"digest_matches": False,
|
|
"elapsed_ms": round((time.monotonic() - started) * 1000),
|
|
"error": f"{type(exc).__name__}: {str(exc)[:180]}",
|
|
}
|
|
|
|
target = candidate.model
|
|
target_without_latest = target.removesuffix(":latest")
|
|
matched: Mapping[str, Any] = {}
|
|
if any(not isinstance(item, Mapping) for item in models):
|
|
return {
|
|
**candidate.as_public_dict(),
|
|
"ok": False,
|
|
"digest": None,
|
|
"digest_matches": False,
|
|
"parameter_size": "",
|
|
"quantization_level": "",
|
|
"elapsed_ms": round((time.monotonic() - started) * 1000),
|
|
"error": "malformed_models_schema",
|
|
}
|
|
if any(
|
|
"details" in item and not isinstance(item.get("details"), Mapping)
|
|
for item in models
|
|
):
|
|
return {
|
|
**candidate.as_public_dict(),
|
|
"ok": False,
|
|
"digest": None,
|
|
"digest_matches": False,
|
|
"parameter_size": "",
|
|
"quantization_level": "",
|
|
"elapsed_ms": round((time.monotonic() - started) * 1000),
|
|
"error": "malformed_model_details",
|
|
}
|
|
for item in models:
|
|
names = {
|
|
str(item.get("name") or "").strip(),
|
|
str(item.get("model") or "").strip(),
|
|
}
|
|
if target in names or target_without_latest in names:
|
|
matched = item
|
|
break
|
|
digest = str(matched.get("digest") or "").strip().lower()
|
|
raw_details = matched.get("details", {}) if matched else {}
|
|
malformed_details = bool(matched) and not isinstance(raw_details, Mapping)
|
|
details = raw_details if isinstance(raw_details, Mapping) else {}
|
|
digest_matches = (
|
|
bool(digest)
|
|
and digest == candidate.expected_digest
|
|
and not malformed_details
|
|
)
|
|
error = None
|
|
if not matched:
|
|
error = f"model_not_found:{target}"
|
|
elif malformed_details:
|
|
error = "malformed_model_details"
|
|
elif not digest_matches:
|
|
error = "model_digest_mismatch"
|
|
return {
|
|
**candidate.as_public_dict(),
|
|
"ok": bool(matched) and digest_matches,
|
|
"digest": digest or None,
|
|
"digest_matches": digest_matches,
|
|
"parameter_size": str(details.get("parameter_size") or ""),
|
|
"quantization_level": str(details.get("quantization_level") or ""),
|
|
"elapsed_ms": round((time.monotonic() - started) * 1000),
|
|
"error": error,
|
|
}
|
|
|
|
|
|
__all__ = [
|
|
"FALLBACK_ATTEMPT_TIMEOUT_SEC",
|
|
"FALLBACK_EXPECTED_DIGEST",
|
|
"FALLBACK_MODEL",
|
|
"GCP_ATTEMPT_TIMEOUT_SEC",
|
|
"IDENTITY_TIMEOUT_SEC",
|
|
"NemotronRuntimeCandidate",
|
|
"NEMOTRON_FALLBACK_NUM_CTX",
|
|
"NEMOTRON_FALLBACK_NUM_PREDICT",
|
|
"PRIMARY_EXPECTED_DIGEST",
|
|
"PRIMARY_MODEL",
|
|
"build_nemotron_runtime_candidates",
|
|
"inspect_nemotron_model_identity",
|
|
]
|