Add file inspect service function

src/agentscope/service/inspector/file_inspector.py

@@ -0,0 +1,217 @@
+# -*- coding: utf-8 -*-
+"""
+The ServiceFunction utils for inspecting files.
+Convert different file formats to LLM understandable text.
+"""
+from io import StringIO
+import os
+from typing import Optional, Union, Dict, Callable
+
+import mammoth
+import puremagic
+import markdownify
+from bs4 import BeautifulSoup
+import pandas as pd
+import pdfminer
+import pdfminer.high_level
+from agentscope.service import ServiceResponse, ServiceExecStatus
+
+
+def _guess_ext_magic(file_path: str) -> Union[str, None]:
+ """
+ Use puremagic to guess a file's extension based on the first few bytes.
+ """
+ guesses = puremagic.magic_file(file_path)
+ if len(guesses) > 0:
+ ext = guesses[0].extension.strip()
+ if len(ext) > 0:
+ return ext
+
+ return None
+
+
+# currently supported file types
+# plain_text, docx, xlsx, pdf, html, py, txt
+# TODO support more types: image, audio, etc.
+def inspect_file_as_text(file_path: str) -> ServiceResponse:
+ """Inspect common-types files to markdown style text,
+ so the llm can inspect such files.
+ Currently support '.docx', '.xlsx', '.pdf', '.txt', '.py', '.html', etc.
+ Args:
+ file_path (str):
+ the path of the file to inspect
+ """
+ try:
+ if not os.path.isfile(file_path):
+ raise FileNotFoundError(f"No such file: '{file_path}'")
+
+ _, ext = os.path.splitext(file_path)
+ ext = ext.lower()
+
+ if not ext:
+ ext = _guess_ext_magic(file_path)
+ if ext:
+ ext = "." + ext.lower()
+
+ inspect_functions: Dict[str, Callable[[str], ServiceResponse]] = {
+ ".docx": inspect_docx_as_text,
+ ".xlsx": inspect_xlsx_as_text,
+ ".pdf": inspect_pdf_as_text,
+ ".html": inspect_html_as_text,
+ ".txt": inspect_raw_local_file,
+ ".py": inspect_raw_local_file,
+ }
+
+ if ext in inspect_functions:
+ return inspect_functions[ext](file_path)
+ else:
+ return ServiceResponse(
+ status=ServiceExecStatus.ERROR,
+ content=f"Unsupported file extension: {ext}",
+ )
+ except FileNotFoundError as e:
+ return ServiceResponse(
+ status=ServiceExecStatus.ERROR,
+ content=str(e),
+ )
+ except Exception as e:
+ return ServiceResponse(
+ status=ServiceExecStatus.ERROR,
+ content=f"An error occurred: {str(e)}",
+ )
+
+
+def parse_html_to_markdown(html_content: str) -> ServiceResponse:
+ """Parse the html content to markdown format.
+ Args:
+ html_content (str):
+ the html content to be parsed
+ """
+ try:
+ soup = BeautifulSoup(html_content, "html.parser")
+ # Remove javascript and style blocks that may be too messy
+ for block in soup(["script", "style"]):
+ block.extract()
+ # if there are main content, find the main content only
+ body_elm = soup.find("body")
+ content_text = ""
+ if body_elm:
+ content_text = markdownify.MarkdownConverter().convert_soup(
+ body_elm,
+ )
+ else:
+ content_text = markdownify.MarkdownConverter().convert_soup(soup)
+ content_title = ""
+ if soup.title:
+ content_title = soup.title.string + "\n"
+ return ServiceResponse(
+ status=ServiceExecStatus.SUCCESS,
+ content=content_title + content_text,
+ )
+ except Exception as e:
+ return ServiceResponse(status=ServiceExecStatus.ERROR, content=str(e))
+
+
+def inspect_docx_as_text(file_path: str) -> ServiceResponse:
+ """Inspect the text content in the docx file.
+ Args:
+ file_path (str):
+ the path of the docx file
+ """
+ try:
+ with open(file_path, "rb") as docx_file:
+ result = mammoth.convert_to_markdown(docx_file)
+ markdown = result.value
+ return ServiceResponse(
+ status=ServiceExecStatus.SUCCESS,
+ content=markdown,
+ )
+ except Exception as e:
+ return ServiceResponse(status=ServiceExecStatus.ERROR, content=str(e))
+
+
+def inspect_html_as_text(file_path: str) -> ServiceResponse:
+ """Inspect the text content in the html file.
+ Args:
+ file_path (str):
+ the path of the html file
+ """
+ try:
+ with open(file_path, "rt") as file: # pylint: disable=W1514
+ parsed = parse_html_to_markdown(file.read())
+ return parsed
+ except Exception as e:
+ return ServiceResponse(status=ServiceExecStatus.ERROR, content=str(e))
+
+
+def inspect_xlsx_as_text(file_path: str) -> ServiceResponse:
+ """Inspect the content in the xlsx file.
+ Args:
+ file_path (str):
+ the path of the xlsx file
+ """
+ try:
+ sheets = pd.read_excel(file_path)
+ md_buffer = StringIO()
+ for sheet_name, sheet_data in sheets.items():
+ md_buffer.write(f"## {sheet_name}\n")
+ html_content = sheet_data.to_html(index=False)
+ md_buffer.write(
+ parse_html_to_markdown(html_content).content.strip(),
+ )
+ md_buffer.write("\n\n")
+ md_content = md_buffer.getvalue().strip()
+ return ServiceResponse(
+ status=ServiceExecStatus.SUCCESS,
+ content=md_content,
+ )
+ except FileNotFoundError as e:
+ return ServiceResponse(status=ServiceExecStatus.ERROR, content=str(e))
+
+
+def inspect_pdf_as_text(file_path: str) -> ServiceResponse:
+ """Inspect the text content in the pdf file.
+ Args:
+ file_path (str):
+ the path of the pdf file
+ """
+ # TODO we could consider using pdf conversion repos such as
+ # https://github.com/VikParuchuri/marker for better performance.
+ # However, such module would require
+ # heavy local computations and dependencies.
+ try:
+ pdf_content = pdfminer.high_level.extract_text(file_path)
+ return ServiceResponse(
+ status=ServiceExecStatus.SUCCESS,
+ content=pdf_content,
+ )
+ except Exception as e:
+ return ServiceResponse(status=ServiceExecStatus.ERROR, content=str(e))
+
+
+def inspect_raw_local_file(
+ file_path: str,
+ set_nu: Optional[bool] = True,
+) -> ServiceResponse:
+ """Inspect the content in the local file.
+ Useful for '.py' and '.txt' files.
+
+ Args:
+ file_path (str):
+ the path of the file
+ set_nu (bool, optional):
+ whether to show the line number in given the content.
+ Defaults to True.
+ """
+ try:
+ with open(file_path, "rt") as file: # pylint: disable=W1514
+ content = file.readlines()
+ if set_nu:
+ content = [f"{i+1}: {line}" for i, line in enumerate(content)]
+ text_content = "".join(content)
+ return ServiceResponse(
+ status=ServiceExecStatus.SUCCESS,
+ content=text_content,
+ )
+ except Exception as e:
+ return ServiceResponse(status=ServiceExecStatus.ERROR, content=str(e))