From e06bf61a940430f8dc74ec3946d3df02c5b0c36d Mon Sep 17 00:00:00 2001
From: Howard Gil <howardbgil@gmail.com>
Date: Mon, 16 Sep 2024 17:14:32 -0700
Subject: [PATCH] Fixed Time Travel output for multithreading and clarity
 (#388)

* Fixed for multithreading and clarity

* dont need threading

* leftover line

* fixed for threading

* cleaning up time travel output

* Added prepend string
---
 agentops/time_travel.py | 53 +++++++++++++----------------------------
 1 file changed, 17 insertions(+), 36 deletions(-)

diff --git a/agentops/time_travel.py b/agentops/time_travel.py
index 14e8b2af..6c1c7588 100644
--- a/agentops/time_travel.py
+++ b/agentops/time_travel.py
@@ -5,6 +5,8 @@
 from .exceptions import ApiServerException
 from .singleton import singleton
 
+ttd_prepend_string = "🖇️ Agentops: ⏰ Time Travel |"
+
 
 @singleton
 class TimeTravel:
@@ -47,9 +49,9 @@ def fetch_time_travel_id(ttd_id):
 
         set_time_travel_active_state(True)
     except ApiServerException as e:
-        manage_time_travel_state(activated=False, error=e)
+        print(f"{ttd_prepend_string} Error - {e}")
     except Exception as e:
-        manage_time_travel_state(activated=False, error=e)
+        print(f"{ttd_prepend_string} Error - {e}")
 
 
 def fetch_completion_override_from_time_travel_cache(kwargs):
@@ -64,14 +66,14 @@ def fetch_completion_override_from_time_travel_cache(kwargs):
 def find_cache_hit(prompt_messages, completion_overrides):
     if not isinstance(prompt_messages, (list, tuple)):
         print(
-            "Time Travel Error - unexpected type for prompt_messages. Expected 'list' or 'tuple'. Got ",
+            f"{ttd_prepend_string} Error - unexpected type for prompt_messages. Expected 'list' or 'tuple'. Got ",
             type(prompt_messages),
         )
         return None
 
     if not isinstance(completion_overrides, dict):
         print(
-            "Time Travel Error - unexpected type for completion_overrides. Expected 'dict'. Got ",
+            f"{ttd_prepend_string} Error - unexpected type for completion_overrides. Expected 'dict'. Got ",
             type(completion_overrides),
         )
         return None
@@ -80,7 +82,7 @@ def find_cache_hit(prompt_messages, completion_overrides):
             completion_override_dict = eval(key)
             if not isinstance(completion_override_dict, dict):
                 print(
-                    "Time Travel Error - unexpected type for completion_override_dict. Expected 'dict'. Got ",
+                    f"{ttd_prepend_string} Error - unexpected type for completion_override_dict. Expected 'dict'. Got ",
                     type(completion_override_dict),
                 )
                 continue
@@ -88,7 +90,7 @@ def find_cache_hit(prompt_messages, completion_overrides):
             cached_messages = completion_override_dict.get("messages")
             if not isinstance(cached_messages, list):
                 print(
-                    "Time Travel Error - unexpected type for cached_messages. Expected 'list'. Got ",
+                    f"{ttd_prepend_string} Error - unexpected type for cached_messages. Expected 'list'. Got ",
                     type(cached_messages),
                 )
                 continue
@@ -105,10 +107,12 @@ def find_cache_hit(prompt_messages, completion_overrides):
                 return value
         except (SyntaxError, ValueError, TypeError) as e:
             print(
-                f"Time Travel Error - Error processing completion_overrides item: {e}"
+                f"{ttd_prepend_string} Error - Error processing completion_overrides item: {e}"
             )
         except Exception as e:
-            print(f"Time Travel Error - Unexpected error in find_cache_hit: {e}")
+            print(
+                f"{ttd_prepend_string} Error - Unexpected error in find_cache_hit: {e}"
+            )
     return None
 
 
@@ -120,14 +124,10 @@ def check_time_travel_active():
     try:
         with open(config_file_path, "r") as config_file:
             config = yaml.safe_load(config_file)
-            if config.get("Time_Travel_Debugging_Active", True):
-                manage_time_travel_state(activated=True)
-                return True
+            return config.get("Time_Travel_Debugging_Active", False)
     except FileNotFoundError:
         return False
 
-    return False
-
 
 def set_time_travel_active_state(is_active: bool):
     config_path = ".agentops_time_travel.yaml"
@@ -144,30 +144,11 @@ def set_time_travel_active_state(is_active: bool):
             yaml.dump(config, config_file)
         except:
             print(
-                f"🖇 AgentOps: Unable to write to {config_path}. Time Travel not activated"
+                f"{ttd_prepend_string} Error - Unable to write to {config_path}. Time Travel not activated"
             )
             return
 
-        if is_active:
-            manage_time_travel_state(activated=True)
-            print("🖇 AgentOps: Time Travel Activated")
-        else:
-            manage_time_travel_state(activated=False)
-            print("🖇 AgentOps: Time Travel Deactivated")
-
-
-def add_time_travel_terminal_indicator():
-    print(f"🖇️ ⏰ | ", end="")
-
-
-def reset_terminal():
-    print("\033[0m", end="")
-
-
-def manage_time_travel_state(activated=False, error=None):
-    if activated:
-        add_time_travel_terminal_indicator()
+    if is_active:
+        print(f"{ttd_prepend_string} Activated")
     else:
-        reset_terminal()
-        if error is not None:
-            print(f"🖇 Deactivating Time Travel. Error with configuration: {error}")
+        print(f"{ttd_prepend_string} Deactivated")