feat(llm): add GroqClient support to set any Groq model as default LLMClient (#1977)

Author: Kunjan Shah

giskard/llm/__init__.py

@@ -10,4 +10,5 @@
     "set_llm_api",
     "set_default_embedding",
     "set_embedding_model",
+    "GroqClient"
 ]

giskard/llm/client/__init__.py

@@ -6,6 +6,10 @@
 
 from .base import ChatMessage, LLMClient
 from .logger import LLMLogger
+from .groq_client import GroqClient
+from ..config import LLMConfigurationError
+
+
 
 logger = logging.getLogger(__name__)
 
@@ -36,12 +40,15 @@ def get_default_llm_api() -> str:
     global _default_llm_api
     if _default_llm_api is None:
         _default_llm_api = os.getenv(
-            "GSK_LLM_API", "azure" if "AZURE_OPENAI_API_KEY" in os.environ else "openai"
+            "GSK_LLM_API", 
+            "azure" if "AZURE_OPENAI_API_KEY" in os.environ 
+            else "groq" if "GROQ_API_KEY" in os.environ
+            else "openai"
         ).lower()
 
-        if _default_llm_api not in {"azure", "openai"}:
+        if _default_llm_api not in {"azure", "openai", "groq"}:
             logging.warning(
-                f"LLM-based evaluation is only working with `azure` and `openai`. Found {_default_llm_api} in GSK_LLM_API, falling back to `openai`"
+                f"LLM-based evaluation is only working with `azure`, `openai`, 'groq'. Found {_default_llm_api} in GSK_LLM_API, falling back to `openai`"
             )
             _default_llm_api = "openai"
 
@@ -89,27 +96,34 @@ def get_default_client() -> LLMClient:
     global _default_llm_api
     global _default_llm_model
     global _disable_structured_output
+ 
 
     if _default_client is not None:
         return _default_client
 
     try:
-        from .litellm import LiteLLMClient
-
-        if (
-            _default_llm_api is not None
-            and "/" in _default_llm_model
-            and not _default_llm_model.startswith(f"{_default_llm_api}/")
-        ):
-            raise ValueError(
-                f"Model {_default_llm_model} is not compatible with {_default_llm_api}: https://docs.giskard.ai/en/latest/open_source/setting_up/index.html "
-            )
-        if _default_llm_api is not None and "/" not in _default_llm_model:
-            _default_llm_model = f"{_default_llm_api}/{_default_llm_model}"
-
-        _default_client = LiteLLMClient(_default_llm_model, _disable_structured_output, _default_completion_params)
-    except ImportError:
-        raise ValueError(f"LLM scan using {_default_llm_model} requires litellm")
+        if _default_llm_api == "groq":
+            groq_api_key = os.getenv("GROQ_API_KEY")
+            if not groq_api_key:
+                raise LLMConfigurationError("GROQ_API_KEY environment variable is not set")
+            _default_client = GroqClient(_default_llm_model)
+        else:
+            from .litellm import LiteLLMClient
+
+            if (
+                _default_llm_api is not None
+                and "/" in _default_llm_model
+                and not _default_llm_model.startswith(f"{_default_llm_api}/")
+            ):
+                raise ValueError(
+                    f"Model {_default_llm_model} is not compatible with {_default_llm_api}: https://docs.giskard.ai/en/latest/open_source/setting_up/index.html "
+                )
+            if _default_llm_api is not None and "/" not in _default_llm_model:
+                _default_llm_model = f"{_default_llm_api}/{_default_llm_model}"
+
+            _default_client = LiteLLMClient(_default_llm_model, _disable_structured_output, _default_completion_params)
+    except ImportError as e:
+        raise ValueError(f"LLM scan using {_default_llm_model} requires appropriate client library") from e
 
     return _default_client
 
@@ -122,4 +136,5 @@ def get_default_client() -> LLMClient:
     "set_llm_model",
     "get_default_llm_api",
     "set_llm_api",
+    "GroqClient",
 ]

giskard/llm/client/groq_client.py

