96 lines
3.3 KiB
Python
96 lines
3.3 KiB
Python
"""Tool implementations used by the Deep Agent.
|
||
|
||
The tools are simple wrappers around the virtual file system and a
|
||
web‑search helper. They expose a ``name`` and ``description`` that are
|
||
used by the LangChain agent.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from typing import Dict, Any
|
||
|
||
from langchain_core.tools import BaseTool
|
||
from langchain_ollama import OllamaEmbeddings
|
||
from langchain_community.tools.tavily_search import TavilySearchResults
|
||
|
||
from virtual_fs import VirtualFileSystem
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# WebSearch tool
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class WebSearch(BaseTool):
|
||
"""Search the web using Tavily.
|
||
|
||
The tool returns a short string containing the top results.
|
||
"""
|
||
|
||
name: str = "web_search"
|
||
description: str = "Search the web for information. Input should be a query."
|
||
|
||
def _run(self, query: str) -> str:
|
||
tavily = TavilySearchResults(max_results=3)
|
||
results = tavily.run(query)
|
||
return "\n".join(f"{i+1}. {r['title']} – {r['url']}" for i, r in enumerate(results))
|
||
|
||
def _arun(self, query: str) -> str: # pragma: no cover
|
||
return self._run(query)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# CreateVirtualFile tool
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class CreateVirtualFile(BaseTool):
|
||
"""Create or overwrite a virtual file.
|
||
|
||
Input format: ``filename:content``. The tool splits on the first
|
||
colon and stores the content in the shared virtual file system.
|
||
"""
|
||
|
||
name: str = "create_virtual_file"
|
||
description: str = (
|
||
"Create or overwrite a virtual file. Input format: 'filename:content'."
|
||
)
|
||
|
||
def __init__(self, vfs: VirtualFileSystem):
|
||
super().__init__()
|
||
self.vfs = vfs
|
||
|
||
def _run(self, input_text: str) -> str:
|
||
if ":" not in input_text:
|
||
raise ValueError("Input must be in the form 'filename:content'")
|
||
name, content = input_text.split(":", 1)
|
||
self.vfs.create_file(name.strip(), content.strip())
|
||
return f"File '{name.strip()}' created with {len(content.strip())} characters."
|
||
|
||
def _arun(self, input_text: str) -> str: # pragma: no cover
|
||
return self._run(input_text)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# ExportVirtualFiles tool
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class ExportVirtualFiles(BaseTool):
|
||
"""Export all virtual files to the real filesystem.
|
||
|
||
Input can optionally specify a target directory. If omitted, the
|
||
current working directory is used.
|
||
"""
|
||
|
||
name: str = "export_virtual_files"
|
||
description: str = (
|
||
"Export all virtual files to disk. Input is an optional directory path."
|
||
)
|
||
|
||
def __init__(self, vfs: VirtualFileSystem):
|
||
super().__init__()
|
||
self.vfs = vfs
|
||
|
||
def _run(self, input_text: str | None = None) -> str:
|
||
target = input_text.strip() if input_text else "."
|
||
self.vfs.export_to_disk(target)
|
||
return f"Exported {len(self.vfs.files)} files to {target}"
|
||
|
||
def _arun(self, input_text: str | None = None) -> str: # pragma: no cover
|
||
return self._run(input_text)
|