118 lines
4.0 KiB
Python
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()
|
|
|