Gemini clent test

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!"