dict structure change

giskard/llm/client/gemini.py

@@ -5,6 +5,8 @@
 from ..errors import LLMImportError
 from . import LLMClient
 from .base import ChatMessage
+from google.generativeai.types import ContentDict
+import os
 
 try:
  import google.generativeai as genai
@@ -13,6 +15,16 @@
  flavor="llm", msg="To use Gemini models, please install the `genai` package with `pip install google-generativeai`"
  ) from err
 
+genai.configure(api_key=os.environ["GOOGLE_API_KEY"])
+
+AUTH_ERROR_MESSAGE = (
+ "Could not authenticate with Gemini API. Please make sure you have configured the API key by "
+ "setting GOOGLE_API_KEY in the environment."
+)
+
+def _to_content_dict(messages: Sequence[ChatMessage]) -> Sequence[ContentDict]: 
+ return [ContentDict(role=m.role, parts=m.content) for m in messages]
+
 class GeminiClient(LLMClient):
  def __init__(self, model: str = "gemini-pro", _client=None ):
  self.model = model
@@ -37,21 +49,24 @@ def complete(
 
  try:
  completion = self._client.generate_content(
- messages=[asdict(m) for m in messages],
+ messages=_to_content_dict(messages),
  temperature=temperature,
  max_tokens=max_tokens,
  **extra_params,
  )
  except RuntimeError as err:
  raise LLMConfigurationError("Could not get response from Gemini API") from err
 
+ prompt_tokens = self._client.count_tokens(messages)
+ completion_tokens = self._client.count_tokens(completion.text)
+
  self.logger.log_call(
- prompt_tokens=completion.usage.prompt_tokens,
- sampled_tokens=completion.usage.completion_tokens,
+ prompt_tokens=prompt_tokens,
+ sampled_tokens=completion_tokens,
  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.candidates[0].content.parts)
+ return ChatMessage(role=completion.candidates[0].content.role, content=completion.text)