Update app.py

app.py

@@ -1,177 +1,89 @@
-# app.py
-import os
-import gradio as gr
-import requests
-import pandas as pd
-from typing import Dict, List
-
-# custom imports
-from agents import Agent
-from tool import get_tools
-from model import get_model
-
-# --- Constants ---
-DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
-MODEL_ID = "vertex_ai/gemini-1.0-pro-001"
-
-
-# --- Async Question Processing ---
-async def process_question(agent, question: str, task_id: str) -> Dict:
-    """Process a single question and return both answer AND full log entry"""
-    try:
-        answer = agent(question)
-        return {
-            "submission": {"task_id": task_id, "submitted_answer": answer},
-            "log": {"Task ID": task_id, "Question": question, "Submitted Answer": answer}
-        }
-    except Exception as e:
-        error_msg = f"ERROR: {str(e)}"
-        return {
-            "submission": {"task_id": task_id, "submitted_answer": error_msg},
-            "log": {"Task ID": task_id, "Question": question, "Submitted Answer": error_msg}
-        }
-
-async def run_questions_async(agent, questions_data: List[Dict]) -> tuple:
-    """Process questions sequentially"""
-    submissions = []
-    logs = []
-    
-    total = len(questions_data)
-    for idx, q in enumerate(questions_data):
-        print(f"Processing {idx+1}/{total}: {q['question'][:80]}...")
-        result = await process_question(agent, q["question"], q["task_id"])
-        submissions.append(result["submission"])
-        logs.append(result["log"])
-    
-    return submissions, logs
-
+"""Model wrapper for LiteLLM"""
 
-async def run_and_submit_all(profile: gr.OAuthProfile | None):
-    """
-    Fetches all questions, runs the Agent on them, submits all answers,
-    and displays the results.
-    """
-    space_id = os.getenv("SPACE_ID")
+import os
+from typing import List, Dict, Any, Optional
 
-    if profile:
-        username = f"{profile.username}"
-        print(f"User logged in: {username}")
-    else:
-        print("User not logged in.")
-        return "Please Login to Hugging Face with the button.", None
+try:
+    import litellm
+except ImportError:
+    print("⚠️ litellm not installed. Install with: pip install litellm")
+    litellm = None
 
-    api_url = DEFAULT_API_URL
-    questions_url = f"{api_url}/questions"
-    submit_url = f"{api_url}/submit"
 
-    # 1. Instantiate Agent
-    try:
-        agent = Agent(
-            model=get_model("LiteLLMModel", MODEL_ID),
-            tools=get_tools()
-        )
-    except Exception as e:
-        print(f"Error instantiating agent: {e}")
-        return f"Error initializing agent: {e}", None
+class LiteLLMModel:
+    """Wrapper for LiteLLM models"""
     
