Update utils.py

utils.py

@@ -32,8 +32,6 @@ from langchain.schema.language_model import BaseLanguageModel
 from langchain.utilities.sql_database import SQLDatabase
 from langchain_experimental.pydantic_v1 import Extra, Field, root_validator
 
-import streamlit as st
-
 emb_model = SentenceTransformer("all-MiniLM-L6-v2")
 
 class EmbeddingsSearch:
@@ -182,6 +180,50 @@ warnings.filterwarnings('ignore', message="pandas only supports SQLAlchemy conne
 
 intermediate_steps_KEY = "intermediate_steps"
 
+template = """
+
+            You are a database expert. Given an input question, first create a syntactically correct {dialect} query to run, then look at the results of the query and return the answer.
+            The final answer should be in a concise natural language. 
+            
+            Use the history if you can not understand the question.
+            
+            If the question is in another language, translate it to English before proceeding.
+            
+            Do not repeat the question while generating the SQL query.
+            
+            Only generate a correct {dialect} query.
+            
+            Once the SQLResult is available, generate the final answer in natural language format. Do not regenerate the question or SQL query in the final answer.
+            
+            If the question asks any information for any particular number of days, use the lookback from the maximum date in the table, not from today's date.
+            
+            Please note that MSSQL does not use LIMIT, but uses TOP clause.
+            
+            You may also need to resolve the column name, as per the metadata. For instance, if the user asks about families and the column name is family, you should use family in the generated SQL.
+            
+            Make sure that the column names are present in the table, by looking at the metadata.
+            
+            If a question asks about availability over a period of time, you need to use SUM to calculate the total availability over that time period.
+            
+            If a question mentions SKU, then use SKU column for filter, do not use any other column like comodity
+            
+            If a question asks about AV of shortage, do not use AV in the SQL query as AV is not a valid column name. AV is the key in the Shortage column.
+            
+            In the OpenOrderShotage table, the column Item should be used to extract the part ids, to answer questions related to shortage.
+            
+            In the OpenOrderShotage table, Customer_Part_Name column is equivalent to SKU.
+            
+            Use the following format:
+            Question: Question here
+            SQLQuery: SQL Query to run
+            SQLResult: Result of the SQLQuery
+            Answer: Final answer here.
+            
+            Only use the following tables:
+            {table_info}
+            Question: {input}
+            """
+
 def num_tokens_from_string(string: str, encoding_name: str) -> int:
     encoding = tiktoken.get_encoding(encoding_name)
     num_tokens = len(encoding.encode(string))
@@ -276,7 +318,7 @@ class SQLDatabaseChainPatched(SQLDatabaseChain):
         history = inputs['history'].copy()
         history.reverse()
         
-        inputs[self.input_key] = translate_to_english(self.llms['4k'], inputs[self.input_key])
+        #inputs[self.input_key] = translate_to_english(self.llms['4k'], inputs[self.input_key])
         input_text = f"{inputs[self.input_key]} \nHistory: {history} \nSQLQuery:"
         _run_manager.on_text(input_text, verbose=self.verbose)
         # If not present, then defaults to None which is all tables.
@@ -306,50 +348,84 @@ class SQLDatabaseChainPatched(SQLDatabaseChain):
         # list to store estimated num of tokens
         self.intermediate_steps['n_tokens_list'] = []
         input_text_bkp = input_text
-
-        # get sql
-        self.llm_chain, n_tokens1 = self.prepare_llm(llm_inputs, chain=self.llm_chain)
-        # self.intermediate_steps['n_tokens_list'].append(n_tokens1)
-        sql_cmd = self.llm_chain.predict(
-            callbacks=_run_manager.get_child(),
-            **llm_inputs,
-        ).strip()
-
-        st.markdown(sql_cmd)
+        try:
+            # get sql
+            self.llm_chain, n_tokens1 = self.prepare_llm(llm_inputs, chain=self.llm_chain)
+            # self.intermediate_steps['n_tokens_list'].append(n_tokens1)
+            sql_cmd = 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['sql_cmd_unchecked'] = sql_cmd
         
-        # self.llm_chain = self.revert_to_small_model(chain=self.llm_chain)
-        self.intermediate_steps['sql_cmd_unchecked'] = sql_cmd
-    
-        self.intermediate_steps['sql_cmd'] = clean_sql(sql_cmd)
+            self.intermediate_steps['sql_cmd'] = clean_sql(sql_cmd)
 
-        # run sql            
-        sql_data = self.database._execute(self.intermediate_steps['sql_cmd'], fetch='all')
-        self.intermediate_steps['sql_data'] = sql_data
-        
-        # provide human answer
-        input_text += f"{sql_cmd}\nSQLResult: {str(sql_data)}\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
-        
-        if 'result' in self.intermediate_steps:
-            self.intermediate_steps['translated_result'] = translate(self.llms['4k'], orig_question, self.intermediate_steps['result'])
-        
+            # run sql            
+            sql_data = self.database._execute(self.intermediate_steps['sql_cmd'], fetch='all')
+            self.intermediate_steps['sql_data'] = sql_data
+            
+            # provide human answer
+            input_text += f"{sql_cmd}\nSQLResult: {str(sql_data)}\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
+            
+            #if 'result' in self.intermediate_steps:
+            #    self.intermediate_steps['translated_result'] = translate(self.llms['4k'], orig_question, self.intermediate_steps['result'])
+            
+        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
+            #if 'result' in self.intermediate_steps:
+            #    self.intermediate_steps['translated_result'] = translate(self.llms['4k'], orig_question, self.intermediate_steps['result'])
+            
+            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
+                self.intermediate_steps['result'] = "I don't know the answer for this."
+                #self.intermediate_steps['translated_result'] = "I don't know the answer for this."
+
+