From b7437147af1fd9d12d6aacdff44b769a7f694509 Mon Sep 17 00:00:00 2001 From: mrT23 Date: Sat, 19 Oct 2024 11:32:45 +0300 Subject: [PATCH] fix: correct model type extraction for O1 model handling in litellm_ai_handler.py --- pr_agent/algo/ai_handlers/litellm_ai_handler.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pr_agent/algo/ai_handlers/litellm_ai_handler.py b/pr_agent/algo/ai_handlers/litellm_ai_handler.py index 12bdb4f7..3127849e 100644 --- a/pr_agent/algo/ai_handlers/litellm_ai_handler.py +++ b/pr_agent/algo/ai_handlers/litellm_ai_handler.py @@ -188,7 +188,8 @@ class LiteLLMAIHandler(BaseAiHandler): # Currently O1 does not support separate system and user prompts O1_MODEL_PREFIX = 'o1-' - if model.startswith(O1_MODEL_PREFIX): + model_type = model.split('/')[-1] # 'azure/o1-' or 'o1-' + if model_type.startswith(O1_MODEL_PREFIX): user = f"{system}\n\n\n{user}" system = "" get_logger().info(f"Using O1 model, combining system and user prompts")