Commit acf93257 authored by Aneeb Imamdin's avatar Aneeb Imamdin

table support added in api

parent 1ba254e7
...@@ -7,8 +7,12 @@ app = Flask(__name__) ...@@ -7,8 +7,12 @@ app = Flask(__name__)
@app.route('/vector-search', methods=['GET']) @app.route('/vector-search', methods=['GET'])
def get_data(): def get_data():
query = request.args.get('query', default='Can i work from home ?') query = request.args.get('query', default='Can i work from home ?')
table = request.args.get('table', default='products_embeddings')
if table != 'products_embeddings' and table != 'product_sentence_embeddings':
return 'Not a valid table.'
query_embeddings = embed_text([query], "RETRIEVAL_DOCUMENT", 256) query_embeddings = embed_text([query], "RETRIEVAL_DOCUMENT", 256)
data = vector_search_in_bigquery(query_embeddings[0]) data = vector_search_in_bigquery(query_embeddings[0], table)
# Return the data as a JSON response # Return the data as a JSON response
return jsonify(data) return jsonify(data)
......
...@@ -99,11 +99,11 @@ def save_to_bq(data): ...@@ -99,11 +99,11 @@ def save_to_bq(data):
print("Errors:", errors) print("Errors:", errors)
def vector_search_in_bigquery(query_embedding): def vector_search_in_bigquery(query_embedding, table):
sql_query = f""" sql_query = f"""
SELECT base.text, distance SELECT base.text, distance
FROM VECTOR_SEARCH( FROM VECTOR_SEARCH(
TABLE ai_practice_dataset.ai_poc_data , 'embeddings', TABLE DATSET.{table} , 'embeddings',
(SELECT {query_embedding} as embed) , top_k => 5, distance_type => 'COSINE') (SELECT {query_embedding} as embed) , top_k => 5, distance_type => 'COSINE')
""" """
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment