From 02732ae906be636436d10ae757bf1dcd86bd7ac5 Mon Sep 17 00:00:00 2001 From: DavidePaglieri Date: Tue, 25 Nov 2025 11:29:03 +0000 Subject: [PATCH 1/5] feat: google-genai api --- balrog/client.py | 71 ++++++++++++++------------------------- balrog/config/config.yaml | 1 + setup.py | 2 +- 3 files changed, 27 insertions(+), 47 deletions(-) diff --git a/balrog/client.py b/balrog/client.py index df387f48..eb16ccf2 100644 --- a/balrog/client.py +++ b/balrog/client.py @@ -8,7 +8,9 @@ from collections import namedtuple from io import BytesIO -import google.generativeai as genai +from google import genai +from google.genai import types + from anthropic import Anthropic from google.generativeai import caching from openai import OpenAI @@ -230,7 +232,8 @@ def __init__(self, client_config): def _initialize_client(self): """Initialize the Generative AI client if not already initialized.""" if not self._initialized: - self.model = genai.GenerativeModel(self.model_id) + self.client = genai.Client() + self.model = None # Create kwargs dictionary for GenerationConfig client_kwargs = { @@ -241,71 +244,46 @@ def _initialize_client(self): temperature = self.client_kwargs.get("temperature") if temperature is not None: client_kwargs["temperature"] = temperature + + thinking_budget = self.client_kwargs.get("thinking_budget", -1) - self.generation_config = genai.types.GenerationConfig(**client_kwargs) + self.generation_config = genai.types.GenerateContentConfig( + **client_kwargs, + thinking_config=types.ThinkingConfig(thinking_budget=) + ) self._initialized = True def convert_messages(self, messages): - """Convert messages to the format expected by the Generative AI API. + """Convert messages to the format expected by the new Google GenAI SDK. Args: messages (list): A list of message objects. Returns: - list: A list of messages formatted for the Generative AI API. + list[types.Content]: A list of Content objects formatted for the API. """ - # Convert standard Message objects to Gemini's format converted_messages = [] + for msg in messages: parts = [] + role = msg.role if role == "assistant": role = "model" elif role == "system": role = "user" + if msg.content: - parts.append(msg.content) + parts.append(types.Part(text=msg.content)) + if msg.attachment is not None: - parts.append(msg.attachment) + parts.append(types.Part(image=msg.attachment)) + converted_messages.append( - { - "role": role, - "parts": parts, - } + types.Content(role=role, parts=parts) ) return converted_messages - def get_completion(self, converted_messages, max_retries=5, delay=5): - """Get the completion from the model with retries upon failure. - - Args: - converted_messages (list): Messages formatted for the Generative AI API. - max_retries (int, optional): Maximum number of retries. Defaults to 5. - delay (int, optional): Delay between retries in seconds. Defaults to 5. - - Returns: - Response object from the API. - - Raises: - Exception: If the API call fails after the maximum number of retries. - """ - retries = 0 - while retries < max_retries: - try: - response = self.model.generate_content( - converted_messages, - generation_config=self.generation_config, - ) - return response - except Exception as e: - retries += 1 - logger.error(f"Retryable error during generate_content: {e}. Retry {retries}/{max_retries}") - sleep_time = delay * (2 ** (retries - 1)) # Exponential backoff - time.sleep(sleep_time) - - # If maximum retries are reached and still no valid response - raise Exception(f"Failed to get a valid completion after {max_retries} retries.") - def extract_completion(self, response): """Extract the completion text from the API response. @@ -354,9 +332,10 @@ def generate(self, messages): converted_messages = self.convert_messages(messages) def api_call(): - response = self.model.generate_content( - converted_messages, - generation_config=self.generation_config, + response = self.client.models.generate_content( + model=self.model_id, + contents=converted_messages, + config=self.generation_config, ) # Attempt to extract completion immediately after API call completion = self.extract_completion(response) diff --git a/balrog/config/config.yaml b/balrog/config/config.yaml index 89d5ed72..3bde26ff 100644 --- a/balrog/config/config.yaml +++ b/balrog/config/config.yaml @@ -32,6 +32,7 @@ client: generate_kwargs: temperature: 1.0 # Sampling temperature. If null the API default temperature is used instead max_tokens: 4096 # Max tokens to generate in the response + thinking_budget: null # Thinking budget. Set to a number of tokens to control. timeout: 60 # Timeout for API requests in seconds max_retries: 5 # Max number of retries for failed API calls delay: 2 # Exponential backoff factor between retries in seconds diff --git a/setup.py b/setup.py index 268f7a86..b972bf50 100644 --- a/setup.py +++ b/setup.py @@ -21,7 +21,7 @@ install_requires=[ "openai", "anthropic", - "google-generativeai", + "google-genai", "hydra-core", "opencv-python-headless", "wandb", From e0afca7bf142212175f23bd3a3fbb2156142a572 Mon Sep 17 00:00:00 2001 From: DavidePaglieri Date: Tue, 25 Nov 2025 11:35:26 +0000 Subject: [PATCH 2/5] chore: add thinking_budget --- balrog/client.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/balrog/client.py b/balrog/client.py index eb16ccf2..ff01b023 100644 --- a/balrog/client.py +++ b/balrog/client.py @@ -249,7 +249,7 @@ def _initialize_client(self): self.generation_config = genai.types.GenerateContentConfig( **client_kwargs, - thinking_config=types.ThinkingConfig(thinking_budget=) + thinking_config=types.ThinkingConfig(thinking_budget=thinking_budget) ) self._initialized = True From 8e8a0d0d646fbfa30a7ea8d66bfbd2690225a9e0 Mon Sep 17 00:00:00 2001 From: DavidePaglieri Date: Tue, 25 Nov 2025 11:38:20 +0000 Subject: [PATCH 3/5] remove old import --- balrog/client.py | 1 - 1 file changed, 1 deletion(-) diff --git a/balrog/client.py b/balrog/client.py index ff01b023..032ba678 100644 --- a/balrog/client.py +++ b/balrog/client.py @@ -12,7 +12,6 @@ from google.genai import types from anthropic import Anthropic -from google.generativeai import caching from openai import OpenAI LLMResponse = namedtuple( From 21554970daff70f2d1c0f99982429ec5d90955bf Mon Sep 17 00:00:00 2001 From: DavidePaglieri Date: Tue, 25 Nov 2025 11:42:47 +0000 Subject: [PATCH 4/5] fix import --- balrog/utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/balrog/utils.py b/balrog/utils.py index ad625875..13467c0b 100644 --- a/balrog/utils.py +++ b/balrog/utils.py @@ -7,7 +7,7 @@ from collections import defaultdict from pathlib import Path -import google.generativeai as genai +from google import genai import openai @@ -206,7 +206,7 @@ def setup_environment( """ secrets = load_secrets(os.path.join(original_cwd, "SECRETS")) if secrets[gemini_tag]: - genai.configure(api_key=secrets[gemini_tag]) + os.environ["GEMINI_API_KEY"] = secrets[gemini_tag] if secrets[anthropic_tag]: os.environ["ANTHROPIC_API_KEY"] = secrets[anthropic_tag] if secrets[openai_tag]: From b851e2052d9ead6134d5de2969c43fc0fb33be68 Mon Sep 17 00:00:00 2001 From: DavidePaglieri Date: Tue, 25 Nov 2025 11:43:21 +0000 Subject: [PATCH 5/5] remove import --- balrog/utils.py | 1 - 1 file changed, 1 deletion(-) diff --git a/balrog/utils.py b/balrog/utils.py index 13467c0b..7d84816c 100644 --- a/balrog/utils.py +++ b/balrog/utils.py @@ -7,7 +7,6 @@ from collections import defaultdict from pathlib import Path -from google import genai import openai