xiaxinhang/pr-agent
1
0
Fork 0
mirror of https://github.com/qodo-ai/pr-agent.git synced 2025-07-02 11:50:37 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #1537 from KennyDizi/main

Browse Source 
Refactoring `LiteLLMAIHandler` class for better using reasoning models
This commit is contained in:
Tal
2025-02-16 19:50:03 +02:00
committed by GitHub
parent 778d00d1a0 adfc2a6b69
commit 52883fb1b5
2 changed files with 18 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
pr_agent/algo/__init__.py
View File

@ -85,11 +85,19 @@ MAX_TOKENS = {
}
USER_MESSAGE_ONLY_MODELS = [
"deepseek/deepseek-reasoner",
"o1-mini",
"o1-mini-2024-09-12",
"o1-preview"
]
NO_SUPPORT_TEMPERATURE_MODELS = [
"deepseek/deepseek-reasoner",
"o1-mini",
"o1-mini-2024-09-12",
"o1",
"o1-2024-12-17",
"o3-mini",
"o3-mini-2025-01-31"
"o3-mini-2025-01-31",
"o1-preview"
]

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

@ -6,7 +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 import NO_SUPPORT_TEMPERATURE_MODELS, 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
@ -98,6 +98,9 @@ class LiteLLMAIHandler(BaseAiHandler):
# Models that only use user meessage
self.user_message_only_models = USER_MESSAGE_ONLY_MODELS
# Model that doesn't support temperature argument
self.no_support_temperature_models = NO_SUPPORT_TEMPERATURE_MODELS
def prepare_logs(self, response, system, user, resp, finish_reason):
response_log = response.dict().copy()
response_log['system'] = system
@ -202,7 +205,7 @@ class LiteLLMAIHandler(BaseAiHandler):
{"type": "image_url", "image_url": {"url": img_path}}]
# 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):
if model 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")
@ -219,11 +222,14 @@ class LiteLLMAIHandler(BaseAiHandler):
"model": model,
"deployment_id": deployment_id,
"messages": messages,
"temperature": temperature,
"timeout": get_settings().config.ai_timeout,
"api_base": self.api_base,
}
# Add temperature only if model supports it
if model not in self.no_support_temperature_models:
kwargs["temperature"] = temperature
if get_settings().litellm.get("enable_callbacks", False):
kwargs = self.add_litellm_callbacks(kwargs)