Refactor model selection logic: streamline default model handling in ChatService and WorkbenchService, and enhance LLM selection tests

Signed-off-by: Andre Bossard <anbossar@microsoft.com>

Author: abossard

backend/agent_builder/chat_service.py

@@ -32,6 +32,14 @@ def _env_int(name: str, default: int) -> int:
         return default
 
 
+def _default_model(explicit_model: str = "") -> str:
+    if explicit_model:
+        return explicit_model
+    if os.getenv("AGENT_BACKEND", "").strip().lower() == "openai":
+        return os.getenv("OPENAI_MODEL", "gpt-4o-mini")
+    return os.getenv("LITELLM_MODEL", "github_copilot/gpt-4o")
+
+
 class ChatService:
     """
     Runs one-shot chat agent interactions.
@@ -45,12 +53,12 @@ def __init__(
         self,
         tool_registry: ToolRegistry,
         openai_api_key: str = "",
-        openai_model: str = "gpt-4o-mini",
+        openai_model: str = "",
         openai_base_url: str = "",
     ) -> None:
         self._registry = tool_registry
         self._api_key = openai_api_key or os.getenv("OPENAI_API_KEY", "")
-        self._model = openai_model or os.getenv("OPENAI_MODEL", "gpt-4o-mini")
+        self._model = _default_model(openai_model)
         self._base_url = openai_base_url or os.getenv("OPENAI_BASE_URL", "")
         self._llm: Any = None
         self._recursion_limit = max(3, _env_int("REACT_AGENT_RECURSION_LIMIT", 8))

backend/agent_builder/engine/react_runner.py

@@ -18,6 +18,11 @@
 logger = logging.getLogger(__name__)
 
 
+def _force_openai_backend() -> bool:
+    """Use OpenAI only when explicitly requested, matching backend/agents.py."""
+    return os.getenv("AGENT_BACKEND", "").strip().lower() == "openai"
+
+
 @dataclass
 class RunResult:
     """Immutable result of a ReAct agent invocation."""
@@ -35,8 +40,8 @@ def build_llm(
     max_tokens: int = 0,
     reasoning_effort: str = "low",
 ) -> Any:
-    """Construct an LLM instance — ChatOpenAI when api_key is provided, ChatLiteLLM otherwise."""
-    if api_key:
+    """Construct an LLM instance — LiteLLM by default, OpenAI only when forced."""
+    if _force_openai_backend() and api_key:
         from langchain_openai import ChatOpenAI
         kwargs: dict[str, Any] = {
             "model": model,
@@ -51,7 +56,7 @@ def build_llm(
         return ChatOpenAI(**kwargs)
     else:
         from langchain_litellm import ChatLiteLLM
-        litellm_model = os.getenv("LITELLM_MODEL", "github_copilot/gpt-4o")
+        litellm_model = model or os.getenv("LITELLM_MODEL", "github_copilot/gpt-4o")
         return ChatLiteLLM(model=litellm_model, temperature=temperature)
 
 

backend/agent_builder/service.py

@@ -17,10 +17,15 @@
 from time import perf_counter
 from typing import Any, Optional
 
-from .engine.prompt_builder import append_output_instructions, resolve_output_schema
-from .engine.react_runner import build_llm, build_react_agent, extract_tools_used, make_tool_logging_callback
 from .engine.callbacks import make_streaming_callback
 from .engine.event_bus import AgentEvent, agent_event_bus
+from .engine.prompt_builder import append_output_instructions, resolve_output_schema
+from .engine.react_runner import (
+    build_llm,
+    build_react_agent,
+    extract_tools_used,
+    make_tool_logging_callback,
+)
 from .evaluator import compute_score
 from .evaluator import evaluate_run as _evaluate_criteria
 from .models import (
@@ -41,6 +46,14 @@
 logger = logging.getLogger(__name__)
 
 
+def _default_model(explicit_model: str = "") -> str:
+    if explicit_model:
+        return explicit_model
+    if os.getenv("AGENT_BACKEND", "").strip().lower() == "openai":
+        return os.getenv("OPENAI_MODEL", "gpt-4o-mini")
+    return os.getenv("LITELLM_MODEL", "github_copilot/gpt-4o")
+
+
 # ============================================================================
 # CALCULATIONS (pure — no side effects, no I/O)
 # ============================================================================
@@ -182,13 +195,13 @@ def __init__(
         tool_registry: ToolRegistry,
         db_path: Optional[Path] = None,
         openai_api_key: str = "",
-        openai_model: str = "gpt-4o-mini",
+        openai_model: str = "",
         openai_base_url: str = "",
         recursion_limit: int = 10,
     ) -> None:
         self._registry = tool_registry
         self._api_key = openai_api_key or os.getenv("OPENAI_API_KEY", "")
-        self._model = openai_model or os.getenv("OPENAI_MODEL", "gpt-4o-mini")
+        self._model = _default_model(openai_model)
         self._base_url = openai_base_url or os.getenv("OPENAI_BASE_URL", "")
         self._recursion_limit = recursion_limit
         self._db_path = db_path or (

backend/agent_builder/tests/test_engine.py

@@ -1,5 +1,7 @@
 """Tests for engine helpers — extract_tools_used and ReactRunner."""
 
+from unittest.mock import patch
+
 from agent_builder.engine.react_runner import extract_tools_used
 
 
@@ -70,3 +72,27 @@ def test_mixed_sources(self):
         result = extract_tools_used(messages)
         assert "tool_a" in result
         assert "tool_b" in result
+
+
+class TestBuildLlmSelection:
+    def test_defaults_to_litellm_even_with_api_key(self, monkeypatch):
+        monkeypatch.delenv("AGENT_BACKEND", raising=False)
+
+        with patch("langchain_litellm.ChatLiteLLM") as mock_litellm:
+            from agent_builder.engine.react_runner import build_llm
+
+            build_llm("openai/nvidia/nemotron-3-nano-4b", api_key="test-key")
+
+        mock_litellm.assert_called_once()
+        assert mock_litellm.call_args.kwargs["model"] == "openai/nvidia/nemotron-3-nano-4b"
+
+    def test_uses_openai_when_explicitly_forced(self, monkeypatch):
+        monkeypatch.setenv("AGENT_BACKEND", "openai")
+
+        with patch("langchain_openai.ChatOpenAI") as mock_openai:
+            from agent_builder.engine.react_runner import build_llm
+
+            build_llm("gpt-4o-mini", api_key="test-key", base_url="http://localhost:1234/v1")
+
+        mock_openai.assert_called_once()
+        assert mock_openai.call_args.kwargs["model"] == "gpt-4o-mini"

backend/agent_builder/tests/test_service.py

@@ -35,6 +35,20 @@ def _make_service(tmp_path: Path) -> WorkbenchService:
 
 
 class TestWorkbenchServiceCRUD:
+    def test_defaults_to_litellm_model_when_not_forced(self, monkeypatch):
+        monkeypatch.delenv("AGENT_BACKEND", raising=False)
+        monkeypatch.setenv("LITELLM_MODEL", "openai/nvidia/nemotron-3-nano-4b")
+        monkeypatch.setenv("OPENAI_MODEL", "gpt-4o-mini")
+
+        with TemporaryDirectory() as tmp:
+            registry = ToolRegistry()
+            svc = WorkbenchService(
+                tool_registry=registry,
+                db_path=Path(tmp) / "test.db",
+                openai_api_key="test-key",
+            )
+            assert svc._model == "openai/nvidia/nemotron-3-nano-4b"
+
     def test_create_and_get_agent(self):
         with TemporaryDirectory() as tmp:
             svc = _make_service(Path(tmp))

backend/workbench_integration.py

@@ -14,7 +14,6 @@
 
 # Ensure operations are loaded so @operation decorators run
 import operations  # noqa: F401
-
 from agent_builder import ChatService, ToolRegistry, WorkbenchService
 from api_decorators import get_langchain_tools
 
@@ -56,14 +55,12 @@ def _build_registry() -> ToolRegistry:
     tool_registry=_tool_registry,
     db_path=Path(__file__).parent / "data" / "workbench.db",
     openai_api_key=os.getenv("OPENAI_API_KEY", ""),
-    openai_model=os.getenv("OPENAI_MODEL", "gpt-4o-mini"),
     openai_base_url=os.getenv("OPENAI_BASE_URL", ""),
 )
 
 chat_service = ChatService(
     tool_registry=_tool_registry,
     openai_api_key=os.getenv("OPENAI_API_KEY", ""),
-    openai_model=os.getenv("OPENAI_MODEL", "gpt-4o-mini"),
     openai_base_url=os.getenv("OPENAI_BASE_URL", ""),
 )