add main.py
This commit is contained in:
@@ -0,0 +1,246 @@
|
||||
import os
|
||||
import json
|
||||
import re
|
||||
from typing import Dict, List
|
||||
|
||||
import requests
|
||||
from bs4 import BeautifulSoup
|
||||
from openai import OpenAI
|
||||
from rich.console import Console
|
||||
from rich.table import Table
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Виртуальная файловая система
|
||||
# --------------------------------------------------------------------------- #
|
||||
class VirtualFileSystem:
|
||||
"""Хранит файлы в памяти. Позволяет создавать и экспортировать их."""
|
||||
|
||||
def __init__(self):
|
||||
self.files: Dict[str, str] = {}
|
||||
|
||||
def write(self, filename: str, content: str) -> None:
|
||||
self.files[filename] = content
|
||||
|
||||
def read(self, filename: str) -> str | None:
|
||||
return self.files.get(filename)
|
||||
|
||||
def list_files(self) -> List[str]:
|
||||
return list(self.files.keys())
|
||||
|
||||
def export_all(self, directory: str = ".") -> None:
|
||||
os.makedirs(directory, exist_ok=True)
|
||||
for name, content in self.files.items():
|
||||
path = os.path.join(directory, name)
|
||||
with open(path, "w", encoding="utf-8") as f:
|
||||
f.write(content)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Инструменты
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def web_search(query: str) -> str:
|
||||
"""Поиск в интернете через DuckDuckGo, возвращает первые 3 результата."""
|
||||
url = "https://duckduckgo.com/html/"
|
||||
params = {"q": query}
|
||||
try:
|
||||
resp = requests.get(url, params=params, timeout=10,
|
||||
headers={"User-Agent": "Mozilla/5.0"})
|
||||
resp.raise_for_status()
|
||||
except Exception as e:
|
||||
return f"Ошибка при поиске: {e}"
|
||||
|
||||
soup = BeautifulSoup(resp.text, "html.parser")
|
||||
results = []
|
||||
for a in soup.select("a.result__a")[:3]:
|
||||
title = a.get_text(strip=True)
|
||||
href = a["href"]
|
||||
match = re.search(r"uddg=([^&]+)", href)
|
||||
if match:
|
||||
href = requests.utils.unquote(match.group(1))
|
||||
results.append(f"{title}\n{href}")
|
||||
return "\n\n".join(results) or "Ничего не найдено."
|
||||
|
||||
|
||||
def create_file(vfs: VirtualFileSystem, filename: str, content: str) -> str:
|
||||
"""Создаёт или обновляет виртуальный файл."""
|
||||
vfs.write(filename, content)
|
||||
return f"Файл '{filename}' успешно создан/обновлён."
|
||||
|
||||
|
||||
def read_file(vfs: VirtualFileSystem, filename: str) -> str:
|
||||
"""Читает файл из виртуальной ФС."""
|
||||
content = vfs.read(filename)
|
||||
if content is None:
|
||||
return f"Файл '{filename}' не найден."
|
||||
return content
|
||||
|
||||
|
||||
def list_files(vfs: VirtualFileSystem) -> str:
|
||||
"""Список файлов в виртуальной ФС."""
|
||||
files = vfs.list_files()
|
||||
if not files:
|
||||
return "Виртуальная ФС пуста."
|
||||
return "\n".join(files)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Описание инструментов для LLM
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
TOOLS_DESCRIPTION = """Ты — автономный поисковый агент. Работай в цикле Thought/Action/Observation.
|
||||
|
||||
Доступные инструменты:
|
||||
1. web_search — поиск в интернете. Action Input: поисковый запрос (строка)
|
||||
2. create_file — создать файл в виртуальной ФС. Action Input: JSON {"filename": "...", "content": "..."}
|
||||
3. read_file — прочитать файл из виртуальной ФС. Action Input: имя файла (строка)
|
||||
4. list_files — список файлов в виртуальной ФС. Action Input: (пусто)
|
||||
5. finish — завершить работу. Action Input: финальный ответ пользователю
|
||||
|
||||
Формат каждого шага — СТРОГО такой (три строки, никаких отклонений):
|
||||
Thought: <что ты думаешь, что нужно сделать>
|
||||
Action: <одно из: web_search | create_file | read_file | list_files | finish>
|
||||
Action Input: <аргумент>
|
||||
|
||||
Правила:
|
||||
- Сначала ищи информацию через web_search.
|
||||
- После сбора информации сохрани результат через create_file.
|
||||
- Заверши работу командой finish только после того, как файл создан.
|
||||
- Никогда не выдумывай информацию — используй только то, что получил из web_search.
|
||||
"""
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Агентный цикл (ReAct: Reason + Act)
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def run_agent(task: str, vfs: VirtualFileSystem, console: Console) -> str:
|
||||
"""
|
||||
Агентный цикл ReAct from scratch:
|
||||
LLM сам решает когда и какой инструмент вызвать,
|
||||
получает результат (Observation) и продолжает рассуждение.
|
||||
"""
|
||||
client = OpenAI()
|
||||
|
||||
messages = [
|
||||
{"role": "system", "content": TOOLS_DESCRIPTION},
|
||||
{"role": "user", "content": task},
|
||||
]
|
||||
|
||||
for step in range(15):
|
||||
# --- LLM генерирует следующий шаг ---
|
||||
try:
|
||||
response = client.chat.completions.create(
|
||||
model="gpt-4o-mini",
|
||||
messages=messages,
|
||||
temperature=0,
|
||||
)
|
||||
except Exception as e:
|
||||
return f"Ошибка при вызове OpenAI: {e}"
|
||||
|
||||
raw = response.choices[0].message.content.strip()
|
||||
console.print(f"\n[bold yellow][Шаг {step + 1}][/]\n{raw}")
|
||||
messages.append({"role": "assistant", "content": raw})
|
||||
|
||||
# --- Парсинг Action и Action Input ---
|
||||
action_match = re.search(r"^Action:\s*(\w+)", raw, re.MULTILINE)
|
||||
input_match = re.search(r"^Action Input:\s*(.+)", raw, re.MULTILINE | re.DOTALL)
|
||||
|
||||
if not action_match:
|
||||
break
|
||||
|
||||
action = action_match.group(1).strip()
|
||||
action_input = input_match.group(1).strip() if input_match else ""
|
||||
|
||||
# --- Выполнение инструмента ---
|
||||
if action == "finish":
|
||||
return action_input
|
||||
|
||||
elif action == "web_search":
|
||||
observation = web_search(action_input)
|
||||
|
||||
elif action == "create_file":
|
||||
try:
|
||||
args = json.loads(action_input)
|
||||
observation = create_file(vfs, args["filename"], args["content"])
|
||||
except (json.JSONDecodeError, KeyError) as e:
|
||||
observation = f"Ошибка: неверный формат JSON для create_file: {e}"
|
||||
|
||||
elif action == "read_file":
|
||||
observation = read_file(vfs, action_input)
|
||||
|
||||
elif action == "list_files":
|
||||
observation = list_files(vfs)
|
||||
|
||||
else:
|
||||
observation = (
|
||||
f"Неизвестный инструмент: '{action}'. "
|
||||
"Используй только: web_search, create_file, read_file, list_files, finish."
|
||||
)
|
||||
|
||||
# --- Передаём результат обратно в LLM ---
|
||||
preview = observation[:300] + ("..." if len(observation) > 300 else "")
|
||||
console.print(f"[dim]Observation: {preview}[/]")
|
||||
messages.append({"role": "user", "content": f"Observation: {observation}"})
|
||||
|
||||
return "Агент завершил работу (достигнут лимит шагов)."
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Выгрузка виртуальных файлов на диск
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def flush_virtual_fs(vfs: VirtualFileSystem, console: Console) -> None:
|
||||
"""Экспортирует все файлы из виртуальной ФС в реальную файловую систему."""
|
||||
if not vfs.list_files():
|
||||
console.print("[red]Нет виртуальных файлов для выгрузки.[/]")
|
||||
return
|
||||
output_dir = "output"
|
||||
vfs.export_all(output_dir)
|
||||
console.print(f"\n[green]Файлы выгружены в '{os.path.abspath(output_dir)}':[/]")
|
||||
for name in vfs.list_files():
|
||||
console.print(f" ✓ {name}")
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Точка входа
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def main():
|
||||
console = Console()
|
||||
vfs = VirtualFileSystem()
|
||||
|
||||
console.print("[bold cyan]═══ Deep Search Agent ═══[/]\n")
|
||||
task = input("Введите задачу для агента: ").strip()
|
||||
if not task:
|
||||
console.print("[red]Задача не задана. Выход.[/]")
|
||||
return
|
||||
|
||||
console.print(f"\n[bold]Задача:[/] {task}")
|
||||
console.print("[bold cyan]\nЗапуск агентного цикла...[/]")
|
||||
|
||||
# --- Запуск агента ---
|
||||
result = run_agent(task, vfs, console)
|
||||
|
||||
# --- Финальный ответ ---
|
||||
console.print("\n[bold green]═══ Финальный ответ агента ═══[/]")
|
||||
console.print(result)
|
||||
|
||||
# --- Таблица виртуальных файлов ---
|
||||
console.print("\n[bold magenta]Виртуальные файлы, созданные агентом:[/]")
|
||||
table = Table(show_header=True, header_style="bold blue")
|
||||
table.add_column("Имя файла", style="dim", min_width=20)
|
||||
table.add_column("Размер (байт)")
|
||||
for fname in vfs.list_files():
|
||||
size = str(len(vfs.read(fname) or ""))
|
||||
table.add_row(fname, size)
|
||||
console.print(table)
|
||||
|
||||
# --- Выгрузка на диск ---
|
||||
console.print("\n[bold cyan]Выгрузка виртуальных файлов в реальную ФС...[/]")
|
||||
flush_virtual_fs(vfs, console)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user