188 lines
6.5 KiB
Python
188 lines
6.5 KiB
Python
"""Deep Agent implementation based on LangGraph.
|
|
|
|
This agent can:
|
|
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.
|
|
"""
|
|
|
|
import pathlib
|
|
from typing import Dict, List, Any, Optional
|
|
|
|
# LangChain imports
|
|
from langchain_ollama import ChatOllama, OllamaEmbeddings
|
|
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
|
from langchain_chroma import Chroma
|
|
|
|
# LangGraph imports
|
|
from langgraph.graph import Graph, State
|
|
|
|
# Requests for web search
|
|
import requests
|
|
|
|
# -----------------------------
|
|
# 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)
|
|
|
|
# -----------------------------
|
|
# Helper functions
|
|
# -----------------------------
|
|
|
|
def search_web(query: str) -> str:
|
|
"""Return a short summary of the search results using DuckDuckGo."""
|
|
params = {
|
|
"q": query,
|
|
"format": "json",
|
|
"no_html": 1,
|
|
"skip_disambig": 1,
|
|
}
|
|
try:
|
|
r = requests.get(SEARCH_URL, params=params, timeout=10)
|
|
r.raise_for_status()
|
|
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}"
|
|
|
|
# -----------------------------
|
|
# State definition
|
|
# -----------------------------
|
|
class AgentState(State):
|
|
query: str
|
|
virtual_files: Dict[str, str]
|
|
history: List[Dict[str, str]]
|
|
answer: Optional[str] = None
|
|
|
|
# -----------------------------
|
|
# Tool: write_file
|
|
# -----------------------------
|
|
|
|
def write_file_tool(state: AgentState, file_name: str, content: str) -> AgentState:
|
|
new_files = state.virtual_files.copy()
|
|
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")
|
|
|
|
return graph
|
|
|
|
# -----------------------------
|
|
# Executor helper
|
|
# -----------------------------
|
|
|
|
def create_agent_executor() -> Graph:
|
|
return create_agent()
|
|
|
|
# -----------------------------
|
|
# Main execution
|
|
# -----------------------------
|
|
if __name__ == "__main__":
|
|
import argparse
|
|
parser = argparse.ArgumentParser(description="Run the deep agent.")
|
|
parser.add_argument("query", type=str, help="User query to process")
|
|
args = parser.parse_args()
|
|
|
|
graph = create_agent_executor()
|
|
initial_state = AgentState(query=args.query, virtual_files={}, history=[{"role": "user", "content": args.query}])
|
|
|
|
final_state = graph.invoke(initial_state)
|
|
|
|
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}")
|