use st.status instead of st.expander; support output of langchain's A…

…gent

README.md

@@ -13,7 +13,8 @@ This package(>=1.0.0) will focus on wrapper of official chat elements to make ch
 A Streamlit component to show chat messages.
 It's basiclly a wrapper of streamlit officeial elements including the chat elemnts.
 
-![demo](https://github.com/liunux4odoo/streamlit-chatbox/blob/master/demo.gif?raw=true)
+![](demo.gif)
+![](demo_agent.gif)
 
 ## Features
 
@@ -83,10 +84,10 @@ if query := st.chat_input('input your question here'):
         text = ""
         for x, docs in generator:
             text += x
-            chat_box.update_msg(text, 0, streaming=True)
-            chat_box.update_msg("\n\n".join(docs), 1, streaming=False)
+            chat_box.update_msg(text, element_index=0, streaming=True)
         # update the element without focus
-        chat_box.update_msg(text, 0, streaming=False)
+        chat_box.update_msg(text, element_index=0, streaming=False, state="complete")
+        chat_box.update_msg("\n\n".join(docs), element_index=1, streaming=False, state="complete")
     else:
         text, docs = llm.chat(query)
         chat_box.ai_say(
@@ -98,7 +99,8 @@ if query := st.chat_input('input your question here'):
             ]
         )
 
-if st.button('show me the multimedia'):
+cols = st.columns(2)
+if cols[0].button('show me the multimedia'):
     chat_box.ai_say(Image(
         'https://tse4-mm.cn.bing.net/th/id/OIP-C.cy76ifbr2oQPMEs2H82D-QHaEv?w=284&h=181&c=7&r=0&o=5&dpr=1.5&pid=1.7'))
     time.sleep(0.5)
@@ -108,6 +110,29 @@ if st.button('show me the multimedia'):
     chat_box.ai_say(
         Audio('https://sample-videos.com/video123/mp4/720/big_buck_bunny_720p_1mb.mp4'))
 
