diff --git a/src/crewai/cli/constants.py b/src/crewai/cli/constants.py index 62d161c71..a192e7830 100644 --- a/src/crewai/cli/constants.py +++ b/src/crewai/cli/constants.py @@ -123,7 +123,7 @@ MODELS = { "claude-3-haiku-20240307", ], "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-pro", "gemini/gemini-gemma-2-9b-it", diff --git a/src/crewai/llm.py b/src/crewai/llm.py index 661ed9426..2b127da25 100644 --- a/src/crewai/llm.py +++ b/src/crewai/llm.py @@ -71,21 +71,26 @@ class FilteredStream: 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 = { # openai "gpt-4": 8192, - "gpt-4o": 128000, - "gpt-4o-mini": 128000, - "gpt-4-turbo": 128000, - "o1-preview": 128000, - "o1-mini": 128000, - "o3-mini": 200000, # Based on official o3-mini specifications + "gpt-4o": GPT4_TURBO_CONTEXT_SIZE, + "gpt-4o-mini": GPT4_TURBO_CONTEXT_SIZE, + "gpt-4-turbo": GPT4_TURBO_CONTEXT_SIZE, + "o1-preview": GPT4_TURBO_CONTEXT_SIZE, + "o1-mini": GPT4_TURBO_CONTEXT_SIZE, + "o3-mini": CLAUDE_LARGE_CONTEXT_SIZE, # Based on official o3-mini specifications # gemini - "gemini-2.5-pro-exp-03-25": 2097152, - "gemini-2.0-flash": 1048576, - "gemini-1.5-pro": 2097152, - "gemini-1.5-flash": 1048576, - "gemini-1.5-flash-8b": 1048576, + "gemini-2.5-pro-exp-03-25": PRO_CONTEXT_SIZE, + "gemini-2.0-flash": FLASH_CONTEXT_SIZE, + "gemini-1.5-pro": PRO_CONTEXT_SIZE, + "gemini-1.5-flash": FLASH_CONTEXT_SIZE, + "gemini-1.5-flash-8b": FLASH_CONTEXT_SIZE, # deepseek "deepseek-chat": 128000, # groq diff --git a/tests/test_gemini_2_5_support.py b/tests/test_gemini_2_5_support.py index 1d843a5b0..906e7716a 100644 --- a/tests/test_gemini_2_5_support.py +++ b/tests/test_gemini_2_5_support.py @@ -1,6 +1,6 @@ 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(): @@ -11,6 +11,31 @@ def test_get_custom_llm_provider_gemini_2_5(): def test_gemini_2_5_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") - from crewai.llm import CONTEXT_WINDOW_USAGE_RATIO - expected_size = int(2097152 * CONTEXT_WINDOW_USAGE_RATIO) + expected_size = int(PRO_CONTEXT_SIZE * CONTEXT_WINDOW_USAGE_RATIO) 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()