Files
ai-novel-engine/core/runtime/orchestration.py
T
2026-04-06 11:49:28 +02:00

150 lines
4.8 KiB
Python

from __future__ import annotations
from dataclasses import dataclass
import json
from typing import Any, Callable
from core.runtime.config import openai_base_url_for_model, runtime_probe_profile
from core.runtime.health import current_apple_model, probe_runtime_health, runtime_model_ids, wait_for_expected_apple_model
from core.runtime.profiles import runtime_probe_name
JsonFetcher = Callable[[str, float], Any]
@dataclass(frozen=True)
class RuntimeExecutionPlan:
model: str
openai_base_url: str
timeout_seconds: int
requires_native_ollama_preflight: bool
native_preflight_timeout_seconds: float
probe_profile_name: str
@dataclass(frozen=True)
class RuntimeCheckpointSignals:
core_health_ok: bool
apple_model_active: str | None
ollama_openai_runtime_ready: bool
def runtime_timeout_for_model(model: str, *, smoke_timeout_seconds: int) -> int:
if model.startswith("apple-coreml:"):
return max(600, smoke_timeout_seconds)
return max(120, smoke_timeout_seconds)
def build_runtime_execution_plan(
model: str,
*,
core_base_url: str,
ollama_runtime: str,
ollama_openai_base_url: str,
smoke_timeout_seconds: int,
) -> RuntimeExecutionPlan:
openai_base_url = openai_base_url_for_model(
model,
core_base_url=core_base_url,
ollama_runtime=ollama_runtime,
ollama_openai_base_url=ollama_openai_base_url,
)
timeout_seconds = runtime_timeout_for_model(model, smoke_timeout_seconds=smoke_timeout_seconds)
requires_native_preflight = model.startswith("ollama:") and ollama_runtime == "native"
probe_kind = "ollama_openai" if openai_base_url == ollama_openai_base_url and model.startswith("ollama:") else "core"
return RuntimeExecutionPlan(
model=model,
openai_base_url=openai_base_url,
timeout_seconds=timeout_seconds,
requires_native_ollama_preflight=requires_native_preflight,
native_preflight_timeout_seconds=min(45.0, float(timeout_seconds)),
probe_profile_name=runtime_probe_name(probe_kind),
)
def collect_checkpoint_runtime_signals(
model: str,
*,
core_base_url: str,
apple_runtime_url: str,
apple_model_ready_timeout_seconds: float,
apple_model_poll_interval_seconds: float,
ollama_runtime: str,
ollama_openai_base_url: str,
json_fetcher: JsonFetcher,
) -> RuntimeCheckpointSignals:
core_profile = runtime_probe_profile(core_base_url, name=runtime_probe_name("core"))
core_health = probe_runtime_health(
core_profile,
health_fetcher=json_fetcher,
json_fetcher=json_fetcher,
)
apple_model = None
if model.startswith("apple-coreml:"):
apple_model = wait_for_expected_apple_model(
apple_runtime_url,
model.split(":", 1)[1],
json_fetcher=json_fetcher,
timeout_seconds=apple_model_ready_timeout_seconds,
poll_interval_seconds=apple_model_poll_interval_seconds,
)
ollama_ready = False
if model.startswith("ollama:") and ollama_runtime == "openai_compatible":
ollama_ready = model in openai_runtime_model_ids(
ollama_openai_base_url,
json_fetcher=json_fetcher,
profile_name=runtime_probe_name("ollama_openai"),
)
return RuntimeCheckpointSignals(
core_health_ok=core_health.ok,
apple_model_active=apple_model,
ollama_openai_runtime_ready=ollama_ready,
)
def openai_runtime_model_ids(base_url: str, *, json_fetcher: JsonFetcher, profile_name: str) -> set[str]:
profile = runtime_probe_profile(base_url, name=profile_name)
health = probe_runtime_health(
profile,
health_fetcher=json_fetcher,
json_fetcher=json_fetcher,
)
if not health.ok or not health.available_models:
return set()
try:
model_ids = runtime_model_ids(profile, json_fetcher=json_fetcher)
except (OSError, json.JSONDecodeError, ValueError):
return set(health.available_models)
return model_ids or set(health.available_models)
def missing_ollama_models(
required_models: list[str],
*,
tags_url: str,
json_fetcher: JsonFetcher,
) -> list[str]:
try:
payload = json_fetcher(tags_url, 10.0)
except (OSError, json.JSONDecodeError, ValueError):
return []
models = payload.get("models") if isinstance(payload, dict) else None
if not isinstance(models, list):
return []
names = {
str(item.get("name", "")).strip()
for item in models
if isinstance(item, dict) and str(item.get("name", "")).strip()
}
return [model for model in required_models if model not in names]
def read_current_apple_model(apple_runtime_url: str, *, json_fetcher: JsonFetcher) -> str | None:
return current_apple_model(
apple_runtime_url,
json_fetcher=json_fetcher,
)