Merge pull request agentstack-ai#122 from tcdent/inputs-file

Store inputs in project YAML file, accept CLI flag input data

Author: tcdent

agentstack/__init__.py

@@ -1,8 +1,14 @@
+"""
+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.exceptions import ValidationError
+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

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

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)

agentstack/cli/cli.py

@@ -1,21 +1,16 @@
-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 requests
+
+import json
+import shutil
 import itertools
 
 from art import text2art
 import inquirer
-import os
-import importlib.resources
 from cookiecutter.main import cookiecutter
-from dotenv import load_dotenv
-import subprocess
-from packaging.metadata import Metadata
 
 from .agentstack_data import (
     FrameworkData,
@@ -28,12 +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
 
 
@@ -162,27 +156,6 @@ def configure_default_model(path: Optional[str] = None):
         agentstack_config.default_model = model
 
 
-def run_project(framework: str, path: str = ''):
-    """Validate that the project is ready to run and then run it."""
-    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:
-        print(term_color("Project validation failed:", 'red'))
-        print(e)
-        sys.exit(1)
-
-    load_dotenv(Path.home() / '.env')  # load the user's .env file
-    load_dotenv(_path / '.env', override=True)  # load the project's .env file
-    print("Running your agent...")
-    subprocess.run(['python', 'src/main.py'], env=os.environ)
-
-
 def ask_framework() -> str:
     framework = "CrewAI"
     # framework = inquirer.list_input(
@@ -401,7 +374,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,
@@ -537,21 +510,16 @@ 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,
         method="sequential",  # TODO this needs to be stored in the project somewhere
         agents=agents,
         tasks=tasks,
         tools=tools,
-        inputs=inputs,
+        inputs=inputs.get_inputs(),
     )
 
     try:

agentstack/cli/run.py

@@ -0,0 +1,70 @@
+from typing import Optional
+import sys
+from pathlib import Path
+import importlib.util
+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.
+
+    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
+    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(f"Project validation failed:\n{e}", 'red'))
+        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)
+
+    load_dotenv(Path.home() / '.env')  # load the user's .env file
+    load_dotenv(_path / '.env', override=True)  # load the project's .env file
+
+    # 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
+    print("Running your agent...")
+    return getattr(project_main, command)()

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

agentstack/inputs.py

@@ -0,0 +1,92 @@
+from typing import Optional
+import os
+from pathlib import Path
+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
+
+# run_inputs are set at the beginning of the run and are not saved
+run_inputs: dict[str, str] = {}
+
+
+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:
+    """
+    Get the inputs configuration file and override with run_inputs.
+    """
+    path = path if path else Path()
+    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

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
 
@@ -43,7 +42,30 @@ 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(
+        "--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
     generate_parser = subparsers.add_parser("generate", aliases=["g"], help="Generate agents or tasks")
@@ -94,8 +116,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 +137,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(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: