Gemini client

giskard/llm/client/gemini.py

@@ -0,0 +1,69 @@
+from typing import Optional, Sequence
+
+from logging import warning
+
+from google.generativeai.types import ContentDict
+
+from ..config import LLMConfigurationError
+from ..errors import LLMImportError
+from . import LLMClient
+from .base import ChatMessage
+
+try:
+ import google.generativeai as genai
+except ImportError as err:
+ raise LLMImportError(
+ flavor="llm",
+ msg="To use Gemini models, please install the `genai` package with `pip install google-generativeai`",
+ ) from err
+
+AUTH_ERROR_MESSAGE = (
+ "Could not get Response from Gemini API. Please make sure you have configured the API key by "
+ "setting GOOGLE_API_KEY in the environment."
+)
+
+
+class GeminiClient(LLMClient):
+ def __init__(self, model: str = "gemini-pro", _client=None):
+ self.model = model
+ self._client = _client or genai.GenerativeModel(self.model)
+
+ 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:
+ extra_params = dict()
+ if seed is not None:
+ extra_params["seed"] = seed
+
+ if format:
+ warning(f"Unsupported format '{format}', ignoring.")
+ format = None
+
+ try:
+ completion = self._client.generate_content(
+ contents=[ContentDict(role=m.role, parts=m.content) for m in messages],
+ generation_config=genai.types.GenerationConfig(
+ temperature=temperature,
+ max_output_tokens=max_tokens,
+ **extra_params,
+ ),
+ )
+ except RuntimeError as err:
+ raise LLMConfigurationError(AUTH_ERROR_MESSAGE) from err
+
+ self.logger.log_call(
+ prompt_tokens=self._client.count_tokens([m.content for m in messages]),
+ sampled_tokens=self._client.count_tokens(completion.text),
+ model=self.model,
+ client_class=self.__class__.__name__,
+ caller_id=caller_id,
+ )
+
+ # Assuming the response structure is similar to the ChatMessage structure
+ return ChatMessage(role=completion.candidates[0].content.role, content=completion.text)

tests/llm/test_llm_client.py

@@ -7,11 +7,13 @@
 from openai.types import CompletionUsage
 from openai.types.chat import ChatCompletion, ChatCompletionMessage
 from openai.types.chat.chat_completion import Choice
+from google.generativeai.types import ContentDict
 
 from giskard.llm.client import ChatMessage
 from giskard.llm.client.bedrock import ClaudeBedrockClient
 from giskard.llm.client.mistral import MistralClient
 from giskard.llm.client.openai import OpenAIClient
+from giskard.llm.client.gemini import GeminiClient
 
 DEMO_OPENAI_RESPONSE = ChatCompletion(
  id="chatcmpl-abc123",
@@ -119,3 +121,33 @@ def test_claude_bedrock_client():
  # Assert that the response is a ChatMessage and has the correct content
  assert isinstance(res, ChatMessage)
  assert res.content == "This is a test!"
+
+def test_gemini_client():
+ # Mock the Gemini client
+ gemini_api_client = Mock()
+ gemini_api_client.generate_content = MagicMock(
+ return_value=Mock(
+ text="This is a test!",
+ candidates=[Mock(content=Mock(role="assistant"))]
+ )
+ )
+ gemini_api_client.count_tokens = MagicMock(
+ side_effect=lambda text: sum(len(t.split()) for t in text) if isinstance(text, list) else len(text.split())
+ )
+
+ # Initialize the GeminiClient with the mocked gemini_api_client
+ client = GeminiClient(model="gemini-pro", _client=gemini_api_client)
+
+ # Call the complete method
+ res = client.complete([ChatMessage(role="user", content="Hello")], temperature=0.11, max_tokens=12)
+ print(res)
+
+ # Assert that the generate_content method was called with the correct arguments
+ gemini_api_client.generate_content.assert_called_once()
+ assert gemini_api_client.generate_content.call_args[1]["contents"] == ([ContentDict(role="user", parts="Hello")])
+ assert gemini_api_client.generate_content.call_args[1]["generation_config"].temperature == 0.11
+ assert gemini_api_client.generate_content.call_args[1]["generation_config"].max_output_tokens == 12
+
+ # Assert that the response is a ChatMessage and has the correct content
+ assert isinstance(res, ChatMessage)
+ assert res.content == "This is a test!"