Skip to content

fix(models): add Gemini 2.5 token limits #1140

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
amirshahzadhashmi7145 wants to merge 1 commit into
base: pre/beta
Choose a base branch
from
+106 −66
Open

fix(models): add Gemini 2.5 token limits #1140

Show file tree
Hide file tree
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
9 changes: 9 additions & 0 deletions scrapegraphai/helpers/models_tokens.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -140,6 +140,11 @@
"gemini-2.0-flash-latest": 1000000,
"gemini-2.0-flash-exp": 1000000,
"gemini-2.0-pro-exp": 2000000,
"gemini-2.5-flash": 1000000,
"gemini-2.5-flash-latest": 1000000,
"gemini-2.5-flash-lite": 1000000,
"gemini-2.5-pro": 1000000,
"gemini-flash-latest": 1000000,
"models/embedding-001": 2048,
},
"google_vertexai": {
Expand All @@ -150,6 +155,10 @@
"gemini-2.0-flash-exp": 1048576,
"gemini-2.0-pro": 2000000,
"gemini-2.0-pro-exp": 2000000,
"gemini-2.5-flash": 1048576,
"gemini-2.5-flash-lite": 1048576,
"gemini-2.5-pro": 1048576,
"gemini-flash-latest": 1048576,
},
"ollama": {
"command-r": 12800,
Expand Down
163 changes: 97 additions & 66 deletions tests/test_models_tokens.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,9 +7,9 @@ class TestModelsTokens:
def test_openai_tokens(self):
"""Test that the 'openai' provider exists and its tokens are valid positive integers."""
openai_models = models_tokens.get("openai")
assert openai_models is not None, (
"'openai' key should be present in models_tokens"
)
assert (
openai_models is not None
), "'openai' key should be present in models_tokens"
for model, token in openai_models.items():
assert isinstance(model, str), "Model name should be a string"
assert isinstance(token, int), "Token limit should be an integer"
Expand All @@ -30,19 +30,50 @@ def test_google_providers(self):
assert google_genai is not None, "'google_genai' key should be present"
assert google_vertexai is not None, "'google_vertexai' key should be present"
# Check a specific key from google_genai
assert "gemini-pro" in google_genai, (
"'gemini-pro' should be in google_genai models"
)
assert (
"gemini-pro" in google_genai
), "'gemini-pro' should be in google_genai models"
# Validate token values types
for provider in [google_genai, google_vertexai]:
for token in provider.values():
assert isinstance(token, int), "Token limit must be an integer"

def test_gemini_2_5_models_are_registered(self):
"""Gemini 2.5 / flash-latest must be in the table so they are not truncated to 8192.

#1121: an unknown model silently falls back to an 8192-token window.
google_genai/gemini-2.5-flash is the reported case; gemini-flash-latest
is the current flash alias. Both have a 1M input context.
"""
google_genai = models_tokens["google_genai"]
google_vertexai = models_tokens["google_vertexai"]

for model in (
"gemini-2.5-flash",
"gemini-2.5-flash-latest",
"gemini-2.5-flash-lite",
"gemini-2.5-pro",
"gemini-flash-latest",
):
assert (
google_genai.get(model) == 1000000
), f"Expected 1M context for {model} in google_genai"

for model in (
"gemini-2.5-flash",
"gemini-2.5-flash-lite",
"gemini-2.5-pro",
"gemini-flash-latest",
):
assert (
google_vertexai.get(model) == 1048576
), f"Expected 1M context for {model} in google_vertexai"

def test_non_existent_provider(self):
"""Test that a non-existent provider returns None."""
assert models_tokens.get("non_existent") is None, (
"Non-existent provider should return None"
)
assert (
models_tokens.get("non_existent") is None
), "Non-existent provider should return None"

def test_total_model_keys(self):
"""Test that the total number of models across all providers is above an expected count."""
Expand All @@ -59,136 +90,136 @@ def test_non_empty_model_keys(self):
"""Ensure that model token names are non-empty strings."""
for provider, model_dict in models_tokens.items():
for model in model_dict.keys():
assert model != "", (
f"Model name in provider '{provider}' should not be empty."
)
assert (
model != ""
), f"Model name in provider '{provider}' should not be empty."

def test_token_limits_range(self):
"""Test that token limits for all models fall within a plausible range (e.g., 1 to 300000)."""
for provider, model_dict in models_tokens.items():
for model, token in model_dict.items():
assert 1 <= token <= 1100000, (
f"Token limit for {model} in provider {provider} is out of plausible range."
)
assert (
1 <= token <= 1100000
), f"Token limit for {model} in provider {provider} is out of plausible range."

def test_provider_structure(self):
"""Test that every provider in models_tokens has a dictionary as its value."""
for provider, models in models_tokens.items():
assert isinstance(models, dict), (
f"Provider {provider} should map to a dictionary, got {type(models).__name__}"
)
assert isinstance(
models, dict
), f"Provider {provider} should map to a dictionary, got {type(models).__name__}"

def test_non_empty_provider(self):
"""Test that each provider dictionary is not empty."""
for provider, models in models_tokens.items():
assert len(models) > 0, (
f"Provider {provider} should contain at least one model."
)
assert (
len(models) > 0
), f"Provider {provider} should contain at least one model."

def test_specific_model_token_values(self):
"""Test specific expected token values for selected models from various providers."""
# Verify a token for a selected model from the 'openai' provider
openai = models_tokens.get("openai")
assert openai.get("gpt-3.5-turbo-0125") == 16385, (
"Expected token limit for gpt-3.5-turbo-0125 in openai to be 16385"
)
assert (
openai.get("gpt-3.5-turbo-0125") == 16385
), "Expected token limit for gpt-3.5-turbo-0125 in openai to be 16385"

# Verify a token for a selected model from the 'azure_openai' provider
azure = models_tokens.get("azure_openai")
assert azure.get("gpt-3.5") == 4096, (
"Expected token limit for gpt-3.5 in azure_openai to be 4096"
)
assert (
azure.get("gpt-3.5") == 4096
), "Expected token limit for gpt-3.5 in azure_openai to be 4096"

# Verify a token for a selected model from the 'anthropic' provider
anthropic = models_tokens.get("anthropic")
assert anthropic.get("claude_instant") == 100000, (
"Expected token limit for claude_instant in anthropic to be 100000"
)
assert (
anthropic.get("claude_instant") == 100000
), "Expected token limit for claude_instant in anthropic to be 100000"

def test_providers_count(self):
"""Test that the total number of providers is as expected (at least 15)."""
assert len(models_tokens) >= 15, (
"Expected at least 15 providers in models_tokens"
)
assert (
len(models_tokens) >= 15
), "Expected at least 15 providers in models_tokens"

def test_non_existent_model(self):
"""Test that a non-existent model within a valid provider returns None."""
openai = models_tokens.get("openai")
assert openai.get("non_existent_model") is None, (
"Non-existent model should return None from a valid provider."
)
assert (
openai.get("non_existent_model") is None
), "Non-existent model should return None from a valid provider."

def test_no_whitespace_in_model_names(self):
"""Test that model names do not contain leading or trailing whitespace."""
for provider, model_dict in models_tokens.items():
for model in model_dict.keys():
# Assert that stripping whitespace does not change the model name
assert model == model.strip(), (
f"Model name '{model}' in provider '{provider}' contains leading or trailing whitespace."
)
assert (
model == model.strip()
), f"Model name '{model}' in provider '{provider}' contains leading or trailing whitespace."

def test_specific_models_additional(self):
"""Test specific token values for additional models across various providers."""
# Check some models in the 'ollama' provider
ollama = models_tokens.get("ollama")
assert ollama.get("llama2") == 4096, (
"Expected token limit for 'llama2' in ollama to be 4096"
)
assert ollama.get("llama2:70b") == 4096, (
"Expected token limit for 'llama2:70b' in ollama to be 4096"
)
assert (
ollama.get("llama2") == 4096
), "Expected token limit for 'llama2' in ollama to be 4096"
assert (
ollama.get("llama2:70b") == 4096
), "Expected token limit for 'llama2:70b' in ollama to be 4096"

# Check a specific model from the 'mistralai' provider
mistralai = models_tokens.get("mistralai")
assert mistralai.get("open-codestral-mamba") == 256000, (
"Expected token limit for 'open-codestral-mamba' in mistralai to be 256000"
)
assert (
mistralai.get("open-codestral-mamba") == 256000
), "Expected token limit for 'open-codestral-mamba' in mistralai to be 256000"

# Check a specific model from the 'deepseek' provider
deepseek = models_tokens.get("deepseek")
assert deepseek.get("deepseek-chat") == 28672, (
"Expected token limit for 'deepseek-chat' in deepseek to be 28672"
)
assert (
deepseek.get("deepseek-chat") == 28672
), "Expected token limit for 'deepseek-chat' in deepseek to be 28672"

# Check a model from the 'ernie' provider
ernie = models_tokens.get("ernie")
assert ernie.get("ernie-bot") == 4096, (
"Expected token limit for 'ernie-bot' in ernie to be 4096"
)
assert (
ernie.get("ernie-bot") == 4096
), "Expected token limit for 'ernie-bot' in ernie to be 4096"

def test_nvidia_specific(self):
"""Test specific token value for 'meta/codellama-70b' in the nvidia provider."""
nvidia = models_tokens.get("nvidia")
assert nvidia is not None, "'nvidia' provider should exist"
# Verify token for 'meta/codellama-70b' equals 16384 as defined in the nvidia dictionary
assert nvidia.get("meta/codellama-70b") == 16384, (
"Expected token limit for 'meta/codellama-70b' in nvidia to be 16384"
)
assert (
nvidia.get("meta/codellama-70b") == 16384
), "Expected token limit for 'meta/codellama-70b' in nvidia to be 16384"

def test_groq_specific(self):
"""Test specific token value for 'claude-3-haiku-20240307\'' in the groq provider."""
groq = models_tokens.get("groq")
assert groq is not None, "'groq' provider should exist"
# Note: The model name has an embedded apostrophe at the end in its name.
assert groq.get("claude-3-haiku-20240307'") == 8192, (
"Expected token limit for 'claude-3-haiku-20240307\\'' in groq to be 8192"
)
assert (
groq.get("claude-3-haiku-20240307'") == 8192
), "Expected token limit for 'claude-3-haiku-20240307\\'' in groq to be 8192"

def test_togetherai_specific(self):
"""Test specific token value for 'meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo' in the toghetherai provider."""
togetherai = models_tokens.get("toghetherai")
assert togetherai is not None, "'toghetherai' provider should exist"
expected = 128000
model_name = "meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo"
assert togetherai.get(model_name) == expected, (
f"Expected token limit for '{model_name}' in toghetherai to be {expected}"
)
assert (
togetherai.get(model_name) == expected
), f"Expected token limit for '{model_name}' in toghetherai to be {expected}"

def test_ernie_all_values(self):
"""Test that all models in the 'ernie' provider have token values exactly 4096."""
ernie = models_tokens.get("ernie")
assert ernie is not None, "'ernie' provider should exist"
for model, token in ernie.items():
assert token == 4096, (
f"Expected token limit for '{model}' in ernie to be 4096, got {token}"
)
assert (
token == 4096
), f"Expected token limit for '{model}' in ernie to be 4096, got {token}"