Merge pull request #149 from tcdent/fancy-init

Migrate to `uv` for package/venv management in user projects

agentstack/cli/__init__.py

@@ -1,3 +1,4 @@
-from .cli import init_project_builder, configure_default_model, export_template
+from .cli import init_project_builder, configure_default_model, export_template, welcome_message
+from .init import init_project
 from .tools import list_tools, add_tool
 from .run import run_project

agentstack/cli/cli.py

@@ -86,15 +86,13 @@ def init_project_builder(
         tools = [tools.model_dump() for tools in template_data.tools]
 
     elif use_wizard:
-        welcome_message()
         project_details = ask_project_details(slug_name)
         welcome_message()
         framework = ask_framework()
         design = ask_design()
         tools = ask_tools()
 
     else:
-        welcome_message()
         # the user has started a new project; let's give them something to work with
         default_project = TemplateConfig.from_template_name('hello_alex')
         project_details = {
@@ -115,9 +113,6 @@ def init_project_builder(
     log.debug(f"project_details: {project_details}" f"framework: {framework}" f"design: {design}")
     insert_template(project_details, framework, design, template_data)
 
-    # we have an agentstack.json file in the directory now
-    conf.set_path(project_details['name'])
-
     for tool_data in tools:
         generation.add_tool(tool_data['name'], agents=tool_data['agents'])
 
@@ -410,14 +405,14 @@ def insert_template(
         f'{template_path}/{"{{cookiecutter.project_metadata.project_slug}}"}/.env',
     )
 
-    if os.path.isdir(project_details['name']):
-        print(
-            term_color(
-                f"Directory {template_path} already exists. Please check this and try again",
-                "red",
-            )
-        )
-        sys.exit(1)
+    # if os.path.isdir(project_details['name']):
+    #     print(
+    #         term_color(
+    #             f"Directory {template_path} already exists. Please check this and try again",
+    #             "red",
+    #         )
+    #     )
+    #     sys.exit(1)
 
     cookiecutter(str(template_path), no_input=True, extra_context=None)
 
@@ -431,26 +426,6 @@ def insert_template(
     except:
         print("Failed to initialize git repository. Maybe you're already in one? Do this with: git init")
 
-    # TODO: check if poetry is installed and if so, run poetry install in the new directory
-    # os.system("poetry install")
-    # os.system("cls" if os.name == "nt" else "clear")
-    # TODO: add `agentstack docs` command
-    print(
-        "\n"
-        "🚀 \033[92mAgentStack project generated successfully!\033[0m\n\n"
-        "  Next, run:\n"
-        f"    cd {project_metadata.project_slug}\n"
-        "    python -m venv .venv\n"
-        "    source .venv/bin/activate\n\n"
-        "  Make sure you have the latest version of poetry installed:\n"
-        "    pip install -U poetry\n\n"
-        "  You'll need to install the project's dependencies with:\n"
-        "    poetry install\n\n"
-        "  Finally, try running your agent with:\n"
-        "    agentstack run\n\n"
-        "  Run `agentstack quickstart` or `agentstack docs` for next steps.\n"
-    )
-
 
 def export_template(output_filename: str):
     """

agentstack/cli/init.py

@@ -0,0 +1,72 @@
+import os, sys
+from typing import Optional
+from pathlib import Path
+from agentstack import conf
+from agentstack import packaging
+from agentstack.cli import welcome_message, init_project_builder
+from agentstack.utils import term_color
+
+
+# TODO move the rest of the CLI init tooling into this file
+
+
+def require_uv():
+    try:
+        uv_bin = packaging.get_uv_bin()
+        assert os.path.exists(uv_bin)
+    except (AssertionError, ImportError):
+        print(term_color("Error: uv is not installed.", 'red'))
+        print("Full installation instructions at: https://docs.astral.sh/uv/getting-started/installation")
+        match sys.platform:
+            case 'linux' | 'darwin':
+                print("Hint: run `curl -LsSf https://astral.sh/uv/install.sh | sh`")
+            case _:
+                pass
+        sys.exit(1)
+
+
+def init_project(
+    slug_name: Optional[str] = None,
+    template: Optional[str] = None,
+    use_wizard: bool = False,
+):
+    """
+    Initialize a new project in the current directory.
+
+    - create a new virtual environment
+    - copy project skeleton
+    - install dependencies
+    """
+    require_uv()
+
+    # TODO prevent the user from passing the --path arguent to init
+    if slug_name:
+        conf.set_path(conf.PATH / slug_name)
+    else:
+        print("Error: No project directory specified.")
+        print("Run `agentstack init <project_name>`")
+        sys.exit(1)
+
+    if os.path.exists(conf.PATH):  # cookiecutter requires the directory to not exist
+        print(f"Error: Directory already exists: {conf.PATH}")
+        sys.exit(1)
+
+    welcome_message()
+    print(term_color("🦾 Creating a new AgentStack project...", 'blue'))
+    print(f"Using project directory: {conf.PATH.absolute()}")
+
+    # copy the project skeleton, create a virtual environment, and install dependencies
+    init_project_builder(slug_name, template, use_wizard)
+    packaging.create_venv()
+    packaging.install_project()
+
+    print(
+        "\n"
+        "🚀 \033[92mAgentStack project generated successfully!\033[0m\n\n"
+        "  To get started, activate the virtual environment with:\n"
+        f"    cd {conf.PATH}\n"
+        "    source .venv/bin/activate\n\n"
+        "  Run your new agent with:\n"
+        "    agentstack run\n\n"
+        "  Or, run `agentstack quickstart` or `agentstack docs` for more next steps.\n"
+    )

agentstack/exceptions.py

@@ -5,3 +5,12 @@ class ValidationError(Exception):
     """
 
     pass
+
+
+class EnvironmentError(Exception):
+    """
+    Raised when an error occurs in the execution environment ie. a command is
+    not present or the environment is not configured as expected.
+    """
+
+    pass

agentstack/main.py

@@ -4,7 +4,7 @@
 
 from agentstack import conf, auth
 from agentstack.cli import (
-    init_project_builder,
+    init_project,
     add_tool,
     list_tools,
     configure_default_model,
@@ -167,7 +167,7 @@ def main():
         elif args.command in ["templates"]:
             webbrowser.open("https://docs.agentstack.sh/quickstart")
         elif args.command in ["init", "i"]:
-            init_project_builder(args.slug_name, args.template, args.wizard)
+            init_project(args.slug_name, args.template, args.wizard)
         elif args.command in ["tools", "t"]:
             if args.tools_command in ["list", "l"]:
                 list_tools()

agentstack/packaging.py

@@ -1,18 +1,173 @@
-import os
-from typing import Optional
+import os, sys
+from typing import Optional, Callable
+from pathlib import Path
+import re
+import subprocess
+import select
+from agentstack import conf
 
-PACKAGING_CMD = "poetry"
 
+DEFAULT_PYTHON_VERSION = "3.12"
+VENV_DIR_NAME: Path = Path(".venv")
 
-def install(package: str, path: Optional[str] = None):
-    if path:
-        os.chdir(path)
-    os.system(f"{PACKAGING_CMD} add {package}")
+# filter uv output by these words to only show useful progress messages
+RE_UV_PROGRESS = re.compile(r'^(Resolved|Prepared|Installed|Uninstalled|Audited)')
+
+
+# When calling `uv` we explicitly specify the --python executable to use so that
+# the packages are installed into the correct virtual environment.
+# In testing, when this was not set, packages could end up in the pyenv's
+# site-packages directory; it's possible an environemnt variable can control this.
+
+
+def install(package: str):
+    """Install a package with `uv` and add it to pyproject.toml."""
+
+    def on_progress(line: str):
+        if RE_UV_PROGRESS.match(line):
+            print(line.strip())
+
+    def on_error(line: str):
+        print(f"uv: [error]\n {line.strip()}")
+
+    _wrap_command_with_callbacks(
+        [get_uv_bin(), 'add', '--python', '.venv/bin/python', package],
+        on_progress=on_progress,
+        on_error=on_error,
+    )
+
+
+def install_project():
+    """Install all dependencies for the user's project."""
+
+    def on_progress(line: str):
+        if RE_UV_PROGRESS.match(line):
+            print(line.strip())
+
+    def on_error(line: str):
+        print(f"uv: [error]\n {line.strip()}")
+
+    _wrap_command_with_callbacks(
+        [get_uv_bin(), 'pip', 'install', '--python', '.venv/bin/python', '.'],
+        on_progress=on_progress,
+        on_error=on_error,
+    )
 
 
 def remove(package: str):
-    os.system(f"{PACKAGING_CMD} remove {package}")
+    """Uninstall a package with `uv`."""
+
+    # TODO it may be worth considering removing unused sub-dependencies as well
+    def on_progress(line: str):
+        if RE_UV_PROGRESS.match(line):
+            print(line.strip())
+
+    def on_error(line: str):
+        print(f"uv: [error]\n {line.strip()}")
+
+    _wrap_command_with_callbacks(
+        [get_uv_bin(), 'remove', '--python', '.venv/bin/python', package],
+        on_progress=on_progress,
+        on_error=on_error,
+    )
 
 
 def upgrade(package: str):
-    os.system(f"{PACKAGING_CMD} add {package}")
+    """Upgrade a package with `uv`."""
+
+    # TODO should we try to update the project's pyproject.toml as well?
+    def on_progress(line: str):
+        if RE_UV_PROGRESS.match(line):
+            print(line.strip())
+
+    def on_error(line: str):
+        print(f"uv: [error]\n {line.strip()}")
+
+    _wrap_command_with_callbacks(
+        [get_uv_bin(), 'pip', 'install', '-U', '--python', '.venv/bin/python', package],
+        on_progress=on_progress,
+        on_error=on_error,
+    )
+
+
+def create_venv(python_version: str = DEFAULT_PYTHON_VERSION):
+    """Intialize a virtual environment in the project directory of one does not exist."""
+    if os.path.exists(conf.PATH / VENV_DIR_NAME):
+        return  # venv already exists
+
+    RE_VENV_PROGRESS = re.compile(r'^(Using|Creating)')
+
+    def on_progress(line: str):
+        if RE_VENV_PROGRESS.match(line):
+            print(line.strip())
+
+    def on_error(line: str):
+        print(f"uv: [error]\n {line.strip()}")
+
+    _wrap_command_with_callbacks(
+        [get_uv_bin(), 'venv', '--python', python_version],
+        on_progress=on_progress,
+        on_error=on_error,
+    )
+
+
+def get_uv_bin() -> str:
+    """Find the path to the uv binary."""
+    try:
+        import uv
+
+        return uv.find_uv_bin()
+    except ImportError as e:
+        raise e
+
+
+def _setup_env() -> dict[str, str]:
+    """Copy the current environment and add the virtual environment path for use by a subprocess."""
+    env = os.environ.copy()
+    env["VIRTUAL_ENV"] = str(conf.PATH / VENV_DIR_NAME.absolute())
+    env["UV_INTERNAL__PARENT_INTERPRETER"] = sys.executable
+    return env
+
+
+def _wrap_command_with_callbacks(
+    command: list[str],
+    on_progress: Callable[[str], None] = lambda x: None,
+    on_complete: Callable[[str], None] = lambda x: None,
+    on_error: Callable[[str], None] = lambda x: None,
+) -> None:
+    """Run a command with progress callbacks."""
+    try:
+        all_lines = ''
+        process = subprocess.Popen(
+            command,
+            cwd=conf.PATH.absolute(),
+            env=_setup_env(),
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+        )
+        assert process.stdout and process.stderr  # appease type checker
+
+        readable = [process.stdout, process.stderr]
+        while readable:
+            ready, _, _ = select.select(readable, [], [])
+            for fd in ready:
+                line = fd.readline()
+                if not line:
+                    readable.remove(fd)
+                    continue
+
+                on_progress(line)
+                all_lines += line
+
+        if process.wait() == 0:  # return code: success
+            on_complete(all_lines)
+        else:
+            on_error(all_lines)
+    except Exception as e:
+        on_error(str(e))
+    finally:
+        try:
+            process.terminate()
+        except:
+            pass

agentstack/telemetry.py

@@ -76,7 +76,7 @@ def collect_machine_telemetry(command: str):
 
 
 def track_cli_command(command: str, args: Optional[str] = None):
-    if bool(os.environ['AGENTSTATCK_IS_TEST_ENV']):
+    if bool(os.getenv('AGENTSTACK_IS_TEST_ENV')):
         return
 
     try:
@@ -91,7 +91,7 @@ def track_cli_command(command: str, args: Optional[str] = None):
         pass
 
 def update_telemetry(id: int, result: int, message: Optional[str] = None):
-    if bool(os.environ['AGENTSTATCK_IS_TEST_ENV']):
+    if bool(os.getenv('AGENTSTACK_IS_TEST_ENV')):
         return
 
     try: