diff --git a/private_gpt/components/engines/chat/resumable_runner.py b/private_gpt/components/engines/chat/resumable_runner.py index 38f0b84a..11bee659 100644 --- a/private_gpt/components/engines/chat/resumable_runner.py +++ b/private_gpt/components/engines/chat/resumable_runner.py @@ -34,6 +34,7 @@ if TYPE_CHECKING: from private_gpt.components.engines.chat.async_chat_engine import AsyncChatEngine from private_gpt.components.engines.chat.models.chat_state import ChatState + from private_gpt.components.tools.remote_execution import ToolExecutionResponse from private_gpt.events.models import Event @@ -135,7 +136,8 @@ class ResumableChatRunner: return channel = BrokerEventChannel(self._events, execution_id) try: - responses = list((await self._state.get_results(execution_id)).values()) + results = await self._state.get_results(execution_id) + responses = self._ordered_results(saved, results) request_data = dict(saved.request_data) request_data["messages"] = [ *list(request_data.get("messages", [])), @@ -250,6 +252,17 @@ class ResumableChatRunner: await self._state.cleanup(execution_id) await self._events.finish(execution_id) + @staticmethod + def _ordered_results( + checkpoint: ChatCheckpoint, + results: dict[str, ToolExecutionResponse], + ) -> list[ToolExecutionResponse]: + return [ + results[tool_id] + for tool_id in checkpoint.checkpoint_payload.pending_async_tools + if tool_id in results + ] + @staticmethod def _request(request_data: dict[str, Any]) -> ResolvedChatRequest: request = ResolvedChatRequest.model_validate(request_data) diff --git a/tests/engines/test_resumable_runtime.py b/tests/engines/test_resumable_runtime.py index 38421097..e9b95e37 100644 --- a/tests/engines/test_resumable_runtime.py +++ b/tests/engines/test_resumable_runtime.py @@ -264,6 +264,48 @@ def test_runner_restores_additional_kwargs_after_serialization() -> None: assert isinstance(additional_kwargs["tool_calls"][0], ToolSelection) +def test_runner_orders_results_by_pending_tool_order() -> None: + from private_gpt.components.engines.chat.resumable_runner import ResumableChatRunner + + results = { + "tool-2": ToolExecutionResponse( + tool_name="second", + tool_id="tool-2", + result_content=[TextBlock(text="second")], + tool_message={ + "role": "tool", + "content": "second", + "additional_kwargs": {"tool_call_id": "tool-2"}, + }, + ), + "tool-1": ToolExecutionResponse( + tool_name="first", + tool_id="tool-1", + result_content=[TextBlock(text="first")], + tool_message={ + "role": "tool", + "content": "first", + "additional_kwargs": {"tool_call_id": "tool-1"}, + }, + ), + } + checkpoint = ChatCheckpoint( + correlation_id="execution-order", + request_data={}, + stream_type="chat_completion", + metadata={}, + iteration=0, + checkpoint="tools", + checkpoint_payload=IterationCheckpointPayload( + pending_async_tools={"tool-1": "job-1", "tool-2": "job-2"} + ), + ) + + ordered = ResumableChatRunner._ordered_results(checkpoint, results) + + assert [response.tool_id for response in ordered] == ["tool-1", "tool-2"] + + def test_resolved_chat_request_is_json_roundtrip_serializable() -> None: import json