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 #1279 from Codium-ai/tr/o1_system

Browse Source 
feat: add support for O1 model by combining system and user prompt
This commit is contained in:
Tal
2024-10-09 08:59:40 +03:00
committed by GitHub
parent 5d9d48dc82 e6c56c7355
commit 3ac1ddc5d7
1 changed files with 24 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

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