feat: solution for '8. Самописный поисковый агент на основе deep agents from scratch'
This commit is contained in:
+282
-70
@@ -1,76 +1,288 @@
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
# Ensure the virtual_files package is importable
|
||||
sys.path.append(str(Path(__file__).resolve().parent))
|
||||
|
||||
from langchain import OpenAI
|
||||
from langchain.agents import initialize_agent
|
||||
from langchain.tools import DuckDuckGoSearchRun, Tool
|
||||
from src.virtual_files import VirtualFileSystem
|
||||
|
||||
def main():
|
||||
# Initialize the virtual file system
|
||||
vfs = VirtualFileSystem()
|
||||
|
||||
# Define a custom tool to write to the virtual file system
|
||||
def write_file_tool(input_str: str) -> str:
|
||||
"""
|
||||
Expected input format: filename|content
|
||||
Example: python_history.txt|Python was created by Guido van Rossum...
|
||||
"""
|
||||
if "|" not in input_str:
|
||||
return "Error: Input must be in the format 'filename|content'."
|
||||
filename, content = input_str.split("|", 1)
|
||||
filename = filename.strip()
|
||||
content = content.strip()
|
||||
if not filename:
|
||||
return "Error: Filename cannot be empty."
|
||||
vfs.write_file(filename, content)
|
||||
return f"File '{filename}' written successfully."
|
||||
|
||||
write_tool = Tool(
|
||||
name="WriteFile",
|
||||
func=write_file_tool,
|
||||
description=(
|
||||
"Writes content to a virtual file. "
|
||||
"Use the format: filename|content. "
|
||||
"The file will be stored in the virtual file system and exported at the end."
|
||||
),
|
||||
)
|
||||
|
||||
# Search tool
|
||||
search_tool = DuckDuckGoSearchRun()
|
||||
|
||||
# LLM configuration
|
||||
llm = OpenAI(temperature=0)
|
||||
|
||||
# Initialize the agent with the tools
|
||||
agent_executor = initialize_agent(
|
||||
tools=[search_tool, write_tool],
|
||||
llm=llm,
|
||||
agent="zero-shot-react-description",
|
||||
verbose=True,
|
||||
)
|
||||
|
||||
# Example task: gather information about Python programming language
|
||||
task = """
|
||||
You are a research assistant. Your task is to gather information about the Python programming language, including its history, key features, and popular libraries.
|
||||
Create a virtual file named 'python_history.txt' containing the history, a file named 'python_features.txt' containing key features, and a file named 'python_libraries.txt' containing a list of popular libraries.
|
||||
Use the web search tool to find reliable information. After gathering the data, write each section to the corresponding virtual file using the WriteFile tool.
|
||||
Finally, return a summary of what you have done.
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Deep Agent implementation based on the Deep Agents from Scratch template.
|
||||
The agent can search the web, create virtual files, and export them to the real filesystem.
|
||||
"""
|
||||
|
||||
# Run the agent
|
||||
result = agent_executor.run(task)
|
||||
print("\nAgent finished. Result:")
|
||||
print(result)
|
||||
import os
|
||||
import re
|
||||
import ast
|
||||
import argparse
|
||||
import requests
|
||||
from bs4 import BeautifulSoup
|
||||
|
||||
from langchain.llms import OpenAI
|
||||
from langchain.tools import Tool
|
||||
from langchain.prompts import PromptTemplate
|
||||
from langchain.chains import LLMChain
|
||||
from langchain.memory import ConversationBufferMemory
|
||||
from langchain.agents import BaseAgent
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Virtual File System
|
||||
# --------------------------------------------------------------------------- #
|
||||
class VirtualFileSystem:
|
||||
"""In-memory virtual file system."""
|
||||
|
||||
def __init__(self):
|
||||
self.files = {}
|
||||
|
||||
def write_file(self, filename: str, content: str):
|
||||
self.files[filename] = content
|
||||
|
||||
def read_file(self, filename: str) -> str:
|
||||
return self.files.get(filename, "")
|
||||
|
||||
def export_to_disk(self, directory: str):
|
||||
os.makedirs(directory, exist_ok=True)
|
||||
for filename, content in self.files.items():
|
||||
path = os.path.join(directory, filename)
|
||||
with open(path, "w", encoding="utf-8") as f:
|
||||
f.write(content)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Tools
|
||||
# --------------------------------------------------------------------------- #
|
||||
class SearchTool(Tool):
|
||||
"""Search the web using DuckDuckGo."""
|
||||
|
||||
name = "Search"
|
||||
description = (
|
||||
"Search the web for information. Input: query string. Output: search results as text."
|
||||
)
|
||||
|
||||
def __init__(self):
|
||||
super().__init__(name=self.name, description=self.description, func=self.run)
|
||||
|
||||
def run(self, query: str) -> str:
|
||||
url = "https://duckduckgo.com/html/"
|
||||
params = {"q": query}
|
||||
try:
|
||||
response = requests.get(url, params=params, timeout=10)
|
||||
response.raise_for_status()
|
||||
except Exception as e:
|
||||
return f"Error during search: {e}"
|
||||
|
||||
soup = BeautifulSoup(response.text, "html.parser")
|
||||
results = []
|
||||
for a in soup.select("a.result__a"):
|
||||
title = a.get_text()
|
||||
href = a.get("href")
|
||||
results.append(f"{title}\n{href}")
|
||||
if not results:
|
||||
return "No results found."
|
||||
return "\n\n".join(results[:5]) # Return top 5 results
|
||||
|
||||
|
||||
class VirtualFileSystemTool(Tool):
|
||||
"""Create a virtual file with given filename and content."""
|
||||
|
||||
name = "CreateFile"
|
||||
description = (
|
||||
"Create a virtual file with given filename and content. "
|
||||
"Input format: filename|content"
|
||||
)
|
||||
|
||||
def __init__(self, vfs: VirtualFileSystem):
|
||||
self.vfs = vfs
|
||||
super().__init__(name=self.name, description=self.description, func=self.run)
|
||||
|
||||
def run(self, args: str) -> str:
|
||||
parts = args.split("|", 1)
|
||||
if len(parts) != 2:
|
||||
return "Error: expected format 'filename|content'"
|
||||
filename, content = parts
|
||||
filename = filename.strip()
|
||||
content = content.strip()
|
||||
self.vfs.write_file(filename, content)
|
||||
return f"File '{filename}' created."
|
||||
|
||||
|
||||
class ExportTool(Tool):
|
||||
"""Export all virtual files to the specified directory."""
|
||||
|
||||
name = "ExportFiles"
|
||||
description = "Export all virtual files to the specified directory. Input: directory path"
|
||||
|
||||
def __init__(self, vfs: VirtualFileSystem):
|
||||
self.vfs = vfs
|
||||
super().__init__(name=self.name, description=self.description, func=self.run)
|
||||
|
||||
def run(self, directory: str) -> str:
|
||||
directory = directory.strip()
|
||||
self.vfs.export_to_disk(directory)
|
||||
return f"Exported {len(self.vfs.files)} files to '{directory}'."
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Deep Agent
|
||||
# --------------------------------------------------------------------------- #
|
||||
class DeepAgent(BaseAgent):
|
||||
"""Deep Agent following the Deep Agents from Scratch template."""
|
||||
|
||||
def __init__(self, llm, tools, memory, verbose=False):
|
||||
super().__init__()
|
||||
self.llm = llm
|
||||
self.tools = tools
|
||||
self.memory = memory
|
||||
self.verbose = verbose
|
||||
|
||||
self.tool_names = [tool.name for tool in tools]
|
||||
self.tool_map = {tool.name: tool for tool in tools}
|
||||
|
||||
# Prompt templates
|
||||
self.plan_prompt = PromptTemplate(
|
||||
input_variables=["input", "agent_scratchpad"],
|
||||
template=(
|
||||
"You are a helpful assistant. Your task is to answer the user query: {input}\n"
|
||||
"You have access to the following tools: {tool_names}\n"
|
||||
"You can use the tools to gather information.\n"
|
||||
"Plan your steps. Each step should be a single tool call in the format: {tool_name}({arguments})\n"
|
||||
"If you have enough information, provide the final answer.\n"
|
||||
"Your plan:\n{agent_scratchpad}"
|
||||
),
|
||||
)
|
||||
|
||||
self.tool_prompt = PromptTemplate(
|
||||
input_variables=["tool_input", "agent_scratchpad"],
|
||||
template=(
|
||||
"You are a tool. The user wants: {tool_input}\n"
|
||||
"You have the following context: {agent_scratchpad}\n"
|
||||
"You should produce the tool output.\n"
|
||||
"Tool output:"
|
||||
),
|
||||
)
|
||||
|
||||
self.final_prompt = PromptTemplate(
|
||||
input_variables=["agent_scratchpad"],
|
||||
template=(
|
||||
"You are a helpful assistant. Based on the previous steps, provide the final answer.\n"
|
||||
"Answer:\n{agent_scratchpad}"
|
||||
),
|
||||
)
|
||||
|
||||
self.plan_chain = LLMChain(llm=self.llm, prompt=self.plan_prompt)
|
||||
self.tool_chain = LLMChain(llm=self.llm, prompt=self.tool_prompt)
|
||||
self.final_chain = LLMChain(llm=self.llm, prompt=self.final_prompt)
|
||||
|
||||
def _get_scratchpad(self) -> str:
|
||||
"""Return the conversation history as a string."""
|
||||
history = self.memory.load_memory_variables({})["history"]
|
||||
if isinstance(history, list):
|
||||
return "\n".join([msg.content if hasattr(msg, "content") else str(msg) for msg in history])
|
||||
return str(history)
|
||||
|
||||
def plan(self, input_text: str) -> str:
|
||||
scratchpad = self._get_scratchpad()
|
||||
plan = self.plan_chain.run(
|
||||
input=input_text,
|
||||
agent_scratchpad=scratchpad,
|
||||
tool_names=", ".join(self.tool_names),
|
||||
)
|
||||
return plan
|
||||
|
||||
def parse_plan(self, plan_text: str) -> list:
|
||||
"""Parse the plan into individual steps."""
|
||||
lines = [line.strip() for line in plan_text.splitlines() if line.strip()]
|
||||
return lines
|
||||
|
||||
def parse_step(self, step_text: str) -> tuple:
|
||||
"""Parse a single step into tool name and arguments."""
|
||||
match = re.match(r"(\w+)\((.*)\)", step_text)
|
||||
if not match:
|
||||
raise ValueError(f"Invalid step format: {step_text}")
|
||||
tool_name = match.group(1)
|
||||
args_str = match.group(2).strip()
|
||||
# Try to parse arguments as a tuple
|
||||
try:
|
||||
args_tuple = ast.literal_eval(f"({args_str},)")
|
||||
except Exception:
|
||||
args_tuple = (args_str,)
|
||||
return tool_name, args_tuple
|
||||
|
||||
def run(self, input_text: str) -> str:
|
||||
"""Run the agent on the given input."""
|
||||
self.memory.save_context({"input": input_text}, {})
|
||||
plan = self.plan(input_text)
|
||||
if self.verbose:
|
||||
print("\n=== PLAN ===")
|
||||
print(plan)
|
||||
print("============\n")
|
||||
|
||||
steps = self.parse_plan(plan)
|
||||
for step in steps:
|
||||
tool_name, args_tuple = self.parse_step(step)
|
||||
if tool_name not in self.tool_map:
|
||||
raise ValueError(f"Unknown tool: {tool_name}")
|
||||
tool = self.tool_map[tool_name]
|
||||
# Use the first argument as the tool input
|
||||
tool_input = args_tuple[0] if args_tuple else ""
|
||||
if self.verbose:
|
||||
print(f"\n=== TOOL CALL: {tool_name} ===")
|
||||
print(f"Input: {tool_input}")
|
||||
tool_output = tool.run(tool_input)
|
||||
if self.verbose:
|
||||
print(f"Output: {tool_output}\n")
|
||||
self.memory.save_context({"tool_output": tool_output}, {})
|
||||
|
||||
final_answer = self.final_chain.run(agent_scratchpad=self._get_scratchpad())
|
||||
return final_answer
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Main
|
||||
# --------------------------------------------------------------------------- #
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Deep Agent Demo")
|
||||
parser.add_argument(
|
||||
"--prompt",
|
||||
type=str,
|
||||
required=True,
|
||||
help="The user query for the agent to answer.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
type=str,
|
||||
default="output_files",
|
||||
help="Directory to export virtual files.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--verbose",
|
||||
action="store_true",
|
||||
help="Print detailed logs.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
# Ensure OpenAI API key is set
|
||||
if "OPENAI_API_KEY" not in os.environ:
|
||||
raise RuntimeError("Please set the OPENAI_API_KEY environment variable.")
|
||||
|
||||
# Initialize components
|
||||
llm = OpenAI(temperature=0, model_name="gpt-3.5-turbo")
|
||||
vfs = VirtualFileSystem()
|
||||
tools = [
|
||||
SearchTool(),
|
||||
VirtualFileSystemTool(vfs),
|
||||
ExportTool(vfs),
|
||||
]
|
||||
memory = ConversationBufferMemory(memory_key="history", return_messages=True)
|
||||
|
||||
agent = DeepAgent(llm=llm, tools=tools, memory=memory, verbose=args.verbose)
|
||||
|
||||
# Run the agent
|
||||
print("\n=== RUNNING AGENT ===")
|
||||
final_answer = agent.run(args.prompt)
|
||||
print("\n=== FINAL ANSWER ===")
|
||||
print(final_answer)
|
||||
|
||||
# Export files (in case the agent didn't call ExportFiles)
|
||||
if not any(tool.name == "ExportFiles" for tool in tools):
|
||||
export_tool = ExportTool(vfs)
|
||||
export_output = export_tool.run(args.output_dir)
|
||||
print("\n=== EXPORT OUTPUT ===")
|
||||
print(export_output)
|
||||
|
||||
# Export virtual files to disk
|
||||
output_dir = Path(__file__).resolve().parent / "output_files"
|
||||
vfs.export_to_disk(str(output_dir))
|
||||
print(f"\nVirtual files exported to {output_dir}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user