diff --git a/graphiti_core/llm_client/gemini_client.py b/graphiti_core/llm_client/gemini_client.py index 53aed7f8..2acd3866 100644 --- a/graphiti_core/llm_client/gemini_client.py +++ b/graphiti_core/llm_client/gemini_client.py @@ -29,15 +29,7 @@ from .errors import RateLimitError logger = logging.getLogger(__name__) -DEFAULT_MODEL = 'gemini-2.0-flash' -DEFAULT_THINKING_BUDGET = 0 - -# Gemini models that support thinking capabilities -GEMINI_THINKING_MODELS = [ - 'gemini-2.5-pro', - 'gemini-2.5-flash', - 'gemini-2.5-flash-lite', -] +DEFAULT_MODEL = 'gemini-2.5-flash' class GeminiClient(LLMClient): @@ -140,17 +132,6 @@ class GeminiClient(LLMClient): types.Content(role=m.role, parts=[types.Part.from_text(text=m.content)]) ) - # Determine the model to be used - model_to_use = self.model or DEFAULT_MODEL - - # Conditionally create thinking_config for models that support thinking - thinking_config_arg = None - if model_to_use in GEMINI_THINKING_MODELS: - thinking_config_arg = types.ThinkingConfig( - include_thoughts=False, - thinking_budget=self.thinking_budget, - ) - # Create generation config generation_config = types.GenerateContentConfig( temperature=self.temperature, @@ -163,7 +144,7 @@ class GeminiClient(LLMClient): # Generate content using the simple string approach response = await self.client.aio.models.generate_content( - model=model_to_use, + model=self.model or DEFAULT_MODEL, contents=gemini_messages, config=generation_config, )