Abstracted callbacks and code cleanup.
Modified to match the new public api of docgen