RahulSinghPundir/SQL_Wizard

1import pandas as pd
2import json
3from sentence_transformers import SentenceTransformer
4
5# Functianal Tools implement function which helps in getting table details 
6class FunctionalTools():
7    @staticmethod
8    def getDatabaseDetails():
9        # Load table descriptions from a CSV file
10        description_df = pd.read_csv("database_table_descriptions.csv")
11        description_text = ""
12        for _, row in description_df.iterrows():
13            description_text += f"Table Name: {row['Table']}\nTable Description: {row['Description']}\n\n"
14        return description_text
15
16    @staticmethod
17    def getTableDetails(recommended_tables):
18        description_text = ""
19        description_df = pd.read_csv("database_table_descriptions.csv")
20        description_df[description_df["Table"]=="aes"]["Table"]
21        for i in recommended_tables:
22            description_text+=f"Table Name: {description_df[description_df['Table']==i]['Table']}\nTable Description: {description_df[description_df['Table']==i]['Description']}\n\n"
23        return description_text
24    
25    @staticmethod
26    def getDatabaseDetailsJson():
27        # Load the large JSON file
28        with open('static\database_details.json', 'r') as file:
29            data = json.load(file)
30
31        # Initialize the model
32        model = SentenceTransformer('all-MiniLM-L6-v2')
33
34        def vectorize_table(table):
35            table_desc = f"{table['name']}: {table['description']}. Columns: " + ", ".join(
36                [f"{col['name']} ({col['type']})" for col in table['columns']]
37            )
38            # print(table_desc)
39            return model.encode(table_desc)
40
41        # Vectorize the table details
42        table_vectors = {table['name']: vectorize_table(table) for table in data['Tables']}
43
44        def create_prompt(tables):
45            prompt = "Here are the table relationships:\n\n"
46            for table, vector in tables.items():
47                prompt += f"Table: {table}\nDetails: {vector.tolist()}\n\n"
48            return prompt
49
50        # Assuming the API can handle limited text length, divide the prompts
51        table_names = list(table_vectors.keys())
52        chunk_size = 5  # Adjust based on API limits
53        prompts = [create_prompt({name: table_vectors[name] for name in table_names[i:i + chunk_size]}) 
54                for i in range(0, len(table_names), chunk_size)]
55
56        return prompts
57    
58    @staticmethod
59    def getTableDetailsJson(recommended_tables):
60        # Load the large JSON file
61        with open('static\database_details.json', 'r') as file:
62            data = json.load(file)
63
64        # Initialize the model
65        model = SentenceTransformer('all-MiniLM-L6-v2')
66
67        def vectorize_table(table):
68            table_desc = f"{table['name']}: {table['description']}. Columns: " + ", ".join(
69                [f"{col['name']} ({col['type']})" for col in table['columns']]
70            )
71            # print(table_desc)
72            return model.encode(table_desc)
73
74        # Vectorize the table details
75        table_vectors = {table['name']: vectorize_table(table) for table in data['Tables'] if (table in recommended_tables)}
76
77        def create_prompt(tables):
78            prompt = "Here are the table relationships:\n\n"
79            for table, vector in tables.items():
80                prompt += f"Table: {table}\nDetails: {vector.tolist()}\n\n"
81            return prompt
82
83        # Assuming the API can handle limited text length, divide the prompts
84        table_names = list(table_vectors.keys())
85        chunk_size = 5  # Adjust based on API limits
86        prompts = [create_prompt({name: table_vectors[name] for name in table_names[i:i + chunk_size]}) 
87                for i in range(0, len(table_names), chunk_size)]
88
89        return prompts
90    
91    @staticmethod
92    def getDatabaseJoins():
93        with open('static\database_joins.json', 'r') as file:
94            tables_data = json.load(file)
95
96        # Initialize the sentence-transformers model
97        model = SentenceTransformer('all-MiniLM-L6-v2')
98
99        # Vectorize the table details
100        table_vectors = {}
101        for table in tables_data["Joins"]:
102            table_name = table["name"]
103            table_desc = f"{table_name}: {table['joinType']} with {table['models']} based on {table['condition']}"
104            # print(table_desc)
105            table_vectors[table_name] = model.encode(table_desc)
106
107        # Create a prompt for the Gemini API
108        prompt = "Here are the table relationships:\n\n"
109        for table, vector in table_vectors.items():
110            prompt += f"Table: {table}\nDetails: {vector.tolist()}\n\n"
111
112        return prompt
113    
114    @staticmethod
115    def getTablesJoins(recommended_tables):
116        with open('static\database_joins.json', 'r') as file:
117            tables_data = json.load(file)
118
119        # Initialize the sentence-transformers model
120        model = SentenceTransformer('all-MiniLM-L6-v2')
121
122        # Vectorize the table details
123        table_vectors = {}
124        for table in tables_data["Joins"]:
125            table_name = table["name"]
126            if(table_name not in recommended_tables):
127                continue
128            table_desc = f"{table_name}: {table['joinType']} with {table['models']} based on {table['condition']}"
129            # print(table_desc)
130            table_vectors[table_name] = model.encode(table_desc)
131
132        # Create a prompt for the Gemini API
133        prompt = "Here are the table relationships:\n\n"
134        for table, vector in table_vectors.items():
135            prompt += f"Table: {table}\nDetails: {vector.tolist()}\n\n"
136
137        return prompt