|
| 1 | +import argparse |
| 2 | +import glob |
| 3 | +import json |
| 4 | +import os |
| 5 | + |
| 6 | +import openai |
| 7 | +import yaml |
| 8 | + |
| 9 | + |
| 10 | +def read_missing_properties(path): |
| 11 | + path = os.path.join(path, "missing_fields.json") |
| 12 | + with open(path, "r", encoding="utf-8") as f: |
| 13 | + missing_properties = json.load(f) |
| 14 | + |
| 15 | + return missing_properties |
| 16 | + |
| 17 | + |
| 18 | +def gen_values(missing_values, path, api_key, operator): |
| 19 | + openai.api_key = api_key |
| 20 | + |
| 21 | + context = f"You are a expert of the {operator} of the Kubernetes ecosystem. You are tasked with providing values for properties of the {operator} CRD" |
| 22 | + |
| 23 | + for i in range(len(missing_values)): |
| 24 | + p = missing_values[i] |
| 25 | + |
| 26 | + if p[0].endswith("ITEM"): |
| 27 | + continue |
| 28 | + |
| 29 | + prop = f"- {p[0]}\n description: {p[1]}\n type: {p[2]}\n structure: {p[3]}\n" |
| 30 | + |
| 31 | + prompt = "Here is the property that need values:\n" |
| 32 | + prompt += f"{prop}\n" |
| 33 | + |
| 34 | + prompt += "\nThe property has a datatype and description provided above, please make sure the generated value satisfies the datatype and description.\n" |
| 35 | + |
| 36 | + prompt += "\n If the property has structure that indicating subfields, make sure to generate all the subfields for the property as a whole.\n" |
| 37 | + |
| 38 | + prompt += "\nProvide three values for the property and please follow the cr yaml format. Directly give me the yaml file without any other message, for example\n" |
| 39 | + |
| 40 | + format = "spec:\n" |
| 41 | + format += f" {p[0]}: value\n" |
| 42 | + format += "---" |
| 43 | + format += "spec:\n" |
| 44 | + format += f" {p[0]}: value\n" |
| 45 | + |
| 46 | + prompt += format |
| 47 | + |
| 48 | + prompt += 'If the property has `ITEM` in the property path, that means the property should be an item in an array. For example, for "spec.pdms.ITEM.config:" the format should be:\n' |
| 49 | + array_format = "spec:\n" |
| 50 | + array_format += f" property:\n" |
| 51 | + array_format += " - subproperty: value\n" |
| 52 | + array_format += " - subproperty: value\n" |
| 53 | + array_format += " - subproperty: value\n" |
| 54 | + |
| 55 | + prompt += array_format |
| 56 | + |
| 57 | + completion = openai.chat.completions.create( |
| 58 | + model="gpt-4o", |
| 59 | + messages=[ |
| 60 | + {"role": "system", "content": context}, |
| 61 | + {"role": "user", "content": prompt}, |
| 62 | + ], |
| 63 | + ) |
| 64 | + |
| 65 | + result_text = completion.choices[0].message.content |
| 66 | + |
| 67 | + result = result_text.split("```yaml\n")[1].split("```")[0] |
| 68 | + |
| 69 | + output_file = os.path.join(path, f"value_{p[0]}.yaml") |
| 70 | + with open(output_file, "w") as f: |
| 71 | + f.write(result) |
| 72 | + |
| 73 | + |
| 74 | +def store_to_examples(path): |
| 75 | + main_results = [] |
| 76 | + |
| 77 | + for file in glob.glob(os.path.join(path, "*.yaml"), recursive=True): |
| 78 | + if file == "examples.yaml": |
| 79 | + continue |
| 80 | + with open(file, "r") as f: |
| 81 | + content = yaml.safe_load_all(f) |
| 82 | + for doc in content: |
| 83 | + main_results.append(doc) |
| 84 | + with open(os.path.join(path, "examples.yaml"), "w") as f: |
| 85 | + yaml.dump_all(main_results, f) |
| 86 | + |
| 87 | + |
| 88 | +if __name__ == "__main__": |
| 89 | + parser = argparse.ArgumentParser() |
| 90 | + parser.add_argument( |
| 91 | + "--path", |
| 92 | + type=str, |
| 93 | + help="Path to the file containing the missing properties", |
| 94 | + ) |
| 95 | + parser.add_argument( |
| 96 | + "--api_key", type=str, help="API key for the OpenAI API" |
| 97 | + ) |
| 98 | + parser.add_argument("--operator", type=str, help="Name of the operator") |
| 99 | + args = parser.parse_args() |
| 100 | + |
| 101 | + missing_properties = read_missing_properties(args.path) |
| 102 | + |
| 103 | + gen_values(missing_properties, args.path, args.api_key, args.operator) |
| 104 | + store_to_examples(args.path) |
0 commit comments