122 lines
4.9 KiB
Python
122 lines
4.9 KiB
Python
# solution.py
|
|
|
|
"""
|
|
Самописный поисковый агент на основе LangChain.
|
|
Собирает информацию по запросу, создаёт виртуальные файлы,
|
|
выводит их содержимое в консоль и сохраняет на диск.
|
|
"""
|
|
|
|
from langchain.tools import tool
|
|
from langchain_openai import ChatOpenAI
|
|
from langchain.agents import AgentExecutor
|
|
import questionary
|
|
from rich.console import Console
|
|
|
|
# ------------------------------------------------------------------
|
|
# 1. Инструмент поиска
|
|
# ------------------------------------------------------------------
|
|
@tool("search", "Поиск в интернете по запросу")
|
|
def search_tool(query: str) -> list[str]:
|
|
"""
|
|
Возвращает список URL‑ов, найденных по запросу.
|
|
Для демонстрации используется простая имитация поисковой выдачи.
|
|
"""
|
|
# В реальном приложении здесь можно использовать API поисковика.
|
|
# Здесь просто формируем фиктивные ссылки на основе запроса.
|
|
return [
|
|
f"https://example.com/search?q={query.replace(' ', '+')}",
|
|
f"https://another-site.org/{query.replace(' ', '-')}",
|
|
]
|
|
|
|
# ------------------------------------------------------------------
|
|
# 2. Виртуальная файловая система
|
|
# ------------------------------------------------------------------
|
|
class VirtualFileSystem:
|
|
"""
|
|
Хранит файлы в памяти и умеет сохранять их на диск.
|
|
"""
|
|
|
|
def __init__(self):
|
|
# словарь: имя файла -> содержимое
|
|
self.files = {}
|
|
|
|
def add_file(self, name: str, content: str) -> None:
|
|
"""Добавляет файл в память."""
|
|
self.files[name] = content
|
|
|
|
def list_files(self) -> list[str]:
|
|
"""Возвращает список имён файлов."""
|
|
return list(self.files.keys())
|
|
|
|
def get_content(self, name: str) -> str | None:
|
|
"""Получить содержимое файла по имени."""
|
|
return self.files.get(name)
|
|
|
|
def save_to_disk(self, path: str = "./output") -> None:
|
|
"""
|
|
Сохраняет все файлы на диск в указанную папку.
|
|
Если папка не существует – создаёт её.
|
|
"""
|
|
import os
|
|
|
|
os.makedirs(path, exist_ok=True)
|
|
for name, content in self.files.items():
|
|
file_path = os.path.join(path, name)
|
|
with open(file_path, "w", encoding="utf-8") as f:
|
|
f.write(content)
|
|
|
|
# ------------------------------------------------------------------
|
|
# 3. Агент
|
|
# ------------------------------------------------------------------
|
|
def create_agent() -> AgentExecutor:
|
|
"""
|
|
Создаёт агента LangChain с LLM и инструментами.
|
|
"""
|
|
llm = ChatOpenAI(temperature=0) # можно заменить на другую модель
|
|
|
|
tools = [search_tool]
|
|
agent_executor = AgentExecutor.from_llm_and_tools(llm, tools)
|
|
return agent_executor
|
|
|
|
# ------------------------------------------------------------------
|
|
# 4. Основной цикл
|
|
# ------------------------------------------------------------------
|
|
def main() -> None:
|
|
console = Console()
|
|
|
|
# Запрос пользователя
|
|
query = questionary.text("Введите запрос для поиска:").ask()
|
|
if not query:
|
|
console.print("[red]Запрос не введён. Завершение работы.[/]")
|
|
return
|
|
|
|
# Создаём виртуальную файловую систему
|
|
vfs = VirtualFileSystem()
|
|
|
|
# Запускаем агента
|
|
agent = create_agent()
|
|
result = agent.invoke({"input": query})
|
|
|
|
# Предположим, что агент вернёт список URL‑ов и контент для каждого файла.
|
|
# Для демонстрации будем использовать результат инструмента поиска.
|
|
urls: list[str] = search_tool(query)
|
|
|
|
for i, url in enumerate(urls, start=1):
|
|
file_name = f"result_{i}.txt"
|
|
content = f"URL: {url}\n\n(Содержимое страницы не загружено в демо-версии.)"
|
|
vfs.add_file(file_name, content)
|
|
|
|
# Выводим содержимое файлов
|
|
console.print("\n[bold green]Созданные файлы:[/]")
|
|
for name in vfs.list_files():
|
|
console.print(f"\n[underline]{name}[/]")
|
|
console.print(vfs.get_content(name))
|
|
|
|
# Сохраняем на диск
|
|
output_dir = "./output"
|
|
vfs.save_to_disk(output_dir)
|
|
console.print(f"\n[bold blue]Файлы сохранены в папке: {output_dir}[/]")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |