From 70d09e85016b61f27032a63b39f47ad39c723c93 Mon Sep 17 00:00:00 2001
From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com>
Date: Mon, 3 Jun 2024 12:27:07 -0400
Subject: [PATCH] Restore human in the loop

---
 src/controlflow/llm/handlers.py        | 11 ++++++++---
 src/controlflow/tools/talk_to_human.py |  5 +++--
 2 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/src/controlflow/llm/handlers.py b/src/controlflow/llm/handlers.py
index 389c9fb9..a25f93fb 100644
--- a/src/controlflow/llm/handlers.py
+++ b/src/controlflow/llm/handlers.py
@@ -95,21 +95,26 @@ def on_tool_result_done(self, message: ToolMessage):
 
 
 class PrintHandler(CompletionHandler):
+    def __init__(self):
+        self.messages: dict[str, ControlFlowMessage] = {}
+        self.live = Live(auto_refresh=False)
+
     def on_start(self):
-        self.live = Live(refresh_per_second=12)
         self.live.start()
-        self.messages: dict[str, ControlFlowMessage] = {}
 
     def on_end(self):
         self.live.stop()
 
+    def on_exception(self, exc: Exception):
+        self.live.stop()
+
     def update_live(self):
         messages = sorted(self.messages.values(), key=lambda m: m.timestamp)
         content = []
         for message in messages:
             content.append(format_message(message))
 
-        self.live.update(Group(*content))
+        self.live.update(Group(*content), refresh=True)
 
     def on_message_delta(self, delta: AssistantMessage, snapshot: AssistantMessage):
         self.messages[snapshot.id] = snapshot
diff --git a/src/controlflow/tools/talk_to_human.py b/src/controlflow/tools/talk_to_human.py
index 028483b9..6620d492 100644
--- a/src/controlflow/tools/talk_to_human.py
+++ b/src/controlflow/tools/talk_to_human.py
@@ -4,6 +4,7 @@
 
 from prefect.context import FlowRunContext
 from prefect.input.run_input import receive_input
+from rich.prompt import Prompt
 
 import controlflow
 from controlflow.llm.tools import tool
@@ -16,8 +17,8 @@ async def get_terminal_input():
     # as a convenience, we wait for human input on the local terminal
     # this is not necessary for the flow to run, but can be useful for testing
     loop = asyncio.get_event_loop()
-    user_input = await loop.run_in_executor(None, input, "Type your response: ")
-    # user_input = await loop.run_in_executor(None, Prompt.ask, "Type your response")
+    # user_input = await loop.run_in_executor(None, input, "Type your response: ")
+    user_input = await loop.run_in_executor(None, Prompt.ask, "Type your response")
     return user_input