Pandas Apply Key Error
As answered by EdChum in the comments. The issue is that apply works column wise by default (see the docs). Therefore, the column names cannot be accessed.
To specify that it should be applied to each row instead, axis=1 must be passed:
test.apply(lambda x: find_max(x,test,'document_id','confidence_level','category_id'), axis=1)