gaia_agent.py

"""
Улучшенный агент GAIA с интеграцией LLM для курса Hugging Face
"""

import os
import gradio as gr
import requests
import pandas as pd
import json
import time
from typing import List, Dict, Any, Optional, Callable, Union
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

# --- Константы ---
DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
DEFAULT_MODEL = "google/flan-t5-small"  # Меньшая модель для быстрой загрузки
MAX_RETRIES = 3  # Максимальное количество попыток отправки
RETRY_DELAY = 5  # Задержка между попытками в секундах

class LLMGAIAAgent:
    """
    Улучшенный агент GAIA, использующий языковую модель для генерации ответов.
    """
    
    def __init__(self, model_name=DEFAULT_MODEL):
        """Инициализация агента с языковой моделью."""
        print(f"Инициализация LLMGAIAAgent с моделью: {model_name}")
        try:
            self.tokenizer = AutoTokenizer.from_pretrained(model_name)
            self.model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
            self.model_name = model_name
            print(f"Успешно загружена модель: {model_name}")
        except Exception as e:
            print(f"Ошибка загрузки модели: {e}")
            print("Переход к шаблонным ответам")
            self.model = None
            self.tokenizer = None
            self.model_name = None
    
    def __call__(self, question: str) -> str:
        """Обработка вопроса и возврат ответа с использованием языковой модели."""
        print(f"Обработка вопроса: {question}")
        
        if self.model is None or self.tokenizer is None:
            return self._fallback_response(question)
        
        try:
            prompt = self._prepare_prompt(question)
            inputs = self.tokenizer(prompt, return_tensors="pt", max_length=512, truncation=True)
            outputs = self.model.generate(
                inputs["input_ids"],
                max_length=150,
                min_length=20,
                temperature=0.7,
                top_p=0.9,
                do_sample=True,
                num_return_sequences=1
            )
            response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
            response = self._clean_response(response)
            return response
        except Exception as e:
            print(f"Ошибка генерации ответа: {e}")
            return self._fallback_response(question)
    
    def _prepare_prompt(self, question: str) -> str:
        """Подготовка подходящего запроса на основе типа вопроса."""
        question_lower = question.lower()
        if any(keyword in question_lower for keyword in [
            "calculate", "compute", "sum", "difference", 
            "product", "divide", "plus", "minus", "times"
        ]):
            return f"Решите эту математическую задачу шаг за шагом: {question}"
        elif any(keyword in question_lower for keyword in [
            "image", "picture", "photo", "graph", "chart", "diagram"
        ]):
            return f"Опишите, что может быть изображено на картинке, связанной с этим вопросом: {question}"
        elif any(keyword in question_lower for keyword in [
            "who", "what", "where", "when", "why", "how"
        ]):
            return f"Дайте краткий и точный ответ на этот фактический вопрос: {question}"
        else:
            return f"Дайте краткий, информативный ответ на этот вопрос: {question}"
    
    def _clean_response(self, response: str) -> str:
        """Очистка ответа модели для получения чистого текста."""
        prefixes = [
            "Answer:", "Response:", "A:", "The answer is:", 
            "It is:", "I think it is:", "The result is:", 
            "Based on the image:", "In the image:", 
            "The image shows:", "From the image:"
        ]
        for prefix in prefixes:
            if response.lower().startswith(prefix.lower()):
                response = response[len(prefix):].strip()
        if len(response) < 10:
            return self._fallback_response("general")
        return response.strip()
    
    def _fallback_response(self, question: str) -> str:
        """Резервный ответ, если модель не сработала."""
        question_lower = question.lower() if isinstance(question, str) else ""
        if "who" in question_lower:
            return "Известная личность в этой области."
        elif "when" in question_lower:
            return "Это произошло в значительный исторический период."
        elif "where" in question_lower:
            return "Место известно своей культурной значимостью."
        elif "what" in question_lower:
            return "Это важное понятие или объект."
        elif "why" in question_lower:
            return "Это произошло из-за ряда факторов."
        elif "how" in question_lower:
            return "Процесс включает несколько ключевых шагов."
        return "Ответ включает несколько важных факторов."

