From 479f9a0b938916a701ed723fe62f7a6c8b26e530 Mon Sep 17 00:00:00 2001
From: Travis Dent <tcdent@gmail.com>
Date: Tue, 10 Dec 2024 17:53:19 -0800
Subject: [PATCH 1/5] Break inputs out into a separate file, update project
 templates to write inputs, begin agentstack public API for use in projects.

---
 agentstack/__init__.py                        | 14 +++-
 agentstack/cli/cli.py                         |  2 +-
 agentstack/inputs.py                          | 76 +++++++++++++++++++
 agentstack/proj_templates.py                  |  2 +-
 .../src/config/inputs.yaml                    |  4 +
 .../src/main.py                               | 22 +-----
 .../proj_templates/content_creator.json       |  6 +-
 7 files changed, 104 insertions(+), 22 deletions(-)
 create mode 100644 agentstack/inputs.py
 create mode 100644 agentstack/templates/crewai/{{cookiecutter.project_metadata.project_slug}}/src/config/inputs.yaml

diff --git a/agentstack/__init__.py b/agentstack/__init__.py
index e645be5c..a2e17af1 100644
--- a/agentstack/__init__.py
+++ b/agentstack/__init__.py
@@ -1,8 +1,20 @@
+"""
+This it the beginning of the agentstack public API. 
 
+Methods that have been imported into this file are expected to be used by the
+end user inside of their project. 
+"""
+from agentstack.inputs import get_inputs
+
+___all___ = [
+    "ValidationError", 
+    "get_inputs", 
+]
 
 class ValidationError(Exception):
     """
     Raised when a validation error occurs ie. a file does not meet the required
     format or a syntax error is found.
     """
-    pass
\ No newline at end of file
+    pass
+
diff --git a/agentstack/cli/cli.py b/agentstack/cli/cli.py
index 3d426746..63a0e230 100644
--- a/agentstack/cli/cli.py
+++ b/agentstack/cli/cli.py
@@ -389,7 +389,7 @@ def insert_template(
     project_structure = ProjectStructure()
     project_structure.agents = design["agents"]
     project_structure.tasks = design["tasks"]
-    project_structure.set_inputs(design["inputs"])
+    project_structure.inputs = design["inputs"]
 
     cookiecutter_data = CookiecutterData(
         project_metadata=project_metadata,
diff --git a/agentstack/inputs.py b/agentstack/inputs.py
new file mode 100644
index 00000000..a8cd58bf
--- /dev/null
+++ b/agentstack/inputs.py
@@ -0,0 +1,76 @@
+from typing import Optional
+import os
+from pathlib import Path
+import pydantic
+from ruamel.yaml import YAML, YAMLError
+from ruamel.yaml.scalarstring import FoldedScalarString
+from agentstack import ValidationError
+
+
+INPUTS_FILENAME: Path = Path("src/config/inputs.yaml")
+
+yaml = YAML()
+yaml.preserve_quotes = True  # Preserve quotes in existing data
+
+
+class InputsConfig:
+    """
+    Interface for interacting with inputs configuration.
+
+    Use it as a context manager to make and save edits:
+    ```python
+    with InputsConfig() as inputs:
+        inputs.topic = "Open Source Aritifical Intelligence"
+    ```
+    """
+
+    _attributes: dict[str, str]
+
+    def __init__(self, path: Optional[Path] = None):
+        self.path = path if path else Path()
+        filename = self.path / INPUTS_FILENAME
+        
+        if not os.path.exists(filename):
+            os.makedirs(filename.parent, exist_ok=True)
+            filename.touch()
+
+        try:
+            with open(filename, 'r') as f:
+                self._attributes = yaml.load(f) or {}
+        except YAMLError as e:
+            # TODO format MarkedYAMLError lines/messages
+            raise ValidationError(f"Error parsing inputs file: {filename}\n{e}")
+
+    def __getitem__(self, key: str) -> str:
+        return self._attributes[key]
+
+    def __setitem__(self, key: str, value: str):
+        self._attributes[key] = value
+
+    def __contains__(self, key: str) -> bool:
+        return key in self._attributes
+
+    def to_dict(self) -> dict[str, str]:
+        return self._attributes
+
+    def model_dump(self) -> dict:
+        dump = {}
+        for key, value in self._attributes.items():
+            dump[key] = FoldedScalarString(value)
+        return dump
+
+    def write(self):
+        with open(self.path / INPUTS_FILENAME, 'w') as f:
+            yaml.dump(self.model_dump(), f)
+
+    def __enter__(self) -> 'InputsConfig':
+        return self
+
+    def __exit__(self, *args):
+        self.write()
+
+
+def get_inputs(path: Optional[Path] = None) -> dict:
+    path = path if path else Path()
+    config = InputsConfig(path)
+    return config.to_dict()
diff --git a/agentstack/proj_templates.py b/agentstack/proj_templates.py
index c96000c5..35717497 100644
--- a/agentstack/proj_templates.py
+++ b/agentstack/proj_templates.py
@@ -43,7 +43,7 @@ class TemplateConfig(pydantic.BaseModel):
     agents: list[dict]
     tasks: list[dict]
     tools: list[dict]
-    inputs: list[str]
+    inputs: dict[str, str]
 
     @classmethod
     def from_template_name(cls, name: str) -> 'TemplateConfig':
diff --git a/agentstack/templates/crewai/{{cookiecutter.project_metadata.project_slug}}/src/config/inputs.yaml b/agentstack/templates/crewai/{{cookiecutter.project_metadata.project_slug}}/src/config/inputs.yaml
new file mode 100644
index 00000000..fcac19a9
--- /dev/null
+++ b/agentstack/templates/crewai/{{cookiecutter.project_metadata.project_slug}}/src/config/inputs.yaml
@@ -0,0 +1,4 @@
+{%- for key, value in cookiecutter.structure.inputs.items() %}
+{{key}}: >
+    {{value}}
+{%- endfor %}
\ No newline at end of file
diff --git a/agentstack/templates/crewai/{{cookiecutter.project_metadata.project_slug}}/src/main.py b/agentstack/templates/crewai/{{cookiecutter.project_metadata.project_slug}}/src/main.py
index 291a6177..264a2eda 100644
--- a/agentstack/templates/crewai/{{cookiecutter.project_metadata.project_slug}}/src/main.py
+++ b/agentstack/templates/crewai/{{cookiecutter.project_metadata.project_slug}}/src/main.py
@@ -1,6 +1,7 @@
 #!/usr/bin/env python
 import sys
 from crew import {{cookiecutter.project_metadata.project_name|replace('-', '')|replace('_', '')|capitalize}}Crew
+import agentstack
 import agentops
 
 agentops.init(default_tags=['crewai', 'agentstack'])
@@ -10,25 +11,15 @@ def run():
     """
     Run the crew.
     """
-    inputs = {
-{%- for input in cookiecutter.structure.inputs %}
-        "{{input}}": "",
-{%- endfor %}
-    }
-    {{cookiecutter.project_metadata.project_name|replace('-', '')|replace('_', '')|capitalize}}Crew().crew().kickoff(inputs=inputs)
+    {{cookiecutter.project_metadata.project_name|replace('-', '')|replace('_', '')|capitalize}}Crew().crew().kickoff(inputs=agentstack.get_inputs())
 
 
 def train():
     """
     Train the crew for a given number of iterations.
     """
-    inputs = {
-{%- for input in cookiecutter.structure.inputs %}
-        "{{input}}": "",
-{%- endfor %}
-    }
     try:
-        {{cookiecutter.project_metadata.project_name|replace('-', '')|replace('_', '')|capitalize}}Crew().crew().train(n_iterations=int(sys.argv[1]), filename=sys.argv[2], inputs=inputs)
+        {{cookiecutter.project_metadata.project_name|replace('-', '')|replace('_', '')|capitalize}}Crew().crew().train(n_iterations=int(sys.argv[1]), filename=sys.argv[2], inputs=agentstack.get_inputs())
 
     except Exception as e:
         raise Exception(f"An error occurred while training the crew: {e}")
@@ -49,13 +40,8 @@ def test():
     """
     Test the crew execution and returns the results.
     """
-    inputs = {
-{%- for input in cookiecutter.structure.inputs %}
-        "{{input}}": "",
-{%- endfor %}
-    }
     try:
-        {{cookiecutter.project_metadata.project_name|replace('-', '')|replace('_', '')|capitalize}}Crew().crew().test(n_iterations=int(sys.argv[1]), openai_model_name=sys.argv[2], inputs=inputs)
+        {{cookiecutter.project_metadata.project_name|replace('-', '')|replace('_', '')|capitalize}}Crew().crew().test(n_iterations=int(sys.argv[1]), openai_model_name=sys.argv[2], inputs=agentstack.get_inputs())
 
     except Exception as e:
         raise Exception(f"An error occurred while replaying the crew: {e}")
diff --git a/agentstack/templates/proj_templates/content_creator.json b/agentstack/templates/proj_templates/content_creator.json
index 53247ca9..8bb81e72 100644
--- a/agentstack/templates/proj_templates/content_creator.json
+++ b/agentstack/templates/proj_templates/content_creator.json
@@ -46,5 +46,9 @@
     "agents": ["researcher"]
   }],
   "method": "sequential",
-  "inputs": ["topic", "audience", "content_type"]
+  "inputs": {
+    "topic": "", 
+    "audience": "", 
+    "content_type": ""
+  }
 }
\ No newline at end of file

From 6f0d2891430e57301bef192be33d6062a7a91689 Mon Sep 17 00:00:00 2001
From: Travis Dent <tcdent@gmail.com>
Date: Tue, 10 Dec 2024 20:54:10 -0800
Subject: [PATCH 2/5] Inputs config file tests. Versioned project templates.
 project main.py cleanup.

---
 agentstack/__init__.py                        |  8 +--
 agentstack/exceptions.py                      |  7 ++
 agentstack/inputs.py                          |  2 +-
 agentstack/proj_templates.py                  | 70 ++++++++++++++-----
 .../src/main.py                               | 20 ++++--
 .../proj_templates/content_creator.json       |  2 +-
 tests/fixtures/inputs_max.yaml                |  4 ++
 tests/fixtures/inputs_min.yaml                |  0
 tests/test_inputs_config.py                   | 29 ++++++++
 tests/test_templates_config.py                |  2 +-
 10 files changed, 109 insertions(+), 35 deletions(-)
 create mode 100644 agentstack/exceptions.py
 create mode 100644 tests/fixtures/inputs_max.yaml
 create mode 100644 tests/fixtures/inputs_min.yaml
 create mode 100644 tests/test_inputs_config.py

