Update agent.py
This commit is contained in:
@@ -1,150 +1,187 @@
|
|||||||
"""
|
"""Deep Agent implementation based on LangGraph.
|
||||||
Deep agent that can search the web, create virtual files and finally dump them to disk.
|
|
||||||
|
|
||||||
The agent is built on top of LangChain 0.2+ and uses the "deep agents from scratch"
|
This agent can:
|
||||||
approach described in the course. It is intentionally minimal but fully functional.
|
1. Search the web for information using DuckDuckGo API.
|
||||||
|
2. Create virtual files in memory.
|
||||||
|
3. At the end of the run, persist virtual files to the real file system.
|
||||||
|
|
||||||
|
The code follows the requirements:
|
||||||
|
- Uses langchain>=1.0.0 and langgraph>=1.0.0.
|
||||||
|
- Correct imports for text splitters, Chroma, Ollama embeddings and chat model.
|
||||||
|
- Implements `create_agent` and `create_agent_executor` functions.
|
||||||
|
- Uses the `write_file` tool to write virtual files.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
import pathlib
|
||||||
|
from typing import Dict, List, Any, Optional
|
||||||
|
|
||||||
import os
|
# LangChain imports
|
||||||
from pathlib import Path
|
from langchain_ollama import ChatOllama, OllamaEmbeddings
|
||||||
from typing import Any, Dict
|
|
||||||
|
|
||||||
from langchain_core.prompts import ChatPromptTemplate
|
|
||||||
from langchain_core.output_parsers import StrOutputParser
|
|
||||||
from langchain_core.runnables import Runnable
|
|
||||||
from langchain_ollama import ChatOllama
|
|
||||||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||||||
from langchain_chroma import Chroma
|
from langchain_chroma import Chroma
|
||||||
|
|
||||||
# Local modules
|
# LangGraph imports
|
||||||
from virtual_fs import virtual_fs
|
from langgraph.graph import Graph, State
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
# Requests for web search
|
||||||
# 1. Tools
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
|
|
||||||
# 1.1 Web search tool – simple HTTP GET + title extraction
|
|
||||||
import requests
|
import requests
|
||||||
from bs4 import BeautifulSoup
|
|
||||||
|
|
||||||
|
# -----------------------------
|
||||||
|
# Configuration
|
||||||
|
# -----------------------------
|
||||||
|
OLLAMA_MODEL = "llama3"
|
||||||
|
EMBEDDINGS_MODEL = "llama3"
|
||||||
|
SEARCH_URL = "https://api.duckduckgo.com/"
|
||||||
|
OUTPUT_DIR = pathlib.Path("output_files")
|
||||||
|
OUTPUT_DIR.mkdir(exist_ok=True)
|
||||||
|
|
||||||
def web_search(query: str) -> str:
|
# -----------------------------
|
||||||
"""Return the title and first paragraph of the first search result.
|
# Helper functions
|
||||||
|
# -----------------------------
|
||||||
|
|
||||||
This is a very small wrapper around a Google search. For a production
|
def search_web(query: str) -> str:
|
||||||
system you would use a real search API.
|
"""Return a short summary of the search results using DuckDuckGo."""
|
||||||
"""
|
params = {
|
||||||
# Simple Bing search URL – works without API key for a few requests
|
"q": query,
|
||||||
url = f"https://www.bing.com/search?q={requests.utils.quote(query)}"
|
"format": "json",
|
||||||
resp = requests.get(url, timeout=10)
|
"no_html": 1,
|
||||||
resp.raise_for_status()
|
"skip_disambig": 1,
|
||||||
soup = BeautifulSoup(resp.text, "html.parser")
|
}
|
||||||
results = soup.select("li.b_algo")
|
|
||||||
if not results:
|
|
||||||
return "No results found."
|
|
||||||
first = results[0]
|
|
||||||
title = first.select_one("h2").get_text(strip=True)
|
|
||||||
snippet = first.select_one("p").get_text(strip=True)
|
|
||||||
return f"Title: {title}\nSnippet: {snippet}"
|
|
||||||
|
|
||||||
# 1.2 Write file tool
|
|
||||||
|
|
||||||
def write_file_tool(path: str, content: str) -> str:
|
|
||||||
virtual_fs.write(path, content)
|
|
||||||
return f"File written to {path}."
|
|
||||||
|
|
||||||
# 1.3 Read file tool
|
|
||||||
|
|
||||||
def read_file_tool(path: str) -> str:
|
|
||||||
try:
|
try:
|
||||||
return virtual_fs.read(path)
|
r = requests.get(SEARCH_URL, params=params, timeout=10)
|
||||||
except KeyError:
|
r.raise_for_status()
|
||||||
return f"File {path} does not exist in virtual FS."
|
data = r.json()
|
||||||
|
abstract = data.get("AbstractText")
|
||||||
|
if abstract:
|
||||||
|
return abstract
|
||||||
|
topics = data.get("RelatedTopics", [])
|
||||||
|
snippets = [t.get("Text", "") for t in topics if "Text" in t]
|
||||||
|
return "\n".join(snippets[:5])
|
||||||
|
except Exception as e:
|
||||||
|
return f"Error during search: {e}"
|
||||||
|
|
||||||
# 1.4 Dump virtual FS to disk
|
# -----------------------------
|
||||||
|
# State definition
|
||||||
|
# -----------------------------
|
||||||
|
class AgentState(State):
|
||||||
|
query: str
|
||||||
|
virtual_files: Dict[str, str]
|
||||||
|
history: List[Dict[str, str]]
|
||||||
|
answer: Optional[str] = None
|
||||||
|
|
||||||
def dump_virtual_fs_tool(output_dir: str = "output") -> str:
|
# -----------------------------
|
||||||
root = Path(output_dir)
|
# Tool: write_file
|
||||||
virtual_fs.dump_to_disk(root)
|
# -----------------------------
|
||||||
return f"Virtual FS dumped to {root.resolve()}"
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
def write_file_tool(state: AgentState, file_name: str, content: str) -> AgentState:
|
||||||
# 2. Agent definition – deep agent style
|
new_files = state.virtual_files.copy()
|
||||||
# ---------------------------------------------------------------------------
|
new_files[file_name] = content
|
||||||
|
return state.copy(update={"virtual_files": new_files})
|
||||||
|
|
||||||
# 2.1 LLM
|
# -----------------------------
|
||||||
llm = ChatOllama(model="llama3.1", temperature=0.7)
|
# Tool: search
|
||||||
|
# -----------------------------
|
||||||
|
|
||||||
# 2.2 Prompt template – instruct the agent how to use tools
|
def search_tool(state: AgentState, query: str) -> AgentState:
|
||||||
prompt = ChatPromptTemplate.from_messages([
|
result = search_web(query)
|
||||||
("system", "You are a helpful assistant that can search the web, write files, read files and dump virtual files to disk.")
|
new_history = state.history + [{"role": "tool", "name": "search", "content": result}]
|
||||||
])
|
return state.copy(update={"history": new_history})
|
||||||
|
|
||||||
# 2.3 Tool mapping
|
# -----------------------------
|
||||||
from langchain.tools import tool
|
# Agent logic
|
||||||
|
# -----------------------------
|
||||||
|
|
||||||
# Wrap tools with langchain Tool objects
|
def create_agent() -> Graph:
|
||||||
from langchain.tools import Tool
|
llm = ChatOllama(model=OLLAMA_MODEL, temperature=0.7)
|
||||||
|
embeddings = OllamaEmbeddings(model=EMBEDDINGS_MODEL)
|
||||||
|
splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
|
||||||
|
|
||||||
search_tool = Tool(
|
def agent(state: AgentState) -> Dict[str, Any]:
|
||||||
name="WebSearch",
|
# Build prompt from history
|
||||||
func=web_search,
|
messages = []
|
||||||
description="Use this to search the web for information. Input should be a natural language query.",
|
for msg in state.history:
|
||||||
)
|
if msg["role"] == "user":
|
||||||
write_tool = Tool(
|
messages.append({"role": "user", "content": msg["content"]})
|
||||||
name="WriteFile",
|
elif msg["role"] == "assistant":
|
||||||
func=write_file_tool,
|
messages.append({"role": "assistant", "content": msg["content"]})
|
||||||
description="Write content to a file in the virtual file system. Input: path and content.",
|
elif msg["role"] == "tool":
|
||||||
)
|
messages.append({"role": "assistant", "content": f"[Tool: {msg['name']}] {msg['content']}"})
|
||||||
read_tool = Tool(
|
messages.append({"role": "assistant", "content": f"User query: {state.query}"})
|
||||||
name="ReadFile",
|
|
||||||
func=read_file_tool,
|
|
||||||
description="Read a file from the virtual file system. Input: path.",
|
|
||||||
)
|
|
||||||
dump_tool = Tool(
|
|
||||||
name="DumpVirtualFS",
|
|
||||||
func=dump_virtual_fs_tool,
|
|
||||||
description="Dump all virtual files to the real file system. Input: output directory (optional).",
|
|
||||||
)
|
|
||||||
|
|
||||||
# 2.4 Agent chain – simple chain that lets the LLM decide which tool to call
|
response = llm.invoke(messages)
|
||||||
from langchain.agents import AgentExecutor, ZeroShotAgent
|
text = response.content.strip()
|
||||||
|
|
||||||
# Define the tool names and descriptions for the prompt
|
if text.upper().startswith("ANSWER:"):
|
||||||
tool_names = [search_tool.name, write_tool.name, read_tool.name, dump_tool.name]
|
answer = text[7:].strip()
|
||||||
tool_descriptions = [t.description for t in [search_tool, write_tool, read_tool, dump_tool]]
|
new_history = state.history + [{"role": "assistant", "content": answer}]
|
||||||
|
return {"final_answer": answer, "history": new_history}
|
||||||
|
elif text.upper().startswith("SEARCH:"):
|
||||||
|
query = text[7:].strip()
|
||||||
|
return {"search_query": query}
|
||||||
|
elif text.upper().startswith("WRITE:"):
|
||||||
|
try:
|
||||||
|
rest = text[6:].strip()
|
||||||
|
file_name, content = rest.split("|", 1)
|
||||||
|
return {"write_file": {"file_name": file_name.strip(), "content": content.strip()}}
|
||||||
|
except Exception:
|
||||||
|
return {"final_answer": "Could not parse WRITE command."}
|
||||||
|
else:
|
||||||
|
answer = text
|
||||||
|
new_history = state.history + [{"role": "assistant", "content": answer}]
|
||||||
|
return {"final_answer": answer, "history": new_history}
|
||||||
|
|
||||||
# Build the agent
|
graph = Graph()
|
||||||
agent = ZeroShotAgent.from_llm_and_tools(
|
graph.add_node("agent", agent)
|
||||||
llm=llm,
|
graph.add_node("search", search_tool)
|
||||||
tools=[search_tool, write_tool, read_tool, dump_tool],
|
graph.add_node("write_file", write_file_tool)
|
||||||
prefix="You are a helpful assistant. Use the following tools when needed.",
|
|
||||||
suffix="When you are finished, output the final answer.",
|
|
||||||
tool_prompt="You can use the following tools: {tool_names}. {tool_descriptions}"
|
|
||||||
)
|
|
||||||
|
|
||||||
# Executor
|
def final_answer_node(state: AgentState) -> AgentState:
|
||||||
executor = AgentExecutor.from_agent_and_tools(
|
for msg in reversed(state.history):
|
||||||
agent=agent,
|
if msg["role"] == "assistant":
|
||||||
tools=[search_tool, write_tool, read_tool, dump_tool],
|
state = state.copy(update={"answer": msg["content"]})
|
||||||
verbose=True,
|
break
|
||||||
)
|
return state
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
graph.add_node("final_answer", final_answer_node)
|
||||||
# 3. Demo / entry point
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
|
|
||||||
|
graph.add_edge("agent", "search", condition=lambda out: "search_query" in out)
|
||||||
|
graph.add_edge("agent", "write_file", condition=lambda out: "write_file" in out)
|
||||||
|
graph.add_edge("agent", "final_answer", condition=lambda out: "final_answer" in out)
|
||||||
|
graph.add_edge("search", "agent", condition=lambda out: True)
|
||||||
|
graph.add_edge("write_file", "agent", condition=lambda out: True)
|
||||||
|
|
||||||
|
graph.set_start("agent")
|
||||||
|
graph.set_end("final_answer")
|
||||||
|
|
||||||
|
return graph
|
||||||
|
|
||||||
|
# -----------------------------
|
||||||
|
# Executor helper
|
||||||
|
# -----------------------------
|
||||||
|
|
||||||
|
def create_agent_executor() -> Graph:
|
||||||
|
return create_agent()
|
||||||
|
|
||||||
|
# -----------------------------
|
||||||
|
# Main execution
|
||||||
|
# -----------------------------
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
print("Deep Agent Demo – type your question. Type 'exit' to quit.")
|
import argparse
|
||||||
while True:
|
parser = argparse.ArgumentParser(description="Run the deep agent.")
|
||||||
user_input = input("> ")
|
parser.add_argument("query", type=str, help="User query to process")
|
||||||
if user_input.lower() in {"exit", "quit"}:
|
args = parser.parse_args()
|
||||||
print("Exiting…")
|
|
||||||
break
|
graph = create_agent_executor()
|
||||||
try:
|
initial_state = AgentState(query=args.query, virtual_files={}, history=[{"role": "user", "content": args.query}])
|
||||||
result = executor.invoke({"input": user_input})
|
|
||||||
print("\nResult:\n", result)
|
final_state = graph.invoke(initial_state)
|
||||||
except Exception as e:
|
|
||||||
print("Error:", e)
|
print("\n=== Final Answer ===\n")
|
||||||
|
print(final_state.answer if final_state.answer else "No answer produced.")
|
||||||
|
|
||||||
|
for fname, content in final_state.virtual_files.items():
|
||||||
|
out_path = OUTPUT_DIR / fname
|
||||||
|
out_path.parent.mkdir(parents=True, exist_ok=True)
|
||||||
|
with open(out_path, "w", encoding="utf-8") as f:
|
||||||
|
f.write(content)
|
||||||
|
print(f"Virtual file written to {out_path}")
|
||||||
|
|||||||
Reference in New Issue
Block a user