Create mistral_api.py

mistral_api.py

@@ -0,0 +1,50 @@
+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_completion_tokens=800) -> str:
+    """Chat completion using Mistral models.
+    https://docs.mistral.ai/capabilities/completion/ """
+    response = client.chat.complete(
+        model=model,
+        messages=messages,
+        max_completion_tokens=max_completion_tokens,
+        temperature=temperature,
+        # stream=True
+    )
+    return response.choices[0].message.content