From 4dac1337756478dc6d7f7aae858ee12f15c413c7 Mon Sep 17 00:00:00 2001 From: Sita Vemuri Date: Thu, 9 Jul 2026 00:12:26 -0700 Subject: [PATCH 1/2] adapt for gemini --- README.md | 8 +- llmx/configs/config.default.yml | 16 +++ llmx/generators/text/gemini_textgen.py | 130 +++++++++++++++++++++++++ llmx/generators/text/textgen.py | 9 +- pyproject.toml | 1 + tests/test_generators.py | 10 ++ 6 files changed, 171 insertions(+), 3 deletions(-) create mode 100644 llmx/generators/text/gemini_textgen.py diff --git a/README.md b/README.md index 23908a3..7a24eee 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@ [![PyPI version](https://badge.fury.io/py/llmx.svg)](https://badge.fury.io/py/llmx) -A simple python package that provides a unified interface to several LLM providers of chat fine-tuned models [OpenAI, AzureOpenAI, PaLM, Cohere and local HuggingFace Models]. +A simple python package that provides a unified interface to several LLM providers of chat fine-tuned models [OpenAI, AzureOpenAI, PaLM, Gemini, Cohere, Anthropic and local HuggingFace Models]. > **Note** > llmx wraps multiple api providers and its interface _may_ change as the providers as well as the general field of LLMs evolve. @@ -16,6 +16,7 @@ from llmx import llm gen = llm(provider="openai") # support azureopenai models too. gen = llm(provider="palm") # or google +gen = llm(provider="gemini") gen = llm(provider="cohere") # or palm gen = llm(provider="hf", model="HuggingFaceH4/zephyr-7b-beta", device_map="auto") # run huggingface model locally ``` @@ -83,6 +84,9 @@ export PALM_API_KEY= export PALM_SERVICE_ACCOUNT_KEY_FILE= export PALM_PROJECT_ID= export PALM_PROJECT_LOCATION= + +# for Gemini, get an api key from Google AI Studio (https://aistudio.google.com/app/apikey) +export GEMINI_API_KEY= ``` You can also set the default provider and list of supported providers via a config file. Use the yaml format in this [sample `config.default.yml` file](llmx/configs/config.default.yml) and set the `LLMX_CONFIG_PATH` to the path of the config file. @@ -127,7 +131,9 @@ hfgen_gen = llm( - Supported models - [x] OpenAI - [x] PaLM ([MakerSuite](https://developers.generativeai.google/api/rest/generativelanguage), [Vertex AI](https://cloud.google.com/vertex-ai/docs/generative-ai/learn/models)) + - [x] Gemini - [x] Cohere + - [x] Anthropic - [x] HuggingFace (local) ## Caveats diff --git a/llmx/configs/config.default.yml b/llmx/configs/config.default.yml index c13730c..96751d4 100644 --- a/llmx/configs/config.default.yml +++ b/llmx/configs/config.default.yml @@ -16,6 +16,22 @@ providers: provider: anthropic parameters: model: claude-3-5-sonnet-20240620 + gemini: + name: Gemini + description: Google's Gemini models (via the Gemini API / AI Studio). + models: + - name: gemini-1.5-pro + max_tokens: 8192 + model: + provider: gemini + parameters: + model: gemini-1.5-pro + - name: gemini-1.5-flash + max_tokens: 8192 + model: + provider: gemini + parameters: + model: gemini-1.5-flash openai: name: OpenAI description: OpenAI's and AzureOpenAI GPT-3 and GPT-4 models. diff --git a/llmx/generators/text/gemini_textgen.py b/llmx/generators/text/gemini_textgen.py new file mode 100644 index 0000000..053d08e --- /dev/null +++ b/llmx/generators/text/gemini_textgen.py @@ -0,0 +1,130 @@ +from typing import Union, List, Dict +import os +from google import genai +from google.genai import types +from dataclasses import asdict + +from .base_textgen import TextGenerator +from ...datamodel import TextGenerationConfig, TextGenerationResponse, Message +from ...utils import cache_request, get_models_maxtoken_dict, num_tokens_from_messages + + +class GeminiTextGenerator(TextGenerator): + def __init__( + self, + api_key: str = None, + provider: str = "gemini", + model: str = None, + models: Dict = None, + ): + super().__init__(provider=provider) + api_key = api_key or os.environ.get( + "GEMINI_API_KEY", os.environ.get("GOOGLE_API_KEY", None) + ) + if api_key is None: + raise ValueError( + "Gemini API key is not set. Please set the GEMINI_API_KEY environment variable." + ) + self.client = genai.Client(api_key=api_key) + self.model_max_token_dict = get_models_maxtoken_dict(models) + self.model_name = model or "gemini-1.5-flash" + + def format_messages(self, messages): + system_message = None + formatted_messages = [] + for message in messages: + content = message["content"].strip() + if message["role"] == "system": + system_message = content if system_message is None else system_message + "\n" + content + else: + role = "model" if message["role"] == "assistant" else "user" + formatted_messages.append( + types.Content(role=role, parts=[types.Part.from_text(text=content)]) + ) + return system_message, formatted_messages + + def generate( + self, + messages: Union[List[Dict], str], + config: TextGenerationConfig = TextGenerationConfig(), + **kwargs, + ) -> TextGenerationResponse: + use_cache = config.use_cache + model = config.model or self.model_name + self.model_name = model + + system_message, formatted_messages = self.format_messages(messages) + if not formatted_messages: + raise ValueError("At least one message is required") + + prompt_tokens = num_tokens_from_messages(messages) + max_tokens = max( + self.model_max_token_dict.get(model, 8192) - prompt_tokens - 10, 200 + ) + + stop_sequences = config.stop if isinstance(config.stop, list) else ( + [config.stop] if config.stop else None + ) + max_output_tokens = config.max_tokens or max_tokens + + cache_key_params = { + "model": model, + "messages": messages, + "system_message": system_message, + "generation_config": { + "candidate_count": config.n, + "max_output_tokens": max_output_tokens, + "temperature": config.temperature, + "top_p": config.top_p, + "top_k": config.top_k, + "stop_sequences": stop_sequences, + }, + } + + if use_cache: + response = cache_request(cache=self.cache, params=cache_key_params) + if response: + return TextGenerationResponse(**response) + + generation_config = types.GenerateContentConfig( + system_instruction=system_message, + candidate_count=config.n, + max_output_tokens=max_output_tokens, + temperature=config.temperature, + top_p=config.top_p, + top_k=config.top_k, + stop_sequences=stop_sequences, + ) + + gemini_response = self.client.models.generate_content( + model=model, contents=formatted_messages, config=generation_config + ) + + response_text = [ + Message(role="assistant", content=candidate.content.parts[0].text) + for candidate in gemini_response.candidates + ] + + usage = {} + if gemini_response.usage_metadata: + usage = { + "prompt_tokens": gemini_response.usage_metadata.prompt_token_count, + "completion_tokens": gemini_response.usage_metadata.candidates_token_count, + "total_tokens": gemini_response.usage_metadata.total_token_count, + } + + response = TextGenerationResponse( + text=response_text, + logprobs=[], + config=cache_key_params["generation_config"], + usage=usage, + response=gemini_response, + ) + + cache_request( + cache=self.cache, params=cache_key_params, values=asdict(response) + ) + return response + + def count_tokens(self, text) -> int: + return num_tokens_from_messages(text) diff --git a/llmx/generators/text/textgen.py b/llmx/generators/text/textgen.py index 3d86002..40de7d0 100644 --- a/llmx/generators/text/textgen.py +++ b/llmx/generators/text/textgen.py @@ -3,6 +3,7 @@ from .palm_textgen import PalmTextGenerator from .cohere_textgen import CohereTextGenerator from .anthropic_textgen import AnthropicTextGenerator +from .gemini_textgen import GeminiTextGenerator import logging logger = logging.getLogger("llmx") @@ -19,9 +20,11 @@ def sanitize_provider(provider: str): return "hf" elif provider.lower() == "anthropic" or provider.lower() == "claude": return "anthropic" + elif provider.lower() == "gemini": + return "gemini" else: raise ValueError( - f"Invalid provider '{provider}'. Supported providers are 'openai', 'hf', 'palm', 'cohere', and 'anthropic'." + f"Invalid provider '{provider}'. Supported providers are 'openai', 'hf', 'palm', 'cohere', 'anthropic', and 'gemini'." ) @@ -58,6 +61,8 @@ def llm(provider: str = None, **kwargs): return CohereTextGenerator(**kwargs) elif provider.lower() == "anthropic": return AnthropicTextGenerator(**kwargs) + elif provider.lower() == "gemini": + return GeminiTextGenerator(**kwargs) elif provider.lower() == "hf": try: import transformers @@ -80,5 +85,5 @@ def llm(provider: str = None, **kwargs): else: raise ValueError( - f"Invalid provider '{provider}'. Supported providers are 'openai', 'hf', 'palm', 'cohere', and 'anthropic'." + f"Invalid provider '{provider}'. Supported providers are 'openai', 'hf', 'palm', 'cohere', 'anthropic', and 'gemini'." ) \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 4c8c56b..ee78864 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,6 +26,7 @@ dependencies = [ "cohere", "google.auth", "anthropic", + "google-genai", "typer", "pyyaml", ] diff --git a/tests/test_generators.py b/tests/test_generators.py index 4f4e59c..87feccc 100644 --- a/tests/test_generators.py +++ b/tests/test_generators.py @@ -50,6 +50,16 @@ def test_google(): # assert len(google_response.text) == 2 palm may chose to return 1 or 2 responses +def test_gemini(): + gemini_gen = llm(provider="gemini", api_key=os.environ.get("GEMINI_API_KEY", None)) + config.model = "gemini-1.5-flash" + gemini_response = gemini_gen.generate(messages, config=config) + answer = gemini_response.text[0].content + print(gemini_response.text[0].content) + + assert ("paris" in answer.lower()) + + def test_cohere(): cohere_gen = llm(provider="cohere") config.model = "command" From 8fd9b31089a7c968dec57711be4ab2a62e26aa08 Mon Sep 17 00:00:00 2001 From: Sita Vemuri Date: Thu, 9 Jul 2026 00:38:30 -0700 Subject: [PATCH 2/2] config gemini settings --- llmx/configs/config.default.yml | 12 ++++++------ llmx/generators/text/gemini_textgen.py | 6 +++++- 2 files changed, 11 insertions(+), 7 deletions(-) diff --git a/llmx/configs/config.default.yml b/llmx/configs/config.default.yml index 96751d4..adab0b5 100644 --- a/llmx/configs/config.default.yml +++ b/llmx/configs/config.default.yml @@ -1,8 +1,8 @@ # Sets the the default model to use for llm() when no provider parameter is set. model: - provider: openai + provider: gemini parameters: - api_key: null + model: gemini-2.5-flash # list of supported providers. providers: @@ -20,18 +20,18 @@ providers: name: Gemini description: Google's Gemini models (via the Gemini API / AI Studio). models: - - name: gemini-1.5-pro + - name: gemini-2.5-pro max_tokens: 8192 model: provider: gemini parameters: - model: gemini-1.5-pro - - name: gemini-1.5-flash + model: gemini-2.5-pro + - name: gemini-2.5-flash max_tokens: 8192 model: provider: gemini parameters: - model: gemini-1.5-flash + model: gemini-2.5-flash openai: name: OpenAI description: OpenAI's and AzureOpenAI GPT-3 and GPT-4 models. diff --git a/llmx/generators/text/gemini_textgen.py b/llmx/generators/text/gemini_textgen.py index 053d08e..293b6a4 100644 --- a/llmx/generators/text/gemini_textgen.py +++ b/llmx/generators/text/gemini_textgen.py @@ -37,7 +37,11 @@ def format_messages(self, messages): if message["role"] == "system": system_message = content if system_message is None else system_message + "\n" + content else: - role = "model" if message["role"] == "assistant" else "user" + # lida sends its instruction/prompt content labeled as "assistant" + # (there's no genuine prior model turn to preserve), so treat + # everything non-system as a user turn to satisfy Gemini's + # requirement that single-turn requests end with a user role. + role = "user" formatted_messages.append( types.Content(role=role, parts=[types.Part.from_text(text=content)]) )