Update utils.py

utils.py

@@ -60,17 +60,23 @@ xls = pd.ExcelFile('SmartClever table explanations.xlsx')
 metadata_df = pd.DataFrame()
 i = 0
 sheet_to_df_map = {}
-for sheet_name in xls.sheet_names:
-    sheet_to_df_map[sheet_name.strip()] = xls.parse(sheet_name, header=None)
-    sheet_to_df_map[sheet_name.strip()].columns = sheet_to_df_map[sheet_name.strip()].iloc[1]
-    sheet_to_df_map[sheet_name.strip()] = sheet_to_df_map[sheet_name.strip()].iloc[:1].fillna('')
-    sheet_to_df_map[sheet_name.strip()]['metadata'] = sheet_to_df_map[sheet_name.strip()].apply(lambda x: \
-                                                        ". ".join([x[col] for col in sheet_to_df_map[sheet_name.strip()].columns]), axis=1)
-    
-    metadata_df.loc[i, "table"] = sheet_name.strip()
-    metadata_df.loc[i, "desc"] = sheet_to_df_map[sheet_name.strip()]['metadata'].iloc[0]
-    
-    i += 1
+for k, sheet_name in enumerate(xls.sheet_names):
+    if k > 0:
+        sheet_to_df_map[sheet_name.strip()] = xls.parse(sheet_name, header=None)
+        sheet_to_df_map[sheet_name.strip()].columns = sheet_to_df_map[sheet_name.strip()].iloc[1]
+        sheet_to_df_map[sheet_name.strip()] = sheet_to_df_map[sheet_name.strip()].iloc[:1].fillna('')
+        sheet_to_df_map[sheet_name.strip()]['metadata'] = sheet_to_df_map[sheet_name.strip()].apply(lambda x: \
+                                                            ". ".join([x[col] for col in sheet_to_df_map[sheet_name.strip()].columns]), axis=1)
+
+        metadata_df.loc[i, "table"] = sheet_name.strip()
+        metadata_df.loc[i, "desc"] = sheet_to_df_map[sheet_name.strip()]['metadata'].iloc[0]
+
+        i += 1
+
+metadata_df2 = xls.parse('Table explanations',header=1).dropna(axis=0,how='all').dropna(axis=1,how='all')
+metadata_df2.columns = ['table','metadata']
+metadata_df2.table = metadata_df2.table.apply(lambda x: x.strip())
+metadata_df = pd.merge(metadata_df, metadata_df2, how='inner')
 
 table_search = EmbeddingsSearch(metadata_df=metadata_df, emb_model=emb_model)
 
@@ -93,7 +99,24 @@ def extract_question_type(llm, query):
         return 'specific'
     else:
         return 'unknown'
-        
+
+def extract_table_name(query):
+    messages = [
+    (
+        "system",
+        """
+        You are an AI assistant that determines the most relevant table name given a user query. Following is the metadata information you need to use to determine the most relevant table.\
+        {}.""".format(metadata_df[['table','metadata']].to_string()),
+    ),
+    ("human", query),
+    ]
+    output = llm.invoke(messages)
+    pred = output.content
+    
+    for table in metadata_df.table.unique():
+        if table in pred:
+            return table
+            
 warnings.filterwarnings('ignore', message="pandas only supports SQLAlchemy connectable.*", category=UserWarning, module='chain')
 
 intermediate_steps_KEY = "intermediate_steps"
@@ -248,7 +271,7 @@ def clean_sql(s: str) -> str:
         s = s.replace("TOP 1","").strip()
         s = s.replace("SELECT","SELECT TOP 1")
     return s
-    
+
 class SQLDatabaseChainPatched(SQLDatabaseChain):
     intermediate_steps: List[Any] = Field(default_factory=list)
     llms: Dict[str, Any] = Field(default_factory=dict)
@@ -270,6 +293,7 @@ class SQLDatabaseChainPatched(SQLDatabaseChain):
         # get number of tokens in the input prompt
         selected_inputs = {k: inputs[k] for k in chain.prompt.input_variables}
         prompt = chain.prompt.format_prompt(**selected_inputs)
+        #print (prompt)
         # https://stackoverflow.com/questions/75804599/openai-api-how-do-i-count-tokens-before-i-send-an-api-request
         n_tokens = num_tokens_from_string(string=prompt.text, encoding_name='cl100k_base')
         print(f"N tokens in input: {n_tokens}")
@@ -297,6 +321,7 @@ class SQLDatabaseChainPatched(SQLDatabaseChain):
         # If not present, then defaults to None which is all tables.
         table_names_to_use = inputs.get("table_names_to_use")
         table_info = self.database.get_table_info(table_names=table_names_to_use)
+        table_info += get_metadata_info(metadata_df, table_names_to_use)
         llm_inputs = {
             "input": input_text,
             "history": inputs["history"],
@@ -319,6 +344,7 @@ class SQLDatabaseChainPatched(SQLDatabaseChain):
             
         # list to store estimated num of tokens
         self.intermediate_steps['n_tokens_list'] = []
+        input_text_bkp = input_text
         try:
             # get sql
             self.llm_chain, n_tokens1 = self.prepare_llm(llm_inputs, chain=self.llm_chain)
@@ -360,10 +386,34 @@ class SQLDatabaseChainPatched(SQLDatabaseChain):
             # self.llm_chain = self.revert_to_small_model(chain=self.llm_chain)
             self.intermediate_steps['query_explanation'] = explanation
             
-        except Exception as exc:
-            # Append intermediate steps to exception, to aid in logging and later
-            # improvement of few shot prompt seeds
-            exc.intermediate_steps = self.intermediate_steps  # type: ignore
-            raise exc
-
+        except:
+            try:
+                sql_data_new = sql_data[-20:] + sql_data[:20]
+                input_text = input_text_bkp + f"{sql_cmd}\nSQLResult: {str(sql_data_new)}\nAnswer:"
+                llm_inputs["input"] = input_text
+                self.llm_chain, n_tokens3 = self.prepare_llm(llm_inputs, chain=self.llm_chain)
+                # self.intermediate_steps['n_tokens_list'].append(n_tokens3)
+                final_result = self.llm_chain.predict(
+                    callbacks=_run_manager.get_child(),
+                    **llm_inputs,
+                ).strip()
+                # self.llm_chain = self.revert_to_small_model(chain=self.llm_chain)
+                self.intermediate_steps['result'] = final_result
+
+                # provide explanation
+                input_text += f"{final_result}\nExplanation:"
+                llm_inputs["input"] = input_text
+                self.llm_chain, n_tokens4 = self.prepare_llm(llm_inputs, chain=self.llm_chain)
+                # self.intermediate_steps['n_tokens_list'].append(n_tokens3)
+                explanation = self.llm_chain.predict(
+                    callbacks=_run_manager.get_child(),
+                    **llm_inputs,
+                ).strip()
+                # self.llm_chain = self.revert_to_small_model(chain=self.llm_chain)
+                self.intermediate_steps['query_explanation'] = explanation
+            except Exception as exc:
+                # Append intermediate steps to exception, to aid in logging and later
+                # improvement of few shot prompt seeds
+                exc.intermediate_steps = self.intermediate_steps  # type: ignore
+                raise exc