Merge pull request #25 from basedosdados/chore/handle-new-streaming

chore: handle new streaming

Author: vrtornisiello

frontend/components/chat_page.py

@@ -330,7 +330,7 @@ def render(self):
 
                 with st.status(label=label, state=state) as status:
                     for step in message_pair.safe_steps:
-                        st.caption(step.content)
+                        display_streaming_step(step)
 
                 if message_pair.assistant_message:
                     st.write(message_pair.assistant_message)
@@ -370,7 +370,7 @@ def render(self):
                         if streaming_status == "running":
                             step: Step = message
                             status.update(label=step.label)
-                            st.caption(step.content)
+                            display_streaming_step(step)
                         elif streaming_status == "complete":
                             message_pair: MessagePair = message
                             if message_pair.assistant_message:
@@ -438,3 +438,18 @@ def toggle_flag(flag_id: str):
         flag_id (str): The flag identifier.
     """
     st.session_state[flag_id] = not st.session_state[flag_id]
+
+def display_streaming_step(step: Step):
+    """Display a streaming step.
+
+    Args:
+        step (Step): The step to display, containing either a string
+            or a list of `StepContent` items.
+    """
+    if isinstance(step.content, str):
+        st.caption(step.content)
+    else:
+        for content in step.content:
+            if content.title:
+                st.caption(f"##### :green[{content.title}]")
+            st.caption(content.body.strip())

frontend/components/typewriter.py

@@ -12,11 +12,11 @@
 
         @keyframes blink-caret {
             from, to { border-color: transparent; }
-            50% { border-color: white; }
+            50% { border-color: #31333f; }
         }
 
         @keyframes disappear {
-            0% { border-color: white; }
+            0% { border-color: #31333f; }
             100% { border-color: transparent; }
         }
 
@@ -28,7 +28,7 @@
             overflow: hidden; /* Ensures the content is not revealed until the animation */
             padding-right: 0.1em; /* Adds space between the text and the cursor */
             margin: 0 auto; /* Centering effect */
-            border-right: 0.125em solid white; /* The typewriter cursor */
+            border-right: 0.125em solid #31333f; /* The typewriter cursor */
             animation:
                 typing 1s steps(80, end),
                 blink-caret 0.75s step-end infinite,

frontend/datatypes/datatypes.py

@@ -16,9 +16,13 @@ class UserMessage(BaseModel):
     id: str = Field(default_factory=lambda: str(uuid.uuid4()))
     content: str
 
+class StepContent(BaseModel):
+    title: str|None
+    body: str
+
 class Step(BaseModel):
     label: str
-    content: str
+    content: str|list[StepContent] # str allowed for backward compatibility
 
 class MessagePair(BaseModel):
     id: UUID4 = Field(default_factory=uuid.uuid4)