diff --git a/agent/src/stats_collection.h b/agent/src/stats_collection.h
index 743d0be..ab9f174 100644
--- a/agent/src/stats_collection.h
+++ b/agent/src/stats_collection.h
@@ -22,13 +22,13 @@
 
 extern "C"
 {
-#include <bal_model_types.h>
+#include <bcmolt_api_model_supporting_structs.h>
 }
 
 void init_stats();
 void stop_collecting_statistics();
 openolt::PortStatistics* get_default_port_statistics();
-openolt::PortStatistics* collectPortStatistics(bcmbal_interface_key key);
+openolt::PortStatistics* collectPortStatistics(bcmolt_interface_id intf_id, bcmolt_interface_type intf_type);
 #if 0
 openolt::FlowStatistics* get_default_flow_statistics();
 openolt::FlowStatistics* collectFlowStatistics(bcmbal_flow_id flow_id, bcmbal_flow_type flow_type);