Update app.py

app.py

@@ -3,10 +3,8 @@ import gradio as gr
 import requests
 import inspect
 import pandas as pd
-# from google import genai
-# from google.genai import types
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from google import genai
+from google.genai import types
 
 # (Keep Constants as is)
 # --- Constants ---
@@ -23,49 +21,49 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 #         print(f"Agent returning fixed answer: {fixed_answer}")
 #         return fixed_answer
 
-# class BasicAgent:
-#     def __init__(self):
-#         print("CustomAgent (using Gemini 2.0) initialized.")
-
-#         # Set the environment variable (important for Hugging Face Spaces)
-#         api_key = os.environ.get("GEMINI_API_KEY")
-#         if not api_key:
-#             raise ValueError("GEMINI_API_KEY not found in environment variables.")
-
-#         os.environ["GOOGLE_API_KEY"] = api_key  # Needed for google-genai Client
-
-#         # Initialize the Gemini client
-#         self.client = genai.Client()
-
-#         # Set model ID (you can choose another if needed)
-#         self.model_id = "gemini-2.0-flash-exp"
-
-#         # (Optional) Define generation config
-#         self.generation_config = types.GenerateContentConfig(
-#             temperature=0.4,
-#             top_p=0.95,
-#             top_k=20,
-#             candidate_count=1,
-#             seed=5,
-#             presence_penalty=0.0,
-#             frequency_penalty=0.0,
-#         )
+class BasicAgent:
+    def __init__(self):
+        print("CustomAgent using Gemini 2.0 initialized.")
+
+        # Set up Gemini API
+        api_key = os.environ.get("GEMINI_API_KEY")
+        if not api_key:
+            raise ValueError("GEMINI_API_KEY not found in environment variables.")
+        
+        os.environ["GOOGLE_API_KEY"] = api_key  # Required by google-genai
+
+        self.client = genai.Client()
+        self.model_id = "gemini-2.0-flash-exp"  # Or "gemini-1.5-flash-002" if you want faster
+
+        self.generation_config = types.GenerateContentConfig(
+            temperature=0.4,
+            top_p=0.9,
+            top_k=40,
+            candidate_count=1,
+            seed=42,
+            presence_penalty=0.0,
+            frequency_penalty=0.0,
+        )
 
-#     def __call__(self, question: str) -> str:
-#         print(f"Agent received question (first 50 chars): {question[:50]}...")
+    def __call__(self, question: str) -> str:
+        print(f"Agent received question (first 50 chars): {question[:50]}...")
 
-#         try:
-#             response = self.client.models.generate_content(
-#                 model=self.model_id,
-#                 contents=f"Answer the following question clearly and concisely: {question}",
-#                 config=self.generation_config
-#             )
-#             answer = response.text.strip()
-#             print(f"Agent returning answer (first 100 chars): {answer[:100]}")
-#             return answer
-#         except Exception as e:
-#             print(f"Error during Gemini API call: {str(e)}")
-#             return f"Error: {str(e)}"
+        try:
+            response = self.client.models.generate_content(
+                model=self.model_id,
+                contents=f"""You are a smart, factual assistant. Answer clearly and concisely:\n\n{question}\n\nProvide only the final answer without extra commentary.""",
+                config=self.generation_config,
+            )
+            answer = response.text.strip()
+
+            # ✨ Add a short sleep to avoid hitting rate limits
+            time.sleep(7)  # Wait 7 seconds after each question
+            print(f"Returning answer (first 100 chars): {answer[:100]}")
+            return answer
+
+        except Exception as e:
+            print(f"Error during Gemini call: {str(e)}")
+            return f"Error: {str(e)}"
 
 # class BasicAgent(ReActAgent):
 #     def __init__(self):
@@ -100,39 +98,40 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
 #         # Extract only the answer part
 #         return answer.split("Answer:")[-1].strip()
-class BasicAgent:
-    def __init__(self):
-        print("BasicAgent using local LLM initialized.")
-
-        # Load a small Hugging Face model
-        model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"  # Change if you want
-        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
-        self.model = AutoModelForCausalLM.from_pretrained(
-            model_name,
-            torch_dtype=torch.float16,
-            device_map="auto"  # Use GPU if available
-        )
 
-    def __call__(self, task: str) -> str:
-        """Answer a question."""
-        prompt = f"Answer the following question clearly and concisely:\n\n{task}\n\nAnswer:"
-        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
-
-        with torch.no_grad():
-            outputs = self.model.generate(
-                **inputs,
-                max_new_tokens=256,
-                do_sample=True,
-                temperature=0.7,
-                top_p=0.9,
-                top_k=50,
-            )
-        decoded = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+# class BasicAgent:
+#     def __init__(self):
+#         print("BasicAgent using local LLM initialized.")
+
+#         # Load a small Hugging Face model
+#         model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"  # Change if you want
+#         self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+#         self.model = AutoModelForCausalLM.from_pretrained(
+#             model_name,
+#             torch_dtype=torch.float16,
+#             device_map="auto"  # Use GPU if available
+#         )
+
+#     def __call__(self, task: str) -> str:
+#         """Answer a question."""
+#         prompt = f"Answer the following question clearly and concisely:\n\n{task}\n\nAnswer:"
+#         inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
+
+#         with torch.no_grad():
+#             outputs = self.model.generate(
+#                 **inputs,
+#                 max_new_tokens=256,
+#                 do_sample=True,
+#                 temperature=0.7,
+#                 top_p=0.9,
+#                 top_k=50,
+#             )
+#         decoded = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-        # Extract the answer part
-        if "Answer:" in decoded:
-            return decoded.split("Answer:")[-1].strip()
-        return decoded.strip()
+#         # Extract the answer part
+#         if "Answer:" in decoded:
+#             return decoded.split("Answer:")[-1].strip()
+#         return decoded.strip()
 
 def run_and_submit_all( profile: gr.OAuthProfile | None):
     """