import gevent.monkey # noqa gevent.monkey.patch_all() # noqa # place it under the main gevent patch import grpc.experimental.gevent # noqa grpc.experimental.gevent.init_gevent() # noqa import time from errors.errors import BadRequestCausedAlgorithmException, BadRequestException, NotFoundException, PermissionDeniedException, parse_exception from werkzeug.exceptions import InternalServerError from flask import Flask, g, jsonify, request, Response import sentry_sdk from sentry_sdk.integrations.flask import FlaskIntegration from prometheus_client import ( Histogram, Gauge, multiprocess, generate_latest, CollectorRegistry, CONTENT_TYPE_LATEST, ) from flasgger import APISpec, Schema, Swagger, fields from apispec.ext.marshmallow import MarshmallowPlugin from apispec_webframeworks.flask import FlaskPlugin from handlers.base import fail, get_current_user_email from handlers.context import context_bp from handlers.chat2query import chat2query_bp from handlers.inner import inner_bp from handlers.edatask import ( edatask_bp, create_data_context, create_session_context, handle_chat2data, handle_chat2data_v3, handle_create_data_summaries, handle_create_knowledge_base_data, handle_create_knowledge_base, handle_create_session_context, handle_db_exploring, handle_delete_knowledge_base_data, handle_execute_query, handle_generate_suggest_questions, handle_get_data_summaries, handle_get_data_summary, handle_get_knowledge_base, handle_get_session_context, handle_get_session_contexts, handle_list_knowledge_base_data, handle_list_knowledge_bases, handle_multi_round_chat2sql, handle_query_job_detail, handle_question_breakdown, handle_reset_session_context, handle_sql_refine, handle_text2sql, get_edatask_job_detail, handle_like_or_not, handle_update_column_description, handle_update_data_summary, handle_update_knowledge_base_data, handle_update_knowledge_base, handle_update_session_context, handle_update_table_description, handle_execute_sql, handle_update_sql, edatask_v2_bp, edatask_v3_bp, edatask_v1_bp, ) from handlers.sqleditor import sqleditor_v1_bp from models.base import Session from config import config from schema.schema import ( CreateDataContextSchema, CreateDataSummarySchema, CreateSessionContextSchema, LikeOrNotSchema, TaskExecutionSchema, QuestionBreakdownSchema, SQLEditSchema, Text2SQLSchema, UpdateColumnDescriptionSchema, UpdateDataDescriptionSchema, UpdateTableDescriptionSchema, CheckConnectionSchema, Chat2DataSchema, SQLExecutionSchema ) from tasks.utils import get_default_queue_length from utils.logging import logging, trace_id_vars app = Flask(__name__) app.config['SWAGGER'] = { 'openapi': "3.0.3", } sentry_sdk.init( dsn=config.SENTRY_DSN, integrations=[ FlaskIntegration(), ], environment=config.ENVIRONMENT, enable_tracing=False, traces_sample_rate=0, ignore_errors=[BadRequestCausedAlgorithmException, BadRequestException], ) REQUEST_LATENCY = Histogram( 'request_latency_seconds', 'request latency in seconds(support float number)', labelnames=("base_url", "method", "status_code") ) QUEUE_LENGTH = Gauge("queue_length", "celery queue length", labelnames=("name",)) app.register_blueprint(context_bp, url_prefix="/api/v1beta/apps/eda/context") # app.register_blueprint(credential_bp, url_prefix="/api/v1/edaservice/credential") # app.register_blueprint(prompt_bp, url_prefix="/api/v1/edaservice/prompt") app.register_blueprint(edatask_bp, url_prefix="/api/v1beta/apps/") app.register_blueprint(edatask_v1_bp, url_prefix="/inner/eda/api/v1beta/app//endpoint/v1") app.register_blueprint(edatask_v2_bp, url_prefix="/inner/eda/api/v1beta/app//endpoint/v2") app.register_blueprint(edatask_v3_bp, url_prefix="/inner/eda/api/v1beta/app//endpoint/v3") app.register_blueprint(chat2query_bp, url_prefix="/inner/eda/api/v1/dataservices/central/orgs//projects//apps//chat2query") app.register_blueprint(inner_bp, url_prefix="/inner/internal") app.register_blueprint(sqleditor_v1_bp, url_prefix="/inner/eda/api/v1/eda/orgs//projects/") # Create an APISpec spec = APISpec( title="Chat2Query V2", version='0.1', openapi_version='2.0', plugins=[ FlaskPlugin(), MarshmallowPlugin(), ], ) def real_panic(foo): raise Exception("panic") @app.route("/panic") def panic(): hello = "world" real_panic(hello) @app.before_request def before_request(): g.session = Session() request.start_time = time.time() b3_trace_id = request.headers.get("B3") if b3_trace_id: trace_id_vars.set(b3_trace_id) # @app.errorhandler(werkzeug.exceptions.BadRequest) # @app.errorhandler(BadRequestException) # def bad_request(e): # return fail(400, "bad request", None) @app.errorhandler(404) @app.errorhandler(NotFoundException) def page_not_found(e): return fail(404, "not found", None) @app.errorhandler(500) @app.errorhandler(Exception) def internal_server_error(e): msg, code = parse_exception(e) if code >= 500: logging.error("exception found %s", e) msg = "internal server error" return fail(code, msg, None) @app.after_request def after_request(response): try: g.session.commit() except Exception as e: logging.info("failed to commit, start to rollback: %s", e) g.session.rollback() finally: g.session.close() latency = time.time() - request.start_time REQUEST_LATENCY.labels(request.url_rule, request.method, response.status_code).observe(latency) return response @app.route("/metrics") def metrics(): length = get_default_queue_length() QUEUE_LENGTH.labels("celery").set(length) registry = CollectorRegistry() multiprocess.MultiProcessCollector(registry) data = generate_latest(registry) status = '200 OK' headers = [ ('Content-type', CONTENT_TYPE_LATEST), ('Content-Length', str(len(data))) ] return Response(data, status, headers) @app.get("/healthz") def healthz(): return "I'm OK", 200 def rule_filter(rule) -> bool: return True def model_filter(tag) -> bool: return True handler_paths = [ handle_create_data_summaries, handle_chat2data, handle_query_job_detail, handle_execute_query, handle_get_data_summaries, handle_create_data_summaries, handle_get_data_summary, handle_update_data_summary, handle_update_table_description, handle_update_column_description, handle_chat2data_v3, handle_get_session_contexts, handle_create_session_context, handle_get_session_context, handle_reset_session_context, handle_update_session_context, handle_sql_refine, handle_create_knowledge_base, handle_update_knowledge_base, handle_list_knowledge_bases, handle_get_knowledge_base, handle_list_knowledge_base_data, handle_create_knowledge_base_data, handle_update_knowledge_base_data, handle_delete_knowledge_base_data, handle_multi_round_chat2sql, handle_generate_suggest_questions, ] template = spec.to_flasgger(app, paths=handler_paths) # start Flasgger using a template from apispec schemas = { "SuggestQuestion": { "type": "object", "properties": { "cluster_id": { "type": "string", "description": "ID of your cluster in TiDB Cloud", }, "database": { "type": "string", "description": "Name of your database", }, }, }, "SuggestQuestionResponse": { "type": "object", "properties": { "job_id": { "type": "string", "description": "ID of the async generation job", }, }, }, "CreateDataSummary": { "type": "object", "properties": { "cluster_id": { "type": "string", "description": "ID of your cluster in TiDB Cloud", }, "database": { "type": "string", "description": "Name of your database", }, "description": { "type": "string", "description": "Description of the data summary", }, "reuse": { "type": "boolean", "description": "Whether to reuse the existing data summary, by default false", }, "default": { "type": "boolean", "description": "Set the data summary as the default one after explored, by default false", } }, "required": ["cluster_id", "database", "description"], }, "CreateDataSummaryResponse": { "type": "object", "properties": { "data_summary_id": { "type": "integer", "description": "ID of the data summary", }, "job_id": { "type": "string", "description": "ID of the async generation job", }, }, "required": ["data_summary_id", "job_id"], }, "ListDataSummaryV2Response": { "type": "object", "properties": { "data_summary_id": { "type": "integer", "description": "ID of the data summary", }, "status": { "type": "string", "description": "The current status of the data summary", }, }, }, "ListDataSummaryResponse": { "type": "object", "properties": { "data_summaries": { "type": "array", "items": { "type": "object", "properties": { "data_summary_id": { "type": "integer", "description": "ID of the data summary", }, "status": { "type": "string", "description": "The current status of the data summary", }, "description": { "type": "string", "description": "Description of the data summary", }, "default": { "type": "boolean", "description": "Whether the data summary is the default one", }, }, }, } } }, "DataSummaryDetail": { "type": "object", "properties": { "data_summary_id": { "type": "integer", "description": "ID of the data summary", }, "status": { "type": "string", "description": "The current status of the data summary", }, "default": { "type": "boolean", "description": "Whether the data summary is the default one", }, "description": { "type": "string", "description": "Description of the data summary", }, "tables": { "type": "object", "description": "The details of the tables, represented by a JSON object", "properties": { "name": { "type": "string", "description": "Name of the table", }, "description": { "type": "string", "description": "Description of the table", }, "columns": { "type": "object", "description": "The details of the columns, represented by a JSON object, with key as the column name and value as column description", } } } }, }, "UpdateDataSummary": { "type": "object", "properties": { "description": { "type": "string", "description": "Description of the data summary", }, "default": { "type": "boolean", "description": "Whether the data summary is the default one", }, }, }, "UpdateTableDescription": { "type": "object", "properties": { "description": { "type": "string", "description": "Description of the table", }, }, }, "UpdateColumnsDescription": { "type": "object", "properties": { "columns": { "type": "array", "description": "The details of the columns, represented by a JSON object, with key as the column name and value as column description", "items": { "type": "object", "properties": { "name": { "type": "string", "description": "Name of the column", }, "description": { "type": "string", "description": "Description of the column", }, } } } }, }, "CreateKnowledgeBase": { "type": "object", "properties": { "cluster_id": { "type": "string", "description": "ID of your cluster in TiDB Cloud", }, "database": { "type": "string", "description": "Name of your database", }, "description": { "type": "string", "description": "Description of the data summary", }, }, }, "CreateKnowledgeBaseResponse": { "type": "object", "properties": { "knowledge_base_id": { "type": "integer", "description": "ID of the knowledge base", }, }, "required": ["knowledge_base_id"], }, "UpdateKnowledgeBase": { "type": "object", "properties": { "description": { "type": "string", "description": "Description of the knowledge base", }, "default": { "type": "boolean", "description": "Whether the knowledge base is the default one", }, }, }, "ListKnowledgeBasesResponse": { "type": "object", "properties": { "knowledge_bases": { "type": "array", "items": { "type": "object", "properties": { "knowledge_base_id": { "type": "integer", "description": "ID of the knowledge base", }, "description": { "type": "string", "description": "Description of the knowledge base", }, "default": { "type": "boolean", "description": "Whether the knowledge base is the default one", }, }, }, } }, }, "KnowledgeBase": { "type": "object", "properties": { "knowledge_base_id": { "type": "integer", "description": "ID of the knowledge base", }, "description": { "type": "string", "description": "Description of the knowledge base", }, "default": { "type": "boolean", "description": "Whether the knowledge base is the default one", }, }, }, "Chat2Data": { "type": "object", "properties": { "cluster_id": { "type": "string", "description": "ID of the data summary", }, "database": { "type": "string", "description": "Name of your database", }, "question": { "type": "string", "description": "Question to be resolved", }, "sql_generate_mode": { "type": "string", "description": "SQL generation mode, support mode: 'direct' or 'auto_breakdown', by default 'direct'", }, }, "required": ["cluster_id", "database", "question"], }, "Chat2DataResponse": { "type": "object", "properties": { "job_id": { "type": "string", "description": "ID of the async generation job", }, }, "required": ["job_id"], }, "JobDetail": { "type": "object", "properties": { "job_id": { "type": "string", "description": "ID of the async generation job", }, "status": { "type": "string", "description": "The current status reflects the job's progress. A 'done' status signifies successful job completion, while a 'failed' status indicates unsuccessful completion. Any other status suggests that the job is ongoing", }, "result": { "type": "object", "description": "The details of the job, represented by a JSON object", }, "ended_at": { "type": "string", "description": "The time when the job ended", }, }, "required": ["job_id", "status", "result", "ended_at"], }, "ExecuteQuery": { "type": "object", "properties": { "cluster_id": { "type": "string", "description": "ID of your cluster in TiDB Cloud", }, "database": { "type": "string", "description": "Name of your database", }, "sql": { "type": "string", "description": "SQL to be executed", }, }, "required": ["cluster_id", "database", "sql"], }, "ExecuteQueryResponse": { "type": "object", "properties": { "code": { "type": "integer", "description": "The status code of the response", }, "msg": { "type": "string", "description": "The message of the response", }, "result": { "type": "object", "description": "The details of the execution, represented by a JSON object", } }, "required": ["code", "msg", "result"], }, "UserFeedback": { "type": "object", "properties": { "session_id": { "type": "integer", "description": "ID of the session context", }, "question_id": { "type": "string", "description": "ID of the question", }, "feedback": { "type": "string", "description": "Feedback to be sent, support feedbac: 'like' or 'dislike'", }, "comment": { "type": "string", "description": "Comment to be sent", }, }, "required": ["session_id", "question_id", "feedback"], }, "UserFeedbackResponse": { "type": "object", "properties": { "code": { "type": "integer", "description": "The status code of the response", }, "msg": { "type": "string", "description": "The message of the response", }, "result": { "type": "object", "description": "The details of the response, represented by a JSON object", } }, "required": ["code", "msg", "result"], }, "ListKnowledgeBaseDataResponse": { "type": "object", "properties": { "data": { "type": "array", "items": { "type": "object", "properties": { "knowledge_data_id": { "type": "integer", "description": "ID of the knowledge base data", }, "raw_data": { "type": "object", "description": "The details of the knowledge base data, represented by a JSON object", }, "type": { "type": "string", "description": "Type of the knowledge base data", }, "app_id": { "type": "integer", "description": "ID of the app", }, "meta_data": { "type": "object", "description": "The meta data of the knowledge base data, represented by a JSON object", }, }, }, }, }, }, "CreateKnowledgeBaseData": { "type": "object", "properties": { "type": { "type": "string", "description": "Type of the knowledge base data", }, "raw_data": { "type": "object", "description": "The details of the knowledge base data, represented by a JSON object", }, "meta_data": { "type": "object", "description": "The meta data of the knowledge base data, represented by a JSON object", }, }, "required": ["app_id", "type", "raw_data"], }, "CreateKnowledgeBaseDataResponse": { "type": "object", "properties": { "knowledge_data_id": { "type": "integer", "description": "ID of the knowledge base data", }, }, "required": ["knowledge_data_id"], }, "UpdateKnowledgeBaseData": { "type": "object", "properties": { "raw_data": { "type": "object", "description": "The details of the knowledge base data, represented by a JSON object", }, "meta_data": { "type": "object", "description": "The meta data of the knowledge base data, represented by a JSON object", }, }, "required": ["type", "raw_data"], }, "UpdateKnowledgeBaseDataResponse": { "type": "object", "properties": { "knowledge_data_id": { "type": "integer", "description": "ID of the knowledge base data", }, "raw_data": { "type": "object", "description": "The details of the knowledge base data, represented by a JSON object", }, "type": { "type": "string", "description": "Type of the knowledge base data", }, "app_id": { "type": "integer", "description": "ID of the app", }, "meta_data": { "type": "object", "description": "The meta data of the knowledge base data, represented by a JSON object", }, }, "required": ["knowledge_data_id", "raw_data", "type", "app_id"], }, "DeleteKnowledgeBaseDataResponse": { "type": "object", "properties": {}, }, "MultiRoundChat2SQL": { "type": "object", "properties": { "question": { "type": "string", "description": "Question to be resolved", }, "feedback_answer_id": { "type": "string", "description": "ID of the feedback answer", }, "feedback_task_id": { "type": "string", "description": "ID of the feedback task", }, "sql_generate_mode": { "type": "string", "description": "SQL generation mode, support mode: 'direct' or 'auto_breakdown', by default 'direct'", }, }, "required": ["question"], }, "MultiRoundChat2SQLResponse": { "type": "object", "properties": { "job_id": { "type": "string", "description": "ID of the async generation job", }, "session_id": { "type": "number", "description": "ID of the session", }, }, "required": ["job_id", "session_id"], }, "MultiRoundFeedback": { "type": "object", "properties": { "session_id": { "type": "number", "description": "ID of the session", }, "feedback": { "type": "string", "description": "Feedback to be sent, support feedback: 'like' or 'dislike'", }, "query_id": { "type": "string", "description": "ID of the query", }, }, "required": ["session_id", "feedback", "query_id"], }, "MultiRoundFeedbackResponse": { "type": "object", "properties": { "job_id": { "type": "string", "description": "ID of the async generation job", }, "session_id": { "type": "number", "description": "ID of the session", }, }, "required": ["job_id", "session_id"], }, "SessionContexts": { "type": "object", "properties": { "sessions": { "type": "array", "items": { "type": "object", "properties": { "session_id": { "type": "number", "description": "ID of the session", }, "meta": { "type": "object", "description": "The meta data of the session, represented by a JSON object", "properties": { "name": { "type": "string", "description": "Name of the session", }, "created_at": { "type": "integer", "description": "Creation time of the session", }, "updated_at": { "type": "integer", "description": "Update time of the session", }, } }, }, }, "count": { "type": "number", "description": "The number of sessions", }, } }, "required": ["session_id", "session_context"], }, "CreateSession": { "type": "object", "properties": { "cluster_id": { "type": "string", "description": "ID of your cluster in TiDB Cloud", }, "database": { "type": "string", "description": "Name of your database", }, "name": { "type": "string", "description": "Name of the session", }, }, "required": ["cluster_id", "database", "session_id"], }, "SessionContextDetails": { "type": "object", "properties": { "session_id": { "type": "number", "description": "ID of the session", }, "meta": { "type": "object", "description": "The meta data of the session, represented by a JSON object", "properties": { "name": { "type": "string", "description": "Name of the session", }, "created_at": { "type": "integer", "description": "Creation time of the session", }, "updated_at": { "type": "integer", "description": "Update time of the session", }, "org_id": { "type": "string", "description": "ID of the organization", }, "creator": { "type": "string", "description": "Creator of the session", }, } }, "messages": { "type": "array", "items": { "type": "object", "properties": { "created_at": { "type": "number", "description": "Creation time of the message", }, "question": { "type": "object", "description": "The question of the message, represented by a JSON object", "properties": { "question": { "type": "string", "description": "Content of the question", }, "feedback_answer_id": { "type": "string", "description": "ID of the feedback answer", }, "feedback_task_id": { "type": "string", "description": "ID of the feedback task", }, } }, "answers": { "type": "array", "items": { "type": "object", "properties": { "agent": { "type": "string", "description": "Agent of the answer", }, "status": { "type": "string", "description": "Status of the answer", }, "answer_id": { "type": "string", "description": "ID of the answer", }, "tasks": { "type": "object", "description": "The details of the tasks, represented by a JSON object", "properties": { "type": { "type": "string", "description": "Type of the task", }, "clarified_task": { "type": "string", "description": "Clarified task", }, "description": { "type": "string", "description": "Description of the task", }, "data": { "type": "object", "description": "The details of the data, represented by a JSON object", "properties": { "columns": { "type": "array", "description": "The details of the columns, represented by an array", "items": { "type": "object", "properties": { "col": { "type": "string", "description": "Name of the column", }, } } }, "rows": { "type": "array", "description": "The details of the rows, represented by an array", "items": { "type": "any", }, }, } }, "sql": { "type": "string", "description": "SQL of the task", }, "task_id": { "type": "string", "description": "ID of the task", }, "sub_tasks": { "type": "array", "description": "The details of the sub tasks, represented by an array", "items": { "type": "any", "description": "The details of the sub task, represented by a JSON object", } }, } } } } } } } }, }, "required": ["session_id", "session_context"], }, "UpdateSession": { "type": "object", "properties": { "name": { "type": "string", "description": "Name of the session", }, }, "required": ["name"], }, "RefineSQL": { "type": "object", "properties": { "session_id": { "type": "number", "description": "ID of the session", }, "answer_id": { "type": "string", "description": "ID of the answer", }, "task_id": { "type": "string", "description": "ID of the task", }, "sql": { "type": "string", "description": "SQL to be refined", }, "feedback": { "type": "string", "description": "Feedback to be sent, support feedback: 'like' or 'dislike'", }, }, "required": ["session_id", "sql"], }, "RefineSQLResponse": { "type": "object", "properties": { "corrected_sql": { "type": "string", "description": "Corrected SQL", }, "purpose": { "type": "string", "description": "Purpose of the SQL", }, "rewritten_sql": { "type": "string", "description": "Rewritten SQL", }, "solution": { "type": "string", "description": "Solution of the SQL", }, "session_id": { "type": "number", "description": "ID of the session", }, }, "required": ["job_id", "session_id"], }, } swagger_config = { "headers": [ ], "optional_fields": ['components'], "specs": [ { "endpoint": 'edaservice_api', "route": '/eda/apidocs/edaservice_api.json', "rule_filter": rule_filter, "model_filter": model_filter, } ], "static_url_path": "/eda/apidocs/static", # "static_folder": "static", # must be set by user "swagger_ui": False, "specs_route": "/eda/apidocs/", "components": { "securitySchemes": { "basicAuth": { "type": "http", "scheme": "basic", "description": "Enter your public key for the username field and private key for the password field", } }, "schemas": schemas, } } swag = Swagger(app, template=template, config=swagger_config) @app.get("/eda/apidocs/swagger.json") def get_swagger_json(): user_email = get_current_user_email() server_url = request.args.get("server_url", "") logging.info("user_email: %s, server_url: %s", user_email, server_url) data = swag.get_apispecs("edaservice_api") paths = data["paths"] new_paths = {} ignore_prefixs = ["/inner/eda/api/v1beta/app/{data_app_id}/endpoint"] ignore_paths = { "/v2/feedback", "/v2/executeQuery", "/v3/multiRoundFeedback", } for k, v in paths.items(): for ignore_prefix in ignore_prefixs: if k.startswith(ignore_prefix): k = k.removeprefix(ignore_prefix) if k in ignore_paths: continue new_paths[k] = v # pop schemas if not pingcap employee components = data.get("components") if components: schemas = components.get("schemas") if schemas: schemas.pop("ExecuteQuery", None) schemas.pop("ExecuteQueryResponse", None) schemas.pop("UserFeedback", None) schemas.pop("UserFeedbackResponse", None) schemas.pop("MultiRoundFeedback", None) schemas.pop("MultiRoundFeedbackResponse", None) data["paths"] = new_paths data["servers"] = [{"url": server_url}] data["security"] = [{"basicAuth": []}] data.pop("swagger", None) return data, 200 if __name__ == "__main__": app.run()