diff --git a/agentstack/__init__.py b/agentstack/__init__.py
index a2e17af1..e4328361 100644
--- a/agentstack/__init__.py
+++ b/agentstack/__init__.py
@@ -4,6 +4,7 @@
 Methods that have been imported into this file are expected to be used by the
 end user inside of their project. 
 """
+from agentstack.exceptions import ValidationError
 from agentstack.inputs import get_inputs
 
 ___all___ = [
@@ -11,10 +12,3 @@
     "get_inputs", 
 ]
 
-class ValidationError(Exception):
-    """
-    Raised when a validation error occurs ie. a file does not meet the required
-    format or a syntax error is found.
-    """
-    pass
-
diff --git a/agentstack/exceptions.py b/agentstack/exceptions.py
new file mode 100644
index 00000000..c0e95569
--- /dev/null
+++ b/agentstack/exceptions.py
@@ -0,0 +1,7 @@
+class ValidationError(Exception):
+    """
+    Raised when a validation error occurs ie. a file does not meet the required
+    format or a syntax error is found.
+    """
+
+    pass
diff --git a/agentstack/inputs.py b/agentstack/inputs.py
index a8cd58bf..70b2ff14 100644
--- a/agentstack/inputs.py
+++ b/agentstack/inputs.py
@@ -29,7 +29,7 @@ class InputsConfig:
     def __init__(self, path: Optional[Path] = None):
         self.path = path if path else Path()
         filename = self.path / INPUTS_FILENAME
-        
+
         if not os.path.exists(filename):
             os.makedirs(filename.parent, exist_ok=True)
             filename.touch()
diff --git a/agentstack/proj_templates.py b/agentstack/proj_templates.py
index 35717497..299f74b8 100644
--- a/agentstack/proj_templates.py
+++ b/agentstack/proj_templates.py
@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Optional, Literal
 import os, sys
 from pathlib import Path
 import pydantic
@@ -7,6 +7,31 @@
 from agentstack.utils import get_package_path, open_json_file, term_color
 
 
+class TemplateConfig_v1(pydantic.BaseModel):
+    name: str
+    description: str
+    template_version: Literal[1]
+    framework: str
+    method: str
+    agents: list[dict]
+    tasks: list[dict]
+    tools: list[dict]
+    inputs: list[str]
+
+    def to_v2(self) -> 'TemplateConfig':
+        return TemplateConfig(
+            name=self.name,
+            description=self.description,
+            template_version=2,
+            framework=self.framework,
+            method=self.method,
+            agents=self.agents,
+            tasks=self.tasks,
+            tools=self.tools,
+            inputs={key: "" for key in self.inputs},
+        )
+
+
 class TemplateConfig(pydantic.BaseModel):
     """
     Interface for interacting with template configuration files.
@@ -19,7 +44,7 @@ class TemplateConfig(pydantic.BaseModel):
         The name of the project.
     description: str
         A description of the template.
-    template_version: str
+    template_version: int
         The version of the template.
     framework: str
         The framework the template is for.
@@ -31,13 +56,13 @@ class TemplateConfig(pydantic.BaseModel):
         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.
+    inputs: dict[str, str]
+        A list of inputs and values used by the project.
     """
 
     name: str
     description: str
-    template_version: int
+    template_version: Literal[2]
     framework: str
     method: str
     agents: list[dict]
@@ -51,19 +76,11 @@ def from_template_name(cls, name: str) -> 'TemplateConfig':
         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)
+        return cls.from_file(path)
 
     @classmethod
-    def from_json(cls, path: Path) -> 'TemplateConfig':
-        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)
+    def from_file(cls, path: Path) -> 'TemplateConfig':
+        return cls.from_json(open_json_file(path))
 
     @classmethod
     def from_url(cls, url: str) -> 'TemplateConfig':
@@ -72,7 +89,24 @@ def from_url(cls, url: str) -> 'TemplateConfig':
         response = requests.get(url)
         if response.status_code != 200:
             raise ValidationError(f"Failed to fetch template from {url}")
-        return cls(**response.json())
+        return cls.from_json(response.json())
+
+    @classmethod
+    def from_json(cls, data: dict) -> 'TemplateConfig':
+        try:
+            match data.get('template_version'):
+                case 1:
+                    return TemplateConfig_v1(**data).to_v2()
+                case 2:
+                    return cls(**data)  # current version
+                case _:
+                    raise ValidationError(f"Unsupported template version: {data.get('template_version')}")
+        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)
 
 
 def get_all_template_paths() -> list[Path]:
@@ -89,4 +123,4 @@ def get_all_template_names() -> list[str]:
 
 
 def get_all_templates() -> list[TemplateConfig]:
-    return [TemplateConfig.from_json(path) for path in get_all_template_paths()]
+    return [TemplateConfig.from_file(path) for path in get_all_template_paths()]
diff --git a/agentstack/templates/crewai/{{cookiecutter.project_metadata.project_slug}}/src/main.py b/agentstack/templates/crewai/{{cookiecutter.project_metadata.project_slug}}/src/main.py
index 264a2eda..98f7987e 100644
--- a/agentstack/templates/crewai/{{cookiecutter.project_metadata.project_slug}}/src/main.py
+++ b/agentstack/templates/crewai/{{cookiecutter.project_metadata.project_slug}}/src/main.py
@@ -6,12 +6,13 @@
 
 agentops.init(default_tags=['crewai', 'agentstack'])
 
+instance = {{cookiecutter.project_metadata.project_name|replace('-', '')|replace('_', '')|capitalize}}Crew().crew()
 
 def run():
     """
     Run the crew.
     """
-    {{cookiecutter.project_metadata.project_name|replace('-', '')|replace('_', '')|capitalize}}Crew().crew().kickoff(inputs=agentstack.get_inputs())
+    instance.kickoff(inputs=agentstack.get_inputs())
 
 
 def train():
@@ -19,8 +20,11 @@ def train():
     Train the crew for a given number of iterations.
     """
     try:
-        {{cookiecutter.project_metadata.project_name|replace('-', '')|replace('_', '')|capitalize}}Crew().crew().train(n_iterations=int(sys.argv[1]), filename=sys.argv[2], inputs=agentstack.get_inputs())
-
+        instance.train(
+            n_iterations=int(sys.argv[1]), 
+            filename=sys.argv[2], 
+            inputs=agentstack.get_inputs(), 
+        )
     except Exception as e:
         raise Exception(f"An error occurred while training the crew: {e}")
 
