Final_Assignment / agent.py
benjipeng's picture
Create agent.py
e04605e verified
raw
history blame
2.45 kB
import os
import google.generativeai as genai
class GeminiAgent:
"""
An agent that uses the Gemini-1.5-Pro model to answer questions.
"""
def __init__(self):
"""
Initializes the agent, configures the Gemini API key, and sets up the model.
Raises a ValueError if the GEMINI_API_KEY is not found in the environment secrets.
"""
print("Initializing GeminiAgent...")
# 1. Get API Key from Hugging Face Secrets
api_key = os.getenv("GEMINI_API_KEY")
if not api_key:
raise ValueError("GEMINI_API_KEY secret not found! Please set it in your Space's settings.")
# 2. Configure the Generative AI client
genai.configure(api_key=api_key)
# 3. Initialize the Gemini 1.5 Pro model
self.model = genai.GenerativeModel('gemini-1.5-pro-latest')
print("GeminiAgent initialized successfully.")
def __call__(self, question: str) -> str:
"""
Processes a question by sending it to the Gemini model and returns the stripped text answer.
The prompt is engineered to request a direct, exact-match answer as required by the competition.
"""
print(f"Agent received question (first 80 chars): {question[:80]}...")
# Prompt engineered for the "EXACT MATCH" requirement.
# It instructs the model to provide only the answer and nothing else.
prompt = f"""You are an expert problem-solving agent. Your goal is to answer the following question as accurately as possible.
The evaluation system requires an EXACT MATCH. Therefore, you must provide only the final answer and nothing else.
Do not include any introductory text, explanations, or the phrase "FINAL ANSWER".
For example, if the question asks for a specific year, your response should be just "2023". If it's a name, just "John Doe". If it is a number, just "42".
Question: {question}
Final Answer:"""
try:
# 4. Call the model
response = self.model.generate_content(prompt)
# 5. Extract and clean the answer
final_answer = response.text.strip()
print(f"Agent returning answer: {final_answer}")
return final_answer
except Exception as e:
print(f"An error occurred while calling the Gemini API: {e}")
return f"Error processing question: {e}"