refac: convert_payload_openai_to_ollama

This commit is contained in:
Timothy J. Baek 2024-10-21 04:10:28 -07:00
parent f99facf383
commit b4e73c7f19

View File

@ -88,6 +88,50 @@ def apply_model_params_to_body_ollama(params: dict, form_data: dict) -> dict:
return form_data
def convert_messages_openai_to_ollama(messages: list[dict]) -> list[dict]:
ollama_messages = []
for message in messages:
# Initialize the new message structure with the role
new_message = {"role": message["role"]}
content = message.get("content", [])
# Check if the content is a string (just a simple message)
if isinstance(content, str):
# If the content is a string, it's pure text
new_message["content"] = content
else:
# Otherwise, assume the content is a list of dicts, e.g., text followed by an image URL
content_text = ""
images = []
# Iterate through the list of content items
for item in content:
# Check if it's a text type
if item.get("type") == "text":
content_text += item.get("text", "")
# Check if it's an image URL type
elif item.get("type") == "image_url":
img_url = item.get("image_url", {}).get("url", "")
if img_url:
images.append(img_url)
# Add content text (if any)
if content_text:
new_message["content"] = content_text.strip()
# Add images (if any)
if images:
new_message["images"] = images
# Append the new formatted message to the result
ollama_messages.append(new_message)
return ollama_messages
def convert_payload_openai_to_ollama(openai_payload: dict) -> dict:
"""
Converts a payload formatted for OpenAI's API to be compatible with Ollama's API endpoint for chat completions.
@ -102,7 +146,9 @@ def convert_payload_openai_to_ollama(openai_payload: dict) -> dict:
# Mapping basic model and message details
ollama_payload["model"] = openai_payload.get("model")
ollama_payload["messages"] = openai_payload.get("messages")
ollama_payload["messages"] = convert_messages_openai_to_ollama(
openai_payload.get("messages")
)
ollama_payload["stream"] = openai_payload.get("stream", False)
# If there are advanced parameters in the payload, format them in Ollama's options field