coder/main_helpers.py at main · rednebmas/coder · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
import os
import argparse
from planner.plan_files import plan_files
from planner.refine_plan import refine_plan
from planner.refactor_plan import refactor_plan
from modification import modify_files, validate_output_multi
from summarization import generate_summary
from custom_print import custom_print
from notification import play_chime, show_notification
from file_utils import update_file


def parse_arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument('--project', default='.', help='Project directory')
    parser.add_argument('--prompt', required=True, help='Modification instructions')
    parser.add_argument('--entry', required=True, help='Entry point file (relative to project directory)')
    parser.add_argument('--max_attempts', type=int, default=1, help='Maximum modification attempts')
    parser.add_argument('--refactor', action='store_true', help='Run in refactor mode for structural reorganization (no logic modifications)')
    return parser.parse_args()


def extract_global_summary(plan):
    if isinstance(plan, list) and plan and isinstance(plan[0], dict) and "global_summary" in plan[0]:
        return plan[0]["global_summary"], plan[1:]
    return "", plan


def print_plan_summary(refined):
    global_sum, file_plans = extract_global_summary(refined)
    custom_print("⏩ Global Plan Summary:", global_sum, "\n")
    custom_print("File-specific Plans:", file_plans, "\n")
    play_chime()
    show_notification("Plan ready for verification!")
    return global_sum, file_plans


def build_included_files(plan, project):
    included = {}
    for item in plan:
        if item.get('action') in ('modify', 'reference'):
            path = os.path.join(project, item['filename'])
            included[item['filename']] = open(path, 'r').read() if os.path.exists(path) else ""
    return included


def interactive_refinement(prompt, included):
    refined = refine_plan(prompt, included)
    print_plan_summary(refined)
    show_notification("Plan ready for verification!")
    feedback = ''
    while True:
        most_recent_feedback = input("Choose an option: (1) continue (2) refine further: ").strip()
        feedback += '\n' + most_recent_feedback
        if most_recent_feedback in ("continue", "1", ""):
            print("")
            break
        updated_prompt = prompt + "\nAdditional feedback on the previous plan: " + feedback
        refined = refine_plan(updated_prompt, included, previous_plan=str(refined))
        print_plan_summary(refined)
    return refined


def update_file_content(filename, content, project):
    path = os.path.join(project, filename)
    update_file(path, content)
    print(f"Updated {filename}")


def update_file_summary(filename, project):
    path = os.path.join(project, filename)
    summary = generate_summary(path)
    update_file(path + ".summary", summary)


def process_modifications(prompt, refined_plan, project):
    mods = modify_files(prompt, refined_plan, project)
    custom_print("⏩ Modifications:", mods, "\n")
    for fname, content in mods.items():
        update_file_content(fname, content, project)
    return mods


def validate_project(prompt, entry, project, config):
    print(f"\nValidating project compilation using '{config['compile_command']}'...\n")
    result = validate_output_multi(prompt, entry, config['compile_command'], project)
    print(f"\n⏩ Validation result:\n{result}\n")
    if "Success" in result:
        print("Test passed.")
    else:
        print("Test failed")
    return result


def process_prompt(prompt, project, refactor=False, entry_file=None):
    if refactor:
        initial_plan = refactor_plan(prompt, entry_file=entry_file)
        custom_print("⏩ Initial refactor plan:", initial_plan, "\n")
    else:
        initial_plan = plan_files(prompt)
        custom_print("⏩ Initial plan:", initial_plan, "\n")
    included = build_included_files(initial_plan, project)
    refined = interactive_refinement(prompt, included)
    return refined