modification.py

import os
import sys
import json
import subprocess
import shlex
from file_utils import move_file, get_file_contents, update_file
from custom_print import custom_print, pretty_obj, print_chat_messages
import llm


openai_api_key = os.getenv("OPENAI_API_KEY")
if not openai_api_key:
    print("Set OPENAI_API_KEY in your environment.")
    sys.exit(1)


def process_file_item(item, project_dir):
    fn, act = item['filename'], item['action']
    orig = os.path.join(project_dir, fn)
    if act in ('modify','reference'):
        if item.get('new_filepath'):
            newpath = os.path.join(project_dir, item['new_filepath'])
            move_file(orig, newpath)
            path, key = newpath, item['new_filepath']
        else:
            path, key = orig, fn
        print(f"Reading file: {path}")
        return key, get_file_contents(path) if os.path.exists(path) else ""
    elif act == 'create':
        return fn, ""


def get_files_data(final_plan, project_dir):
    data = {}
    for item in final_plan:
        if 'filename' not in item: continue
        key, content = process_file_item(item, project_dir)
        data[key] = content
    return data


def build_messages(mod_prompt, final_plan, full_text, concise_text):
    user_template = (
        f"User modification prompt: {mod_prompt}\n\n"
        f"Plan:\n{json.dumps(final_plan, indent=2)}\n\n"
        "Here are the current contents of the files to be modified/referenced/created:\n"
        "```$files_text```\n\n"
        "Please output only JSON where each key is a filename and its value is the full new contents for that file."
    )
    system_msg = ("You are an expert software engineer tasked with modifying a multi-file project. "
                  "Please write elegant and modular clean code. Keep functions under 10 lines, files under 100 lines, "
                  "and you can refactor existing code if necessary.")
    api_messages = [
        {"role": "system", "content": system_msg},
        {"role": "user", "content": user_template.replace("$files_text", full_text)}
    ]
    console_messages = [
        {"role": "system", "content": system_msg},
        {"role": "user", "content": user_template.replace("$files_text", concise_text)}
    ]
    return api_messages, console_messages


def parse_response(response):
    try:
        return json.loads(response.choices[0].message.content.strip())
    except json.JSONDecodeError:
        print("Failed to parse modifications JSON. Response was:")
        print(response.choices[0].message.content)
        sys.exit(1)


def modify_files(mod_prompt, final_plan, project_dir):
    files_data = get_files_data(final_plan, project_dir)
    full_text = "\n".join([f"{fname}:\n```{content}```" for fname, content in files_data.items()])
    concise = ", ".join(files_data.keys())
    api_msgs, console_msgs = build_messages(mod_prompt, final_plan, full_text, concise)
    print_chat_messages(console_msgs)
    print("🧠 Thinking")
    response = llm.create_completion(model="o3-mini", messages=api_msgs, reasoning_effort='high')
    return parse_response(response)


def validate_output_multi(mod_prompt, entry_file, compile_command, project_dir):
    cmd = shlex.split(compile_command) + [project_dir]
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode == 0:
        return "Success: Project compiled without errors."
    else:
        return f"Failure: Compilation error:\n{result.stderr}"