(self, session: Session, res: str, operate: OperationEnum)
| 941 | token_usage=token_usage) |
| 942 | |
| 943 | def check_sql(self, session: Session, res: str, operate: OperationEnum) -> tuple[str, Optional[list]]: |
| 944 | json_str = extract_nested_json(res) |
| 945 | |
| 946 | log = self.current_logs[operate] |
| 947 | |
| 948 | if json_str is None: |
| 949 | trigger_log_error(session, log) |
| 950 | raise SingleMessageError(orjson.dumps({'message': 'SQL answer is not a valid json object', |
| 951 | 'traceback': "SQL answer is not a valid json object:\n" + res}).decode()) |
| 952 | sql: str |
| 953 | data: dict |
| 954 | try: |
| 955 | data = orjson.loads(json_str) |
| 956 | |
| 957 | if data['success']: |
| 958 | sql = data['sql'] |
| 959 | else: |
| 960 | message = data['message'] |
| 961 | raise SingleMessageError(message) |
| 962 | except SingleMessageError as e: |
| 963 | trigger_log_error(session, log) |
| 964 | raise e |
| 965 | except Exception: |
| 966 | trigger_log_error(session, log) |
| 967 | raise SingleMessageError(orjson.dumps({'message': 'Cannot parse sql from answer', |
| 968 | 'traceback': "Cannot parse sql from answer:\n" + res}).decode()) |
| 969 | |
| 970 | if sql.strip() == '': |
| 971 | trigger_log_error(session, log) |
| 972 | raise SingleMessageError("SQL query is empty") |
| 973 | return sql, data.get('tables') |
| 974 | |
| 975 | @staticmethod |
| 976 | def get_chart_type_from_sql_answer(res: str) -> Optional[str]: |
no test coverage detected