fix: make agent callback logging opt-in for CLIs

Author: DavidJBianco

peak_assistant/data_assistant/__main__.py

@@ -87,6 +87,11 @@ def main() -> None:
         action="store_true",
         help="Skip user feedback and automatically accept the generated data discovery report"
     )
+    parser.add_argument(
+        "--debug-agents",
+        action="store_true",
+        help="Enable agent debug logging to msgs.txt and results.txt"
+    )
     args = parser.parse_args()
 
     # Load environment variables
@@ -156,6 +161,16 @@ def main() -> None:
             exit(1)
 
     messages: List[TextMessage] = list()
+    debug_agents_opts = dict()
+
+    if args.debug_agents:
+        debug_agents_opts = {
+            "msg_preprocess_callback": preprocess_messages_logging,
+            "msg_preprocess_kwargs": {"agent_id": "data-discovery"},
+            "msg_postprocess_callback": postprocess_messages_logging,
+            "msg_postprocess_kwargs": {"agent_id": "data-discovery"},
+        }
+
     while True:
         # Run the hypothesizer asynchronously
         data_sources = asyncio.run(
@@ -167,10 +182,7 @@ def main() -> None:
                 local_context=local_context,
                 verbose=args.verbose,
                 previous_run=messages,
-                msg_preprocess_callback=preprocess_messages_logging,
-                msg_preprocess_kwargs={"agent_id": "data-discovery"},
-                msg_postprocess_callback=postprocess_messages_logging,
-                msg_postprocess_kwargs={"agent_id": "data-discovery"},
+                **debug_agents_opts,
             )
         )
 

peak_assistant/hypothesis_assistant/hypothesis_refiner_cli.py

@@ -416,6 +416,11 @@ def main() -> None:
         help="Skip user feedback and automatically accept the refined hypothesis",
         default=False,
     )
+    parser.add_argument(
+        "--debug-agents",
+        action="store_true",
+        help="Enable agent debug logging to msgs.txt and results.txt",
+    )
 
     # Parse the arguments
     args = parser.parse_args()
@@ -480,6 +485,16 @@ def main() -> None:
 
     messages: List[TextMessage] = list()
     current_hypothesis = args.hypothesis
+    debug_agents_opts = dict()
+
+    if args.debug_agents:
+        debug_agents_opts = {
+            "msg_preprocess_callback": preprocess_messages_logging,
+            "msg_preprocess_kwargs": {"agent_id": "hypothesis-refiner"},
+            "msg_postprocess_callback": postprocess_messages_logging,
+            "msg_postprocess_kwargs": {"agent_id": "hypothesis-refiner"},
+        }
+
     while True:
         # Run the hypothesizer asynchronously
         response = asyncio.run(
@@ -490,10 +505,7 @@ def main() -> None:
                 local_data_document=local_data or "",
                 verbose=args.verbose,
                 previous_run=messages,
-                msg_preprocess_callback=preprocess_messages_logging,
-                msg_preprocess_kwargs={"agent_id": "hypothesis-refiner"},
-                msg_postprocess_callback=postprocess_messages_logging,
-                msg_postprocess_kwargs={"agent_id": "hypothesis-refiner"},
+                **debug_agents_opts,
             )
         )
 

peak_assistant/planning_assistant/__main__.py

@@ -95,6 +95,11 @@ def main() -> None:
         action="store_true",
         help="Skip user feedback and automatically accept the generated hunt plan"
     )
+    parser.add_argument(
+        "--debug-agents",
+        action="store_true",
+        help="Enable agent debug logging to msgs.txt and results.txt"
+    )
     args = parser.parse_args()
 
     # Load environment variables
@@ -177,6 +182,16 @@ def main() -> None:
             exit(1)
 
     messages: List[TextMessage] = list()
+    debug_agents_opts = dict()
+
+    if args.debug_agents:
+        debug_agents_opts = {
+            "msg_preprocess_callback": preprocess_messages_logging,
+            "msg_preprocess_kwargs": {"agent_id": "hunt-planner"},
+            "msg_postprocess_callback": postprocess_messages_logging,
+            "msg_postprocess_kwargs": {"agent_id": "hunt-planner"},
+        }
+
     while True:
         # Run the hypothesizer asynchronously
         data_sources = asyncio.run(
@@ -189,10 +204,7 @@ def main() -> None:
                 local_context=local_context or "",
                 verbose=args.verbose,
                 previous_run=messages,
-                msg_preprocess_callback=preprocess_messages_logging,
-                msg_preprocess_kwargs={"agent_id": "hunt-planner"},
-                msg_postprocess_callback=postprocess_messages_logging,
-                msg_postprocess_kwargs={"agent_id": "hunt-planner"},
+                **debug_agents_opts,
             )
         )