Update app.py

app.py

@@ -1,75 +1,106 @@
 import os
 import time
+import json
 import gradio as gr
-from huggingface_hub import InferenceClient
+from openai import OpenAI
 from dotenv import load_dotenv
 
 # Load API keys from .env file
 load_dotenv()
-HF_API_KEY = os.getenv("HF_API_KEY")  # Hugging Face API Key
-TOGETHER_API_KEY = os.getenv("TOGETHER_API_KEY")  # Together AI API Key
-
-# Initialize clients
-gemma_client = InferenceClient(provider="hf-inference", api_key=HF_API_KEY)
-together_client = InferenceClient(provider="together", api_key=TOGETHER_API_KEY)
-
-# Function to query Hugging Face (Gemma)
-def query_gemma(user_input):
-    messages = [{"role": "user", "content": user_input}]
-    completion = gemma_client.chat.completions.create(
-        model="google/gemma-2-27b-it",
-        messages=messages,
-        max_tokens=500
-    )
-    return completion.choices[0].message["content"]
-
-# Function to query Together (Llama)
-def query_llama(user_input):
-    messages = [{"role": "user", "content": user_input}]
-    completion = together_client.chat.completions.create(
-        model="meta-llama/Llama-3.3-70B-Instruct",
-        messages=messages,
-        max_tokens=500
-    )
-    return completion.choices[0].message["content"]
-
-# Function to query Together (DeepSeek)
-def query_deepseek(user_input):
-    messages = [{"role": "user", "content": user_input}]
-    completion = together_client.chat.completions.create(
-        model="deepseek-ai/DeepSeek-V3",
-        messages=messages,
-        max_tokens=500
-    )
-    return completion.choices[0].message["content"]
-
-# Function to refine responses using DeepSeek-V3
+OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY")  # OpenRouter API Key
+
+# Initialize OpenRouter Client
+openrouter_client = OpenAI(
+    base_url="https://openrouter.ai/api/v1",
+    api_key=OPENROUTER_API_KEY
+)
+
+# Query OpenRouter (Gemma-2-9B)
+def query_gemma_openrouter(user_input):
+    try:
+        completion = openrouter_client.chat.completions.create(
+            model="google/gemma-2-9b-it:free",
+            messages=[{"role": "user", "content": user_input}]
+        )
+        return completion.choices[0].message.content
+    except Exception as e:
+        return f"Error querying Gemma-2-9B: {str(e)}"
+
+# Query OpenRouter (Llama-3.2-3B)
+def query_llama_openrouter(user_input):
+    try:
+        completion = openrouter_client.chat.completions.create(
+            model="meta-llama/llama-3.2-3b-instruct:free",
+            messages=[{"role": "user", "content": user_input}]
+        )
+        return completion.choices[0].message.content
+    except Exception as e:
+        return f"Error querying Llama-3.2-3B: {str(e)}"
+
+# Query OpenRouter (DeepSeek-R1)
+def query_deepseek_openrouter(user_input):
+    try:
+        completion = openrouter_client.chat.completions.create(
+            model="deepseek/deepseek-r1:free",
+            messages=[{"role": "user", "content": user_input}]
+        )
+        return completion.choices[0].message.content
+    except Exception as e:
+        return f"Error querying DeepSeek-R1: {str(e)}"
+
+# Function to refine responses using DeepSeek-R1
 def refine_response(user_input):
     try:
-        # Get responses
-        gemma_response = query_gemma(user_input)
-        llama_response = query_llama(user_input)
-        deepseek_response = query_deepseek(user_input)
+        # Get responses from all three models
+        gemma_response = query_gemma_openrouter(user_input)
+        llama_response = query_llama_openrouter(user_input)
+        deepseek_response = query_deepseek_openrouter(user_input)
+
+        # If any response is missing, return the available ones
+        responses = {
+            "Gemma": gemma_response.strip(),
+            "Llama": llama_response.strip(),
+            "DeepSeek-R1": deepseek_response.strip()
+        }
+        valid_responses = {k: v for k, v in responses.items() if v}
+
+        if len(valid_responses) < 2:
+            return "\n\n".join(f"{k} Response: {v}" for k, v in valid_responses.items())
 
         # Prepare refinement prompt
         improvement_prompt = f"""
         Here are three AI-generated responses:
 
         Response 1 (Gemma): {gemma_response}
-        Response 2 (Llama 3.3): {llama_response}
-        Response 3 (DeepSeek-V3): {deepseek_response}
+        Response 2 (Llama 3.2): {llama_response}
+        Response 3 (DeepSeek-R1): {deepseek_response}
 
-        Please combine the best elements of all three and provide an improved answer.
+        Please combine the best elements of all three, improve clarity, and provide a final refined answer.
         """
 
-        messages = [{"role": "user", "content": improvement_prompt}]
-        refined_completion = together_client.chat.completions.create(
-            model="deepseek-ai/DeepSeek-V3",
-            messages=messages,
-            max_tokens=500
-        )
+        # Retry loop for DeepSeek-R1 refinement
+        max_retries = 3
+        for attempt in range(max_retries):
+            try:
+                messages = [{"role": "user", "content": improvement_prompt}]
+                refined_completion = openrouter_client.chat.completions.create(
+                    model="deepseek/deepseek-r1:free",
+                    messages=messages
+                )
+
+                refined_content = refined_completion.choices[0].message.content
+
+                if refined_content.strip():
+                    return refined_content
+                else:
+                    print("Received empty response from DeepSeek-R1, retrying...")
+                    time.sleep(2)
+
+            except Exception as e:
+                print(f"Error on attempt {attempt + 1}: {str(e)}")
+                time.sleep(2)
 
-        return refined_completion.choices[0].message["content"]
+        return f"Refinement failed. Here’s the best available response:\n\n{max(valid_responses.values(), key=len)}"
 
     except Exception as e:
         return f"Error refining response: {str(e)}"
@@ -80,7 +111,7 @@ iface = gr.Interface(
     inputs=gr.Textbox(lines=2, placeholder="Ask me anything..."), 
     outputs="text",
     title="Multi-Model AI Enhancer",
-    description="Get responses from Gemma, Llama, and DeepSeek. Then receive an improved answer."
+    description="Get responses from Gemma, Llama-3.2, and DeepSeek-R1. Then receive an improved final answer."
 )
 
 # Launch app