216 lines
8.9 KiB
Python
216 lines
8.9 KiB
Python
#!/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()
|