class EvaluationRunner:
    """
    Управление процессом оценки: получение вопросов, запуск агента и отправка ответов.
    """
    
    def __init__(self, api_url: str = DEFAULT_API_URL):
        """Инициализация с конечными точками API."""
        self.api_url = api_url
        self.questions_url = f"{api_url}/questions"
        self.submit_url = f"{api_url}/submit"
    
    def run_evaluation(self, 
                      agent: Callable[[str], str], 
                      username: str, 
                      agent_code_url: str) -> tuple[str, pd.DataFrame]:
        """Запуск полного процесса оценки."""
        questions_data = self._fetch_questions()
        if isinstance(questions_data, str):
            return questions_data, None
        
        results_log, answers_payload = self._run_agent_on_questions(agent, questions_data)
        if not answers_payload:
            return "Агент не дал ответов для отправки.", pd.DataFrame(results_log)
        
        submission_result = self._submit_answers_with_retry(username, agent_code_url, answers_payload)
        return submission_result, pd.DataFrame(results_log)
    
    def _fetch_questions(self) -> Union[List[Dict[str, Any]], str]:
        """Получение вопросов с сервера оценки."""
        print(f"Получение вопросов с: {self.questions_url}")
        try:
            response = requests.get(self.questions_url, timeout=15)
            response.raise_for_status()
            questions_data = response.json()
            if not questions_data:
                return "Список вопросов пуст или некорректен."
            print(f"Успешно получено {len(questions_data)} вопросов.")
            return questions_data
        except Exception as e:
            return f"Ошибка получения вопросов: {e}"
    
    def _run_agent_on_questions(self, 
                               agent: Callable[[str], str], 
                               questions_data: List[Dict[str, Any]]) -> tuple[List[Dict[str, Any]], List[Dict[str, Any]]]:
        """Запуск агента на всех вопросах."""
        results_log = []
        answers_payload = []
        print(f"Запуск агента на {len(questions_data)} вопросах...")
        for item in questions_data:
            task_id = item.get("task_id")
            question_text = item.get("question")
            if not task_id or question_text is None:
                continue
            try:
                submitted_answer = agent(question_text)
                answers_payload.append({"task_id": task_id, "submitted_answer": submitted_answer})
                results_log.append({"Task ID": task_id, "Question": question_text, "Submitted Answer": submitted_answer})
            except Exception as e:
                results_log.append({"Task ID": task_id, "Question": question_text, "Submitted Answer": f"ОШИБКА: {e}"})
        return results_log, answers_payload
    
    def _submit_answers_with_retry(self, 
                                  username: str, 
                                  agent_code_url: str, 
                                  answers_payload: List[Dict[str, Any]]) -> str:
        """Отправка ответов с логикой повтора."""
        submission_data = {
            "username": username.strip(),
            "agent_code_url": agent_code_url,  # Исправленный ключ
            "answers": answers_payload
        }
        print(f"Отправка {len(answers_payload)} ответов для пользователя '{username}'...")
        for attempt in range(1, MAX_RETRIES + 1):
            try:
                print(f"Попытка {attempt} из {MAX_RETRIES}...")
                response = requests.post(self.submit_url, json=submission_data, timeout=60)
                response.raise_for_status()
                result_data = response.json()
                final_status = (
                    f"Отправка успешна!\n"
                    f"Пользователь: {result_data.get('username')}\n"
                    f"Общий балл: {result_data.get('overall_score', 'N/A')}\n"
                    f"Правильные ответы: {result_data.get('correct_answers', 'N/A')}\n"
                    f"Всего вопросов: {result_data.get('total_questions', 'N/A')}\n"
                )
                if all(result_data.get(key, "N/A") == "N/A" for key in ["overall_score", "correct_answers", "total_questions"]):
                    final_status += (
                        "\nПримечание: Результаты показывают 'N/A'. Возможные причины:\n"
                        "- Ограничения активности аккаунта\n"
                        "- Задержка обработки\n"
                        "- Проблема с API\n"
                        f"Проверьте статус: {DEFAULT_API_URL}/results?username={username}"
                    )
                print(final_status)
                return final_status
            except Exception as e:
                if attempt < MAX_RETRIES:
                    time.sleep(RETRY_DELAY)
                else:
                    return f"Ошибка отправки после {MAX_RETRIES} попыток: {e}"

def run_and_submit_all(profile: gr.OAuthProfile | None, *args):
    """Основная функция для запуска через Gradio."""
    if not profile:
        return "Пожалуйста, войдите в Hugging Face.", None
    username = profile.username
    space_id = os.getenv("SPACE_ID")
    agent_code_url = f"https://huggingface.co/spaces/{space_id}/tree/main"
    print(f"URL кода агента: {agent_code_url}")
    try:
        agent = LLMGAIAAgent()
        runner = EvaluationRunner()
        return runner.run_evaluation(agent, username, agent_code_url)
    except Exception as e:
        return f"Ошибка инициализации: {e}", None

# --- Интерфейс Gradio ---
with gr.Blocks() as demo:
    gr.Markdown("# Оценка агента GAIA (с улучшенным LLM)")
    gr.Markdown("## Инструкции:")
    gr.Markdown("1. Войдите в аккаунт Hugging Face.")
    gr.Markdown("2. Нажмите 'Запустить оценку и отправить все ответы'.")
    gr.Markdown("3. Посмотрите результаты в разделе вывода.")
    with gr.Row():
        login_button = gr.LoginButton(value="Войти через Hugging Face")
    with gr.Row():
        submit_button = gr.Button("Запустить оценку и отправить все ответы")
    with gr.Row():
        output_status = gr.Textbox(label="Результат отправки", lines=10)
        output_results = gr.Dataframe(label="Вопросы и ответы агента")
    submit_button.click(run_and_submit_all, inputs=[login_button], outputs=[output_status, output_results])

# --- Локальная тестовая функция ---
def test_agent():
    """Тестирование агента с примерами вопросов."""
    agent = LLMGAIAAgent()
    test_questions = [
        "What is 2 + 2?",
        "Who is the first president of the USA?",
        "What is the capital of France?"
    ]
    for question in test_questions:
        answer = agent(question)
        print(f"Вопрос: {question}")
        print(f"Ответ: {answer}")
        print("---")

if __name__ == "__main__":
    test_agent()
    # demo.launch()