mistral_api.py

import os
from openai import OpenAI
from mistralai import Mistral
from file_utils import load_file

# Embedding API using OpenAI
# Ensure you set your API token in your environment variables or pass it directly
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")  # Make sure it's set in your environment

if OPENAI_API_KEY is None:
    raise ValueError("Hugging Face API token is missing. Set it as an environment variable: OPENAI_API_KEY")

# Initialize the OpenAI client using the API keys
embedder = OpenAI(api_key=OPENAI_API_KEY)

EMBEDDING_MODEL = "text-embedding-3-large" # "text-embedding-3-small" 

def get_embedding(text: str, model: str=EMBEDDING_MODEL) -> list[float]:
    """Get text embeddings from OpenAI."""
    result = embedder.embeddings.create(
      model=model,
      input=text
    )
    print("Got to get_embedding")
    return result.data[0].embedding

# Completions API using HF InferenceClient
# Ensure you set your API token in your environment variables or pass it directly
MISTRAL_API_KEY = os.getenv("MISTRAL_API_KEY")  # Make sure it's set in your environment
# api_key = os.environ["MISTRAL_API_KEY"]

if MISTRAL_API_KEY is None:
    raise ValueError("Hugging Face API token is missing. Set it as an environment variable: MISTRAL_API_KEY")

client = Mistral(api_key=MISTRAL_API_KEY)

COMPLETIONS_MODEL = "mistral-large-latest"

def get_response(messages: list[dict], model: str=COMPLETIONS_MODEL, 
                temperature=0, max_tokens=800) -> str:
    """Chat completion using Mistral models.
    https://docs.mistral.ai/capabilities/completion/ 
    https://docs.mistral.ai/api/#tag/chat
    """
    response = client.chat.complete(
        model=model,
        messages=messages,
        max_tokens=max_tokens,
        temperature=temperature,
        # stream=True
    )
    return response.choices[0].message.content