github/privateGPT
1
0
Fork 1
mirror of https://github.com/imartinez/privateGPT.git synced 2026-07-18 05:02:30 +00:00
Code Issues Packages Projects Releases Wiki Activity

...

Browse Source
This commit is contained in:
Javier Martinez
2026-07-13 15:43:40 +02:00
parent e2d9f842b3
commit f238979b57
9 changed files with 106 additions and 93 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
private_gpt/arq/enqueue.py
View File

@@ -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,

1
private_gpt/arq/iteration_state.py
View File

@@ -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]

28
private_gpt/arq/tasks/chat/resume.py
View File

@@ -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,
)

20
private_gpt/arq/tasks/chat/start.py
View File

@@ -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,
)

66
private_gpt/components/engines/chat/execution/chat_step.py Normal file
View File

@@ -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,
)

5
private_gpt/components/engines/chat/schedulers/iteration_scheduler.py
View File

@@ -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,

12
private_gpt/components/streaming/stream/stream_manager.py
View File

@@ -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)

1
private_gpt/server/chat_async/chat_async_service.py
View File

@@ -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,

60
tests/components/streaming/test_chat_scheduler.py
View File

@@ -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