@@ -30,8 +34,7 @@ def replay():
     Replay the crew execution from a specific task.
     """
     try:
-        {{cookiecutter.project_metadata.project_name|replace('-', '')|replace('_', '')|capitalize}}Crew().crew().replay(task_id=sys.argv[1])
-
+        instance.replay(task_id=sys.argv[1])
     except Exception as e:
         raise Exception(f"An error occurred while replaying the crew: {e}")
 
@@ -41,8 +44,11 @@ def test():
     Test the crew execution and returns the results.
     """
     try:
-        {{cookiecutter.project_metadata.project_name|replace('-', '')|replace('_', '')|capitalize}}Crew().crew().test(n_iterations=int(sys.argv[1]), openai_model_name=sys.argv[2], inputs=agentstack.get_inputs())
-
+        instance.test(
+            n_iterations=int(sys.argv[1]), 
+            openai_model_name=sys.argv[2], 
+            inputs=agentstack.get_inputs(), 
+        )
     except Exception as e:
         raise Exception(f"An error occurred while replaying the crew: {e}")
 
diff --git a/agentstack/templates/proj_templates/content_creator.json b/agentstack/templates/proj_templates/content_creator.json
index 8bb81e72..3e498ced 100644
--- a/agentstack/templates/proj_templates/content_creator.json
+++ b/agentstack/templates/proj_templates/content_creator.json
@@ -1,7 +1,7 @@
 {
   "name": "content_creator",
   "description": "Multi-agent system for creating high-quality content",
-  "template_version": 1,
+  "template_version": 2,
   "framework": "crewai",
   "agents": [{
     "name": "researcher",
diff --git a/tests/fixtures/inputs_max.yaml b/tests/fixtures/inputs_max.yaml
new file mode 100644
index 00000000..4de77bbe
--- /dev/null
+++ b/tests/fixtures/inputs_max.yaml
@@ -0,0 +1,4 @@
+input_name: >-
+    This in an input
+input_name_2: >-
+    This is another input
\ No newline at end of file
diff --git a/tests/fixtures/inputs_min.yaml b/tests/fixtures/inputs_min.yaml
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/test_inputs_config.py b/tests/test_inputs_config.py
new file mode 100644
index 00000000..1f20ace2
--- /dev/null
+++ b/tests/test_inputs_config.py
@@ -0,0 +1,29 @@
+import os
+import shutil
+import unittest
+from pathlib import Path
+from agentstack.inputs import InputsConfig
+
+BASE_PATH = Path(__file__).parent
+
+
+class InputsConfigTest(unittest.TestCase):
+    def setUp(self):
+        self.project_dir = BASE_PATH / "tmp/inputs_config"
+        os.makedirs(self.project_dir)
+        os.makedirs(self.project_dir / "src/config")
+
+    def tearDown(self):
+        shutil.rmtree(self.project_dir)
+
+    def test_minimal_input_config(self):
+        shutil.copy(BASE_PATH / "fixtures/inputs_min.yaml", self.project_dir / "src/config/inputs.yaml")
+        config = InputsConfig(self.project_dir)
+        assert config.to_dict() == {}
+
+    def test_maximal_input_config(self):
+        shutil.copy(BASE_PATH / "fixtures/inputs_max.yaml", self.project_dir / "src/config/inputs.yaml")
+        config = InputsConfig(self.project_dir)
+        assert config['input_name'] == "This in an input"
+        assert config['input_name_2'] == "This is another input"
+        assert config.to_dict() == {'input_name': "This in an input", 'input_name_2': "This is another input"}
diff --git a/tests/test_templates_config.py b/tests/test_templates_config.py
index 10077c80..05396da1 100644
--- a/tests/test_templates_config.py
+++ b/tests/test_templates_config.py
@@ -15,6 +15,6 @@ def test_all_configs_from_template_name(self):
 
     def test_all_configs_from_template_path(self):
         for path in get_all_template_paths():
-            config = TemplateConfig.from_json(path)
+            config = TemplateConfig.from_file(path)
             assert config.name == path.stem
             # We can assume that pydantic validation caught any other issues

From 0defaa00d20b8093712f31bd11f33afc6fa4101a Mon Sep 17 00:00:00 2001
From: Travis Dent <tcdent@gmail.com>
Date: Wed, 11 Dec 2024 17:17:21 -0800
Subject: [PATCH 3/5] Modify `run_project` imports to share scope with the
 `agentstack` app. Allow patching `inputs` with command line flags. Update
 tests.

---
 agentstack/cli/agentstack_data.py |  2 +-
 agentstack/cli/cli.py             | 51 ++++++++++++++++++++++---------
 agentstack/inputs.py              | 21 +++++++++++--
 agentstack/main.py                | 22 ++++++++++---
 agentstack/proj_templates.py      | 16 +++++-----
 tests/test_cli_templates.py       |  6 ++--
 tests/test_project_run.py         |  7 +++--
 7 files changed, 91 insertions(+), 34 deletions(-)

diff --git a/agentstack/cli/agentstack_data.py b/agentstack/cli/agentstack_data.py
index c51e3dd7..ec540deb 100644
--- a/agentstack/cli/agentstack_data.py
+++ b/agentstack/cli/agentstack_data.py
@@ -54,7 +54,7 @@ class ProjectStructure:
     def __init__(self):
         self.agents = []
         self.tasks = []
-        self.inputs = []
+        self.inputs = {}
 
     def add_agent(self, agent):
         self.agents.append(agent)
diff --git a/agentstack/cli/cli.py b/agentstack/cli/cli.py
index 7f9ee7b7..b920a72c 100644
--- a/agentstack/cli/cli.py
+++ b/agentstack/cli/cli.py
@@ -12,10 +12,12 @@
 import inquirer
 import os
 import importlib.resources
+import importlib.util
+from importlib import import_module
 from cookiecutter.main import cookiecutter
 from dotenv import load_dotenv
 import subprocess
-from packaging.metadata import Metadata
+
 
 from .agentstack_data import (
     FrameworkData,
@@ -30,6 +32,7 @@
 from agentstack import frameworks
 from agentstack import packaging
 from agentstack import generation
+from agentstack import inputs
 from agentstack.agents import get_all_agents
 from agentstack.tasks import get_all_tasks
 from agentstack.utils import open_json_file, term_color, is_snake_case, get_framework
@@ -110,7 +113,7 @@ def init_project_builder(
 
         framework = "crewai"  # TODO: if --no-wizard, require a framework flag
 
-        design = {'agents': [], 'tasks': [], 'inputs': []}
+        design = {'agents': [], 'tasks': [], 'inputs': {}}
 
         tools = []
 
@@ -155,14 +158,15 @@ def configure_default_model(path: Optional[str] = None):
         agentstack_config.default_model = model
 
 
-def run_project(framework: str, path: str = ''):
+def run_project(path: Optional[str] = None, cli_args: Optional[str] = None):
     """Validate that the project is ready to run and then run it."""
+    _path = Path(path) if path else Path.cwd()
+    framework = get_framework(_path)
+
     if framework not in frameworks.SUPPORTED_FRAMEWORKS:
         print(term_color(f"Framework {framework} is not supported by agentstack.", 'red'))
         sys.exit(1)
 
-    _path = Path(path)
-
     try:
         frameworks.validate_project(framework, _path)
     except frameworks.ValidationError as e:
@@ -170,8 +174,32 @@ def run_project(framework: str, path: str = ''):
         print(e)
         sys.exit(1)
 
-    load_dotenv(_path / '.env')  # explicitly load the project's .env file
-    subprocess.run(['python', 'src/main.py'], env=os.environ)
+    # Parse extra --input-* arguments for runtime overrides of the project's inputs
+    if cli_args:
+        for arg in cli_args:
+            if not arg.startswith('--input-'):
+                continue
+            key, value = arg[len('--input-') :].split('=')
+            inputs.add_input_for_run(key, value)
+
+    # explicitly load the project's .env file
+    load_dotenv(_path / '.env')
+
+    # import src/main.py from the project path
+    try:
+        spec = importlib.util.spec_from_file_location("main", str(_path / "src/main.py"))
+        assert spec is not None  # appease type checker
+        assert spec.loader is not None  # appease type checker
+        project_entrypoint = importlib.util.module_from_spec(spec)
+        sys.path.append(str(_path / "src"))
+        spec.loader.exec_module(project_entrypoint)
+    except ImportError as e:
+        print(term_color(f"Failed to import project. Does 'src/main.py' exist?:\n{e}", 'red'))
+        sys.exit(1)
+
+    # run the project's main function
+    # TODO try/except this and print detailed information with a --debug flag
+    project_entrypoint.run()
 
 
 def ask_framework() -> str:
@@ -527,13 +555,8 @@ def export_template(output_filename: str, path: str = ''):
             )
         )
 
-    inputs: list[str] = []
-    # TODO extract inputs from project
-    # for input in frameworks.get_input_names():
-    #     inputs.append(input)
-
     template = TemplateConfig(
-        template_version=1,
+        template_version=2,
         name=metadata.project_name,
         description=metadata.project_description,
         framework=framework,
@@ -541,7 +564,7 @@ def export_template(output_filename: str, path: str = ''):
         agents=agents,
         tasks=tasks,
         tools=tools,
-        inputs=inputs,
+        inputs=inputs.get_inputs(),
     )
 
     try:
diff --git a/agentstack/inputs.py b/agentstack/inputs.py
index 70b2ff14..b969b28c 100644
--- a/agentstack/inputs.py
+++ b/agentstack/inputs.py
@@ -12,6 +12,9 @@
 yaml = YAML()
 yaml.preserve_quotes = True  # Preserve quotes in existing data
 
+# run_inputs are set at the beginning of the run and are not saved
+run_inputs: dict[str, str] = {}
+
 
 class InputsConfig:
     """
@@ -71,6 +74,20 @@ def __exit__(self, *args):
 
 
 def get_inputs(path: Optional[Path] = None) -> dict:
+    """
+    Get the inputs configuration file and override with run_inputs.
+    """
     path = path if path else Path()
-    config = InputsConfig(path)
-    return config.to_dict()
+    config = InputsConfig(path).to_dict()
+    # run_inputs override saved inputs
+    for key, value in run_inputs.items():
+        config[key] = value
+    return config
+
+
+def add_input_for_run(key: str, value: str):
+    """
+    Add an input override for the current run.
+    This is used by the CLI to allow inputs to be set at runtime.
+    """
+    run_inputs[key] = value
diff --git a/agentstack/main.py b/agentstack/main.py
index 3b99ad3b..86fdb09d 100644
--- a/agentstack/main.py
+++ b/agentstack/main.py
@@ -43,7 +43,20 @@ def main():
     init_parser.add_argument("--template", "-t", help="Agent template to use")
 
     # 'run' command
-    _ = subparsers.add_parser("run", aliases=["r"], help="Run your agent")
+    run_parser = subparsers.add_parser(
+        "run",
+        aliases=["r"],
+        help="Run your agent",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog='''
+  --input-<key>=VALUE   Specify inputs to be passed to the run. 
+                        These will override the inputs in the project's inputs.yaml file.
+                        Examples: --input-topic=Sports --input-content-type=News
+    ''',
+    )
+    run_parser.add_argument(
+        "--path", "-p", help="Path to the project directory, defaults to current directory"
+    )
 
     # 'generate' command
     generate_parser = subparsers.add_parser("generate", aliases=["g"], help="Generate agents or tasks")
@@ -94,8 +107,8 @@ def main():
 
     update = subparsers.add_parser('update', aliases=['u'], help='Check for updates')
 
-    # Parse arguments
-    args = parser.parse_args()
+    # Parse known args and store unknown args in extras; some commands use them later on
+    args, extra_args = parser.parse_known_args()
 
     # Handle version
     if args.version:
@@ -115,8 +128,7 @@ def main():
     elif args.command in ["init", "i"]:
         init_project_builder(args.slug_name, args.template, args.wizard)
     elif args.command in ["run", "r"]:
-        framework = get_framework()
-        run_project(framework)
+        run_project(args.path, cli_args=extra_args)
     elif args.command in ['generate', 'g']:
         if args.generate_command in ['agent', 'a']:
             if not args.llm:
diff --git a/agentstack/proj_templates.py b/agentstack/proj_templates.py
index bd7cb591..a3cc1123 100644
--- a/agentstack/proj_templates.py
+++ b/agentstack/proj_templates.py
@@ -26,9 +26,9 @@ def to_v2(self) -> 'TemplateConfig':
             template_version=2,
             framework=self.framework,
             method=self.method,
-            agents=self.agents,
-            tasks=self.tasks,
-            tools=self.tools,
+            agents=[TemplateConfig.Agent(**agent) for agent in self.agents],
+            tasks=[TemplateConfig.Task(**task) for task in self.tasks],
+            tools=[TemplateConfig.Tool(**tool) for tool in self.tools],
             inputs={key: "" for key in self.inputs},
         )
 
@@ -99,14 +99,14 @@ def write_to_file(self, filename: Path):
     @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):
-            raise ValidationError(f"Template {name} not found.")
+        if not name in get_all_template_names():
+            raise ValidationError(f"Template {name} not bundled with agentstack.")
         return cls.from_file(path)
 
     @classmethod
     def from_file(cls, path: Path) -> 'TemplateConfig':
         if not os.path.exists(path):
-            raise ValidationError(f"Template {name} not found.")
+            raise ValidationError(f"Template {path} not found.")
         with open(path, 'r') as f:
             return cls.from_json(json.load(f))
 
@@ -130,12 +130,12 @@ def from_json(cls, data: dict) -> 'TemplateConfig':
                 case _:
                     raise ValidationError(f"Unsupported template version: {data.get('template_version')}")
         except pydantic.ValidationError as e:
-            err_msg = "Error validating template config JSON: \n    {path}\n\n"
+            err_msg = "Error validating template config JSON:\n"
             for error in e.errors():
                 err_msg += f"{' '.join([str(loc) for loc in error['loc']])}: {error['msg']}\n"
             raise ValidationError(err_msg)
         except json.JSONDecodeError as e:
-            raise ValidationError(f"Error decoding template JSON from URL:\n    {url}\n\n{e}")
+            raise ValidationError(f"Error decoding template JSON.\n{e}")
 
 
 def get_all_template_paths() -> list[Path]:
diff --git a/tests/test_cli_templates.py b/tests/test_cli_templates.py
index e98edf8f..1c8fb614 100644
--- a/tests/test_cli_templates.py
+++ b/tests/test_cli_templates.py
@@ -46,6 +46,8 @@ def test_export_template_v1(self):
         self.assertEqual(result.returncode, 0)
 
         result = self._run_cli('export', 'test_template.json')
+        print(result.stdout)
+        print(result.stderr)
         self.assertEqual(result.returncode, 0)
         self.assertTrue((self.project_dir / 'test_project/test_template.json').exists())
         template_str = (self.project_dir / 'test_project/test_template.json').read_text()
@@ -55,7 +57,7 @@ def test_export_template_v1(self):
             """{
     "name": "test_project",
     "description": "New agentstack project",
-    "template_version": 1,
+    "template_version": 2,
     "framework": "crewai",
     "method": "sequential",
     "agents": [
@@ -83,6 +85,6 @@ def test_export_template_v1(self):
             ]
         }
     ],
-    "inputs": []
+    "inputs": {}
 }""",
         )
diff --git a/tests/test_project_run.py b/tests/test_project_run.py
index 787589da..fec70f55 100644
--- a/tests/test_project_run.py
+++ b/tests/test_project_run.py
@@ -20,6 +20,9 @@ def setUp(self):
         os.makedirs(self.project_dir / 'src')
         (self.project_dir / 'src' / '__init__.py').touch()
 
+        with open(self.project_dir / 'src' / 'main.py', 'w') as f:
+            f.write('def run(): pass')
+
         # set the framework in agentstack.json
         shutil.copy(BASE_PATH / 'fixtures' / 'agentstack.json', self.project_dir / 'agentstack.json')
         with ConfigFile(self.project_dir) as config:
@@ -36,11 +39,11 @@ def tearDown(self):
         shutil.rmtree(self.project_dir)
 
     def test_run_project(self):
-        run_project(self.framework, self.project_dir)
+        run_project(self.project_dir)
 
     def test_env_is_set(self):
         """
         After running a project, the environment variables should be set from project_dir/.env.
         """
-        run_project(self.framework, self.project_dir)
+        run_project(self.project_dir)
         assert os.getenv('ENV_VAR1') == 'value1'

From b5c1a080ffeb2475b3a854a0504b8c72c6a929f2 Mon Sep 17 00:00:00 2001
From: Travis Dent <tcdent@gmail.com>
Date: Wed, 11 Dec 2024 20:14:23 -0800
Subject: [PATCH 4/5] Allow specifying function to run from main.py in
 `agentstack run`. Move run function to run.py in cli package for cleaner
 module import context.

---
 agentstack/cli/__init__.py |  3 +-
 agentstack/cli/cli.py      | 60 +++------------------------------
 agentstack/cli/run.py      | 68 ++++++++++++++++++++++++++++++++++++++
 agentstack/main.py         | 14 ++++++--
 tests/test_project_run.py  |  4 +--
 5 files changed, 89 insertions(+), 60 deletions(-)
 create mode 100644 agentstack/cli/run.py

diff --git a/agentstack/cli/__init__.py b/agentstack/cli/__init__.py
index 47adf18c..94f53ebf 100644
--- a/agentstack/cli/__init__.py
+++ b/agentstack/cli/__init__.py
@@ -1 +1,2 @@
-from .cli import init_project_builder, list_tools, configure_default_model, run_project, export_template
+from .cli import init_project_builder, list_tools, configure_default_model, export_template
+from .run import run_project
\ No newline at end of file
diff --git a/agentstack/cli/cli.py b/agentstack/cli/cli.py
index b920a72c..d33a85f3 100644
--- a/agentstack/cli/cli.py
+++ b/agentstack/cli/cli.py
@@ -1,23 +1,17 @@
-import json
-import shutil
-import sys
+from typing import Optional
+import os, sys
 import time
 from datetime import datetime
-from typing import Optional
 from pathlib import Path
+
+import json
+import shutil
 import requests
 import itertools
 
 from art import text2art
 import inquirer
-import os
-import importlib.resources
-import importlib.util
-from importlib import import_module
 from cookiecutter.main import cookiecutter
-from dotenv import load_dotenv
-import subprocess
-
 
 from .agentstack_data import (
     FrameworkData,
@@ -158,50 +152,6 @@ def configure_default_model(path: Optional[str] = None):
         agentstack_config.default_model = model
 
 
-def run_project(path: Optional[str] = None, cli_args: Optional[str] = None):
-    """Validate that the project is ready to run and then run it."""
-    _path = Path(path) if path else Path.cwd()
-    framework = get_framework(_path)
-
-    if framework not in frameworks.SUPPORTED_FRAMEWORKS:
-        print(term_color(f"Framework {framework} is not supported by agentstack.", 'red'))
-        sys.exit(1)
-
-    try:
-        frameworks.validate_project(framework, _path)
-    except frameworks.ValidationError as e:
-        print(term_color("Project validation failed:", 'red'))
-        print(e)
-        sys.exit(1)
-
-    # Parse extra --input-* arguments for runtime overrides of the project's inputs
-    if cli_args:
-        for arg in cli_args:
-            if not arg.startswith('--input-'):
-                continue
-            key, value = arg[len('--input-') :].split('=')
-            inputs.add_input_for_run(key, value)
-
-    # explicitly load the project's .env file
-    load_dotenv(_path / '.env')
-
-    # import src/main.py from the project path
-    try:
-        spec = importlib.util.spec_from_file_location("main", str(_path / "src/main.py"))
-        assert spec is not None  # appease type checker
-        assert spec.loader is not None  # appease type checker
-        project_entrypoint = importlib.util.module_from_spec(spec)
-        sys.path.append(str(_path / "src"))
-        spec.loader.exec_module(project_entrypoint)
-    except ImportError as e:
-        print(term_color(f"Failed to import project. Does 'src/main.py' exist?:\n{e}", 'red'))
-        sys.exit(1)
-
-    # run the project's main function
-    # TODO try/except this and print detailed information with a --debug flag
-    project_entrypoint.run()
-
-
 def ask_framework() -> str:
     framework = "CrewAI"
     # framework = inquirer.list_input(
diff --git a/agentstack/cli/run.py b/agentstack/cli/run.py
new file mode 100644
index 00000000..7e3eef4f
--- /dev/null
+++ b/agentstack/cli/run.py
@@ -0,0 +1,68 @@
+from typing import Optional
+import os, sys
+from pathlib import Path
+import importlib.resources
+import importlib.util
+from importlib import import_module
+from cookiecutter.main import cookiecutter
+from dotenv import load_dotenv
+
+from agentstack import ValidationError
+from agentstack import inputs
+from agentstack import frameworks
+from agentstack.utils import term_color, get_framework
+
+MAIN_FILENAME: Path = Path("src/main.py")
+MAIN_MODULE_NAME = "main"
+
+
+def _import_project_module(path: Path):
+    """Import `main` from the project path."""
+    spec = importlib.util.spec_from_file_location(MAIN_MODULE_NAME, str(path / MAIN_FILENAME))
+
+    assert spec is not None  # appease type checker
+    assert spec.loader is not None  # appease type checker
+
+    project_module = importlib.util.module_from_spec(spec)
+    sys.path.append(str((path / MAIN_FILENAME).parent))
+    spec.loader.exec_module(project_module)
+    return project_module
+
+
+def run_project(command: str = 'run', path: Optional[str] = None, cli_args: Optional[str] = None):
+    """Validate that the project is ready to run and then run it."""
+    _path = Path(path) if path else Path.cwd()
+    framework = get_framework(_path)
+
+    if framework not in frameworks.SUPPORTED_FRAMEWORKS:
+        print(term_color(f"Framework {framework} is not supported by agentstack.", 'red'))
+        sys.exit(1)
+
+    try:
+        frameworks.validate_project(framework, _path)
+    except ValidationError as e:
+        print(term_color("Project validation failed:", 'red'))
+        print(e)
+        sys.exit(1)
+
+    # Parse extra --input-* arguments for runtime overrides of the project's inputs
+    if cli_args:
+        for arg in cli_args:
+            if not arg.startswith('--input-'):
+                continue
+            key, value = arg[len('--input-') :].split('=')
+            inputs.add_input_for_run(key, value)
+
+    # explicitly load the project's .env file
+    load_dotenv(_path / '.env')
+
+    # import src/main.py from the project path
+    try:
+        project_main = _import_project_module(_path)
+    except ImportError as e:
+        print(term_color(f"Failed to import project. Does '{MAIN_FILENAME}' exist?:\n{e}", 'red'))
+        sys.exit(1)
+
+    # run `command` from the project's main.py
+    # TODO try/except this and print detailed information with a --debug flag
+    return getattr(project_main, command)()
diff --git a/agentstack/main.py b/agentstack/main.py
index 86fdb09d..952118c1 100644
--- a/agentstack/main.py
+++ b/agentstack/main.py
@@ -55,7 +55,17 @@ def main():
     ''',
     )
     run_parser.add_argument(
-        "--path", "-p", help="Path to the project directory, defaults to current directory"
+        "--function",
+        "-f",
+        help="Function to call in main.py, defaults to 'run'",
+        default="run",
+        dest="function",
+    )
+    run_parser.add_argument(
+        "--path",
+        "-p",
+        help="Path to the project directory, defaults to current working directory",
+        dest="path",
     )
 
     # 'generate' command
@@ -128,7 +138,7 @@ def main():
     elif args.command in ["init", "i"]:
         init_project_builder(args.slug_name, args.template, args.wizard)
     elif args.command in ["run", "r"]:
-        run_project(args.path, cli_args=extra_args)
+        run_project(command=args.function, path=args.path, cli_args=extra_args)
     elif args.command in ['generate', 'g']:
         if args.generate_command in ['agent', 'a']:
             if not args.llm:
diff --git a/tests/test_project_run.py b/tests/test_project_run.py
index fec70f55..0740aad4 100644
--- a/tests/test_project_run.py
+++ b/tests/test_project_run.py
@@ -39,11 +39,11 @@ def tearDown(self):
         shutil.rmtree(self.project_dir)
 
     def test_run_project(self):
-        run_project(self.project_dir)
+        run_project(path=self.project_dir)
 
     def test_env_is_set(self):
         """
         After running a project, the environment variables should be set from project_dir/.env.
         """
-        run_project(self.project_dir)
+        run_project(path=self.project_dir)
         assert os.getenv('ENV_VAR1') == 'value1'

From 952a51c81511ea2c735ab0e1a569097f77a10de3 Mon Sep 17 00:00:00 2001
From: Travis Dent <tcdent@gmail.com>
Date: Wed, 11 Dec 2024 20:41:53 -0800
Subject: [PATCH 5/5] Cleanup imports, docstrings

---
 agentstack/cli/cli.py |  3 ---
 agentstack/cli/run.py | 15 ++++++++-------
 agentstack/inputs.py  |  1 -
 agentstack/main.py    |  3 +--
 4 files changed, 9 insertions(+), 13 deletions(-)

diff --git a/agentstack/cli/cli.py b/agentstack/cli/cli.py
index d33a85f3..89baceb3 100644
--- a/agentstack/cli/cli.py
+++ b/agentstack/cli/cli.py
@@ -6,7 +6,6 @@
 
 import json
 import shutil
-import requests
 import itertools
 
 from art import text2art
@@ -24,13 +23,11 @@
 from agentstack.tools import get_all_tools
 from agentstack.generation.files import ConfigFile, ProjectFile
 from agentstack import frameworks
-from agentstack import packaging
 from agentstack import generation
 from agentstack import inputs
 from agentstack.agents import get_all_agents
 from agentstack.tasks import get_all_tasks
 from agentstack.utils import open_json_file, term_color, is_snake_case, get_framework
-from agentstack.update import AGENTSTACK_PACKAGE
 from agentstack.proj_templates import TemplateConfig
 
 
diff --git a/agentstack/cli/run.py b/agentstack/cli/run.py
index 7e3eef4f..9f38f74d 100644
--- a/agentstack/cli/run.py
+++ b/agentstack/cli/run.py
@@ -1,10 +1,7 @@
 from typing import Optional
-import os, sys
+import sys
 from pathlib import Path
-import importlib.resources
 import importlib.util
-from importlib import import_module
-from cookiecutter.main import cookiecutter
 from dotenv import load_dotenv
 
 from agentstack import ValidationError
@@ -17,7 +14,12 @@
 
 
 def _import_project_module(path: Path):
-    """Import `main` from the project path."""
+    """
+    Import `main` from the project path.
+
+    We do it this way instead of spawning a subprocess so that we can share
+    state with the user's project.
+    """
     spec = importlib.util.spec_from_file_location(MAIN_MODULE_NAME, str(path / MAIN_FILENAME))
 
     assert spec is not None  # appease type checker
@@ -41,8 +43,7 @@ def run_project(command: str = 'run', path: Optional[str] = None, cli_args: Opti
     try:
         frameworks.validate_project(framework, _path)
     except ValidationError as e:
-        print(term_color("Project validation failed:", 'red'))
-        print(e)
+        print(term_color(f"Project validation failed:\n{e}", 'red'))
         sys.exit(1)
 
     # Parse extra --input-* arguments for runtime overrides of the project's inputs
diff --git a/agentstack/inputs.py b/agentstack/inputs.py
index b969b28c..209d5a52 100644
--- a/agentstack/inputs.py
+++ b/agentstack/inputs.py
@@ -1,7 +1,6 @@
 from typing import Optional
 import os
 from pathlib import Path
-import pydantic
 from ruamel.yaml import YAML, YAMLError
 from ruamel.yaml.scalarstring import FoldedScalarString
 from agentstack import ValidationError
diff --git a/agentstack/main.py b/agentstack/main.py
index 952118c1..e5e004f8 100644
--- a/agentstack/main.py
+++ b/agentstack/main.py
@@ -1,5 +1,4 @@
 import argparse
-import os
 import sys
 
 from agentstack.cli import (
@@ -10,7 +9,7 @@
     export_template,
 )
 from agentstack.telemetry import track_cli_command
-from agentstack.utils import get_version, get_framework
+from agentstack.utils import get_version
 from agentstack import generation
 from agentstack.update import check_for_updates