xiaxinhang/pr-agent
1
0
Fork 0
mirror of https://github.com/qodo-ai/pr-agent.git synced 2025-07-03 12:20:38 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #1509 from KennyDizi/main

Browse Source 
Define user_message_only_models list for using user prompt only model
This commit is contained in:
Tal
2025-02-02 16:14:05 +02:00
committed by GitHub
parent dc21f71ee2 a6482c37ce
commit d8fba02b48
2 changed files with 18 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
pr_agent/algo/__init__.py
View File

@ -26,6 +26,8 @@ MAX_TOKENS = {
'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 @@ MAX_TOKENS = {
"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"
]

10
pr_agent/algo/ai_handlers/litellm_ai_handler.py
View File

@ -6,6 +6,7 @@ import requests
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 @@ class LiteLLMAIHandler(BaseAiHandler):
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 @@ class LiteLLMAIHandler(BaseAiHandler):
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")