feat: don't let others spend my Gemini money T_T

app.py

@@ -66,21 +66,7 @@ def answer_one(agent, question_data):
             "Question": question_text,
             "Submitted Answer": submitted_answer or agent_error,
         }
-        return payload, log_entry
-
-def _submit_answers_to_file(answers_payload, file_path):
-    """
-    Submits the answers to a local file.
-    """
-    
-    try:
-        with open(file_path, "w") as file:
-            json.dump(answers_payload, file, indent=4)
-        submit_status = (f"Answers successfully written to {file_path}")
-    except Exception as e:
-        submit_status = (f"Error writing answers to file: {e}")
-    print(submit_status)
-    return submit_status
+        return payload, log_entry    
     
 def _submit_all(username, agent_code, answers_payload, submit_url):
     # Prepare Submission 
@@ -130,7 +116,7 @@ def _submit_all(username, agent_code, answers_payload, submit_url):
 def prepare_agent(api_key=None):  
     # 1. Instantiate Agent ( modify this part to create your agent)
     try:
-        agent = GeminiAgentContainer()
+        agent = GeminiAgentContainer(api_key=api_key)
         print(agent.system_prompt)
     except Exception as e:
         print(f"Error instantiating agent: {e}")
@@ -138,7 +124,27 @@ def prepare_agent(api_key=None):
     
     return agent
 
-def _run_all(api_key: str | None = None):
+def save_answers_to_file():
+    """
+    Submits the answers to a local file named with the current epoch time.
+    """
+    if not answers_by_task:
+        return ("Nothing to save, no answers found.")
+    answers_payload = list(answers_by_task.values())
+
+    file_path = f"answers-{int(time.time())}.json"
+    print(f"Saving answers to file: {file_path}")
+    try:
+        with open(file_path, "w") as file:
+            json.dump(answers_payload, file, indent=4)
+        submit_status = (f"Answers successfully written to {file_path}")
+    except Exception as e:
+        submit_status = (f"Error writing answers to file: {e}")
+    print(submit_status)
+    return submit_status
+
+
+def run_all(api_key: str | None = None):
     """
     Fetches all questions, runs the BasicAgent on them,
     """
@@ -165,41 +171,30 @@ def _run_all(api_key: str | None = None):
     print(final_status)
     return final_status, pd.DataFrame(results_log)
 
-
-def run_all():
-    run_status, results_df = _run_all(os.environ.get("GOOGLE_API_KEY"))
-    return run_status, results_df
-
-def submit_all( profile: gr.OAuthProfile | None, to_file = True):
+def submit_all( profile: gr.OAuthProfile | None):
     """
     Submits all answers and displays the results.
     """
     submit_url = f"{DEFAULT_API_URL}/submit"
 
-    if not to_file:
-        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."
-    
+    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."
+
     # --- Determine HF Space Runtime URL and Repo URL ---
     space_id = os.getenv("SPACE_ID") # Get the SPACE_ID for sending link to the code
 
     if not answers_by_task:
         submit_status = "No answers to submit."
     else:
-
         # 4. Submit all answers
-        answers_payload = list(answers_by_task.values())
         # In the case of an app running as a hugging Face space, this link points toward your codebase ( usefull for others so please keep it public)
         agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main"
-        
-        if to_file:
-            submit_status = _submit_answers_to_file(answers_payload, f"run-{int(time.time())}.json")
-        else:
-            submit_status = _submit_all(username, agent_code, answers_payload, submit_url)
+
+        submit_status = _submit_all(username, agent_code, list(answers_by_task.values()), submit_url)
 
     return submit_status
 
@@ -219,7 +214,16 @@ with gr.Blocks() as demo:
 
     gr.LoginButton()
 
+    api_key_input = gr.Textbox(
+        label="Gemini API Key",
+        placeholder="Enter your Gemini API key here",
+        type="password",
+        lines=1,
+        visible=True
+    )
+
     run_button = gr.Button("Run Evaluation")
+    save_button = gr.Button("Save Answers to File")
     submit_button = gr.Button("Submit All Answers")
 
     status_output = gr.Textbox(label="Run Status / Submission Result", lines=5, interactive=False)
@@ -227,9 +231,15 @@ with gr.Blocks() as demo:
 
     run_button.click(
         fn=run_all,
+        inputs=[api_key_input],
         outputs=[status_output, results_table]
     )
 
