264 lines
9.8 KiB
Python
264 lines
9.8 KiB
Python
# SPDX-FileCopyrightText: 2026 Art Sea LLC
|
|
#
|
|
# SPDX-License-Identifier: AGPL-3.0-only
|
|
|
|
"""First-Aider: decompose a large prompt into sequential Band-Aider tasks."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from band_aider.ai_client.completion_client import CompletionClient
|
|
from band_aider.ai_client import json_completion
|
|
from band_aider.ai_client.openai_compatible_client import OpenAICompatibleClient
|
|
from band_aider.env_aliases import FIRST_AIDER_PROVIDER_ALIAS_FAMILIES, resolve_provider_alias_family
|
|
from band_aider.judgement.prompt_formatter import (
|
|
MarkdownPromptFormatterBuilder,
|
|
PromptFormatterBuilder,
|
|
format_prompt,
|
|
)
|
|
from band_aider.judgement.repo_exec import DockerContainerRepoExec, HostRepoExec
|
|
|
|
__all__ = ["build_parser", "decompose", "main"]
|
|
|
|
|
|
def build_parser() -> argparse.ArgumentParser:
|
|
"""Build the first-aider command parser."""
|
|
|
|
parser = argparse.ArgumentParser(prog="first-aider")
|
|
parser.add_argument("prompt", nargs="?", help="Larger user prompt")
|
|
parser.add_argument("--task", help="Larger user prompt")
|
|
parser.add_argument("--repo-root")
|
|
parser.add_argument("--repo-container")
|
|
parser.add_argument("--repo-python-bin")
|
|
parser.add_argument("--src-root", default=".")
|
|
parser.add_argument("--tests-root", default="tests")
|
|
parser.add_argument("--dry-run", action="store_true")
|
|
parser.add_argument("--strict-decompose", action="store_true")
|
|
parser.add_argument("--model")
|
|
parser.add_argument("--model-api-base")
|
|
parser.add_argument("--model-api-key")
|
|
parser.add_argument("--aider-cmd")
|
|
parser.add_argument("--aider-wrapper-mode", nargs="?", const="true")
|
|
parser.add_argument("--aider-model")
|
|
parser.add_argument("--aider-api-base")
|
|
parser.add_argument("--aider-api-key")
|
|
parser.add_argument("--aider-edit-format")
|
|
parser.add_argument("--aider-extra-arg", action="append", default=[])
|
|
return parser
|
|
|
|
|
|
def _client_from_env(args: argparse.Namespace) -> OpenAICompatibleClient:
|
|
"""Build the decomposition model client from CLI arguments and environment."""
|
|
|
|
resolved = resolve_provider_alias_family(os.environ, FIRST_AIDER_PROVIDER_ALIAS_FAMILIES)
|
|
return OpenAICompatibleClient(
|
|
model=args.model or (resolved.model if resolved is not None else None) or "default",
|
|
api_base=(
|
|
args.model_api_base
|
|
or (resolved.api_base if resolved is not None else None)
|
|
or "https://iona.local.arpeggio.one/api"
|
|
),
|
|
api_key=(
|
|
args.model_api_key
|
|
or (resolved.api_key if resolved is not None else None)
|
|
),
|
|
)
|
|
|
|
|
|
def _model_is_configured(client: OpenAICompatibleClient) -> bool:
|
|
"""Return whether the client has enough provider configuration to run."""
|
|
|
|
return bool(client.model and client.api_base and client.api_key)
|
|
|
|
|
|
def _envelope_prompt(
|
|
prompt: str,
|
|
*,
|
|
stage: str,
|
|
end_of_batch: bool,
|
|
prompt_formatter_builder: PromptFormatterBuilder,
|
|
) -> str:
|
|
"""Return one explicitly sectioned decomposition request."""
|
|
|
|
return format_prompt(
|
|
prompt_formatter_builder,
|
|
"FIRST-AIDER DECOMPOSITION",
|
|
lambda: "\n\n".join((
|
|
format_prompt(
|
|
prompt_formatter_builder,
|
|
"STAGE",
|
|
f"Stage: {stage}\nEnd-Of-Batch: {'true' if end_of_batch else 'false'}",
|
|
depth=2,
|
|
),
|
|
format_prompt(
|
|
prompt_formatter_builder,
|
|
"TASK",
|
|
prompt,
|
|
depth=2,
|
|
),
|
|
)),
|
|
escape=False,
|
|
)
|
|
|
|
|
|
def decompose(
|
|
model: CompletionClient,
|
|
prompt: str,
|
|
*,
|
|
prompt_formatter_builder: PromptFormatterBuilder,
|
|
) -> list[dict[str, Any]] | None:
|
|
"""Ask the model to decompose a large prompt into Band-Aider tasks."""
|
|
|
|
system_prompt = (
|
|
"Decompose the following large task into a sequence of smaller, independent, and actionable Band-Aider tasks.\n"
|
|
"Return strict JSON with a 'tasks' array.\n"
|
|
"Each task should have a 'title', 'task', and 'reason'.\n"
|
|
"Rules:\n"
|
|
"- If the prompt is already one cohesive change, return exactly one task using the original prompt.\n"
|
|
"- Split only when there are independent axes of work.\n"
|
|
"- Prefer 1 to 5 tasks.\n"
|
|
"- Keep each task independently actionable.\n"
|
|
"- Preserve ordering when later tasks depend on earlier tasks.\n"
|
|
"- Do not create microscopic tasks for style, formatting, or trivia.\n"
|
|
"- Do not invent requirements not present in the original prompt.\n"
|
|
"- Each generated task should be suitable to pass directly to band-aider --task.\n"
|
|
)
|
|
|
|
try:
|
|
data = json_completion.load_json_completion(
|
|
model,
|
|
_envelope_prompt(
|
|
f"{system_prompt}\nTask: {prompt}",
|
|
stage="global_decompose",
|
|
end_of_batch=False,
|
|
prompt_formatter_builder=prompt_formatter_builder,
|
|
),
|
|
prompt_formatter_builder=prompt_formatter_builder,
|
|
references=[{"tasks": [{"title": "", "task": "", "reason": ""}]}],
|
|
validate=_validate_tasks,
|
|
)
|
|
except Exception as error:
|
|
print(f"Warning: Decomposition failed: {error}", file=sys.stderr)
|
|
return None
|
|
tasks = data.get("tasks")
|
|
return tasks if isinstance(tasks, list) else None
|
|
|
|
|
|
def _validate_tasks(data: dict) -> None:
|
|
"""Require actionable decomposition before launching child tasks."""
|
|
# Decomposition cannot consume the original task without assigning any work.
|
|
if not data["tasks"]:
|
|
raise ValueError("tasks must contain at least one task; retain the original task when no split is needed.")
|
|
# Each child needs an executable task and an explanation of its boundary.
|
|
for task in data["tasks"]:
|
|
# Empty required strings cannot become child work.
|
|
if not all(value.strip() for value in task.values()):
|
|
raise ValueError("Every task needs nonempty title, task and reason strings.")
|
|
|
|
|
|
def main(argv: list[str] | None = None) -> None:
|
|
"""Run first-aider from CLI arguments."""
|
|
|
|
parser = build_parser()
|
|
args = parser.parse_args(argv)
|
|
|
|
if args.prompt and args.task:
|
|
print("Error: Both positional prompt and --task provided.", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
prompt = args.prompt or args.task
|
|
if not prompt:
|
|
parser.print_help()
|
|
sys.exit(1)
|
|
|
|
model = _client_from_env(args)
|
|
prompt_formatter_builder = MarkdownPromptFormatterBuilder()
|
|
tasks: list[dict[str, Any]] | None = None
|
|
if _model_is_configured(model):
|
|
tasks = decompose(
|
|
model,
|
|
prompt,
|
|
prompt_formatter_builder=prompt_formatter_builder,
|
|
)
|
|
|
|
if tasks is None:
|
|
if args.strict_decompose:
|
|
print("Error: Decomposition failed in strict mode.", file=sys.stderr)
|
|
sys.exit(1)
|
|
if _model_is_configured(model):
|
|
print("Warning: Decomposition failed, falling back to original prompt.")
|
|
tasks = [{"title": "Original Task", "task": prompt, "reason": "Decomposition failed or model not configured"}]
|
|
|
|
if args.dry_run:
|
|
print(f"Decomposed tasks for: {prompt}")
|
|
for index, task in enumerate(tasks, 1):
|
|
print(f"{index}. {task.get('title')}: {task.get('task')}")
|
|
sys.exit(0)
|
|
|
|
for index, task in enumerate(tasks, 1):
|
|
print(f"\n--- Task {index}: {task.get('title')} ---")
|
|
print(f"Goal: {task.get('task')}")
|
|
cmd = ["band-aider", "--task", str(task.get("task") or "")]
|
|
if args.repo_root is not None:
|
|
cmd.extend(["--repo-root", args.repo_root])
|
|
if args.repo_container is not None:
|
|
cmd.extend(["--repo-container", args.repo_container])
|
|
if args.repo_python_bin is not None:
|
|
cmd.extend(["--repo-python-bin", args.repo_python_bin])
|
|
cmd.extend(["--src-root", args.src_root])
|
|
cmd.extend(["--tests-root", args.tests_root])
|
|
if args.model:
|
|
cmd.extend(["--model", args.model])
|
|
if args.model_api_base:
|
|
cmd.extend(["--model-api-base", args.model_api_base])
|
|
if args.model_api_key:
|
|
cmd.extend(["--model-api-key", args.model_api_key])
|
|
if args.aider_cmd:
|
|
cmd.extend(["--aider-cmd", args.aider_cmd])
|
|
if args.aider_wrapper_mode:
|
|
cmd.extend(["--aider-wrapper-mode", args.aider_wrapper_mode])
|
|
if args.aider_model:
|
|
cmd.extend(["--aider-model", args.aider_model])
|
|
if args.aider_api_base:
|
|
cmd.extend(["--aider-api-base", args.aider_api_base])
|
|
if args.aider_api_key:
|
|
cmd.extend(["--aider-api-key", args.aider_api_key])
|
|
if args.aider_edit_format:
|
|
cmd.extend(["--aider-edit-format", args.aider_edit_format])
|
|
for extra_arg in args.aider_extra_arg:
|
|
cmd.extend(["--aider-extra-arg", extra_arg])
|
|
|
|
if args.repo_root is None:
|
|
print("Error: --repo-root is required for project task execution.", file=sys.stderr)
|
|
sys.exit(2)
|
|
repo_exec = (
|
|
DockerContainerRepoExec(
|
|
container=args.repo_container,
|
|
repo_root=Path(args.repo_root),
|
|
python_bin=args.repo_python_bin,
|
|
)
|
|
if args.repo_container
|
|
else HostRepoExec(Path(args.repo_root))
|
|
)
|
|
completed = repo_exec.run_cli(
|
|
cmd,
|
|
cwd=Path(args.repo_root),
|
|
capture_output=False,
|
|
text=True,
|
|
check=False,
|
|
)
|
|
returncode = completed.returncode
|
|
if returncode != 0:
|
|
print(f"Error: Task {index} failed with exit code {returncode}", file=sys.stderr)
|
|
sys.exit(returncode)
|
|
|
|
sys.exit(0)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|