From f238979b57ff2d8590b374094b12414eb42531db Mon Sep 17 00:00:00 2001 From: Javier Martinez Date: Mon, 13 Jul 2026 15:43:40 +0200 Subject: [PATCH] ... --- private_gpt/arq/enqueue.py | 6 +- private_gpt/arq/iteration_state.py | 1 - private_gpt/arq/tasks/chat/resume.py | 28 ++++---- private_gpt/arq/tasks/chat/start.py | 20 +++--- .../engines/chat/execution/chat_step.py | 66 +++++++++++++++++++ .../chat/schedulers/iteration_scheduler.py | 5 +- .../streaming/stream/stream_manager.py | 12 ++++ .../server/chat_async/chat_async_service.py | 1 - .../streaming/test_chat_scheduler.py | 60 +---------------- 9 files changed, 106 insertions(+), 93 deletions(-) create mode 100644 private_gpt/components/engines/chat/execution/chat_step.py diff --git a/private_gpt/arq/enqueue.py b/private_gpt/arq/enqueue.py index f8ddcc5c..d896658f 100644 --- a/private_gpt/arq/enqueue.py +++ b/private_gpt/arq/enqueue.py @@ -15,7 +15,7 @@ from private_gpt.settings.settings import settings as _settings async def enqueue_start_chat_job( *, - body: dict[str, Any], + request_data: dict[str, Any], correlation_id: str, stream_type: str, metadata: dict[str, Any], @@ -26,7 +26,7 @@ async def enqueue_start_chat_job( try: await redis.enqueue_job( START_CHAT_TASK_NAME, - body, + request_data, correlation_id, stream_type, metadata, @@ -40,7 +40,6 @@ async def enqueue_start_chat_job( async def enqueue_resume_iteration_job( *, - body: dict[str, Any], correlation_id: str, stream_type: str, metadata: dict[str, Any], @@ -56,7 +55,6 @@ async def enqueue_resume_iteration_job( try: await redis.enqueue_job( RESUME_ITERATION_TASK_NAME, - body, correlation_id, stream_type, metadata, diff --git a/private_gpt/arq/iteration_state.py b/private_gpt/arq/iteration_state.py index c1c7df14..c94c371e 100644 --- a/private_gpt/arq/iteration_state.py +++ b/private_gpt/arq/iteration_state.py @@ -20,7 +20,6 @@ if TYPE_CHECKING: class IterationContext(BaseModel): correlation_id: str - body: dict[str, Any] request_data: dict[str, Any] stream_type: str metadata: dict[str, Any] diff --git a/private_gpt/arq/tasks/chat/resume.py b/private_gpt/arq/tasks/chat/resume.py index fd077749..a74f8480 100644 --- a/private_gpt/arq/tasks/chat/resume.py +++ b/private_gpt/arq/tasks/chat/resume.py @@ -5,7 +5,9 @@ from private_gpt.arq.event_channel import RedisEventChannel from private_gpt.arq.iteration_state import IterationStateService from private_gpt.arq.settings import RESUME_ITERATION_TASK_NAME, TOOL_RESUME_TASK_NAME from private_gpt.arq.tasks import arq_task -from private_gpt.components.chat.models.chat_config_models import ResolvedChatRequest +from private_gpt.components.engines.chat.execution.chat_step import ( + execute_chat_resume_step, +) from private_gpt.components.engines.chat.models.chat_state import ChatStatus from private_gpt.components.engines.chat.models.execution_hooks import ( ExecutionHooks, @@ -64,7 +66,6 @@ async def tool_resume_job( request_data["messages"] = messages await enqueue_resume_iteration_job( - body=saved.body, correlation_id=correlation_id, stream_type=saved.stream_type, metadata=saved.metadata, @@ -83,7 +84,6 @@ async def tool_resume_job( @arq_task(name=RESUME_ITERATION_TASK_NAME) async def resume_iteration_job( ctx: dict[Any, Any], - body: dict[str, Any], correlation_id: str, stream_type: str, metadata: dict[str, Any], @@ -106,22 +106,16 @@ async def resume_iteration_job( if saved is None: await iteration_state_service.append_done(correlation_id) return - request = ResolvedChatRequest.model_validate(request_data) - engine = chat_service.build_engine() - state = await engine.resume( - pause_type, - request, + state = await execute_chat_resume_step( + chat_service=chat_service, + request_data=request_data, + pause_type=pause_type, + tool_results=tool_results, iteration=iteration, next_block_count=next_block_count, + iteration_state_service=iteration_state_service, + correlation_id=correlation_id, hooks=_ARQ_HOOKS, - checkpoint_payload=saved.checkpoint_payload.model_copy( - update={ - "tool_responses": [ - ToolExecutionResponse.model_validate(item) - for item in tool_results - ] - } - ), channel=channel, ) except Exception as exc: @@ -136,7 +130,7 @@ async def resume_iteration_job( scheduler = IterationScheduler( iteration_state_service=iteration_state_service, correlation_id=correlation_id, - body=body, + request_data=request_data, stream_type=stream_type, metadata=metadata, ) diff --git a/private_gpt/arq/tasks/chat/start.py b/private_gpt/arq/tasks/chat/start.py index 953a66d8..17afffb0 100644 --- a/private_gpt/arq/tasks/chat/start.py +++ b/private_gpt/arq/tasks/chat/start.py @@ -4,6 +4,9 @@ from private_gpt.arq.event_channel import RedisEventChannel from private_gpt.arq.iteration_state import IterationStateService from private_gpt.arq.settings import START_CHAT_TASK_NAME from private_gpt.arq.tasks import arq_task +from private_gpt.components.engines.chat.execution.chat_step import ( + execute_chat_start_step, +) from private_gpt.components.engines.chat.models.chat_state import ChatStatus from private_gpt.components.engines.chat.models.execution_hooks import ( ExecutionHooks, @@ -14,8 +17,6 @@ from private_gpt.components.engines.chat.schedulers.iteration_scheduler import ( ) from private_gpt.di import get_global_injector from private_gpt.events.event_serializer import StreamingEventHandler -from private_gpt.server.chat.chat_models import ChatBody -from private_gpt.server.chat.chat_request_mapper import ChatRequestMapper from private_gpt.server.chat.chat_service import ChatService _ARQ_HOOKS = ExecutionHooks( @@ -30,7 +31,7 @@ _ARQ_HOOKS = ExecutionHooks( @arq_task(name=START_CHAT_TASK_NAME) async def start_chat_job( ctx: dict[Any, Any], - body: dict[str, Any], + request_data: dict[str, Any], correlation_id: str, stream_type: str, metadata: dict[str, Any], @@ -38,17 +39,18 @@ async def start_chat_job( injector = ctx.get("injector") or get_global_injector( allow_to_generate_new_injectors=True ) - chat_request_mapper = injector.get(ChatRequestMapper) chat_service = injector.get(ChatService) iteration_state_service = injector.get(IterationStateService) event_handler = StreamingEventHandler() channel = RedisEventChannel(iteration_state_service, correlation_id, event_handler) try: - chat_body = ChatBody.model_validate(body) - request = await chat_request_mapper.create_request_from_body(chat_body) - engine = chat_service.build_engine() - state = await engine.execute(request, hooks=_ARQ_HOOKS, channel=channel) + state = await execute_chat_start_step( + chat_service=chat_service, + request_data=request_data, + hooks=_ARQ_HOOKS, + channel=channel, + ) except Exception as exc: await iteration_state_service.append_event( correlation_id, @@ -61,7 +63,7 @@ async def start_chat_job( scheduler = IterationScheduler( iteration_state_service=iteration_state_service, correlation_id=correlation_id, - body=body, + request_data=request_data, stream_type=stream_type, metadata=metadata, ) diff --git a/private_gpt/components/engines/chat/execution/chat_step.py b/private_gpt/components/engines/chat/execution/chat_step.py new file mode 100644 index 00000000..73f62bbc --- /dev/null +++ b/private_gpt/components/engines/chat/execution/chat_step.py @@ -0,0 +1,66 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +from private_gpt.components.chat.models.chat_config_models import ResolvedChatRequest +from private_gpt.components.tools.remote_execution import ToolExecutionResponse + +if TYPE_CHECKING: + from private_gpt.arq.iteration_state import IterationStateService + from private_gpt.components.engines.chat.async_chat_engine import ( + EventChannel, + IterationCheckpointPayload, + ) + from private_gpt.components.engines.chat.models.chat_state import ChatState + from private_gpt.components.engines.chat.models.execution_hooks import ( + ExecutionHooks, + ) + from private_gpt.server.chat.chat_service import ChatService + + +async def execute_chat_start_step( + *, + chat_service: ChatService, + request_data: dict[str, object], + hooks: ExecutionHooks | None, + channel: EventChannel, +) -> ChatState: + request = ResolvedChatRequest.model_validate(request_data) + engine = chat_service.build_engine() + return await engine.execute(request, hooks=hooks, channel=channel) + + +async def execute_chat_resume_step( + *, + chat_service: ChatService, + request_data: dict[str, object], + pause_type: str, + tool_results: list[dict[str, object]], + iteration: int, + next_block_count: int, + iteration_state_service: IterationStateService, + correlation_id: str, + hooks: ExecutionHooks | None, + channel: EventChannel, +) -> ChatState: + request = ResolvedChatRequest.model_validate(request_data) + engine = chat_service.build_engine() + saved = await iteration_state_service.load(correlation_id) + checkpoint_payload: IterationCheckpointPayload | None = None + if saved is not None: + checkpoint_payload = saved.checkpoint_payload.model_copy( + update={ + "tool_responses": [ + ToolExecutionResponse.model_validate(item) for item in tool_results + ] + } + ) + return await engine.resume( + pause_type, + request, + iteration=iteration, + next_block_count=next_block_count, + hooks=hooks, + checkpoint_payload=checkpoint_payload, + channel=channel, + ) diff --git a/private_gpt/components/engines/chat/schedulers/iteration_scheduler.py b/private_gpt/components/engines/chat/schedulers/iteration_scheduler.py index de825590..4162f71c 100644 --- a/private_gpt/components/engines/chat/schedulers/iteration_scheduler.py +++ b/private_gpt/components/engines/chat/schedulers/iteration_scheduler.py @@ -17,13 +17,13 @@ class IterationScheduler: *, iteration_state_service: IterationStateService, correlation_id: str, - body: dict[str, Any], + request_data: dict[str, Any], stream_type: str, metadata: dict[str, Any], ) -> None: self._iteration_state_service = iteration_state_service self._correlation_id = correlation_id - self._body = body + self._request_data = request_data self._stream_type = stream_type self._metadata = metadata self._was_waiting = False @@ -38,7 +38,6 @@ class IterationScheduler: await self._iteration_state_service.save( IterationContext( correlation_id=self._correlation_id, - body=self._body, request_data=state.input.request.model_dump(mode="json"), stream_type=self._stream_type, metadata=self._metadata, diff --git a/private_gpt/components/streaming/stream/stream_manager.py b/private_gpt/components/streaming/stream/stream_manager.py index 7b128b45..67bba093 100644 --- a/private_gpt/components/streaming/stream/stream_manager.py +++ b/private_gpt/components/streaming/stream/stream_manager.py @@ -59,6 +59,18 @@ class StreamManager: return correlation_id + async def create_stream( + self, + stream_type: str, + correlation_id: str | None = None, + metadata: dict[str, Any] | None = None, + ) -> str: + return await self.stream_service.create_stream( + stream_type=stream_type, + correlation_id=correlation_id, + metadata=metadata, + ) + async def cancel_stream(self, correlation_id: str) -> bool: """Cancel a stream.""" return await self.processor.cancel_stream_processing(correlation_id) diff --git a/private_gpt/server/chat_async/chat_async_service.py b/private_gpt/server/chat_async/chat_async_service.py index 912d1d74..5d517291 100644 --- a/private_gpt/server/chat_async/chat_async_service.py +++ b/private_gpt/server/chat_async/chat_async_service.py @@ -39,7 +39,6 @@ class ChatAsyncService: stream_type="chat_completion", event_generator=event_generator, correlation_id=message_id, - # For debugging purposes, we can include metadata about the request metadata={ "message_count": len(request.messages), "thinking_enabled": request.thinking.enabled, diff --git a/tests/components/streaming/test_chat_scheduler.py b/tests/components/streaming/test_chat_scheduler.py index e2280061..9941a78d 100644 --- a/tests/components/streaming/test_chat_scheduler.py +++ b/tests/components/streaming/test_chat_scheduler.py @@ -78,66 +78,10 @@ def test_chat_scheduler_factory_selects_arq_mode(injector: Injector) -> None: def test_chat_scheduler_factory_raises_on_unknown_mode(injector: Injector) -> None: - with pytest.raises(ValueError, match="Unknown"): + with pytest.raises(ValueError, match=r"Unknown scheduler\.chat\.mode"): ChatSchedulerFactory( settings=SimpleNamespace( - scheduler=SimpleNamespace(chat=SimpleNamespace(mode="unknown")) + scheduler=SimpleNamespace(chat=SimpleNamespace(mode="missing")) ), injector=injector, ) - - -@pytest.mark.anyio -async def test_arq_chat_scheduler_cancel_returns_false_when_abort_fails( - monkeypatch: pytest.MonkeyPatch, -) -> None: - abort_chat_job = AsyncMock(return_value=False) - monkeypatch.setattr( - "private_gpt.components.streaming.tasks.chat_scheduler.abort_chat_job", - abort_chat_job, - ) - - scheduler = ArqChatScheduler() - cancelled = await scheduler.cancel("msg-arq-not-found") - - assert cancelled is False - abort_chat_job.assert_awaited_once_with(correlation_id="msg-arq-not-found") - - -@pytest.mark.anyio -async def test_local_chat_scheduler_cancel_marks_task_as_cancelled() -> None: - async def _work() -> None: - await asyncio.sleep(100) - - task = asyncio.create_task(_work(), name="chat_msg-local-mark") - - scheduler = LocalChatScheduler() - result = await scheduler.cancel("msg-local-mark") - assert result is True - - with pytest.raises(asyncio.CancelledError): - await task - - assert task.cancelled() - - -@pytest.mark.anyio -async def test_local_chat_scheduler_cancel_only_targets_the_matching_task() -> None: - async def _work() -> None: - await asyncio.sleep(100) - - task_target = asyncio.create_task(_work(), name="chat_target-cid") - task_other = asyncio.create_task(_work(), name="chat_other-cid") - - scheduler = LocalChatScheduler() - result = await scheduler.cancel("target-cid") - assert result is True - - with pytest.raises(asyncio.CancelledError): - await task_target - assert task_target.cancelled() - - assert not task_other.done() - task_other.cancel() - with pytest.raises(asyncio.CancelledError): - await task_other