in dbconnectors/BQConnector.py [0:0]
def retrieve_matches(self, mode, user_grouping, qe, similarity_threshold, limit):
"""
This function retrieves the most similar table_schema and column_schema.
Modes can be either 'table', 'column', or 'example'
"""
matches = []
if mode == 'table':
sql = '''select base.content as tables_content from vector_search(
(SELECT * FROM `{}.table_details_embeddings` WHERE user_grouping = '{}'), "embedding",
(SELECT {} as qe), top_k=> {},distance_type=>"COSINE") where 1-distance > {} '''
elif mode == 'column':
sql='''select base.content as columns_content from vector_search(
(SELECT * FROM `{}.tablecolumn_details_embeddings` WHERE user_grouping = '{}'), "embedding",
(SELECT {} as qe), top_k=> {}, distance_type=>"COSINE") where 1-distance > {} '''
elif mode == 'example':
sql='''select base.example_user_question, base.example_generated_sql from vector_search (
(SELECT * FROM `{}.example_prompt_sql_embeddings` WHERE user_grouping = '{}'), "embedding",
(select {} as qe), top_k=> {}, distance_type=>"COSINE") where 1-distance > {} '''
else:
ValueError("No valid mode. Must be either table, column, or example")
name_txt = ''
results=self.client.query_and_wait(sql.format('{}.{}'.format(self.project_id,self.opendataqna_dataset),user_grouping,qe,limit,similarity_threshold)).to_dataframe()
# CHECK RESULTS
if len(results) == 0:
print(f"Did not find any results for {mode}. Adjust the query parameters.")
else:
print(f"Found {len(results)} similarity matches for {mode}.")
if mode == 'table':
name_txt = ''
for _ , r in results.iterrows():
name_txt=name_txt+r["tables_content"]+"\n"
elif mode == 'column':
name_txt = ''
for _ ,r in results.iterrows():
name_txt=name_txt+r["columns_content"]+"\n"
elif mode == 'example':
name_txt = ''
for _ , r in results.iterrows():
example_user_question=r["example_user_question"]
example_sql=r["example_generated_sql"]
name_txt = name_txt + "\n Example_question: "+example_user_question+ "; Example_SQL: "+example_sql
else:
ValueError("No valid mode. Must be either table, column, or example")
name_txt = ''
matches.append(name_txt)
return matches