import asyncio import argparse import uuid from agents import EmbedderAgent, BuildSQLAgent, DebugSQLAgent, ValidateSQLAgent, ResponseAgent,VisualizeAgent from utilities import (PROJECT_ID, PG_REGION, BQ_REGION, EXAMPLES, LOGGING, VECTOR_STORE, BQ_OPENDATAQNA_DATASET_NAME, USE_SESSION_HISTORY) from dbconnectors import bqconnector, pgconnector, firestoreconnector from embeddings.store_embeddings import add_sql_embedding #Based on VECTOR STORE in config.ini initialize vector connector and region if VECTOR_STORE=='bigquery-vector': region=BQ_REGION vector_connector = bqconnector call_await = False elif VECTOR_STORE == 'cloudsql-pgvector': region=PG_REGION vector_connector = pgconnector call_await=True else: raise ValueError("Please specify a valid Data Store. Supported are either 'bigquery-vector' or 'cloudsql-pgvector'") def generate_uuid(): """Generates a random UUID (Universally Unique Identifier) Version 4. Returns: str: A string representation of the UUID in the format xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx. """ return str(uuid.uuid4()) ############################ #_____GET ALL DATABASES_____# ############################ def get_all_databases(): """Retrieves a list of all distinct databases (with source type) from the vector store. This function queries the vector store (BigQuery or PostgreSQL) to fetch a list of unique databases, including their source type. The source type indicates whether the database is a BigQuery dataset or a PostgreSQL schema. Returns: tuple: A tuple containing two elements: - result (str or list): A JSON-formatted string containing the list of databases and their source types, or an error message if an exception occurs. - invalid_response (bool): A flag indicating whether an error occurred during retrieval (True) or if the response is valid (False). Raises: Exception: If there is an issue connecting to or querying the vector store. The exception message will be included in the returned `result`. """ try: if VECTOR_STORE=='bigquery-vector': final_sql=f'''SELECT DISTINCT user_grouping AS table_schema FROM `{PROJECT_ID}.{BQ_OPENDATAQNA_DATASET_NAME}.table_details_embeddings`''' else: final_sql="""SELECT DISTINCT user_grouping AS table_schema FROM table_details_embeddings""" result = vector_connector.retrieve_df(final_sql) result = result.to_json(orient='records') invalid_response=False except Exception as e: result="Issue was encountered while extracting databases in vector store:: " + str(e) invalid_response=True return result,invalid_response ############################ #_____GET SOURCE TYPE_____## ############################ def get_source_type(user_grouping): """Retrieves the source type of a specified database from the vector store. This function queries the vector store (BigQuery or PostgreSQL) to determine whether the given database is a BigQuery dataset ('bigquery') or a PostgreSQL schema ('postgres'). Args: user_grouping (str): The name of the database to look up. Returns: tuple: A tuple containing two elements: - result (str): The source type of the database ('bigquery' or 'postgres'), or an error message if not found or an exception occurs. - invalid_response (bool): A flag indicating whether an error occurred during retrieval (True) or if the response is valid (False). Raises: Exception: If there is an issue connecting to or querying the vector store. The exception message will be included in the returned `result`. """ try: if VECTOR_STORE=='bigquery-vector': sql=f'''SELECT DISTINCT source_type FROM `{PROJECT_ID}.{BQ_OPENDATAQNA_DATASET_NAME}.table_details_embeddings` where user_grouping='{user_grouping}' ''' else: sql=f'''SELECT DISTINCT source_type FROM table_details_embeddings where user_grouping='{user_grouping}' ''' result = vector_connector.retrieve_df(sql) result = (str(result.iloc[0, 0])).lower() invalid_response=False except Exception as e: result="Error at finding the datasource :: "+str(e) invalid_response=True return result,invalid_response ############################ ###_____GENERATE SQL_____### ############################ async def generate_sql(session_id, user_question, user_grouping, RUN_DEBUGGER, DEBUGGING_ROUNDS, LLM_VALIDATION, Embedder_model, SQLBuilder_model, SQLChecker_model, SQLDebugger_model, num_table_matches, num_column_matches, table_similarity_threshold, column_similarity_threshold, example_similarity_threshold, num_sql_matches, user_id="opendataqna-user@google.com"): """Generates an SQL query based on a user's question and database. This asynchronous function orchestrates a pipeline to generate an SQL query from a natural language question. It leverages various agents for embedding, SQL building, validation, and debugging. Args: session_id (str): Session ID to identify the chat conversation user_question (str): The user's natural language question. user_grouping (str): The name of the database to query. RUN_DEBUGGER (bool): Whether to run the SQL debugger. DEBUGGING_ROUNDS (int): The number of debugging rounds to perform. LLM_VALIDATION (bool): Whether to use LLM for validation. Embedder_model (str): The name of the embedding model. SQLBuilder_model (str): The name of the SQL builder model. SQLChecker_model (str): The name of the SQL checker model. SQLDebugger_model (str): The name of the SQL debugger model. num_table_matches (int): The number of table matches to retrieve. num_column_matches (int): The number of column matches to retrieve. table_similarity_threshold (float): The similarity threshold for table matching. column_similarity_threshold (float): The similarity threshold for column matching. example_similarity_threshold (float): The similarity threshold for example matching. num_sql_matches (int): The number of similar SQL queries to retrieve. Returns: tuple: A tuple containing: - final_sql (str): The final generated SQL query, or an error message if generation failed. - invalid_response (bool): True if the response is invalid (e.g., due to an error), False otherwise. """ try: if session_id is None or session_id=="": print("This is a new session") session_id=generate_uuid() ## LOAD AGENTS print("Loading Agents.") embedder = EmbedderAgent(Embedder_model) SQLBuilder = BuildSQLAgent(SQLBuilder_model) SQLChecker = ValidateSQLAgent(SQLChecker_model) SQLDebugger = DebugSQLAgent(SQLDebugger_model) re_written_qe=user_question print("Getting the history for the session.......\n") session_history = firestoreconnector.get_chat_logs_for_session(session_id) if USE_SESSION_HISTORY else None print("Grabbed history for the session:: "+ str(session_history)) if session_history is None or not session_history: print("No records for the session. Not rewriting the question\n") else: concated_questions,re_written_qe=SQLBuilder.rewrite_question(user_question,session_history) found_in_vector = 'N' # if an exact query match was found final_sql='Not Generated Yet' # final generated SQL process_step='Not Started' error_msg='' corrected_sql = '' DATA_SOURCE = 'Yet to determine' DATA_SOURCE,src_invalid = get_source_type(user_grouping) if src_invalid: raise ValueError(DATA_SOURCE) #vertexai.init(project=PROJECT_ID, location=region) #aiplatform.init(project=PROJECT_ID, location=region) print("Source selected as : "+ str(DATA_SOURCE) + "\nSchema or Dataset Name is : "+ str(user_grouping)) print("Vector Store selected as : "+ str(VECTOR_STORE)) # Reset AUDIT_TEXT AUDIT_TEXT = 'Creating embedding for given question' # Fetch the embedding of the user's input question embedded_question = embedder.create(re_written_qe) AUDIT_TEXT = AUDIT_TEXT + "\nUser Question : " + str(user_question) + "\nUser Database : " + str(user_grouping) process_step = "\n\nGet Exact Match: " # Look for exact matches in known questions IF kgq is enabled if EXAMPLES: exact_sql_history = vector_connector.getExactMatches(user_question) else: exact_sql_history = None # If exact user query has been found, retrieve the SQL and skip Generation Pipeline if exact_sql_history is not None: found_in_vector = 'Y' final_sql = exact_sql_history invalid_response = False AUDIT_TEXT = AUDIT_TEXT + "\nExact match has been found! Going to retrieve the SQL query from cache and serve!" else: # No exact match found. Proceed looking for similar entries in db IF kgq is enabled if EXAMPLES: AUDIT_TEXT = AUDIT_TEXT + process_step + "\nNo exact match found in query cache, retrieving relevant schema and known good queries for few shot examples using similarity search...." process_step = "\n\nGet Similar Match: " if call_await: similar_sql = await vector_connector.getSimilarMatches('example', user_grouping, embedded_question, num_sql_matches, example_similarity_threshold) else: similar_sql = vector_connector.getSimilarMatches('example', user_grouping, embedded_question, num_sql_matches, example_similarity_threshold) else: similar_sql = "No similar SQLs provided..." process_step = "\n\nGet Table and Column Schema: " # Retrieve matching tables and columns if call_await: table_matches = await vector_connector.getSimilarMatches('table', user_grouping, embedded_question, num_table_matches, table_similarity_threshold) column_matches = await vector_connector.getSimilarMatches('column', user_grouping, embedded_question, num_column_matches, column_similarity_threshold) else: table_matches = vector_connector.getSimilarMatches('table', user_grouping, embedded_question, num_table_matches, table_similarity_threshold) column_matches = vector_connector.getSimilarMatches('column', user_grouping, embedded_question, num_column_matches, column_similarity_threshold) AUDIT_TEXT = AUDIT_TEXT + process_step + "\nRetrieved Similar Known Good Queries, Table Schema and Column Schema: \n" + '\nRetrieved Tables: \n' + str(table_matches) + '\n\nRetrieved Columns: \n' + str(column_matches) + '\n\nRetrieved Known Good Queries: \n' + str(similar_sql) # If similar table and column schemas found: if len(table_matches.replace('Schema(values):','').replace(' ','')) > 0 or len(column_matches.replace('Column name(type):','').replace(' ','')) > 0 : # GENERATE SQL process_step = "\n\nBuild SQL: " generated_sql = SQLBuilder.build_sql(DATA_SOURCE,user_grouping,user_question,session_history,table_matches,column_matches,similar_sql) final_sql=generated_sql AUDIT_TEXT = AUDIT_TEXT + process_step + "\nGenerated SQL : " + str(generated_sql) if 'unrelated_answer' in generated_sql : invalid_response=True final_sql="This is an unrelated question or you are not asking a valid query" # If agent assessment is valid, proceed with checks else: invalid_response=False if RUN_DEBUGGER: generated_sql, invalid_response, AUDIT_TEXT = SQLDebugger.start_debugger(DATA_SOURCE,user_grouping, generated_sql, user_question, SQLChecker, table_matches, column_matches, AUDIT_TEXT, similar_sql, DEBUGGING_ROUNDS, LLM_VALIDATION) # AUDIT_TEXT = AUDIT_TEXT + '\n Feedback from Debugger: \n' + feedback_text final_sql=generated_sql AUDIT_TEXT = AUDIT_TEXT + "\nFinal SQL after Debugger: \n" +str(final_sql) # No matching table found else: invalid_response=True print('No tables found in Vector ...') AUDIT_TEXT = AUDIT_TEXT + "\nNo tables have been found in the Vector DB. The question cannot be answered with the provide data source!" # print(f'\n\n AUDIT_TEXT: \n {AUDIT_TEXT}') if LOGGING: bqconnector.make_audit_entry(DATA_SOURCE, user_grouping, SQLBuilder_model, user_question, final_sql, found_in_vector, "", process_step, error_msg,AUDIT_TEXT) except Exception as e: error_msg=str(e) final_sql="Error generating the SQL Please check the logs. "+str(e) invalid_response=True AUDIT_TEXT=AUDIT_TEXT+ "\nException at SQL generation" print("Error :: "+str(error_msg)) if LOGGING: bqconnector.make_audit_entry(DATA_SOURCE, user_grouping, SQLBuilder_model, user_question, final_sql, found_in_vector, "", process_step, error_msg,AUDIT_TEXT) if USE_SESSION_HISTORY and not invalid_response: firestoreconnector.log_chat(session_id,user_question,final_sql,user_id) print("Session history persisted") return final_sql,session_id,invalid_response ############################ ###_____GET RESULTS_____#### ############################ def get_results(user_grouping, final_sql, invalid_response=False, EXECUTE_FINAL_SQL=True): """Executes the final SQL query (if valid) and retrieves the results. This function first determines the data source (BigQuery or PostgreSQL) based on the provided database name. If the SQL query is valid and execution is enabled, it fetches the results using the appropriate connector. Args: user_grouping (str): The name of the database to query. final_sql (str): The final SQL query to execute. invalid_response (bool, optional): A flag indicating whether the SQL query is invalid. Defaults to False. EXECUTE_FINAL_SQL (bool, optional): Whether to execute the final SQL query. Defaults to True. Returns: tuple: A tuple containing: - result_df (pandas.DataFrame or str): The results of the SQL query as a DataFrame, or an error message if the query is invalid or execution failed. - invalid_response (bool): True if the response is invalid (e.g., due to an error), False otherwise. Raises: ValueError: If the data source is invalid or not supported. Exception: If there's an error executing the SQL query or retrieving the results. """ try: DATA_SOURCE,src_invalid = get_source_type(user_grouping) if not src_invalid: ## SET DATA SOURCE if DATA_SOURCE=='bigquery': src_connector = bqconnector else: src_connector = pgconnector else: raise ValueError(DATA_SOURCE) if not invalid_response: try: if EXECUTE_FINAL_SQL is True: final_exec_result_df=src_connector.retrieve_df(final_sql.replace("```sql","").replace("```","").replace("EXPLAIN ANALYZE ","")) result_df = final_exec_result_df else: # Do not execute final SQL print("Not executing final SQL since EXECUTE_FINAL_SQL variable is False\n ") result_df = "Please enable the Execution of the final SQL so I can provide an answer" invalid_response = True except ValueError: result_df= "Error has been encountered :: " + str(e) invalid_response=True else: # Do not execute final SQL result_df = "Not executing final SQL as it is invalid, please debug!" except Exception as e: print(f"An error occured. Aborting... Error Message: {e}") result_df="Error has been encountered :: " + str(e) invalid_response=True return result_df,invalid_response def get_response(session_id,user_question,result_df,Responder_model='gemini-1.0-pro'): try: Responder = ResponseAgent(Responder_model) if session_id is None or session_id=="": print("This is a new session") else: session_history =firestoreconnector.get_chat_logs_for_session(session_id) if USE_SESSION_HISTORY else None if session_history is None or not session_history: print("No records for the session. Not rewriting the question\n") else: concated_questions,re_written_qe=Responder.rewrite_question(user_question,session_history) user_question=re_written_qe _resp=Responder.run(user_question, result_df) invalid_response=False except Exception as e: print(f"An error occured. Aborting... Error Message: {e}") _resp= "Error has been encountered :: " + str(e) invalid_response=True return _resp,invalid_response ############################ ###_____RUN PIPELINE_____### ############################ async def run_pipeline(session_id, user_question, user_grouping, RUN_DEBUGGER=True, EXECUTE_FINAL_SQL=True, DEBUGGING_ROUNDS = 2, LLM_VALIDATION=False, Embedder_model='vertex', SQLBuilder_model= 'gemini-1.5-pro', SQLChecker_model= 'gemini-1.0-pro', SQLDebugger_model= 'gemini-1.0-pro', Responder_model= 'gemini-1.0-pro', num_table_matches = 5, num_column_matches = 10, table_similarity_threshold = 0.3, column_similarity_threshold = 0.3, example_similarity_threshold = 0.3, num_sql_matches=3): """Orchestrates the end-to-end SQL generation and response pipeline. This asynchronous function manages the entire process of generating an SQL query from a user's question, executing the query (if valid), and formulating a natural language response based on the results. Args: user_question (str): The user's natural language question. user_grouping (str): The name of the user grouping to query. RUN_DEBUGGER (bool, optional): Whether to run the SQL debugger. Defaults to True. EXECUTE_FINAL_SQL (bool, optional): Whether to execute the final SQL query. Defaults to True. DEBUGGING_ROUNDS (int, optional): The number of debugging rounds to perform. Defaults to 2. LLM_VALIDATION (bool, optional): Whether to use LLM for validation. Defaults to True. Embedder_model (str, optional): The name of the embedding model. Defaults to 'vertex'. SQLBuilder_model (str, optional): The name of the SQL builder model. Defaults to 'gemini-1.5-pro'. SQLChecker_model (str, optional): The name of the SQL checker model. Defaults to 'gemini-1.0-pro'. SQLDebugger_model (str, optional): The name of the SQL debugger model. Defaults to 'gemini-1.0-pro'. Responder_model (str, optional): The name of the responder model. Defaults to 'gemini-1.0-pro'. num_table_matches (int, optional): The number of table matches to retrieve. Defaults to 5. num_column_matches (int, optional): The number of column matches to retrieve. Defaults to 10. table_similarity_threshold (float, optional): The similarity threshold for table matching. Defaults to 0.3. column_similarity_threshold (float, optional): The similarity threshold for column matching. Defaults to 0.3. example_similarity_threshold (float, optional): The similarity threshold for example matching. Defaults to 0.3. num_sql_matches (int, optional): The number of similar SQL queries to retrieve. Defaults to 3. Returns: tuple: A tuple containing: - final_sql (str): The final generated SQL query, or an error message if generation failed. - results_df (pandas.DataFrame or str): The results of the SQL query as a DataFrame, or an error message if the query is invalid or execution failed. - _resp (str): The generated natural language response based on the results, or an error message if response generation failed. """ final_sql,session_id, invalid_response = await generate_sql(session_id, user_question, user_grouping, RUN_DEBUGGER, DEBUGGING_ROUNDS, LLM_VALIDATION, Embedder_model, SQLBuilder_model, SQLChecker_model, SQLDebugger_model, num_table_matches, num_column_matches, table_similarity_threshold, column_similarity_threshold, example_similarity_threshold, num_sql_matches) if not invalid_response: results_df, invalid_response = get_results(user_grouping, final_sql, invalid_response=invalid_response, EXECUTE_FINAL_SQL=EXECUTE_FINAL_SQL) if not invalid_response: _resp,invalid_response=get_response(session_id,user_question,results_df.to_json(orient='records'),Responder_model=Responder_model) else: _resp=results_df else: results_df=final_sql _resp=final_sql return final_sql, results_df, _resp ############################ #####_____GET KGQ_____###### ############################ def get_kgq(user_grouping): """Retrieves known good SQL queries (KGQs) for a specific database from the vector store. This function queries the vector store (BigQuery or PostgreSQL) to fetch a limited number of distinct user questions and their corresponding generated SQL queries that are relevant to the specified database. These KGQs can be used as examples or references for generating new SQL queries. Args: user_grouping (str): The name of the user grouping for which to retrieve KGQs. Returns: tuple: A tuple containing two elements: - result (str): A JSON-formatted string containing the list of KGQs (user questions and SQL queries), or an error message if an exception occurs. - invalid_response (bool): A flag indicating whether an error occurred during retrieval (True) or if the response is valid (False). Raises: Exception: If there is an issue connecting to or querying the vector store. The exception message will be included in the returned `result`. """ try: if VECTOR_STORE=='bigquery-vector': sql=f'''SELECT distinct example_user_question, example_generated_sql FROM `{PROJECT_ID}.{BQ_OPENDATAQNA_DATASET_NAME}.example_prompt_sql_embeddings` where user_grouping='{user_grouping}' LIMIT 5 ''' else: sql="""select distinct example_user_question, example_generated_sql from example_prompt_sql_embeddings where user_grouping = '{user_grouping}' LIMIT 5""".format(user_grouping=user_grouping) result = vector_connector.retrieve_df(sql) result = result.to_json(orient='records') invalid_response = False except Exception as e: result="Issue was encountered while extracting known good sqls in vector store:: " + str(e) invalid_response=True return result,invalid_response ############################ ####_____EMBED SQL_____##### ############################ async def embed_sql(session_id,user_grouping,user_question,generate_sql): """Embeds a generated SQL query into the vector store as an example. This asynchronous function takes a user's question, a generated SQL query, and a database name as input. It calls the `add_sql_embedding` function to create an embedding of the SQL query and store it in the vector store, potentially for future reference as a known good query (KGQ). Args: user_grouping (str): The name of the grouping associated with the query. user_question (str): The user's original question. generate_sql (str): The SQL query generated from the user's question. Returns: tuple: A tuple containing two elements: - embedded (str or None): The embedded SQL query if successful, or an error message if an exception occurs. - invalid_response (bool): A flag indicating whether an error occurred during embedding (True) or if the response is valid (False). Raises: Exception: If there is an issue with the embedding process. The exception message will be included in the returned `embedded` value. """ try: Rewriter=ResponseAgent('gemini-1.5-pro') if session_id is None or session_id=="": print("This is a new session") else: session_history =firestoreconnector.get_chat_logs_for_session(session_id) if USE_SESSION_HISTORY else None if session_history is None or not session_history: print("No records for the session. Not rewriting the question\n") else: concated_questions,re_written_qe=Rewriter.rewrite_question(user_question,session_history) user_question=re_written_qe embedded = await add_sql_embedding(user_question, generate_sql,user_grouping) invalid_response=False except Exception as e: embedded="Issue was encountered while embedding the SQL as example." + str(e) invalid_response=True return embedded,invalid_response def visualize(session_id,user_question,generated_sql,sql_results): try: Rewriter=ResponseAgent('gemini-1.5-pro') if session_id is None or session_id=="": print("This is a new session") else: session_history =firestoreconnector.get_chat_logs_for_session(session_id) if USE_SESSION_HISTORY else None if session_history is None or not session_history: print("No records for the session. Not rewriting the question\n") else: concated_questions,re_written_qe=Rewriter.rewrite_question(user_question,session_history) user_question=re_written_qe _viz=VisualizeAgent() js_chart = _viz.generate_charts(user_question, generate_sql,sql_results) invalid_response=False except Exception as e: js_chart="Issue was encountered while Generating Charts ::" + str(e) invalid_response=True return js_chart,invalid_response ############################ #######_____MAIN_____####### ############################ if __name__ == '__main__': # user_question = "How many movies have review ratings above 5?" # user_grouping='MovieExplorer-bigquery' parser = argparse.ArgumentParser(description="Open Data QnA SQL Generation") parser.add_argument("--session_id", type=str, required=True, help="Session Id") parser.add_argument("--user_question", type=str, required=True, help="The user's question.") parser.add_argument("--user_grouping", type=str, required=True, help="The user grouping specificed in the source list CSV file") # Optional Arguments for run_pipeline Parameters parser.add_argument("--run_debugger", action="store_true", help="Enable the debugger (default: False)") parser.add_argument("--execute_final_sql", action="store_true", help="Execute the final SQL (default: False)") parser.add_argument("--debugging_rounds", type=int, default=2, help="Number of debugging rounds (default: 2)") parser.add_argument("--llm_validation", action="store_true", help="Enable LLM validation (default: False)") parser.add_argument("--embedder_model", type=str, default='vertex', help="Embedder model name (default: 'vertex')") parser.add_argument("--sqlbuilder_model", type=str, default='gemini-1.5-pro', help="SQL builder model name (default: 'gemini-1.0-pro')") parser.add_argument("--sqlchecker_model", type=str, default='gemini-1.5-pro', help="SQL checker model name (default: 'gemini-1.0-pro')") parser.add_argument("--sqldebugger_model", type=str, default='gemini-1.5-pro', help="SQL debugger model name (default: 'gemini-1.0-pro')") parser.add_argument("--responder_model", type=str, default='gemini-1.5-pro', help="Responder model name (default: 'gemini-1.0-pro')") parser.add_argument("--num_table_matches", type=int, default=5, help="Number of table matches (default: 5)") parser.add_argument("--num_column_matches", type=int, default=10, help="Number of column matches (default: 10)") parser.add_argument("--table_similarity_threshold", type=float, default=0.1, help="Threshold for table similarity (default: 0.1)") parser.add_argument("--column_similarity_threshold", type=float, default=0.1, help="Threshold for column similarity (default: 0.1)") parser.add_argument("--example_similarity_threshold", type=float, default=0.1, help="Threshold for example similarity (default: 0.1)") parser.add_argument("--num_sql_matches", type=int, default=3, help="Number of SQL matches (default: 3)") args = parser.parse_args() # Use Argument Values in run_pipeline final_sql, response, _resp = asyncio.run(run_pipeline(args.session_id, args.user_question, args.user_grouping, RUN_DEBUGGER=args.run_debugger, EXECUTE_FINAL_SQL=args.execute_final_sql, DEBUGGING_ROUNDS=args.debugging_rounds, LLM_VALIDATION=args.llm_validation, Embedder_model=args.embedder_model, SQLBuilder_model=args.sqlbuilder_model, SQLChecker_model=args.sqlchecker_model, SQLDebugger_model=args.sqldebugger_model, Responder_model=args.responder_model, num_table_matches=args.num_table_matches, num_column_matches=args.num_column_matches, table_similarity_threshold=args.table_similarity_threshold, column_similarity_threshold=args.column_similarity_threshold, example_similarity_threshold=args.example_similarity_threshold, num_sql_matches=args.num_sql_matches )) # user_question = "How many +18 movies have a rating above 4?" # final_sql, response, _resp = asyncio.run(run_pipeline(user_question, # 'imdb', # RUN_DEBUGGER=True, # EXECUTE_FINAL_SQL=True, # DEBUGGING_ROUNDS = 2, # LLM_VALIDATION=True, # Embedder_model='vertex', # SQLBuilder_model= 'gemini-1.0-pro', # SQLChecker_model= 'gemini-1.0-pro', # SQLDebugger_model= 'gemini-1.0-pro', # Responder_model= 'gemini-1.0-pro', # num_table_matches = 5, # num_column_matches = 10, # table_similarity_threshold = 0.1, # column_similarity_threshold = 0.1, # example_similarity_threshold = 0.1, # num_sql_matches=3)) print("*"*50 +"\nGenerated SQL\n"+"*"*50+"\n"+final_sql) print("\n"+"*"*50 +"\nResults\n"+"*"*50) print(response) print("*"*50 +"\nNatural Response\n"+"*"*50+"\n"+_resp)