Added tests for ai_models

esbmc_ai/ai_models.py

@@ -318,6 +318,7 @@ def is_valid_ai_model(
     name: str = ai_model.name if isinstance(ai_model, AIModel) else ai_model
 
     # Try accessing openai api and checking if there is a model defined.
+    # NOTE: This is not tested as no way to mock API currently.
     if api_keys and api_keys.openai:
         try:
             from openai import Client
@@ -339,7 +340,6 @@ def _get_openai_model_max_tokens(name: str) -> int:
     # https://platform.openai.com/docs/models
     tokens = {
         "gpt-4o": 128000,
-        "gpt-4-turbo": 128000,
         "gpt-4": 8192,
         "gpt-3.5-turbo": 16385,
         "gpt-3.5-turbo-instruct": 4096,
@@ -348,11 +348,16 @@ def _get_openai_model_max_tokens(name: str) -> int:
     # Split into - segments and remove each section from the end to find out
     # which one matches the most.
 
+    # Base Case
+    if name in tokens:
+        return tokens[name]
+
+    # Step Case
     name_split: list[str] = name.split("-")
-    for i in range(name.count("-")):
+    for i in range(1, name.count("-")):
         subname: str = "-".join(name_split[:-i])
         if subname in tokens:
-            return tokens[name]
+            return tokens[subname]
 
     raise ValueError(f"Could not figure out max tokens for model: {name}")
 

tests/test_ai_models.py

@@ -16,13 +16,15 @@
     _AIModels,
     get_ai_model_by_name,
     AIModelTextGen,
+    _get_openai_model_max_tokens,
 )
 
+"""TODO Find a way to mock the OpenAI API and test GPT LLM code."""
+
 
 def test_is_valid_ai_model() -> None:
     assert is_valid_ai_model(_AIModels.FALCON_7B.value)
-    assert is_valid_ai_model(_AIModels.GPT_3_16K.value)
-    assert is_valid_ai_model("gpt-3.5-turbo")
+    assert is_valid_ai_model(_AIModels.STARCHAT_BETA.value)
     assert is_valid_ai_model("falcon-7b")
 
 
@@ -56,7 +58,7 @@ def test_add_custom_ai_model() -> None:
 
 def test_get_ai_model_by_name() -> None:
     # Try with first class AI
-    assert get_ai_model_by_name("gpt-3.5-turbo")
+    assert get_ai_model_by_name("falcon-7b")
 
     # Try with custom AI.
     # Add custom AI model if not added by previous tests.
@@ -141,3 +143,15 @@ def test_escape_messages() -> None:
     assert result[3] == filtered[3]
     assert result[4] == filtered[4]
     assert result[5] == filtered[5]
+
+
+def test__get_openai_model_max_tokens() -> None:
+    assert _get_openai_model_max_tokens("gpt-4o") == 128000
+    assert _get_openai_model_max_tokens("gpt-4-turbo") == 8192
+    assert _get_openai_model_max_tokens("gpt-3.5-turbo") == 16385
+    assert _get_openai_model_max_tokens("gpt-3.5-turbo-instruct") == 4096
+    assert _get_openai_model_max_tokens("gpt-3.5-turbo-aaaaaa") == 16385
+    assert _get_openai_model_max_tokens("gpt-3.5-turbo-instruct-bbb") == 4096
+
+    with raises(ValueError):
+        _get_openai_model_max_tokens("aaaaa")