diff --git a/docs/swarms/artifacts/artifact.md b/docs/swarms/artifacts/artifact.md index f9551b2fa..d7727f8e3 100644 --- a/docs/swarms/artifacts/artifact.md +++ b/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()) -``` \ No newline at end of file +``` diff --git a/pyproject.toml b/pyproject.toml index e0f2469a5..8a2735585 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -78,6 +78,7 @@ aiofiles = "*" swarm-models = "*" clusterops = "*" chromadb = "*" +reportlab = "*" [tool.poetry.scripts] swarms = "swarms.cli.main:main" diff --git a/requirements.txt b/requirements.txt index 953de5178..8f9df9b91 100644 --- a/requirements.txt +++ b/requirements.txt @@ -33,4 +33,5 @@ swarms-memory pre-commit aiofiles swarm-models -clusterops \ No newline at end of file +clusterops +reportlab diff --git a/swarms/artifacts/main_artifact.py b/swarms/artifacts/main_artifact.py index ed9b9d939..a59b6fb94 100644 --- a/swarms/artifacts/main_artifact.py +++ b/swarms/artifacts/main_artifact.py @@ -243,6 +243,60 @@ def from_dict(cls, data: Dict[str, Any]) -> "Artifact": 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 @@ def from_dict(cls, data: Dict[str, Any]) -> "Artifact": # # # 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 diff --git a/tests/artifacts/test_artifact_output_types.py b/tests/artifacts/test_artifact_output_types.py new file mode 100644 index 000000000..8fe37dd71 --- /dev/null +++ b/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()