app.py

import os
import gradio as gr
import requests
import json
import logging
import google.generativeai as genai
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# API Keys configuration
COHERE_API_KEY = os.getenv("COHERE_API_KEY")
MISTRAL_API_KEY = os.getenv("MISTRAL_API_KEY")
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")

if not all([COHERE_API_KEY, MISTRAL_API_KEY, GEMINI_API_KEY]):
    raise ValueError("Missing required API keys in environment variables")

# Configure Gemini
genai.configure(api_key=GEMINI_API_KEY)

# 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 = "onramps"

# Model configurations
MODELS = {
    "Cohere": {
        "name": "command-r-plus-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
    },
    "Gemini": {
        "name": "gemini-1.5-pro",
        "model": genai.GenerativeModel('gemini-1.5-pro'),
        "api_key": GEMINI_API_KEY
    }
}

def search_document(query, k):
    try:
        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)
        response.raise_for_status()
        return response.json(), "", k
    except requests.exceptions.RequestException as e:
        logging.error(f"Error in search: {e}")
        return {"error": str(e)}, query, k

def generate_answer_cohere(question, context, citations):
    headers = {
        "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": MODELS['Cohere']['name'],
        "preamble": "You are an AI-assistant chatbot. Provide thorough responses with citations.",
        "chat_history": []
    }
    
    try:
        response = requests.post(MODELS['Cohere']['api_url'], headers=headers, json=payload)
        response.raise_for_status()
        answer = response.json()['text']
        
        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_cohere: {e}")
        return f"An error occurred: {str(e)}"

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 and any pre-trained knowledge. 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 generate_answer_gemini(question, context, citations):
    prompt = f"Context: {context}\n\nQuestion: {question}\n\nAnswer the question based on the given context. Include citations as [1], [2], etc.:"
    
    try:
        model = MODELS['Gemini']['model']
        response = model.generate_content(
            prompt,
            generation_config=genai.types.GenerationConfig(
                temperature=1.0,
                top_k=40,
                top_p=0.95,
                max_output_tokens=8192,
            )
        )
        
        answer = response.text
        
        answer += "\n\nSources:"
        for i, citation in enumerate(citations, 1):
            answer += f"\n[{i}] {citation}"
        
        return answer
    except Exception as e:
        logging.error(f"Error in generate_answer_gemini: {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)
    
    # Extract and combine the retrieved contexts
    if "results" in search_results:
        contexts = []
        citations = []
        for item in search_results['results']:
            contexts.append(item['metadata']['content'])
            citations.append(f"{item['metadata'].get('title', 'Unknown Source')} - {item['metadata'].get('source', 'No source provided')}")
        combined_context = " ".join(contexts)
    else:
        logging.error(f"Error in database search or no results found: {search_results}")
        combined_context = ""
        citations = []
    
    # Generate answer using the selected model
    if model_choice == "Cohere":
        return generate_answer_cohere(question, combined_context, citations)
    elif model_choice == "Mistral":
        return generate_answer_mistral(question, combined_context, citations)
    else:
        return generate_answer_gemini(question, combined_context, citations)

def chatbot(message, history, model_choice):
    response = answer_question(message, model_choice)
    return response

# Example questions with default model choice
EXAMPLE_QUESTIONS = [
    ["Why was Anne Hutchinson banished from Massachusetts?", "Cohere"],
    ["What were the major causes of World War I?", "Mistral"],
    ["Who was the first President of the United States?", "Gemini"],
    ["What was the significance of the Industrial Revolution?", "Cohere"]
]

# Create Gradio interface
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", "Gemini"],
            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=EXAMPLE_QUESTIONS,
        cache_examples=False,
        retry_btn=None,
        undo_btn="Delete Previous",
        clear_btn="Clear",
    )

# Launch the app
if __name__ == "__main__":
    iface.launch()