first commit
This commit is contained in:
DigiJ
@@ -0,0 +1,3 @@
|
||||
from ._file_surfer import FileSurfer
|
||||
|
||||
__all__ = ["FileSurfer"]
|
||||
@@ -0,0 +1,208 @@
|
||||
import json
|
||||
import os
|
||||
import traceback
|
||||
from typing import List, Sequence, Tuple
|
||||
|
||||
from agentdhal_agentchat.agents import BaseChatAgent
|
||||
from agentdhal_agentchat.base import Response
|
||||
from agentdhal_agentchat.messages import (
|
||||
BaseChatMessage,
|
||||
TextMessage,
|
||||
)
|
||||
from agentdhal_agentchat.utils import remove_images
|
||||
from agentdhal_core import CancellationToken, Component, ComponentModel, FunctionCall
|
||||
from agentdhal_core.models import (
|
||||
AssistantMessage,
|
||||
ChatCompletionClient,
|
||||
LLMMessage,
|
||||
SystemMessage,
|
||||
UserMessage,
|
||||
)
|
||||
from pydantic import BaseModel
|
||||
from typing_extensions import Self
|
||||
|
||||
from ._markdown_file_browser import MarkdownFileBrowser
|
||||
|
||||
# from typing_extensions import Annotated
|
||||
from ._tool_definitions import (
|
||||
TOOL_FIND_NEXT,
|
||||
TOOL_FIND_ON_PAGE_CTRL_F,
|
||||
TOOL_OPEN_PATH,
|
||||
TOOL_PAGE_DOWN,
|
||||
TOOL_PAGE_UP,
|
||||
)
|
||||
|
||||
|
||||
class FileSurferConfig(BaseModel):
|
||||
"""Configuration for FileSurfer agent"""
|
||||
|
||||
name: str
|
||||
model_client: ComponentModel
|
||||
description: str | None = None
|
||||
|
||||
|
||||
class FileSurfer(BaseChatAgent, Component[FileSurferConfig]):
|
||||
"""An agent, used by MagenticOne, that acts as a local file previewer. FileSurfer can open and read a variety of common file types, and can navigate the local file hierarchy.
|
||||
|
||||
Installation:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
pip install "agentdhal-ext[file-surfer]"
|
||||
|
||||
Args:
|
||||
name (str): The agent's name
|
||||
model_client (ChatCompletionClient): The model to use (must be tool-use enabled)
|
||||
description (str): The agent's description used by the team. Defaults to DEFAULT_DESCRIPTION
|
||||
base_path (str): The base path to use for the file browser. Defaults to the current working directory.
|
||||
|
||||
"""
|
||||
|
||||
component_config_schema = FileSurferConfig
|
||||
component_provider_override = "agentdhal_extensions.agents.file_surfer.FileSurfer"
|
||||
|
||||
DEFAULT_DESCRIPTION = "An agent that can handle local files."
|
||||
|
||||
DEFAULT_SYSTEM_MESSAGES = [
|
||||
SystemMessage(
|
||||
content="""
|
||||
You are a helpful AI Assistant.
|
||||
When given a user query, use available functions to help the user with their request."""
|
||||
),
|
||||
]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
model_client: ChatCompletionClient,
|
||||
description: str = DEFAULT_DESCRIPTION,
|
||||
base_path: str = os.getcwd(),
|
||||
) -> None:
|
||||
super().__init__(name, description)
|
||||
self._model_client = model_client
|
||||
self._chat_history: List[LLMMessage] = []
|
||||
self._browser = MarkdownFileBrowser(viewport_size=1024 * 5, base_path=base_path)
|
||||
|
||||
@property
|
||||
def produced_message_types(self) -> Sequence[type[BaseChatMessage]]:
|
||||
return (TextMessage,)
|
||||
|
||||
async def on_messages(self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken) -> Response:
|
||||
for chat_message in messages:
|
||||
self._chat_history.append(chat_message.to_model_message())
|
||||
try:
|
||||
_, content = await self._generate_reply(cancellation_token=cancellation_token)
|
||||
self._chat_history.append(AssistantMessage(content=content, source=self.name))
|
||||
return Response(chat_message=TextMessage(content=content, source=self.name))
|
||||
|
||||
except BaseException:
|
||||
content = f"File surfing error:\n\n{traceback.format_exc()}"
|
||||
self._chat_history.append(AssistantMessage(content=content, source=self.name))
|
||||
return Response(chat_message=TextMessage(content=content, source=self.name))
|
||||
|
||||
async def on_reset(self, cancellation_token: CancellationToken) -> None:
|
||||
self._chat_history.clear()
|
||||
|
||||
def _get_browser_state(self) -> Tuple[str, str]:
|
||||
"""
|
||||
Get the current state of the browser, including the header and content.
|
||||
"""
|
||||
header = f"Path: {self._browser.path}\n"
|
||||
|
||||
if self._browser.page_title is not None:
|
||||
header += f"Title: {self._browser.page_title}\n"
|
||||
|
||||
current_page = self._browser.viewport_current_page
|
||||
total_pages = len(self._browser.viewport_pages)
|
||||
header += f"Viewport position: Showing page {current_page+1} of {total_pages}.\n"
|
||||
|
||||
return (header, self._browser.viewport)
|
||||
|
||||
async def _generate_reply(self, cancellation_token: CancellationToken) -> Tuple[bool, str]:
|
||||
history = self._chat_history[0:-1]
|
||||
last_message = self._chat_history[-1]
|
||||
assert isinstance(last_message, UserMessage)
|
||||
|
||||
task_content = last_message.content # the last message from the sender is the task
|
||||
|
||||
assert self._browser is not None
|
||||
|
||||
context_message = UserMessage(
|
||||
source="user",
|
||||
content=f"Your file viewer is currently open to the file or directory '{self._browser.page_title}' with path '{self._browser.path}'.",
|
||||
)
|
||||
|
||||
task_message = UserMessage(
|
||||
source="user",
|
||||
content=task_content,
|
||||
)
|
||||
|
||||
create_result = await self._model_client.create(
|
||||
messages=self._get_compatible_context(history + [context_message, task_message]),
|
||||
tools=[
|
||||
TOOL_OPEN_PATH,
|
||||
TOOL_PAGE_DOWN,
|
||||
TOOL_PAGE_UP,
|
||||
TOOL_FIND_NEXT,
|
||||
TOOL_FIND_ON_PAGE_CTRL_F,
|
||||
],
|
||||
cancellation_token=cancellation_token,
|
||||
)
|
||||
|
||||
response = create_result.content
|
||||
|
||||
if isinstance(response, str):
|
||||
# Answer directly.
|
||||
return False, response
|
||||
|
||||
elif isinstance(response, list) and all(isinstance(item, FunctionCall) for item in response):
|
||||
function_calls = response
|
||||
for function_call in function_calls:
|
||||
tool_name = function_call.name
|
||||
|
||||
try:
|
||||
arguments = json.loads(function_call.arguments)
|
||||
except json.JSONDecodeError as e:
|
||||
error_str = f"File surfer encountered an error decoding JSON arguments: {e}"
|
||||
return False, error_str
|
||||
|
||||
if tool_name == "open_path":
|
||||
path = arguments["path"]
|
||||
self._browser.open_path(path)
|
||||
elif tool_name == "page_up":
|
||||
self._browser.page_up()
|
||||
elif tool_name == "page_down":
|
||||
self._browser.page_down()
|
||||
elif tool_name == "find_on_page_ctrl_f":
|
||||
search_string = arguments["search_string"]
|
||||
self._browser.find_on_page(search_string)
|
||||
elif tool_name == "find_next":
|
||||
self._browser.find_next()
|
||||
header, content = self._get_browser_state()
|
||||
final_response = header.strip() + "\n=======================\n" + content
|
||||
return False, final_response
|
||||
|
||||
final_response = "TERMINATE"
|
||||
return False, final_response
|
||||
|
||||
def _get_compatible_context(self, messages: List[LLMMessage]) -> List[LLMMessage]:
|
||||
"""Ensure that the messages are compatible with the underlying client, by removing images if needed."""
|
||||
if self._model_client.model_info["vision"]:
|
||||
return messages
|
||||
else:
|
||||
return remove_images(messages)
|
||||
|
||||
def _to_config(self) -> FileSurferConfig:
|
||||
return FileSurferConfig(
|
||||
name=self.name,
|
||||
model_client=self._model_client.dump_component(),
|
||||
description=self.description,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _from_config(cls, config: FileSurferConfig) -> Self:
|
||||
return cls(
|
||||
name=config.name,
|
||||
model_client=ChatCompletionClient.load_component(config.model_client),
|
||||
description=config.description or cls.DEFAULT_DESCRIPTION,
|
||||
)
|
||||
@@ -0,0 +1,317 @@
|
||||
# ruff: noqa: E722
|
||||
import datetime
|
||||
import io
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from typing import List, Optional, Tuple, Union
|
||||
|
||||
# TODO: Fix unfollowed import
|
||||
from markitdown import FileConversionException, MarkItDown, UnsupportedFormatException # type: ignore
|
||||
|
||||
|
||||
class MarkdownFileBrowser:
|
||||
"""
|
||||
(In preview) An extremely simple Markdown-powered file browser.
|
||||
"""
|
||||
|
||||
# TODO: Fix unfollowed import
|
||||
def __init__( # type: ignore
|
||||
self,
|
||||
viewport_size: Union[int, None] = 1024 * 8,
|
||||
base_path: str | None = os.getcwd(),
|
||||
cwd: str | None = None,
|
||||
):
|
||||
"""
|
||||
Instantiate a new MarkdownFileBrowser.
|
||||
|
||||
Arguments:
|
||||
viewport_size: Approximately how many *characters* fit in the viewport. Viewport dimensions are adjusted dynamically to avoid cutting off words (default: 8192).
|
||||
base_path: The base path to use for the file browser. Files outside this path cannot be accessed. Defaults to the current working directory.
|
||||
cwd: The browser's current working directory. Defaults to the system's current working directory.
|
||||
"""
|
||||
self.viewport_size = viewport_size # Applies only to the standard uri types
|
||||
self.history: List[Tuple[str, float]] = list()
|
||||
self.page_title: Optional[str] = None
|
||||
self.viewport_current_page = 0
|
||||
self.viewport_pages: List[Tuple[int, int]] = list()
|
||||
self._markdown_converter = MarkItDown()
|
||||
self._base_path = None if base_path is None else os.path.realpath(base_path)
|
||||
self._page_content: str = ""
|
||||
self._find_on_page_query: Union[str, None] = None
|
||||
self._find_on_page_last_result: Union[int, None] = None # Location of the last result
|
||||
|
||||
# Set the working directory
|
||||
if cwd is None:
|
||||
if self._validate_path(os.getcwd()):
|
||||
# Use the current working directory if it's in the base path
|
||||
cwd = os.path.realpath(os.getcwd())
|
||||
elif self._base_path is not None:
|
||||
# Otherwise, use the base path
|
||||
cwd = os.path.realpath(self._base_path)
|
||||
else:
|
||||
raise ValueError("No valid working directory (cwd) provided.")
|
||||
elif not self._validate_path(cwd):
|
||||
# A cwd was provided, but it is not valid
|
||||
raise ValueError(f"Working directory (cwd) '{cwd}' is not valid. It must be within the base path.")
|
||||
|
||||
# Populate the history with the current working directory
|
||||
self.set_path(os.path.realpath(cwd))
|
||||
|
||||
@property
|
||||
def path(self) -> str:
|
||||
"""Return the path of the current page."""
|
||||
assert len(self.history) > 0
|
||||
return self.history[-1][0]
|
||||
|
||||
def _validate_path(self, path: str) -> bool:
|
||||
"""Validates the path to ensure it is within the base path.
|
||||
|
||||
Arguments:
|
||||
path: The path to validate.
|
||||
Returns:
|
||||
True if the path is valid, False otherwise.
|
||||
"""
|
||||
if self._base_path is None:
|
||||
return True
|
||||
|
||||
# Normalize the paths
|
||||
path = os.path.realpath(path)
|
||||
base = os.path.realpath(self._base_path)
|
||||
|
||||
# Check if the path is within the base path
|
||||
if os.path.commonpath([path, base]) != base:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def set_path(self, path: str) -> None:
|
||||
"""Sets the path of the current page.
|
||||
This will result in the file being opened for reading.
|
||||
|
||||
Arguments:
|
||||
path: An absolute or relative path of the file or directory to open."
|
||||
"""
|
||||
|
||||
# Handle relative paths
|
||||
path = os.path.expanduser(path)
|
||||
if not os.path.isabs(path):
|
||||
if os.path.isfile(self.path):
|
||||
path = os.path.abspath(os.path.join(os.path.dirname(self.path), path))
|
||||
elif os.path.isdir(self.path):
|
||||
path = os.path.abspath(os.path.join(self.path, path))
|
||||
# If neither a file or a directory, take it verbatim
|
||||
|
||||
# Validating the path wrt. the base path is done in _open_path
|
||||
path = os.path.realpath(path)
|
||||
|
||||
self.history.append((path, time.time()))
|
||||
self._open_path(path)
|
||||
self.viewport_current_page = 0
|
||||
self.find_on_page_query = None
|
||||
self.find_on_page_viewport = None
|
||||
|
||||
@property
|
||||
def viewport(self) -> str:
|
||||
"""Return the content of the current viewport."""
|
||||
bounds = self.viewport_pages[self.viewport_current_page]
|
||||
return self.page_content[bounds[0] : bounds[1]]
|
||||
|
||||
@property
|
||||
def page_content(self) -> str:
|
||||
"""Return the full contents of the current page."""
|
||||
return self._page_content
|
||||
|
||||
def _set_page_content(self, content: str, split_pages: bool = True) -> None:
|
||||
"""Sets the text content of the current page."""
|
||||
self._page_content = content
|
||||
|
||||
if split_pages:
|
||||
self._split_pages()
|
||||
else:
|
||||
self.viewport_pages = [(0, len(self._page_content))]
|
||||
|
||||
if self.viewport_current_page >= len(self.viewport_pages):
|
||||
self.viewport_current_page = len(self.viewport_pages) - 1
|
||||
|
||||
def page_down(self) -> None:
|
||||
"""Move the viewport down one page, if possible."""
|
||||
self.viewport_current_page = min(self.viewport_current_page + 1, len(self.viewport_pages) - 1)
|
||||
|
||||
def page_up(self) -> None:
|
||||
"""Move the viewport up one page, if possible."""
|
||||
self.viewport_current_page = max(self.viewport_current_page - 1, 0)
|
||||
|
||||
def find_on_page(self, query: str) -> Union[str, None]:
|
||||
"""Searches for the query from the current viewport forward, looping back to the start if necessary."""
|
||||
|
||||
# Did we get here via a previous find_on_page search with the same query?
|
||||
# If so, map to find_next
|
||||
if query == self._find_on_page_query and self.viewport_current_page == self._find_on_page_last_result:
|
||||
return self.find_next()
|
||||
|
||||
# Ok it's a new search start from the current viewport
|
||||
self._find_on_page_query = query
|
||||
viewport_match = self._find_next_viewport(query, self.viewport_current_page)
|
||||
if viewport_match is None:
|
||||
self._find_on_page_last_result = None
|
||||
return None
|
||||
else:
|
||||
self.viewport_current_page = viewport_match
|
||||
self._find_on_page_last_result = viewport_match
|
||||
return self.viewport
|
||||
|
||||
def find_next(self) -> Union[str, None]:
|
||||
"""Scroll to the next viewport that matches the query"""
|
||||
|
||||
if self._find_on_page_query is None:
|
||||
return None
|
||||
|
||||
starting_viewport = self._find_on_page_last_result
|
||||
if starting_viewport is None:
|
||||
starting_viewport = 0
|
||||
else:
|
||||
starting_viewport += 1
|
||||
if starting_viewport >= len(self.viewport_pages):
|
||||
starting_viewport = 0
|
||||
|
||||
viewport_match = self._find_next_viewport(self._find_on_page_query, starting_viewport)
|
||||
if viewport_match is None:
|
||||
self._find_on_page_last_result = None
|
||||
return None
|
||||
else:
|
||||
self.viewport_current_page = viewport_match
|
||||
self._find_on_page_last_result = viewport_match
|
||||
return self.viewport
|
||||
|
||||
def _find_next_viewport(self, query: Optional[str], starting_viewport: int) -> Union[int, None]:
|
||||
"""Search for matches between the starting viewport looping when reaching the end."""
|
||||
|
||||
if query is None:
|
||||
return None
|
||||
|
||||
# Normalize the query, and convert to a regular expression
|
||||
nquery = re.sub(r"\*", "__STAR__", query)
|
||||
nquery = " " + (" ".join(re.split(r"\W+", nquery))).strip() + " "
|
||||
nquery = nquery.replace(" __STAR__ ", "__STAR__ ") # Merge isolated stars with prior word
|
||||
nquery = nquery.replace("__STAR__", ".*").lower()
|
||||
|
||||
if nquery.strip() == "":
|
||||
return None
|
||||
|
||||
idxs: List[int] = list()
|
||||
idxs.extend(range(starting_viewport, len(self.viewport_pages)))
|
||||
idxs.extend(range(0, starting_viewport))
|
||||
|
||||
for i in idxs:
|
||||
bounds = self.viewport_pages[i]
|
||||
content = self.page_content[bounds[0] : bounds[1]]
|
||||
|
||||
# TODO: Remove markdown links and images
|
||||
ncontent = " " + (" ".join(re.split(r"\W+", content))).strip().lower() + " "
|
||||
if re.search(nquery, ncontent):
|
||||
return i
|
||||
|
||||
return None
|
||||
|
||||
def open_path(self, path: str) -> str:
|
||||
"""Open a file or directory in the file surfer."""
|
||||
self.set_path(path)
|
||||
return self.viewport
|
||||
|
||||
def _split_pages(self) -> None:
|
||||
"""Split the page contents into pages that are approximately the viewport size. Small deviations are permitted to ensure words are not broken."""
|
||||
# Handle empty pages
|
||||
if len(self._page_content) == 0:
|
||||
self.viewport_pages = [(0, 0)]
|
||||
return
|
||||
|
||||
# Break the viewport into pages
|
||||
self.viewport_pages = []
|
||||
start_idx = 0
|
||||
while start_idx < len(self._page_content):
|
||||
end_idx = min(start_idx + self.viewport_size, len(self._page_content)) # type: ignore[operator]
|
||||
# Adjust to end on a space
|
||||
while end_idx < len(self._page_content) and self._page_content[end_idx - 1] not in [" ", "\t", "\r", "\n"]:
|
||||
end_idx += 1
|
||||
self.viewport_pages.append((start_idx, end_idx))
|
||||
start_idx = end_idx
|
||||
|
||||
def _open_path(
|
||||
self,
|
||||
path: str,
|
||||
) -> None:
|
||||
"""Open a file for reading, converting it to Markdown in the process.
|
||||
|
||||
Arguments:
|
||||
path: The path of the file or directory to open.
|
||||
"""
|
||||
|
||||
if not self._validate_path(path):
|
||||
# Not robust to TOCTOU issues.
|
||||
# Mitigate by running with limited permissions, or use a sandbox.
|
||||
self.page_title = "FileNotFoundError"
|
||||
self._set_page_content(f"# FileNotFoundError\n\nFile not found: {path}")
|
||||
else:
|
||||
try:
|
||||
if os.path.isdir(path): # TODO: Fix markdown_converter types
|
||||
res = self._markdown_converter.convert_stream( # type: ignore
|
||||
io.BytesIO(self._fetch_local_dir(path).encode("utf-8")), file_extension=".txt"
|
||||
)
|
||||
assert self._validate_path(path)
|
||||
self.page_title = res.title
|
||||
self._set_page_content(res.text_content, split_pages=False)
|
||||
else:
|
||||
res = self._markdown_converter.convert_local(path)
|
||||
assert self._validate_path(path)
|
||||
self.page_title = res.title
|
||||
self._set_page_content(res.text_content)
|
||||
except UnsupportedFormatException:
|
||||
self.page_title = "UnsupportedFormatException"
|
||||
self._set_page_content(f"# UnsupportedFormatException\n\nCannot preview '{path}' as Markdown.")
|
||||
except FileConversionException:
|
||||
self.page_title = "FileConversionException."
|
||||
self._set_page_content(f"# FileConversionException\n\nError converting '{path}' to Markdown.")
|
||||
except FileNotFoundError:
|
||||
self.page_title = "FileNotFoundError"
|
||||
self._set_page_content(f"# FileNotFoundError\n\nFile not found: {path}")
|
||||
|
||||
def _fetch_local_dir(self, local_path: str) -> str:
|
||||
"""Render a local directory listing in HTML to assist with local file browsing via the "file://" protocol.
|
||||
Through rendered in HTML, later parts of the pipeline will convert the listing to Markdown.
|
||||
|
||||
Arguments:
|
||||
local_path: A path to the local directory whose contents are to be listed.
|
||||
|
||||
Returns:
|
||||
A directory listing, rendered in HTML.
|
||||
"""
|
||||
listing = f"""
|
||||
# Index of {local_path}
|
||||
|
||||
| Name | Size | Date Modified |
|
||||
| ---- | ---- | ------------- |
|
||||
| .. (parent directory) | | |
|
||||
"""
|
||||
for entry in os.listdir(local_path):
|
||||
size = ""
|
||||
full_path = os.path.join(local_path, entry)
|
||||
|
||||
mtime = ""
|
||||
try:
|
||||
mtime = datetime.datetime.fromtimestamp(os.path.getmtime(full_path)).strftime("%Y-%m-%d %H:%M")
|
||||
except Exception as e:
|
||||
# Handles PermissionError, etc.
|
||||
mtime = f"N/A: {type(e).__name__}"
|
||||
|
||||
if os.path.isdir(full_path):
|
||||
entry = entry + os.path.sep
|
||||
else:
|
||||
try:
|
||||
size = str(os.path.getsize(full_path))
|
||||
except Exception as e:
|
||||
# Handles PermissionError, etc.
|
||||
size = f"N/A: {type(e).__name__}"
|
||||
|
||||
listing += f"| {entry} | {size} | {mtime} |\n"
|
||||
return listing
|
||||
@@ -0,0 +1,50 @@
|
||||
from agentdhal_core.tools import ParametersSchema, ToolSchema
|
||||
|
||||
TOOL_OPEN_PATH = ToolSchema(
|
||||
name="open_path",
|
||||
description="Open a local file or directory at a path in the text-based file browser and return current viewport content.",
|
||||
parameters=ParametersSchema(
|
||||
type="object",
|
||||
properties={
|
||||
"path": {
|
||||
"type": "string",
|
||||
"description": "The relative or absolute path of a local file to visit.",
|
||||
},
|
||||
},
|
||||
required=["path"],
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
TOOL_PAGE_UP = ToolSchema(
|
||||
name="page_up",
|
||||
description="Scroll the viewport UP one page-length in the current file and return the new viewport content.",
|
||||
)
|
||||
|
||||
|
||||
TOOL_PAGE_DOWN = ToolSchema(
|
||||
name="page_down",
|
||||
description="Scroll the viewport DOWN one page-length in the current file and return the new viewport content.",
|
||||
)
|
||||
|
||||
|
||||
TOOL_FIND_ON_PAGE_CTRL_F = ToolSchema(
|
||||
name="find_on_page_ctrl_f",
|
||||
description="Scroll the viewport to the first occurrence of the search string. This is equivalent to Ctrl+F.",
|
||||
parameters=ParametersSchema(
|
||||
type="object",
|
||||
properties={
|
||||
"search_string": {
|
||||
"type": "string",
|
||||
"description": "The string to search for on the page. This search string supports wildcards like '*'",
|
||||
},
|
||||
},
|
||||
required=["search_string"],
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
TOOL_FIND_NEXT = ToolSchema(
|
||||
name="find_next",
|
||||
description="Scroll the viewport to next occurrence of the search string.",
|
||||
)
|
||||
Reference in New Issue
Block a user