Merge pull request #292 from ag-ui-protocol/fix/include-payload-with-regenerated-message

fix: include payload with regenerated message

Author: ranst91

typescript-sdk/integrations/langgraph/python/ag_ui_langgraph/agent.py

@@ -685,7 +685,15 @@ async def get_checkpoint_before_message(self, message_id: str, thread_id: str):
                     empty_snapshot = snapshot
                     empty_snapshot.values["messages"] = []
                     return empty_snapshot
-                return history_list[idx - 1]  # return one snapshot *before* the one that includes the message
+
+                snapshot_values_without_messages = snapshot.values.copy()
+                del snapshot_values_without_messages["messages"]
+                checkpoint = history_list[idx - 1]
+
+                merged_values = {**checkpoint.values, **snapshot_values_without_messages}
+                checkpoint = checkpoint._replace(values=merged_values)
+
+                return checkpoint
 
         raise ValueError("Message ID not found in history")
 

typescript-sdk/integrations/langgraph/python/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "ag-ui-langgraph"
-version = "0.0.5"
+version = "0.0.6"
 description = "Implementation of the AG-UI protocol for LangGraph."
 authors = ["Ran Shem Tov <[email protected]>"]
 readme = "README.md"

typescript-sdk/integrations/langgraph/src/agent.ts

@@ -992,7 +992,9 @@ export class LangGraphAgent extends AbstractAgent {
 
     const targetStateIndex = reversed.indexOf(targetState);
 
-    return reversed[targetStateIndex - 1] ?? { ...targetState, values: {} };
+    const { messages, ...targetStateValuesWithoutMessages } = targetState.values as State
+    const selectedCheckpoint = reversed[targetStateIndex - 1] ?? { ...targetState, values: {} }
+    return { ...selectedCheckpoint, values: { ...selectedCheckpoint.values, ...targetStateValuesWithoutMessages } };
   }
 }