From 3615a3598f5c08e4d406845bcf42b1b714a8ea2d Mon Sep 17 00:00:00 2001 From: Elena Zherdeva <107525751+elenazherdeva@users.noreply.github.com> Date: Fri, 20 Sep 2024 09:23:10 -0700 Subject: [PATCH] feat: support markdown plans (#79) --- src/goose/cli/main.py | 23 ++++++++++++- src/goose/toolkit/utils.py | 37 ++++++++++++++++++++- tests/toolkit/test_utils.py | 65 +++++++++++++++++++++++++++++++++++++ 3 files changed, 123 insertions(+), 2 deletions(-) create mode 100644 tests/toolkit/test_utils.py diff --git a/src/goose/cli/main.py b/src/goose/cli/main.py index e4e11238..898ad2a6 100644 --- a/src/goose/cli/main.py +++ b/src/goose/cli/main.py @@ -8,6 +8,7 @@ from ruamel.yaml import YAML from goose.cli.config import SESSIONS_PATH from goose.cli.session import Session +from goose.toolkit.utils import render_template, parse_plan from goose.utils import load_plugins from goose.utils.session_file import list_sorted_session_files @@ -73,11 +74,31 @@ def session_start(profile: str, plan: Optional[str] = None) -> None: _plan = yaml.load(f) else: _plan = None - session = Session(profile=profile, plan=_plan) session.run() +def parse_args(ctx: click.Context, param: click.Parameter, value: str) -> dict[str, str]: + if not value: + return {} + args = {} + for item in value.split(","): + key, val = item.split(":") + args[key.strip()] = val.strip() + + return args + + +@session.command(name="planned") +@click.option("--plan", type=click.Path(exists=True)) +@click.option("-a", "--args", callback=parse_args, help="Args in the format arg1:value1,arg2:value2") +def session_planned(plan: str, args: Optional[dict[str, str]]) -> None: + plan_templated = render_template(Path(plan), context=args) + _plan = parse_plan(plan_templated) + session = Session(plan=_plan) + session.run() + + @session.command(name="resume") @click.argument("name", required=False) @click.option("--profile") diff --git a/src/goose/toolkit/utils.py b/src/goose/toolkit/utils.py index 61632b77..ad97360f 100644 --- a/src/goose/toolkit/utils.py +++ b/src/goose/toolkit/utils.py @@ -1,5 +1,5 @@ from pathlib import Path -from typing import Optional +from typing import Optional, Dict from pygments.lexers import get_lexer_for_filename from pygments.util import ClassNotFound @@ -42,3 +42,38 @@ def render_template(template_path: Path, context: Optional[dict] = None) -> str: env = Environment(loader=FileSystemLoader(template_path.parent)) template = env.get_template(template_path.name) return template.render(context or {}) + + +def find_last_task_group_index(input_str: str) -> int: + lines = input_str.splitlines() + last_group_start_index = -1 + current_group_start_index = -1 + + for i, line in enumerate(lines): + line = line.strip() + if line.startswith("-"): + # If this is the first line of a new group, mark its start + if current_group_start_index == -1: + current_group_start_index = i + else: + # If we encounter a non-hyphenated line and had a group, update last group start + if current_group_start_index != -1: + last_group_start_index = current_group_start_index + current_group_start_index = -1 # Reset for potential future groups + + # If the input ended in a task group, update the last group index + if current_group_start_index != -1: + last_group_start_index = current_group_start_index + return last_group_start_index + + +def parse_plan(input_plan_str: str) -> Dict: + last_group_start_index = find_last_task_group_index(input_plan_str) + if last_group_start_index == -1: + return {"kickoff_message": input_plan_str, "tasks": []} + + kickoff_message_list = input_plan_str.splitlines()[:last_group_start_index] + kickoff_message = "\n".join(kickoff_message_list).strip() + tasks_list = input_plan_str.splitlines()[last_group_start_index:] + tasks_list_output = [s[1:] for s in tasks_list if s.strip()] # filter leading - + return {"kickoff_message": kickoff_message, "tasks": tasks_list_output} diff --git a/tests/toolkit/test_utils.py b/tests/toolkit/test_utils.py new file mode 100644 index 00000000..b5b45ac9 --- /dev/null +++ b/tests/toolkit/test_utils.py @@ -0,0 +1,65 @@ +from goose.toolkit.utils import parse_plan + + +def test_parse_plan_simple(): + plan_str = ( + "Here is python repo\n" + "-use uv\n" + "-do not use poetry\n\n" + "Now you should:\n\n" + "-Open a file\n" + "-Run a test" + ) + expected_result = { + "kickoff_message": "Here is python repo\n-use uv\n-do not use poetry\n\nNow you should:", + "tasks": ["Open a file", "Run a test"], + } + assert expected_result == parse_plan(plan_str) + + +def test_parse_plan_multiple_groups(): + plan_str = ( + "Here is python repo\n" + "-use uv\n" + "-do not use poetry\n\n" + "Now you should:\n\n" + "-Open a file\n" + "-Run a test\n\n" + "Now actually follow the steps:\n" + "-Step1\n" + "-Step2" + ) + expected_result = { + "kickoff_message": ( + "Here is python repo\n" + "-use uv\n" + "-do not use poetry\n\n" + "Now you should:\n\n" + "-Open a file\n" + "-Run a test\n\n" + "Now actually follow the steps:" + ), + "tasks": ["Step1", "Step2"], + } + assert expected_result == parse_plan(plan_str) + + +def test_parse_plan_empty_tasks(): + plan_str = "Here is python repo" + expected_result = {"kickoff_message": "Here is python repo", "tasks": []} + assert expected_result == parse_plan(plan_str) + + +def test_parse_plan_empty_kickoff_message(): + plan_str = "-task1\n-task2" + expected_result = {"kickoff_message": "", "tasks": ["task1", "task2"]} + assert expected_result == parse_plan(plan_str) + + +def test_parse_plan_with_numbers(): + plan_str = "Here is python repo\n" "Now you should:\n\n" "-1 Open a file\n" "-2 Run a test" + expected_result = { + "kickoff_message": "Here is python repo\nNow you should:", + "tasks": ["1 Open a file", "2 Run a test"], + } + assert expected_result == parse_plan(plan_str)