fix: predictive state examples

Fix predictive state examples by renaming the `write_document` tool to
`write_document_local` and updating the metadata accordingly. This
ensures that the examples correctly reflect the intended functionality
of the predictive state updates, while still allowing the Pydantic AI
example to leverage the AG-UI defined `write_document` tool.

Restore the `confirm_changes` tool to enable the confirmation of
changes for existing examples.

Author: stevenh

typescript-sdk/.gitignore

@@ -41,4 +41,7 @@ yarn-error.log*
 packages/proto/src/generated
 
 # LangGraph API
-**/**/.langgraph_api
+**/**/.langgraph_api
+
+# Python
+__pycache__/

typescript-sdk/apps/dojo/src/app/[integrationId]/feature/predictive_state_updates/page.tsx

@@ -127,6 +127,28 @@ const DocumentEditor = () => {
     }
   }, [text]);
 
+  // TODO(steve): Remove this when all agents have been updated to use write_document tool.
+  useCopilotAction({
+    name: "confirm_changes",
+    renderAndWaitForResponse: ({ args, respond, status }) => (
+      <ConfirmChanges
+        args={args}
+        respond={respond}
+        status={status}
+        onReject={() => {
+          editor?.commands.setContent(fromMarkdown(currentDocument));
+          setAgentState({ document: currentDocument });
+        }}
+        onConfirm={() => {
+          editor?.commands.setContent(fromMarkdown(agentState?.document || ""));
+          setCurrentDocument(agentState?.document || "");
+          setAgentState({ document: agentState?.document || "" });
+        }}
+      />
+    ),
+  });
+
+  // Action to write the document.
   useCopilotAction({
     name: "write_document",
     description: `Present the proposed changes to the user for review`,

typescript-sdk/integrations/crewai/python/ag_ui_crewai/examples/predictive_state_updates.py

@@ -16,7 +16,7 @@
 WRITE_DOCUMENT_TOOL = {
     "type": "function",
     "function": {
-        "name": "write_document",
+        "name": "write_document_local",
         "description": " ".join("""
             Write a document. Use markdown formatting to format the document.
             It's good to format the document extensively so it's easy to read.
@@ -63,19 +63,19 @@ async def chat(self):
         Standard chat node.
         """
         system_prompt = f"""
-        You are a helpful assistant for writing documents. 
-        To write the document, you MUST use the write_document tool.
+        You are a helpful assistant for writing documents.
+        To write the document, you MUST use the write_document_local tool.
         You MUST write the full document, even when changing only a few words.
         When you wrote the document, DO NOT repeat it as a message. 
         Just briefly summarize the changes you made. 2 sentences max.
         This is the current state of the document: ----\n {self.state.document}\n-----
         """
 
-        # 1. Here we specify that we want to stream the tool call to write_document
+        # 1. Here we specify that we want to stream the tool call to write_document_local
         #    to the frontend as state.
         await copilotkit_predict_state({
             "document": {
-                "tool_name": "write_document",
+                "tool_name": "write_document_local",
                 "tool_argument": "document"
             }
         })
@@ -122,7 +122,7 @@ async def chat(self):
             tool_call_name = tool_call["function"]["name"]
             tool_call_args = json.loads(tool_call["function"]["arguments"])
 
-            if tool_call_name == "write_document":
+            if tool_call_name == "write_document_local":
                 self.state.document = tool_call_args["document"]
 
                 # 4.1 Append the result to the messages in state

typescript-sdk/integrations/langgraph/examples/agents/predictive_state_updates/agent.py

@@ -18,7 +18,7 @@
 WRITE_DOCUMENT_TOOL = {
     "type": "function",
     "function": {
-        "name": "write_document",
+        "name": "write_document_local",
         "description": " ".join("""
             Write a document. Use markdown formatting to format the document.
             It's good to format the document extensively so it's easy to read.
@@ -64,8 +64,8 @@ async def chat_node(state: AgentState, config: RunnableConfig):
     """
 
     system_prompt = f"""
-    You are a helpful assistant for writing documents. 
-    To write the document, you MUST use the write_document tool.
+    You are a helpful assistant for writing documents.
+    To write the document, you MUST use the write_document_local tool.
     You MUST write the full document, even when changing only a few words.
     When you wrote the document, DO NOT repeat it as a message. 
     Just briefly summarize the changes you made. 2 sentences max.
@@ -79,10 +79,10 @@ async def chat_node(state: AgentState, config: RunnableConfig):
     if config is None:
         config = RunnableConfig(recursion_limit=25)
 
-    # Use "predict_state" metadata to set up streaming for the write_document tool
+    # Use "predict_state" metadata to set up streaming for the write_document_local tool
     config["metadata"]["predict_state"] = [{
         "state_key": "document",
-        "tool": "write_document",
+        "tool": "write_document_local",
         "tool_argument": "document"
     }]
 
@@ -120,7 +120,7 @@ async def chat_node(state: AgentState, config: RunnableConfig):
             tool_call_name = tool_call.name
             tool_call_args = tool_call.args
 
-        if tool_call_name == "write_document":
+        if tool_call_name == "write_document_local":
             # Add the tool response to messages
             tool_response = {
                 "role": "tool",

typescript-sdk/integrations/langgraph/python/ag_ui_langgraph/examples/agents/predictive_state_updates.py

@@ -19,7 +19,7 @@
 WRITE_DOCUMENT_TOOL = {
     "type": "function",
     "function": {
-        "name": "write_document",
+        "name": "write_document_local",
         "description": " ".join("""
             Write a document. Use markdown formatting to format the document.
             It's good to format the document extensively so it's easy to read.
@@ -66,7 +66,7 @@ async def chat_node(state: AgentState, config: RunnableConfig):
 
     system_prompt = f"""
     You are a helpful assistant for writing documents. 
-    To write the document, you MUST use the write_document tool.
+    To write the document, you MUST use the write_document_local tool.
     You MUST write the full document, even when changing only a few words.
     When you wrote the document, DO NOT repeat it as a message. 
     Just briefly summarize the changes you made. 2 sentences max.
@@ -80,10 +80,10 @@ async def chat_node(state: AgentState, config: RunnableConfig):
     if config is None:
         config = RunnableConfig(recursion_limit=25)
 
-    # Use "predict_state" metadata to set up streaming for the write_document tool
+    # Use "predict_state" metadata to set up streaming for the write_document_local tool
     config["metadata"]["predict_state"] = [{
         "state_key": "document",
-        "tool": "write_document",
+        "tool": "write_document_local",
         "tool_argument": "document"
     }]
 
@@ -121,7 +121,7 @@ async def chat_node(state: AgentState, config: RunnableConfig):
             tool_call_name = tool_call.name
             tool_call_args = tool_call.args
 
-        if tool_call_name == "write_document":
+        if tool_call_name == "write_document_local":
             # Add the tool response to messages
             tool_response = {
                 "role": "tool",

typescript-sdk/integrations/server-starter-all-features/server/python/example_server/predictive_state_updates.py

@@ -80,7 +80,7 @@ def make_story(name: str) -> str:
 async def send_tool_call_events():
     """Send tool call events with predictive state and incremental story generation"""
     tool_call_id = str(uuid.uuid4())
-    tool_call_name = "write_document"
+    tool_call_name = "write_document_local"
 
     # Generate a random story
     story = make_story(random.choice(dog_names))
@@ -93,13 +93,13 @@ async def send_tool_call_events():
         value=[
             {
                 "state_key": "document",
-                "tool": "write_document",
+                "tool": "write_document_local",
                 "tool_argument": "document"
             }
         ]
     )
 
-    # First tool call: write_document
+    # First tool call: write_document_local
     yield ToolCallStartEvent(
         type=EventType.TOOL_CALL_START,
         tool_call_id=tool_call_id,