add deep_search_agent.py

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": "<tool_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()