app.py

"""
GAIA Agent - Мощный ИИ-агент для прохождения теста GAIA

Этот агент реализует архитектуру, основанную на подходе ReAct (Reasoning + Acting),
с поддержкой мультимодальности, многошагового рассуждения и строгого форматирования ответов.
"""

import os
import json
import re
import requests
from typing import Dict, List, Any, Optional, Union, Tuple
from enum import Enum
import base64
from PIL import Image
from io import BytesIO

# Определение типов сообщений и ролей
class Role(str, Enum):
    SYSTEM = "system"
    USER = "user"
    ASSISTANT = "assistant"
    TOOL = "tool"

class Message:
    def __init__(self, role: Role, content: str, name: Optional[str] = None):
        self.role = role
        self.content = content
        self.name = name
    
    def to_dict(self) -> Dict[str, Any]:
        message = {"role": self.role, "content": self.content}
        if self.name:
            message["name"] = self.name
        return message

# Определение инструментов
class Tool:
    def __init__(self, name: str, description: str, function):
        self.name = name
        self.description = description
        self.function = function
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "type": "function",
            "function": {
                "name": self.name,
                "description": self.description
            }
        }

# Класс для обработки изображений
class ImageProcessor:
    @staticmethod
    def encode_image(image_path: str) -> str:
        """Кодирует изображение в base64 для отправки в API."""
        with open(image_path, "rb") as image_file:
            return base64.b64encode(image_file.read()).decode('utf-8')
    
    @staticmethod
    def decode_image(base64_string: str) -> Image.Image:
        """Декодирует base64-строку в объект изображения."""
        image_data = base64.b64decode(base64_string)
        return Image.open(BytesIO(image_data))
    
    @staticmethod
    def analyze_image(image_path: str) -> Dict[str, Any]:
        """Анализирует изображение и возвращает описание содержимого."""
        # Здесь будет код для анализа изображения с использованием API или локальной модели
        # В реальной реализации это может быть вызов внешнего API для анализа изображений
        return {"description": "Описание изображения будет здесь"}

# Класс для веб-поиска и извлечения информации
class WebSearchTool:
    def __init__(self, api_key: Optional[str] = None):
        self.api_key = api_key
    
    def search(self, query: str) -> List[Dict[str, Any]]:
        """Выполняет веб-поиск по запросу."""
        # Здесь будет код для выполнения веб-поиска
        # В реальной реализации это может быть вызов API поисковой системы
        return [{"title": "Результат поиска", "url": "https://example.com", "snippet": "Фрагмент текста"}]
    
    def extract_content(self, url: str) -> str:
        """Извлекает содержимое веб-страницы."""
        # Здесь будет код для извлечения содержимого веб-страницы
        # В реальной реализации это может быть использование библиотеки requests и BeautifulSoup
        return "Содержимое веб-страницы"

# Класс для форматирования ответов GAIA
class GAIAResponseFormatter:
    @staticmethod
    def extract_format_requirements(question: str) -> Dict[str, Any]:
        """Извлекает требования к формату ответа из вопроса."""
        format_requirements = {
            "list_format": False,
            "ordering": None,
            "separator": None,
            "special_format": None
        }
        
        # Проверка на список
        if re.search(r"(list|comma-separated|semicolon-separated)", question, re.IGNORECASE):
            format_requirements["list_format"] = True
            
            # Определение разделителя
            if re.search(r"comma-separated", question, re.IGNORECASE):
                format_requirements["separator"] = ","
            elif re.search(r"semicolon-separated", question, re.IGNORECASE):
                format_requirements["separator"] = ";"
            else:
                format_requirements["separator"] = ","  # По умолчанию запятая
            
            # Определение порядка
            if re.search(r"(alphabetical|alphabetically)", question, re.IGNORECASE):
                format_requirements["ordering"] = "alphabetical"
            elif re.search(r"(chronological|chronologically)", question, re.IGNORECASE):
                format_requirements["ordering"] = "chronological"
            elif re.search(r"(clockwise)", question, re.IGNORECASE):
                format_requirements["ordering"] = "clockwise"
            
        # Проверка на специальные форматы (например, даты, числа)
        if re.search(r"(YYYY-MM-DD|MM/DD/YYYY)", question, re.IGNORECASE):
            format_requirements["special_format"] = "date"
        
        return format_requirements
    
    @staticmethod
    def format_response(answer: Any, format_requirements: Dict[str, Any]) -> str:
        """Форматирует ответ в соответствии с требованиями."""
        if not format_requirements["list_format"]:
            return str(answer)
        
        if isinstance(answer, list):
            # Применение сортировки, если требуется
            if format_requirements["ordering"] == "alphabetical":
                answer = sorted(answer)
            # Другие типы сортировки могут быть добавлены здесь
            
            # Форматирование списка с указанным разделителем
            separator = format_requirements["separator"] + " " if format_requirements["separator"] else ", "
            return separator.join(answer)
        
        return str(answer)

