consective_months.py

import pandas as pd

df = pd.DataFrame({'agent_id': [1, 1, 1, 1, 2, 2, 2, 2],
                   'date': ["2021-01-01", "2021-04-01", "2021-05-01", "2021-06-01", "2021-01-01", "2021-02-01", "2021-03-01", "2021-06-01"], 
                   'txn_amount': [100, 200, 100, 200, 100, 200, 100, 200],
                   'txn_status': ["Failure", "Success", "Failure", "Success", "Failure", "Success", "Failure", "Success"]})
df['date'] = pd.to_datetime(df['date'], format='%Y-%m-%d')

df1 = df.set_index('date')

df1['month'] = df1.index.month
df1['year'] = df1.index.year

df1.groupby([df1['agent_id'] , df1.year]).apply(is_at_least_three_consec)

def is_at_least_three_consec(grouped_df):
    """
    Function to check for 3 consecutive months
    Input: Grouped df
    Output: Subset of the df which has continuous month values
    """
    month_diff = grouped_df['month'].diff().values.tolist()
    consec_count = 0
    for index , val in enumerate(month_diff):
        if index != 0 and val == 1:
                consec_count += 1
                if consec_count == 2:
                    return grouped_df[index-2: index+1]
        else:
            consec_count = 0
    return None