+if cols[1].button('run agent'):
+    chat_box.user_say('run agent')
+    agent = FakeAgent()
+    text = ""
+
+    if streaming:
+    # streaming:
+    chat_box.ai_say() # generate a blank placeholder to render messages
+    for d in agent.run_stream():
+        if d["type"] == "complete":
+            chat_box.update_msg(expanded=False, state="complete")
+            chat_box.insert_msg(d["llm_output"])
+            break
+
+        if d["status"] == 1:
+            chat_box.update_msg(expanded=False, state="complete")
+            text = ""
+            chat_box.insert_msg(Markdown(text, title=d["text"], in_expander=True, expanded=True))
+        elif d["status"] == 2:
+            text += d["llm_output"]
+            chat_box.update_msg(text, streaming=True)
+        else:
+            chat_box.update_msg(text, streaming=False)
 
 btns.download_button(
     "Export Markdown",
@@ -130,7 +155,6 @@ if btns.button("clear history"):
 
 if show_history:
     st.write(chat_box.history)
-
 ```
 
 ## Todos
@@ -160,3 +184,5 @@ if show_history:
 	- [x] export to markdown
 	- [x] export to json
     - [x] import json
+
+- [x] support output of langchain' Agent.

example.py

@@ -47,10 +47,10 @@
         text = ""
         for x, docs in generator:
             text += x
-            chat_box.update_msg(text, 0, streaming=True)
-            chat_box.update_msg("\n\n".join(docs), 1, streaming=False)
+            chat_box.update_msg(text, element_index=0, streaming=True)
         # update the element without focus
-        chat_box.update_msg(text, 0, streaming=False)
+        chat_box.update_msg(text, element_index=0, streaming=False, state="complete")
+        chat_box.update_msg("\n\n".join(docs), element_index=1, streaming=False, state="complete")
     else:
         text, docs = llm.chat(query)
         chat_box.ai_say(
@@ -62,7 +62,8 @@
             ]
         )
 
-if st.button('show me the multimedia'):
+cols = st.columns(2)
+if cols[0].button('show me the multimedia'):
     chat_box.ai_say(Image(
         'https://tse4-mm.cn.bing.net/th/id/OIP-C.cy76ifbr2oQPMEs2H82D-QHaEv?w=284&h=181&c=7&r=0&o=5&dpr=1.5&pid=1.7'))
     time.sleep(0.5)
@@ -72,6 +73,28 @@
     chat_box.ai_say(
         Audio('https://sample-videos.com/video123/mp4/720/big_buck_bunny_720p_1mb.mp4'))
 
+if cols[1].button('run agent'):
+    chat_box.user_say('run agent')
+    agent = FakeAgent()
+    text = ""
+
+    # streaming:
+    chat_box.ai_say() # generate a blank placeholder to render messages
+    for d in agent.run_stream():
+        if d["type"] == "complete":
+            chat_box.update_msg(expanded=False, state="complete")
+            chat_box.insert_msg(d["llm_output"])
+            break
+
+        if d["status"] == 1:
+            chat_box.update_msg(expanded=False, state="complete")
+            text = ""
+            chat_box.insert_msg(Markdown(text, title=d["text"], in_expander=True, expanded=True))
+        elif d["status"] == 2:
+            text += d["llm_output"]
+            chat_box.update_msg(text, streaming=True)
+        else:
+            chat_box.update_msg(text, streaming=False)
 
 btns.download_button(
     "Export Markdown",

setup.py

@@ -9,7 +9,7 @@ def readme():
 
 setuptools.setup(
     name='streamlit-chatbox',
-    version='1.1.7',
+    version='1.1.8',
     author='liunux',
     author_email='[email protected]',
     description='A chat box and some helpful tools used to build chatbot app with streamlit',
@@ -21,7 +21,7 @@ def readme():
     classifiers=[],
     python_requires='>=3.8',
     install_requires=[
-        'streamlit>=1.24.0',
+        'streamlit>=1.26.0',
         'simplejson',
     ]
 )

streamlit_chatbox/__init__.py

@@ -14,6 +14,7 @@
     "Video",
     "OutputElement",
     "FakeLLM",
+    "FakeAgent",
 ]
 
 

streamlit_chatbox/elements.py

@@ -1,6 +1,6 @@
 from typing import *
 import streamlit as st
-from pydantic import BaseModel, Field
+# from pydantic import BaseModel, Field
 
 
 class Element:
@@ -9,12 +9,13 @@ class Element:
     '''
 
     def __init__(self,
+                 *,
                  output_method: str = "markdown",
-                 *args: Any,
+                 metadata: Dict = {},
                  **kwargs: Any,
                  ) -> None:
         self._output_method = output_method
-        self._args = args
+        self._metadata = metadata
         self._kwargs = kwargs
         self._defualt_kwargs = {
             "markdown": {
@@ -39,15 +40,28 @@ def _set_default_kwargs(self) -> None:
             for k, v in default.items():
                 self._kwargs.setdefault(k, v)
 
-    def __call__(self) -> st._DeltaGenerator:
+    def __call__(self, render_to: st._DeltaGenerator=None) -> st._DeltaGenerator:
         # assert self._dg is None, "Every element can be rendered once only."
-        self._place_holder = st.empty()
+        render_to = render_to or st
+        self._place_holder = render_to.empty()
         output_method = getattr(self._place_holder, self._output_method)
         assert callable(
             output_method), f"The attribute st.{self._output_mehtod} is not callable."
         self._dg = output_method(*self._args, **self._kwargs)
         return self._dg
 
+    @property
+    def dg(self) -> st._DeltaGenerator:
+        return self._dg
+
+    @property
+    def place_holder(self) -> st._DeltaGenerator:
+        return self._place_holder
+
+    @property
+    def metadata(self) -> Dict:
+        return self._metadata
+
 
 class OutputElement(Element):
     def __init__(self,
@@ -56,25 +70,31 @@ def __init__(self,
                  title: str = "",
                  in_expander: bool = False,
                  expanded: bool = False,
+                 state: Literal["running", "complete", "error"] = "running",
                  **kwargs: Any,
                  ) -> None:
         super().__init__(output_method=output_method, **kwargs)
         self._content = content
         self._title = title
         self._in_expander = in_expander
         self._expanded = expanded
+        self._state = state
+        self._attrs = ["_content", "_output_method", "_kwargs", "_metadata",
+                       "_title", "_in_expander", "_expanded", "_state",]
 
-    def __call__(self) -> st._DeltaGenerator:
-        self._args = (self._content,)
-        if self._in_expander:
-            with st.expander(self._title, self._expanded):
-                return super().__call__()
-        else:
-            return super().__call__()
+    def clone(self) -> "OutputElement":
+        obj = type(self)()
+        for n in self._attrs:
+            setattr(obj, n, getattr(self, n))
+        return obj
+
+    @property
+    def content(self) -> Union[str, bytes]:
+        return self._content
 
     def __repr__(self) -> str:
         method = self._output_method.capitalize()
-        return f"{method} Element:\n{self._content}"
+        return f"{method} Element:\n{self.content}"
 
     def to_dict(self) -> Dict:
         return {
@@ -83,6 +103,8 @@ def to_dict(self) -> Dict:
             "title": self._title,
             "in_expander": self._in_expander,
             "expanded": self._expanded,
+            "state": self._state,
+            "metadata": self._metadata,
             "kwargs": self._kwargs,
         }
 
@@ -101,25 +123,61 @@ def from_dict(cls, d: Dict) -> "OutputElement":
             title=d.get("title"),
             in_expander=d.get("in_expander"),
             expanded=d.get("expanded"),
+            state=d.get("state"),
+            metadata=d.get("metadata", {}),
             **d.get("kwargs", {}),
             )
         if factory_cls is cls:
             kwargs["output_method"] = d.get("output_method")
 
         return factory_cls(**kwargs)
 
+    def __call__(self, render_to: st._DeltaGenerator=None, direct: bool=False) -> st._DeltaGenerator:
+        if render_to is None:
+            if self._place_holder is None:
+                self._place_holder = st.empty()
+        else:
+            if direct:
+                self._place_holder = render_to
+            else:
+                self._place_holder = render_to.empty()
+        temp_dg = self._place_holder
+
+        if self._in_expander:
+            temp_dg = self._place_holder.status(self._title, expanded=self._expanded, state=self._state)
+        output_method = getattr(temp_dg, self._output_method)
+        assert callable(
+            output_method), f"The attribute st.{self._output_mehtod} is not callable."
+        self._dg = output_method(self._content, **self._kwargs)
+        return self._dg
+
     def update_element(
         self,
-        element: "OutputElement",
-        streaming: Optional[bool] = None,
+        element: "OutputElement" = None,
+        *,
+        title: str = None,
+        expanded: bool = None,
+        state: bool = None,
     ) -> st._DeltaGenerator:
-        assert self._place_holder is not None, "You must render the element before setting new element."
-        with self._place_holder:
-            self._dg = element()
+        assert self.place_holder is not None, f"You must render the element {self} before setting new element."
+        attrs = {}
+        if title is not None:
+            attrs["_title"] = title
+        if expanded is not None:
+            attrs["_expanded"] = expanded
+        if state is not None:
+            attrs["_state"] = state
+
+        if element is None:
+            element = self
+        for k, v in attrs.items():
+            setattr(element, k, v)
+
+        element(self.place_holder, direct=True)
         return self._dg
 
-    def attrs_from(self, target):
-        for attr in ["_in_expander", "_expanded", "_title"]:
+    def status_from(self, target):
+        for attr in ["_in_expander", "_expanded", "_title", "_state"]:
             setattr(self, attr, getattr(target, attr))
 
 
@@ -128,24 +186,60 @@ class InputElement(Element):
 
 
 class Markdown(OutputElement):
-    def __init__(self, content: Union[str, bytes] = "", title: str = "", in_expander: bool = False, expanded: bool = False, **kwargs: Any) -> None:
+    def __init__(
+        self,
+        content: Union[str, bytes] = "",
+        title: str = "",
+        in_expander: bool = False,
+        expanded: bool = False,
+        state: Literal["running", "complete", "error"] = "running",
+        **kwargs: Any,
+    ) -> None:
         super().__init__(content, output_method="markdown", title=title,
-                         in_expander=in_expander, expanded=expanded, **kwargs)
+                         in_expander=in_expander, expanded=expanded,
+                         state=state, **kwargs)
 
 
 class Image(OutputElement):
-    def __init__(self, content: Union[str, bytes] = "", title: str = "", in_expander: bool = False, expanded: bool = False, **kwargs: Any) -> None:
+    def __init__(
+        self,
+        content: Union[str, bytes] = "",
+        title: str = "",
+        in_expander: bool = False,
+        expanded: bool = False,
+        state: Literal["running", "complete", "error"] = "running",
+        **kwargs: Any,
+    ) -> None:
         super().__init__(content, output_method="image", title=title,
-                         in_expander=in_expander, expanded=expanded, **kwargs)
+                         in_expander=in_expander, expanded=expanded,
+                         state=state, **kwargs)
 
 
 class Audio(OutputElement):
-    def __init__(self, content: Union[str, bytes] = "", title: str = "", in_expander: bool = False, expanded: bool = False, **kwargs: Any) -> None:
+    def __init__(
+        self,
+        content: Union[str, bytes] = "",
+        title: str = "",
+        in_expander: bool = False,
+        expanded: bool = False,
+        state: Literal["running", "complete", "error"] = "running",
+        **kwargs: Any,
+    ) -> None:
         super().__init__(content, output_method="audio", title=title,
-                         in_expander=in_expander, expanded=expanded, **kwargs)
+                         in_expander=in_expander, expanded=expanded,
+                         state=state, **kwargs)
 
 
 class Video(OutputElement):
-    def __init__(self, content: Union[str, bytes] = "", title: str = "", in_expander: bool = False, expanded: bool = False, **kwargs: Any) -> None:
+    def __init__(
+        self,
+        content: Union[str, bytes] = "",
+        title: str = "",
+        in_expander: bool = False,
+        expanded: bool = False,
+        state: Literal["running", "complete", "error"] = "running",
+        **kwargs: Any,
+    ) -> None:
         super().__init__(content, output_method="video", title=title,
-                         in_expander=in_expander, expanded=expanded, **kwargs)
+                         in_expander=in_expander, expanded=expanded,
+                         state=state, **kwargs)