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

feat: add support for O1 model by combining system and user prompts in litellm_ai_handler

Browse Source
This commit is contained in:
mrT23
2024-10-09 08:53:34 +03:00
parent 8e8062fefc
commit 727b08fde3
1 changed files with 23 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@ -171,6 +171,7 @@ class LiteLLMAIHandler(BaseAiHandler):
get_logger().warning(
"Empty system prompt for claude model. Adding a newline character to prevent OpenAI API error.")
messages = [{"role": "system", "content": system}, {"role": "user", "content": user}]
if img_path:
try:
# check if the image link is alive
@ -185,14 +186,28 @@ class LiteLLMAIHandler(BaseAiHandler):
messages[1]["content"] = [{"type": "text", "text": messages[1]["content"]},
{"type": "image_url", "image_url": {"url": img_path}}]
kwargs = {
"model": model,
"deployment_id": deployment_id,
"messages": messages,
"temperature": temperature,
"timeout": get_settings().config.ai_timeout,
"api_base": self.api_base,
}
# Currently O1 does not support separate system and user prompts
if model.startswith('o1-'):
user = f"{system}\n\n\n{user}"
system = ""
get_logger().info(f"Using O1 model, combining system and user prompts")
messages = [{"role": "user", "content": user}]
kwargs = {
"model": model,
"deployment_id": deployment_id,
"messages": messages,
"timeout": get_settings().config.ai_timeout,
"api_base": self.api_base,
}
else:
kwargs = {
"model": model,
"deployment_id": deployment_id,
"messages": messages,
"temperature": temperature,
"timeout": get_settings().config.ai_timeout,
"api_base": self.api_base,
}
if get_settings().litellm.get("enable_callbacks", False):
kwargs = self.add_litellm_callbacks(kwargs)