測試 Mistral API
這篇文章是快速測試 Mistral API 的。我使用 mistral-small-2501 模型進行這次測試。下面的程式碼顯示如何呼叫 API 並獲取回應。
import os
import requests
from dotenv import load_dotenv
import argparse
load_dotenv()
def call_mistral_api(prompt, model="mistral-small-2501"):
api_key = os.environ.get("MISTRAL_API_KEY")
if not api_key:
print("錯誤：MISTRAL_API_KEY 環境變數未設定。")
return None
url = "https://api.mistral.ai/v1/chat/completions"
headers = {
"Content-Type": "application/json",
"Accept": "application/json",
"Authorization": f"Bearer {api_key}"
}
data = {
"model": model,
"messages": [
{
"role": "user",
"content": prompt
}
]
}
try:
response = requests.post(url, headers=headers, json=data)
response.raise_for_status()
response_json = response.json()
print(response_json)
if response_json and response_json['choices']:
content = response_json['choices'][0]['message']['content']
return content
else:
print(f"Mistral API 錯誤：無效的回應格式：{response_json}")
return None
except requests.exceptions.RequestException as e:
print(f"Mistral API 錯誤：{e}")
if e.response:
print(f"回應狀態碼：{e.response.status_code}")
print(f"回應內容：{e.response.text}")
return None
def call_codestral_api(prompt, model="codestral-latest"):
api_key = os.environ.get("MISTRAL_API_KEY")
if not api_key:
print("錯誤：MISTRAL_API_KEY 環境變數未設定。")
return None
url = "https://api.mistral.ai/v1/fim/completions"
headers = {
"Content-Type": "application/json",
"Accept": "application/json",
"Authorization": f"Bearer {api_key}"
}
data = {
"model": model,
"prompt": prompt,
"suffix": "return a + b",
"max_tokens": 64,
"temperature": 0
}
try:
response = requests.post(url, headers=headers, json=data)
response.raise_for_status()
response_json = response.json()
print(response_json)
if response_json and response_json['choices']:
content = response_json['choices'][0]['message']['content']
return content
else:
print(f"Codestral API 錯誤：無效的回應格式：{response_json}")
return None
except requests.exceptions.RequestException as e:
print(f"Codestral API 錯誤：{e}")
if e.response:
print(f"回應狀態碼：{e.response.status_code}")
print(f"回應內容：{e.response.text}")
return None
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="測試 Mistral 和 Codestral API。")
parser.add_argument("--type", type=str, default="mistral", choices=["mistral", "codestral"], help="呼叫的 API 類型 (mistral 或 codestral)")
args = parser.parse_args()
if args.type == "mistral":
prompt = "what's your cutting point of knowledge"
response = call_mistral_api(prompt)
if response:
print(f"回應：{response}")
elif args.type == "codestral":
prompt = "def f("
response = call_codestral_api(prompt, model="codestral-latest")
if response:
print(f"回應：{response}")