task3-deep-search-agent/agent.py

"""
Deep Search Agent — самописный поисковый агент на основе deep agents from scratch.

Возможности:
1. Поиск информации в интернете (DuckDuckGo)
2. Виртуальная файловая система в памяти
3. Выгрузка виртуальных файлов в реальную файловую систему

Основан на: https://github.com/langchain-ai/deep-agents-from-scratch
"""

import os
import json
import datetime
from typing import Optional
from pydantic import BaseModel, Field
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langgraph.prebuilt import create_react_agent


# ═══════════════════════════════════════════════════════════════════════════════
# Виртуальная файловая система
# ═══════════════════════════════════════════════════════════════════════════════

class VirtualFileSystem:
    """Простая файловая система в памяти."""

    def __init__(self):
        self.files: dict[str, str] = {}

    def read(self, file_path: str) -> str:
        """Прочитать файл из виртуальной ФС."""
        if file_path not in self.files:
            return f"Error: File '{file_path}' not found."
        return self.files[file_path]

    def write(self, file_path: str, content: str) -> str:
        """Записать (или перезаписать) файл в виртуальной ФС."""
        self.files[file_path] = content
        return f"Successfully wrote {len(content)} bytes to '{file_path}'."

    def list_files(self) -> str:
        """Список всех файлов в виртуальной ФС."""
        if not self.files:
            return "No files in virtual filesystem."
        lines = [f"  {path} ({len(content)} bytes)"
                 for path, content in self.files.items()]
        return "Virtual filesystem:\n" + "\n".join(lines)

    def export_to_disk(self, output_dir: str = "./output") -> str:
        """Выгрузить все виртуальные файлы на реальный диск."""
        os.makedirs(output_dir, exist_ok=True)
        exported = []
        for path, content in self.files.items():
            full_path = os.path.join(output_dir, path)
            os.makedirs(os.path.dirname(full_path), exist_ok=True)
            with open(full_path, "w", encoding="utf-8") as f:
                f.write(content)
            exported.append(full_path)
        return f"Exported {len(exported)} files to '{output_dir}/':\n" + \
               "\n".join(f"  {p}" for p in exported)


# Глобальный экземпляр виртуальной ФС
vfs = VirtualFileSystem()


# ═══════════════════════════════════════════════════════════════════════════════
# Инструменты агента
# ═══════════════════════════════════════════════════════════════════════════════

@tool
def search_web(query: str) -> str:
    """
    Поиск информации в интернете через DuckDuckGo.
    Используй для поиска фактов, новостей, документации и любой информации.
    Вход: поисковый запрос (строка).
    """
    try:
        from duckduckgo_search import DDGS
        with DDGS() as ddgs:
            results = list(ddgs.text(query, max_results=5))
        if not results:
            return "No results found."

        output = []
        for i, r in enumerate(results, 1):
            output.append(
                f"[{i}] {r.get('title', 'No title')}\n"
                f"    URL: {r.get('href', 'N/A')}\n"
                f"    {r.get('body', 'No description')}"
            )
        return "\n\n".join(output)
    except ImportError:
        return ("Error: duckduckgo_search not installed. "
                "Run: pip install duckduckgo-search")


@tool
def write_file(file_path: str, content: str) -> str:
    """
    Записать файл в виртуальную файловую систему.
    Вход: file_path — путь файла (например, 'report.md', 'data/notes.txt'),
           content — содержимое файла.
    """
    return vfs.write(file_path, content)


@tool
def read_file(file_path: str) -> str:
    """
    Прочитать файл из виртуальной файловой системы.
    Вход: file_path — путь файла.
    """
    return vfs.read(file_path)


@tool
def list_files() -> str:
    """Показать список всех файлов в виртуальной файловой системе."""
    return vfs.list_files()


@tool
def export_files(output_dir: str = "./output") -> str:
    """
    Выгрузить все виртуальные файлы на реальный диск.
    Вход: output_dir — директория для выгрузки (по умолчанию './output').
    """
    return vfs.export_to_disk(output_dir)


TOOLS = [search_web, write_file, read_file, list_files, export_files]


# ═══════════════════════════════════════════════════════════════════════════════
# Сборка агента
# ═══════════════════════════════════════════════════════════════════════════════

SYSTEM_PROMPT = """Ты — Deep Search Agent, автономный поисковый агент.

Твои возможности:
- Искать информацию в интернете (search_web)
- Создавать и редактировать файлы в виртуальной файловой системе (write_file, read_file)
- Просматривать список файлов (list_files)
- Выгружать файлы на реальный диск (export_files)

Правила работы:
1. Сначала исследуй тему через search_web, собери информацию
2. Сохраняй полезные находки в виртуальную файловую систему через write_file
3. При необходимости читай ранее сохранённые файлы через read_file
4. В конце работы ВСЕГДА выгружай файлы на диск через export_files
5. Дай краткий итог выполненной работы

Будь методичен: ищи → анализируй → сохраняй → выгружай.
"""


def build_agent() -> tuple:
    """Создать и вернуть агента + виртуальную ФС."""
    api_base = os.environ.get("OPENAI_API_BASE", "https://api.openai.com/v1")
    api_key = os.environ.get("OPENAI_API_KEY", "")
    model_name = os.environ.get("OPENAI_MODEL", "gpt-4o-mini")

    llm = ChatOpenAI(
        model=model_name,
        openai_api_base=api_base,
        openai_api_key=api_key,
        temperature=0,
    )

    agent = create_react_agent(
        model=llm,
        tools=TOOLS,
        prompt=SYSTEM_PROMPT,
    )

    return agent, vfs


def run_agent(task: str, output_dir: str = "./output") -> dict:
    """
    Запустить агента на выполнение задачи.
    Возвращает словарь с результатом и списком файлов.
    """
    agent, vfs = build_agent()

    # Добавляем инструкцию о выгрузке в конец задачи
    full_task = (
        f"{task}\n\n"
        f"После завершения исследования ОБЯЗАТЕЛЬНО:\n"
        f"1. Вызови list_files() чтобы увидеть все созданные файлы\n"
        f"2. Вызови export_files('{output_dir}') чтобы выгрузить файлы на диск"
    )

    result = agent.invoke(
        {"messages": [{"role": "user", "content": full_task}]}
    )

    # Собираем ответ
    messages = result.get("messages", [])
    last_ai_msg = None
    for msg in reversed(messages):
        if hasattr(msg, "content") and msg.type == "ai":
            last_ai_msg = msg.content
            break

    return {
        "response": last_ai_msg,
        "files_created": list(vfs.files.keys()),
        "output_dir": output_dir,
    }