Skip to content

fix: allow custom Gemini models instead of strict list check #94

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
bluzername wants to merge 1 commit into
base: main
Choose a base branch
from
+10 −6
Open

fix: allow custom Gemini models instead of strict list check #94

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 10 additions & 6 deletions server.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -122,6 +122,10 @@ def format(self, record):
"gemini-2.5-pro"
]

def is_gemini_model(model_name: str) -> bool:
"""Check if a model name is a Gemini model (known list or starts with 'gemini')."""
return model_name in GEMINI_MODELS or model_name.startswith("gemini")

# Helper function to clean schema for Gemini
def clean_gemini_schema(schema: Any) -> Any:
"""Recursively removes unsupported fields from a JSON schema for Gemini."""
Expand Down Expand Up @@ -222,7 +226,7 @@ def validate_model_field(cls, v, info): # Renamed to avoid conflict

# Map Haiku to SMALL_MODEL based on provider preference
elif 'haiku' in clean_v.lower():
if PREFERRED_PROVIDER == "google" and SMALL_MODEL in GEMINI_MODELS:
if PREFERRED_PROVIDER == "google" and is_gemini_model(SMALL_MODEL):
new_model = f"gemini/{SMALL_MODEL}"
mapped = True
else:
Expand All @@ -231,7 +235,7 @@ def validate_model_field(cls, v, info): # Renamed to avoid conflict

# Map Sonnet to BIG_MODEL based on provider preference
elif 'sonnet' in clean_v.lower():
if PREFERRED_PROVIDER == "google" and BIG_MODEL in GEMINI_MODELS:
if PREFERRED_PROVIDER == "google" and is_gemini_model(BIG_MODEL):
new_model = f"gemini/{BIG_MODEL}"
mapped = True
else:
Expand All @@ -240,7 +244,7 @@ def validate_model_field(cls, v, info): # Renamed to avoid conflict

# Add prefixes to non-mapped models if they match known lists
elif not mapped:
if clean_v in GEMINI_MODELS and not v.startswith('gemini/'):
if is_gemini_model(clean_v) and not v.startswith('gemini/'):
new_model = f"gemini/{clean_v}"
mapped = True # Technically mapped to add prefix
elif clean_v in OPENAI_MODELS and not v.startswith('openai/'):
Expand Down Expand Up @@ -295,7 +299,7 @@ def validate_model_token_count(cls, v, info): # Renamed to avoid conflict
mapped = False
# Map Haiku to SMALL_MODEL based on provider preference
if 'haiku' in clean_v.lower():
if PREFERRED_PROVIDER == "google" and SMALL_MODEL in GEMINI_MODELS:
if PREFERRED_PROVIDER == "google" and is_gemini_model(SMALL_MODEL):
new_model = f"gemini/{SMALL_MODEL}"
mapped = True
else:
Expand All @@ -304,7 +308,7 @@ def validate_model_token_count(cls, v, info): # Renamed to avoid conflict

# Map Sonnet to BIG_MODEL based on provider preference
elif 'sonnet' in clean_v.lower():
if PREFERRED_PROVIDER == "google" and BIG_MODEL in GEMINI_MODELS:
if PREFERRED_PROVIDER == "google" and is_gemini_model(BIG_MODEL):
new_model = f"gemini/{BIG_MODEL}"
mapped = True
else:
Expand All @@ -313,7 +317,7 @@ def validate_model_token_count(cls, v, info): # Renamed to avoid conflict

# Add prefixes to non-mapped models if they match known lists
elif not mapped:
if clean_v in GEMINI_MODELS and not v.startswith('gemini/'):
if is_gemini_model(clean_v) and not v.startswith('gemini/'):
new_model = f"gemini/{clean_v}"
mapped = True # Technically mapped to add prefix
elif clean_v in OPENAI_MODELS and not v.startswith('openai/'):
Expand Down