Update app.py

app.py

@@ -8,21 +8,36 @@ from dotenv import load_dotenv
 # Load environment variables
 load_dotenv()
 
-# Cohere API configuration
+# API Keys configuration
 COHERE_API_KEY = os.getenv("COHERE_API_KEY")
-if not COHERE_API_KEY:
-    raise ValueError("COHERE_API_KEY not found in environment variables")
+MISTRAL_API_KEY = os.getenv("MISTRAL_API_KEY")
 
-COHERE_API_URL = "https://api.cohere.ai/v1/chat"
-MODEL_NAME = "command-r-08-2024"
+if not COHERE_API_KEY or not MISTRAL_API_KEY:
+    raise ValueError("Missing required API keys in environment variables")
 
-# Vector database configuration
-API_URL = "https://sendthat.cc"
+# API endpoints configuration
+COHERE_API_URL = "https://api.cohere.ai/v1/chat"
+MISTRAL_API_URL = "https://api.mistral.ai/v1/chat/completions"
+VECTOR_API_URL = "https://sendthat.cc"
 HISTORY_INDEX = "history"
 
+# Model configurations
+MODELS = {
+    "Cohere": {
+        "name": "command-r-08-2024",
+        "api_url": COHERE_API_URL,
+        "api_key": COHERE_API_KEY
+    },
+    "Mistral": {
+        "name": "ft:open-mistral-nemo:ef730d29:20241022:2a0e7d46",
+        "api_url": MISTRAL_API_URL,
+        "api_key": MISTRAL_API_KEY
+    }
+}
+
 def search_document(query, k):
     try:
-        url = f"{API_URL}/search/{HISTORY_INDEX}"
+        url = f"{VECTOR_API_URL}/search/{HISTORY_INDEX}"
         payload = {"text": query, "k": k}
         headers = {"Content-Type": "application/json"}
         response = requests.post(url, json=payload, headers=headers)
@@ -32,37 +47,69 @@ def search_document(query, k):
         logging.error(f"Error in search: {e}")
         return {"error": str(e)}, query, k
 
-def generate_answer(question, context, citations):
-    prompt = f"Context: {context}\n\nQuestion: {question}\n\nAnswer the question based on the given context. At the end of your answer, provide citations for the sources you used, referencing them as [1], [2], etc.:"
-    
+def generate_answer_cohere(question, context, citations):
     headers = {
-        "Authorization": f"Bearer {COHERE_API_KEY}",
+        "Authorization": f"Bearer {MODELS['Cohere']['api_key']}",
         "Content-Type": "application/json"
     }
     
+    prompt = f"Context: {context}\n\nQuestion: {question}\n\nAnswer the question based on the given context. Include citations as [1], [2], etc.:"
+    
     payload = {
         "message": prompt,
-        "model": MODEL_NAME,
-        "preamble": "You are an AI-assistant chatbot. You are trained to assist users by providing thorough and helpful responses to their queries based on the given context. Always include citations at the end of your answer.",
-        "chat_history": []  # You can add chat history here if needed
+        "model": MODELS['Cohere']['name'],
+        "preamble": "You are an AI-assistant chatbot. Provide thorough responses with citations.",
+        "chat_history": []
     }
     
     try:
-        response = requests.post(COHERE_API_URL, headers=headers, json=payload)
+        response = requests.post(MODELS['Cohere']['api_url'], headers=headers, json=payload)
         response.raise_for_status()
         answer = response.json()['text']
         
-        # Append citations to the answer
         answer += "\n\nSources:"
         for i, citation in enumerate(citations, 1):
             answer += f"\n[{i}] {citation}"
         
         return answer
     except requests.exceptions.RequestException as e:
-        logging.error(f"Error in generate_answer: {e}")
+        logging.error(f"Error in generate_answer_cohere: {e}")
         return f"An error occurred: {str(e)}"
 
-def answer_question(question, k=3):
+def generate_answer_mistral(question, context, citations):
+    headers = {
+        "Authorization": f"Bearer {MODELS['Mistral']['api_key']}",
+        "Content-Type": "application/json",
+        "Accept": "application/json"
+    }
+    
+    prompt = f"Context: {context}\n\nQuestion: {question}\n\nAnswer the question based on the given context. Include citations as [1], [2], etc.:"
+    
+    payload = {
+        "model": MODELS['Mistral']['name'],
+        "messages": [
+            {
+                "role": "user",
+                "content": prompt
+            }
+        ]
+    }
+    
+    try:
+        response = requests.post(MODELS['Mistral']['api_url'], headers=headers, json=payload)
+        response.raise_for_status()
+        answer = response.json()['choices'][0]['message']['content']
+        
+        answer += "\n\nSources:"
+        for i, citation in enumerate(citations, 1):
+            answer += f"\n[{i}] {citation}"
+        
+        return answer
+    except requests.exceptions.RequestException as e:
+        logging.error(f"Error in generate_answer_mistral: {e}")
+        return f"An error occurred: {str(e)}"
+
+def answer_question(question, model_choice, k=3):
     # Search the vector database
     search_results, _, _ = search_document(question, k)
     
@@ -79,33 +126,45 @@ def answer_question(question, k=3):
         combined_context = ""
         citations = []
     
-    # Generate answer using the Cohere LLM
-    answer = generate_answer(question, combined_context, citations)
-    return answer
+    # Generate answer using the selected model
+    if model_choice == "Cohere":
+        return generate_answer_cohere(question, combined_context, citations)
+    else:
+        return generate_answer_mistral(question, combined_context, citations)
 
-def chatbot(message, history):
-    response = answer_question(message)
+def chatbot(message, history, model_choice):
+    response = answer_question(message, model_choice)
     return response
 
 # Create Gradio interface
-iface = gr.ChatInterface(
-    chatbot,
-    chatbot=gr.Chatbot(height=300),
-    textbox=gr.Textbox(placeholder="Ask a question about history...", container=False, scale=7),
-    title="History Chatbot",
-    description="Ask me anything about history, and I'll provide answers with citations!",
-    theme="soft",
-    examples=[
-        "Why was Anne Hutchinson banished from Massachusetts?",
-        "What were the major causes of World War I?",
-        "Who was the first President of the United States?",
-        "What was the significance of the Industrial Revolution?"
-    ],
-    cache_examples=False,
-    retry_btn=None,
-    undo_btn="Delete Previous",
-    clear_btn="Clear",
-)
+with gr.Blocks(theme="soft") as iface:
+    gr.Markdown("# History Chatbot")
+    gr.Markdown("Ask me anything about history, and I'll provide answers with citations!")
+    
+    with gr.Row():
+        model_choice = gr.Radio(
+            choices=["Cohere", "Mistral"],
+            value="Cohere",
+            label="Choose LLM Model",
+            info="Select which AI model to use for generating responses"
+        )
+    
+    chatbot_interface = gr.ChatInterface(
+        fn=lambda message, history, model: chatbot(message, history, model),
+        additional_inputs=[model_choice],
+        chatbot=gr.Chatbot(height=300),
+        textbox=gr.Textbox(placeholder="Ask a question about history...", container=False, scale=7),
+        examples=[
+            "Why was Anne Hutchinson banished from Massachusetts?",
+            "What were the major causes of World War I?",
+            "Who was the first President of the United States?",
+            "What was the significance of the Industrial Revolution?"
+        ],
+        cache_examples=False,
+        retry_btn=None,
+        undo_btn="Delete Previous",
+        clear_btn="Clear",
+    )
 
 # Launch the app
 if __name__ == "__main__":