scripts/text_inspector_tool.py

#!/usr/bin/env python
# coding=utf-8
from typing import Optional

from smolagents import Tool
from smolagents.models import MessageRole, Model

from .mdconvert import MarkdownConverter


class TextInspectorTool(Tool):
    """
    Tool for converting various file types to text and answering questions about their contents.

    Supported file types include:
    - Text documents (.txt, .md)
    - Web documents (.html, .htm)
    - Office documents (.docx, .xlsx, .pptx)
    - Audio files (.wav, .mp3, .flac)
    - PDF documents (.pdf)

    Images are not supported and should be processed with a visualizer tool instead.
    """

    name = "view_file"
    description = """
    You cannot load files yourself: instead call this tool to read a file as markdown text and ask questions about it.
    This tool handles the following file extensions: [".html", ".htm", ".md", ".txt", ".xlsx", ".pptx", ".wav", ".mp3", ".flac", ".pdf", ".docx"], and all other types of text files. IT DOES NOT HANDLE IMAGES.
    """

    inputs = {
        "file_path": {
            "description": "The path to the file you want to read as text. Must be a '.something' file, like '.pdf'. If it is an image, use the visualizer tool instead! DO NOT use this tool for an HTML webpage: use the web_search tool instead!",
            "type": "string",
        },
        "question": {
            "description": "[Optional]: Your question, as a natural language sentence. Provide as much context as possible. Do not pass this parameter if you just want to directly return the content of the file.",
            "type": "string",
            "nullable": True,
        },
    }
    output_type = "string"
    md_converter = MarkdownConverter()

    def __init__(self, model: Model, text_limit: int):
        """
        Initialize the TextInspectorTool with a model to use for generating text and a limit for the amount of text to generate.
        """
        super().__init__()
        self.model = model
        self.text_limit = text_limit

    def forward_initial_exam_mode(self, file_path, question):
        """
        This is used for generating code for the initial exam, and is not used for the final exam.
        """
        result = self.md_converter.convert(file_path)

        if file_path[-4:] in [".png", ".jpg", ".webp"]:
            raise Exception(
                "Cannot use inspect_file_as_text tool with images: use visualizer instead!"
            )

        if ".zip" in file_path:
            return result.text_content

        if not question:
            return result.text_content

        if len(result.text_content) < 4000:
            return "Document content: " + result.text_content

        messages = [
            {
                "role": MessageRole.SYSTEM,
                "content": [
                    {
                        "type": "text",
                        "text": "Here is a file:\n### "
                        + str(result.title)
                        + "\n\n"
                        + result.text_content[: self.text_limit],
                    }
                ],
            },
            {
                "role": MessageRole.USER,
                "content": [
                    {
                        "type": "text",
                        "text": "Now please write a short, 5 sentence caption for this document, that could help someone asking this question: "
                        + question
                        + "\n\nDon't answer the question yourself! Just provide useful notes on the document",
                    }
                ],
            },
        ]
        return self.model(messages).content

    def forward(self, file_path: str, question: Optional[str] = None) -> str:
        """
        Process a file and optionally answer a question about its contents.

        Args:
            file_path: Path to the file to be processed. Must be a supported file type.
            question: Optional question to answer about the file contents.
                    If None, returns the raw file content.

        Returns:
            Either the raw file content if no question is provided, or the model's
            response to the question based on the file contents.

        Raises:
            Exception: If the file is an image file or has an unsupported format.
        """
        result = self.md_converter.convert(file_path)

        if file_path[-4:] in [".png", ".jpg"]:
            raise Exception(
                "Cannot use inspect_file_as_text tool with images: use visualizer instead!"
            )

        if ".zip" in file_path:
            return result.text_content

        if not question:
            return result.text_content

        messages = [
            {
                "role": MessageRole.SYSTEM,
                "content": [
                    {
                        "type": "text",
                        "text": "You will have to write a short caption for this file, then answer this question:"
                        + question,
                    }
                ],
            },
            {
                "role": MessageRole.USER,
                "content": [
                    {
                        "type": "text",
                        "text": "Here is the complete file:\n### "
                        + str(result.title)
                        + "\n\n"
                        + result.text_content[: self.text_limit],
                    }
                ],
            },
            {
                "role": MessageRole.USER,
                "content": [
                    {
                        "type": "text",
                        "text": "Now answer the question below. Use these three headings: '1. Short answer', '2. Extremely detailed answer', '3. Additional Context on the document and question asked'."
                        + question,
                    }
                ],
            },
        ]
        return self.model(messages).content


__all__ = ["TextInspectorTool"]