Browse Source

Added LangChain handler to the examples.

Modeled after the [Streamlit handler](https://github.com/langchain-ai/langchain/blob/f4bec9686d6a2d07a0f0540b34f9313b0cd5ae89/libs/langchain/langchain/callbacks/streamlit/streamlit_callback_handler.py#L225C4-L225C4).
Cam 1 year ago
parent
commit
0114f1d56b
2 changed files with 104 additions and 23 deletions
  1. 54 0
      examples/chat_with_ai/log_callback_handler.py
  2. 50 23
      examples/chat_with_ai/main.py

+ 54 - 0
examples/chat_with_ai/log_callback_handler.py

@@ -0,0 +1,54 @@
+from langchain.callbacks.base import BaseCallbackHandler
+from langchain.schema import AgentAction, AgentFinish
+from typing import Dict, Any, Optional
+from nicegui.element import Element
+
+
+class NiceGuiLogElementCallbackHandler(BaseCallbackHandler):
+    """Callback Handler that writes to the log element of NicGui."""
+
+    def __init__(self, element: Element) -> None:
+        """Initialize callback handler."""
+        self.element = element
+
+    def print_text(self, message: str) -> None:
+        self.element.push(message)
+
+    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...",
+        )
+
+    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}")
+
+    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:
+        """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)
+        if llm_prefix is not None:
+            self.print_text(f"\n{llm_prefix}")
+
+    def on_text(self, text: str, **kwargs: Any) -> None:
+        """Run when agent ends."""
+        self.print_text(text)
+
+    def on_agent_finish(self, finish: AgentFinish, **kwargs: Any) -> None:
+        """Run on agent end."""
+        self.print_text(finish.log)

+ 50 - 23
examples/chat_with_ai/main.py

@@ -5,10 +5,13 @@ from langchain.chains import ConversationChain
 from langchain.chat_models import ChatOpenAI
 
 from nicegui import Client, ui
+from log_callback_handler import NiceGuiLogElementCallbackHandler
 
-OPENAI_API_KEY = 'not-set'  # TODO: set your OpenAI API key here
+OPENAI_API_KEY = "not-set"  # TODO: set your OpenAI API key here
 
-llm = ConversationChain(llm=ChatOpenAI(model_name='gpt-3.5-turbo', openai_api_key=OPENAI_API_KEY))
+llm = ConversationChain(
+    llm=ChatOpenAI(model_name="gpt-3.5-turbo", openai_api_key=OPENAI_API_KEY),
+)
 
 messages: List[Tuple[str, str]] = []
 thinking: bool = False
@@ -17,41 +20,65 @@ thinking: bool = False
 @ui.refreshable
 async def chat_messages() -> None:
     for name, text in messages:
-        ui.chat_message(text=text, name=name, sent=name == 'You')
+        ui.chat_message(text=text, name=name, sent=name == "You")
     if thinking:
-        ui.spinner(size='3rem').classes('self-center')
-    await ui.run_javascript('window.scrollTo(0, document.body.scrollHeight)', respond=False)
+        ui.spinner(size="3rem").classes("self-center")
+    await ui.run_javascript(
+        "window.scrollTo(0, document.body.scrollHeight)", respond=False
+    )
 
 
-@ui.page('/')
+@ui.page("/")
 async def main(client: Client):
+    with ui.tabs().classes("w-full") as tabs:
+        chat_panel = ui.tab("Chat")
+        logs_panel = ui.tab("Logs")
+    with ui.tab_panels(tabs).classes("w-full"):
+        with ui.tab_panel(logs_panel):
+            log = ui.log().classes("w-full h-full")
+
     async def send() -> None:
         global thinking
         message = text.value
-        messages.append(('You', text.value))
+        messages.append(("You", text.value))
         thinking = True
-        text.value = ''
+        text.value = ""
         chat_messages.refresh()
 
-        response = await llm.arun(message)
-        messages.append(('Bot', response))
+        response = await llm.arun(
+            message, callbacks=[NiceGuiLogElementCallbackHandler(log)]
+        )
+        messages.append(("Bot", response))
         thinking = False
         chat_messages.refresh()
 
-    anchor_style = r'a:link, a:visited {color: inherit !important; text-decoration: none; font-weight: 500}'
-    ui.add_head_html(f'<style>{anchor_style}</style>')
+    anchor_style = r"a:link, a:visited {color: inherit !important; text-decoration: none; font-weight: 500}"
+    ui.add_head_html(f"<style>{anchor_style}</style>")
     await client.connected()
 
-    with ui.column().classes('w-full max-w-2xl mx-auto items-stretch'):
-        await chat_messages()
+    with ui.tab_panels(tabs, value=chat_panel).classes("w-full"):
+        with ui.tab_panel(chat_panel):
+            with ui.column().classes("w-full max-w-2xl mx-auto items-stretch"):
+                await chat_messages()
+
+    with ui.footer().classes("bg-white"), ui.column().classes(
+        "w-full max-w-3xl mx-auto my-6"
+    ):
+        with ui.row().classes("w-full no-wrap items-center"):
+            placeholder = (
+                "message"
+                if OPENAI_API_KEY != "not-set"
+                else "Please provide your OPENAI key in the Python script first!"
+            )
+            text = (
+                ui.input(placeholder=placeholder)
+                .props("rounded outlined input-class=mx-3")
+                .classes("w-full self-center")
+                .on("keydown.enter", send)
+            )
+        ui.markdown("simple chat app built with [NiceGUI](https://nicegui.io)").classes(
+            "text-xs self-end mr-8 m-[-1em] text-primary"
+        )
 
-    with ui.footer().classes('bg-white'), ui.column().classes('w-full max-w-3xl mx-auto my-6'):
-        with ui.row().classes('w-full no-wrap items-center'):
-            placeholder = 'message' if OPENAI_API_KEY != 'not-set' else \
-                'Please provide your OPENAI key in the Python script first!'
-            text = ui.input(placeholder=placeholder).props('rounded outlined input-class=mx-3') \
-                .classes('w-full self-center').on('keydown.enter', send)
-        ui.markdown('simple chat app built with [NiceGUI](https://nicegui.io)') \
-            .classes('text-xs self-end mr-8 m-[-1em] text-primary')
 
-ui.run(title='Chat with GPT-3 (example)')
+ui.run(title="Chat with GPT-3 (example)")