151 lines
5.0 KiB
Python
151 lines
5.0 KiB
Python
"""
|
||
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"
|
||
approach described in the course. It is intentionally minimal but fully functional.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
from pathlib import Path
|
||
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_chroma import Chroma
|
||
|
||
# Local modules
|
||
from virtual_fs import virtual_fs
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 1. Tools
|
||
# ---------------------------------------------------------------------------
|
||
|
||
# 1.1 Web search tool – simple HTTP GET + title extraction
|
||
import requests
|
||
from bs4 import BeautifulSoup
|
||
|
||
|
||
def web_search(query: str) -> str:
|
||
"""Return the title and first paragraph of the first search result.
|
||
|
||
This is a very small wrapper around a Google search. For a production
|
||
system you would use a real search API.
|
||
"""
|
||
# Simple Bing search URL – works without API key for a few requests
|
||
url = f"https://www.bing.com/search?q={requests.utils.quote(query)}"
|
||
resp = requests.get(url, timeout=10)
|
||
resp.raise_for_status()
|
||
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:
|
||
return virtual_fs.read(path)
|
||
except KeyError:
|
||
return f"File {path} does not exist in virtual FS."
|
||
|
||
# 1.4 Dump virtual FS to disk
|
||
|
||
def dump_virtual_fs_tool(output_dir: str = "output") -> str:
|
||
root = Path(output_dir)
|
||
virtual_fs.dump_to_disk(root)
|
||
return f"Virtual FS dumped to {root.resolve()}"
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 2. Agent definition – deep agent style
|
||
# ---------------------------------------------------------------------------
|
||
|
||
# 2.1 LLM
|
||
llm = ChatOllama(model="llama3.1", temperature=0.7)
|
||
|
||
# 2.2 Prompt template – instruct the agent how to use tools
|
||
prompt = ChatPromptTemplate.from_messages([
|
||
("system", "You are a helpful assistant that can search the web, write files, read files and dump virtual files to disk.")
|
||
])
|
||
|
||
# 2.3 Tool mapping
|
||
from langchain.tools import tool
|
||
|
||
# Wrap tools with langchain Tool objects
|
||
from langchain.tools import Tool
|
||
|
||
search_tool = Tool(
|
||
name="WebSearch",
|
||
func=web_search,
|
||
description="Use this to search the web for information. Input should be a natural language query.",
|
||
)
|
||
write_tool = Tool(
|
||
name="WriteFile",
|
||
func=write_file_tool,
|
||
description="Write content to a file in the virtual file system. Input: path and content.",
|
||
)
|
||
read_tool = Tool(
|
||
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
|
||
from langchain.agents import AgentExecutor, ZeroShotAgent
|
||
|
||
# Define the tool names and descriptions for the prompt
|
||
tool_names = [search_tool.name, write_tool.name, read_tool.name, dump_tool.name]
|
||
tool_descriptions = [t.description for t in [search_tool, write_tool, read_tool, dump_tool]]
|
||
|
||
# Build the agent
|
||
agent = ZeroShotAgent.from_llm_and_tools(
|
||
llm=llm,
|
||
tools=[search_tool, write_tool, read_tool, dump_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
|
||
executor = AgentExecutor.from_agent_and_tools(
|
||
agent=agent,
|
||
tools=[search_tool, write_tool, read_tool, dump_tool],
|
||
verbose=True,
|
||
)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 3. Demo / entry point
|
||
# ---------------------------------------------------------------------------
|
||
|
||
if __name__ == "__main__":
|
||
print("Deep Agent Demo – type your question. Type 'exit' to quit.")
|
||
while True:
|
||
user_input = input("> ")
|
||
if user_input.lower() in {"exit", "quit"}:
|
||
print("Exiting…")
|
||
break
|
||
try:
|
||
result = executor.invoke({"input": user_input})
|
||
print("\nResult:\n", result)
|
||
except Exception as e:
|
||
print("Error:", e)
|