From f2297147dd74345ff6a391cea9a8cfd0b54eb34c Mon Sep 17 00:00:00 2001
From: Veronica Valeros <vero.valeros@gmail.com>
Date: Mon, 4 Nov 2024 19:07:19 +0100
Subject: [PATCH] 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
---
 bin/aip                         | 56 ++++++++++++++++-----------------
 lib/aip/utils/knowledge_base.py | 11 +++----
 2 files changed, 32 insertions(+), 35 deletions(-)

diff --git a/bin/aip b/bin/aip
index 0c746a7..f3879b9 100755
--- 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
index 4f32c1c..8e970ab 100644
--- a/lib/aip/utils/knowledge_base.py
+++ b/lib/aip/utils/knowledge_base.py
@@ -75,7 +75,7 @@ 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.
     """
@@ -83,13 +83,12 @@ def _rebuild(start_date, log_level=logging.ERROR):
     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():