Create text_inspector_tool.py

text_inspector_tool.py

@@ -0,0 +1,122 @@
+from typing import Optional
+
+from smolagents import Tool
+from smolagents.models import MessageRole, Model
+
+from .mdconvert import MarkdownConverter
+
+
+class TextInspectorTool(Tool):
+    name = "inspect_file_as_text"
+    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", ".xlsx", ".pptx", ".wav", ".mp3", ".m4a", ".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):
+        super().__init__()
+        self.model = model
+        self.text_limit = text_limit
+
+    def forward_initial_exam_mode(self, file_path, question):
+        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
+
+        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, question: Optional[str] = None) -> str:
+        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