Simplify model resolution and refactor tool discovery via MCPServer

- Remove --router CLI flag from `jarvis ask`; use config fallback chain
  (default_model -> first engine model -> fallback_model) instead
- Simplify SDK _resolve_model() to match the same fallback chain
- Reclassify OrchestratorSFTPolicy and OrchestratorGRPOPolicy as
  IntelligenceLearningPolicy (they fine-tune LM weights, not agent logic)
- Add MCPServer.get_tools() for SystemBuilder integration
- Refactor SystemBuilder._resolve_tools() to discover tools via MCPServer
  with dependency injection, replacing manual per-tool-name imports
- Update CLI router tests to test the new fallback chain

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Jon Saad-Falcon
2026-02-23 18:41:07 +00:00
co-authored by Claude Opus 4.6
parent 8d538cd1b0
commit c775d620af
7 changed files with 135 additions and 150 deletions
+8 -29
View File
@@ -18,8 +18,6 @@ from openjarvis.engine import (
get_engine,
)
from openjarvis.intelligence import (
HeuristicRouter,
build_routing_context,
merge_discovered_models,
register_builtin_models,
)
@@ -170,10 +168,6 @@ def _run_agent(
"--tools", "tool_names", default=None,
help="Comma-separated tool names to enable (e.g. calculator,think).",
)
@click.option(
"--router", "router_policy", default=None,
help="Router policy for model selection (heuristic, grpo).",
)
def ask(
query: tuple[str, ...],
model_name: str | None,
@@ -185,7 +179,6 @@ def ask(
no_context: bool,
agent_name: str | None,
tool_names: str | None,
router_policy: str | None,
) -> None:
"""Ask Jarvis a question."""
console = Console(stderr=True)
@@ -227,33 +220,19 @@ def ask(
for ek, model_ids in all_models.items():
merge_discovered_models(ek, model_ids)
# Route to model
# Resolve model via config fallback chain
if model_name is None:
policy_key = router_policy or config.learning.default_policy
from openjarvis.core.registry import RouterPolicyRegistry
from openjarvis.learning import ensure_registered as _ensure_learning
_ensure_learning()
if RouterPolicyRegistry.contains(policy_key):
router_cls = RouterPolicyRegistry.get(policy_key)
else:
router_cls = HeuristicRouter # fallback
router = router_cls(
available_models=all_models.get(engine_name, []),
default_model=config.intelligence.default_model,
fallback_model=config.intelligence.fallback_model,
)
ctx = build_routing_context(query_text)
model_name = router.select_model(ctx)
model_name = config.intelligence.default_model
if not model_name:
# Last resort: pick first available model from the engine
# Try first available from engine
engine_models = all_models.get(engine_name, [])
if engine_models:
model_name = engine_models[0]
else:
console.print("[red]No model available on engine.[/red]")
sys.exit(1)
if not model_name:
model_name = config.intelligence.fallback_model
if not model_name:
console.print("[red]No model available on engine.[/red]")
sys.exit(1)
# Agent mode
if agent_name is not None:
@@ -30,7 +30,7 @@ except ImportError:
F = None # type: ignore[assignment]
from openjarvis.core.registry import LearningRegistry
from openjarvis.learning._stubs import AgentLearningPolicy
from openjarvis.learning._stubs import IntelligenceLearningPolicy
# ---------------------------------------------------------------------------
# Config
@@ -474,7 +474,7 @@ def _ensure_registered() -> None:
return
@LearningRegistry.register("orchestrator_grpo")
class OrchestratorGRPOPolicy(AgentLearningPolicy):
class OrchestratorGRPOPolicy(IntelligenceLearningPolicy):
"""Wrapper that registers the GRPO trainer as a learning policy."""
def update(
@@ -25,7 +25,7 @@ except ImportError:
DataLoader = None # type: ignore[assignment,misc]
from openjarvis.core.registry import LearningRegistry
from openjarvis.learning._stubs import AgentLearningPolicy
from openjarvis.learning._stubs import IntelligenceLearningPolicy
# ---------------------------------------------------------------------------
# Config
@@ -377,7 +377,7 @@ def _ensure_registered() -> None:
return
@LearningRegistry.register("orchestrator_sft")
class OrchestratorSFTPolicy(AgentLearningPolicy):
class OrchestratorSFTPolicy(IntelligenceLearningPolicy):
"""Wrapper that registers the SFT trainer as a learning policy."""
def update(
+4
View File
@@ -161,6 +161,10 @@ class MCPServer:
return tools
def get_tools(self) -> List[BaseTool]:
"""Return all tool instances (for use by SystemBuilder)."""
return list(self._tools.values())
def handle(self, request: MCPRequest) -> MCPResponse:
"""Dispatch an MCP request and return a response."""
if request.method == "initialize":
+9 -33
View File
@@ -305,41 +305,17 @@ class Jarvis:
}
def _resolve_model(self, query: str) -> Optional[str]:
"""Use router policy to select a model."""
"""Resolve model using config fallback chain."""
if self._config.intelligence.default_model:
return self._config.intelligence.default_model
# Try first available from engine
try:
from openjarvis.core.registry import RouterPolicyRegistry
from openjarvis.engine._discovery import discover_engines, discover_models
from openjarvis.intelligence import (
HeuristicRouter,
build_routing_context,
merge_discovered_models,
register_builtin_models,
)
from openjarvis.learning import ensure_registered as _ensure_learning
register_builtin_models()
_ensure_learning()
all_engines = discover_engines(self._config)
all_models = discover_models(all_engines)
for ek, model_ids in all_models.items():
merge_discovered_models(ek, model_ids)
policy_key = self._config.learning.default_policy
if RouterPolicyRegistry.contains(policy_key):
router_cls = RouterPolicyRegistry.get(policy_key)
else:
router_cls = HeuristicRouter
router = router_cls(
available_models=all_models.get(self._resolved_engine_key, []),
default_model=self._config.intelligence.default_model,
fallback_model=self._config.intelligence.fallback_model,
)
ctx = build_routing_context(query)
return router.select_model(ctx)
models = self._engine.list_models()
if models:
return models[0]
except Exception:
return None
pass
return self._config.intelligence.fallback_model or None
def _run_agent(
self,
+77 -51
View File
@@ -498,65 +498,91 @@ class SystemBuilder:
def _resolve_tools(self, config, engine, model, memory_backend,
channel_backend=None):
"""Resolve tool instances."""
tool_names_str = self._tool_names
if tool_names_str is None:
# Use config default or agent defaults
"""Resolve tool instances via MCPServer (primary) + external MCP servers."""
from openjarvis.mcp.server import MCPServer
# 1. Build internal MCPServer with all auto-discovered tools
internal_server = MCPServer()
# 2. Inject runtime dependencies into tools that need them
for tool in internal_server.get_tools():
self._inject_tool_deps(tool, engine, model, memory_backend, channel_backend)
# 3. Determine which tool names to include
tool_names = self._tool_names
if tool_names is None:
raw = config.tools.enabled or config.agent.default_tools
if raw:
tool_names_str = [n.strip() for n in raw.split(",") if n.strip()]
tool_names = [n.strip() for n in raw.split(",") if n.strip()]
else:
tool_names_str = []
tool_names = []
tools: List[BaseTool] = []
for name in tool_names_str:
# 4. Filter to requested tool names (if specified)
if tool_names:
all_tools = {t.spec.name: t for t in internal_server.get_tools()}
tools = [all_tools[n] for n in tool_names if n in all_tools]
else:
tools = []
# 5. Discover external MCP server tools
if config.tools.mcp.servers:
try:
if name == "retrieval" and memory_backend:
from openjarvis.tools.retrieval import RetrievalTool
tools.append(RetrievalTool(memory_backend))
elif name == "llm":
from openjarvis.tools.llm_tool import LLMTool
tools.append(LLMTool(engine, model=model))
elif name in (
"memory_store",
"memory_retrieve",
"memory_search",
"memory_index",
):
from openjarvis.core.registry import ToolRegistry
from openjarvis.tools import storage_tools # noqa: F401
if ToolRegistry.contains(name):
tool = ToolRegistry.create(name, backend=memory_backend)
tools.append(tool)
elif name in (
"channel_send",
"channel_list",
"channel_status",
):
import openjarvis.tools.channel_tools # noqa: F401
from openjarvis.core.registry import ToolRegistry
if ToolRegistry.contains(name):
tool = ToolRegistry.create(name, channel=channel_backend)
tools.append(tool)
else:
from openjarvis.core.registry import ToolRegistry
if ToolRegistry.contains(name):
tools.append(ToolRegistry.create(name))
else:
# Try direct import for tools that need ensure_registered
import openjarvis.tools # noqa: F401
if ToolRegistry.contains(name):
tools.append(ToolRegistry.create(name))
except Exception:
import json
server_list = json.loads(config.tools.mcp.servers)
if isinstance(server_list, list):
for server_cfg in server_list:
try:
external_tools = self._discover_external_mcp(server_cfg)
if tool_names:
external_tools = [
t for t in external_tools
if t.spec.name in tool_names
]
tools.extend(external_tools)
except Exception:
pass
except (json.JSONDecodeError, TypeError):
pass
return tools
@staticmethod
def _inject_tool_deps(tool, engine, model, memory_backend, channel_backend):
"""Inject runtime dependencies into tools that need them."""
name = tool.spec.name
if name == "llm":
if hasattr(tool, "_engine"):
tool._engine = engine
if hasattr(tool, "_model"):
tool._model = model
elif name == "retrieval":
if hasattr(tool, "_backend"):
tool._backend = memory_backend
elif name.startswith("memory_"):
if hasattr(tool, "_backend"):
tool._backend = memory_backend
elif name.startswith("channel_"):
if hasattr(tool, "_channel"):
tool._channel = channel_backend
@staticmethod
def _discover_external_mcp(server_cfg) -> List[BaseTool]:
"""Discover tools from an external MCP server configuration."""
import json
from openjarvis.mcp.client import MCPClient
from openjarvis.mcp.transport import StdioTransport
from openjarvis.tools.mcp_adapter import MCPToolProvider
cfg = json.loads(server_cfg) if isinstance(server_cfg, str) else server_cfg
command = cfg.get("command", "")
args = cfg.get("args", [])
if not command:
return []
transport = StdioTransport(command=command, args=args)
client = MCPClient(transport)
provider = MCPToolProvider(client)
return provider.discover()
__all__ = ["JarvisSystem", "SystemBuilder"]
+33 -33
View File
@@ -1,4 +1,4 @@
"""Tests for the --router CLI option in jarvis ask."""
"""Tests for model resolution fallback chain in jarvis ask."""
from __future__ import annotations
@@ -48,8 +48,9 @@ def _patch_engine(engine):
)
class TestAskRouter:
def test_default_uses_heuristic(self) -> None:
class TestAskModelResolution:
def test_default_model_from_config(self) -> None:
"""When no -m flag, uses config.intelligence.default_model."""
engine = _mock_engine()
patches = _patch_engine(engine)
with patches[0], patches[1], patches[2], patches[3], patches[4], patches[5]:
@@ -57,33 +58,8 @@ class TestAskRouter:
assert result.exit_code == 0
assert "Hello!" in result.output
def test_explicit_heuristic(self) -> None:
engine = _mock_engine()
patches = _patch_engine(engine)
with patches[0], patches[1], patches[2], patches[3], patches[4], patches[5]:
result = CliRunner().invoke(cli, ["ask", "--router", "heuristic", "Hello"])
assert result.exit_code == 0
assert "Hello!" in result.output
def test_grpo_raises_error(self) -> None:
engine = _mock_engine()
patches = _patch_engine(engine)
with patches[0], patches[1], patches[2], patches[3], patches[4], patches[5]:
result = CliRunner().invoke(cli, ["ask", "--router", "grpo", "Hello"])
# GRPO raises NotImplementedError which should surface as an error
assert result.exit_code != 0
def test_unknown_router_falls_back(self) -> None:
engine = _mock_engine()
patches = _patch_engine(engine)
with patches[0], patches[1], patches[2], patches[3], patches[4], patches[5]:
result = CliRunner().invoke(
cli, ["ask", "--router", "nonexistent", "Hello"],
)
assert result.exit_code == 0
assert "Hello!" in result.output
def test_model_flag_bypasses_router(self) -> None:
def test_explicit_model_flag(self) -> None:
"""The -m flag directly selects a model, bypassing fallback chain."""
engine = _mock_engine()
patches = _patch_engine(engine)
with patches[0], patches[1], patches[2], patches[3], patches[4], patches[5]:
@@ -93,7 +69,8 @@ class TestAskRouter:
assert result.exit_code == 0
assert "Hello!" in result.output
def test_config_default_policy_respected(self) -> None:
def test_fallback_to_engine_models(self) -> None:
"""When default_model is empty, falls back to first engine model."""
engine = _mock_engine()
patches = _patch_engine(engine)
with (
@@ -104,9 +81,32 @@ class TestAskRouter:
):
cfg = mock_config.return_value
cfg.telemetry.enabled = False
cfg.intelligence.default_model = "test-model"
cfg.intelligence.default_model = ""
cfg.intelligence.fallback_model = ""
cfg.learning.default_policy = "heuristic"
cfg.memory.context_injection = False
result = CliRunner().invoke(cli, ["ask", "Hello"])
assert result.exit_code == 0
def test_fallback_to_fallback_model(self) -> None:
"""When default_model is empty and no engine models, uses fallback_model."""
engine = _mock_engine()
patches = _patch_engine(engine)
# Override discover_models to return empty list
with (
patches[0], patches[1],
mock.patch.object(
_ask_mod, "discover_models",
return_value={"mock": []},
),
patches[3], patches[4], patches[5],
mock.patch.object(
_ask_mod, "load_config",
) as mock_config,
):
cfg = mock_config.return_value
cfg.telemetry.enabled = False
cfg.intelligence.default_model = ""
cfg.intelligence.fallback_model = "fallback-model"
cfg.memory.context_injection = False
result = CliRunner().invoke(cli, ["ask", "Hello"])
assert result.exit_code == 0