-    agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main"
-    print(f"Agent code: {agent_code}")
-
-    # 2. Fetch Questions
-    print(f"Fetching questions from: {questions_url}")
-    try:
-        response = requests.get(questions_url, timeout=15)
-        response.raise_for_status()
-        questions_data = response.json()
-        if not questions_data:
-            print("Fetched questions list is empty.")
-            return "Fetched questions list is empty or invalid format.", None
-        print(f"Fetched {len(questions_data)} questions.")
-        # Remove this line to process all questions: questions_data = questions_data[:2]
-    except Exception as e:
-        print(f"Error fetching questions: {e}")
-        return f"Error fetching questions: {e}", None
-
-    # 3. Run Agent
-    print(f"Running agent on {len(questions_data)} questions...")
-    answers_payload, results_log = await run_questions_async(agent, questions_data)
-
-    if not answers_payload:
-        print("Agent did not produce any answers to submit.")
-        return "Agent did not produce any answers to submit.", pd.DataFrame(results_log)
-
-    # 4. Prepare Submission 
-    submission_data = {
-        "username": username.strip(),
-        "agent_code": agent_code,
-        "answers": answers_payload
-    }
-    print(f"Submitting {len(answers_payload)} answers for user '{username}'...")
-
-    # 5. Submit
-    try:
-        response = requests.post(submit_url, json=submission_data, timeout=60)
-        response.raise_for_status()
-        result_data = response.json()
-        final_status = (
-            f"✅ Submission Successful!\n\n"
-            f"User: {result_data.get('username')}\n"
-            f"Overall Score: {result_data.get('score', 'N/A')}% "
-            f"({result_data.get('correct_count', '?')}/{result_data.get('total_attempted', '?')} correct)\n\n"
-            f"Message: {result_data.get('message', 'No message received.')}\n\n"
-            f"Leaderboard: {api_url}/leaderboard"
-        )
-        print("Submission successful.")
-        results_df = pd.DataFrame(results_log)
-        return final_status, results_df
-    except Exception as e:
-        status_message = f"❌ Submission Failed: {e}"
-        print(status_message)
-        results_df = pd.DataFrame(results_log)
-        return status_message, results_df
-
-
-# --- Build Gradio Interface ---
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🤖 GAIA Agent Evaluation")
-    gr.Markdown(
-        """
-        **Instructions:**
-        1. Log in to your Hugging Face account using the button below
-        2. Click 'Run Evaluation & Submit' to test your agent
-        3. The agent will use web search and other tools to answer questions
+    def __init__(self, model_id: str):
+        self.model_id = model_id
         
-        **Current Setup:**
-        - Model: Gemini 2.5 Flash (via LiteLLM)
-        - Tools: Web search, Wikipedia, calculation, and more
-        """
-    )
-
-    gr.LoginButton()
-
-    run_button = gr.Button("🚀 Run Evaluation & Submit", variant="primary")
-
-    status_output = gr.Textbox(label="📊 Status / Results", lines=8, interactive=False)
-    results_table = gr.DataFrame(label="📋 Questions and Answers", wrap=True, max_height=400)
-
-    run_button.click(
-        fn=run_and_submit_all,
-        outputs=[status_output, results_table]
-    )
-
-if __name__ == "__main__":
-    print("\n" + "="*70)
-    print("🤖 GAIA Agent Starting")
-    print("="*70)
+        if "gemini" in model_id.lower():
+            if not os.getenv("GEMINI_API_KEY"):
+                print("⚠️ GEMINI_API_KEY not set in environment")
     
-    space_host = os.getenv("SPACE_HOST")
-    space_id = os.getenv("SPACE_ID")
+    def generate(self, messages: List[Dict], tools: Optional[List] = None) -> Dict:
+        if not litellm:
+            return {"content": "Unknown - litellm not installed"}
+        
+        try:
+            formatted_tools = None
+            if tools:
+                formatted_tools = [
+                    {
+                        "type": "function",
+                        "function": {
+                            "name": tool.name,
+                            "description": tool.description,
+                            "parameters": tool.parameters
+                        }
+                    }
+                    for tool in tools
+                ]
+            
+            if "gemini" in self.model_id.lower():
+                api_key = os.getenv("GEMINI_API_KEY")
+                if not api_key:
+                    raise RuntimeError("GEMINI_API_KEY not set in environment")
+                
+                print(f"DEBUG: Using model id: {self.model_id}")
+                
+                response = litellm.completion(
+                    model=self.model_id,
+                    api_key=api_key,
+                    messages=messages,
+                    tools=formatted_tools,
+                    temperature=0.1
+                )
+            else:
+                response = litellm.completion(
+                    model=self.model_id,
+                    messages=messages,
+                    tools=formatted_tools,
+                    temperature=0.1
+                )
+            
+            message = response.choices[0].message
+            result = {
+                "content": message.content or ""
+            }
+            
+            if hasattr(message, 'tool_calls') and message.tool_calls:
+                result["tool_calls"] = [
+                    {
+                        "name": tc.function.name,
+                        "arguments": eval(tc.function.arguments) if isinstance(tc.function.arguments, str) else tc.function.arguments
+                    }
+                    for tc in message.tool_calls
+                ]
+            
+            return result
+        
+        except Exception as e:
+            print(f"Model error: {e}")
+            return {"content": "Unknown"}
 
-    if space_host:
-        print(f"✅ Runtime URL: https://{space_host}.hf.space")
-    if space_id:
-        print(f"✅ Repo URL: https://huggingface.co/spaces/{space_id}")
-    
-    print("="*70 + "\n")
-    demo.launch(debug=True, share=False)
 
+def get_model(model_type: str, model_id: str):
+    if model_type == "LiteLLMModel":
+        return LiteLLMModel(model_id)
+    else:
+        raise ValueError(f"Unknown model type: {model_type}")