agent.py

import io
from contextlib import redirect_stdout

from smolagents import (
    CodeAgent, 
    LiteLLMModel, 
    InferenceClientModel,
    DuckDuckGoSearchTool, 
    VisitWebpageTool, 
    WikipediaSearchTool
)

from config import (
    USE_LOCAL_MODEL,
    OLLAMA_MODEL_ID, OLLAMA_API_BASE, OLLAMA_API_KEY,
    HF_MODEL_ID, HF_TOKEN,
    MAX_STEPS, VERBOSITY_LEVEL, AUTHORIZED_IMPORTS,
    QUESTION_TYPES
)
from tools import smart_visit, get_youtube_info
from utils import clean_answer, clean_ansi_codes


class EnhancedAgent:
    """Enhanced agent with question-type specific strategies."""
    
    def __init__(self):
        print(f"   🤖 Initializing agent...")
        
        if USE_LOCAL_MODEL:
            # Usar Ollama local
            self.model = LiteLLMModel(
                model_id=OLLAMA_MODEL_ID, 
                api_base=OLLAMA_API_BASE,
                api_key=OLLAMA_API_KEY
            )
            print(f"   📦 Model: {OLLAMA_MODEL_ID} (local)")
        else:
            # Use HuggingFace API
            self.model = InferenceClientModel(
                model_id=HF_MODEL_ID,
                token=HF_TOKEN
            )
            print(f"   ☁️  Model: {HF_MODEL_ID} (HuggingFace)")

        search_tool = DuckDuckGoSearchTool()
        visit_tool = VisitWebpageTool()
        wiki_tool = WikipediaSearchTool()

        self.agent = CodeAgent(
            tools=[search_tool, visit_tool, wiki_tool, smart_visit, get_youtube_info], 
            model=self.model,
            max_steps=MAX_STEPS,
            verbosity_level=VERBOSITY_LEVEL,
            additional_authorized_imports=AUTHORIZED_IMPORTS
        )
    
    def build_prompt(self, question, local_file, question_type):
        """Construye prompt optimizado según el tipo de pregunta."""
        
        base_context = f"""TASK: You are solving a GAIA benchmark question. Be precise and methodical.

QUESTION: {question}
"""
        
        strategies = {
            QUESTION_TYPES['YOUTUBE_VIDEO']: """
STRATEGY - YouTube Video:
1. Extract the video ID from the URL in the question
2. Use get_youtube_info tool to get context
3. Search DuckDuckGo for: "[video_id] transcript" or "[video_id] [keywords_from_question]"
4. Look for Reddit threads, forums, or blogs discussing this video
5. Find the specific information requested

IMPORTANT: You CANNOT watch the video. Search for transcripts or discussions online.
""",
            
            QUESTION_TYPES['IMAGE_FILE']: f"""
STRATEGY - Image File:
1. File '{local_file}' is in current directory
2. You CANNOT read image files directly with Python
3. Search online for: "{local_file}" OR search for keywords from the question
4. Look for discussions, analysis, or descriptions of this image online
5. For chess positions: search "[piece positions] chess position solution"

IMPORTANT: Do NOT attempt cv2, PIL, or any image processing. Search online instead.
""",
            
            QUESTION_TYPES['AUDIO_FILE']: f"""
STRATEGY - Audio File:
1. File '{local_file}' is in current directory
2. You CANNOT play or transcribe audio with Python
3. Search online for: "{local_file}" OR the exact question text
4. Look for transcripts, Reddit threads, or forums discussing this audio

IMPORTANT: Do NOT attempt librosa, soundfile, or audio processing. Search online.
""",
            
            QUESTION_TYPES['DATA_FILE']: f"""
STRATEGY - Data File (Excel/CSV):
1. File '{local_file}' is in current directory
2. Use pandas to read: pd.read_excel('{local_file}') or pd.read_csv('{local_file}')
3. Explore columns with df.columns and df.head()
4. Filter and sum/count as needed
5. Double-check calculations

CODE TEMPLATE:
```python
import pandas as pd
df = pd.read_excel('{local_file}')  # or read_csv
print(df.columns)
print(df.head())
# ... your analysis
```
""",
            
            QUESTION_TYPES['CODE_FILE']: f"""
STRATEGY - Code File:
1. File '{local_file}' is in current directory
2. Read it with open('{local_file}', 'r').read()
3. Analyze the code logic carefully
4. If needed, execute it: exec(open('{local_file}').read())
5. Return the requested output

IMPORTANT: Read and understand before executing.
""",
            
            QUESTION_TYPES['WIKIPEDIA']: """
STRATEGY - Wikipedia Search:
1. Identify the exact topic/entity from the question
2. Use web_search to find the correct Wikipedia article URL
3. Use smart_visit to read the Wikipedia page content
4. Extract the specific information requested (dates, numbers, names, etc.)
5. For counting tasks: CREATE A PYTHON LIST with each item, then count with len()

TIPS:
- Search: "[topic] Wikipedia 2022" for latest version
- For discographies: look for "Discography" section or table
- For featured articles: search "Wikipedia Featured Article [topic] [date]"
- ALWAYS create a list and count programmatically, don't count manually
""",
            
            QUESTION_TYPES['COUNTING']: """
STRATEGY - Counting Task:
1. Research and LIST all items first (don't just count)
2. Use smart_visit to get complete data from Wikipedia or official sources
3. Store items in a Python list: items = []
4. Count with len(items) and verify manually
5. Double-check you haven't missed anything

IMPORTANT: First collect ALL items, THEN count. Show your work.
""",
            
            QUESTION_TYPES['TEXT_MANIPULATION']: """
STRATEGY - Text Manipulation:
1. Read the question VERY carefully
2. If text is backwards, reverse it: text[::-1]
3. If asking for opposite: use logic (left ↔ right, up ↔ down, etc.)
4. Return ONLY the answer, no explanation

EXAMPLE: ".rewsna eht sa 'tfel' drow..." 
→ Reverse to read: "...word 'left' as the answer."
→ Opposite of "left" is "right"
""",
            
            QUESTION_TYPES['GENERAL']: """
STRATEGY - General Research:
1. Break down the question into sub-tasks
2. Use web_search for initial research
3. Use smart_visit to read relevant pages in detail
4. Cross-reference multiple sources if needed
5. Extract the precise answer requested

TIPS:
- Be specific in searches: include years, full names, exact terms
- Read carefully - answers are often in tables, lists, or footnotes
"""
        }
        
        strategy = strategies.get(question_type, strategies[QUESTION_TYPES['GENERAL']])
        
        output_format = """
FINAL OUTPUT FORMAT:
Return ONLY the answer value. No markdown, no "The answer is", no explanations.

Examples of GOOD answers:
- "3"
- "right"
- "Ian Rose"
- "14.50"
- "d5, e2"

Examples of BAD answers:
- "The answer is 3"
- "**3**"
- "Based on my research, the answer is 3."
"""
        
        return base_context + strategy + output_format
    
    def solve(self, question, local_file=None, question_type=None):
        """
        Solve a question using an optimized strategy.
        
        Args:
            question: The question text
            local_file: Path to attached file (optional)
            question_type: Detected question type
            
        Returns:
            tuple: (answer, execution logs)
        """
        if question_type is None:
            question_type = QUESTION_TYPES['GENERAL']
            
        prompt = self.build_prompt(question, local_file, question_type)
        
        log_capture = io.StringIO()
        final_answer = "Error"
        
        try:
            with redirect_stdout(log_capture):
                answer = self.agent.run(prompt)
                final_answer = clean_answer(answer)
                
                # Si está vacío después de limpiar, buscar en logs
                if not final_answer or final_answer == "Error":
                    logs = log_capture.getvalue()
                    for line in reversed(logs.split('\n')):
                        if line.strip() and not any(x in line for x in ['===', '---', 'Step', 'Tool']):
                            potential_answer = line.strip()
                            if len(potential_answer) < 200:
                                final_answer = potential_answer
                                break
                
        except Exception as e:
            log_capture.write(f"\n❌ CRITICAL ERROR: {e}\n")
            final_answer = "Error"
        
        return final_answer, clean_ansi_codes(log_capture.getvalue())