Skip to content

Commit

Permalink
feat: Support markdown plans
Browse files Browse the repository at this point in the history
  • Loading branch information
elenazherdeva committed Sep 19, 2024
1 parent 71863a1 commit 652a240
Show file tree
Hide file tree
Showing 3 changed files with 113 additions and 2 deletions.
23 changes: 22 additions & 1 deletion src/goose/cli/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@

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

Expand Down Expand Up @@ -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_params(ctx: click.Context, param: click.Parameter, value: str) -> Dict[str, str]:
if not value:
return {}
params = {}
for item in value.split(","):
key, val = item.split(":")
params[key.strip()] = val.strip()

return params


@session.command(name="planned")
@click.option("--plan", type=click.Path(exists=True))
@click.option("-p", "--params", callback=parse_params, help="Parameters in the format param1:value1,param2:value2")
def session_planned(plan: str, params: Optional[Dict[str, str]]) -> None:
plan_templated = render_template(Path(plan), context=params)
_plan = parse_plan(plan_templated)
session = Session(plan=_plan)
session.run()


@session.command(name="resume")
@click.argument("name", required=False)
@click.option("--profile")
Expand Down
36 changes: 35 additions & 1 deletion src/goose/toolkit/utils.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -42,3 +42,37 @@ 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:]
return {"kickoff_message": kickoff_message, "tasks": list(filter(str.strip, tasks_list))}
56 changes: 56 additions & 0 deletions tests/toolkit/test_utils.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,56 @@
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)

0 comments on commit 652a240

Please sign in to comment.