From 68b5d24c3e1f86934263db17c6cbc4d52f197175 Mon Sep 17 00:00:00 2001 From: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> Date: Fri, 15 Aug 2025 11:47:08 +0200 Subject: [PATCH 1/3] fix(prompts): prevent IndexError when LLM provided via constructor with empty models config - Add check in get_task_model to handle empty _models list gracefully - Return None instead of throwing IndexError when no models match - Add comprehensive test coverage for various model configuration scenarios Fixes the issue where providing an LLM object directly to LLMRails constructor would fail if the YAML config had an empty models list. --- nemoguardrails/llm/prompts.py | 3 +- tests/test_llm_task_manager.py | 58 +++++++++++++++++++++++++++++++++- tests/test_llmrails.py | 36 +++++++++++++++++++++ 3 files changed, 95 insertions(+), 2 deletions(-) diff --git a/nemoguardrails/llm/prompts.py b/nemoguardrails/llm/prompts.py index 390bbeeb1..03d09be2d 100644 --- a/nemoguardrails/llm/prompts.py +++ b/nemoguardrails/llm/prompts.py @@ -139,7 +139,8 @@ def get_task_model(config: RailsConfig, task: Union[str, Task]) -> Model: if not _models: _models = [model for model in config.models if model.type == "main"] - return _models[0] + if _models: + return _models[0] return None diff --git a/tests/test_llm_task_manager.py b/tests/test_llm_task_manager.py index 9afd48914..c833fb709 100644 --- a/tests/test_llm_task_manager.py +++ b/tests/test_llm_task_manager.py @@ -20,7 +20,7 @@ from nemoguardrails import RailsConfig from nemoguardrails.llm.filters import conversation_to_events -from nemoguardrails.llm.prompts import get_prompt +from nemoguardrails.llm.prompts import get_prompt, get_task_model from nemoguardrails.llm.taskmanager import LLMTaskManager from nemoguardrails.llm.types import Task @@ -457,3 +457,59 @@ def test_reasoning_traces_not_included_in_prompt_history(): "Hi there!" in rendered_prompt or "I don't have access to real-time weather information." in rendered_prompt ) + + +def test_get_task_model_with_empty_models(): + """Test that get_task_model returns None when models list is empty. + + This tests the fix for the IndexError that occurred when the models list was empty. + """ + config = RailsConfig.parse_object({"models": []}) + + result = get_task_model(config, "main") + assert result is None + + result = get_task_model(config, Task.GENERAL) + assert result is None + + +def test_get_task_model_with_no_matching_models(): + """Test that get_task_model returns None when no models match the requested type.""" + config = RailsConfig.parse_object( + { + "models": [ + { + "type": "embeddings", + "engine": "openai", + "model": "text-embedding-ada-002", + } + ] + } + ) + + result = get_task_model(config, "main") + assert result is None + + +def test_get_task_model_with_main_model(): + """Test that get_task_model returns the main model when present.""" + config = RailsConfig.parse_object( + {"models": [{"type": "main", "engine": "openai", "model": "gpt-3.5-turbo"}]} + ) + + result = get_task_model(config, "main") + assert result is not None + assert result.type == "main" + assert result.engine == "openai" + assert result.model == "gpt-3.5-turbo" + + +def test_get_task_model_fallback_to_main(): + """Test that get_task_model falls back to main model when specific task model not found.""" + config = RailsConfig.parse_object( + {"models": [{"type": "main", "engine": "openai", "model": "gpt-3.5-turbo"}]} + ) + + result = get_task_model(config, "some_other_task") + assert result is not None + assert result.type == "main" diff --git a/tests/test_llmrails.py b/tests/test_llmrails.py index 96ff01b67..2451dd7b2 100644 --- a/tests/test_llmrails.py +++ b/tests/test_llmrails.py @@ -856,6 +856,42 @@ async def test_other_models_honored(mock_init, llm_config_with_multiple_models): assert any(event.get("intent") == "express greeting" for event in new_events) +@pytest.mark.asyncio +async def test_llm_constructor_with_empty_models_config(): + """Test that LLMRails can be initialized with constructor LLM when config has empty models list. + + This tests the fix for the IndexError that occurred when providing an LLM via constructor + but having an empty models list in the config. + """ + config = RailsConfig.parse_object( + { + "models": [], + "user_messages": { + "express greeting": ["Hello!"], + }, + "flows": [ + { + "elements": [ + {"user": "express greeting"}, + {"bot": "express greeting"}, + ] + }, + ], + "bot_messages": { + "express greeting": ["Hello! How are you?"], + }, + } + ) + + injected_llm = FakeLLM(responses=["express greeting"]) + llm_rails = LLMRails(config=config, llm=injected_llm) + assert llm_rails.llm == injected_llm + + events = [{"type": "UtteranceUserActionFinished", "final_transcript": "Hello!"}] + new_events = await llm_rails.runtime.generate_events(events) + assert any(event.get("intent") == "express greeting" for event in new_events) + + @pytest.mark.asyncio @patch( "nemoguardrails.rails.llm.llmrails.init_llm_model", From f9b5e4bacb7edb63f14da31d77161d59d153f891 Mon Sep 17 00:00:00 2001 From: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> Date: Tue, 19 Aug 2025 12:20:52 +0200 Subject: [PATCH 2/3] update get_task_model return type to Optional --- nemoguardrails/llm/prompts.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/nemoguardrails/llm/prompts.py b/nemoguardrails/llm/prompts.py index 03d09be2d..8f00b2b55 100644 --- a/nemoguardrails/llm/prompts.py +++ b/nemoguardrails/llm/prompts.py @@ -16,7 +16,7 @@ """Prompts for the various steps in the interaction.""" import os -from typing import List, Union +from typing import List, Optional, Union import yaml @@ -129,7 +129,7 @@ def _get_prompt( raise ValueError(f"Could not find prompt for task {task_name} and model {model}") -def get_task_model(config: RailsConfig, task: Union[str, Task]) -> Model: +def get_task_model(config: RailsConfig, task: Union[str, Task]) -> Optional[Model]: """Return the model for the given task in the current config.""" # Fetch current task parameters like name, models to use, and the prompting mode task_name = str(task.value) if isinstance(task, Task) else task From 3d8935363bd6126cc3d2063652bec75666889077 Mon Sep 17 00:00:00 2001 From: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> Date: Tue, 19 Aug 2025 12:23:50 +0200 Subject: [PATCH 3/3] expand models in main model test --- tests/test_llm_task_manager.py | 21 ++++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/tests/test_llm_task_manager.py b/tests/test_llm_task_manager.py index c833fb709..7897e55b6 100644 --- a/tests/test_llm_task_manager.py +++ b/tests/test_llm_task_manager.py @@ -494,7 +494,26 @@ def test_get_task_model_with_no_matching_models(): def test_get_task_model_with_main_model(): """Test that get_task_model returns the main model when present.""" config = RailsConfig.parse_object( - {"models": [{"type": "main", "engine": "openai", "model": "gpt-3.5-turbo"}]} + { + "models": [ + { + "type": "embeddings", + "engine": "openai", + "model": "text-embedding-ada-002", + }, + { + "type": "custom_task", + "engine": "anthropic", + "model": "claude-4.1-opus", + }, + { + "type": "fact_checking", + "engine": "openai", + "model": "gpt-4", + }, + {"type": "main", "engine": "openai", "model": "gpt-3.5-turbo"}, + ] + } ) result = get_task_model(config, "main")