celsius_csrd_chatbot/chains/esrs_categorization.py

import re


def make_esrs_categorization_node():

    def categorize_message(state):
        query = state["query"]
        pattern = r"ESRS \d+[A-Z0-9]*|ESRS [A-Z]+\d+[A-Z0-9]*"
        esrs_truth = [
            "ESRS 1",
            "ESRS 2",
            "ESRS E1",
            "ESRS E2",
            "ESRS E3",
            "ESRS E4",
            "ESRS E5",
            "ESRS S1",
            "ESRS S2",
            "ESRS S3",
            "ESRS S4",
            "ESRS G1",
        ]

        matches = re.findall(pattern, query.upper())
        if matches:
            true_matches = [match for match in matches if match in esrs_truth]
            output = {"esrs_type": true_matches if true_matches else "wrong_esrs"}
        else:
            output = {"esrs_type": "none"}

        return output

    return categorize_message