Testing Mistral API
This post is a quick test of the Mistral API. I’m using the mistral-small-2501 model for this test. The code below shows how to call the API and get a response.
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("Error: MISTRAL_API_KEY environment variable not set.")
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 Error: Invalid response format: {response_json}")
return None
except requests.exceptions.RequestException as e:
print(f"Mistral API Error: {e}")
if e.response:
print(f"Response status code: {e.response.status_code}")
print(f"Response content: {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("Error: MISTRAL_API_KEY environment variable not set.")
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 Error: Invalid response format: {response_json}")
return None
except requests.exceptions.RequestException as e:
print(f"Codestral API Error: {e}")
if e.response:
print(f"Response status code: {e.response.status_code}")
print(f"Response content: {e.response.text}")
return None
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Test Mistral and Codestral APIs.")
parser.add_argument("--type", type=str, default="mistral", choices=["mistral", "codestral"], help="Type of API to call (mistral or 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: {response}")
elif args.type == "codestral":
prompt = "def f("
response = call_codestral_api(prompt, model="codestral-latest")
if response:
print(f"Response: {response}")