Skip to content

Commit

Permalink
Register Google Gemini Models (#403)
Browse files Browse the repository at this point in the history
* Register Google Gemini Models

* update base url

* modified:   src/ell/models/__init__.py
	modified:   src/ell/models/google.py

* modified:   src/ell/models/google.py

* added gemini2.0-flash support

---------

Co-authored-by: aryamanpandya <[email protected]>
Co-authored-by: Aryaman Pandya <[email protected]>
  • Loading branch information
3 people authored Dec 16, 2024
1 parent 36ca5ee commit 92e43a0
Show file tree
Hide file tree
Showing 2 changed files with 73 additions and 2 deletions.
5 changes: 3 additions & 2 deletions src/ell/models/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
"""

from ell.models import openai, anthropic, ollama, groq, bedrock, xai
from ell.models import openai, anthropic, ollama, groq, bedrock, xai, google

__all__ = ["openai", "anthropic", "ollama", "groq", "bedrock", "xai", "google"]

__all__ = ["openai", "anthropic", "ollama", "groq", "bedrock", "xai"]
70 changes: 70 additions & 0 deletions src/ell/models/google.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,70 @@
"""
This module handles the registration of Google Gemini models within the ell framework.
It provides functionality to register various Google Gemini models with a given Google Gemini client,
making them available for use throughout the system. The module also sets up a default
client behavior for unregistered models.
Key features:
1. Registration of specific Google models with their respective types (system, google, google-internal).
2. Utilization of a default Google client for any unregistered models,
The default client behavior ensures that even if a specific model is not explicitly
registered, the system can still attempt to use it with the default Google client.
This fallback mechanism provides flexibility in model usage while maintaining a
structured approach to model registration.
Note: The actual model availability may depend on your Google account's access and the
current offerings from Google.
In this file, we use google's openai compatible client to perform the registration.
"""

import os
from ell.configurator import config
import openai

import logging
import colorama

logger = logging.getLogger(__name__)

def register(client: openai.Client):
"""
Register OpenAI models with the provided client.
This function takes an OpenAI client and registers various OpenAI models
with the global configuration. It allows the system to use these models
for different AI tasks.
Args:
client (openai.Client): An instance of the OpenAI client to be used
for model registration.
Note:
The function doesn't return anything but updates the global
configuration with the registered models.
"""
standard_models = [
'gemini-2.0-flash-exp'
'gemini-1.5-flash',
'gemini-1.5-flash-8b',
'gemini-1.5-pro',
'gemini-1.0-pro',
]
for model_id in standard_models:
config.register_model(model_id, client)


default_client = None
try:
gemini_api_key = os.environ.get("GEMINI_API_KEY")
if not gemini_api_key:
raise openai.OpenAIError("GEMINI_API_KEY not found in environment variables")
default_client = openai.Client(base_url="https://generativelanguage.googleapis.com/v1beta/openai/", api_key=gemini_api_key)
except openai.OpenAIError as e:
pass

register(default_client)

config.default_client = default_client

0 comments on commit 92e43a0

Please sign in to comment.