gmakstutis/graphiti
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

merge

Browse source
This commit is contained in:
Daniel Chalef 2025-06-27 18:21:40 -07:00
parent 4cfd7baf36
commit d57a3bb6c4
1 changed files with 2 additions and 21 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

23
graphiti_core/llm_client/gemini_client.py
View file

@ -29,15 +29,7 @@ from .errors import RateLimitError
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
DEFAULT_MODEL = 'gemini-2.0-flash' DEFAULT_MODEL = 'gemini-2.5-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',
]
class GeminiClient(LLMClient): class GeminiClient(LLMClient):
@ -140,17 +132,6 @@ class GeminiClient(LLMClient):
types.Content(role=m.role, parts=[types.Part.from_text(text=m.content)]) 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 # Create generation config
generation_config = types.GenerateContentConfig( generation_config = types.GenerateContentConfig(
temperature=self.temperature, temperature=self.temperature,
@ -163,7 +144,7 @@ class GeminiClient(LLMClient):
# Generate content using the simple string approach # Generate content using the simple string approach
response = await self.client.aio.models.generate_content( response = await self.client.aio.models.generate_content(
model=model_to_use, model=self.model or DEFAULT_MODEL,
contents=gemini_messages, contents=gemini_messages,
config=generation_config, config=generation_config,
) )