Update app.py

app.py

@@ -25,6 +25,9 @@ import requests
 import random
 import datetime
 from groq import Groq
+import os
+from mistralai import Mistral
+from dotenv import load_dotenv
 
 # Automatically get the current year
 current_year = datetime.datetime.now().year
@@ -53,6 +56,10 @@ GROQ_API_KEY = os.getenv("GROQ_API_KEY")
 # Initialize Groq client
 groq_client = Groq(api_key=GROQ_API_KEY)
 
+# Initialize Mistral client
+MISTRAL_API_KEY = os.getenv("MISTRAL_API_KEY")
+mistral_client = Mistral(api_key=MISTRAL_API_KEY)
+
 # Initialize the similarity model
 similarity_model = SentenceTransformer('all-MiniLM-L6-v2')
 
@@ -387,7 +394,18 @@ Instructions:
                 stream=False
             )
             return response.choices[0].message.content.strip()
-        else:
+        elif model == "mistral":
+            response = mistral_client.chat.complete(
+                model="Mistral-Nemo-Instruct-2407",
+                messages=messages,
+                max_tokens=5500,
+                temperature=temperature,
+                top_p=0.9,
+                presence_penalty=1.2,
+                stream=False
+            )
+            return response.choices[0].message.content.strip()
+        else:  # huggingface
             response = client.chat_completion(
                 messages=messages,
                 max_tokens=10000,
@@ -620,7 +638,7 @@ iface = gr.ChatInterface(
         gr.Slider(0, 2, value=2, step=1, label="Safe Search Level"),
         gr.Radio(["GET", "POST"], value="POST", label="HTTP Method"),
         gr.Slider(0, 1, value=0.2, step=0.1, label="LLM Temperature"),
-        gr.Dropdown(["huggingface", "groq"], value="huggingface", label="LLM Model"),
+        gr.Dropdown(["huggingface", "groq", "mistral"], value="huggingface", label="LLM Model"),
         gr.Checkbox(label="Use PyPDF2 for PDF scraping", value=True),
     ],
     additional_inputs_accordion=gr.Accordion("⚙️ Advanced Parameters", open=True),