def generate_sql_and_query()

in genai-for-marketing/backend_apis/app/utils_codey.py [0:0]

• 33 lines of code
• 2 McCabe index (conditional complexity)

def generate_sql_and_query(
        llm,
        datacatalog_client,
        prompt_template: str,
        query_metadata: str,
        question: str,
        project_id: str,
        dataset_id: str,
        tag_template_name: str,
        bqclient: bigquery.Client):
    """Generates a GoogleSQL query and executes it against a BigQuery dataset.

    Args:
        query: 
            The initial query text.
        project_id: 
            The ID of the BigQuery project.
        dataset_id: 
            The ID of the BigQuery dataset.
        tag_template_name: 
            The name of the tag template to use for the query.
        bqclient: 
            A BigQuery client object.

    Returns:
        A DataFrame containing the results of the query.

    Raises:
        NotFoundError: If the dataset or table is not found.
        BadRequestError: If the query is invalid.
    """

    metadata = get_metadata_from_dataset(
        bqclient=bqclient,
        datacatalog_client=datacatalog_client,
        query=query_metadata, 
        project_id=project_id, 
        dataset_id=dataset_id,
        tag_template_name=tag_template_name)

    prompt = generate_prompt(
        question, 
        metadata,
        prompt_template,
        project_id)

    gen_code = llm.predict(
        prompt = prompt,
        max_output_tokens = 1024,
        temperature=0.3
    ).text.replace("```","")
    gen_code = gen_code[gen_code.find("SELECT"):]
    result = []
    result_job = bqclient.query(gen_code)
    for row in result_job:
        result.append(dict(row.items()))
    return result, gen_code, prompt