MABELab-RS/porting/tools/translate_one.py

118 lines
4.0 KiB
Python

#!/usr/bin/env python3
import argparse, os, sys, subprocess, yaml
from llm_clients import LLMClient
from prompt_inject import render_template
from chunker import chunk_text
def read(path):
with open(path, "r", encoding="utf-8") as f:
return f.read()
def write(path, content):
os.makedirs(os.path.dirname(path), exist_ok=True)
with open(path, "w", encoding="utf-8") as f:
f.write(content)
def load_cfg():
cfg_path = "porting/CONFIG.yml"
if not os.path.exists(cfg_path):
cfg_path = "porting/CONFIG.example.yml"
with open(cfg_path, "r", encoding="utf-8") as f:
return yaml.safe_load(f)
def main():
p = argparse.ArgumentParser()
p.add_argument("--mode", required=True, choices=["header-to-traits", "impl-pass", "unit-tests", "review-fixit"])
p.add_argument("--input", required=True, help="Path to C++ source (or error log for review-fixit)")
p.add_argument("--skeleton", default=None, help="Existing Rust module to guide/patch")
p.add_argument("--out", required=True, help="Output Rust file (or tests module for unit-tests)")
p.add_argument("--test-spec", default=None, help="Text spec or path to spec")
args = p.parse_args()
cfg = load_cfg()
client = LLMClient(cfg, cfg["logs"]["dir"])
glossary = read(cfg["project"]["glossary_path"])
style = read(cfg["project"]["style_path"])
determinism = read(cfg["project"]["determinism_path"])
def load_prompt(name):
return read(cfg["prompts"][name])
system = ""
if args.mode == "header-to-traits":
tmpl = load_prompt("header_to_traits")
src = read(args.input)
skel = read(args.skeleton) if args.skeleton and os.path.exists(args.skeleton) else ""
user = render_template(
tmpl,
GLOSSARY=glossary,
STYLE=style,
DETERMINISM=determinism,
SKELETON=skel,
SOURCE_CHUNK=src,
)
resp = client.chat(system, user)
write(args.out, resp)
elif args.mode == "impl-pass":
tmpl = load_prompt("impl_pass")
src = read(args.input)
skel = read(args.skeleton) if args.skeleton and os.path.exists(args.skeleton) else ""
test_spec = read(args.test_spec) if args.test_spec and os.path.exists(args.test_spec) else (args.test_spec or "")
user = render_template(
tmpl,
GLOSSARY=glossary,
STYLE=style,
DETERMINISM=determinism,
SKELETON=skel,
SOURCE_CHUNK=src,
TEST_SPEC=test_spec,
)
resp = client.chat(system, user)
write(args.out, resp)
elif args.mode == "unit-tests":
tmpl = load_prompt("unit_tests")
skel = read(args.skeleton) if args.skeleton and os.path.exists(args.skeleton) else ""
test_spec = read(args.test_spec) if args.test_spec and os.path.exists(args.test_spec) else (args.test_spec or "")
user = render_template(
tmpl,
DETERMINISM=determinism,
SKELETON=skel,
TEST_SPEC=test_spec,
)
resp = client.chat(system, user)
# Append or create tests in the same file:
if os.path.exists(args.out):
write(args.out, read(args.out) + "\n\n" + resp)
else:
write(args.out, resp)
elif args.mode == "review-fixit":
tmpl = load_prompt("review_fixit")
errors = read(args.input)
skel = read(args.skeleton) if args.skeleton and os.path.exists(args.skeleton) else ""
user = render_template(
tmpl,
STYLE=style,
DETERMINISM=determinism,
SNIPPET=skel,
ERRORS=errors,
)
resp = client.chat(system, user)
write(args.out, resp)
else:
print(f"Unknown mode: {args.mode}", file=sys.stderr)
sys.exit(2)
# Compile after generation if configured
bcmd = cfg["compile"].get("build_cmd")
if bcmd:
subprocess.run(bcmd, shell=True)
if __name__ == "__main__":
main()