From 6c159d39d497dcf1b6bf527bc7b28be79dd52282 Mon Sep 17 00:00:00 2001 From: Ariel Frischer Date: Sat, 15 Mar 2025 09:48:52 -0700 Subject: [PATCH] feat(agent_utils.py): add get_model_name_from_chat_model function to improve model handling refactor(build_agent_kwargs): simplify state modifier logic by using model name instead of model attribute --- ra_aid/agent_utils.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/ra_aid/agent_utils.py b/ra_aid/agent_utils.py index d952c2c..e4be3e8 100644 --- a/ra_aid/agent_utils.py +++ b/ra_aid/agent_utils.py @@ -52,6 +52,7 @@ from ra_aid.database.repositories.human_input_repository import ( from ra_aid.database.repositories.trajectory_repository import get_trajectory_repository from ra_aid.database.repositories.config_repository import get_config_repository from ra_aid.anthropic_token_limiter import ( + get_model_name_from_chat_model, sonnet_35_state_modifier, state_modifier, get_model_token_limit, @@ -102,11 +103,10 @@ def build_agent_kwargs( ): def wrapped_state_modifier(state: AgentState) -> list[BaseMessage]: - if not hasattr(model, 'model'): - return state_modifier(state, model, max_input_tokens=max_input_tokens) + model_name = get_model_name_from_chat_model(model) if any( - pattern in model.model + pattern in model_name for pattern in ["claude-3.5", "claude3.5", "claude-3-5"] ): return sonnet_35_state_modifier(