Knowledge base rebuild take start and end dates via arguments (#72)

* Catch ValueError once not every time * Rename st_time to start_time * _rebuild: add end_date argument to * Change logging level to debug
stratosphereips · Nov 4, 2024 · f229714 · f229714
1 parent 235b570
commit f229714
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 35 deletions.
diff --git a/bin/aip b/bin/aip
@@ -78,43 +78,41 @@ def main():
     log_fmt = '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
     logging.basicConfig(level=args.log_level, format=log_fmt)
 
-    if args.rebuild_kb:
-        try:
+    try: 
+        # Rebuild knowledge base
+        if args.rebuild_kb:
             rebuild_date = validate_and_convert_date(args.rebuild_kb)
-        except ValueError as err:
-            logger.error(err)
+            logger.info(f"Rebuilding knowledge base from {rebuild_date} onwards")
+            _rebuild(start_date=rebuild_date, log_level=args.log_level)
             sys.exit()
-        logger.info(f"Rebuilding knowledge base from {rebuild_date} onwards")
-        _rebuild(start_date=rebuild_date, log_level=args.log_level)
-        sys.exit()
 
-    # Validate input date
-    try:
-        run_date_day = validate_and_convert_date(args.date)
-    except ValueError as err:
-        logger.error(err)
-        sys.exit()
+        # Validate input date
+        if args.date: 
+            run_date_day = validate_and_convert_date(args.date)
 
-    logging.info(f"Running {args.model} model(s) for date {run_date_day}.")
-    # Run Alpha Model
-    if args.model in ['Alpha', 'all']:
-        run_model('Alpha', Alpha(), run_date_day, args.log_level)
+        logging.info(f"Running {args.model} model(s) for date {run_date_day}.")
+        # Run Alpha Model
+        if args.model in ['Alpha', 'all']:
+            run_model('Alpha', Alpha(), run_date_day, args.log_level)
 
-    # Alpha 7 Model
-    if args.model in ['Alpha7', 'all']:
-        run_model('Alpha7', Alpha7(), run_date_day, args.log_level)
+        # Alpha 7 Model
+        if args.model in ['Alpha7', 'all']:
+            run_model('Alpha7', Alpha7(), run_date_day, args.log_level)
 
-    # Prioritize New Model
-    if args.model in ['Prioritize_New', 'all']:
-        run_model('Prioritize_New', New(), run_date_day, args.log_level)
+        # Prioritize New Model
+        if args.model in ['Prioritize_New', 'all']:
+            run_model('Prioritize_New', New(), run_date_day, args.log_level)
 
-    # Prioritize Consistent Model
-    if args.model in ['Prioritize_Consistent', 'all']:
-        run_model('Prioritize_Consistent', Consistent(), run_date_day, args.log_level)
+        # Prioritize Consistent Model
+        if args.model in ['Prioritize_Consistent', 'all']:
+            run_model('Prioritize_Consistent', Consistent(), run_date_day, args.log_level)
 
-    # Prioritize Random Forest Model
-    if args.model in ['Random_Forest', 'all']:
-        run_model('Random_Forest', RandomForest(), run_date_day, args.log_level)
+        # Prioritize Random Forest Model
+        if args.model in ['Random_Forest', 'all']:
+            run_model('Random_Forest', RandomForest(), run_date_day, args.log_level)
+    except ValueError as err:
+        logger.error(err)
+        sys.exit()
 
 if __name__ == '__main__':
     main()
diff --git a/lib/aip/utils/knowledge_base.py b/lib/aip/utils/knowledge_base.py
@@ -75,21 +75,20 @@ def _build_knowledge(start, end, log_level=logging.ERROR):
     return last_knowledge
 
 
-def _rebuild(start_date, log_level=logging.ERROR):
+def _rebuild(start_date, end_date=date.today(), log_level=logging.ERROR):
     """
     Wrapper to rebuild the knowledge base from the specified start date to today.
     """
     # Configure logger
     logger = logging.getLogger(__name__)
     logging.basicConfig(level=log_level)
 
-    st_time = time.time()
-    current_date = date.today()
-    logger.info(f"_rebuild - Rebuilding knowledge base from {start_date} to {current_date}")
+    start_time = time.time()
+    logger.debug(f"_rebuild - Rebuilding knowledge base from {start_date} to {end_date}")
 
-    _build_knowledge(start=start_date, end=current_date, log_level=log_level)
+    _build_knowledge(start=start_date, end=end_date, log_level=log_level)
 
-    logger.info(f'_rebuild - Rebuilding knowledge base finished in {(time.time() - st_time)/60} minutes')
+    logger.debug(f'_rebuild - Rebuilding knowledge base finished in {(time.time() - start_time)/60} minutes')
 
 
 class Knowledgebase():