From 13a6e335fb163b932ed037562fcedbc269f0d5a5 Mon Sep 17 00:00:00 2001 From: Travis Dent Date: Mon, 9 Dec 2024 15:37:46 -0800 Subject: [PATCH 1/3] Project template validation. Add tests for initializing project from all bundled templates. --- agentstack/cli/cli.py | 52 +++++------ agentstack/proj_templates.py | 92 +++++++++++++++++++ .../proj_templates/content_creator.json | 2 +- tests/test_cli_init.py | 43 +++++++++ tests/test_cli_loads.py | 18 ---- tests/test_templates_config.py | 20 ++++ 6 files changed, 177 insertions(+), 50 deletions(-) create mode 100644 agentstack/proj_templates.py create mode 100644 tests/test_cli_init.py create mode 100644 tests/test_templates_config.py diff --git a/agentstack/cli/cli.py b/agentstack/cli/cli.py index bc5648c..85c5654 100644 --- a/agentstack/cli/cli.py +++ b/agentstack/cli/cli.py @@ -31,6 +31,7 @@ from agentstack import generation from agentstack.utils import open_json_file, term_color, is_snake_case from agentstack.update import AGENTSTACK_PACKAGE +from agentstack.proj_templates import TemplateConfig PREFERRED_MODELS = [ @@ -57,45 +58,34 @@ def init_project_builder( template_data = None if template is not None: - url_start = "https://" - if template[: len(url_start)] == url_start: - # template is a url - response = requests.get(template) - if response.status_code == 200: - template_data = response.json() - else: - print( - term_color( - f"Failed to fetch template data from {template}. Status code: {response.status_code}", - 'red', - ) - ) + if template.startswith("https://"): + try: + template_data = TemplateConfig.from_url(template) + except Exception as e: + print(term_color(f"Failed to fetch template data from {template}", 'red')) sys.exit(1) else: - with importlib.resources.path( - 'agentstack.templates.proj_templates', f'{template}.json' - ) as template_path: - if template_path is None: - print(term_color(f"No such template {template} found", 'red')) - sys.exit(1) - template_data = open_json_file(template_path) + try: + template_data = TemplateConfig.from_template_name(template) + except Exception as e: + print(term_color(f"Failed to load template {template}", 'red')) + sys.exit(1) if template_data: project_details = { - "name": slug_name or template_data['name'], + "name": slug_name or template_data.name, "version": "0.0.1", - "description": template_data['description'], + "description": template_data.description, "author": "Name ", "license": "MIT", } - framework = template_data['framework'] + framework = template_data.framework design = { - 'agents': template_data['agents'], - 'tasks': template_data['tasks'], - 'inputs': template_data['inputs'], + 'agents': template_data.agents, + 'tasks': template_data.tasks, + 'inputs': template_data.inputs, } - - tools = template_data['tools'] + tools = template_data.tools elif use_wizard: welcome_message() @@ -390,7 +380,7 @@ def insert_template( project_details: dict, framework_name: str, design: dict, - template_data: Optional[dict] = None, + template_data: Optional[TemplateConfig] = None, ): framework = FrameworkData(framework_name.lower()) project_metadata = ProjectMetadata( @@ -400,8 +390,8 @@ def insert_template( version="0.0.1", license="MIT", year=datetime.now().year, - template=template_data['name'] if template_data else 'none', - template_version=template_data['template_version'] if template_data else '0', + template=template_data.name if template_data else 'none', + template_version=template_data.template_version if template_data else '0', ) project_structure = ProjectStructure() diff --git a/agentstack/proj_templates.py b/agentstack/proj_templates.py new file mode 100644 index 0000000..1ace2fe --- /dev/null +++ b/agentstack/proj_templates.py @@ -0,0 +1,92 @@ +from typing import Optional +import os, sys +from pathlib import Path +import pydantic +import requests +from agentstack import ValidationError +from agentstack.utils import get_package_path, open_json_file, term_color + + +class TemplateConfig(pydantic.BaseModel): + """ + Interface for interacting with template configuration files. + + Templates are read-only. + + Template Schema + ------------- + name: str + The name of the project. + description: str + A description of the template. + template_version: str + The version of the template. + framework: str + The framework the template is for. + method: str + The method used by the project. ie. "sequential" + agents: list[dict] + A list of agents used by the project. TODO vaidate this against an agent schema + tasks: list[dict] + A list of tasks used by the project. TODO validate this against a task schema + tools: list[dict] + A list of tools used by the project. TODO validate this against a tool schema + inputs: list[str] + A list of inputs used by the project. + """ + + name: str + description: str + template_version: int + framework: str + method: str + agents: list[dict] + tasks: list[dict] + tools: list[dict] + inputs: list[str] + + @classmethod + def from_template_name(cls, name: str) -> 'TemplateConfig': + path = get_package_path() / f'templates/proj_templates/{name}.json' + if not os.path.exists(path): # TODO raise exceptions and handle message/exit in cli + print(term_color(f'No known agentstack tool: {name}', 'red')) + sys.exit(1) + return cls.from_json(path) + + @classmethod + def from_json(cls, path: Path) -> 'ToolConfig': + data = open_json_file(path) + try: + return cls(**data) + except pydantic.ValidationError as e: + # TODO raise exceptions and handle message/exit in cli + print(term_color(f"Error validating template config JSON: \n{path}", 'red')) + for error in e.errors(): + print(f"{' '.join([str(loc) for loc in error['loc']])}: {error['msg']}") + sys.exit(1) + + @classmethod + def from_url(cls, url: str) -> 'TemplateConfig': + if not url.startswith("https://"): + raise ValidationError(f"Invalid URL: {url}") + response = requests.get(url) + if response.status_code != 200: + raise ValidationError(f"Failed to fetch template from {url}") + return cls(**response.json()) + + +def get_all_template_paths() -> list[Path]: + paths = [] + templates_dir = get_package_path() / 'templates/proj_templates' + for file in templates_dir.iterdir(): + if file.suffix == '.json': + paths.append(file) + return paths + + +def get_all_template_names() -> list[str]: + return [path.stem for path in get_all_template_paths()] + + +def get_all_templates() -> list[TemplateConfig]: + return [TemplateConfig.from_json(path) for path in get_all_template_paths()] diff --git a/agentstack/templates/proj_templates/content_creator.json b/agentstack/templates/proj_templates/content_creator.json index bf63a47..53247ca 100644 --- a/agentstack/templates/proj_templates/content_creator.json +++ b/agentstack/templates/proj_templates/content_creator.json @@ -1,5 +1,5 @@ { - "name": "content_creation", + "name": "content_creator", "description": "Multi-agent system for creating high-quality content", "template_version": 1, "framework": "crewai", diff --git a/tests/test_cli_init.py b/tests/test_cli_init.py new file mode 100644 index 0000000..2b0aa74 --- /dev/null +++ b/tests/test_cli_init.py @@ -0,0 +1,43 @@ +import subprocess +import os, sys +import unittest +from parameterized import parameterized +from pathlib import Path +import shutil +from agentstack.proj_templates import get_all_template_names + +BASE_PATH = Path(__file__).parent +CLI_ENTRY = [ + sys.executable, + "-m", + "agentstack.main", +] + + +class CLIInitTest(unittest.TestCase): + def setUp(self): + self.project_dir = Path(BASE_PATH / 'tmp/cli_init') + os.makedirs(self.project_dir) + + def tearDown(self): + shutil.rmtree(self.project_dir) + + def _run_cli(self, *args): + """Helper method to run the CLI with arguments.""" + return subprocess.run([*CLI_ENTRY, *args], capture_output=True, text=True) + + def test_init_command(self): + """Test the 'init' command to create a project directory.""" + os.chdir(self.project_dir) + result = self._run_cli('init', str(self.project_dir)) + self.assertEqual(result.returncode, 0) + self.assertTrue(self.project_dir.exists()) + + @parameterized.expand([(x, ) for x in get_all_template_names()]) + def test_init_command_for_template(self, template_name): + """Test the 'init' command to create a project directory with a template.""" + os.chdir(self.project_dir) + result = self._run_cli('init', str(self.project_dir), '--template', template_name) + self.assertEqual(result.returncode, 0) + self.assertTrue(self.project_dir.exists()) + diff --git a/tests/test_cli_loads.py b/tests/test_cli_loads.py index 7dfaf0b..819983a 100644 --- a/tests/test_cli_loads.py +++ b/tests/test_cli_loads.py @@ -8,7 +8,6 @@ class TestAgentStackCLI(unittest.TestCase): - # Replace with your actual CLI entry point if different CLI_ENTRY = [ sys.executable, "-m", @@ -32,23 +31,6 @@ def test_invalid_command(self): self.assertNotEqual(result.returncode, 0) self.assertIn("usage:", result.stderr) - def test_init_command(self): - """Test the 'init' command to create a project directory.""" - test_dir = Path(BASE_PATH / 'tmp/test_project') - - # Ensure the directory doesn't exist from previous runs - if test_dir.exists(): - shutil.rmtree(test_dir) - os.makedirs(test_dir) - - os.chdir(test_dir) - result = self.run_cli("init", str(test_dir)) - self.assertEqual(result.returncode, 0) - self.assertTrue(test_dir.exists()) - - # Clean up - shutil.rmtree(test_dir) - def test_run_command_invalid_project(self): """Test the 'run' command on an invalid project.""" test_dir = Path(BASE_PATH / 'tmp/test_project') diff --git a/tests/test_templates_config.py b/tests/test_templates_config.py new file mode 100644 index 0000000..10077c8 --- /dev/null +++ b/tests/test_templates_config.py @@ -0,0 +1,20 @@ +import json +import unittest +from pathlib import Path +from agentstack.proj_templates import TemplateConfig, get_all_template_names, get_all_template_paths + +BASE_PATH = Path(__file__).parent + + +class TemplateConfigTest(unittest.TestCase): + def test_all_configs_from_template_name(self): + for template_name in get_all_template_names(): + config = TemplateConfig.from_template_name(template_name) + assert config.name == template_name + # We can assume that pydantic validation caught any other issues + + def test_all_configs_from_template_path(self): + for path in get_all_template_paths(): + config = TemplateConfig.from_json(path) + assert config.name == path.stem + # We can assume that pydantic validation caught any other issues From fe39a508b647f9ec512d89e040802d9472fcf50f Mon Sep 17 00:00:00 2001 From: Travis Dent Date: Mon, 9 Dec 2024 15:42:44 -0800 Subject: [PATCH 2/3] Type check, lint. --- agentstack/cli/agentstack_data.py | 4 ++-- agentstack/cli/cli.py | 2 +- agentstack/proj_templates.py | 2 +- tests/test_cli_init.py | 3 +-- 4 files changed, 5 insertions(+), 6 deletions(-) diff --git a/agentstack/cli/agentstack_data.py b/agentstack/cli/agentstack_data.py index bc4760d..58600b5 100644 --- a/agentstack/cli/agentstack_data.py +++ b/agentstack/cli/agentstack_data.py @@ -1,6 +1,6 @@ import json from datetime import datetime -from typing import Optional +from typing import Optional, Union from agentstack.utils import clean_input, get_version from agentstack.logger import log @@ -17,7 +17,7 @@ def __init__( license: str = "", year: int = datetime.now().year, template: str = "none", - template_version: str = "0", + template_version: int = 0, ): self.project_name = clean_input(project_name) if project_name else "myagent" self.project_slug = clean_input(project_slug) if project_slug else self.project_name diff --git a/agentstack/cli/cli.py b/agentstack/cli/cli.py index d26b413..1601ffb 100644 --- a/agentstack/cli/cli.py +++ b/agentstack/cli/cli.py @@ -384,7 +384,7 @@ def insert_template( license="MIT", year=datetime.now().year, template=template_data.name if template_data else 'none', - template_version=template_data.template_version if template_data else '0', + template_version=template_data.template_version if template_data else 0, ) project_structure = ProjectStructure() diff --git a/agentstack/proj_templates.py b/agentstack/proj_templates.py index 1ace2fe..c96000c 100644 --- a/agentstack/proj_templates.py +++ b/agentstack/proj_templates.py @@ -54,7 +54,7 @@ def from_template_name(cls, name: str) -> 'TemplateConfig': return cls.from_json(path) @classmethod - def from_json(cls, path: Path) -> 'ToolConfig': + def from_json(cls, path: Path) -> 'TemplateConfig': data = open_json_file(path) try: return cls(**data) diff --git a/tests/test_cli_init.py b/tests/test_cli_init.py index 2b0aa74..80b7faa 100644 --- a/tests/test_cli_init.py +++ b/tests/test_cli_init.py @@ -33,11 +33,10 @@ def test_init_command(self): self.assertEqual(result.returncode, 0) self.assertTrue(self.project_dir.exists()) - @parameterized.expand([(x, ) for x in get_all_template_names()]) + @parameterized.expand([(x,) for x in get_all_template_names()]) def test_init_command_for_template(self, template_name): """Test the 'init' command to create a project directory with a template.""" os.chdir(self.project_dir) result = self._run_cli('init', str(self.project_dir), '--template', template_name) self.assertEqual(result.returncode, 0) self.assertTrue(self.project_dir.exists()) - From 9e781bc9cf44278752f90674c0660be502d7b642 Mon Sep 17 00:00:00 2001 From: Travis Dent Date: Mon, 9 Dec 2024 15:50:54 -0800 Subject: [PATCH 3/3] Unnecessary import --- agentstack/cli/agentstack_data.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/agentstack/cli/agentstack_data.py b/agentstack/cli/agentstack_data.py index 58600b5..c51e3dd 100644 --- a/agentstack/cli/agentstack_data.py +++ b/agentstack/cli/agentstack_data.py @@ -1,6 +1,6 @@ import json from datetime import datetime -from typing import Optional, Union +from typing import Optional from agentstack.utils import clean_input, get_version from agentstack.logger import log