Create frontmatter_tool.py

scripts/frontmatter_tool.py

@@ -0,0 +1,448 @@
+"""
+Frontmatter Generator Tool for Smolagents
+
+This tool helps generate consistent YAML frontmatter for documents,
+useful for RAG systems, static site generators, and document organization.
+"""
+
+import re
+import yaml
+from datetime import datetime
+from typing import Dict, List, Optional, Any, Union
+from smolagents import Tool
+
+
+class FrontmatterGeneratorTool(Tool):
+    """Tool for generating and manipulating YAML frontmatter in documents."""
+
+    name = "frontmatter_generator"
+    description = """
+    Generates or extracts YAML frontmatter for documents. Frontmatter provides structured 
+    metadata for documents including title, author, date, description, and tags.
+    Useful for document organization, RAG systems, and static site generators.
+    """
+
+    inputs = {
+        "content": {
+            "type": "string",
+            "description": "Document content (with or without existing frontmatter)",
+        },
+        "title": {"type": "string", "description": "Document title", "optional": True},
+        "author": {
+            "type": "string",
+            "description": "Document author(s)",
+            "optional": True,
+        },
+        "date": {
+            "type": "string",
+            "description": "Document date in YYYY-MM-DD format (defaults to today if not provided)",
+            "optional": True,
+        },
+        "date_format": {
+            "type": "string",
+            "description": "Format string for the document date (e.g., '%Y-%m-%d', '%d/%m/%Y'). Defaults to '%Y-%m-%d'",
+            "optional": True,
+            "default": "%Y-%m-%d",
+        },
+        "description": {
+            "type": "string",
+            "description": "Brief description of the document",
+            "optional": True,
+        },
+        "tags": {
+            "type": "string",
+            "description": "Comma-separated list of tags",
+            "optional": True,
+        },
+        "additional_fields": {
+            "type": "string",
+            "description": "JSON string with additional frontmatter fields",
+            "optional": True,
+        },
+        "mode": {
+            "type": "string",
+            "description": "Operation mode: 'generate' (create new), 'extract' (get existing), 'update' (modify existing), or 'strip' (remove)",
+            "default": "generate",
+        },
+    }
+    output_type = "string"
+
+    # Regular expression to detect and extract YAML frontmatter
+    FRONTMATTER_PATTERN = r"^---\s*\n(.*?)\n---\s*\n"
+
+    def forward(
+        self,
+        content: str,
+        title: Optional[str] = None,
+        author: Optional[str] = None,
+        date: Optional[str] = None,
+        date_format: Optional[str] = "%Y-%m-%d",
+        description: Optional[str] = None,
+        tags: Optional[str] = None,
+        additional_fields: Optional[str] = None,
+        mode: str = "generate",
+    ) -> str:
+        """
+        Process document content based on specified mode.
+
+        Args:
+            content: Document content with or without frontmatter
+            title: Document title
+            author: Document author(s)
+            date: Document date (YYYY-MM-DD)
+            date_format: strftime format string
+            description: Brief document description
+            tags: Comma-separated list of tags
+            additional_fields: JSON string with additional fields
+            mode: Operation mode (generate, extract, update, strip)
+
+        Returns:
+            Processed document or extracted frontmatter
+        """
+        # Validate inputs
+        if not isinstance(content, str):
+            return "Error: Content must be a string"
+        if title and not isinstance(title, str):
+            return "Error: Title must be a string"
+        if author and not isinstance(author, str):
+            return "Error: Author must be a string"
+        if date and not isinstance(date, str):
+            return "Error: Date must be a string"
+        if description and not isinstance(description, str):
+            return "Error: Description must be a string"
+        if tags and not isinstance(tags, str):
+            return "Error: Tags must be a string"
+        if additional_fields and not isinstance(additional_fields, str):
+            return "Error: Additional_fields must be a string"
+        if not isinstance(mode, str):
+            return "Error: Mode must be a string"
+
+        # Validate mode
+        valid_modes = ["generate", "extract", "update", "strip"]
+        if mode not in valid_modes:
+            return f"Error: Invalid mode '{mode}'. Valid options are: {', '.join(valid_modes)}"
+
+        # Handle empty content
+        if not content or not content.strip():
+            if mode == "generate":
+                # We can still generate frontmatter from provided fields
+                content = ""
+            else:
+                return "Error: Empty content provided"
+
+        # Process based on mode
+        try:
+            if mode == "extract":
+                return self._extract_frontmatter(content)
+            elif mode == "strip":
+                return self._strip_frontmatter(content)
+            elif mode == "update":
+                return self._update_frontmatter(
+                    content,
+                    title,
+                    author,
+                    date,
+                    description,
+                    tags,
+                    additional_fields,
+                    date_format,
+                )
+            else:  # generate
+                return self._generate_frontmatter(
+                    content,
+                    title,
+                    author,
+                    date,
+                    description,
+                    tags,
+                    additional_fields,
+                    date_format,
+                )
+        except Exception as e:
+            return f"Error processing frontmatter: {str(e)}"
+
+    def _extract_frontmatter(self, content: str) -> str:
+        """Extract and return existing frontmatter as formatted YAML."""
+        match = re.search(self.FRONTMATTER_PATTERN, content, re.DOTALL)
+        if not match:
+            return "No frontmatter found in the document"
+
+        try:
+            yaml_content = match.group(1)
+            # Parse and reformat for consistency
+            frontmatter_dict = yaml.safe_load(yaml_content)
+            return f"Extracted frontmatter:\n\n```yaml\n{yaml.dump(frontmatter_dict, sort_keys=False, default_flow_style=False)}```"
+        except yaml.YAMLError:
+            return "Found frontmatter but failed to parse it as valid YAML"
+
+    def _strip_frontmatter(self, content: str) -> str:
+        """Remove frontmatter from document and return clean content."""
+        result = re.sub(self.FRONTMATTER_PATTERN, "", content, count=1, flags=re.DOTALL)
+
+        # Check if anything was actually removed
+        if result == content:
+            return "No frontmatter found to strip. Content unchanged."
+
+        return result.strip()
+
+    def _parse_additional_fields(self, additional_fields: str) -> Dict[str, Any]:
+        """Parse the additional_fields JSON string into a dictionary."""
+        if not additional_fields:
+            return {}
+
+        try:
+            import json
+
+            return json.loads(additional_fields)
+        except json.JSONDecodeError:
+            raise ValueError("additional_fields must be a valid JSON string")
+
+    def _infer_title_from_content(self, content: str) -> Optional[str]:
+        """Attempt to infer document title from content."""
+        # Try to find the first heading
+        heading_match = re.search(r"^#\s+(.+)$", content, re.MULTILINE)
+        if heading_match:
+            return heading_match.group(1).strip()
+
+        # Try to find the first non-empty line
+        lines = content.split("\n")
+        for line in lines:
+            if line.strip():
+                # Limit to a reasonable title length
+                return line.strip()[:100]
+
+        return None
+
+    def _parse_tags(self, tags_string: str) -> List[str]:
+        """Parse comma-separated tags into a list."""
+        if not tags_string:
+            return []
+
+        # Split by comma and clean each tag
+        tag_list = [tag.strip() for tag in tags_string.split(",")]
+        # Remove any empty tags
+        return [tag for tag in tag_list if tag]
+
+    def _parse_flexible_date(
+        self, date_str: str, date_format: Optional[str] = None
+    ) -> str:
+        """
+        Try to parse dates in various formats and convert to YYYY-MM-DD.
+
+        Args:
+            date_str: The date string to parse
+            date_format: Optional preferred format to try first
+
+        Returns:
+            Formatted date as string (YYYY-MM-DD by default)
+        """
+        if not date_str:
+            return datetime.now().strftime("%Y-%m-%d")
+
+        # If a specific format is provided, try it first
+        if date_format:
+            try:
+                parsed_date = datetime.strptime(date_str, date_format)
+                return parsed_date.strftime("%Y-%m-%d")
+            except ValueError:
+                # If it fails, continue with other formats
+                pass
+
+        # Common formats to try
+        formats = [
+            "%Y-%m-%d",  # 2013-03-13
+            "%d %B %Y",  # 13 March 2013
+            "%B %Y",  # September 2013
+            "%Y",  # 1958
+            "%d/%m/%Y",  # 13/03/2013
+            "%m/%d/%Y",  # 03/13/2013
+            "%d-%m-%Y",  # 13-03-2013
+            "%m-%d-%Y",  # 03-13-2013
+            "%Y/%m/%d",  # 2013/03/13
+        ]
+
+        for fmt in formats:
+            try:
+                parsed_date = datetime.strptime(date_str, fmt)
+                return parsed_date.strftime("%Y-%m-%d")
+            except ValueError:
+                continue
+
+        # If no format matched, return the original string
+        return date_str
+
+    def _update_frontmatter(
+        self,
+        content: str,
+        title: Optional[str] = None,
+        author: Optional[str] = None,
+        date: Optional[str] = None,
+        description: Optional[str] = None,
+        tags: Optional[str] = None,
+        additional_fields: Optional[str] = None,
+        date_format: Optional[str] = None,
+    ) -> str:
+        """Update existing frontmatter with new values."""
+        # Check if frontmatter exists
+        match = re.search(self.FRONTMATTER_PATTERN, content, re.DOTALL)
+        if not match:
+            # If no frontmatter exists, generate new one
+            return self._generate_frontmatter(
+                content,
+                title,
+                author,
+                date,
+                description,
+                tags,
+                additional_fields,
+                date_format,
+            )
+
+        # Parse existing frontmatter
+        yaml_content = match.group(1)
+        try:
+            frontmatter_dict = yaml.safe_load(yaml_content) or {}
+        except yaml.YAMLError:
+            frontmatter_dict = {}
+
+        # Update with new values if provided
+        if title:
+            frontmatter_dict["title"] = title
+        if author:
+            frontmatter_dict["author"] = author
+        if date:
+            # Try to parse the date with the flexible parser
+            frontmatter_dict["date"] = self._parse_flexible_date(date, date_format)
+        if description:
+            frontmatter_dict["description"] = description
+        if tags:
+            frontmatter_dict["tags"] = self._parse_tags(tags)
+
+        # Add additional fields
+        if additional_fields:
+            additional_dict = self._parse_additional_fields(additional_fields)
+            frontmatter_dict.update(additional_dict)
+
+        # Generate new frontmatter
+        new_frontmatter = yaml.dump(
+            frontmatter_dict, sort_keys=False, default_flow_style=False
+        )
+        new_frontmatter = f"---\n{new_frontmatter}---\n\n"
+
+        # Replace old frontmatter with new one
+        return re.sub(
+            self.FRONTMATTER_PATTERN, new_frontmatter, content, count=1, flags=re.DOTALL
+        )
+
+    def _generate_frontmatter(
+        self,
+        content: str,
+        title: Optional[str] = None,
+        author: Optional[str] = None,
+        date: Optional[str] = None,
+        description: Optional[str] = None,
+        tags: Optional[str] = None,
+        additional_fields: Optional[str] = None,
+        date_format: Optional[str] = None,
+    ) -> str:
+        """Generate new frontmatter and prepend to content."""
+        # Strip any existing frontmatter
+        clean_content = (
+            self._strip_frontmatter(content) if isinstance(content, str) else ""
+        )
+
+        # Build frontmatter dictionary
+        frontmatter_dict = {}
+
+        # Try to infer title if not provided
+        if title:
+            frontmatter_dict["title"] = title
+        else:
+            inferred_title = self._infer_title_from_content(clean_content)
+            if inferred_title:
+                frontmatter_dict["title"] = inferred_title
+
+        # Add other fields if provided
+        if author:
+            frontmatter_dict["author"] = author
+
+        # Process date with flexible parser
+        if date:
+            frontmatter_dict["date"] = self._parse_flexible_date(date, date_format)
+        else:
+            # Use current date with provided format or default
+            format_to_use = date_format or "%Y-%m-%d"
+            frontmatter_dict["date"] = datetime.now().strftime(format_to_use)
+
+        if description:
+            frontmatter_dict["description"] = description
+
+        if tags:
+            frontmatter_dict["tags"] = self._parse_tags(tags)
+
+        # Add additional fields
+        if additional_fields:
+            additional_dict = self._parse_additional_fields(additional_fields)
+            frontmatter_dict.update(additional_dict)
+
+        # Generate YAML frontmatter
+        frontmatter_yaml = yaml.dump(
+            frontmatter_dict, sort_keys=False, default_flow_style=False
+        )
+        frontmatter = f"---\n{frontmatter_yaml}---\n\n"
+
+        # Combine frontmatter with content
+        return frontmatter + clean_content
+
+
+# Example usage
+def example_usage():
+    """Simple example of how to use the FrontmatterGeneratorTool"""
+    tool = FrontmatterGeneratorTool()
+
+    # Example document
+    document = """# The Role of the Seat in International Arbitration
+
+This paper examines the significance of the seat in international arbitration proceedings,
+with a focus on recent developments in Australian law."""
+
+    # Generate frontmatter
+    result = tool.forward(
+        content=document,
+        title="The Role of the Seat in International Arbitration",
+        author="Matthew Barry",
+        date="13 March 2013",  # Now handles this format correctly
+        description="A legal analysis of the role of the seat in international arbitration.",
+        tags="international arbitration, enforcement, australian courts",
+        mode="generate",
+    )
+
+    print(result)
+    return result
+
+
+# Define what gets imported with "from frontmatter_generator import *"
+__all__ = ["FrontmatterGeneratorTool", "example_usage"]
+
+# Example usage
+if __name__ == "__main__":
+    tool = FrontmatterGeneratorTool()
+
+#     # Example document
+#     document = """# The Role of the Seat in International Arbitration
+
+# This paper examines the significance of the seat in international arbitration proceedings,
+# with a focus on recent developments in Australian law."""
+
+#     # Generate frontmatter
+#     result = tool.forward(
+#         content=document,
+#         title="The Role of the Seat in International Arbitration",
+#         author="Matthew Barry",
+#         date="2013-12-13",
+#         description="A legal analysis of the role of the seat in international arbitration.",
+#         tags="international arbitration, enforcement, australian courts",
+#         mode="generate"
+#     )
+
+#     print(result)