generator: Groq API

README.md

@@ -162,6 +162,12 @@ Private Replicate endpoints:
 * `--model_name` (optional, `command` by default) - The specific Cohere model you'd like to test
 * set the `COHERE_API_KEY` environment variable to your Cohere API key, e.g. "aBcDeFgHiJ123456789"; see https://dashboard.cohere.ai/api-keys when logged in
 
+### Groq
+
+* `--model_type groq`
+* `--model_name` - The name of the model to access via the Groq API
+* set the `GROQ_API_KEY` environment variable to your Groq API key, see https://console.groq.com/docs/quickstart for details on creating an API key
+
 ### ggml
 
 * `--model_type ggml`

docs/source/garak.generators.groq.rst

@@ -0,0 +1,8 @@
+garak.generators.groq
+====================
+
+.. automodule:: garak.generators.groq
+   :members:
+   :undoc-members:
+   :show-inheritance:   
+

docs/source/generators.rst

@@ -13,6 +13,7 @@ For a detailed oversight into how a generator operates, see :ref:`garak.generato
    garak.generators.cohere
    garak.generators.function
    garak.generators.ggml
+   garak.generators.groq
    garak.generators.guardrails
    garak.generators.huggingface
    garak.generators.langchain

garak/generators/groq.py

@@ -0,0 +1,72 @@
+"""GroqChat API support"""
+
+import random
+from typing import List, Union
+
+import openai
+
+from garak.generators.openai import OpenAICompatible
+
+
+class GroqChat(OpenAICompatible):
+    """Wrapper for Groq-hosted LLM models. Expects GROQ_API_KEY environment variable.
+    See https://console.groq.com/docs/quickstart for more info on how to set up a Groq API key
+    Uses the [OpenAI-compatible API](https://console.groq.com/docs/openai)
+    """
+
+    # per https://console.groq.com/docs/openai
+    # 2024.09.04, `n>1` is not supported
+    ENV_VAR = "GROQ_API_KEY"
+    DEFAULT_PARAMS = OpenAICompatible.DEFAULT_PARAMS | {
+        "temperature": 0.7,
+        "top_p": 1.0,
+        "uri": "https://api.groq.com/openai/v1",
+        "vary_seed_each_call": True,  # encourage variation when generations>1
+        "vary_temp_each_call": True,  # encourage variation when generations>1
+        "suppressed_params": {
+            "n",
+            "frequency_penalty",
+            "presence_penalty",
+            "logprobs",
+            "logit_bias",
+            "top_logprobs",
+        },
+    }
+    active = True
+    supports_multiple_generations = False
+    generator_family_name = "Groq"
+
+    timeout = 60
+
+    def _load_client(self):
+        self.client = openai.OpenAI(base_url=self.uri, api_key=self.api_key)
+        if self.name in ("", None):
+            raise ValueError(
+                "Groq API requires model name to be set, e.g. --model_name llama-3.1-8b-instant \nCurrent models:\n"
+                + "\n - ".join(
+                    sorted([entry.id for entry in self.client.models.list().data])
+                )
+            )
+        self.generator = self.client.chat.completions
+
+    def _clear_client(self):
+        self.generator = None
+        self.client = None
+
+    def _call_model(
+        self, prompt: str | List[dict], generations_this_call: int = 1
+    ) -> List[Union[str, None]]:
+        assert (
+            generations_this_call == 1
+        ), "generations_per_call / n > 1 is not supported"
+
+        if self.vary_seed_each_call:
+            self.seed = random.randint(0, 65535)
+
+        if self.vary_temp_each_call:
+            self.temperature = random.random()
+
+        return super()._call_model(prompt, generations_this_call)
+
+
+DEFAULT_CLASS = "GroqChat"

tests/generators/test_groq.py

@@ -0,0 +1,31 @@
+import os
+import pytest
+
+import garak.cli
+from garak.generators.groq import GroqChat
+
+
+@pytest.mark.skipif(
+    os.getenv(GroqChat.ENV_VAR, None) is None,
+    reason=f"GroqChat API key is not set in {GroqChat.ENV_VAR}",
+)
+def test_groq_instantiate():
+    g = GroqChat(name="llama3-8b-8192")
+
+
+@pytest.mark.skipif(
+    os.getenv(GroqChat.ENV_VAR, None) is None,
+    reason=f"GroqChat  API key is not set in {GroqChat.ENV_VAR}",
+)
+def test_groq_generate_1():
+    g = GroqChat(name="llama3-8b-8192")
+    result = g._call_model("this is a test", generations_this_call=1)
+    assert isinstance(result, list), "GroqChat _call_model should return a list"
+    assert len(result) == 1, "GroqChat _call_model result list should have one item"
+    assert isinstance(result[0], str), "GroqChat _call_model should return a list"
+    result = g.generate("this is a test", generations_this_call=1)
+    assert isinstance(result, list), "GroqChat generate() should return a list"
+    assert (
+        len(result) == 1
+    ), "GroqChat generate() result list should have one item when generations_this_call=1"
+    assert isinstance(result[0], str), "GroqChat generate() should return a list"