Update tools.py
This commit is contained in:
@@ -1,84 +1,95 @@
|
||||
# tools.py
|
||||
import os
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import List
|
||||
from langchain_core.messages import ToolMessage
|
||||
"""Tool implementations used by the Deep Agent.
|
||||
|
||||
# Global in-memory file store
|
||||
_virtual_files = {}
|
||||
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.
|
||||
"""
|
||||
|
||||
# Define the search tool (placeholder implementation)
|
||||
# In a real scenario, this would query an external search API
|
||||
# For demonstration, we return a static result
|
||||
from __future__ import annotations
|
||||
|
||||
def search(query: str) -> str:
|
||||
"""Simulate a web search and return a short result string."""
|
||||
# For simplicity, just return a canned response
|
||||
return f"Search result for '{query}': This is a placeholder summary."
|
||||
from typing import Dict, Any
|
||||
|
||||
# Define the write_file tool
|
||||
from langchain_core.tools import BaseTool
|
||||
from langchain_ollama import OllamaEmbeddings
|
||||
from langchain_community.tools.tavily_search import TavilySearchResults
|
||||
|
||||
def write_file(file_path: str, content: str) -> str:
|
||||
"""Write content to an in-memory virtual file.
|
||||
from virtual_fs import VirtualFileSystem
|
||||
|
||||
Parameters
|
||||
----------
|
||||
file_path: str
|
||||
Path relative to the workspace root.
|
||||
content: str
|
||||
Text content to write.
|
||||
# ---------------------------------------------------------------------------
|
||||
# WebSearch tool
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
Returns
|
||||
-------
|
||||
str
|
||||
Confirmation message.
|
||||
class WebSearch(BaseTool):
|
||||
"""Search the web using Tavily.
|
||||
|
||||
The tool returns a short string containing the top results.
|
||||
"""
|
||||
_virtual_files[file_path] = content
|
||||
return f"File '{file_path}' written successfully."
|
||||
|
||||
# Define the export_files tool
|
||||
name: str = "web_search"
|
||||
description: str = "Search the web for information. Input should be a query."
|
||||
|
||||
def export_files(export_dir: str) -> str:
|
||||
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.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
export_dir: str
|
||||
Directory path where files should be created.
|
||||
|
||||
Returns
|
||||
-------
|
||||
str
|
||||
Summary of exported files.
|
||||
Input can optionally specify a target directory. If omitted, the
|
||||
current working directory is used.
|
||||
"""
|
||||
Path(export_dir).mkdir(parents=True, exist_ok=True)
|
||||
exported = []
|
||||
for file_path, content in _virtual_files.items():
|
||||
full_path = Path(export_dir) / file_path
|
||||
full_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
full_path.write_text(content, encoding="utf-8")
|
||||
exported.append(str(full_path))
|
||||
return f"Exported {len(exported)} files: {', '.join(exported)}"
|
||||
|
||||
# Expose tool metadata for LangGraph
|
||||
search_tool = {
|
||||
"name": "search",
|
||||
"description": "Search the web for a query and return a short summary.",
|
||||
"func": search,
|
||||
}
|
||||
name: str = "export_virtual_files"
|
||||
description: str = (
|
||||
"Export all virtual files to disk. Input is an optional directory path."
|
||||
)
|
||||
|
||||
write_file_tool = {
|
||||
"name": "write_file",
|
||||
"description": "Write content to a virtual file.",
|
||||
"func": write_file,
|
||||
}
|
||||
def __init__(self, vfs: VirtualFileSystem):
|
||||
super().__init__()
|
||||
self.vfs = vfs
|
||||
|
||||
export_files_tool = {
|
||||
"name": "export_files",
|
||||
"description": "Export all virtual files to the real filesystem.",
|
||||
"func": export_files,
|
||||
}
|
||||
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}"
|
||||
|
||||
# List of tools for LangGraph
|
||||
TOOLS = [search_tool, write_file_tool, export_files_tool]
|
||||
def _arun(self, input_text: str | None = None) -> str: # pragma: no cover
|
||||
return self._run(input_text)
|
||||
|
||||
Reference in New Issue
Block a user