Update app.py

app.py

@@ -4,44 +4,51 @@ import os
 import gradio as gr
 import requests
 import pandas as pd
+from PIL import Image
+import base64
+import io
 
-from smolagents import LiteLLMModel, CodeAgent, DuckDuckGoSearchTool
+import google.generativeai as genai
+from smolagents import CodeAgent, DuckDuckGoSearchTool, LiteLLMModel
 
-# System prompt for the agent
+# System prompt used by the agent
 SYSTEM_PROMPT = """You are a general AI assistant. I will ask you a question.
-Report your thoughts, and finish your answer with just the answer — no prefixes like "FINAL ANSWER:".
+Report your thoughts, and finish your answer with just the answer — no prefixes like \"FINAL ANSWER:\".
 Your answer should be a number OR as few words as possible OR a comma-separated list of numbers and/or strings.
 If you're asked for a number, don’t use commas or units like $ or %, unless specified.
 If you're asked for a string, don’t use articles or abbreviations (e.g. for cities), and write digits in plain text unless told otherwise."""
 
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
-# Agent wrapper using LiteLLMModel
+# Load GEMINI_API_KEY from environment
+GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
+
+# Agent wrapper
 class MyAgent:
     def __init__(self):
-        gemini_api_key = os.getenv("GEMINI_API_KEY")
-        if not gemini_api_key:
-            raise ValueError("GEMINI_API_KEY not set in environment variables.")
-        
-        # Instantiate LiteLLMModel with Gemini API key and model id
-        self.model = LiteLLMModel(
-            model_id="gemini/gemini-2.0-flash-lite",
-            api_key=gemini_api_key,
-            system_prompt=SYSTEM_PROMPT
-        )
-        
-        # Create the CodeAgent with optional base tools and DuckDuckGo search
-        self.agent = CodeAgent(
-            tools=[DuckDuckGoSearchTool()],
-            model=self.model,
-            add_base_tools=True,
-        )
-
-    def __call__(self, question: str) -> str:
-        return self.agent.run(question)
+        model = LiteLLMModel(model_id="gemini/gemini-1.5-flash", api_key=GEMINI_API_KEY)
+        self.agent = CodeAgent(tools=[DuckDuckGoSearchTool()], model=model)
+
+    def __call__(self, question: str, code: str | None = None, excel_df: pd.DataFrame | None = None, image: Image.Image | None = None) -> str:
+        if excel_df is not None:
+            preview = excel_df.head().to_csv(index=False)
+            context = f"This is a preview of the attached Excel sales data:\n\n{preview}"
+            prompt = f"{question}\n\n{context}"
+            return self.agent.run(prompt)
+        elif image is not None:
+            buffered = io.BytesIO()
+            image.save(buffered, format="JPEG")
+            img_b64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
+            prompt = f"{question}\n\nThis is the attached image (base64 JPEG):\n\n{img_b64}"
+            return self.agent.run(prompt)
+        elif code:
+            formatted = f"{question}\n\nThoughts: Let's analyze the attached code.\nCode:\n```py\n{code}\n```<end_code>"
+            return self.agent.run(formatted)
+        else:
+            return self.agent.run(question)
 
 # Main evaluation function
-def run_and_submit_all(profile: gr.OAuthProfile | None):
+def run_and_submit_all(profile: gr.OAuthProfile | None, uploaded_code: list[gr.File] | None, uploaded_excel: list[gr.File] | None, uploaded_image: list[gr.File] | None):
     space_id = os.getenv("SPACE_ID")
 
     if profile:
@@ -66,6 +73,30 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
     except Exception as e:
         return f"Error fetching questions: {e}", None
 
+    uploaded_code_str = ""
+    if uploaded_code:
+        try:
+            uploaded_file = uploaded_code[0]
+            uploaded_code_str = uploaded_file.read().decode("utf-8")
+        except Exception as e:
+            uploaded_code_str = f"# Failed to load uploaded code: {e}"
+
+    uploaded_excel_df = None
+    if uploaded_excel:
+        try:
+            uploaded_excel_df = pd.read_excel(uploaded_excel[0].name)
+        except Exception as e:
+            print(f"Error reading Excel: {e}")
+            uploaded_excel_df = None
+
+    uploaded_image_obj = None
+    if uploaded_image:
+        try:
+            uploaded_image_obj = Image.open(uploaded_image[0].name)
+        except Exception as e:
+            print(f"Error loading image: {e}")
+            uploaded_image_obj = None
+
     results_log = []
     answers_payload = []
 
@@ -75,9 +106,14 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
         if not task_id or question_text is None:
             continue
         try:
-            submitted_answer = agent(question_text)
-            answers_payload.append({"task_id": task_id, "submitted_answer": submitted_answer})
-            results_log.append({"Task ID": task_id, "Question": question_text, "Submitted Answer": submitted_answer})
+            answer = agent(
+                question_text,
+                uploaded_code_str if "code" in question_text.lower() else None,
+                uploaded_excel_df if "excel" in question_text.lower() or "spreadsheet" in question_text.lower() else None,
+                uploaded_image_obj if "image" in question_text.lower() or "photo" in question_text.lower() or "jpg" in question_text.lower() else None
+            )
+            answers_payload.append({"task_id": task_id, "submitted_answer": answer})
+            results_log.append({"Task ID": task_id, "Question": question_text, "Submitted Answer": answer})
         except Exception as e:
             results_log.append({
                 "Task ID": task_id,
@@ -116,19 +152,26 @@ with gr.Blocks() as demo:
     **Instructions:**
     1. Clone this space and configure your Gemini API key.
     2. Log in to Hugging Face.
-    3. Run your agent on evaluation tasks and submit answers.
+    3. Optionally upload Python code, Excel file, or image used by the questions.
+    4. Run your agent on evaluation tasks and submit answers.
     """)
 
     gr.LoginButton()
+    code_upload = gr.File(label="Upload Python code file", file_types=[".py"])
+    excel_upload = gr.File(label="Upload Excel file", file_types=[".xls", ".xlsx"])
+    image_upload = gr.File(label="Upload Image file", file_types=[".jpg", ".jpeg"])
     run_button = gr.Button("Run Evaluation & Submit All Answers")
     status_output = gr.Textbox(label="Submission Result", lines=5, interactive=False)
     results_table = gr.DataFrame(label="Results", wrap=True)
 
-    run_button.click(fn=run_and_submit_all, outputs=[status_output, results_table])
+    run_button.click(
+        fn=run_and_submit_all,
+        inputs=[gr.OAuthProfile(), code_upload, excel_upload, image_upload],
+        outputs=[status_output, results_table]
+    )
 
 if __name__ == "__main__":
     print("🔧 App starting...")
     demo.launch(debug=True, share=False)
 
 
-