feat: Prune endpoint

libs/langgraph/langgraph/pregel/remote.py

@@ -631,26 +631,6 @@ class RemoteGraph(PregelProtocol):
         )
         return self._get_config(response["checkpoint"])
 
-    def _prepare_run_input(
-        self,
-        input: dict[str, Any] | Any,
-        config: RunnableConfig | None,
-    ) -> tuple[RunnableConfig, dict[str, Any] | Any, CommandSDK | None, str | None]:
-        """Prepare input for run calls.
-
-        Returns:
-            Tuple of (sanitized_config, input, command, thread_id)
-        """
-        merged_config = merge_configs(self.config, config)
-        sanitized_config = self._sanitize_config(merged_config)
-        if isinstance(input, Command):
-            command: CommandSDK | None = cast(CommandSDK, asdict(input))
-            input = None
-        else:
-            command = None
-        thread_id = sanitized_config.get("configurable", {}).pop("thread_id", None)
-        return sanitized_config, input, command, thread_id
-
     def _get_stream_modes(
         self,
         stream_mode: StreamMode | list[StreamMode] | None,
@@ -735,12 +715,17 @@ class RemoteGraph(PregelProtocol):
             The output of the graph.
         """
         sync_client = self._validate_sync_client()
-        sanitized_config, input, command, thread_id = self._prepare_run_input(
-            input, config
-        )
+        merged_config = merge_configs(self.config, config)
+        sanitized_config = self._sanitize_config(merged_config)
         stream_modes, requested, req_single, stream = self._get_stream_modes(
             stream_mode, config
         )
+        if isinstance(input, Command):
+            command: CommandSDK | None = cast(CommandSDK, asdict(input))
+            input = None
+        else:
+            command = None
+        thread_id = sanitized_config.get("configurable", {}).pop("thread_id", None)
 
         for chunk in sync_client.runs.stream(
             thread_id=thread_id,
@@ -840,12 +825,17 @@ class RemoteGraph(PregelProtocol):
             The output of the graph.
         """
         client = self._validate_client()
-        sanitized_config, input, command, thread_id = self._prepare_run_input(
-            input, config
-        )
+        merged_config = merge_configs(self.config, config)
+        sanitized_config = self._sanitize_config(merged_config)
         stream_modes, requested, req_single, stream = self._get_stream_modes(
             stream_mode, config
         )
+        if isinstance(input, Command):
+            command: CommandSDK | None = cast(CommandSDK, asdict(input))
+            input = None
+        else:
+            command = None
+        thread_id = sanitized_config.get("configurable", {}).pop("thread_id", None)
 
         async for chunk in client.runs.stream(
             thread_id=thread_id,
@@ -947,31 +937,27 @@ class RemoteGraph(PregelProtocol):
             interrupt_before: Interrupt the graph before these nodes.
             interrupt_after: Interrupt the graph after these nodes.
             headers: Additional headers to pass to the request.
-            **kwargs: Additional params to pass to client.runs.wait.
+            **kwargs: Additional params to pass to RemoteGraph.stream.
 
         Returns:
             The output of the graph.
         """
-        sync_client = self._validate_sync_client()
-        sanitized_config, input, command, thread_id = self._prepare_run_input(
-            input, config
-        )
-
-        return sync_client.runs.wait(  # type: ignore
-            thread_id=thread_id,
-            assistant_id=self.assistant_id,
-            input=input,
-            command=command,
-            config=sanitized_config,
+        for chunk in self.stream(
+            input,
+            config=config,
             interrupt_before=interrupt_before,
             interrupt_after=interrupt_after,
-            if_not_exists="create",
-            headers=(
-                _merge_tracing_headers(headers) if self.distributed_tracing else headers
-            ),
+            headers=headers,
+            stream_mode="values",
             params=params,
             **kwargs,
-        )
+        ):
+            pass
+        try:
+            return chunk
+        except UnboundLocalError:
+            logger.warning("No events received from remote graph")
+            return None
 
     async def ainvoke(
         self,
@@ -992,31 +978,27 @@ class RemoteGraph(PregelProtocol):
             interrupt_before: Interrupt the graph before these nodes.
             interrupt_after: Interrupt the graph after these nodes.
             headers: Additional headers to pass to the request.
-            **kwargs: Additional params to pass to client.runs.wait.
+            **kwargs: Additional params to pass to RemoteGraph.astream.
 
         Returns:
             The output of the graph.
         """
-        client = self._validate_client()
-        sanitized_config, input, command, thread_id = self._prepare_run_input(
-            input, config
-        )
-
-        return await client.runs.wait(
-            thread_id=thread_id,
-            assistant_id=self.assistant_id,
-            input=input,
-            command=command,
-            config=sanitized_config,
+        async for chunk in self.astream(
+            input,
+            config=config,
             interrupt_before=interrupt_before,
             interrupt_after=interrupt_after,
-            if_not_exists="create",
-            headers=(
-                _merge_tracing_headers(headers) if self.distributed_tracing else headers
-            ),
+            headers=headers,
+            stream_mode="values",
             params=params,
             **kwargs,
-        )
+        ):
+            pass
+        try:
+            return chunk
+        except UnboundLocalError:
+            logger.warning("No events received from remote graph")
+            return None
 
 
 def _merge_tracing_headers(headers: dict[str, str] | None) -> dict[str, str] | None:

libs/langgraph/tests/test_remote_graph.py

@@ -818,9 +818,13 @@ async def test_astream():
 def test_invoke():
     # set up test
     mock_sync_client = MagicMock()
-    mock_sync_client.runs.wait.return_value = {
-        "messages": [{"type": "human", "content": "world"}]
-    }
+    mock_sync_client.runs.stream.return_value = [
+        StreamPart(event="values", data={"chunk": "data1"}),
+        StreamPart(event="values", data={"chunk": "data2"}),
+        StreamPart(
+            event="values", data={"messages": [{"type": "human", "content": "world"}]}
+        ),
+    ]
 
     # call method / assertions
     remote_pregel = RemoteGraph(
@@ -834,19 +838,13 @@ def test_invoke():
     )
 
     assert result == {"messages": [{"type": "human", "content": "world"}]}
-    # verify runs.wait was called with expected args
-    assert mock_sync_client.runs.wait.called
-    _, kwargs = mock_sync_client.runs.wait.call_args
-    assert kwargs.get("thread_id") == "thread_1"
-    assert kwargs.get("assistant_id") == "test_graph_id"
-    assert kwargs.get("if_not_exists") == "create"
 
 
 def test_invoke_sanitizes_thread_id():
     # Ensure that invoking with thread_id passes thread_id as a top-level arg
     # and removes it from the config body.
     mock_sync_client = MagicMock()
-    mock_sync_client.runs.wait.return_value = {}
+    mock_sync_client.runs.stream.return_value = []
     remote_pregel = RemoteGraph("test_graph_id", sync_client=mock_sync_client)
 
     config = {"configurable": {"thread_id": "thread_1"}}
@@ -854,8 +852,8 @@ def test_invoke_sanitizes_thread_id():
         {"input": {"messages": [{"type": "human", "content": "hello"}]}}, config
     )
 
-    assert mock_sync_client.runs.wait.called
-    _, kwargs = mock_sync_client.runs.wait.call_args
+    assert mock_sync_client.runs.stream.called
+    _, kwargs = mock_sync_client.runs.stream.call_args
     assert kwargs.get("thread_id") == "thread_1"
     passed_config = kwargs.get("config") or {}
     assert "configurable" in passed_config
@@ -885,10 +883,16 @@ def test_stream_sanitizes_thread_id():
 @pytest.mark.anyio
 async def test_ainvoke():
     # set up test
-    mock_async_client = AsyncMock()
-    mock_async_client.runs.wait.return_value = {
-        "messages": [{"type": "human", "content": "world"}]
-    }
+    mock_async_client = MagicMock()
+    async_iter = MagicMock()
+    async_iter.__aiter__.return_value = [
+        StreamPart(event="values", data={"chunk": "data1"}),
+        StreamPart(event="values", data={"chunk": "data2"}),
+        StreamPart(
+            event="values", data={"messages": [{"type": "human", "content": "world"}]}
+        ),
+    ]
+    mock_async_client.runs.stream.return_value = async_iter
 
     # call method / assertions
     remote_pregel = RemoteGraph(
@@ -902,12 +906,6 @@ async def test_ainvoke():
     )
 
     assert result == {"messages": [{"type": "human", "content": "world"}]}
-    # verify runs.wait was called with expected args
-    assert mock_async_client.runs.wait.called
-    _, kwargs = mock_async_client.runs.wait.call_args
-    assert kwargs.get("thread_id") == "thread_1"
-    assert kwargs.get("assistant_id") == "test_graph_id"
-    assert kwargs.get("if_not_exists") == "create"
 
 
 @pytest.mark.skip(
@@ -1243,18 +1241,12 @@ async def test_include_headers(
     async_iter.__aiter__.return_value = return_value
     astream_mock = mock_async_client.runs.stream
     astream_mock.return_value = async_iter
-    # Mock for ainvoke which uses runs.wait
-    await_mock = AsyncMock(return_value={"chunk": "data1"})
-    mock_async_client.runs.wait = await_mock
 
     mock_sync_client = MagicMock()
     sync_iter = MagicMock()
     sync_iter.__iter__.return_value = return_value
     stream_mock = mock_sync_client.runs.stream
     stream_mock.return_value = async_iter
-    # Mock for invoke which uses runs.wait
-    wait_mock = MagicMock(return_value={"chunk": "data1"})
-    mock_sync_client.runs.wait = wait_mock
 
     remote_pregel = RemoteGraph(
         "test_graph_id",
@@ -1287,12 +1279,8 @@ async def test_include_headers(
         expected["langsmith-trace"] = AnyStr()
         expected["baggage"] = AnyStr("langsmith-metadata=")
 
-    if stream:
-        assert astream_mock.call_args.kwargs["headers"] == expected
-    else:
-        assert await_mock.call_args.kwargs["headers"] == expected
+    assert astream_mock.call_args.kwargs["headers"] == expected
     stream_mock.assert_not_called()
-    wait_mock.assert_not_called()
 
     with ls.tracing_context(enabled=True, client=MagicMock()):
         with ls.trace("foo"):
@@ -1310,7 +1298,4 @@ async def test_include_headers(
                     config,
                     headers=headers,
                 )
-    if stream:
-        assert stream_mock.call_args.kwargs["headers"] == expected
-    else:
-        assert wait_mock.call_args.kwargs["headers"] == expected
+    assert stream_mock.call_args.kwargs["headers"] == expected

libs/sdk-py/langgraph_sdk/client.py

@@ -1523,6 +1523,48 @@ class ThreadsClient:
         """
         await self.http.delete(f"/threads/{thread_id}", headers=headers, params=params)
 
+    async def prune(
+        self,
+        thread_ids: Sequence[str],
+        *,
+        strategy: Literal["delete", "keep_latest"] = "delete",
+        headers: Mapping[str, str] | None = None,
+        params: QueryParamTypes | None = None,
+    ) -> dict[str, int]:
+        """Prune threads by ID.
+
+        Args:
+            thread_ids: List of thread IDs to prune.
+            strategy: Prune strategy. Defaults to "delete".
+                - "delete": Remove threads entirely.
+                - "keep_latest": Prune old checkpoints but keep threads and their
+                  latest state. Requires `FF_USE_CORE_API=true` on the server.
+            headers: Optional custom headers to include with the request.
+            params: Optional query parameters to include with the request.
+
+        Returns:
+            Dictionary with `pruned_count` key indicating how many threads were pruned.
+
+        ???+ example "Example Usage"
+
+            ```python
+            client = get_client(url="http://localhost:2024")
+            result = await client.threads.prune(
+                thread_ids=["thread-1", "thread-2"],
+                strategy="delete"
+            )
+            print(result)  # {"pruned_count": 2}
+            ```
+
+        """
+        payload: dict[str, Any] = {
+            "thread_ids": list(thread_ids),
+            "strategy": strategy,
+        }
+        return await self.http.post(
+            "/threads/prune", json=payload, headers=headers, params=params
+        )
+
     async def search(
         self,
         *,
@@ -4852,6 +4894,48 @@ class SyncThreadsClient:
         """
         self.http.delete(f"/threads/{thread_id}", headers=headers, params=params)
 
+    def prune(
+        self,
+        thread_ids: Sequence[str],
+        *,
+        strategy: Literal["delete", "keep_latest"] = "delete",
+        headers: Mapping[str, str] | None = None,
+        params: QueryParamTypes | None = None,
+    ) -> dict[str, int]:
+        """Prune threads by ID.
+
+        Args:
+            thread_ids: List of thread IDs to prune.
+            strategy: Prune strategy. Defaults to "delete".
+                - "delete": Remove threads entirely.
+                - "keep_latest": Prune old checkpoints but keep threads and their
+                  latest state. Requires `FF_USE_CORE_API=true` on the server.
+            headers: Optional custom headers to include with the request.
+            params: Optional query parameters to include with the request.
+
+        Returns:
+            Dictionary with `pruned_count` key indicating how many threads were pruned.
+
+        ???+ example "Example Usage"
+
+            ```python
+            client = get_sync_client(url="http://localhost:2024")
+            result = client.threads.prune(
+                thread_ids=["thread-1", "thread-2"],
+                strategy="delete"
+            )
+            print(result)  # {"pruned_count": 2}
+            ```
+
+        """
+        payload: dict[str, Any] = {
+            "thread_ids": list(thread_ids),
+            "strategy": strategy,
+        }
+        return self.http.post(
+            "/threads/prune", json=payload, headers=headers, params=params
+        )
+
     def search(
         self,
         *,