def table_filter()

in nl2sql_src/nl2sql_generic.py [0:0]

• 23 lines of code
• 5 McCabe index (conditional complexity)

    def table_filter(self, question):
        """
        This function selects the relevant table(s) to the provided question
        based on their description-keywords.
        It assists in selecting a table from a list of tables based on
        their description-keywords.
        It presents a prompt containing a list of table names along
        with their corresponding description-keywords.
        The function uses a text-based model (text_bison) to analyze
        the prompt and extract the selected table name(s).

        Parameters:
        - question (str): The question for which the relevant table
          need to be identified.

        Returns:
        list: A list of table names most likely relevant to the provided
        question.
        """

        only_tables_info = ""

        for table in self.metadata_json:
            only_tables_info = only_tables_info + f"{table} | \
                {self.metadata_json[table]['Description']}\n"

        prompt = Table_filtering_prompt.format(
            only_tables_info=only_tables_info,
            question=question
            )

        result = self.llm.invoke(prompt)

        segments = result.split(',')
        tables_list = []

        for segment in segments:
            segment = segment.strip()
            if ':' in segment:
                value = segment.split(':')[-1].strip()
                tables_list.append(value.strip())
            elif '\n' in segment:
                value = segment.split('\n')[-1].strip()
                tables_list.append(value.strip())
            else:
                tables_list.append(segment)

        return tables_list