class LanguageModelSqlFilterBuilder {

    private static final Logger log = LoggerFactory.getLogger(LanguageModelSqlFilterBuilder.class);

    private static final PromptTemplate DEFAULT_PROMPT_TEMPLATE = PromptTemplate.from(
            "### Instructions:\n" +
                    "Your task is to convert a question into a SQL query, given a Postgres database schema.\n" +
                    "Adhere to these rules:\n" +
                    "- **Deliberately go through the question and database schema word by word** to appropriately answer the question\n" +
                    "- **Use Table Aliases** to prevent ambiguity. For example, `SELECT table1.col1, table2.col1 FROM table1 JOIN table2 ON table1.id = table2.id`.\n" +
                    "- When creating a ratio, always cast the numerator as float\n" +
                    "\n" +
                    "### Input:\n" +
                    "Generate a SQL query that answers the question `{{query}}`.\n" +
                    "This query will run on a database whose schema is represented in this string:\n" +
                    "{{create_table_statement}}\n" +
                    "\n" +
                    "### Response:\n" +
                    "Based on your instructions, here is the SQL query I have generated to answer the question `{{query}}`:\n" +
                    "```sql"
    );

    protected final ChatLanguageModel chatLanguageModel;
    protected final TableDefinition tableDefinition;
    protected final String createTableStatement;
    protected final PromptTemplate promptTemplate;
    protected final SqlFilterParser sqlFilterParser;

    public LanguageModelSqlFilterBuilder(ChatLanguageModel chatLanguageModel,
                                         TableDefinition tableDefinition) {
        this(chatLanguageModel, tableDefinition, DEFAULT_PROMPT_TEMPLATE, new SqlFilterParser());
    }

    @Builder
    private LanguageModelSqlFilterBuilder(ChatLanguageModel chatLanguageModel,
                                          TableDefinition tableDefinition,
                                          PromptTemplate promptTemplate,
                                          SqlFilterParser sqlFilterParser) {
        this.chatLanguageModel = ensureNotNull(chatLanguageModel, "chatLanguageModel");
        this.tableDefinition = ensureNotNull(tableDefinition, "tableDefinition");
        this.createTableStatement = format(tableDefinition);
        this.promptTemplate = getOrDefault(promptTemplate, DEFAULT_PROMPT_TEMPLATE);
        this.sqlFilterParser = getOrDefault(sqlFilterParser, SqlFilterParser::new);
    }

    public Filter build(Query query) {

        Prompt prompt = createPrompt(query);

        Response<AiMessage> response = chatLanguageModel.generate(prompt.toUserMessage());

        String generatedSql = response.content().text();

        String cleanedSql = clean(generatedSql);
        log.trace("Cleaned SQL: '{}'", cleanedSql);

        try {
            return sqlFilterParser.parse(cleanedSql);
        } catch (Exception e) {
            log.warn("Failed parsing the following SQL: '{}'", cleanedSql, e);
            // TODO implement additional strategies (configurable):
            //  - feed the error to the LLM and retry
            //  - return predefined filter
            //  - return partial filter if the filter is composite and some parts were parsed successfully
            //  - etc
            return fallback(query, generatedSql, cleanedSql, e);
        }
    }

    protected Prompt createPrompt(Query query) {
        Map<String, Object> variables = new HashMap<>();
        variables.put("create_table_statement", createTableStatement);
        variables.put("query", query.text());
        return promptTemplate.apply(variables);
    }

    protected String clean(String sql) {
        return sql.trim();
    }

    protected Filter fallback(Query query, String generatedSql, String cleanedSql, Exception e) {

        String extractedSql = extractSelectStatement(generatedSql);
        if (isNullOrBlank(extractedSql)) {
            log.trace("Cannot extract SQL, giving up");
            return null;
        }

        try {
            log.trace("Extracted SQL: '{}'", extractedSql);
            return sqlFilterParser.parse(extractedSql);
        } catch (Exception e2) {
            log.warn("Failed parsing the following SQL, giving up: '{}'", extractedSql, e2);
            return null;
        }
    }

    protected String extractSelectStatement(String dirtySql) {<FILL_FUNCTION_BODY>}

    protected String format(TableDefinition tableDefinition) {
        StringBuilder createTableStatement = new StringBuilder();
        createTableStatement.append(String.format("CREATE TABLE %s (\n", tableDefinition.name()));
        for (ColumnDefinition columnDefinition : tableDefinition.columns()) {
            createTableStatement.append(String.format("    %s %s,", columnDefinition.name(), columnDefinition.type()));
            if (!isNullOrBlank(columnDefinition.description())) {
                createTableStatement.append(String.format(" -- %s", columnDefinition.description()));
            }
            createTableStatement.append("\n");
        }
        createTableStatement.append(")");
        if (!isNullOrBlank(tableDefinition.description())) {
            createTableStatement.append(String.format(" COMMENT='%s'", tableDefinition.description()));
        }
        createTableStatement.append(";");
        return createTableStatement.toString();
    }
}

        // TODO improve
        if (dirtySql.contains("```sql")) {
            for (String part : dirtySql.split("```sql")) {
                if (part.toUpperCase().contains("SELECT") && part.toUpperCase().contains("WHERE")) {
                    return part.split("```")[0].trim();
                }
            }
        } else if (dirtySql.contains("```")) {
            for (String part : dirtySql.split("```")) {
                if (part.toUpperCase().contains("SELECT") && part.toUpperCase().contains("WHERE")) {
                    return part.split("```")[0].trim();
                }
            }
        } else {
            for (String part : dirtySql.split("SELECT")) {
                if (part.toUpperCase().contains("WHERE")) {
                    if (part.contains("\n")) {
                        for (String part2 : part.split("\n")) {
                            if (part2.toUpperCase().contains("WHERE")) {
                                return "SELECT " + part2.trim();
                            }
                        }
                    } else {
                        return "SELECT " + part.trim();
                    }
                }
            }
        }
        return null;
    
<no_super_class>