Update agent.py
This commit is contained in:
@@ -1,187 +1,121 @@
|
|||||||
"""Deep Agent implementation based on LangGraph.
|
"""Deep Agent implementation based on LangGraph and LangChain.
|
||||||
|
|
||||||
This agent can:
|
The module exposes two public functions:
|
||||||
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:
|
* ``create_agent`` – returns a :class:`langgraph.graph.StateGraph` that can be
|
||||||
- Uses langchain>=1.0.0 and langgraph>=1.0.0.
|
executed.
|
||||||
- Correct imports for text splitters, Chroma, Ollama embeddings and chat model.
|
* ``create_agent_executor`` – returns an :class:`langchain.agents.AgentExecutor`
|
||||||
- Implements `create_agent` and `create_agent_executor` functions.
|
that can be used directly.
|
||||||
- Uses the `write_file` tool to write virtual files.
|
|
||||||
|
Both functions lazily import the heavy LangChain/LangGraph dependencies so
|
||||||
|
that the module can be imported even if those packages are not installed.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import pathlib
|
from __future__ import annotations
|
||||||
from typing import Dict, List, Any, Optional
|
|
||||||
|
|
||||||
# LangChain imports
|
from typing import Any
|
||||||
from langchain_ollama import ChatOllama, OllamaEmbeddings
|
|
||||||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
|
||||||
from langchain_chroma import Chroma
|
|
||||||
|
|
||||||
# LangGraph imports
|
# Lightweight imports – can be imported eagerly
|
||||||
from langgraph.graph import Graph, State
|
from virtual_fs import VirtualFileSystem
|
||||||
|
from tools import WebSearch, CreateVirtualFile, ExportVirtualFiles
|
||||||
|
|
||||||
# Requests for web search
|
# Shared virtual file system instance used by all tools
|
||||||
import requests
|
vfs = VirtualFileSystem()
|
||||||
|
|
||||||
# -----------------------------
|
# Define the tools
|
||||||
# Configuration
|
web_search_tool = WebSearch()
|
||||||
# -----------------------------
|
create_file_tool = CreateVirtualFile(vfs)
|
||||||
OLLAMA_MODEL = "llama3"
|
export_files_tool = ExportVirtualFiles(vfs)
|
||||||
EMBEDDINGS_MODEL = "llama3"
|
|
||||||
SEARCH_URL = "https://api.duckduckgo.com/"
|
|
||||||
OUTPUT_DIR = pathlib.Path("output_files")
|
|
||||||
OUTPUT_DIR.mkdir(exist_ok=True)
|
|
||||||
|
|
||||||
# -----------------------------
|
# ---------------------------------------------------------------------------
|
||||||
# Helper functions
|
# Public API
|
||||||
# -----------------------------
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
def search_web(query: str) -> str:
|
def create_agent() -> Any:
|
||||||
"""Return a short summary of the search results using DuckDuckGo."""
|
"""Return a LangGraph graph.
|
||||||
params = {
|
|
||||||
"q": query,
|
The function performs a lazy import of :mod:`langgraph` and related
|
||||||
"format": "json",
|
dependencies. If the imports fail, a clear ``ImportError`` is raised.
|
||||||
"no_html": 1,
|
"""
|
||||||
"skip_disambig": 1,
|
|
||||||
}
|
|
||||||
try:
|
try:
|
||||||
r = requests.get(SEARCH_URL, params=params, timeout=10)
|
from langgraph.graph import StateGraph
|
||||||
r.raise_for_status()
|
from langgraph.prebuilt import create_react_agent
|
||||||
data = r.json()
|
from langchain_ollama import ChatOllama
|
||||||
abstract = data.get("AbstractText")
|
from langchain_core.prompts import ChatPromptTemplate
|
||||||
if abstract:
|
from langchain_core.output_parsers import StrOutputParser
|
||||||
return abstract
|
except Exception as exc: # pragma: no cover
|
||||||
topics = data.get("RelatedTopics", [])
|
raise ImportError(
|
||||||
snippets = [t.get("Text", "") for t in topics if "Text" in t]
|
"Failed to import LangGraph/LangChain dependencies. Ensure that the\n"
|
||||||
return "\n".join(snippets[:5])
|
"required packages are installed and the environment is correctly\n"
|
||||||
except Exception as e:
|
"configured."
|
||||||
return f"Error during search: {e}"
|
) from exc
|
||||||
|
|
||||||
# -----------------------------
|
# LLM and prompt
|
||||||
# State definition
|
llm = ChatOllama(model="llama3.1")
|
||||||
# -----------------------------
|
prompt = ChatPromptTemplate.from_messages(
|
||||||
class AgentState(State):
|
[
|
||||||
query: str
|
("system", "You are a helpful assistant that can search the web, create virtual files, and export them to disk."),
|
||||||
virtual_files: Dict[str, str]
|
("human", "{input}"),
|
||||||
history: List[Dict[str, str]]
|
]
|
||||||
answer: Optional[str] = None
|
)
|
||||||
|
parser = StrOutputParser()
|
||||||
|
|
||||||
# -----------------------------
|
# Build a simple react agent using LangGraph prebuilt
|
||||||
# Tool: write_file
|
react_agent = create_react_agent(
|
||||||
# -----------------------------
|
llm=llm,
|
||||||
|
tools=[web_search_tool, create_file_tool, export_files_tool],
|
||||||
def write_file_tool(state: AgentState, file_name: str, content: str) -> AgentState:
|
prompt=prompt,
|
||||||
new_files = state.virtual_files.copy()
|
output_parser=parser,
|
||||||
new_files[file_name] = content
|
)
|
||||||
return state.copy(update={"virtual_files": new_files})
|
|
||||||
|
|
||||||
# -----------------------------
|
|
||||||
# Tool: search
|
|
||||||
# -----------------------------
|
|
||||||
|
|
||||||
def search_tool(state: AgentState, query: str) -> AgentState:
|
|
||||||
result = search_web(query)
|
|
||||||
new_history = state.history + [{"role": "tool", "name": "search", "content": result}]
|
|
||||||
return state.copy(update={"history": new_history})
|
|
||||||
|
|
||||||
# -----------------------------
|
|
||||||
# Agent logic
|
|
||||||
# -----------------------------
|
|
||||||
|
|
||||||
def create_agent() -> Graph:
|
|
||||||
llm = ChatOllama(model=OLLAMA_MODEL, temperature=0.7)
|
|
||||||
embeddings = OllamaEmbeddings(model=EMBEDDINGS_MODEL)
|
|
||||||
splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
|
|
||||||
|
|
||||||
def agent(state: AgentState) -> Dict[str, Any]:
|
|
||||||
# Build prompt from history
|
|
||||||
messages = []
|
|
||||||
for msg in state.history:
|
|
||||||
if msg["role"] == "user":
|
|
||||||
messages.append({"role": "user", "content": msg["content"]})
|
|
||||||
elif msg["role"] == "assistant":
|
|
||||||
messages.append({"role": "assistant", "content": msg["content"]})
|
|
||||||
elif msg["role"] == "tool":
|
|
||||||
messages.append({"role": "assistant", "content": f"[Tool: {msg['name']}] {msg['content']}"})
|
|
||||||
messages.append({"role": "assistant", "content": f"User query: {state.query}"})
|
|
||||||
|
|
||||||
response = llm.invoke(messages)
|
|
||||||
text = response.content.strip()
|
|
||||||
|
|
||||||
if text.upper().startswith("ANSWER:"):
|
|
||||||
answer = text[7:].strip()
|
|
||||||
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}
|
|
||||||
|
|
||||||
graph = Graph()
|
|
||||||
graph.add_node("agent", agent)
|
|
||||||
graph.add_node("search", search_tool)
|
|
||||||
graph.add_node("write_file", write_file_tool)
|
|
||||||
|
|
||||||
def final_answer_node(state: AgentState) -> AgentState:
|
|
||||||
for msg in reversed(state.history):
|
|
||||||
if msg["role"] == "assistant":
|
|
||||||
state = state.copy(update={"answer": msg["content"]})
|
|
||||||
break
|
|
||||||
return state
|
|
||||||
|
|
||||||
graph.add_node("final_answer", final_answer_node)
|
|
||||||
|
|
||||||
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")
|
|
||||||
|
|
||||||
|
# Create the graph
|
||||||
|
graph = StateGraph()
|
||||||
|
graph.add_node("react_agent", react_agent)
|
||||||
|
graph.set_entry_point("react_agent")
|
||||||
|
graph.set_finish_point("react_agent")
|
||||||
return graph
|
return graph
|
||||||
|
|
||||||
# -----------------------------
|
|
||||||
# Executor helper
|
|
||||||
# -----------------------------
|
|
||||||
|
|
||||||
def create_agent_executor() -> Graph:
|
def create_agent_executor() -> Any:
|
||||||
return create_agent()
|
"""Return an :class:`langchain.agents.AgentExecutor`.
|
||||||
|
|
||||||
# -----------------------------
|
The function lazily imports the required classes. It is useful for
|
||||||
# Main execution
|
quick experimentation and for environments where the full graph is
|
||||||
# -----------------------------
|
unnecessary.
|
||||||
if __name__ == "__main__":
|
"""
|
||||||
import argparse
|
try:
|
||||||
parser = argparse.ArgumentParser(description="Run the deep agent.")
|
from langchain.agents import AgentExecutor
|
||||||
parser.add_argument("query", type=str, help="User query to process")
|
from langgraph.prebuilt import create_react_agent
|
||||||
args = parser.parse_args()
|
from langchain_ollama import ChatOllama
|
||||||
|
from langchain_core.prompts import ChatPromptTemplate
|
||||||
|
from langchain_core.output_parsers import StrOutputParser
|
||||||
|
except Exception as exc: # pragma: no cover
|
||||||
|
raise ImportError(
|
||||||
|
"Failed to import LangChain dependencies for the AgentExecutor."
|
||||||
|
) from exc
|
||||||
|
|
||||||
graph = create_agent_executor()
|
llm = ChatOllama(model="llama3.1")
|
||||||
initial_state = AgentState(query=args.query, virtual_files={}, history=[{"role": "user", "content": args.query}])
|
prompt = ChatPromptTemplate.from_messages(
|
||||||
|
[
|
||||||
|
("system", "You are a helpful assistant that can search the web, create virtual files, and export them to disk."),
|
||||||
|
("human", "{input}"),
|
||||||
|
]
|
||||||
|
)
|
||||||
|
parser = StrOutputParser()
|
||||||
|
|
||||||
final_state = graph.invoke(initial_state)
|
# Build the react agent
|
||||||
|
react_agent = create_react_agent(
|
||||||
|
llm=llm,
|
||||||
|
tools=[web_search_tool, create_file_tool, export_files_tool],
|
||||||
|
prompt=prompt,
|
||||||
|
output_parser=parser,
|
||||||
|
)
|
||||||
|
|
||||||
print("\n=== Final Answer ===\n")
|
return AgentExecutor.from_agent_and_tools(
|
||||||
print(final_state.answer if final_state.answer else "No answer produced.")
|
agent=react_agent,
|
||||||
|
tools=[web_search_tool, create_file_tool, export_files_tool],
|
||||||
|
verbose=True,
|
||||||
|
)
|
||||||
|
|
||||||
for fname, content in final_state.virtual_files.items():
|
# Expose public names for tests
|
||||||
out_path = OUTPUT_DIR / fname
|
__all__ = ["create_agent", "create_agent_executor", "vfs"]
|
||||||
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