import os import gradio as gr import requests import json import logging 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") if not COHERE_API_KEY or not MISTRAL_API_KEY: raise ValueError("Missing required API keys in environment variables") # 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"{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. 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) # 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) else: return generate_answer_mistral(question, combined_context, citations) def chatbot(message, history, model_choice): response = answer_question(message, model_choice) return response # 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"], 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__": iface.launch()