sql_qa_pair/app/endpoints/generate_qa_pair.py

from fastapi import APIRouter, HTTPException

from app.const import MAX_INPUT_CARD_IDS
from app.endpoints.schema import QARequest, QAResponse, QAResponseBody
from app.pipeline.workflow import Workflow


router = APIRouter()


@router.post("/generate_qa_pair", response_model=QAResponse)
async def generate_qa_pair(request_body: QARequest):
    if len(request_body.card_ids) > MAX_INPUT_CARD_IDS:
        raise HTTPException(
            status_code=400,
            detail=f"Number of card IDs exceeds the maximum limit of {MAX_INPUT_CARD_IDS}.",
        )
    try:
        workflow = Workflow()
        input_args = {
            "dashboard_id": request_body.dashboard_id,
            "card_ids": request_body.card_ids,
            "bbk": request_body.bbk_id,
            "user_request": request_body.user_request or "",
        }
        result = await workflow.execute_workflow(input_args)
        qa_response = QAResponse(
            returnCode="SUCCESS",
            body=QAResponseBody(
                dashboard_id=request_body.dashboard_id,
                card_ids=request_body.card_ids,
                qa_pairs=result,
                bbk_id=request_body.bbk_id,
            ),
        )
        return qa_response
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")