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

This commit is contained in:
mrT23
2024-10-09 08:53:34 +03:00
parent 8e8062fefc
commit 727b08fde3

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,6 +186,20 @@ class LiteLLMAIHandler(BaseAiHandler):
messages[1]["content"] = [{"type": "text", "text": messages[1]["content"]},
{"type": "image_url", "image_url": {"url": img_path}}]
# 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,