diff --git a/lib/crewai/src/crewai/llms/providers/gemini/completion.py b/lib/crewai/src/crewai/llms/providers/gemini/completion.py index 0bdf8ec0ea..85799873d7 100644 --- a/lib/crewai/src/crewai/llms/providers/gemini/completion.py +++ b/lib/crewai/src/crewai/llms/providers/gemini/completion.py @@ -1361,24 +1361,28 @@ def get_context_window_size(self) -> int: f"Context window for {key} must be between {min_context} and {max_context}" ) + # Longest prefix first. Always insert new keys in that order so + # startswith prefers gemini-2.0-flash-thinking over gemini-2.0-flash, etc. context_windows = { "gemini-3-pro-preview": 1048576, # 1M tokens - "gemini-2.0-flash": 1048576, # 1M tokens "gemini-2.0-flash-thinking": 32768, "gemini-2.0-flash-lite": 1048576, + "gemini-2.0-flash": 1048576, # 1M tokens "gemini-2.5-flash": 1048576, "gemini-2.5-pro": 1048576, "gemini-1.5-pro": 2097152, # 2M tokens - "gemini-1.5-flash": 1048576, "gemini-1.5-flash-8b": 1048576, + "gemini-1.5-flash": 1048576, "gemini-1.0-pro": 32768, - "gemma-3-1b": 32000, - "gemma-3-4b": 128000, - "gemma-3-12b": 128000, "gemma-3-27b": 128000, + "gemma-3-12b": 128000, + "gemma-3-4b": 128000, + "gemma-3-1b": 32000, } - for model_prefix, size in context_windows.items(): + for model_prefix, size in sorted( + context_windows.items(), key=lambda item: len(item[0]), reverse=True + ): if self.model.startswith(model_prefix): return int(size * CONTEXT_WINDOW_USAGE_RATIO) diff --git a/lib/crewai/tests/test_llm.py b/lib/crewai/tests/test_llm.py index 51c87cd0ea..bef7d44b92 100644 --- a/lib/crewai/tests/test_llm.py +++ b/lib/crewai/tests/test_llm.py @@ -386,6 +386,18 @@ def test_unrecognized_provider_prefix_is_not_stripped() -> None: ) +def test_gemini_thinking_context_window() -> None: + """Gemini thinking models must not be shadowed by gemini-2.0-flash prefix.""" + from crewai.llms.providers.gemini.completion import GeminiCompletion + + gemini_thinking = GeminiCompletion(model="gemini-2.0-flash-thinking-exp-01-21") + assert gemini_thinking.get_context_window_size() == int(32768 * CONTEXT_WINDOW_USAGE_RATIO) + + gemini_flash = GeminiCompletion(model="gemini-2.0-flash-001") + assert gemini_flash.get_context_window_size() == int(1048576 * CONTEXT_WINDOW_USAGE_RATIO) + + + @pytest.fixture def get_weather_tool_schema(): return {