Gemini API key support

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=<your key>
 export PALM_SERVICE_ACCOUNT_KEY_FILE= <path to your service account key file>
 export PALM_PROJECT_ID=<your gcp project id>
 export PALM_PROJECT_LOCATION=<your project location>
+
+# for Gemini, get an api key from Google AI Studio (https://aistudio.google.com/app/apikey)
+export GEMINI_API_KEY=<your 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

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:
@@ -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-2.5-pro
+        max_tokens: 8192
+        model:
+          provider: gemini
+          parameters:
+            model: gemini-2.5-pro
+      - name: gemini-2.5-flash
+        max_tokens: 8192
+        model:
+          provider: gemini
+          parameters:
+            model: gemini-2.5-flash
   openai:
     name: OpenAI
     description: OpenAI's and AzureOpenAI GPT-3 and GPT-4 models.

llmx/generators/text/gemini_textgen.py

@@ -0,0 +1,134 @@
+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:
+                # 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)])
+                )
+        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)

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'."
         )

pyproject.toml

@@ -26,6 +26,7 @@ dependencies = [
     "cohere", 
     "google.auth",
     "anthropic",
+    "google-genai",
     "typer",
     "pyyaml",
 ]

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"