From 5e26b6f0ff71412bd9f0ff2e0bd6b573e9723ebd Mon Sep 17 00:00:00 2001
From: fyq163 <49147332+fyq163@users.noreply.github.com>
Date: Sat, 15 Feb 2025 01:22:29 +0800
Subject: [PATCH] Update image format correction

Some model has strict requirements matching of  image format and codes
---
 tests/test_llm_api.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tests/test_llm_api.py b/tests/test_llm_api.py
index 9671900..1eb45f4 100644
--- a/tests/test_llm_api.py
+++ b/tests/test_llm_api.py
@@ -22,15 +22,14 @@ class LLMConfig:
 
 def create_message_content(text, image_path=None):
     content = [{"type": "text", "text": text}]
-
+    image_format = "png" if image_path and image_path.endswith(".png") else "jpeg"
     if image_path:
         from src.utils import utils
         image_data = utils.encode_image(image_path)
         content.append({
             "type": "image_url",
-            "image_url": {"url": f"data:image/jpeg;base64,{image_data}"}
+            "image_url": {"url": f"data:image/{image_format};base64,{image_data}"}
         })
-
     return content
 
 def get_env_value(key, provider):
@@ -41,6 +40,7 @@ def get_env_value(key, provider):
         "deepseek": {"api_key": "DEEPSEEK_API_KEY", "base_url": "DEEPSEEK_ENDPOINT"},
         "mistral": {"api_key": "MISTRAL_API_KEY", "base_url": "MISTRAL_ENDPOINT"},
         "alibaba": {"api_key": "ALIBABA_API_KEY", "base_url": "ALIBABA_ENDPOINT"},
+        "moonshot":{"api_key": "MOONSHOT_API_KEY", "base_url": "MOONSHOT_ENDPOINT"},
     }
 
     if provider in env_mappings and key in env_mappings[provider]: