andre/crewAI
1
0
Fork 0
mirror of https://github.com/crewAIInc/crewAI.git synced 2026-01-10 00:28:31 +00:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

Address PR comments: Add constants for context window sizes and expand test coverage

Browse Source 
Co-Authored-By: Joe Moura <joao@crewai.com>
This commit is contained in:
Devin AI
2025-04-02 02:24:30 +00:00
parent 3b378dc5a6
commit baeb0dab48
3 changed files with 45 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/crewai/cli/constants.py
View File

@@ -123,7 +123,7 @@ MODELS = {
"claude-3-haiku-20240307", "claude-3-haiku-20240307",
], ],
"gemini": [ "gemini": [
"gemini/gemini-2.5-pro-exp-03-25", "gemini/gemini-2.5-pro-exp-03-25", # Experimental release - March 2025
"gemini/gemini-1.5-flash", "gemini/gemini-1.5-flash",
"gemini/gemini-1.5-pro", "gemini/gemini-1.5-pro",
"gemini/gemini-gemma-2-9b-it", "gemini/gemini-gemma-2-9b-it",

27
src/crewai/llm.py
View File

@@ -71,21 +71,26 @@ class FilteredStream:
return self._original_stream.flush() return self._original_stream.flush()
PRO_CONTEXT_SIZE = 2097152 # 2M tokens
FLASH_CONTEXT_SIZE = 1048576 # 1M tokens
GPT4_TURBO_CONTEXT_SIZE = 128000
CLAUDE_LARGE_CONTEXT_SIZE = 200000
LLM_CONTEXT_WINDOW_SIZES = { LLM_CONTEXT_WINDOW_SIZES = {
# openai # openai
"gpt-4": 8192, "gpt-4": 8192,
"gpt-4o": 128000, "gpt-4o": GPT4_TURBO_CONTEXT_SIZE,
"gpt-4o-mini": 128000, "gpt-4o-mini": GPT4_TURBO_CONTEXT_SIZE,
"gpt-4-turbo": 128000, "gpt-4-turbo": GPT4_TURBO_CONTEXT_SIZE,
"o1-preview": 128000, "o1-preview": GPT4_TURBO_CONTEXT_SIZE,
"o1-mini": 128000, "o1-mini": GPT4_TURBO_CONTEXT_SIZE,
"o3-mini": 200000, # Based on official o3-mini specifications "o3-mini": CLAUDE_LARGE_CONTEXT_SIZE, # Based on official o3-mini specifications
# gemini # gemini
"gemini-2.5-pro-exp-03-25": 2097152, "gemini-2.5-pro-exp-03-25": PRO_CONTEXT_SIZE,
"gemini-2.0-flash": 1048576, "gemini-2.0-flash": FLASH_CONTEXT_SIZE,
"gemini-1.5-pro": 2097152, "gemini-1.5-pro": PRO_CONTEXT_SIZE,
"gemini-1.5-flash": 1048576, "gemini-1.5-flash": FLASH_CONTEXT_SIZE,
"gemini-1.5-flash-8b": 1048576, "gemini-1.5-flash-8b": FLASH_CONTEXT_SIZE,
# deepseek # deepseek
"deepseek-chat": 128000, "deepseek-chat": 128000,
# groq # groq

31
tests/test_gemini_2_5_support.py
View File

@@ -1,6 +1,6 @@
import pytest import pytest
from crewai.llm import LLM from crewai.llm import LLM, CONTEXT_WINDOW_USAGE_RATIO, PRO_CONTEXT_SIZE
def test_get_custom_llm_provider_gemini_2_5(): def test_get_custom_llm_provider_gemini_2_5():
@@ -11,6 +11,31 @@ def test_get_custom_llm_provider_gemini_2_5():
def test_gemini_2_5_context_window_size(): def test_gemini_2_5_context_window_size():
"""Test that the Gemini 2.5 model has the correct context window size.""" """Test that the Gemini 2.5 model has the correct context window size."""
llm = LLM(model="gemini-2.5-pro-exp-03-25") llm = LLM(model="gemini-2.5-pro-exp-03-25")
from crewai.llm import CONTEXT_WINDOW_USAGE_RATIO expected_size = int(PRO_CONTEXT_SIZE * CONTEXT_WINDOW_USAGE_RATIO)
expected_size = int(2097152 * CONTEXT_WINDOW_USAGE_RATIO)
assert llm.get_context_window_size() == expected_size assert llm.get_context_window_size() == expected_size
def test_gemini_2_5_invalid_model_name():
"""Test handling of invalid model name variations."""
llm = LLM(model="gemini-2.5-wrong")
assert llm._get_custom_llm_provider() != "gemini"
def test_gemini_2_5_model_parameters():
"""Test model initialization with various parameters."""
llm = LLM(
model="gemini/gemini-2.5-pro-exp-03-25",
temperature=0.7,
max_tokens=1000
)
assert llm.model == "gemini/gemini-2.5-pro-exp-03-25"
assert llm.temperature == 0.7
assert llm.max_tokens == 1000
def test_gemini_2_5_with_and_without_prefix():
"""Test that the model works with and without the 'gemini/' prefix."""
llm_with_prefix = LLM(model="gemini/gemini-2.5-pro-exp-03-25")
llm_without_prefix = LLM(model="gemini-2.5-pro-exp-03-25")
assert llm_with_prefix._get_custom_llm_provider() == "gemini"
assert llm_without_prefix._get_custom_llm_provider() == "gemini"
assert llm_with_prefix.get_context_window_size() == llm_without_prefix.get_context_window_size()