mirror of
https://github.com/open-jarvis/OpenJarvis.git
synced 2026-07-30 19:02:16 +00:00
Add four new OpenAI-compatible inference engines across the full stack: Rust backend: - Config structs with serde defaults (ExoEngineConfig, NexaEngineConfig, UzuEngineConfig, AppleFmEngineConfig) in openjarvis-core - Factory constructors on OpenAICompatEngine (exo, nexa, uzu, apple_fm) - Engine enum variants with delegate_engine! macro dispatch - Discovery and resolution in get_engine_static() - PyO3 bridge support in PyEngine Python frontend: - Data-driven engine class registration in openai_compat_engines.py - Config dataclasses with backward-compat host properties - Discovery host map entries - DEFAULT_SEARCH_SPACE updated with new engine backends - Apple FM shim (FastAPI wrapper for python-apple-fm-sdk, macOS only) Config & tests: - TOML config sections for all four engines - Unit tests for factory methods, enum variants, and search space Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
108 lines
3.2 KiB
Python
108 lines
3.2 KiB
Python
"""Engine discovery — probe running engines and aggregate available models."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any, Dict, List, Tuple
|
|
|
|
from openjarvis.core.config import JarvisConfig
|
|
from openjarvis.core.registry import EngineRegistry
|
|
from openjarvis.engine._base import InferenceEngine
|
|
|
|
# Map registry keys to config host attribute (None = no host arg)
|
|
_HOST_MAP: Dict[str, str | None] = {
|
|
"ollama": "ollama_host",
|
|
"vllm": "vllm_host",
|
|
"llamacpp": "llamacpp_host",
|
|
"sglang": "sglang_host",
|
|
"mlx": "mlx_host",
|
|
"lmstudio": "lmstudio_host",
|
|
"exo": "exo_host",
|
|
"nexa": "nexa_host",
|
|
"uzu": "uzu_host",
|
|
"apple_fm": "apple_fm_host",
|
|
"cloud": None,
|
|
"litellm": None,
|
|
}
|
|
|
|
|
|
def _make_engine(key: str, config: JarvisConfig) -> InferenceEngine:
|
|
"""Instantiate a registered engine with the appropriate config host."""
|
|
cls = EngineRegistry.get(key)
|
|
host_attr = _HOST_MAP.get(key)
|
|
if host_attr is not None:
|
|
host = getattr(config.engine, host_attr, None)
|
|
if host:
|
|
return cls(host=host)
|
|
return cls()
|
|
|
|
|
|
def discover_engines(config: JarvisConfig) -> List[Tuple[str, InferenceEngine]]:
|
|
"""Probe registered engines and return ``[(key, instance)]`` for healthy ones.
|
|
|
|
Results are sorted with the config default engine first.
|
|
"""
|
|
healthy: List[Tuple[str, InferenceEngine]] = []
|
|
for key in EngineRegistry.keys():
|
|
try:
|
|
engine = _make_engine(key, config)
|
|
if engine.health():
|
|
healthy.append((key, engine))
|
|
except Exception:
|
|
continue
|
|
|
|
default_key = config.engine.default
|
|
|
|
def sort_key(item: Tuple[str, Any]) -> Tuple[int, str]:
|
|
return (0 if item[0] == default_key else 1, item[0])
|
|
|
|
healthy.sort(key=sort_key)
|
|
return healthy
|
|
|
|
|
|
def discover_models(
|
|
engines: List[Tuple[str, InferenceEngine]],
|
|
) -> Dict[str, List[str]]:
|
|
"""Call ``list_models()`` on each engine and return a dict."""
|
|
result: Dict[str, List[str]] = {}
|
|
for key, engine in engines:
|
|
try:
|
|
result[key] = engine.list_models()
|
|
except Exception:
|
|
result[key] = []
|
|
return result
|
|
|
|
|
|
def get_engine(
|
|
config: JarvisConfig, engine_key: str | None = None
|
|
) -> Tuple[str, InferenceEngine] | None:
|
|
"""Get a specific engine by key, or the default with fallback.
|
|
|
|
Returns ``(key, engine_instance)`` or ``None`` if no engine is available.
|
|
"""
|
|
if engine_key:
|
|
if EngineRegistry.contains(engine_key):
|
|
try:
|
|
engine = _make_engine(engine_key, config)
|
|
if engine.health():
|
|
return (engine_key, engine)
|
|
except Exception:
|
|
pass
|
|
return None
|
|
|
|
# Try default first
|
|
default_key = config.engine.default
|
|
if EngineRegistry.contains(default_key):
|
|
try:
|
|
engine = _make_engine(default_key, config)
|
|
if engine.health():
|
|
return (default_key, engine)
|
|
except Exception:
|
|
pass
|
|
|
|
# Fallback to any healthy engine
|
|
healthy = discover_engines(config)
|
|
return healthy[0] if healthy else None
|
|
|
|
|
|
__all__ = ["discover_engines", "discover_models", "get_engine"]
|