# Основной класс агента GAIA
class GAIAAgent:
    def __init__(self, api_key: str, model: str = "gpt-4"):
        self.api_key = api_key
        self.model = model
        self.messages = []
        self.tools = self._register_tools()
        self.image_processor = ImageProcessor()
        self.web_search = WebSearchTool()
        self.response_formatter = GAIAResponseFormatter()
    
    def _register_tools(self) -> List[Tool]:
        """Регистрирует доступные инструменты."""
        tools = [
            Tool(
                name="web_search",
                description="Search the web for information",
                function=self.web_search.search
            ),
            Tool(
                name="extract_web_content",
                description="Extract content from a web page",
                function=self.web_search.extract_content
            ),
            Tool(
                name="analyze_image",
                description="Analyze the content of an image",
                function=self.image_processor.analyze_image
            )
        ]
        return tools
    
    def add_system_message(self, content: str) -> None:
        """Добавляет системное сообщение в историю."""
        self.messages.append(Message(Role.SYSTEM, content).to_dict())
    
    def add_user_message(self, content: str, image_path: Optional[str] = None) -> None:
        """Добавляет сообщение пользователя в историю, опционально с изображением."""
        if image_path:
            # Если есть изображение, создаем мультимодальное сообщение
            base64_image = self.image_processor.encode_image(image_path)
            content = [
                {"type": "text", "text": content},
                {
                    "type": "image_url",
                    "image_url": {
                        "url": f"data:image/jpeg;base64,{base64_image}"
                    }
                }
            ]
            self.messages.append({"role": Role.USER, "content": content})
        else:
            self.messages.append(Message(Role.USER, content).to_dict())
    
    def add_assistant_message(self, content: str) -> None:
        """Добавляет сообщение ассистента в историю."""
        self.messages.append(Message(Role.ASSISTANT, content).to_dict())
    
    def add_tool_message(self, tool_name: str, content: str) -> None:
        """Добавляет сообщение от инструмента в историю."""
        self.messages.append({"role": Role.TOOL, "name": tool_name, "content": content})
    
    def _call_llm_api(self, messages: List[Dict[str, Any]], tools: Optional[List[Dict[str, Any]]] = None) -> Dict[str, Any]:
        """Вызывает API языковой модели."""
        # Здесь будет код для вызова API языковой модели (например, OpenAI API)
        # В реальной реализации это будет отправка запроса к API
        
        # Заглушка для демонстрации
        return {
            "choices": [
                {
                    "message": {
                        "role": "assistant",
                        "content": "Это ответ языковой модели."
                    }
                }
            ]
        }
    
    def _parse_tool_calls(self, response: Dict[str, Any]) -> List[Dict[str, Any]]:
        """Парсит вызовы инструментов из ответа языковой модели."""
        # Заглушка для демонстрации
        return []
    
    def _execute_tool_call(self, tool_call: Dict[str, Any]) -> str:
        """Выполняет вызов инструмента."""
        tool_name = tool_call.get("function", {}).get("name")
        arguments = json.loads(tool_call.get("function", {}).get("arguments", "{}"))
        
        for tool in self.tools:
            if tool.name == tool_name:
                result = tool.function(**arguments)
                return json.dumps(result)
        
        return json.dumps({"error": f"Tool {tool_name} not found"})
    
    def process_question(self, question: str, image_path: Optional[str] = None) -> str:
        """Обрабатывает вопрос GAIA и возвращает отформатированный ответ."""
        # Очистка истории сообщений для нового вопроса
        self.messages = []
        
        # Добавление системного сообщения с инструкциями
        self.add_system_message("""
        Ты - мощный ИИ-агент для теста GAIA. Твоя задача - отвечать на вопросы точно и в строгом соответствии с требуемым форматом.
        Используй многошаговое рассуждение и доступные инструменты для поиска информации.
        Твой финальный ответ должен быть кратким, точным и соответствовать формату EXACT MATCH.
        """)
        
        # Добавление вопроса пользователя
        self.add_user_message(question, image_path)
        
        # Извлечение требований к формату ответа
        format_requirements = self.response_formatter.extract_format_requirements(question)
        
        # Цикл рассуждения и действия (ReAct)
        max_iterations = 10
        for _ in range(max_iterations):
            # Вызов языковой модели
            response = self._call_llm_api(self.messages, [tool.to_dict() for tool in self.tools])
            
            # Получение сообщения ассистента
            assistant_message = response["choices"][0]["message"]
            
            # Проверка на наличие вызовов инструментов
            tool_calls = self._parse_tool_calls(assistant_message)
            
            if not tool_calls:
                # Если нет вызовов инструментов, это финальный ответ
                raw_answer = assistant_message.get("content", "")
                
                # Форматирование ответа в соответствии с требованиями
                formatted_answer = self.response_formatter.format_response(raw_answer, format_requirements)
                
                return formatted_answer
            
            # Добавление сообщения ассистента с рассуждением
            self.add_assistant_message(assistant_message.get("content", ""))
            
            # Выполнение вызовов инструментов
            for tool_call in tool_calls:
                tool_result = self._execute_tool_call(tool_call)
                self.add_tool_message(tool_call["function"]["name"], tool_result)
        
        # Если достигнуто максимальное количество итераций без финального ответа
        return "Не удалось сформировать ответ в пределах ограничений."

# Пример использования
if __name__ == "__main__":
    # Инициализация агента
    agent = GAIAAgent(api_key="your_api_key_here")
    
    # Пример вопроса GAIA
    question = "Which of the fruits shown in the 2008 painting 'Embroidery from Uzbekistan' were served as part of the October 1949 breakfast menu for the ocean liner that was later used as a floating prop for the film 'The Last Voyage'? Give the items as a comma-separated list, ordering them in clockwise order based on their arrangement in the painting starting from the 12 o'clock position. Use the plural form of each fruit."
    
    # Обработка вопроса
    answer = agent.process_question(question)
    
    print(f"Question: {question}")
    print(f"Answer: {answer}")