Update app.py

app.py

@@ -2,6 +2,8 @@ import os
 import gradio as gr
 import requests
 import pandas as pd
+import json
+import traceback
 from dotenv import load_dotenv
 from typing import List, Dict, Any, Optional
 
@@ -24,28 +26,54 @@ class MessagesState(TypedDict):
 try:
     with open("system_prompt.txt", "r", encoding="utf-8") as f:
         system_prompt = f.read()
+        print("Successfully loaded system prompt from file")
 except FileNotFoundError:
     system_prompt = (
         "You are a helpful AI assistant that uses tools to find information and answer questions.\n"
         "When you don't know something, use the available tools to look up information. Be concise, direct, and provide accurate responses.\n"
         "Always cite your sources when using information from searches or reference materials."
     )
+    print("System prompt file not found, using default prompt")
 
 class AdvancedAgent:
     def __init__(self):
-        print("Initializing AdvancedAgent with LangGraph, Wikipedia, Arxiv, and Gemini 2.0 Flash")
+        print("Initializing AdvancedAgent with LangGraph, Wikipedia, Arxiv, and LLM via OpenRouter")
         load_dotenv()
-        self.graph = self.build_graph()
-        print("Graph successfully built")
+        
+        # Check if OpenRouter API key is set
+        if not os.getenv("OPENROUTER_API_KEY"):
+            print("WARNING: OPENROUTER_API_KEY environment variable not set!")
+            print("Please set this in your HF Space secrets or .env file")
+            
+        try:
+            self.graph = self.build_graph()
+            print("Graph successfully built")
+        except Exception as e:
+            print(f"Error building agent graph: {e}")
+            traceback.print_exc()
+            raise
         
     def build_graph(self):
+        # Add consistent headers for OpenRouter
+        headers = {
+            "HTTP-Referer": "https://huggingface.co/",
+            "X-Title": "HF Agent"
+        }
+        
+    try:
+        # Try alternative OpenAI model if Gemini is causing issues
+        model_name = "google/gemini-2.0-flash-001"
+        # Fallback to OpenAI if needed
+        # model_name = "gpt-3.5-turbo" 
+        
         llm = ChatOpenAI(
-            model="google/gemini-2.0-flash-001",
+            model=model_name,
             temperature=0,
             openai_api_key=os.getenv("OPENROUTER_API_KEY"),
-            openai_api_base="https://openrouter.ai/api/v1"
+            openai_api_base="https://openrouter.ai/api/v1",
+            headers=headers
         )
-        print("LLM initialized: Gemini 2.5 Pro via OpenRouter")
+        print(f"LLM initialized: {model_name} via OpenRouter")
         
         wikipedia_tool = WikipediaQueryRun(api_wrapper=WikipediaAPIWrapper())
         arxiv_tool = ArxivQueryRun()
@@ -86,29 +114,44 @@ class AdvancedAgent:
             SystemMessage(content=system_prompt),
             HumanMessage(content=question)
         ]
+        
+        # Print message details for debugging
+        print(f"Sending {len(messages)} messages to the LLM")
+        
         try:
             # Initial state must be a dict with "messages" key!
+            print("Invoking agent graph...")
             result = self.graph.invoke({"messages": messages})
+            print("Agent graph execution completed")
+            
             final_messages = result["messages"]
+            print(f"Got {len(final_messages)} messages in result")
             
             # Extract only AI messages from the final state
             ai_messages = [msg for msg in final_messages if isinstance(msg, AIMessage)]
+            print(f"Found {len(ai_messages)} AI messages in the final state")
             
             if not ai_messages:
+                print("No AI messages found in the final state")
                 return "I wasn't able to generate a proper response. Please try again."
             
             # Get the last AI message - this should be the final answer
             last_message = ai_messages[-1]
             
             # Return the content of the last AI message
-            return last_message.content if last_message.content else "I wasn't able to generate a proper response. Please try again."
+            if last_message.content:
+                return last_message.content
+            else:
+                print("Last AI message has empty content")
+                return "I wasn't able to generate a proper response. Please try again."
             
         except Exception as e:
             print(f"Error running agent graph: {e}")
+            traceback.print_exc()
             return f"Sorry, I encountered an error while processing your question: {str(e)}"
 
 def run_and_submit_all(profile: gr.OAuthProfile | None):
-    space_id = os.getenv("SPACE_ID")
+    space_id = os.getenv("SPACE_ID", "i-dhilip/Final_Assignment_Template")  # Default to your space ID
     if profile:
         username = f"{profile.username}"
         print(f"User logged in: {username}")
@@ -119,12 +162,31 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
     api_url = DEFAULT_API_URL
     questions_url = f"{api_url}/questions"
     submit_url = f"{api_url}/submit"
+    
+    print(f"Checking OpenRouter API key...")
+    if not os.getenv("OPENROUTER_API_KEY"):
+        return "Error: OPENROUTER_API_KEY not set. Please add this to your Space secrets.", None
 
     try:
         agent = AdvancedAgent()
     except Exception as e:
-        print(f"Error instantiating agent: {e}")
-        return f"Error initializing agent: {e}", None
+        error_msg = f"Error instantiating agent: {e}"
+        print(error_msg)
+        traceback.print_exc()
+        return error_msg, None
+
+    # Generate a simple test question to verify the agent works
+    test_question = "What is the capital of France?"
+    print(f"Testing agent with a simple question: '{test_question}'")
+    try:
+        test_response = agent(test_question)
+        print(f"Test response: {test_response[:100]}...")
+        if "I wasn't able to generate" in test_response or "error" in test_response.lower():
+            print("WARNING: Agent test response indicates potential issues")
+    except Exception as e:
+        print(f"Agent test failed: {e}")
+        traceback.print_exc()
+        return f"Agent test failed: {e}", None
 
     agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main"
     print(agent_code)