fix: predictive state examples

stevenh · stevenh · commit fc29b33ec926 · 2025-07-22T18:27:49.000+01:00
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.
diff --git a/typescript-sdk/.gitignore b/typescript-sdk/.gitignore
@@ -41,4 +41,7 @@ yarn-error.log*
 packages/proto/src/generated
 
 # LangGraph API
-**/**/.langgraph_api
+**/**/.langgraph_api
+
+# Python
+__pycache__/
diff --git a/typescript-sdk/apps/dojo/src/app/[integrationId]/feature/predictive_state_updates/page.tsx b/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`,
diff --git a/typescript-sdk/integrations/crewai/python/ag_ui_crewai/examples/predictive_state_updates.py b/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
diff --git a/typescript-sdk/integrations/langgraph/examples/agents/predictive_state_updates/agent.py b/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,25 +64,25 @@ 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. 
+    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 {state.get('document')}\n-----
     """
 
     # Define the model
     model = ChatOpenAI(model="gpt-4o")
-    
+
     # Define config for the model with emit_intermediate_state to stream tool calls to frontend
     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"
     }]
 
@@ -104,11 +104,11 @@ async def chat_node(state: AgentState, config: RunnableConfig):
 
     # Update messages with the response
     messages = state["messages"] + [response]
-    
+
     # Extract any tool calls from the response
     if hasattr(response, "tool_calls") and response.tool_calls:
         tool_call = response.tool_calls[0]
-        
+
         # Handle tool_call as a dictionary or an object
         if isinstance(tool_call, dict):
             tool_call_id = tool_call["id"]
@@ -120,15 +120,28 @@ 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",
                 "content": "Document written.",
                 "tool_call_id": tool_call_id
             }
 
-            messages = messages + [tool_response]
+            # Add confirmation tool call
+            confirm_tool_call = {
+                "role": "assistant",
+                "content": "",
+                "tool_calls": [{
+                    "id": str(uuid.uuid4()),
+                    "function": {
+                        "name": "confirm_changes",
+                        "arguments": "{}"
+                    }
+                }]
+            }
+
+            messages = messages + [tool_response, confirm_tool_call]
 
             # Return Command to route to end
             return Command(
@@ -138,7 +151,7 @@ async def chat_node(state: AgentState, config: RunnableConfig):
                     "document": tool_call_args["document"]
                 }
             )
-    
+
     # If no tool was called, go to end
     return Command(
         goto=END,
diff --git a/typescript-sdk/integrations/langgraph/python/ag_ui_langgraph/examples/agents/predictive_state_updates.py b/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",
diff --git a/typescript-sdk/integrations/server-starter-all-features/server/python/example_server/predictive_state_updates.py b/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,

Original file line number	Diff line number	Diff line change
`@@ -80,7 +80,7 @@ def make_story(name: str) -> str:`
`80`	`80`	`async def send_tool_call_events():`
`81`	`81`	`"""Send tool call events with predictive state and incremental story generation"""`
`82`	`82`	`tool_call_id = str(uuid.uuid4())`
`83`		`- tool_call_name = "write_document"`
	`83`	`+ tool_call_name = "write_document_local"`
`84`	`84`
`85`	`85`	`# Generate a random story`
`86`	`86`	`story = make_story(random.choice(dog_names))`
`@@ -93,13 +93,13 @@ async def send_tool_call_events():`
`93`	`93`	`value=[`
`94`	`94`	`{`
`95`	`95`	`"state_key": "document",`
`96`		`- "tool": "write_document",`
	`96`	`+ "tool": "write_document_local",`
`97`	`97`	`"tool_argument": "document"`
`98`	`98`	`}`
`99`	`99`	`]`
`100`	`100`	`)`
`101`	`101`
`102`		`- # First tool call: write_document`
	`102`	`+ # First tool call: write_document_local`
`103`	`103`	`yield ToolCallStartEvent(`
`104`	`104`	`type=EventType.TOOL_CALL_START,`
`105`	`105`	`tool_call_id=tool_call_id,`