+    save_button.click(
+        fn=save_answers_to_file,
+        outputs=[status_output]
+    )
+
     submit_button.click(
         fn=submit_all,
         outputs=[status_output]

my_agent.py

@@ -1,30 +1,38 @@
 import os
 import requests
-from smolagents import LiteLLMModel, ToolCallingAgent, tool
+from smolagents import LiteLLMModel, ToolCallingAgent, Tool
 from typing import Optional
 from google import genai
 from google.genai import types
 import wikipedia as wiki
 from markdownify import markdownify as to_markdown
 
-# --- Constants --
-client = genai.Client(api_key=os.getenv("GOOGLE_API_KEY"))
-model_name = "gemini-2.0-flash"
-
-
 # --- Tools ---
-@tool
-def watch_video(video_url: str, user_query:str) -> str:
-    """
-    Watches a video and answers a question about it.
-    Args:
-        video_url (str): The URL of the video to watch.
-        user_query (str): The question to answer about the video.
-    Returns:
-        str: The answer to the question.    
+
+class VideoWatchingTool(Tool):
+    name = "watch_video"
+    description ="""
+    A tool for watching videos and answering questions about them.
     """
-    request_json = {
-        'model': f'models/{model_name}',
+    inputs = {
+        "video_url": {
+            "type": "string",
+            "description": "The URL of the video to watch."
+        },
+        "user_query": {
+            "type": "string",
+            "description": "The question to answer about the video."
+        }
+    }
+    output_type = "string"
+
+    def __init__(self, model_name, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.model_name = model_name
+
+    def forward(self, video_url: str, user_query: str) -> str:
+        request_json = {
+        'model': f'models/{self.model_name}',
         'contents': [{
             "parts": [
                 {
@@ -38,131 +46,157 @@ def watch_video(video_url: str, user_query:str) -> str:
             ]
         }]
     }
-    api_url = f'https://generativelanguage.googleapis.com/v1beta/models/{model_name}:generateContent?key={os.getenv("GOOGLE_API_KEY")}'
-    response = requests.post(
-        api_url, 
-        json=request_json,
-        headers={
-            'Content-Type': 'application/json',
-        }
-    )
-    if response.status_code != 200:
-        return f"Error: {response.status_code} - {response.text}"
-    response_json = response.json()
-    result_parts = response_json['candidates'][0]['content']['parts']
-    result = "".join([_.get('text', '') for _ in result_parts])
-    return result
-
-@tool
-def google_search(query: str) -> str:
-    """
+        api_url = f'https://generativelanguage.googleapis.com/v1beta/models/{self.model_name}:generateContent?key={os.getenv("GOOGLE_API_KEY")}'
+        response = requests.post(
+            api_url, 
+            json=request_json,
+            headers={
+                'Content-Type': 'application/json',
+            }
+        )
+        if response.status_code != 200:
+            return f"Error: {response.status_code} - {response.text}"
+        response_json = response.json()
+        result_parts = response_json['candidates'][0]['content']['parts']
+        result = "".join([_.get('text', '') for _ in result_parts])
+        return result
+    
+class GoogleSearchTool(Tool):
+    name = "google_search"
+    description = """
     Performs a Google search and returns the results.
-    Args:
-        query (str): The search query.
-    Returns:
-        str: The search results.
     """
-    # This is a placeholder for the actual Google search logic
-    # In a real implementation, you would use the Google Search API
-    # to get the results.
-    # For now, we will just return a dummy response.
-    google_search_tool = types.Tool(
-        google_search=types.GoogleSearch()
-    )
-    response = client.models.generate_content(
-        model=model_name,
+    inputs = {
+        "query": {
+            "type": "string",
+            "description": "The search query."
+        }
+    }
+    output_type = "string"
+
+    def __init__(self, client, model_name, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.client = client
+        self.model_name = model_name
+
+    def forward(self, query: str) -> str:
+        google_search_tool = types.Tool(
+            google_search=types.GoogleSearch()
+        )
+        response = self.client.models.generate_content(
+        model=self.model_name,
         contents=f"Please search the internet for: {query}",
         config=types.GenerateContentConfig(
             tools=[google_search_tool],
             response_modalities=['TEXT'],
+            )
         )
-    )
-    result_parts = response.candidates[0].content.parts
-    result = "".join([_.text for _ in result_parts])
-
-    return result
-
-@tool
-
-def check_wikipedia_page_titles(query: str) -> str:
-    """
+        return response.text
+    
+class WikipediaTitleSearchTool(Tool):
+    name = "check_wikipedia_page_titles"
+    description = """
     Searches for Wikipedia pages related to the query and returns the canonical titles of the related pages.
-    Args:
-        query (str): The search query.
-    Returns:
-        str: A comma separated list of canonical Wikipedia page titles that match the query.
-    """
-    response = wiki.search(query)
-    if len(response) > 0:
-        result = ", ".join(response)
-    else:
-        result = "No results found."
-    return result
-   
-@tool
-def get_wikipedia_page(page_title: str) -> str:
     """
-    Gets the content of a we
-    Args:
-        page_title (str): The canonical title of the Wikipedia page.
-    Returns:
-        str: Markdown-formatted content of the Wikipedia page.
-    """
-    # This is a placeholder for the actual Wikipedia search logic
-    # In a real implementation, you would use the Wikipedia API
-    # to get the results.
-    # For now, we will just return a dummy response.
-    try:
-        page = wiki.page(page_title)
-    except wiki.exceptions.PageError:
-        return f"Page '{page_title}' not found."
-    md_content = to_markdown(page.html())
-    return md_content
-
-@tool
-def run_query_with_file(task_id: str, mime_type: str | None, user_query: str) -> str:
+    inputs = {
+        "query": {
+            "type": "string",
+            "description": "The search query."
+        }
+    }
+    output_type = "string"
+
+    def forward(self, query: str) -> str:
+        response = wiki.search(query)
+        if len(response) > 0:
+            result = ", ".join(response)
+        else:
+            result = "No results found."
+        return result
+
+class WikipediaPageTool(Tool):
+    name = "get_wikipedia_page"
+    description = """
+    Gets the content of a Wikipedia page.
     """
+    inputs = {
+        "page_title": {
+            "type": "string",
+            "description": "The canonical title of the Wikipedia page."
+        }
+    }
+    output_type = "string"
+    
+    def forward(self, page_title: str) -> str:
+        # TODO: may need to do caching of the HTML ourselves?
+        try:
+            page = wiki.page(page_title)
+        except wiki.exceptions.PageError:
+            return f"Page '{page_title}' not found."
+        md_content = to_markdown(page.html())
+        return md_content
+   
+class FileAttachmentQueryTool(Tool):
+    name = "run_query_with_file"
+    description = """
     Downloads a file mentioned in a user prompt, adds it to the context, and runs a query on it.
     This assumes the file is 20MB or less.
-    Args:
-        task_id (str): A unique identifier for the task related to this file, used to download it.
-        mime_type (str): The MIME type of the file, or the best guess if unknown
-        user_query (str): The question to answer about the file.
-    Returns:
-        str: The answer to the question.
     """
-    # Download the file
-    file_url = f"https://agents-course-unit4-scoring.hf.space/files/{task_id}"
-    file_response = requests.get(file_url)
-    if file_response.status_code != 200:
-        raise Exception(f"Failed to download file: {file_response.status_code} - {file_response.text}")
-    file_data = file_response.content
-    mime_type = mime_type or file_response.headers.get('Content-Type', 'application/octet-stream')
-    response = client.models.generate_content(
-        model=model_name,
-        contents=[
-            types.Part.from_bytes(
-                data=file_data,
-                mime_type=mime_type,
-            ),
-            user_query,
-        ]
-    )
-
-    return response.text
+    inputs = {
+        "task_id": {
+            "type": "string",
+            "description": "A unique identifier for the task related to this file, used to download it."
+        },
+        "mime_type": {
+            "type": "string",
+            "nullable": True,
+            "description": "The MIME type of the file, or the best guess if unknown."
+        },
+        "user_query": {
+            "type": "string",
+            "description": "The question to answer about the file."
+        }
+    }
+    output_type = "string"
+    def __init__(self, client, model_name, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.client = client
+        self.model_name = model_name
     
+    def forward(self, task_id: str, mime_type: str | None, user_query: str) -> str:
+        # Download the file
+        file_url = f"https://agents-course-unit4-scoring.hf.space/files/{task_id}"
+        file_response = requests.get(file_url)
+        if file_response.status_code != 200:
+            raise Exception(f"Failed to download file: {file_response.status_code} - {file_response.text}")
+        file_data = file_response.content
+        mime_type = mime_type or file_response.headers.get('Content-Type', 'application/octet-stream')
+        response = self.client.models.generate_content(
+            model=self.model_name,
+            contents=[
+                types.Part.from_bytes(
+                    data=file_data,
+                    mime_type=mime_type,
+                ),
+                user_query,
+            ]
+        )
 
+        return response.text
+    
 # --- Agent Management ---
 
 class GeminiAgentContainer:
     """
     A container for the Gemini agent.
     """
-    def __init__(self, model_id: Optional[str] = None, provider: Optional[str] = None):
-        # model_id = model_id or "Qwen/Qwen2.5-Coder-32B-Instruct"
-        # provider = provider or "together"
-        # model = HfApiModel(model_id=model_id, provider=provider)
-        self.model = LiteLLMModel(model_id=f"gemini/{model_name}", api_key=os.getenv("GOOGLE_API_KEY"))
+    # TODO: make it easier to chnge the model
+    MODEL_NAME = "gemini-2.0-flash"
+
+    def __init__(self, api_key: Optional[str] = None):
+        api_key = api_key or os.getenv("GOOGLE_API_KEY")
+        self.model = LiteLLMModel(model_id=f"gemini/{self.MODEL_NAME}", api_key=api_key)
+        self.client = genai.Client(api_key=os.getenv("GOOGLE_API_KEY"))
         system_prompt = """
             You are a general AI assistant. I will ask you a question. 
             YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. 
@@ -180,13 +214,13 @@ class GeminiAgentContainer:
         self.agent = ToolCallingAgent(
             model=self.model, 
             tools = [
-                watch_video, 
-                google_search,
-                run_query_with_file,
-                check_wikipedia_page_titles,
-                get_wikipedia_page,
-                ],
-            max_steps=6,
+                VideoWatchingTool(model_name=self.MODEL_NAME),
+                GoogleSearchTool(client=self.client, model_name=self.MODEL_NAME),
+                WikipediaTitleSearchTool(),
+                WikipediaPageTool(),
+                FileAttachmentQueryTool(client=self.client, model_name=self.MODEL_NAME),
+                ], 
+            max_steps=14,
             planning_interval=2,
         )
         self.system_prompt = system_prompt