def CURATED_FEW_SHOT_COT_PROMPT()

in nl2sql/tasks/table_selection/core.py [0:0]

• 17 lines of code
• 3 McCabe index (conditional complexity)

    def CURATED_FEW_SHOT_COT_PROMPT(self) -> _CoreTableSelectorPrompt:

        def greedy_post_processor(
            intermediate_steps: List[Dict[str, Any]], available_tables: set[str]
        ) -> set[str]:
            """Implements greedy post-processing for yes/no table selection."""
            selected_tables = []

            for step in intermediate_steps:
                if "yes." in step["raw_response"].lower():
                    selected_tables.append(step["table"])
            filtered_tables = set(selected_tables).intersection(available_tables)

            return filtered_tables

        return _CoreTableSelectorPrompt(
            prompt_id="TASK_TABLE_SELECTION_CORE_V1_SPIDER_V1",
            prompt_template=FewShotPrompts.TASK_TABLE_SELECTION_CORE_V1_SPIDER_V1,
            call_for_each_table=True,
            post_processor=lambda x: yes_no_classifier(x) == "True",
            greedy_post_processor=greedy_post_processor,
        )