@@ -30,40 +30,53 @@ def check_llama_version(context: Context, llama_version: str) -> None:
3030
3131@then ("The body of the response has proper model structure" )
3232def check_model_structure (context : Context ) -> None :
33- """Check that the first LLM model has the correct structure and required fields."""
33+ """Check that the expected LLM model has the correct structure and required fields."""
3434 response_json = context .response .json ()
3535 assert response_json is not None , "Response is not valid JSON"
3636
3737 assert "models" in response_json , "Response missing 'models' field"
3838 models = response_json ["models" ]
3939 assert len (models ) > 0 , "Response has empty list of models"
4040
41- # Find first LLM model (same logic as environment.py)
41+ # Get expected values from context (detected in before_all)
42+ expected_model = context .default_model
43+ expected_provider = context .default_provider
44+
45+ # Search for the specific model that was detected in before_all
4246 llm_model = None
4347 for model in models :
44- if model .get ("api_model_type" ) == "llm" :
48+ if (
49+ model .get ("api_model_type" ) == "llm"
50+ and model .get ("provider_id" ) == expected_provider
51+ and model .get ("provider_resource_id" ) == expected_model
52+ ):
4553 llm_model = model
4654 break
4755
48- assert llm_model is not None , "No LLM model found in response"
49-
50- # Get expected values from context
51- expected_model = context .default_model
52- expected_provider = context .default_provider
56+ assert llm_model is not None , (
57+ f"Expected LLM model not found in response. "
58+ f"Looking for provider_id='{ expected_provider } ' and provider_resource_id='{ expected_model } '"
59+ )
5360
5461 # Validate structure and values
55- assert llm_model ["type" ] == "model" , "type should be 'model'"
56- assert llm_model ["api_model_type" ] == "llm" , "api_model_type should be 'llm'"
57- assert llm_model ["model_type" ] == "llm" , "model_type should be 'llm'"
62+ assert (
63+ llm_model ["type" ] == "model"
64+ ), f"type should be 'model', but is { llm_model ["type" ]} "
65+ assert (
66+ llm_model ["api_model_type" ] == "llm"
67+ ), f"api_model_type should be 'llm', but is { llm_model ["api_model_type" ]} "
68+ assert (
69+ llm_model ["model_type" ] == "llm"
70+ ), f"model_type should be 'llm', but is { llm_model ["model_type" ]} "
5871 assert (
5972 llm_model ["provider_id" ] == expected_provider
60- ), f"provider_id should be '{ expected_provider } '"
73+ ), f"provider_id should be '{ expected_provider } ', but is ' { llm_model [ "provider_id" ] } ' "
6174 assert (
6275 llm_model ["provider_resource_id" ] == expected_model
63- ), f"provider_resource_id should be '{ expected_model } '"
76+ ), f"provider_resource_id should be '{ expected_model } ', but is ' { llm_model [ "provider_resource_id" ] } ' "
6477 assert (
6578 llm_model ["identifier" ] == f"{ expected_provider } /{ expected_model } "
66- ), f"identifier should be '{ expected_provider } /{ expected_model } '"
79+ ), f"identifier should be '{ expected_provider } /{ expected_model } ', but is ' { llm_model [ "identifier" ] } ' "
6780
6881
6982@then ("The body of the response has proper shield structure" )
@@ -94,10 +107,10 @@ def check_shield_structure(context: Context) -> None:
94107 ), "provider_id should be 'llama-guard'"
95108 assert (
96109 found_shield ["provider_resource_id" ] == expected_model
97- ), f"provider_resource_id should be '{ expected_model } '"
110+ ), f"provider_resource_id should be '{ expected_model } ', but is ' { found_shield [ "provider_resource_id" ] } ' "
98111 assert (
99112 found_shield ["identifier" ] == "llama-guard-shield"
100- ), "identifier should be 'llama-guard-shield'"
113+ ), f "identifier should be 'llama-guard-shield', but is ' { found_shield [ "identifier" ] } '"
101114
102115
103116@then ("The response contains {count:d} tools listed for provider {provider_name}" )
0 commit comments