@@ -0,0 +1,93 @@
+from typing import Optional, Sequence
+
+from dataclasses import asdict
+from logging import warning
+import logging
+
+from ..config import LLMConfigurationError
+from ..errors import LLMImportError
+from . import LLMClient
+from .base import ChatMessage
+
+try:
+    from groq import Groq
+    import groq
+except ImportError as err:
+    raise LLMImportError(flavor="llm") from err
+
+AUTH_ERROR_MESSAGE = (
+    "Could not authenticate with Groq API. Please make sure you have configured the API key by "
+    "setting GROQ_API_KEY in the environment."
+)
+
+def _supports_json_format(model: str) -> bool:
+    if "llama-3.3-70b-versatile" in model:
+        return True
+
+    if model == "llama-3.1-8b-instant" or model == "gemma2-9b-it":
+        return True
+
+    return False
+
+logger = logging.getLogger(__name__)
+
+class GroqClient(LLMClient):
+    def __init__(
+        self, 
+        model: str = "llama-3.3-70b-versatile",  # Default model for Groq
+        client: Groq = None,
+        json_mode: Optional[bool] = None
+    ):
+        logger.info(f"Initializing GroqClient with model: {model}")
+        self.model = model
+        self._client = client or Groq()
+        self.json_mode = json_mode if json_mode is not None else _supports_json_format(model)
+        logger.info("GroqClient initialized successfully")
+
+    def complete(
+        self,
+        messages: Sequence[ChatMessage],
+        temperature: float = 1.0,
+        max_tokens: Optional[int] = None,
+        caller_id: Optional[str] = None,
+        seed: Optional[int] = None,
+        format=None,
+    ) -> ChatMessage:
+        logger.info(f"GroqClient.complete called with model: {self.model}")
+        logger.info(f"Messages: {messages}")
+        
+        extra_params = dict()
+
+        if seed is not None:
+            extra_params["seed"] = seed
+
+        if self.json_mode:
+            if format not in (None, "json", "json_object"):
+                warning(f"Unsupported format '{format}', ignoring.")
+                format = None
+
+            if format == "json" or format == "json_object":
+                extra_params["response_format"] = {"type": "json_object"}
+
+        try:
+            completion = self._client.chat.completions.create(
+                model=self.model,
+                messages=[asdict(m) for m in messages],
+                temperature=temperature,
+                max_tokens=max_tokens,
+                **extra_params,
+            )
+        except groq.AuthenticationError as err: 
+            raise LLMConfigurationError(AUTH_ERROR_MESSAGE) from err
+
+        self.logger.log_call(
+            prompt_tokens=completion.usage.prompt_tokens,
+            sampled_tokens=completion.usage.completion_tokens,
+            model=self.model,
+            client_class=self.__class__.__name__,
+            caller_id=caller_id,
+        )
+
+        msg = completion.choices[0].message
+
+        return ChatMessage(role=msg.role, content=msg.content)

tests/llm/test_llm_client.py

@@ -12,6 +12,7 @@
 from giskard.llm.client.bedrock import ClaudeBedrockClient
 from giskard.llm.client.gemini import GeminiClient
 from giskard.llm.client.openai import OpenAIClient
+from giskard.llm.client.groq_client import GroqClient #GroqClient 
 
 PYDANTIC_V2 = pydantic.__version__.startswith("2.")
 
@@ -217,3 +218,43 @@ def test_gemini_client():
     # Assert that the response is a ChatMessage and has the correct content
     assert isinstance(res, ChatMessage)
     assert res.content == "This is a test!"
+
+@pytest.mark.skipif(not PYDANTIC_V2, reason="Groq client test requires Pydantic v2")
+def test_groq_client():
+    # Mock the Groq response
+    demo_response = Mock()
+    demo_response.usage = Mock(prompt_tokens=13, completion_tokens=7)
+    demo_response.choices = [
+        Mock(
+            message=Mock(
+                role="assistant",
+                content="This is a test!"
+            )
+        )
+    ]
+
+    # Mock the Groq client
+    mock_client = Mock()
+    mock_client.chat.completions.create.return_value = demo_response
+ 
+    client = GroqClient(model="llama3-8b-8192", client=mock_client)
+
+    # Call the complete method
+    res = client.complete(
+        [ChatMessage(role="user", content="Hello")],
+        temperature=0.7,
+        max_tokens=100
+    )
+
+    # Assert that create was called with correct arguments
+    mock_client.chat.completions.create.assert_called_once()
+    assert mock_client.chat.completions.create.call_args[1]["messages"] == [
+        {"role": "user", "content": "Hello"}
+    ]
+    assert mock_client.chat.completions.create.call_args[1]["temperature"] == 0.7
+    assert mock_client.chat.completions.create.call_args[1]["max_tokens"] == 100
+    assert mock_client.chat.completions.create.call_args[1]["model"] == "llama3-8b-8192"
+
+    # Assert the response is correct
+    assert isinstance(res, ChatMessage)
+    assert res.content == "This is a test!"