code review for #1832

examples/chat_with_ai/log_callback_handler.py

@@ -1,55 +1,48 @@
+from typing import Any, Dict, Optional
+
 from langchain.callbacks.base import BaseCallbackHandler
 from langchain.schema import AgentAction, AgentFinish
-from typing import Dict, Any, Optional
-from nicegui.element import Element
+
+from nicegui import ui
 
 
 class NiceGuiLogElementCallbackHandler(BaseCallbackHandler):
-    """Callback Handler that writes to the log element of NicGui."""
+    """Callback Handler that writes to a log element."""
 
-    def __init__(self, element: Element) -> None:
+    def __init__(self, log_element: ui.log) -> None:
         """Initialize callback handler."""
-        self.element = element
-
-    def print_text(self, message: str) -> None:
-        self.element.push(message)
-        self.element.update()
+        self.log = log_element
 
-    def on_chain_start(
-        self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any
-    ) -> None:
+    def on_chain_start(self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any) -> None:
         """Print out that we are entering a chain."""
-        self.print_text(
-            f"\n\n> Entering new {serialized['id'][-1]} chain...",
-        )
+        self.log.push(f'\n\n> Entering new {serialized["id"][-1]} chain...')
 
     def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> None:
         """Print out that we finished a chain."""
-        self.print_text("\n> Finished chain.")
-        self.print_text(f"\nOutputs: {outputs}")
+        self.log.push('\n> Finished chain.')
+        self.log.push(f'\nOutputs: {outputs}')
 
     def on_agent_action(self, action: AgentAction, **kwargs: Any) -> Any:
         """Run on agent action."""
-        self.print_text(action.log)
-
-    def on_tool_end(
-        self,
-        output: str,
-        observation_prefix: Optional[str] = None,
-        llm_prefix: Optional[str] = None,
-        **kwargs: Any,
-    ) -> None:
+        self.log.push(action.log)
+
+    def on_tool_end(self,
+                    output: str,
+                    observation_prefix: Optional[str] = None,
+                    llm_prefix: Optional[str] = None,
+                    **kwargs: Any,
+                    ) -> None:
         """If not the final action, print out observation."""
         if observation_prefix is not None:
-            self.print_text(f"\n{observation_prefix}")
-        self.print_text(output)
+            self.log.push(f'\n{observation_prefix}')
+        self.log.push(output)
         if llm_prefix is not None:
-            self.print_text(f"\n{llm_prefix}")
+            self.log.push(f'\n{llm_prefix}')
 
     def on_text(self, text: str, **kwargs: Any) -> None:
         """Run when agent ends."""
-        self.print_text(text)
+        self.log.push(text)
 
     def on_agent_finish(self, finish: AgentFinish, **kwargs: Any) -> None:
         """Run on agent end."""
-        self.print_text(finish.log)
+        self.log.push(finish.log)

examples/chat_with_ai/main.py

@@ -5,8 +5,7 @@ from langchain.chains import ConversationChain
 from langchain.chat_models import ChatOpenAI
 from log_callback_handler import NiceGuiLogElementCallbackHandler
 
-import nicegui.globals
-from nicegui import Client, ui
+from nicegui import context, ui
 
 OPENAI_API_KEY = 'not-set'  # TODO: set your OpenAI API key here
 
@@ -22,7 +21,7 @@ def chat_messages() -> None:
         ui.chat_message(text=text, name=name, sent=name == 'You')
     if thinking:
         ui.spinner(size='3rem').classes('self-center')
-    if nicegui.globals.get_client().has_socket_connection:
+    if context.get_client().has_socket_connection:
         ui.run_javascript('window.scrollTo(0, document.body.scrollHeight)')