Merge branch 'feature/litellm-generator' of https://github.com/Tien-C…

…heng/garak into feature/litellm-generator

garak/generators/litellm.py

@@ -138,7 +138,7 @@ def __init__(self, name: str, generations: int = 10):
             )
 
     @backoff.on_exception(backoff.fibo, Exception, max_value=70)
-    def _call_model(self, prompt: Union[str, List[dict]]) -> List[str] | str | None:
+    def _call_model(self, prompt: Union[str, List[dict]]) -> Union[List[str], str, None]:
         if isinstance(prompt, str):
             prompt = [{"role": "user", "content": prompt}]
         elif isinstance(prompt, list):

tests/generators/test_litellm.py

@@ -12,8 +12,9 @@
     reason="OpenAI API key is not set in OPENAI_API_KEY",
 )
 def test_litellm_openai():
-    generator = LiteLLMGenerator(name="gpt-3.5-turbo")
-    assert generator.name == "gpt-3.5-turbo"
+    model_name = "gpt-3.5-turbo"
+    generator = LiteLLMGenerator(name=model_name)
+    assert generator.name == model_name
     assert generator.generations == DEFAULT_GENERATIONS_QTY
     assert isinstance(generator.max_tokens, int)
 
@@ -30,8 +31,9 @@ def test_litellm_openai():
     reason="OpenRouter API key is not set in OPENROUTER_API_KEY",
 )
 def test_litellm_openrouter():
-    generator = LiteLLMGenerator(name="openrouter/google/gemma-7b-it")
-    assert generator.name == "openrouter/google/gemma-7b-it"
+    model_name = "openrouter/google/gemma-7b-it"
+    generator = LiteLLMGenerator(name=model_name)
+    assert generator.name == model_name
     assert generator.generations == DEFAULT_GENERATIONS_QTY
     assert isinstance(generator.max_tokens, int)