fix input_callback (#716)

Author: tastelikefeet

ms_agent/callbacks/input_callback.py

@@ -15,19 +15,18 @@ class InputCallback(Callback):
 
     def __init__(self, config: DictConfig):
         super().__init__(config)
-        self.chat_finished = False
 
-    async def on_generate_response(self, runtime: Runtime,
-                                   messages: List[Message]):
-        if messages[-1].tool_calls or messages[-1].role in ('tool',
-                                                            'user'):  # noqa
+    async def after_tool_call(self, runtime: Runtime, messages: List[Message]):
+        if messages[-1].tool_calls or messages[-1].role in ('tool', 'user'):
             return
 
-        query = input('>>>')
+        while True:
+            query = input('>>> ').strip()
+            if query:
+                break
+
         if not query:
-            self.chat_finished = True
+            runtime.should_stop = True
         else:
+            runtime.should_stop = False
             messages.append(Message(role='user', content=query))
-
-    async def after_tool_call(self, runtime: Runtime, messages: List[Message]):
-        runtime.should_stop = runtime.should_stop and self.chat_finished