feat: solution for '8. Самописный поисковый агент на основе deep agents from scratch'
This commit is contained in:
+137
@@ -0,0 +1,137 @@
|
||||
import os
|
||||
import torch
|
||||
from typing import List, Dict
|
||||
from transformers import AutoTokenizer, AutoModel, AutoModelForCausalLM, pipeline
|
||||
from .utils import bing_search
|
||||
|
||||
class SearchAgent:
|
||||
"""
|
||||
A simple search agent that uses a transformer-based model for natural language
|
||||
understanding and generation, and Bing Web Search API to fetch relevant data.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
nlp_model_name: str = "distilbert-base-uncased",
|
||||
generator_model_name: str = "gpt2",
|
||||
api_key: str = None,
|
||||
):
|
||||
"""
|
||||
Initialize the SearchAgent.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
nlp_model_name : str, optional
|
||||
Hugging Face model name for encoding queries (default: 'distilbert-base-uncased').
|
||||
generator_model_name : str, optional
|
||||
Hugging Face model name for generating summaries (default: 'gpt2').
|
||||
api_key : str, optional
|
||||
Bing Search API key. If not provided, the environment variable
|
||||
BING_API_KEY will be used.
|
||||
"""
|
||||
self.nlp_tokenizer = AutoTokenizer.from_pretrained(nlp_model_name)
|
||||
self.nlp_model = AutoModel.from_pretrained(nlp_model_name)
|
||||
|
||||
self.generator_tokenizer = AutoTokenizer.from_pretrained(generator_model_name)
|
||||
self.generator_model = AutoModelForCausalLM.from_pretrained(generator_model_name)
|
||||
|
||||
device = 0 if torch.cuda.is_available() else -1
|
||||
self.generator = pipeline(
|
||||
"text-generation",
|
||||
model=self.generator_model,
|
||||
tokenizer=self.generator_tokenizer,
|
||||
device=device,
|
||||
)
|
||||
|
||||
self.api_key = api_key or os.getenv("BING_API_KEY")
|
||||
if not self.api_key:
|
||||
raise ValueError(
|
||||
"Bing API key must be provided via parameter or BING_API_KEY env variable"
|
||||
)
|
||||
|
||||
def encode_query(self, query: str):
|
||||
"""
|
||||
Encode the query using the NLP model.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
query : str
|
||||
The natural language query.
|
||||
|
||||
Returns
|
||||
-------
|
||||
torch.Tensor
|
||||
The encoded query representation.
|
||||
"""
|
||||
inputs = self.nlp_tokenizer(query, return_tensors="pt")
|
||||
outputs = self.nlp_model(**inputs)
|
||||
return outputs.last_hidden_state.mean(dim=1)
|
||||
|
||||
def search(self, query: str, count: int = 3) -> List[Dict]:
|
||||
"""
|
||||
Perform a web search using Bing API.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
query : str
|
||||
The search query.
|
||||
count : int, optional
|
||||
Number of results to return (default: 3).
|
||||
|
||||
Returns
|
||||
-------
|
||||
List[Dict]
|
||||
Search results.
|
||||
"""
|
||||
return bing_search(query, self.api_key, count)
|
||||
|
||||
def generate_summary(self, text: str, max_length: int = 150) -> str:
|
||||
"""
|
||||
Generate a summary of the provided text using the generator model.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
text : str
|
||||
Text to summarize.
|
||||
max_length : int, optional
|
||||
Maximum length of the generated summary.
|
||||
|
||||
Returns
|
||||
-------
|
||||
str
|
||||
Generated summary.
|
||||
"""
|
||||
prompt = f"Summarize the following information:\n{text}\nSummary:"
|
||||
outputs = self.generator(prompt, max_length=max_length, num_return_sequences=1)
|
||||
generated = outputs[0]["generated_text"]
|
||||
# Extract the part after "Summary:" if present
|
||||
if "Summary:" in generated:
|
||||
return generated.split("Summary:")[-1].strip()
|
||||
return generated.strip()
|
||||
|
||||
def process_query(self, query: str) -> str:
|
||||
"""
|
||||
Process a user query: search the web and generate a summary.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
query : str
|
||||
The user query.
|
||||
|
||||
Returns
|
||||
-------
|
||||
str
|
||||
The final response to the user.
|
||||
"""
|
||||
results = self.search(query)
|
||||
if not results:
|
||||
return "No results found."
|
||||
|
||||
snippets = "\n".join(
|
||||
[
|
||||
f"{r['name']}\n{r['snippet']}\n{r['url']}"
|
||||
for r in results
|
||||
]
|
||||
)
|
||||
summary = self.generate_summary(snippets)
|
||||
return summary
|
||||
+22
-280
@@ -1,288 +1,30 @@
|
||||
#!/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.
|
||||
"""
|
||||
|
||||
import os
|
||||
import re
|
||||
import ast
|
||||
import argparse
|
||||
import requests
|
||||
from bs4 import BeautifulSoup
|
||||
import click
|
||||
from .agent import SearchAgent
|
||||
|
||||
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
|
||||
@click.command()
|
||||
@click.argument("query", nargs=-1, required=False)
|
||||
def main(query):
|
||||
"""
|
||||
Command-line interface for the Deep Agents search agent.
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Virtual File System
|
||||
# --------------------------------------------------------------------------- #
|
||||
class VirtualFileSystem:
|
||||
"""In-memory virtual file system."""
|
||||
If QUERY is not provided as an argument, the user will be prompted to enter it.
|
||||
"""
|
||||
if not query:
|
||||
query = click.prompt("Enter your query")
|
||||
else:
|
||||
query = " ".join(query)
|
||||
|
||||
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)
|
||||
api_key = os.getenv("BING_API_KEY")
|
||||
if not api_key:
|
||||
click.echo("Error: BING_API_KEY environment variable not set.")
|
||||
return
|
||||
|
||||
agent = SearchAgent(api_key=api_key)
|
||||
click.echo("Searching...")
|
||||
result = agent.process_query(query)
|
||||
click.echo("\nResult:\n")
|
||||
click.echo(result)
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,37 @@
|
||||
import requests
|
||||
from typing import List, Dict
|
||||
|
||||
def bing_search(query: str, api_key: str, count: int = 3) -> List[Dict]:
|
||||
"""
|
||||
Perform a Bing Web Search using the Bing Search API.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
query : str
|
||||
The search query string.
|
||||
api_key : str
|
||||
Bing Search API key.
|
||||
count : int, optional
|
||||
Number of results to return (default is 3).
|
||||
|
||||
Returns
|
||||
-------
|
||||
List[Dict]
|
||||
A list of dictionaries containing 'name', 'url', and 'snippet' for each result.
|
||||
"""
|
||||
endpoint = "https://api.bing.microsoft.com/v7.0/search"
|
||||
headers = {"Ocp-Apim-Subscription-Key": api_key}
|
||||
params = {"q": query, "count": count}
|
||||
response = requests.get(endpoint, headers=headers, params=params, timeout=10)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
results = []
|
||||
for item in data.get("webPages", {}).get("value", []):
|
||||
results.append(
|
||||
{
|
||||
"name": item.get("name"),
|
||||
"url": item.get("url"),
|
||||
"snippet": item.get("snippet"),
|
||||
}
|
||||
)
|
||||
return results
|
||||
Reference in New Issue
Block a user