diff --git a/deep_search_agent.py b/deep_search_agent.py new file mode 100644 index 0000000..1d33959 --- /dev/null +++ b/deep_search_agent.py @@ -0,0 +1,215 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- + +""" +Deep‑Search Agent from Scratch +• ищет в интернете (browser tool) +• генерирует виртуальные файлы (write_file tool) +• после завершения выгружает все файлы в реальную FS +""" + +import os +import json +from pathlib import Path +from typing import Dict, List + +from langchain_core.prompts import ChatPromptTemplate +from langchain_openai import ChatOpenAI +from langgraph.checkpoint.memory import InMemorySaver +from langgraph.constants import START +from langgraph.graph import StateGraph +from langgraph.tools import Tool +from langgraph.types import Command, interrupt +from tqdm import tqdm + +# ────────────────────── 1. Конфиги ─────────────────────── +OPENAI_MODEL = "gpt-4o-mini" +TEMPERATURE = 0.2 +MAX_TOKENS = 800 + +# Папка для виртуальных файлов (внутри памяти) +VIRTUAL_FS_ROOT = Path("virtual_fs") + + +# ────────────────────── 2. Модели и промпты ─────────────────── +class AgentState(dict): + """Состояние агента – список сообщений + виртуальная FS.""" + + messages: List[Dict[str, str]] # chat‑history + virtual_files: Dict[str, str] # filename -> content + + +# --------------------- 2.1. Инструмент «browser» --------------------- +def browser_tool(query: str) -> str: + """ + Поисковый инструмент. + Реализован простым запросом к Bing Search API (или любому публичному). + Для примера – возвращаем фиктивный ответ, но в реальном коде + можно использовать `requests` + `bing-search-api`. + """ + # Здесь просто пример. Замените на настоящий HTTP‑запрос. + return f"Результаты поиска по запросу «{query}»:\n1) ...\n2) ..." + + +# --------------------- 2.2. Инструмент «write_file» --------------------- +def write_file_tool(filename: str, content: str) -> str: + """ + Создаёт виртуальный файл (сохраняет в dict). + Возвращает подтверждение. + """ + AgentState.virtual_files[filename] = content + return f"Файл '{filename}' создан/обновлён." + + +# --------------------- 2.3. Инструмент «finalise» --------------------- +def finalise_tool() -> str: + """Вызывается в конце работы – выгружает виртуальную FS.""" + out_dir = Path("output_files") + out_dir.mkdir(exist_ok=True) + for fname, data in AgentState.virtual_files.items(): + p = out_dir / fname + p.parent.mkdir(parents=True, exist_ok=True) + p.write_text(data, encoding="utf-8") + return f"Виртуальная FS выгружена в {out_dir}" + + +# ────────────────────── 3. Создание графа агента ─────────────────── +def build_agent() -> StateGraph: + llm = ChatOpenAI(temperature=TEMPERATURE, model=OPENAI_MODEL) + + # Промпт с инструкциями к инструментам + prompt_template = """ +Вы – исследовательский агент. Ваша задача: собрать информацию по заданной теме, +создать файлы с результатами и после завершения выгрузить их в файловую систему. + +У вас есть три инструмента: +1️⃣ browser(query) — поиск в интернете. +2️⃣ write_file(filename, content) — создание виртуального файла. +3️⃣ finalise() — вызывается только один раз в конце работы и выгружает файлы. + +Используйте формат JSON для вызова инструментов: + +{ + "name": "", + "arguments": { ... } +} + +Пожалуйста, начните с выполнения необходимого поиска. После завершения +выполните finalise(). + +Тема: {topic} +""" + + prompt = ChatPromptTemplate.from_template(prompt_template) + + # Определяем инструменты как объекты LangGraph Tool + browser = Tool( + name="browser", + func=browser_tool, + description="Ищет информацию в интернете по запросу.", + args_schema=None, # простые строки + ) + write_file = Tool( + name="write_file", + func=write_file_tool, + description="Создаёт виртуальный файл с заданным содержимым.", + args_schema=None, + ) + finalise = Tool( + name="finalise", + func=finalise_tool, + description="Выгружает все виртуальные файлы в реальную FS.", + args_schema=None, + ) + + tools = [browser, write_file, finalise] + + # Структура состояния + def start_state(topic: str) -> AgentState: + return { + "messages": [{"role": "system", "content": prompt_template}], + "virtual_files": {}, + } + + # ── Узлы графа ──────────────────────────────────────── + @build_agent.node() + def llm_node(state: AgentState) -> dict: + """LLM генерирует сообщение (может быть вызовом инструмента).""" + messages = state["messages"] + response = llm.invoke(messages) + # Добавляем ответ LLM в историю + return {"messages": messages + [{"role": "assistant", "content": response.content}]} + + @build_agent.node() + def tool_node(state: AgentState) -> dict: + """Выполняет последний вызов инструмента.""" + last_msg = state["messages"][-1] + # Предполагаем, что LLM вернул JSON‑объект с name/arguments + import json as _json + + try: + payload = _json.loads(last_msg["content"]) + tool_name = payload["name"] + arguments = payload.get("arguments", {}) + except Exception: + return {"messages": state["messages"]} # ничего не делаем + + # Находим нужный инструмент + tool = next(t for t in tools if t.name == tool_name) + result = tool.invoke(arguments) + + # Добавляем результат в историю + return { + "messages": state["messages"] + + [{"role": "tool", "content": result, "name": tool_name}] + } + + # Создаём граф + builder = StateGraph(AgentState) + builder.add_node("llm", llm_node) + builder.add_node("tool", tool_node) + + # Условия перехода: если сообщение содержит вызов инструмента → node «tool», иначе — «llm» + def cond(state: AgentState): + last_msg = state["messages"][-1] + if last_msg.get("role") == "assistant" and "{ \"name\": \"" in last_msg[ + "content" + ]: + return "tool" + else: + return "llm" + + builder.set_entry_point("llm") + builder.add_conditional_edges("llm", cond, {"tool": "tool", "llm": "llm"}) + builder.add_edge("tool", END) + + # Поставляем чекпоинтер (необязательно) + memory = InMemorySaver() + return builder.compile(checkpointer=memory) + + +# ────────────────────── 4. Запуск агента ─────────────────── +def main(): + topic = input("\nВведите тему исследования: ").strip() + agent_graph = build_agent() + + # Инициализируем потоковую генерацию + config = {"configurable": {"thread_id": "search-agent-1"}} + for chunk_type, chunk_data in agent_graph.stream( + {"messages": [], "virtual_files": {}}, # стартовое состояние + config=config, + stream_mode=["messages"], + ): + if chunk_type == "messages": + msg = chunk_data["content"] + print(msg, end="", flush=True) + + print("\n\n=== Виртуальная файловая система: ===") + for fname, content in agent_graph.get_state(config)["virtual_files"].items(): + print(f" - {fname} ({len(content)} символов)") + + # После завершения графа – файлы уже выгружены командой `finalise()` + + +if __name__ == "__main__": + main()