Initial Draft

lib/api/endpoints/VectorStoreAPI.py

@@ -1,139 +1,39 @@
-import pandas as pd
-import numpy as np
-from sentence_transformers import SentenceTransformer
-import faiss
-import os
+from flask.views import MethodView
+from flask import request,Response
 
-# Function to create a vector store
-def fn_create_vector_store(lv_file_name, lv_domain):
-    """
-    Create a vector store by encoding and storing embeddings of column descriptions from an Excel file.
+import json
+import traceback
+import logging
 
-    Args:
-        lv_file_name (str): The path to the Excel file.
-        lv_domain (str): The domain name.
+import lib.api.vector.VectorStore as cv
 
-    Returns:
-        str: A message indicating the status of the operation.
-    """
-    # File Names
-    lv_faiss_file_name = 'db/'+lv_domain+'_index.faiss'
-    lv_rowdata_file_name = 'db/'+lv_domain+'_row_mapping.parquet'
+class VectorStoreAPI(MethodView):
+    lv_logger = logging.getLogger(__name__)
+    
+    def get(self):
+        resp = { "test message": "working"}
+        status = 200
+        return Response(json.dumps(resp), status=status, mimetype='application/json')
 
-    # Check if files exist
-    if os.path.exists(lv_faiss_file_name) and os.path.exists(lv_rowdata_file_name):
-        return "Data Already Exist"
-    else:
+    def post(self):
         try:
-            # Load the pre-trained model
-            lv_model = SentenceTransformer('all-MiniLM-L6-v2')
-
-            # Excel to Dataframe
-            lv_excel_data = pd.read_excel(lv_file_name,sheet_name=None)
-            
-            # Dictionary to store Embeddings, Faiss Index, and Index to Row Mapping
-            lv_embeddings_list = []
-            lv_row_mapping = []
-
-            # Reading each sheet
-            for lv_sheet_name, lv_sheet_data in lv_excel_data.items():        
-                # Creating Embeddings
-                # Details available here -> https://www.sbert.net/docs/pretrained_models.html
-                lv_sheet_data.iloc[:, 1] = lv_sheet_data.iloc[:, 1].apply(lambda x: str(x).replace(u'\xa0', u' '))
-                lv_column_descriptions = lv_sheet_data.iloc[:, 1].astype(str).tolist()
+            # Saving file
+            lv_file = request.files['file']
+            lv_domain = request.form['domain']
+            lv_file_name = 'storage/' + lv_domain + ".xlsx"
+            lv_file.save(lv_file_name)
             
-                lv_embeddings = lv_model.encode(lv_column_descriptions).astype('float32')
-                lv_embeddings_list.append(lv_embeddings)
-
-                # Merging all table, columns, description/hint into table
-                for i, row in enumerate(lv_sheet_data.itertuples(index=False)):
-                    lv_row_mapping.append({
-                        'sheet_name': lv_sheet_name,
-                        'column_name': row[0],
-                        'column_description': row[1]
-                    })
-            
-            # Combine all embeddings into one array
-            lv_merged_embeddings_list = np.vstack(lv_embeddings_list)
-
-            # Create a Faiss index
-            lv_dimension = lv_merged_embeddings_list.shape[1]
-            lv_index = faiss.IndexFlatL2(lv_dimension)
-            lv_index.add(lv_merged_embeddings_list)
-
-            # Saving the Faiss index to a file
-            faiss.write_index(lv_index, lv_faiss_file_name)
-
-            # Saving the Row Data to a file
-            lv_row_mapping_df = pd.DataFrame(lv_row_mapping)
-            lv_row_mapping_df.to_parquet(lv_rowdata_file_name,index=False)
-
-            return "Record Added Successfully"
+            # Processing the file
+            lv_status = cv.fn_create_vector_store(lv_file_name, lv_domain)
+
+            return Response(
+                                json.dumps({"status":lv_status}), 
+                                status=200, 
+                                mimetype='application/json'
+                           )
         except Exception as e:
-            raise e
-
-
-def fn_map_data(lv_saved_file_name,lv_file_name,lv_source_domain):
-
-    # File Names
-    lv_faiss_file_name = 'db/'+lv_source_domain+'_index.faiss'
-    lv_sourcedata_file_name = 'db/'+lv_source_domain+'_row_mapping.parquet'
-    lv_mapping_file_name = 'db/'+lv_source_domain+"_"+lv_file_name
-
-    # Loading Data 
-    if os.path.exists(lv_faiss_file_name) and os.path.exists(lv_sourcedata_file_name):
-        # Load the pre-trained model
-        lv_model = SentenceTransformer('all-MiniLM-L6-v2')
-
-        # Load the Faiss index
-        lv_index = faiss.read_index(lv_faiss_file_name)
-
-        # Load the Row Data
-        lv_source_mapping_df = pd.read_parquet(lv_sourcedata_file_name)
-        lv_source_mapping_df.reindex()
-
-        # Excel to Dataframe
-        lv_excel_data = pd.read_excel(lv_saved_file_name,sheet_name=None)
-
-        # New Mapping Dataframe
-        lv_row_mapping_df = pd.DataFrame(columns=['source_sheet_name','source_column','target_sheet_name','target_column'])
-
-        # Reading each sheet
-        for lv_sheet_name, lv_sheet_data in lv_excel_data.items():        
-            
-            # Processing each row of the sheet
-            for i, row in enumerate(lv_sheet_data.itertuples(index=False)):
-                try:
-                    # Creating Embeddings
-                    # Details available here -> https://www.sbert.net/docs/pretrained_models.html
-                    lv_query = row[1]
-                    lv_query_embedding = lv_model.encode([lv_query])
-
-                    # Search for similar vectors
-                    lv_distances, lv_indices = lv_index.search(np.array(lv_query_embedding), 1)
-                    # print("Rahul Rahul")
-                    # print(lv_indices[0][0])
-
-                    # Mapped Row
-                    lv_row = lv_source_mapping_df.iloc[[lv_indices[0][0]]]
-                    # print(lv_row['sheet_name'])
-                    # print(lv_row['column_name'])
-
-                    lv_new_row = {
-                                    'source_sheet_name': lv_row['sheet_name'].values[0],
-                                    'source_column': lv_row['column_name'].values[0],
-                                    'target_sheet_name': lv_sheet_name,
-                                    'target_column': row[0]
-                                }
-
-                    # Adding to the Dataframe
-                    lv_row_mapping_df = pd.concat([lv_row_mapping_df, pd.DataFrame([lv_new_row])], ignore_index=True)
-                except Exception as e:
-                    pass
-
-        # Saving the Row Data to a file
-        lv_row_mapping_df.to_excel(lv_mapping_file_name,index=False)
+            self.lv_logger.error(e)
+            self.lv_logger.error(type(e))
+            self.lv_logger.error(traceback.format_exc())
 
-        return lv_row_mapping_df.to_json(orient='records')
-    else:
-        raise Exception("Source Domain Data Not Found")
+            return Response(json.dumps({"error_message":str(e)}), status=500, mimetype='application/json')