Update deep_agent.py
This commit is contained in:
+86
-108
@@ -1,118 +1,96 @@
|
||||
"""Deep Agent implementation based on LangGraph.
|
||||
import json
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
This module defines the core functions required by the tests:
|
||||
- create_agent()
|
||||
- create_agent_executor()
|
||||
from langchain_ollama import ChatOllama
|
||||
|
||||
The implementation is a simplified version of the agent logic from
|
||||
`agent.py`. It is intentionally lightweight so that the tests can
|
||||
import the functions without pulling in heavy dependencies.
|
||||
"""
|
||||
from tools import SearchTool, WriteFileTool
|
||||
from virtual_fs import VirtualFileSystem
|
||||
|
||||
from typing import Dict, List, Any, Optional
|
||||
class DeepAgent:
|
||||
"""A deep‑search agent that can search the web, write virtual files and export them.
|
||||
|
||||
# Minimal imports – the actual heavy libraries are imported lazily
|
||||
# inside the functions to avoid import errors during static analysis.
|
||||
|
||||
# -----------------------------
|
||||
# State definition
|
||||
# -----------------------------
|
||||
class AgentState(dict):
|
||||
"""Simple dict‑based state used by the graph.
|
||||
|
||||
The real implementation uses a more sophisticated ``State`` class
|
||||
from LangGraph, but for the purposes of the unit tests a plain
|
||||
dictionary is sufficient.
|
||||
The agent follows a simple plan: the LLM generates a JSON array of steps, each step
|
||||
is either a ``search`` or a ``write`` action. Search results are stored in the
|
||||
virtual file system so that later steps can refer to them.
|
||||
"""
|
||||
|
||||
def __init__(self, query: str, virtual_files: Dict[str, str] = None, history: List[Dict[str, str]] = None):
|
||||
super().__init__()
|
||||
self.update(
|
||||
query=query,
|
||||
virtual_files=virtual_files or {},
|
||||
history=history or [],
|
||||
answer=None,
|
||||
def __init__(self, llm: ChatOllama):
|
||||
self.llm = llm
|
||||
self.vfs = VirtualFileSystem()
|
||||
self.search_tool = SearchTool()
|
||||
self.write_tool = WriteFileTool(self.vfs)
|
||||
|
||||
def _extract_json(self, text: str) -> str:
|
||||
"""Try to extract a JSON array from an arbitrary string.
|
||||
|
||||
The LLM might prepend or append text. We look for the first ``[`` and the
|
||||
matching ``]`` and slice the string. If that fails we fall back to the
|
||||
original text.
|
||||
"""
|
||||
try:
|
||||
start = text.index("[")
|
||||
end = text.rindex("]") + 1
|
||||
return text[start:end]
|
||||
except ValueError:
|
||||
return text
|
||||
|
||||
def generate_plan(self, instruction: str) -> list:
|
||||
"""Ask the LLM to produce a JSON plan for the given instruction.
|
||||
|
||||
Returns a list of step dictionaries.
|
||||
"""
|
||||
prompt = (
|
||||
"You are a helpful assistant that can search the web and write files. "
|
||||
"Given the instruction below, produce a JSON array of steps. "
|
||||
"Each step is an object with \"action\" (\"search\" or \"write\") and "
|
||||
"\"params\". For \"search\", params: \"query\". For \"write\", params: "
|
||||
"\"filename\" and \"content\". Return only the JSON, no explanation.\n\n"
|
||||
f"Instruction: {instruction}"
|
||||
)
|
||||
response = self.llm.invoke(prompt)
|
||||
json_text = self._extract_json(str(response))
|
||||
try:
|
||||
plan = json.loads(json_text)
|
||||
except json.JSONDecodeError as exc:
|
||||
raise ValueError(f"Failed to parse plan JSON: {exc}\nResponse: {response}")
|
||||
if not isinstance(plan, list):
|
||||
raise ValueError("Plan must be a JSON array of steps.")
|
||||
return plan
|
||||
|
||||
def copy(self, update: Dict[str, Any] = None):
|
||||
new = AgentState(self["query"], self["virtual_files"].copy(), self["history"].copy())
|
||||
if update:
|
||||
new.update(update)
|
||||
return new
|
||||
def execute_plan(self, plan: list):
|
||||
for step in plan:
|
||||
action = step.get("action")
|
||||
params = step.get("params", {})
|
||||
if action == "search":
|
||||
query = params.get("query")
|
||||
if not query:
|
||||
continue
|
||||
result = self.search_tool.run(query)
|
||||
# Sanitize filename: replace spaces with underscores
|
||||
safe_query = re.sub(r"[^a-zA-Z0-9_]+", "_", query)
|
||||
self.vfs.write_file(f"search_{safe_query}.txt", result)
|
||||
elif action == "write":
|
||||
filename = params.get("filename")
|
||||
content = params.get("content")
|
||||
if not filename or content is None:
|
||||
continue
|
||||
self.write_tool.run(filename, content)
|
||||
else:
|
||||
# Unknown action – skip
|
||||
continue
|
||||
|
||||
# -----------------------------
|
||||
# Helper functions (stubs)
|
||||
# -----------------------------
|
||||
def run(self, instruction: str, output_dir: str | Path = "./output"):
|
||||
"""Run the agent for a single instruction.
|
||||
|
||||
def search_web(query: str) -> str:
|
||||
"""Stub that returns a deterministic string.
|
||||
|
||||
The real agent performs an HTTP request, but the tests only
|
||||
require that the function exists.
|
||||
"""
|
||||
return f"Search results for '{query}'"
|
||||
|
||||
# -----------------------------
|
||||
# Tool implementations
|
||||
# -----------------------------
|
||||
|
||||
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})
|
||||
|
||||
|
||||
def search_tool(state: AgentState, query: str) -> AgentState:
|
||||
result = search_web(query)
|
||||
new_history = state["history"].copy()
|
||||
new_history.append({"role": "tool", "name": "search", "content": result})
|
||||
return state.copy(update={"history": new_history})
|
||||
|
||||
# -----------------------------
|
||||
# Agent logic
|
||||
# -----------------------------
|
||||
|
||||
def create_agent():
|
||||
"""Return a very small graph object.
|
||||
|
||||
The real implementation uses LangGraph. For the unit tests we
|
||||
return a simple callable that mimics the interface.
|
||||
"""
|
||||
def agent(state: AgentState):
|
||||
# Very naive decision logic – just return an answer.
|
||||
answer = f"Answer to '{state['query']}'"
|
||||
new_history = state["history"].copy()
|
||||
new_history.append({"role": "assistant", "content": answer})
|
||||
return state.copy(update={"history": new_history, "answer": answer})
|
||||
|
||||
# The returned object must have a ``invoke`` method that accepts a state
|
||||
class GraphStub:
|
||||
def __init__(self, func):
|
||||
self.func = func
|
||||
|
||||
def invoke(self, state):
|
||||
return self.func(state)
|
||||
|
||||
return GraphStub(agent)
|
||||
|
||||
# -----------------------------
|
||||
# Executor helper
|
||||
# -----------------------------
|
||||
|
||||
def create_agent_executor():
|
||||
return create_agent()
|
||||
|
||||
# -----------------------------
|
||||
# If run as script
|
||||
# -----------------------------
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser(description="Run the simplified deep agent.")
|
||||
parser.add_argument("query", type=str, help="User query to process")
|
||||
args = parser.parse_args()
|
||||
|
||||
graph = create_agent_executor()
|
||||
state = AgentState(args.query)
|
||||
final_state = graph.invoke(state)
|
||||
print("Answer:", final_state["answer"])
|
||||
Parameters
|
||||
----------
|
||||
instruction: str
|
||||
The natural‑language instruction for the agent.
|
||||
output_dir: str | Path
|
||||
Directory where the virtual files will be flushed to disk.
|
||||
"""
|
||||
plan = self.generate_plan(instruction)
|
||||
self.execute_plan(plan)
|
||||
# Flush virtual FS to disk
|
||||
self.vfs.flush_to_disk(Path(output_dir))
|
||||
Reference in New Issue
Block a user