diff --git a/pr_agent/algo/__init__.py b/pr_agent/algo/__init__.py index 40d4fc075..10b53cb38 100644 --- a/pr_agent/algo/__init__.py +++ b/pr_agent/algo/__init__.py @@ -26,6 +26,8 @@ 'o1-preview-2024-09-12': 128000, # 128K, but may be limited by config.max_model_tokens 'o1-2024-12-17': 204800, # 200K, but may be limited by config.max_model_tokens 'o1': 204800, # 200K, but may be limited by config.max_model_tokens + 'o3-mini': 204800, # 200K, but may be limited by config.max_model_tokens + 'o3-mini-2025-01-31': 204800, # 200K, but may be limited by config.max_model_tokens 'claude-instant-1': 100000, 'claude-2': 100000, 'command-nightly': 4096, @@ -81,3 +83,13 @@ "watsonx/ibm/granite-34b-code-instruct": 8191, "watsonx/mistralai/mistral-large": 32768, } + +USER_MESSAGE_ONLY_MODELS = [ + "deepseek/deepseek-reasoner", + "o1-mini", + "o1-mini-2024-09-12", + "o1", + "o1-2024-12-17", + "o3-mini", + "o3-mini-2025-01-31" +] diff --git a/pr_agent/algo/ai_handlers/litellm_ai_handler.py b/pr_agent/algo/ai_handlers/litellm_ai_handler.py index 50d14b217..c936ebd5f 100644 --- a/pr_agent/algo/ai_handlers/litellm_ai_handler.py +++ b/pr_agent/algo/ai_handlers/litellm_ai_handler.py @@ -6,6 +6,7 @@ from litellm import acompletion from tenacity import retry, retry_if_exception_type, stop_after_attempt +from pr_agent.algo import USER_MESSAGE_ONLY_MODELS from pr_agent.algo.ai_handlers.base_ai_handler import BaseAiHandler from pr_agent.algo.utils import get_version from pr_agent.config_loader import get_settings @@ -94,6 +95,9 @@ def __init__(self): if get_settings().get("DEEPSEEK.KEY", None): os.environ['DEEPSEEK_API_KEY'] = get_settings().get("DEEPSEEK.KEY") + # Models that only use user meessage + self.user_message_only_models = USER_MESSAGE_ONLY_MODELS + def prepare_logs(self, response, system, user, resp, finish_reason): response_log = response.dict().copy() response_log['system'] = system @@ -197,10 +201,8 @@ async def chat_completion(self, model: str, system: str, user: str, temperature: messages[1]["content"] = [{"type": "text", "text": messages[1]["content"]}, {"type": "image_url", "image_url": {"url": img_path}}] - # Currently, model OpenAI o1 series does not support a separate system and user prompts - O1_MODEL_PREFIX = 'o1' - model_type = model.split('/')[-1] if '/' in model else model - if (model_type.startswith(O1_MODEL_PREFIX)) or ("deepseek-reasoner" in model): + # Currently, some models do not support a separate system and user prompts + if self.user_message_only_models and any(entry.lower() in model.lower() for entry in self.user_message_only_models): user = f"{system}\n\n\n{user}" system = "" get_logger().info(f"Using model {model}, combining system and user prompts")