From f9f8c19ec4aef3469b481315d45c67e8ec306403 Mon Sep 17 00:00:00 2001
From: Tat Dat Duong <david@duong.cz>
Date: Thu, 3 Apr 2025 00:41:16 +0200
Subject: [PATCH] Update docstrings

---
 libs/prebuilt/langgraph/prebuilt/ui.py | 97 +++++++++++++++++++++++++-
 1 file changed, 94 insertions(+), 3 deletions(-)

diff --git a/libs/prebuilt/langgraph/prebuilt/ui.py b/libs/prebuilt/langgraph/prebuilt/ui.py
index ab7958f63..57f3b5550 100644
--- a/libs/prebuilt/langgraph/prebuilt/ui.py
+++ b/libs/prebuilt/langgraph/prebuilt/ui.py
@@ -9,6 +9,19 @@ from langgraph.utils.config import get_config, get_stream_writer
 
 
 class UIMessage(TypedDict):
+    """A message type for UI updates in LangGraph.
+
+    This TypedDict represents a UI message that can be sent to update the UI state.
+    It contains information about the UI component to render and its properties.
+
+    Attributes:
+        type: Literal type indicating this is a UI message.
+        id: Unique identifier for the UI message.
+        name: Name of the UI component to render.
+        props: Properties to pass to the UI component.
+        metadata: Additional metadata about the UI message.
+    """
+
     type: Literal["ui"]
     id: str
     name: str
@@ -17,6 +30,16 @@ class UIMessage(TypedDict):
 
 
 class RemoveUIMessage(TypedDict):
+    """A message type for removing UI components in LangGraph.
+
+    This TypedDict represents a message that can be sent to remove a UI component
+    from the current state.
+
+    Attributes:
+        type: Literal type indicating this is a remove-ui message.
+        id: Unique identifier of the UI component to remove.
+    """
+
     type: Literal["remove-ui"]
     id: str
 
@@ -33,7 +56,35 @@ def push_ui_message(
     message: Optional[AnyMessage] = None,
     state_key: str = "ui",
 ) -> UIMessage:
-    """Push a new UI message."""
+    """Push a new UI message to update the UI state.
+
+    This function creates and sends a UI message that will be rendered in the UI.
+    It also updates the graph state with the new UI message.
+
+    Args:
+        name: Name of the UI component to render.
+        props: Properties to pass to the UI component.
+        id: Optional unique identifier for the UI message.
+            If not provided, a random UUID will be generated.
+        metadata: Optional additional metadata about the UI message.
+        message: Optional message object to associate with the UI message.
+        state_key: Key in the graph state where the UI messages are stored.
+            Defaults to "ui".
+
+    Returns:
+        The created UI message.
+
+    Example:
+
+    .. code-block:: python
+
+        message = push_ui_message(
+            name="component-name",
+            props={"content": "Hello world"},
+        )
+
+    """
+
     writer = get_stream_writer()
     config = get_config()
 
@@ -66,7 +117,25 @@ def push_ui_message(
 
 
 def remove_ui_message(id: str, *, state_key: str = "ui") -> RemoveUIMessage:
-    """Delete a UI message by ID."""
+    """Delete a UI message by ID from the UI state.
+
+    This function creates and sends a message to remove a UI component from the current state.
+    It also updates the graph state to remove the UI message.
+
+    Args:
+        id: Unique identifier of the UI component to remove.
+        state_key: Key in the graph state where the UI messages are stored. Defaults to "ui".
+
+    Returns:
+        The remove UI message.
+
+    Example:
+
+    .. code-block:: python
+
+        remove_message = remove_ui_message("message-123")
+
+    """
     writer = get_stream_writer()
     config = get_config()
 
@@ -82,7 +151,29 @@ def reduce_ui_messages(
     left: Union[list[AnyUIMessage], AnyUIMessage],
     right: Union[list[AnyUIMessage], AnyUIMessage],
 ) -> list[AnyUIMessage]:
-    """Merge two lists of UI messages, supporting removing UI messages."""
+    """Merge two lists of UI messages, supporting removing UI messages.
+
+    This function combines two lists of UI messages, handling both regular UI messages
+    and remove-ui messages. When a remove-ui message is encountered, it removes any
+    UI message with the matching ID from the current state.
+
+    Args:
+        left: First list of UI messages or single UI message.
+        right: Second list of UI messages or single UI message.
+
+    Returns:
+        Combined list of UI messages with removals applied.
+
+    Example:
+
+    .. code-block:: python
+
+        messages = reduce_ui_messages(
+            [{"type": "ui", "id": "1", "name": "Chat", "props": {}}],
+            {"type": "remove-ui", "id": "1"}
+        )
+
+    """
     if not isinstance(left, list):
         left = [left]