From c775d620afe94c7b86c08fb62fea242c8ebb1086 Mon Sep 17 00:00:00 2001 From: Jon Saad-Falcon Date: Mon, 23 Feb 2026 18:41:07 +0000 Subject: [PATCH] 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 --- src/openjarvis/cli/ask.py | 37 ++--- .../learning/orchestrator/grpo_trainer.py | 4 +- .../learning/orchestrator/sft_trainer.py | 4 +- src/openjarvis/mcp/server.py | 4 + src/openjarvis/sdk.py | 42 ++---- src/openjarvis/system.py | 128 +++++++++++------- tests/cli/test_ask_router.py | 66 ++++----- 7 files changed, 135 insertions(+), 150 deletions(-) diff --git a/src/openjarvis/cli/ask.py b/src/openjarvis/cli/ask.py index c8a4ceea..9bf58126 100644 --- a/src/openjarvis/cli/ask.py +++ b/src/openjarvis/cli/ask.py @@ -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: diff --git a/src/openjarvis/learning/orchestrator/grpo_trainer.py b/src/openjarvis/learning/orchestrator/grpo_trainer.py index db4a020b..fb4f7ccc 100644 --- a/src/openjarvis/learning/orchestrator/grpo_trainer.py +++ b/src/openjarvis/learning/orchestrator/grpo_trainer.py @@ -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( diff --git a/src/openjarvis/learning/orchestrator/sft_trainer.py b/src/openjarvis/learning/orchestrator/sft_trainer.py index 89998f23..8c2b8e41 100644 --- a/src/openjarvis/learning/orchestrator/sft_trainer.py +++ b/src/openjarvis/learning/orchestrator/sft_trainer.py @@ -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( diff --git a/src/openjarvis/mcp/server.py b/src/openjarvis/mcp/server.py index fceae79e..cf1ab0ba 100644 --- a/src/openjarvis/mcp/server.py +++ b/src/openjarvis/mcp/server.py @@ -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": diff --git a/src/openjarvis/sdk.py b/src/openjarvis/sdk.py index 6ef30f79..3acac3ee 100644 --- a/src/openjarvis/sdk.py +++ b/src/openjarvis/sdk.py @@ -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, diff --git a/src/openjarvis/system.py b/src/openjarvis/system.py index d43c0d61..74f620f9 100644 --- a/src/openjarvis/system.py +++ b/src/openjarvis/system.py @@ -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"] diff --git a/tests/cli/test_ask_router.py b/tests/cli/test_ask_router.py index 6906e673..e9581a04 100644 --- a/tests/cli/test_ask_router.py +++ b/tests/cli/test_ask_router.py @@ -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