diff --git a/python/src/aiconfig/Config.py b/python/src/aiconfig/Config.py index 33ea6e26a..320ff9bae 100644 --- a/python/src/aiconfig/Config.py +++ b/python/src/aiconfig/Config.py @@ -365,7 +365,7 @@ async def run_batch( Raises: IndexError: If the identifier for the prompt doesn't exist in the list of available prompts. - + IndexError: If the model name doesn't exist in the list of available model parsers. Returns: list[Tuple["ExecuteResult", JSONObject | Any, Dict[str, Any]]]: A list of tuples, each tuple consisting of the inference result, the corresponding resolved completion parameters, and the parameters dict used. diff --git a/python/src/aiconfig/registry.py b/python/src/aiconfig/registry.py index fa5937bbc..c055b9b2d 100644 --- a/python/src/aiconfig/registry.py +++ b/python/src/aiconfig/registry.py @@ -61,6 +61,10 @@ def get_model_parser(model_id: str) -> ModelParser: Returns: ModelParser: The retrieved model parser """ + if model_id not in ModelParserRegistry.parser_ids(): + raise IndexError( + f"Model parser '{model_id}' not found in registry, available model parsers are:\n {ModelParserRegistry.parser_ids()}" + ) return ModelParserRegistry._parsers[model_id] @staticmethod diff --git a/python/tests/test_registry.py b/python/tests/test_registry.py index 677d66b10..2e0466780 100644 --- a/python/tests/test_registry.py +++ b/python/tests/test_registry.py @@ -88,7 +88,7 @@ def test_retrieve_model_parser(self): def test_retrieve_nonexistent_model_parser(self): # Attempt to retrieve a model parser that is not registered - with pytest.raises(KeyError): + with pytest.raises(IndexError): ModelParserRegistry.get_model_parser("nonexistent-model") def test_retrieve_model_parser_for_prompt( @@ -134,7 +134,7 @@ def test_retrieve_model_parser_for_prompt_with_nonexistent_model( ai_config_runtime.add_prompt(prompt.name, prompt) # Attempt to retrieve a model parser for the Prompt - with pytest.raises(KeyError): + with pytest.raises(IndexError): ModelParserRegistry.get_model_parser_for_prompt( prompt, ai_config_runtime ).id() @@ -153,7 +153,7 @@ def test_remove_model_parser(self): ModelParserRegistry.remove_model_parser("model-8") # Attempt to retrieve the removed model parser - with pytest.raises(KeyError): + with pytest.raises(IndexError): parser = ModelParserRegistry.get_model_parser("model-8") == None def test_clear_registry(self):