from typing import Dict, Any, List, Set, Optional import json from pydantic import BaseModel from llmops.environment.executor import CompoundExecutor from llmops.config import business_path BUSINESS_PATH = business_path / "tidbcloud_serverless_sut" / "configs" / "sut" class SutAttr(BaseModel): region: str cloud_provider: str project_id: Optional[str] = None # environment: Optional[str] = 'dev' class TidbCloudExpert: def __init__( self, swagger_path: str, executor: CompoundExecutor, sut_attr: SutAttr ): """ Args: swagger_path: OpenAPI/Swagger文件的路径 executor: CompoundExecutor实例 """ self.abs_swagger_path = swagger_path self.executor = executor self.sut_attr = sut_attr if sut_attr.cloud_provider == "aws": cloud_cli_spec_path = BUSINESS_PATH / "aws_cli.json" else: cloud_cli_spec_path = None try: with open(self.abs_swagger_path, "r") as f: self.spec = json.load(f) if cloud_cli_spec_path: with open(cloud_cli_spec_path, "r") as f: self.cloud_cli_spec = json.load(f) with open(BUSINESS_PATH / "mysqlsh.json", "r") as f: self.mysqlsh_spec = json.load(f) except ValueError as e: raise ValueError(f"Failed to load OpenAPI/Swagger file: {e}") @property def operation_ids(self) -> List[str]: openapi_operation_ids = self._openapi_operation_ids() cloud_cli_operation_ids = self._cloud_cli_operation_ids() mysqlsh_operation_ids = self._mysqlsh_operation_ids() return openapi_operation_ids + cloud_cli_operation_ids + mysqlsh_operation_ids def _openapi_operation_ids(self) -> List[str]: """获取所有operationId的列表""" operation_ids = [] for path in self.spec.get("paths", {}).values(): for operation in path.values(): if "operationId" in operation: operation_ids.append(operation["operationId"]) return operation_ids def _cloud_cli_operation_ids(self) -> List[str]: return list(self.cloud_cli_spec.keys()) def _mysqlsh_operation_ids(self) -> List[str]: return list(self.mysqlsh_spec.keys()) @property def extract_ids_and_descriptions(self) -> Dict[str, Dict[str, str]]: openapi_id_descriptions = self._extract_openapi_ids_and_descriptions() cloud_cli_id_descriptions = self._extract_cloud_cli_ids_and_descriptions() mysqlsh_id_descriptions = self._extract_mysqlsh_ids_and_descriptions() return openapi_id_descriptions | cloud_cli_id_descriptions | mysqlsh_id_descriptions def _extract_openapi_ids_and_descriptions( self, ) -> Dict[str, Dict[str, str]]: """获取所有operationId及其描述的映射""" id_descriptions = {} for path in self.spec.get("paths", {}).values(): for operation in path.values(): if "operationId" in operation: id_descriptions[operation["operationId"]] = { "description": operation.get("summary", ""), "type": "TiDB Cloud Serverless OpenAPI", } return id_descriptions def _extract_cloud_cli_ids_and_descriptions(self) -> Dict[str, str]: """获取所有operationId及其描述的映射""" id_descriptions = {} for _id, attr in self.cloud_cli_spec.items(): id_descriptions[_id] = { "description": attr.get("description", ""), "type": f"{self.sut_attr.cloud_provider}_cli", } return id_descriptions def _extract_mysqlsh_ids_and_descriptions(self) -> Dict[str, str]: """获取所有operationId及其描述的映射""" id_descriptions = {} for _id, attr in self.mysqlsh_spec.items(): id_descriptions[_id] = { "description": attr.get("description", ""), "type": "mysqlsh", } return id_descriptions def _collect_refs(self, obj: Any, refs: Set[str]): """递归收集对象中的所有$ref引用 Args: obj: 要检查的对象 refs: 收集到的引用集合 """ if isinstance(obj, dict): if "$ref" in obj: refs.add(obj["$ref"]) for value in obj.values(): self._collect_refs(value, refs) elif isinstance(obj, list): for item in obj: self._collect_refs(item, refs) def _get_schema_name(self, ref: str) -> str: """从$ref字符串中提取schema名称""" return ref.split("/")[-1] def _extract_required_schemas(self, refs: Set[str]) -> Dict[str, Any]: """提取所需的schemas Args: refs: 收集到的$ref引用集合 Returns: 包含所需schemas的字典 """ schemas = {} processed_refs = set() while refs: ref = refs.pop() if ref in processed_refs: continue processed_refs.add(ref) if not ref.startswith("#/definitions/"): continue schema_name = self._get_schema_name(ref) schema = self.spec["definitions"].get(schema_name) if schema: schemas[schema_name] = schema # 收集这个schema中的新引用 new_refs = set() self._collect_refs(schema, new_refs) # 添加新的未处理的引用 refs.update(new_refs - processed_refs) schemas.pop("googlerpcStatus") schemas.pop("protobufAny") return schemas def extract_operation(self, operation_id: str) -> Dict[str, Any]: if self.mysqlsh_spec.get(operation_id) is not None: cli_spec = { operation_id: { "specification": self.mysqlsh_spec[operation_id], "cli_type": "mysqlsh" } } return cli_spec if self.cloud_cli_spec.get(operation_id) is None: return self._extract_openapi_operation(operation_id) cli_spec = { operation_id: { "specification": self.cloud_cli_spec[operation_id], "cli_type": f"{self.sut_attr.cloud_provider}_cli" } } return cli_spec def _extract_openapi_operation(self, operation_id: str) -> Dict[str, Any]: """提取指定operationId的操作信息和相关schemas Args: operation_id: 要提取的操作ID Returns: Dict包含两部分: - paths: 包含该操作的路径信息 - components: 包含该操作依赖的所有schemas """ # HACK: inject region # 查找操作 operation_path = None operation_method = None operation_info = None for path, methods in self.spec.get("paths", {}).items(): for method, operation in methods.items(): if operation.get("operationId") == operation_id: operation_path = path operation_method = method operation_info = operation break if operation_info: break if not operation_info: raise ValueError(f"Operation {operation_id} not found") # 收集所有引用 refs: Set[dict] = set() self._collect_refs(operation_info, refs) # 构建返回结果 result = { "paths": {operation_path: {operation_method: operation_info}}, "components": {"schemas": self._extract_required_schemas(refs)}, } return result # def list_all_resources(self) -> Dict[str, Any]: # """mamually list all resources""" # resources: Dict[str, Any] = {"clusters": {}} # list_clusters_request = { # "method": "GET", # "path": "/clusters", # "auth_type": "digest", # } # clusters_response = self.executor.execute_http_request(**list_clusters_request) # clusters = clusters_response.json.get("clusters") # for cluster in clusters: # resources["clusters"][cluster["clusterId"]] = {} # resources["clusters"][cluster["clusterId"]]["attributes"] = cluster # resources["clusters"][cluster["clusterId"]]["state"] = cluster.get("state") # resources["clusters"][cluster["clusterId"]]["references"] = [] # resources["clusters"][cluster["clusterId"]]["parent_id"] = None # # list_nodegroups_request = { # # "method": "GET", # # "path": f"/api/clusters/{cluster['id']}/nodegroups", # # } # # nodegroups_response = self.executor.execute_http_request(**list_nodegroups_request) # return dict( # clusters=resources["clusters"], # ) def list_all_resources(self) -> Dict[str, Any]: """manually initialize a list of resource types""" return { "clusters": {}, "branches": {}, "endpoints": {}, "export_jobs": {}, "import_jobs": {}, } # def cleanup(self): # """cleanup resources # currently an empty implementation, can be extended later # """ # print("start to clean up resources") # current_state = self.list_all_resources() # print( # "length of current_state: ", # len(current_state["clusters"]), # ) # for cluster_id, _ in current_state["clusters"].items(): # self.executor.execute_http_request( # method="DELETE", path=f"/clusters/{cluster_id}", auth_type="digest" # ) # print("resources cleanup completed") @property def biz_system_prompt_preprocessor(self) -> str: """general biz logic for preprocess""" return """""" @property def biz_system_prompt_operation_selector(self) -> str: """general biz logic for operation_selector""" return """Your task is to select the most appropriate operation for TiDB Cloud resources, primarily focusing on Cluster management. And some auxiliary operations (like cloud cli or mysqlsh client) to prepare resources or data. Analysis Steps: 1. First, identify from `target_operation`: - Which resource type is being operated on - What state transition is desired (from current state to target state) 2. Use the `current_state` as your source of truth to: - Verify the resource's current status - Determine the most effective operation to progress toward the `target_operation` Key Rules: - If the resource is in any transitional state (ending in "ING"): - CREATING - MODIFYING - RESUMING - PAUSING etc. You MUST select a GET-type operation first to obtain the updated resource status Decision Making: - Choose operations that make the most direct progress toward the target state - Prioritize checking current status when resource is in transition - Only proceed with state-changing operations when the resource is in a stable state Example Flow: If current_state shows CREATING: → Select GET operation to check current status If current_state shows ACTIVE and target is PAUSED: → Select PAUSE operation """ @property def biz_system_prompt_request_generator(self) -> str: """general biz logic for request_generator""" return f"""When generating API requests, please follow these rules for the `expect` field: For GET requests: - The `expect` field must be a valid JSON string containing two fields: - `status_code`: Must be set to 200 - `cel`: Must be a valid Common Expression Language (CEL) expression starting with "body." - The CEL expression must explicitly validate the "state" field with an appropriate ENUM value from the response schema - Example: `{{"status_code": 200, "cel": "body.state == 'ACTIVE'"}}` For non-GET requests: - Set the `expect` field to null - Example: `"expect": null` Important: - Ensure the JSON string is properly escaped - The CEL expression must be syntactically valid - Choose the most appropriate state value from the available ENUM options in the response schema If a `region` field is required, please set it as "{self.sut_attr.cloud_provider}-{self.sut_attr.region}". If a `project_id` field is given in the instruction, please set it as "{self.sut_attr.project_id}". TiDB Cloud mysqlsh rules (adhoc, MUST follow when generating mysqlsh requests): 1. Determine whether the target_operation requires using a CLUSTER endpoint or a BRANCH endpoint. - If the text says "branch public endpoint", it is BRANCH. - If the text says "cluster public endpoint", it is CLUSTER. 2. Username/password MUST match the target type: - BRANCH mysqlsh: username MUST be `{{branch_userPrefix}}.root`. Password MUST be the branch password explicitly mentioned in target_operation (e.g. 'pi31415926'). - CLUSTER mysqlsh: username MUST be `{{cluster_userPrefix}}.root`. Password MUST be the cluster password explicitly mentioned in target_operation (e.g. 'pi31415926536'). - NEVER reuse cluster userPrefix/password for branch. - NEVER reuse branch userPrefix/password for cluster. 3. userPrefix/host/port source of truth: - For CLUSTER: extract userPrefix from cluster attributes (`current_state.clusters[*].attributes.userPrefix`) or from the derived public endpoint resource in `current_state.endpoints`. - For BRANCH: extract userPrefix from branch attributes (`current_state.branches[*].attributes.userPrefix`) or from the branch public endpoint nested fields. - If a required field is missing, do NOT invent it. Generate the necessary GET operation (GetCluster / GetBranch) to retrieve it first. 4. mysqlsh URI MUST use the endpoint host/port corresponding to the chosen type (branch vs cluster), and database name from target_operation (e.g. `test`). """ @property def biz_system_prompt_response_analyzer(self) -> str: """general biz logic for preprocess""" return """- Root password won't be shown in response even if succeeded, only check status_code for the changing password operation. - Most attributes in the `target_operation` may not be taken into consideration, such as `displayName`, etc. """ @property def biz_system_prompt_target_matcher(self) -> str: """general biz logic for preprocess""" return """Important fields to compare (may not be top-level fields, could be nested under other fields): - Node specifications (e.g. node_size, instance_type) - Node count and configuration (e.g. node_quantity, replicas) - Storage settings (e.g. storage_size_gib) Fields that do not need comparison: - Display names or labels (e.g. displayName, description) - Metadata fields (e.g. created_at, updated_at) - Internal identifiers (e.g. cluster_id, request_id) """ @property def biz_system_prompt_resource_manager(self) -> str: """general biz logic for resource_manager""" return """## TiDB Cloud Resource Hierarchy and ID Extraction Rules ### Resource Types and Hierarchy TiDB Cloud has the following resource hierarchy: 1. **Cluster** (root resource) - Resource type: `clusters` - Parent: None - Children: endpoints, import_jobs, export_jobs 2. **Endpoint** (child of Cluster) - Resource type: `endpoints` - Parent: cluster - Children: None 3. **Import Job** (references Cluster) - Resource type: `import_jobs` - Parent: cluster (via clusterId reference) - Children: None 4. **Export Job** (references Cluster) - Resource type: `export_jobs` - Parent: cluster (via clusterId reference) - Children: None 5. **Branch** (references Cluster) - Resource type: `branches` - Parent: cluster (via clusterId reference) - Children: None ### ID Extraction Rules (CRITICAL) **For Cluster:** - Primary location: `response.clusterId` OR `response.cluster.clusterId` - Never fabricate cluster IDs - Always use exact value from response **For Endpoint:** - Endpoints don't have explicit IDs in TiDB Cloud API - **Derived ID pattern**: `{clusterId}-{type}-endpoint` - Example: If clusterId is "cls-abc123" and type is "public", endpoint ID = "cls-abc123-public-endpoint" - Extract from: `response.endpoints.public` or `response.endpoints.private` - Fields to track in attributes: - `host`: endpoint hostname - `port`: endpoint port - `type`: "public" or "private" - `userPrefix`: user prefix for authentication (if available) **For Import Job:** - Primary location: `response.importId` OR `response.import.importId` - Parent cluster: extract from `response.clusterId` or request path `/clusters/{clusterId}/imports` - State location: `response.state` (typical values: "PENDING", "IMPORTING", "COMPLETED", "FAILED") **For Export Job:** - Primary location: `response.exportId` OR `response.export.exportId` - Parent cluster: extract from `response.clusterId` or request path `/clusters/{clusterId}/exports` - State location: `response.state` (typical values: "PENDING", "EXPORTING", "SUCCEEDED", "FAILED") ### State Extraction Rules **Cluster states** (from `response.state` or `response.cluster.state`): - CREATING - ACTIVE - MODIFYING - PAUSED - PAUSING - RESUMING - DELETING (rare, usually deleted immediately) **Import/Export Job states** (from `response.state`): - PENDING - IMPORTING / EXPORTING - COMPLETED / SUCCEEDED - FAILED - CANCELED **Endpoints don't have explicit states** - leave state as null ### Special Handling Rules 1. **When creating a cluster with public endpoint:** - CREATE operation returns cluster with `endpoints.public` nested object - You MUST create TWO resources: a. Cluster resource (clusterId from response.clusterId) b. Endpoint resource (derived ID: "{clusterId}-public-endpoint") - Set cluster's `reference_ids` to include the endpoint ID - Set endpoint's `parent_id` to the cluster ID 2. **When deleting a cluster:** - Remove the cluster resource - CASCADE DELETE: Also remove ALL child resources where `parent_id` equals the deleted cluster ID - This includes: endpoints, any associated import/export jobs 3. **When starting an import/export job:** - CREATE operation for import/export - Extract job ID from response (importId/exportId) - Extract parent cluster ID from request path or response - Set job's `parent_id` to the cluster ID - Job's `reference_ids` can include S3 bucket ARN or role ARN if tracked 4. **For mysqlsh CLI operations:** - mysqlsh operations don't create trackable resources - They're side effects on database content (tables, data) - Do NOT create resources for: databases, tables, SQL query results - Update ONLY if the operation affects cluster state (rare) --- **Remember:** - Always cascade delete child resources when parent is deleted - Never track database-level objects (databases, tables, users, SQL results) - Extract state from the response's state field when available - Keep attributes as a complete JSON string of relevant response data """ if __name__ == "__main__": from pathlib import Path from environs import Env current_path = Path(__file__).parent config_path = current_path / "configs/sut" task_path = current_path / "tasks" env = Env() env.read_env(config_path / ".env") api_token = env.str("SUT_TOKEN") executor = CompoundExecutor( base_url="https://serverless.staging.tidbapi.com/v1beta1", api_token=api_token, auth_type="digest", ) sut_attr = SutAttr( region="us-east-1", cloud_provider="aws", project_id="1369847559694040744", ) tidbcloud_expert = TidbCloudExpert(config_path / "openapi.json", executor, sut_attr) resources = tidbcloud_expert.list_all_resources() print("resources: ", resources) # print("operation ids: ", tidbcloud_expert.operation_ids) # print("operation descriptions: ", tidbcloud_expert.extract_ids_and_descriptions) print( "operation: ", json.dumps( tidbcloud_expert.extract_operation("ClusterService_CreateCluster"), indent=2, ), ) print( "operation: ", json.dumps( tidbcloud_expert.extract_operation("BranchService_GetBranch"), indent=2, ), )