Merge pull request #192 from PrefectHQ/ensure-streaming

src/marvin/cli/tui.py

@@ -656,7 +656,11 @@ async def get_bot_response(self, event: Input.Submitted) -> str:
                 ),
             )
 
-            self.query_one("Conversation", Conversation)
+            # call once to populate the response in case there was any trouble
+            # streaming live
+            await self.stream_bot_response(
+                token_buffer=[response.content], response=bot_response
+            )
 
             # if this is one of the first few responses, rename the thread
             # appropriately

src/marvin/utilities/llms.py

@@ -51,10 +51,9 @@ def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
         """Run on new LLM token. Only available when streaming is enabled."""
         self.buffer.append(token)
         if self.on_token_callback is not None:
-            if inspect.iscoroutinefunction(self.on_token_callback):
+            output = self.on_token_callback(self.buffer)
+            if inspect.iscoroutine(output):
                 asyncio.run(self.on_token_callback(self.buffer))
-            else:
-                self.on_token_callback(self.buffer)
 
     def on_llm_end(self, response: LLMResult, **kwargs: Any) -> Any:
         """Run when LLM ends running."""