import asyncio import uuid from collections import OrderedDict from typing import Optional import json import yaml from deepdiff import DeepDiff from agents import Runner, trace, Agent from utils import build_agent from base import TransitionInput, StepJudgeResult, SaveList, ReusableStep4GPT, ReusableStep from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_exponential import openai @retry( stop=stop_after_attempt(10), wait=wait_exponential(multiplier=1, min=4, max=10), retry=retry_if_exception_type(openai.APIConnectionError), reraise=True ) async def run_with_retry(agent: Agent, messages: list[dict[str, str]]): return await Runner.run(agent, messages) class TransitionParser: def __init__(self, log_path_file: str, input_task_file: str, host: str, base_path: str, auth_type: str): self.log_path = log_path_file self.input_task_file = input_task_file self.host = host self.base_path = base_path self.auth_type = auth_type self.judge_agent = build_agent("judge_agent", "transition_judger.md", StepJudgeResult) self.save_builder = build_agent("save_builder", "save_builder.md", SaveList) self.request_parser = build_agent("request_parser", "request_parser.md", str) self.kv = { "region_id": {"regionId": "azure-eastus2"}, "project_id": {"projectId": "1369847559694040595"}, "private_key": {"private_key": "01GPB7L0"}, "public_key": {"public_key": "1a5bc2d1-70db-4621-aa82-91f2314123bc"}, # "subscription_id": {"subscription_id": "85b598b9-a080-4476-bc96-ba89e64f92d1"}, "subnet_id": {"subnet_id": "/subscriptions/85b598b9-a080-4476-bc96-ba89e64f92d1/resourceGroups/openapi-test/providers/Microsoft.Network/virtualNetworks/openapi-test-vnet/subnets/openapi-test-subnet"} } # key-value pairs for variable replacement, region, project, cloud provider are inputs self.responses = [] self.requests = [] with open(self.input_task_file, "r", encoding="utf-8") as f: self.input_task = json.load(f) with open(self.log_path, "r", encoding="utf-8") as f: log = json.load(f) self.transitions = log["transitions"] async def parse(self): self.result = [] candidate_transitions = [] for transition in self.transitions: if transition["request"]["request_type"].startswith("cmd") and (transition["response"]["error"] is not None and transition["response"]["error"] != ""): continue if transition["request"]["request_type"] == "http": if transition["request"].get("expect") and json.loads(transition["request"]["expect"]).get("status_code"): if transition["response"]["status_code"] != json.loads(transition["request"]["expect"]).get("status_code"): continue else: if transition["response"]["status_code"] // 100 != 2: continue transition_input = TransitionInput( operation_id=transition["operation_id"], state_before=transition["state_before"], state_after=transition["state_after"], request_type=transition["request"]["request_type"], request=transition["request"]["request"], response=transition["response"]["json"], error=transition["response"]["error"], expect=transition["request"]["expect"], max_retries=transition["request"]["max_retries"], delay_after=transition["request"]["delay_after"], status_code=transition["response"]["status_code"], operation_selection=transition["operation_selection"]["output"], request_generation=transition["request_generation"]["output"], execution_analysis=transition["execution_analysis"]["output"], target_matching=transition["target_matching"]["output"], resource_manager=transition["resource_manager"]["output"] ) candidate_transitions.append(transition_input) self.requests.append(transition_input.request) for i, candidate in enumerate(candidate_transitions): parsed_transition = await self.parse_transition(candidate) if parsed_transition: reuseable_step = ReusableStep( operation_id=parsed_transition.operation_id, # case_type=parsed_transition.case_type, request_type=parsed_transition.request_type, request=parsed_transition.request, expect=parsed_transition.expect, description=parsed_transition.description, max_retries=parsed_transition.max_retries, delay_after=parsed_transition.delay_after, # step_type=parsed_transition.step_type, name=parsed_transition.name, save=parsed_transition.save, continue_on_unexpected=parsed_transition.continue_on_unexpected, stage=parsed_transition.stage, meta=parsed_transition.meta, generation_id="", order=i ) self.result.append(reuseable_step) async def parse_request(self, request: dict) -> str: content = { "request": request, "cache": self.kv, } abstract_request = await run_with_retry(self.request_parser, [{"content": json.dumps(content, ensure_ascii=False), "role": "user"}]) request_str = abstract_request.final_output request_dict = json.loads(request_str.lstrip("```json\n").rstrip("\n```")) # trick start if request_dict.get("path") and request_dict["path"].startswith("/"): request_dict["path"] = "https://" + self.host + self.base_path + request_dict["path"] if request_dict.get("auth") and request_dict["auth"]["content"] in [None, {}]: if self.auth_type == "digest": request_dict["auth"]["content"] = {"private_key": "placeholder", "public_key": "placeholder"} elif self.auth_type == "bearer": request_dict["auth"]["content"] = {"token": "placeholder"} if request_dict.get("command"): request_dict["name"] = request_dict.pop("command") # trick end return json.dumps(request_dict, ensure_ascii=False, indent=2) async def parse_expect(self, expect: str | None, status_code: int, request_type: str) -> str | None: if expect is None: if request_type == "cli": return json.dumps({"exit_code": status_code, "cel": ""}, indent=2) return json.dumps({"status_code": status_code, "cel": ""}, indent=2) return json.dumps(json.loads(expect), indent=2) async def parse_save(self, request_type: str, response: dict) -> str: content = { "request_type": request_type, "response": response, "requests": self.requests, "cache": self.kv, "input_task": self.input_task } kv_list = await run_with_retry(self.save_builder, [{"content": json.dumps(content, ensure_ascii=False), "role": "user"}]) save = {"placeholder": []} for kv in kv_list.final_output.placeholder: if kv.key is None or kv.eval is None or kv.value is None: continue save["placeholder"].append({"key": kv.key, "eval": kv.eval}) self.kv[kv.key] = {kv.eval: kv.value} if len(save["placeholder"]) == 0: return None return json.dumps(save, ensure_ascii=False, indent=2) # async def generate_name_and_description(self, transition: TransitionInput) -> tuple[str, str]: # return json.dumps(transition.model_dump(), ensure_ascii=False) async def parse_transition(self, transition: TransitionInput) -> Optional[ReusableStep4GPT]: content = { "request": transition.request, "response": transition.response, "state_change": DeepDiff(transition.state_before, transition.state_after).to_json(), "execution_analysis": transition.execution_analysis, "target_matching": transition.target_matching, "input_task": self.input_task } judge_result = await run_with_retry(self.judge_agent, [{"content": json.dumps(content, ensure_ascii=False), "role": "user"}]) if not judge_result.final_output.keep: return None # name, description = await self.generate_name_and_description(transition) name, description = "", "" reuseable_step = ReusableStep4GPT( operation_id=transition.operation_id, request_type=transition.request_type, request=await self.parse_request(transition.request), expect=await self.parse_expect(transition.expect, transition.status_code, transition.request_type), description=description, max_retries=transition.max_retries, delay_after=transition.delay_after, name=name, save=await self.parse_save(transition.request_type, transition.response), continue_on_unexpected=False, stage="", meta="{}" ) return reuseable_step def dump(self, yaml_path: str): class CustomDumper(yaml.SafeDumper): def ignore_aliases(self, data): return True def represent_scalar(self, tag, value, style=None): if isinstance(value, str) and "\n" in value: return yaml.ScalarNode(tag, value, style="|") return super().represent_scalar(tag, value, style) yaml.add_representer( OrderedDict, yaml.representer.SafeRepresenter.represent_dict, Dumper=CustomDumper, ) if hasattr(self.result, 'model_dump'): content = self.result.model_dump() elif isinstance(self.result, list) and hasattr(self.result[0], 'model_dump'): content = [x.model_dump() for x in self.result] else: content = self.result scenario = { "scenario": { "id": str(uuid.uuid4()), "context": { "auth": { "type": self.auth_type }, "base_path": self.base_path, "host": self.host }, "description": self.input_task[list(self.input_task.keys())[0]]["description"], "name": list(self.input_task.keys())[0] }, "steps": content } with open(yaml_path, "w", encoding="utf-8") as yaml_file: yaml.dump( scenario, yaml_file, sort_keys=False, default_flow_style=False, allow_unicode=True, Dumper=CustomDumper, ) async def log2reusable(log_path: str, input_task_path: str, output_path: str, host: str, base_path: str, auth_type: str): with trace("log2reusable_test_0423_v1"): parser = TransitionParser(log_path, input_task_path, host, base_path, auth_type) await parser.parse() parser.dump(output_path) if __name__ == "__main__": import argparse parser = argparse.ArgumentParser(description="Convert log and input_task to reusable YAML.") parser.add_argument('--log', type=str, default="./inputs/log.json", help='Path to the log file') parser.add_argument('--input_task', type=str, default="./inputs/input_task.json", help='Path to the input_task file') parser.add_argument('--output', type=str, default="./outputs/output.yaml", help='Output YAML file path') parser.add_argument('--host', type=str, default="dedicated.staging.tidbapi.com", help='Host for the API') parser.add_argument('--base_path', type=str, default="/v1beta1", help='Base path for the API') parser.add_argument('--auth_type', type=str, default="digest", help='Auth type for the API') args = parser.parse_args() asyncio.run(log2reusable(args.log, args.input_task, args.output, args.host, args.base_path, args.auth_type))