Update utils.py

utils.py

@@ -142,144 +142,97 @@ def extract_question_list(llm, query):
     except:
         return query
 
+def translate_to_english(llm, user_query):
+    sys_prompt = """
+        You are an AI assistant that translates a text to English. \
+        Do not generate any irrelavant text, only return the translation."""
+
+    message1 = SystemMessage(content=sys_prompt)
+
+    message2 = HumanMessage(
+        content=user_query
+    )
+    message_log = [message1, message2]
+    
+    output = llm.invoke(message_log)
+    pred = output.content
+    
+    return pred
+
+def translate(llm, user_query, to_translate):
+    sys_prompt = """
+        You are an AI assistant that determines the language given a user query - {} and translate the provided text in that target language. \
+        Do not generate any irrelavant text, only return the translation. \
+        If the user query is in English, then don't do anything just return the original text, no translation is required there.""".format(user_query)
+
+    message1 = SystemMessage(content=sys_prompt)
+
+    message2 = HumanMessage(
+        content=to_translate
+    )
+    message_log = [message1, message2]
+    
+    output = llm.invoke(message_log)
+    pred = output.content
+    
+    return pred
 
 warnings.filterwarnings('ignore', message="pandas only supports SQLAlchemy connectable.*", category=UserWarning, module='chain')
 
 intermediate_steps_KEY = "intermediate_steps"
 
-class PatchedCreateTable(CreateTable):
-    def __init__(
-        self,
-        element: Table,
-        include_foreign_key_constraints = None,
-        if_not_exists: bool = False,
-        columns_to_ignore: List[str] = None,
-    ):
-        if columns_to_ignore is None:
-            columns_to_ignore = []
-        element.columns = [col for col in element.columns if col.name not in columns_to_ignore]
-        super().__init__(element, if_not_exists=if_not_exists)
+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))
     return num_tokens
 
-
-class PatchedSQLDatabase(SQLDatabase):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-    def get_table_info(self, table_names: Optional[List[str]] = None) -> str:
-        """Get information about specified tables.
-
-        Follows best practices as specified in: Rajkumar et al, 2022
-        (https://arxiv.org/abs/2204.00498)
-
-        If `sample_rows_in_table_info`, the specified number of sample rows will be
-        appended to each table description. This can increase performance as
-        demonstrated in the paper.
-        """
-        all_table_names = self.get_usable_table_names()
-        if table_names is not None:
-            missing_tables = set(table_names).difference(all_table_names)
-            if missing_tables:
-                print('all_table_names', all_table_names)
-                raise ValueError(f"table_names {missing_tables} not found in database")
-            all_table_names = table_names
-
-        meta_tables = [
-            tbl
-            for tbl in self._metadata.sorted_tables
-            if tbl.name in set(all_table_names)
-            and not (self.dialect == "sqlite" and tbl.name.startswith("sqlite_"))
-        ]
-
-        tables = []
-        for table in meta_tables:
-            if self._custom_table_info and table.name in self._custom_table_info:
-                tables.append(self._custom_table_info[table.name])
-                continue
-
-            # Ignore JSON datatyped columns
-            _cols = table.columns
-            if _cols is dict:
-                for k, v in _cols.items():
-                    if type(v.type) is NullType:
-                        table._columns.remove(v)
-
-            # add create table command
-            create_table = str(PatchedCreateTable(
-                table, columns_to_ignore=[]
-            ).compile(self._engine))
-            table_info = ""
-            # table_info += f"{create_table.rstrip()}"
-
-            has_extra_info = (
-                self._indexes_in_table_info or self._sample_rows_in_table_info
-            )
-            if has_extra_info:
-                table_info += "\n\n/*"
-            if self._indexes_in_table_info:
-                table_info += f"\n{self._get_table_indexes(table)}\n"
-            if self._sample_rows_in_table_info:
-                table_info += f"\n{self._get_sample_rows(table)}\n"
-            if has_extra_info:
-                table_info += "*/"
-            table_info += self.get_columns_descriptions(table)
-
-            tables.append(table_info)
-        tables.sort()
-        final_str = "\n\n".join(tables)
-        return final_str
-        
-    def _get_sample_rows(self, table: Table) -> str:
-        # build the select command
-        command = select(table).order_by(func.random()).limit(self._sample_rows_in_table_info)
-
-        # save the columns in string format
-        columns_str = ";".join([f'"{col.name}"' for col in table.columns])
-
-        try:
-            # get the sample rows
-            with self._engine.connect() as connection:
-                sample_rows_result = connection.execute(command)  # type: ignore
-                # shorten values in the sample rows
-                sample_rows = list(
-                    map(lambda ls: [str(i)[:100] for i in ls], sample_rows_result)
-                )
-
-            # save the sample rows in string format
-            sample_rows_str = "\n".join([";".join(row) for row in sample_rows])
-
-        # in some dialects when there are no rows in the table a
-        # 'ProgrammingError' is returned
-        except ProgrammingError:
-            sample_rows_str = ""
-
-        return (
-            f"{self._sample_rows_in_table_info} rows from {table.name} table:\n"
-            f"{columns_str}\n"
-            f"{sample_rows_str}"
-        )
-
-    def _execute(
-        self,
-        command: str,
-        fetch: str = "all",
-    ):
-        with self._engine.begin() as connection:
-            cursor = connection.execute(text(command))
-            print(cursor.__dict__)
-
-            if cursor.returns_rows:
-                fields = list(cursor.keys())
-                result = [dict(zip(fields,row)) for row in cursor.fetchall()]
-                return result
-            return []
-
 def clean_sql(s: str) -> str:
