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 = "history" # 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()