diff --git a/deep_agent.py b/deep_agent.py index 36f99db..8f999d9 100644 --- a/deep_agent.py +++ b/deep_agent.py @@ -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)) \ No newline at end of file