Compare commits

...
9 changed files with 192 additions and 40 deletions
+4 -1
View File
@@ -466,7 +466,10 @@ export function InputArea() {
}
const totalMs = Date.now() - startTime;
const _CLOUD_PREFIXES = ['gpt-', 'o1-', 'o3-', 'o4-', 'claude-', 'gemini-', 'openrouter/', 'MiniMax-', 'chatgpt-'];
const engineLabel = _CLOUD_PREFIXES.some(p => selectedModel.startsWith(p)) ? 'cloud' : 'ollama';
const selectedOwner = useAppStore.getState().models.find((m) => m.id === selectedModel)?.owned_by;
const engineLabel = selectedOwner === 'litellm'
? 'litellm'
: _CLOUD_PREFIXES.some(p => selectedModel.startsWith(p)) ? 'cloud' : 'ollama';
const telemetry: MessageTelemetry = {
engine: engineLabel,
model_id: selectedModel,
+23 -16
View File
@@ -143,7 +143,7 @@ export function CommandPalette() {
}
}, [pullSuccess]);
const handleSelect = async (modelId: string) => {
const handleSelect = async (modelId: string, owner?: string) => {
const previousModel = selectedModel;
setSelectedModel(modelId);
setCommandPaletteOpen(false);
@@ -153,7 +153,7 @@ export function CommandPalette() {
setModelLoading(true);
addLogEntry({ timestamp: Date.now(), level: 'info', category: 'model', message: `Switching to ${modelId}...` });
try {
await preloadModel(modelId);
await preloadModel(modelId, owner);
addLogEntry({ timestamp: Date.now(), level: 'info', category: 'model', message: `${modelId} loaded` });
} catch (e: any) {
addLogEntry({ timestamp: Date.now(), level: 'error', category: 'model', message: `Failed to load ${modelId}: ${e.message}` });
@@ -255,7 +255,8 @@ export function CommandPalette() {
setSelectedIdx((i) => Math.max(i - 1, 0));
} else if (e.key === 'Enter' && tab === 'installed' && filtered.length > 0) {
e.preventDefault();
handleSelect((filtered[selectedIdx] as any).id);
const model = filtered[selectedIdx] as (typeof models)[number];
handleSelect(model.id, model.owned_by);
}
};
@@ -365,11 +366,15 @@ export function CommandPalette() {
onMouseEnter={() => setSelectedIdx(idx)}
>
<button
onClick={() => handleSelect(model.id)}
onClick={() => handleSelect(model.id, model.owned_by)}
className="flex items-center gap-3 flex-1 min-w-0 text-left cursor-pointer"
style={{ background: 'none', border: 'none', padding: 0 }}
>
<Cpu size={16} style={{ color: isActive ? 'var(--color-accent)' : 'var(--color-text-tertiary)' }} />
{model.owned_by === 'litellm' ? (
<Cloud size={16} style={{ color: isActive ? 'var(--color-accent)' : 'var(--color-text-tertiary)' }} />
) : (
<Cpu size={16} style={{ color: isActive ? 'var(--color-accent)' : 'var(--color-text-tertiary)' }} />
)}
<div className="flex-1 min-w-0">
<div className="text-sm truncate" style={{ color: isActive ? 'var(--color-accent)' : 'var(--color-text)', fontWeight: isActive ? 500 : 400 }}>
{model.id}
@@ -381,17 +386,19 @@ export function CommandPalette() {
</span>
)}
</button>
<button
onClick={() => handleDelete(model.id)}
disabled={isDeleting}
className="p-1 rounded transition-colors cursor-pointer"
style={{ color: 'var(--color-text-tertiary)', opacity: 0 }}
title="Delete model"
onMouseEnter={(e) => { e.currentTarget.style.opacity = '1'; e.currentTarget.style.color = 'var(--color-error)'; }}
onMouseLeave={(e) => { e.currentTarget.style.opacity = '0'; e.currentTarget.style.color = 'var(--color-text-tertiary)'; }}
>
{isDeleting ? <Loader2 size={14} className="animate-spin" /> : <Trash2 size={14} />}
</button>
{model.owned_by !== 'litellm' && (
<button
onClick={() => handleDelete(model.id)}
disabled={isDeleting}
className="p-1 rounded transition-colors cursor-pointer"
style={{ color: 'var(--color-text-tertiary)', opacity: 0 }}
title="Delete model"
onMouseEnter={(e) => { e.currentTarget.style.opacity = '1'; e.currentTarget.style.color = 'var(--color-error)'; }}
onMouseLeave={(e) => { e.currentTarget.style.opacity = '0'; e.currentTarget.style.color = 'var(--color-text-tertiary)'; }}
>
{isDeleting ? <Loader2 size={14} className="animate-spin" /> : <Trash2 size={14} />}
</button>
)}
</div>
);
})
+2 -2
View File
@@ -218,9 +218,9 @@ export async function deleteModel(modelName: string): Promise<void> {
const _CLOUD_PREFIXES = ['gpt-', 'o1-', 'o3-', 'o4-', 'claude-', 'gemini-', 'openrouter/'];
export async function preloadModel(modelName: string): Promise<void> {
export async function preloadModel(modelName: string, owner?: string): Promise<void> {
// Cloud models don't need Ollama preloading
if (_CLOUD_PREFIXES.some(p => modelName.startsWith(p))) {
if (owner === 'litellm' || _CLOUD_PREFIXES.some(p => modelName.startsWith(p))) {
return;
}
// Trigger Ollama to load the model into memory (empty prompt, no generation).
+6
View File
@@ -35,6 +35,12 @@ def _make_engine(key: str, config: JarvisConfig) -> InferenceEngine:
"""Instantiate a registered engine with the appropriate config host."""
cls = EngineRegistry.get(key)
# LiteLLM cannot enumerate every model supported by every provider. Its
# list_models() contract therefore advertises the configured default
# model, which must be supplied when discovery constructs the engine.
if key == "litellm":
return cls(default_model=config.intelligence.default_model or None)
# gemma_cpp: pass config fields instead of host
if key == "gemma_cpp":
cfg = config.engine.gemma_cpp
+13 -2
View File
@@ -26,16 +26,19 @@ class MultiEngine(InferenceEngine):
def __init__(self, engines: list[tuple[str, InferenceEngine]]) -> None:
self._engines = engines
self._model_map: Dict[str, InferenceEngine] = {}
self._model_key_map: Dict[str, str] = {}
self._refresh_map()
def _refresh_map(self) -> None:
self._model_map.clear()
for _key, engine in self._engines:
self._model_key_map.clear()
for key, engine in self._engines:
try:
for model_id in engine.list_models():
self._model_map[model_id] = engine
self._model_key_map[model_id] = key
except Exception as exc:
logger.debug("Failed to list models for %s: %s", _key, exc)
logger.debug("Failed to list models for %s: %s", key, exc)
_CLOUD_PREFIXES = ("gpt-", "o1-", "o3-", "o4-", "claude-", "gemini-", "openrouter/")
@@ -117,6 +120,14 @@ class MultiEngine(InferenceEngine):
self._refresh_map()
return list(self._model_map.keys())
def engine_key_for(self, model: str) -> str | None:
"""Return the registry key of the engine advertising *model*."""
key = self._model_key_map.get(model)
if key is not None:
return key
self._refresh_map()
return self._model_key_map.get(model)
def health(self) -> bool:
return any(engine.health() for _key, engine in self._engines)
+62 -18
View File
@@ -336,6 +336,34 @@ def _remember_exchange(
)
def _engine_key_for_model(engine: Any, model: str) -> str | None:
"""Resolve the engine that advertised *model* through wrapper layers."""
from openjarvis.engine.multi import MultiEngine
from openjarvis.security.guardrails import GuardrailsEngine
from openjarvis.telemetry.instrumented_engine import InstrumentedEngine
current = engine
while current is not None:
if isinstance(current, MultiEngine):
return current.engine_key_for(model)
if isinstance(current, InstrumentedEngine):
current = current._inner
continue
if isinstance(current, GuardrailsEngine):
current = current._engine
continue
engine_id = getattr(current, "engine_id", None)
return engine_id if isinstance(engine_id, str) else None
return None
def _uses_direct_cloud_router(engine: Any, model: str) -> bool:
"""Whether *model* should bypass the configured engine for direct cloud."""
from openjarvis.server.cloud_router import is_cloud_model
return is_cloud_model(model) and _engine_key_for_model(engine, model) != "litellm"
def _handle_direct(
engine,
model: str,
@@ -541,12 +569,13 @@ async def _handle_stream_tools(
tool_calls) — identical to the prior plain-stream behaviour, so this never
regresses non-tool-capable engines.
"""
from openjarvis.server.cloud_router import is_cloud_model
messages = _to_messages(req.messages)
messages = _ensure_identity_prompt(messages, app_config)
chunk_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
use_cloud = is_cloud_model(model)
use_cloud = _uses_direct_cloud_router(engine, model)
telemetry_engine = (
"cloud" if use_cloud else (_engine_key_for_model(engine, model) or "ollama")
)
query_text = ""
for _m in reversed(req.messages):
if _m.role == "user" and _m.content:
@@ -626,7 +655,7 @@ async def _handle_stream_tools(
# Tag the finish chunk with the engine label, matching _handle_stream
# so UI/telemetry consumers see the same field on the tools path.
finish_dict.setdefault("telemetry", {})
finish_dict["telemetry"]["engine"] = "cloud" if use_cloud else "ollama"
finish_dict["telemetry"]["engine"] = telemetry_engine
if complexity_info is not None:
finish_dict["complexity"] = complexity_info.model_dump()
yield f"data: {_json.dumps(finish_dict)}\n\n"
@@ -668,11 +697,7 @@ async def _handle_stream(
"""
import time
from openjarvis.server.cloud_router import (
is_cloud_model,
stream_cloud,
stream_local,
)
from openjarvis.server.cloud_router import stream_cloud, stream_local
messages = _to_messages(req.messages)
messages = _ensure_identity_prompt(messages, app_config)
@@ -687,7 +712,10 @@ async def _handle_stream(
# Route directly to the right backend — bypasses engine routing entirely
# so broken MultiEngine state can never misdirect requests.
use_cloud = is_cloud_model(model)
use_cloud = _uses_direct_cloud_router(engine, model)
telemetry_engine = (
"cloud" if use_cloud else (_engine_key_for_model(engine, model) or "ollama")
)
async def generate():
started_at = time.time()
@@ -792,7 +820,7 @@ async def _handle_stream(
query=query_text,
result=full_content,
model=model,
engine="cloud" if use_cloud else "ollama",
engine=telemetry_engine,
started_at=started_at,
ended_at=time.time(),
)
@@ -825,7 +853,7 @@ async def _handle_stream(
# We use the routing decision (use_cloud) directly rather than
# unwrapping the engine chain, which can be in a broken state.
finish_dict.setdefault("telemetry", {})
finish_dict["telemetry"]["engine"] = "cloud" if use_cloud else "ollama"
finish_dict["telemetry"]["engine"] = telemetry_engine
if complexity_info is not None:
finish_dict["complexity"] = complexity_info.model_dump()
@@ -842,24 +870,40 @@ async def _handle_stream(
@router.get("/v1/models")
async def list_models(request: Request) -> ModelListResponse:
"""List locally installed models (Ollama).
"""List selectable engine models for the installed-model picker.
Cloud models are not included here — they live in the Cloud Models tab
of the UI and are selected there, not from this endpoint.
Direct cloud models live in the Cloud Models tab. Models advertised by a
configured LiteLLM engine remain here because LiteLLM owns their routing
and may use provider-qualified IDs that resemble OpenRouter IDs.
"""
from openjarvis.server.cloud_router import is_cloud_model, list_local_models
# Prefer engine.list_models() so mock engines work in tests.
# Filter out any cloud model IDs that may appear via MultiEngine.
# Filter out direct-cloud model IDs that may appear via MultiEngine, but
# retain provider-qualified IDs owned by the configured LiteLLM engine.
# Fall back to direct Ollama query only when the engine returns nothing.
engine = request.app.state.engine
all_ids = await asyncio.to_thread(engine.list_models)
model_ids = [m for m in all_ids if not is_cloud_model(m)]
model_ids = [
m
for m in all_ids
if not is_cloud_model(m) or _engine_key_for_model(engine, m) == "litellm"
]
if not model_ids:
model_ids = await list_local_models()
return ModelListResponse(
data=[ModelObject(id=mid) for mid in model_ids],
data=[
ModelObject(
id=mid,
owned_by=(
"litellm"
if _engine_key_for_model(engine, mid) == "litellm"
else "openjarvis"
),
)
for mid in model_ids
],
)
+20
View File
@@ -8,10 +8,12 @@ from openjarvis.core.config import JarvisConfig
from openjarvis.core.registry import EngineRegistry
from openjarvis.engine._base import InferenceEngine
from openjarvis.engine._discovery import (
_make_engine,
discover_engines,
discover_models,
get_engine,
)
from openjarvis.engine.litellm import LiteLLMEngine
class _FakeEngine(InferenceEngine):
@@ -131,6 +133,24 @@ class TestDiscoverModels:
assert result == {"ollama": ["m1", "m2"], "vllm": ["m3"]}
class TestLiteLLMDiscovery:
def test_configured_default_model_is_advertised(self) -> None:
"""Regression for #713: discovery must configure LiteLLM's model.
LiteLLM cannot enumerate every model supported by every provider, so
``LiteLLMEngine.list_models()`` advertises the configured default
model. Dropping that value while constructing the engine leaves the
API and Web UI with an empty model list.
"""
cfg = JarvisConfig()
cfg.intelligence.default_model = "groq/llama-3.3-70b-versatile"
EngineRegistry.register_value("litellm", LiteLLMEngine)
engine = _make_engine("litellm", cfg)
assert engine.list_models() == ["groq/llama-3.3-70b-versatile"]
class TestGetEngine:
def test_fallback_when_default_unhealthy(self) -> None:
_reg("bad", "bad")
+3
View File
@@ -120,6 +120,9 @@ async def test_multi_routes_stream_full_by_model():
engine_b.list_models = lambda: ["model-b"]
multi = MultiEngine([("a", engine_a), ("b", engine_b)])
assert multi.engine_key_for("model-a") == "a"
assert multi.engine_key_for("model-b") == "b"
assert multi.engine_key_for("missing") is None
# Route to engine A
result_a = []
+59 -1
View File
@@ -3,7 +3,7 @@
from __future__ import annotations
import json
from unittest.mock import MagicMock
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@@ -883,6 +883,64 @@ class TestModelsEndpoint:
data = resp.json()
assert len(data["data"]) == 3
def test_configured_litellm_model_is_listed(self):
"""Regression for #713: LiteLLM models must reach the Web UI."""
model = "groq/llama-3.3-70b-versatile"
engine = _make_engine(models=[model])
engine.engine_id = "litellm"
app = create_app(
engine,
model,
engine_name="litellm",
config=_test_config(),
)
with patch(
"openjarvis.server.cloud_router.list_local_models",
new_callable=AsyncMock,
) as list_local_models:
list_local_models.return_value = []
client = TestClient(app)
resp = client.get("/v1/models")
assert resp.status_code == 200
assert [item["id"] for item in resp.json()["data"]] == [model]
assert resp.json()["data"][0]["owned_by"] == "litellm"
def test_litellm_provider_model_streams_through_active_engine(self):
"""A LiteLLM ``provider/model`` ID must not bypass its engine."""
model = "groq/llama-3.3-70b-versatile"
engine = _make_engine(models=[model])
engine.engine_id = "litellm"
app = create_app(
engine,
model,
engine_name="litellm",
config=_test_config(),
)
async def direct_cloud_tokens():
yield "wrong backend"
with patch(
"openjarvis.server.cloud_router.stream_cloud",
return_value=direct_cloud_tokens(),
) as stream_cloud:
client = TestClient(app)
resp = client.post(
"/v1/chat/completions",
json={
"model": model,
"messages": [{"role": "user", "content": "hello"}],
"stream": True,
},
)
assert resp.status_code == 200
stream_cloud.assert_not_called()
assert "Hello" in resp.text
assert '"engine": "litellm"' in resp.text
# ---------------------------------------------------------------------------
# Health endpoint tests