def call_generate_sql_api()

in UI/utils.py [0:0]

• 59 lines of code
• 10 McCabe index (conditional complexity)

def call_generate_sql_api(question, endpoint) -> tuple[str, str]:
    """
        Common SQL generation function
    """
    # api_url = os.getenv('CORE_EXECUTORS')
    selected_framework = None
    if LITE_API_PART in endpoint:
        selected_framework = "LITE"
        api_url = os.getenv('LITE_EXECUTORS')
        few_shot_gen = False
        if st.session_state.lite_model == FEW_SHOT_GENERATION:
            few_shot_gen = True
        data = {"question": question,
                "execute_sql": st.session_state.execution,
                "few_shot": few_shot_gen}
        st.session_state.sql_generated_by = GEN_BY_LITE
    else:
        selected_framework = "CORE"
        api_url = os.getenv('CORE_EXECUTORS')
        data = {"question": question,
                "execute_sql": st.session_state.execution}
        st.session_state.sql_generated_by = GEN_BY_CORE

    headers = {"Content-type": "application/json",
               "Authorization": f"Bearer {st.session_state.access_token}"}
    api_endpoint = f"{api_url}/{endpoint}"

    logger.info(f"Invoking API : {api_endpoint}")
    logger.info(f"Provided parameters are : Data = {data}")
    api_response = requests.post(api_endpoint,
                                 data=json.dumps(data),
                                 headers=headers,
                                 timeout=None)

    logger.info(f"{str(api_response.text)}")
    # logger.info('-'*100, api_response.text)
    exec_result = ""
    try:
        resp = api_response.json()

        sql = resp['generated_query']
        st.session_state.result_id = resp['result_id']

        if selected_framework == 'CORE' and data["execute_sql"]:
            try:
                # Parse the JSON string to a dictionary
                df_dict = json.loads(resp['df'])
                exec_result_df = pd.DataFrame(df_dict)
                exec_result = ''
            except Exception as e:
                logger.info(f"SQL query execution failed due to {str(e)}")
                exec_result_df = None
                exec_result = ''

        elif selected_framework == 'LITE' and data["execute_sql"]:
            exec_result = resp['sql_result']
            try:
                exec_result_df = run_query(sql)
            except Exception as e:
                logger.info(f"SQL query execution failed due to {str(e)}")
                exec_result_df = None
                exec_result = None
        else:
            exec_result_df = None
            exec_result = ''

    except RuntimeError:
        sql = "Execution Failed ! Error encountered"
        st.session_state.result_id = None

    logger.info(f"Generated SQL = {sql}")
    logger.info(f"Generation ID = {st.session_state.result_id}")
    return sql, exec_result, exec_result_df