+    #s = s.replace("SQL:","").strip()
+    #s = s.replace("Let's execute these queries step-by-step to get the final answer.","").strip()
     s = s.replace("```sql", "")
     for symb in ["'", '"']:
         if s.startswith(symb) and s.endswith(symb):
@@ -296,13 +249,19 @@ def clean_sql(s: str) -> str:
     if s.endswith("TOP 1"):
         s = s.replace("TOP 1","").strip()
         s = s.replace("SELECT","SELECT TOP 1")
+    
+    s = s.split("SQLQuery:")[-1].strip()
+    
     return s
 
 def get_metadata_info(metadata_df, table_names):
     str = ""
     for table in table_names:
-        str += "The following metadata is for the table " + table + "\n"
-        str += metadata_df[metadata_df.table == table].desc.iloc[0]
+        try:
+            str += "The following metadata is for the table " + table + "\n"
+            str += metadata_df[metadata_df.table == table].desc.iloc[0]
+        except:
+            pass
         
     return str
     
@@ -356,7 +315,12 @@ class SQLDatabaseChainPatched(SQLDatabaseChain):
         #print ("input key", self.input_key)
         #print ("===============")
         
-        input_text = f"{inputs[self.input_key]} \nHistory: {inputs['history']} \nSQLQuery:"
+        orig_question = inputs[self.input_key]
+        history = inputs['history'].copy()
+        history.reverse()
+        
+        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.
         table_names_to_use = inputs.get("table_names_to_use")
@@ -364,7 +328,7 @@ class SQLDatabaseChainPatched(SQLDatabaseChain):
         table_info += get_metadata_info(metadata_df, table_names_to_use)
         llm_inputs = {
             "input": input_text,
-            "history": inputs["history"],
+            "history": history,
             "top_k": str(self.top_k),
             "dialect": self.database.dialect,
             "table_info": table_info,
@@ -426,6 +390,9 @@ class SQLDatabaseChainPatched(SQLDatabaseChain):
             # 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]
@@ -451,11 +418,15 @@ class SQLDatabaseChainPatched(SQLDatabaseChain):
                 ).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
-    
+                #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."