add deep_search_agent.py
This commit is contained in:
@@ -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()
|
||||
Reference in New Issue
Block a user