Merge pull request #621 from sambhavnoobcoder/md-py-pdf-output-artifacts

Add Multiple Output Format Support for Artifacts

docs/swarms/artifacts/artifact.md

@@ -41,7 +41,7 @@ The `Artifact` class includes various methods for creating, editing, saving, loa
 artifact = Artifact(file_path="example.txt", file_type="txt")
 artifact.create(initial_content="Initial file content")
 ```
-
+The file type parameter supports the following file types: `.txt`, `.md`, `.py`, `.pdf`.
 #### `edit`
 
 
@@ -240,4 +240,4 @@ new_artifact = Artifact.from_dict(artifact_dict)
 
 # Print the metrics of the new artifact
 print(new_artifact.get_metrics())
-```
+```

pyproject.toml

@@ -78,6 +78,7 @@ aiofiles = "*"
 swarm-models = "*"
 clusterops = "*"
 chromadb = "*"
+reportlab = "*"
 
 [tool.poetry.scripts]
 swarms = "swarms.cli.main:main"

requirements.txt

@@ -33,4 +33,5 @@ swarms-memory
 pre-commit
 aiofiles
 swarm-models
-clusterops
+clusterops
+reportlab

swarms/artifacts/main_artifact.py

@@ -243,6 +243,60 @@ class Artifact(BaseModel):
             logger.error(f"Error creating artifact from dict: {e}")
             raise e
 
+    def save_as(self, output_format: str) -> None:
+        """
+        Saves the artifact's contents in the specified format.
+        
+        Args:
+            output_format (str): The desired output format ('md', 'txt', 'pdf', 'py')
+        
+        Raises:
+            ValueError: If the output format is not supported
+        """
+        supported_formats = {'.md', '.txt', '.pdf', '.py'}
+        if output_format not in supported_formats:
+            raise ValueError(f"Unsupported output format. Supported formats are: {supported_formats}")
+            
+        output_path = os.path.splitext(self.file_path)[0] + output_format
+        
+        if output_format == '.pdf':
+            self._save_as_pdf(output_path)
+        else:
+            with open(output_path, 'w', encoding='utf-8') as f:
+                if output_format == '.md':
+                    # Add markdown formatting if needed
+                    f.write(f"# {os.path.basename(self.file_path)}\n\n")
+                    f.write(self.contents)
+                elif output_format == '.py':
+                    # Add Python file header
+                    f.write('"""\n')
+                    f.write(f'Generated Python file from {self.file_path}\n')
+                    f.write('"""\n\n')
+                    f.write(self.contents)
+                else:  # .txt
+                    f.write(self.contents)
+
+    def _save_as_pdf(self, output_path: str) -> None:
+        """
+        Helper method to save content as PDF using reportlab
+        """
+        try:
+            from reportlab.pdfgen import canvas
+            from reportlab.lib.pagesizes import letter
+            
+            c = canvas.Canvas(output_path, pagesize=letter)
+            # Split content into lines
+            y = 750  # Starting y position
+            for line in self.contents.split('\n'):
+                c.drawString(50, y, line)
+                y -= 15  # Move down for next line
+                if y < 50:  # New page if bottom reached
+                    c.showPage()
+                    y = 750
+            c.save()
+        except ImportError:
+            raise ImportError("reportlab package is required for PDF output. Install with: pip install reportlab")
+
 
 # # Example usage
 # artifact = Artifact(file_path="example.txt", file_type=".txt")
@@ -259,3 +313,15 @@ class Artifact(BaseModel):
 
 # # # Get metrics
 # print(artifact.get_metrics())
+
+
+# Testing saving in different artifact types 
+# Create an artifact
+#artifact = Artifact(file_path="/path/to/file", file_type=".txt",contents="",  edit_count=0  )
+#artifact.create("This is some content\nWith multiple lines")
+
+# Save in different formats
+#artifact.save_as(".md")    # Creates example.md
+#artifact.save_as(".txt")   # Creates example.txt
+#artifact.save_as(".pdf")   # Creates example.pdf
+#artifact.save_as(".py")    # Creates example.py

tests/artifacts/test_artifact_output_types.py

@@ -0,0 +1,108 @@
+import unittest
+import os
+from unittest.mock import patch, mock_open
+import tempfile
+import sys
+from pathlib import Path
+from datetime import datetime
+import json
+from swarms.artifacts.main_artifact import Artifact
+
+class TestArtifactSaveAs(unittest.TestCase):
+    def setUp(self):
+        """Set up test fixtures before each test method."""
+        self.temp_dir = tempfile.mkdtemp()
+        self.test_file_path = os.path.join(self.temp_dir, "test_file.txt")
+        self.test_content = "This is test content\nWith multiple lines"
+        
+        # Create artifact with all required fields
+        self.artifact = Artifact(
+            file_path=self.test_file_path,
+            file_type=".txt",
+            contents=self.test_content,  # Provide initial content
+            edit_count=0
+        )
+        self.artifact.create(self.test_content)
+
+    def tearDown(self):
+        """Clean up test fixtures after each test method."""
+        try:
+            if os.path.exists(self.test_file_path):
+                os.remove(self.test_file_path)
+            # Clean up any potential output files
+            base_path = os.path.splitext(self.test_file_path)[0]
+            for ext in ['.md', '.txt', '.py', '.pdf']:
+                output_file = base_path + ext
+                if os.path.exists(output_file):
+                    os.remove(output_file)
+            os.rmdir(self.temp_dir)
+        except Exception as e:
+            print(f"Cleanup error: {e}")
+
+    def test_save_as_txt(self):
+        """Test saving artifact as .txt file"""
+        output_path = os.path.splitext(self.test_file_path)[0] + '.txt'
+        self.artifact.save_as('.txt')
+        self.assertTrue(os.path.exists(output_path))
+        with open(output_path, 'r', encoding='utf-8') as f:
+            content = f.read()
+        self.assertEqual(content, self.test_content)
+
+    def test_save_as_markdown(self):
+        """Test saving artifact as .md file"""
+        output_path = os.path.splitext(self.test_file_path)[0] + '.md'
+        self.artifact.save_as('.md')
+        self.assertTrue(os.path.exists(output_path))
+        with open(output_path, 'r', encoding='utf-8') as f:
+            content = f.read()
+        self.assertIn(self.test_content, content)
+        self.assertIn('# test_file.txt', content)
+
+    def test_save_as_python(self):
+        """Test saving artifact as .py file"""
+        output_path = os.path.splitext(self.test_file_path)[0] + '.py'
+        self.artifact.save_as('.py')
+        self.assertTrue(os.path.exists(output_path))
+        with open(output_path, 'r', encoding='utf-8') as f:
+            content = f.read()
+        self.assertIn(self.test_content, content)
+        self.assertIn('"""', content)
+        self.assertIn('Generated Python file', content)
+
+    @patch('builtins.open', new_callable=mock_open)
+    def test_file_writing_called(self, mock_file):
+        """Test that file writing is actually called"""
+        self.artifact.save_as('.txt')
+        mock_file.assert_called()
+        mock_file().write.assert_called_with(self.test_content)
+
+    def test_invalid_format(self):
+        """Test saving artifact with invalid format"""
+        with self.assertRaises(ValueError):
+            self.artifact.save_as('.invalid')
+
+    def test_export_import_json(self):
+        """Test exporting and importing JSON format"""
+        json_path = os.path.join(self.temp_dir, "test.json")
+        
+        # Export to JSON
+        self.artifact.export_to_json(json_path)
+        self.assertTrue(os.path.exists(json_path))
+        
+        # Import from JSON and convert timestamp back to string
+        with open(json_path, 'r') as f:
+            data = json.loads(f.read())
+            # Ensure timestamps are strings
+            for version in data.get('versions', []):
+                if isinstance(version.get('timestamp'), str):
+                    version['timestamp'] = version['timestamp']
+        
+        # Import the modified data
+        imported_artifact = Artifact(**data)
+        self.assertEqual(imported_artifact.contents, self.test_content)
+        
+        # Cleanup
+        os.remove(json_path)
+
+if __name__ == '__main__':
+    unittest.main()