96 lines
3.6 KiB
Python
96 lines
3.6 KiB
Python
import json
|
||
import re
|
||
from pathlib import Path
|
||
|
||
from langchain_ollama import ChatOllama
|
||
|
||
from tools import SearchTool, WriteFileTool
|
||
from virtual_fs import VirtualFileSystem
|
||
|
||
class DeepAgent:
|
||
"""A deep‑search agent that can search the web, write virtual files and export them.
|
||
|
||
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, 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 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
|
||
|
||
def run(self, instruction: str, output_dir: str | Path = "./output"):
|
||
"""Run the agent for a single instruction.
|
||
|
||
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)) |