agent.py

import os
import re
import google.generativeai as genai
from tools import web_search, read_file_from_api, python_interpreter

# --- UPGRADED REACT PROMPT ---
REACT_PROMPT = """
You are a state-of-the-art, helpful AI agent designed to solve complex, multi-step problems.

**Your Task:**
Your goal is to answer the user's question with 100% accuracy. You must operate in a loop of Thought, Action, and Observation. Break the problem down into a series of smaller steps.

**Your Tools:**
You have access to the following tools. Choose ONE tool per turn.
1.  `web_search[query]`: Use this to find current information, facts, or to research topics.
2.  `read_file_from_api[task_id]`: Use this ONLY when the question explicitly mentions an attached file.
3.  `python_interpreter[code]`: Use this for all calculations, data processing (with pandas), and complex logic.

**CRITICAL INSTRUCTIONS:**
1.  Your reasoning process is: Thought -> Action -> Observation.
2.  You MUST continue this loop until you are certain of the answer.
3.  When you have the final, definitive answer, your ABSOLUTELY LAST output line MUST be in the format:
`Final Answer: [The single, exact answer]`
4.  Do not output any other text or explanation after the `Final Answer:` line.

---
Here is the problem:
Question: {question}
"""

class GeminiAgent:
    def __init__(self):
        print("Initializing GeminiAgent (Advanced ReAct)...")
        api_key = os.getenv("GEMINI_API_KEY")
        if not api_key:
            raise ValueError("GEMINI_API_KEY secret not found!")
        genai.configure(api_key=api_key)
        self.model = genai.GenerativeModel('gemini-2.5-pro')
        self.tools = {
            "web_search": web_search,
            "read_file_from_api": read_file_from_api,
            "python_interpreter": python_interpreter
        }
        print("GeminiAgent initialized successfully with model 'gemini-2.5-pro'.")

    def __call__(self, question: str, task_id: str) -> str:
        prompt = REACT_PROMPT.format(question=question)
        
        for turn in range(10): # Max 10 turns
            print(f"\n--- Turn {turn + 1} for Task ID: {task_id} ---\n")
            
            response = self.model.generate_content(prompt)
            if not response.parts:
                prompt += "\nObservation: The model returned an empty response. Please try again."
                continue
            
            response_text = response.text
            print(f"LLM Response:\n{response_text}\n")
            
            # Use re.findall to get ALL occurrences of Final Answer
            final_answer_matches = re.findall(r"Final Answer: (.*)", response_text, re.DOTALL)
            if final_answer_matches:
                # The model sometimes outputs multiple 'Final Answer' lines. The last one is the most correct.
                final_answer = final_answer_matches[-1].strip()
                
                # --- NEW: Robust cleaning of the final answer ---
                # Remove common trailing punctuation that isn't part of the answer itself.
                # This handles cases like 'Claus.' but preserves '1759.70'.
                if not final_answer.isnumeric():
                     final_answer = final_answer.rstrip('.?!,')

                print(f"Final Answer extracted and cleaned: '{final_answer}'")
                return final_answer

            action_match = re.search(r"Action: (\w+)\[(.*)\]", response_text, re.DOTALL)
            if not action_match:
                observation = "No valid 'Action:' or 'Final Answer:' found. Please think step-by-step and select a tool or provide the final answer."
            else:
                tool_name = action_match.group(1).strip()
                tool_input = action_match.group(2).strip()
                
                if tool_name not in self.tools:
                    observation = f"Error: Unknown tool '{tool_name}'."
                else:
                    try:
                        observation = self.tools[tool_name](tool_input if tool_name != 'read_file_from_api' else task_id)
                    except Exception as e:
                        observation = f"Error executing tool {tool_name}: {e}"
            
            print(f"Observation:\n{observation}\n")
            prompt += f"{response_text}\nObservation: {observation}\n"

        # Fallback if the agent gets stuck
        last_guess = response_text.split("Final Answer:")[-1].strip()
        print(f"Agent failed to find a 'Final Answer:' signal. Returning last guess: {last_guess}")
        return last_guess