diff --git a/Makefile b/Makefile index d087eab6..50d80dac 100644 --- a/Makefile +++ b/Makefile @@ -1,6 +1,7 @@ # Any args passed to the make script, use with $(call args, default_value) args = `arg="$(filter-out $@,$(MAKECMDGOALS))" && echo $${arg:-${1}}` AUTO_DISCOVER_ARGS ?= +WORKER_ARGS ?= PROD_DIST_DIR ?= .dist PROD_PYTHON ?= .venv/bin/python PROD_BINARY ?= .venv/bin/private-gpt @@ -11,7 +12,7 @@ TEST_LOCAL_DATA_DIR ?= $(TEST_PGPT_HOME)/local_data/tests WIPE_PGPT_HOME := $(if $(PGPT_HOME),$(PGPT_HOME),$(HOME)/.local/share/private-gpt) WIPE_LOCAL_DATA_DIR := $(WIPE_PGPT_HOME)/local_data -.PHONY: test test-coverage black ruff format mypy check auto-discover-models update-openapi-spec run dev-windows dev prod-run api-docs docs ingest wipe celery flower +.PHONY: test test-coverage black ruff format mypy check auto-discover-models update-openapi-spec run dev-windows dev prod-run api-docs docs ingest wipe celery flower celery-worker arq-worker chat-worker tools-worker ######################################################################################################################## # Quality checks @@ -103,7 +104,31 @@ wipe: ######################################################################################################################## celery: - PGPT_WORKER_MODE=worker uv run private-gpt worker + PGPT_WORKER_MODE=celery ./scripts/worker_entrypoint $(WORKER_ARGS) flower: - PGPT_WORKER_MODE=flower uv run private-gpt worker + PGPT_WORKER_MODE=flower ./scripts/worker_entrypoint $(WORKER_ARGS) + +celery-worker: + PGPT_WORKER_MODE=celery ./scripts/worker_entrypoint $(WORKER_ARGS) + +arq-worker: + PGPT_WORKER_MODE=arq ./scripts/worker_entrypoint $(WORKER_ARGS) + +chat-worker: + PGPT_WORKER_MODE=arq \ + PGPT_WORKER_APP_MODULE=private_gpt \ + PGPT_ARQ_QUEUE=chat \ + PGPT_ARQ_TASK_PACKAGES=private_gpt.arq.tasks.chat \ + PGPT_STATEFUL_WORKER_TYPE=chat \ + PGPT_WORKER_WARM_PROFILE=chat \ + ./scripts/worker_entrypoint $(WORKER_ARGS) + +tools-worker: + PGPT_WORKER_MODE=celery \ + PGPT_WORKER_APP_MODULE=private_gpt \ + PGPT_CELERY_QUEUES=tools \ + PGPT_CELERY_TASK_PACKAGES=private_gpt.celery.tasks.tools \ + PGPT_STATEFUL_WORKER_TYPE=tools \ + PGPT_WORKER_WARM_PROFILE=tools \ + ./scripts/worker_entrypoint $(WORKER_ARGS) diff --git a/fern/docs/pages/configuration/cli.mdx b/fern/docs/pages/configuration/cli.mdx index 038abc02..9c0394a7 100644 --- a/fern/docs/pages/configuration/cli.mdx +++ b/fern/docs/pages/configuration/cli.mdx @@ -74,12 +74,12 @@ Starts the Celery worker stack. All behaviour is controlled through environment private-gpt worker ``` -The worker process manages up to three sub-processes depending on `PGPT_WORKER_MODE`: +`PGPT_WORKER_MODE` is required and selects one registered worker runtime: | `PGPT_WORKER_MODE` | What starts | |---|---| -| `mixed` (default) | Celery worker + Flower UI | -| `worker` | Celery worker only | +| `arq` | ARQ worker | +| `celery` | Celery worker | | `flower` | Flower UI only | **Examples:** @@ -87,11 +87,7 @@ The worker process manages up to three sub-processes depending on `PGPT_WORKER_M ```bash - # Default: worker + Flower UI - private-gpt worker - - # Worker only, no Flower - PGPT_WORKER_MODE=worker private-gpt worker + PGPT_WORKER_MODE=celery private-gpt worker # Custom concurrency and queues PGPT_CELERY_QUEUES=ingest,chat PGPT_CELERY_CONCURRENCY=4 private-gpt worker @@ -99,11 +95,7 @@ The worker process manages up to three sub-processes depending on `PGPT_WORKER_M ```powershell - # Default: worker + Flower UI - private-gpt worker - - # Worker only, no Flower - $env:PGPT_WORKER_MODE = "worker" + $env:PGPT_WORKER_MODE = "celery" private-gpt worker # Custom concurrency and queues @@ -114,11 +106,7 @@ The worker process manages up to three sub-processes depending on `PGPT_WORKER_M ```cmd - :: Default: worker + Flower UI - private-gpt worker - - :: Worker only, no Flower - set PGPT_WORKER_MODE=worker + set PGPT_WORKER_MODE=celery private-gpt worker :: Custom concurrency and queues diff --git a/private_gpt/arq/__init__.py b/private_gpt/arq/__init__.py new file mode 100644 index 00000000..876b4748 --- /dev/null +++ b/private_gpt/arq/__init__.py @@ -0,0 +1,3 @@ +from private_gpt.arq.runner import run_arq_worker + +__all__ = ["run_arq_worker"] diff --git a/private_gpt/arq/__main__.py b/private_gpt/arq/__main__.py new file mode 100644 index 00000000..4300e27e --- /dev/null +++ b/private_gpt/arq/__main__.py @@ -0,0 +1,4 @@ +from private_gpt.arq import run_arq_worker + +if __name__ == "__main__": + run_arq_worker() diff --git a/private_gpt/arq/chat/__init__.py b/private_gpt/arq/chat/__init__.py new file mode 100644 index 00000000..5f896047 --- /dev/null +++ b/private_gpt/arq/chat/__init__.py @@ -0,0 +1 @@ +"""ARQ chat worker support.""" diff --git a/private_gpt/arq/chat/iteration_state.py b/private_gpt/arq/chat/iteration_state.py new file mode 100644 index 00000000..e05c2c9c --- /dev/null +++ b/private_gpt/arq/chat/iteration_state.py @@ -0,0 +1,100 @@ +"""Redis-backed state for resumable ARQ chat execution.""" + +from __future__ import annotations + +import json +from typing import Any, cast + +import redis.asyncio as redis # type: ignore[import-untyped] +from injector import inject, singleton + +from private_gpt.arq.settings import get_redis_settings +from private_gpt.components.engines.chat.checkpoint_store import ( + ChatCheckpoint, + ChatCheckpointStore, +) +from private_gpt.components.tools.remote_execution import ToolExecutionResponse +from private_gpt.settings.settings import Settings + + +@singleton +class RedisChatCheckpointStore(ChatCheckpointStore): + """Redis-backed checkpoint storage for multi-process ARQ execution.""" + + @inject + def __init__(self, settings: Settings) -> None: + redis_settings = get_redis_settings(settings) + self._redis = redis.Redis( + host=cast(str, redis_settings.host), + port=redis_settings.port, + db=redis_settings.database, + username=redis_settings.username, + password=redis_settings.password, + decode_responses=True, + ) + self._prefix = "private_gpt:arq:iteration" + self._ttl = settings.scheduler.chat.callback_timeout_seconds + 300 + + def _ctx_key(self, execution_id: str) -> str: + return f"{self._prefix}:ctx:{execution_id}" + + def _results_key(self, execution_id: str) -> str: + return f"{self._prefix}:results:{execution_id}" + + def _resumed_key(self, execution_id: str) -> str: + return f"{self._prefix}:resumed:{execution_id}" + + async def save(self, checkpoint: ChatCheckpoint) -> None: + await self._redis.set( + self._ctx_key(checkpoint.correlation_id), + checkpoint.model_dump_json(), + ex=self._ttl, + ) + await self._redis.delete(self._resumed_key(checkpoint.correlation_id)) + + async def load(self, execution_id: str) -> ChatCheckpoint | None: + data = await self._redis.get(self._ctx_key(execution_id)) + return ChatCheckpoint.model_validate_json(data) if data else None + + async def record_result( + self, execution_id: str, tool_id: str, result: dict[str, Any] + ) -> dict[str, ToolExecutionResponse] | None: + checkpoint = await self.load(execution_id) + results_key = self._results_key(execution_id) + await cast(Any, self._redis.hset(results_key, tool_id, json.dumps(result))) + await cast(Any, self._redis.expire(results_key, self._ttl)) + if checkpoint is None: + return None + expected = set(checkpoint.checkpoint_payload.pending_async_tools) + if tool_id not in expected: + await cast(Any, self._redis.hdel(results_key, tool_id)) + return None + results = await self.get_results(execution_id) + return results if expected.issubset(results) else None + + async def get_results(self, execution_id: str) -> dict[str, ToolExecutionResponse]: + payload = cast( + dict[str, str], + await cast(Any, self._redis.hgetall(self._results_key(execution_id))), + ) + return { + key: ToolExecutionResponse.model_validate_json(value) + for key, value in payload.items() + } + + async def claim_resume(self, execution_id: str) -> bool: + return bool( + await self._redis.set( + self._resumed_key(execution_id), "1", ex=self._ttl, nx=True + ) + ) + + async def cleanup(self, execution_id: str) -> None: + await self._redis.delete( + self._ctx_key(execution_id), + self._results_key(execution_id), + self._resumed_key(execution_id), + ) + + async def close(self) -> None: + await self._redis.aclose() diff --git a/private_gpt/arq/enqueue.py b/private_gpt/arq/enqueue.py new file mode 100644 index 00000000..4b9da379 --- /dev/null +++ b/private_gpt/arq/enqueue.py @@ -0,0 +1,72 @@ +import logging +from typing import Any + +from arq import create_pool +from arq.jobs import Job + +from private_gpt.arq.settings import get_redis_settings +from private_gpt.settings.settings import settings as _settings + +logger = logging.getLogger(__name__) + + +def _log_dispatch( + *, + task_name: str, + queue_name: str, + job_id: str | None, + correlation_id: str, + defer_seconds: int | None = None, +) -> None: + logger.info( + "Dispatching ARQ task=%s queue=%s job_id=%s correlation_id=%s " + "defer_seconds=%s", + task_name, + queue_name, + job_id, + correlation_id, + defer_seconds, + ) + + +async def enqueue_job( + *, + task_name: str, + queue_name: str, + args: tuple[Any, ...] = (), + correlation_id: str, + job_id: str | None = None, + defer_seconds: int | None = None, +) -> None: + current_settings = _settings() + _log_dispatch( + task_name=task_name, + queue_name=queue_name, + job_id=job_id, + correlation_id=correlation_id, + defer_seconds=defer_seconds, + ) + redis = await create_pool(get_redis_settings(current_settings)) + try: + options: dict[str, Any] = {"_queue_name": queue_name} + if job_id is not None: + options["_job_id"] = job_id + if defer_seconds is not None: + options["_defer_by"] = defer_seconds + await redis.enqueue_job(task_name, *args, **options) + finally: + await redis.aclose() + + +async def abort_job(*, job_id: str, queue_name: str) -> bool: + current_settings = _settings() + redis = await create_pool(get_redis_settings(current_settings)) + try: + job = Job( + job_id, + redis=redis, + _queue_name=queue_name, + ) + return await job.abort(timeout=2) + finally: + await redis.aclose() diff --git a/private_gpt/arq/healthcheck.py b/private_gpt/arq/healthcheck.py new file mode 100644 index 00000000..85b6d768 --- /dev/null +++ b/private_gpt/arq/healthcheck.py @@ -0,0 +1,14 @@ +from typing import Any + +from fastapi import FastAPI + +app = FastAPI() + + +@app.get("/health") +async def health() -> dict[str, Any]: + return { + "status": "healthy", + "mode": "arq-worker", + "services": {"worker": "healthy"}, + } diff --git a/private_gpt/arq/hooks.py b/private_gpt/arq/hooks.py new file mode 100644 index 00000000..a93d8112 --- /dev/null +++ b/private_gpt/arq/hooks.py @@ -0,0 +1,11 @@ +import contextlib +from typing import Any + +from private_gpt.di import get_global_injector, set_global_injector + + +async def on_job_end(ctx: dict[Any, Any]) -> None: + del ctx + with contextlib.suppress(Exception): + injector = get_global_injector(allow_to_generate_new_injectors=True) + set_global_injector(injector) diff --git a/private_gpt/arq/lifecycle.py b/private_gpt/arq/lifecycle.py new file mode 100644 index 00000000..2a294131 --- /dev/null +++ b/private_gpt/arq/lifecycle.py @@ -0,0 +1,31 @@ +import os +from typing import Any + +import nest_asyncio # type: ignore + +from private_gpt.di import ( + clean_global_injector, + get_global_injector, + set_global_injector, +) +from private_gpt.eager_loading import warm +from private_gpt.initialize import initialize_globals, initialize_observability +from private_gpt.settings.settings import settings + + +async def startup(ctx: dict[Any, Any]) -> None: + current_settings = settings() + initialize_globals() + initialize_observability(current_settings) + nest_asyncio.apply() + injector = get_global_injector(allow_to_generate_new_injectors=True) + set_global_injector(injector) + warm_profile = os.environ.get("PGPT_WORKER_WARM_PROFILE", "").strip() + if warm_profile: + warm(injector, profile=warm_profile) + ctx["injector"] = injector + + +async def shutdown(ctx: dict[Any, Any]) -> None: + del ctx + await clean_global_injector() diff --git a/private_gpt/arq/runner.py b/private_gpt/arq/runner.py new file mode 100644 index 00000000..198753ad --- /dev/null +++ b/private_gpt/arq/runner.py @@ -0,0 +1,140 @@ +import asyncio +import contextlib +import os +import signal +import subprocess +import sys +from collections.abc import Callable + +from arq.typing import StartupShutdown +from arq.worker import Worker + +from private_gpt.arq.hooks import on_job_end +from private_gpt.arq.lifecycle import shutdown, startup +from private_gpt.arq.settings import get_queue_name, get_redis_settings +from private_gpt.arq.tasks import autodiscover_registered_tasks +from private_gpt.settings.settings import Settings, settings + + +def _default_concurrency() -> int: + cpu_count = os.cpu_count() or 1 + if cpu_count <= 1: + return 1 + return 1 << (cpu_count.bit_length() - 1) + + +def _task_packages() -> tuple[str, ...]: + configured = os.environ.get("PGPT_ARQ_TASK_PACKAGES", "") + task_packages = tuple( + package.strip() for package in configured.split(",") if package.strip() + ) + if not task_packages: + raise ValueError("PGPT_ARQ_TASK_PACKAGES must configure at least one package") + return task_packages + + +def _queue_name() -> str: + queue = os.environ.get("PGPT_ARQ_QUEUE", "").strip() + if not queue: + raise ValueError("PGPT_ARQ_QUEUE must configure a queue") + return get_queue_name(queue) + + +def run_arq_worker( + *, + settings_resolver: Callable[[], Settings] = settings, + startup_hook: StartupShutdown = startup, + shutdown_hook: StartupShutdown = shutdown, +) -> None: + app_module = os.environ.get("PGPT_WORKER_APP_MODULE", "private_gpt") + current_settings = settings_resolver() + task_packages = _task_packages() + queue_name = _queue_name() + max_jobs = int(os.environ.get("PGPT_ARQ_MAX_JOBS", str(_default_concurrency()))) + job_timeout = int(os.environ.get("PGPT_ARQ_JOB_TIMEOUT", "21600")) + api_enabled = os.environ.get("API_ENABLED", "true").lower() == "true" + api_port = os.environ.get("API_PORT", "8091") + healthcheck_app = f"{app_module}.arq.healthcheck:app" + procs: list[subprocess.Popen[bytes]] = [] + + def _cleanup(signum: int = 0, frame: object | None = None) -> None: + del frame + for proc in procs: + proc.terminate() + for proc in procs: + try: + proc.wait(timeout=10) + except subprocess.TimeoutExpired: + proc.kill() + if signum: + raise SystemExit(0) + + signal.signal(signal.SIGTERM, _cleanup) + signal.signal(signal.SIGINT, _cleanup) + + if api_enabled: + print(f"Starting arq worker healthcheck on port {api_port}") + procs.append( + subprocess.Popen( + [ + sys.executable, + "-m", + "uvicorn", + healthcheck_app, + "--host", + "0.0.0.0", + "--port", + api_port, + "--no-access-log", + "--log-level", + "critical", + ] + ) + ) + + async def _main() -> None: + worker = Worker( + functions=autodiscover_registered_tasks(*task_packages), + queue_name=queue_name, + redis_settings=get_redis_settings(current_settings), + on_startup=startup_hook, + on_shutdown=shutdown_hook, + on_job_end=on_job_end, + handle_signals=False, + allow_abort_jobs=True, + max_jobs=max_jobs, + max_tries=1, + retry_jobs=False, + keep_result=0, + job_timeout=job_timeout, + health_check_interval=30, + job_completion_wait=5, + ) + loop = asyncio.get_running_loop() + stop_event = asyncio.Event() + for sig in (signal.SIGINT, signal.SIGTERM): + loop.add_signal_handler(sig, stop_event.set) + worker_task = asyncio.create_task(worker.async_run()) + stop_task = asyncio.create_task(stop_event.wait()) + done, pending = await asyncio.wait( + {worker_task, stop_task}, return_when=asyncio.FIRST_COMPLETED + ) + if stop_task in done and not worker_task.done(): + await worker.close() + with contextlib.suppress(asyncio.CancelledError): + await worker_task + for task in pending: + task.cancel() + + print( + f"Starting arq worker queue={queue_name} task_packages={','.join(task_packages)} " + f"max_jobs={max_jobs} job_timeout={job_timeout}" + ) + try: + asyncio.run(_main()) + finally: + _cleanup() + + +if __name__ == "__main__": + run_arq_worker() diff --git a/private_gpt/arq/settings.py b/private_gpt/arq/settings.py new file mode 100644 index 00000000..24ef0b9e --- /dev/null +++ b/private_gpt/arq/settings.py @@ -0,0 +1,26 @@ +from arq.connections import RedisSettings + +from private_gpt.settings.settings import Settings + +QUEUE_PREFIX = "private_gpt:arq:queue" + + +def get_queue_name(queue: str) -> str: + return f"{QUEUE_PREFIX}:{queue}" + + +def get_redis_settings(settings: Settings) -> RedisSettings: + database = int(settings.redis.database or 0) + 8 + host = settings.redis.host + if ":" in host: + redis_host, redis_port = host.rsplit(":", 1) + else: + redis_host, redis_port = host, "6379" + + return RedisSettings( + host=redis_host, + port=int(redis_port), + database=database, + username=settings.redis.username, + password=settings.redis.password, + ) diff --git a/private_gpt/arq/task_registry.py b/private_gpt/arq/task_registry.py new file mode 100644 index 00000000..dc9b66be --- /dev/null +++ b/private_gpt/arq/task_registry.py @@ -0,0 +1,19 @@ +BASE_TASK_PACKAGES = ("private_gpt.arq.tasks",) + +_EXTERNAL_TASK_PACKAGES: list[str] = [] + + +def register_task_packages(*task_packages: str) -> None: + for task_package in task_packages: + if task_package not in _EXTERNAL_TASK_PACKAGES: + _EXTERNAL_TASK_PACKAGES.append(task_package) + + +def get_task_packages(*task_packages: str) -> tuple[str, ...]: + packages = list(task_packages or BASE_TASK_PACKAGES) + + for task_package in _EXTERNAL_TASK_PACKAGES: + if task_package not in packages: + packages.append(task_package) + + return tuple(packages) diff --git a/private_gpt/arq/tasks/__init__.py b/private_gpt/arq/tasks/__init__.py new file mode 100644 index 00000000..2d9e16c3 --- /dev/null +++ b/private_gpt/arq/tasks/__init__.py @@ -0,0 +1,54 @@ +from __future__ import annotations + +import importlib +import pkgutil +from collections.abc import Callable +from typing import TYPE_CHECKING, Any, TypeVar, cast + +from arq.worker import func + +from private_gpt.arq.task_registry import get_task_packages + +if TYPE_CHECKING: + from arq.worker import Function + +TaskFn = TypeVar("TaskFn", bound=Callable[..., Any]) + + +def arq_task(*, name: str, max_tries: int = 1) -> Callable[[TaskFn], TaskFn]: + def decorator(task_fn: TaskFn) -> TaskFn: + _fn: Any = cast(Any, task_fn) + _fn._arq_task_name = name + _fn._arq_task_max_tries = max_tries + return task_fn + + return decorator + + +def autodiscover_tasks(package_name: str = __name__) -> list[Function]: + package = importlib.import_module(package_name) + discovered: list[Function] = [] + + for module_info in pkgutil.walk_packages(package.__path__, f"{package_name}."): + module = importlib.import_module(module_info.name) + for value in vars(module).values(): + task_name = getattr(value, "_arq_task_name", None) + if task_name is None: + continue + max_tries = cast(int, getattr(value, "_arq_task_max_tries", 1)) + discovered.append( + func( + cast(Callable[..., Any], value), + name=task_name, + max_tries=max_tries, + ) + ) + + return discovered + + +def autodiscover_registered_tasks(*task_packages: str) -> list[Function]: + discovered: list[Function] = [] + for package_name in get_task_packages(*task_packages): + discovered.extend(autodiscover_tasks(package_name)) + return discovered diff --git a/private_gpt/arq/tasks/chat/__init__.py b/private_gpt/arq/tasks/chat/__init__.py new file mode 100644 index 00000000..ac4aeabb --- /dev/null +++ b/private_gpt/arq/tasks/chat/__init__.py @@ -0,0 +1,53 @@ +import asyncio +from typing import TYPE_CHECKING, Any + +from private_gpt.arq.enqueue import abort_job +from private_gpt.arq.tasks.chat.settings import get_queue_name +from private_gpt.settings.settings import settings + +if TYPE_CHECKING: + from private_gpt.arq.tasks.chat.resume import ( + enqueue_chat_timeout_job, + enqueue_resume_iteration_job, + enqueue_tool_resume_job, + ) + from private_gpt.arq.tasks.chat.start import enqueue_start_chat_job + + +async def abort_chat_job(*, correlation_id: str) -> bool: + results = await asyncio.gather( + abort_job( + job_id=f"{correlation_id}:start", + queue_name=get_queue_name(settings()), + ), + abort_job( + job_id=f"{correlation_id}:resume", + queue_name=get_queue_name(settings()), + ), + ) + return any(results) + + +def __getattr__(name: str) -> Any: + if name == "enqueue_start_chat_job": + from private_gpt.arq.tasks.chat.start import enqueue_start_chat_job + + return enqueue_start_chat_job + if name in { + "enqueue_chat_timeout_job", + "enqueue_resume_iteration_job", + "enqueue_tool_resume_job", + }: + from private_gpt.arq.tasks.chat import resume + + return getattr(resume, name) + raise AttributeError(name) + + +__all__ = [ + "abort_chat_job", + "enqueue_chat_timeout_job", + "enqueue_resume_iteration_job", + "enqueue_start_chat_job", + "enqueue_tool_resume_job", +] diff --git a/private_gpt/arq/tasks/chat/callback.py b/private_gpt/arq/tasks/chat/callback.py new file mode 100644 index 00000000..a7adb0e0 --- /dev/null +++ b/private_gpt/arq/tasks/chat/callback.py @@ -0,0 +1,27 @@ +from private_gpt.components.tools.remote_execution import ( + ToolExecutionRequest, + ToolExecutionResponse, +) + + +async def resume_chat_callback( + *, + request: ToolExecutionRequest, + response: ToolExecutionResponse, +) -> None: + correlation_id = request.context.get("correlation_id") + if not correlation_id or not request.tool_id: + return + + from private_gpt.components.engines.chat.execution_scheduler import ( + ChatExecutionSchedulerFactory, + ) + from private_gpt.di import get_global_injector + + injector = get_global_injector(allow_to_generate_new_injectors=True) + scheduler = injector.get(ChatExecutionSchedulerFactory).get() + await scheduler.callback( + execution_id=correlation_id, + tool_id=request.tool_id, + result=response.model_dump(mode="json"), + ) diff --git a/private_gpt/arq/tasks/chat/resume.py b/private_gpt/arq/tasks/chat/resume.py new file mode 100644 index 00000000..dd7a364d --- /dev/null +++ b/private_gpt/arq/tasks/chat/resume.py @@ -0,0 +1,87 @@ +from typing import Any + +from private_gpt.arq.enqueue import enqueue_job +from private_gpt.arq.tasks import arq_task +from private_gpt.arq.tasks.chat.settings import ( + CHAT_TIMEOUT_TASK_NAME, + RESUME_ITERATION_TASK_NAME, + TOOL_RESUME_TASK_NAME, + get_queue_name, +) +from private_gpt.components.engines.chat.async_chat_engine import AsyncChatEngine +from private_gpt.di import get_global_injector +from private_gpt.server.chat.chat_service import ChatService +from private_gpt.settings.settings import settings + + +async def enqueue_resume_iteration_job( + *, correlation_id: str, job_id: str | None = None +) -> None: + await enqueue_job( + task_name=RESUME_ITERATION_TASK_NAME, + queue_name=get_queue_name(settings()), + args=(correlation_id,), + job_id=job_id or f"{correlation_id}:resume", + correlation_id=correlation_id, + ) + + +async def enqueue_chat_timeout_job( + *, + correlation_id: str, + checkpoint_id: str, + delay_seconds: int, + job_id: str, +) -> None: + await enqueue_job( + task_name=CHAT_TIMEOUT_TASK_NAME, + queue_name=get_queue_name(settings()), + args=(correlation_id, checkpoint_id), + job_id=job_id, + correlation_id=correlation_id, + defer_seconds=delay_seconds, + ) + + +async def enqueue_tool_resume_job( + *, correlation_id: str, tool_id: str, result: dict[str, Any] +) -> None: + await enqueue_job( + task_name=TOOL_RESUME_TASK_NAME, + queue_name=get_queue_name(settings()), + args=(correlation_id, tool_id, result), + correlation_id=correlation_id, + ) + + +def _engine(ctx: dict[Any, Any]) -> AsyncChatEngine: + injector = ctx.get("injector") or get_global_injector( + allow_to_generate_new_injectors=True + ) + return injector.get(ChatService).build_async_engine() + + +@arq_task(name=TOOL_RESUME_TASK_NAME) +async def tool_resume_job( + ctx: dict[Any, Any], + correlation_id: str, + tool_id: str, + result: dict[str, Any], +) -> None: + await _engine(ctx).record_callback( + execution_id=correlation_id, tool_id=tool_id, result=result + ) + + +@arq_task(name=RESUME_ITERATION_TASK_NAME) +async def resume_iteration_job(ctx: dict[Any, Any], correlation_id: str) -> None: + await _engine(ctx).execute_scheduled_resume(execution_id=correlation_id) + + +@arq_task(name=CHAT_TIMEOUT_TASK_NAME) +async def timeout_chat_job( + ctx: dict[Any, Any], correlation_id: str, checkpoint_id: str +) -> None: + await _engine(ctx).execute_timeout( + execution_id=correlation_id, checkpoint_id=checkpoint_id + ) diff --git a/private_gpt/arq/tasks/chat/settings.py b/private_gpt/arq/tasks/chat/settings.py new file mode 100644 index 00000000..55b53346 --- /dev/null +++ b/private_gpt/arq/tasks/chat/settings.py @@ -0,0 +1,11 @@ +from private_gpt.arq.settings import get_queue_name as build_queue_name +from private_gpt.settings.settings import Settings + +START_CHAT_TASK_NAME = "private_gpt.chat.start" +RESUME_ITERATION_TASK_NAME = "private_gpt.chat.resume_iteration" +TOOL_RESUME_TASK_NAME = "private_gpt.tool.resume" +CHAT_TIMEOUT_TASK_NAME = "private_gpt.chat.timeout" + + +def get_queue_name(settings: Settings) -> str: + return build_queue_name(settings.scheduler.chat.celery_queue) diff --git a/private_gpt/arq/tasks/chat/start.py b/private_gpt/arq/tasks/chat/start.py new file mode 100644 index 00000000..b2140c86 --- /dev/null +++ b/private_gpt/arq/tasks/chat/start.py @@ -0,0 +1,44 @@ +from typing import Any + +from private_gpt.arq.enqueue import enqueue_job +from private_gpt.arq.tasks import arq_task +from private_gpt.arq.tasks.chat.settings import START_CHAT_TASK_NAME, get_queue_name +from private_gpt.di import get_global_injector +from private_gpt.server.chat.chat_service import ChatService +from private_gpt.settings.settings import settings + + +async def enqueue_start_chat_job( + *, + request_data: dict[str, Any], + correlation_id: str, + stream_type: str, + metadata: dict[str, Any], + job_id: str | None = None, +) -> None: + await enqueue_job( + task_name=START_CHAT_TASK_NAME, + queue_name=get_queue_name(settings()), + args=(request_data, correlation_id, stream_type, metadata), + job_id=job_id or f"{correlation_id}:start", + correlation_id=correlation_id, + ) + + +@arq_task(name=START_CHAT_TASK_NAME) +async def start_chat_job( + ctx: dict[Any, Any], + request_data: dict[str, Any], + correlation_id: str, + stream_type: str, + metadata: dict[str, Any], +) -> None: + injector = ctx.get("injector") or get_global_injector( + allow_to_generate_new_injectors=True + ) + await injector.get(ChatService).build_async_engine().execute_scheduled_start( + execution_id=correlation_id, + request_data=request_data, + stream_type=stream_type, + metadata=metadata, + ) diff --git a/private_gpt/arq/tasks/tools/__init__.py b/private_gpt/arq/tasks/tools/__init__.py new file mode 100644 index 00000000..9263a8cc --- /dev/null +++ b/private_gpt/arq/tasks/tools/__init__.py @@ -0,0 +1 @@ +"""ARQ tool tasks are discovered from their defining modules.""" diff --git a/private_gpt/celery/base.py b/private_gpt/celery/base.py index 59233852..757d4c72 100644 --- a/private_gpt/celery/base.py +++ b/private_gpt/celery/base.py @@ -1,5 +1,6 @@ import asyncio import contextlib +import os import threading from collections.abc import Callable from typing import Any @@ -221,7 +222,14 @@ class _BackgroundTask(Task): # type: ignore super().on_failure(exc, task_id, args, kwargs, einfo) -class AsyncBackgroundTask(_BackgroundTask): +class StatelessBackgroundTask(_BackgroundTask): + """Task that creates a fresh event loop and DI container per invocation. + + Used by ingestion and other short-lived tasks where per-task isolation + is preferred over warm-start performance. The event loop and DI container + are created, used, and destroyed within a single task invocation. + """ + abstract = True async def _run(self, *args: Any, **kwargs: Any) -> Any: @@ -263,12 +271,148 @@ class AsyncBackgroundTask(_BackgroundTask): raise e finally: if loop: - clean_global_injector(loop) + if not loop.is_closed(): + with contextlib.suppress(Exception): + asyncio.run_coroutine_threadsafe( + clean_global_injector(loop), loop + ).result(timeout=10) with contextlib.suppress(RuntimeError): - # Stop the loop if it's still running loop.call_soon_threadsafe(loop.stop) if thr: thr.join(timeout=5.0) - if loop: - clean_global_injector(loop) + with contextlib.suppress(RuntimeError): + loop.close() + + +class StatefulBackgroundTask(_BackgroundTask): + """Long-lived task base that keeps DI and event loop warm across invocations. + + Designed for prefork pools where each child process independently creates + its own persistent event loop and warm DI container. The warm-up is + serialized across children via a file lock in ``bootsteps.py`` to avoid + OOM from simultaneous model loading. + """ + + abstract = True + _loop: asyncio.AbstractEventLoop | None = None + _thread: threading.Thread | None = None + _lock = threading.RLock() + _warmed = False + + @classmethod + def _ensure_runtime(cls) -> asyncio.AbstractEventLoop: + with cls._lock: + if cls._loop is not None and cls._thread is not None: + if cls._thread.is_alive() and not cls._loop.is_closed(): + return cls._loop + + loop = asyncio.new_event_loop() + nest_asyncio.apply(loop) + + def run_loop() -> None: + asyncio.set_event_loop(loop) + loop.run_forever() + + thread = threading.Thread( + target=run_loop, + daemon=True, + name="StatefulWorkerLoop", + ) + thread.start() + cls._loop = loop + cls._thread = thread + cls._warmed = False + return loop + + @classmethod + def run_coroutine(cls, coro: Any) -> Any: + loop = cls._ensure_runtime() + future = asyncio.run_coroutine_threadsafe(coro, loop) + return future.result() + + @classmethod + async def _warm_async(cls) -> None: + from private_gpt.eager_loading import warm + + injector = get_global_injector(allow_to_generate_new_injectors=True) + profile = os.environ.get("PGPT_WORKER_WARM_PROFILE", "").strip() + if not profile: + raise ValueError( + "PGPT_WORKER_WARM_PROFILE is required for stateful workers" + ) + warm(injector, profile=profile) + + @classmethod + def warm_up(cls) -> None: + """Initialize runtime and warm DI. Idempotent across calls.""" + with cls._lock: + if cls._warmed: + return + + cls._ensure_runtime() + cls.run_coroutine(cls._warm_async()) + with cls._lock: + cls._warmed = True + logger.info("StatefulBackgroundTask warm-up complete") + + async def _run(self, *args: Any, **kwargs: Any) -> Any: + """Execute the task logic, handling both sync and async implementations.""" + run_method = self.run + + if run_method is None: + raise NotImplementedError("Subclass must implement 'run' method") + + get_global_injector(allow_to_generate_new_injectors=True) + + if asyncio.iscoroutinefunction(run_method): + result = await run_method(*args, **kwargs) + else: + result = run_method(*args, **kwargs) + + if asyncio.iscoroutine(result): + result = await result + + return result + + @classmethod + async def _shutdown_async(cls) -> None: + await clean_global_injector() + + @classmethod + def shutdown_runtime(cls) -> None: + loop: asyncio.AbstractEventLoop | None = None + thread: threading.Thread | None = None + with cls._lock: + loop = cls._loop + thread = cls._thread + cls._loop = None + cls._thread = None + cls._warmed = False + + if loop is None: + return + + if not loop.is_closed(): + with contextlib.suppress(Exception): + asyncio.run_coroutine_threadsafe(cls._shutdown_async(), loop).result( + timeout=10 + ) + with contextlib.suppress(RuntimeError): + loop.call_soon_threadsafe(loop.stop) + + if thread is not None: + thread.join(timeout=5.0) + + with contextlib.suppress(RuntimeError): + loop.close() + logger.info("%s runtime shut down", cls.__name__) + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + try: + self.warm_up() + return self.run_coroutine(self._run(*args, **kwargs)) + except Exception as e: + if self.rollback_fn: + self.rollback_fn(*args, **kwargs) + raise e diff --git a/private_gpt/celery/bootsteps.py b/private_gpt/celery/bootsteps.py index e930e0af..5417b37e 100644 --- a/private_gpt/celery/bootsteps.py +++ b/private_gpt/celery/bootsteps.py @@ -1,13 +1,25 @@ +import fcntl +import logging +import os from pathlib import Path from typing import Any, ClassVar from celery import bootsteps # type: ignore -from celery.signals import worker_ready, worker_shutdown +from celery.signals import ( + worker_process_init, + worker_process_shutdown, + worker_ready, + worker_shutdown, +) + +logger = logging.getLogger(__name__) # Files for health checks READINESS_FILE = Path("/tmp/celery_ready") HEARTBEAT_FILE = Path("/tmp/celery_worker_heartbeat") +STATEFUL_WARMUP_LOCK = Path("/tmp/stateful_warmup.lock") + class LivenessProbe(bootsteps.StartStopStep): # type: ignore """Liveness probe for Celery worker. @@ -34,12 +46,73 @@ class LivenessProbe(bootsteps.StartStopStep): # type: ignore HEARTBEAT_FILE.touch() +def _is_stateful() -> bool: + return bool(os.getenv("PGPT_STATEFUL_WORKER_TYPE", "").strip()) + + +def _warm_stateful_worker() -> None: + worker_type = os.getenv("PGPT_STATEFUL_WORKER_TYPE", "").strip() + if not worker_type: + return + + from private_gpt.celery.base import StatefulBackgroundTask + + logger.info( + "STATEFUL_WORKER_TYPE=%s: eagerly warming StatefulBackgroundTask", + worker_type, + ) + + STATEFUL_WARMUP_LOCK.parent.mkdir(parents=True, exist_ok=True) + lock_fd = os.open(str(STATEFUL_WARMUP_LOCK), os.O_CREAT | os.O_RDWR) + try: + fcntl.flock(lock_fd, fcntl.LOCK_EX) + StatefulBackgroundTask.warm_up() + finally: + fcntl.flock(lock_fd, fcntl.LOCK_UN) + os.close(lock_fd) + + logger.info("StatefulBackgroundTask DI warmed successfully") + + +def _shutdown_stateful_worker() -> None: + worker_type = os.getenv("PGPT_STATEFUL_WORKER_TYPE", "").strip() + if not worker_type: + return + + from private_gpt.celery.base import StatefulBackgroundTask + + StatefulBackgroundTask.shutdown_runtime() + + @worker_ready.connect def handle_worker_ready(**kwargs: dict[str, Any]) -> None: """Signal handler for worker ready event.""" + if _is_stateful(): + _warm_stateful_worker() + READINESS_FILE.touch() +@worker_process_init.connect +def handle_worker_process_init(**kwargs: dict[str, Any]) -> None: + """Warm each prefork child that will execute stateful tasks. + + The warm-up is serialized via a file lock so children load models + one at a time, avoiding OOM from simultaneous loading. + """ + if _is_stateful(): + _warm_stateful_worker() + + +@worker_process_shutdown.connect +def handle_worker_process_shutdown(**kwargs: dict[str, Any]) -> None: + """Clean up stateful worker loop resources on child shutdown.""" + if not _is_stateful(): + return + + _shutdown_stateful_worker() + + @worker_shutdown.connect def handle_worker_shutdown(**kwargs: dict[str, Any]) -> None: """Signal handler for worker shutdown event.""" diff --git a/private_gpt/celery/callback.py b/private_gpt/celery/callback.py index eb79954f..a6a4c2b4 100644 --- a/private_gpt/celery/callback.py +++ b/private_gpt/celery/callback.py @@ -18,6 +18,11 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) +def _callback_task_name(task: Any) -> str: + name: str = getattr(task, "callback_task_name", task.name) + return name + + def _publish_callback( exchange: str, routing_key: str, @@ -47,18 +52,20 @@ def run_callback( final = False if state == states.SUCCESS: + task_name = _callback_task_name(task) async_response = AsyncResponse( data=result, - type=f"pgpt.{task.name}.done", + type=f"pgpt.{task_name}.done", error=None, callback_properties=callback.properties, ) routing_key = callback_amqp.routing_key_done or async_response.type final = True elif state == custom_states.PROGRESS: + task_name = _callback_task_name(task) async_response = AsyncResponse( data=result.model_dump(), - type=f"pgpt.{task.name}.progress", + type=f"pgpt.{task_name}.progress", error=None, callback_properties=callback.properties, ) @@ -77,7 +84,7 @@ def run_callback( async_response = AsyncResponse( data=None, - type=f"pgpt.{task.name}.error", + type=f"pgpt.{_callback_task_name(task)}.error", error=error_result.dict(), callback_properties=callback.properties, ) @@ -112,7 +119,7 @@ def task_after_return( """ # If the input is not a subclass of BaseCallbackInput, then we don't need to # send a callback - if not issubclass(args[0].__class__, BaseCallbackInput): + if not args or not issubclass(args[0].__class__, BaseCallbackInput): return # If the input is a subclass of BaseCallbackInput, it needs to be the only input diff --git a/private_gpt/celery/celery.py b/private_gpt/celery/celery.py index bf163896..f5b8c32d 100644 --- a/private_gpt/celery/celery.py +++ b/private_gpt/celery/celery.py @@ -1,3 +1,4 @@ +import os from collections.abc import Callable, Sequence from typing import Any @@ -10,6 +11,15 @@ from private_gpt.initialize import initialize_globals, initialize_observability from private_gpt.settings.settings import Settings, settings +def _configured_task_packages(task_packages: Sequence[str]) -> tuple[str, ...]: + if task_packages: + return tuple(task_packages) + configured = os.environ.get("PGPT_CELERY_TASK_PACKAGES", "") + return tuple( + package.strip() for package in configured.split(",") if package.strip() + ) + + def _autodiscover_tasks( celery_app: Celery, task_packages: Sequence[str], *, force: bool = False ) -> None: @@ -52,7 +62,9 @@ def create_app( if before_setup is not None: before_setup() - resolved_task_packages = get_task_packages(*task_packages) + resolved_task_packages = get_task_packages( + *_configured_task_packages(task_packages) + ) celery_app = Celery( app_name, diff --git a/private_gpt/celery/dispatch.py b/private_gpt/celery/dispatch.py new file mode 100644 index 00000000..e25d2548 --- /dev/null +++ b/private_gpt/celery/dispatch.py @@ -0,0 +1,23 @@ +from typing import Any + + +def dispatch_task( + *, + task_name: str, + queue: str, + args: tuple[Any, ...] | list[Any] | None = None, + kwargs: dict[str, Any] | None = None, + task_id: str | None = None, +) -> Any: + from private_gpt.celery.celery import celery_app + + options: dict[str, Any] = {"queue": queue} + if task_id is not None: + options["task_id"] = task_id + + return celery_app.send_task( + task_name, + args=args, + kwargs=kwargs, + **options, + ) diff --git a/private_gpt/celery/healthcheck.py b/private_gpt/celery/healthcheck.py index 2718bd3f..fe42038d 100644 --- a/private_gpt/celery/healthcheck.py +++ b/private_gpt/celery/healthcheck.py @@ -3,7 +3,7 @@ import os import time from typing import Any -import requests +import httpx from fastapi import FastAPI, HTTPException from private_gpt.celery.bootsteps import HEARTBEAT_FILE, READINESS_FILE @@ -11,19 +11,20 @@ from private_gpt.celery.bootsteps import HEARTBEAT_FILE, READINESS_FILE HEALTHCHECK_TIMEOUT = 10 -def check_flower() -> bool: +async def check_flower() -> bool: try: flower_port = os.getenv("PGPT_FLOWER_PORT", "5555") flower_prefix = os.getenv("PGPT_FLOWER_URL_PREFIX", "") url = f"http://localhost:{flower_port}{flower_prefix}/healthcheck" - response = requests.get(url, timeout=HEALTHCHECK_TIMEOUT) - return response.status_code == 200 + async with httpx.AsyncClient(timeout=HEALTHCHECK_TIMEOUT) as client: + response = await client.get(url) + return response.status_code == 200 except Exception as e: logging.critical(f"Error checking Flower health: {e}") return False -def check_worker() -> bool: +async def check_worker() -> bool: try: if not READINESS_FILE.is_file(): return False @@ -43,22 +44,26 @@ def check_worker() -> bool: return False -def health_check() -> dict[str, Any]: - mode = os.getenv("PGPT_WORKER_MODE", "mixed") +async def health_check() -> dict[str, Any]: + mode = os.getenv("PGPT_WORKER_MODE", "").strip().lower() status: dict[str, Any] = {"status": "healthy", "mode": mode, "services": {}} - if mode in ["mixed", "flower"]: - flower_health = check_flower() + if mode == "flower": + flower_health = await check_flower() status["services"]["flower"] = "healthy" if flower_health else "unhealthy" if not flower_health: status["status"] = "unhealthy" - if mode in ["mixed", "worker"]: - worker_health = check_worker() + elif mode == "celery": + worker_health = await check_worker() status["services"]["worker"] = "healthy" if worker_health else "unhealthy" if not worker_health: status["status"] = "unhealthy" + else: + status["status"] = "unhealthy" + status["error"] = f"Unsupported worker mode: {mode or ''}" + return status @@ -67,7 +72,7 @@ app = FastAPI() @app.get("/health") async def health() -> dict[str, Any]: - status = health_check() + status = await health_check() if status["status"] == "unhealthy": raise HTTPException(status_code=503, detail=status) return status diff --git a/private_gpt/celery/task_helper.py b/private_gpt/celery/task_helper.py index 72274019..41865bc0 100644 --- a/private_gpt/celery/task_helper.py +++ b/private_gpt/celery/task_helper.py @@ -69,7 +69,11 @@ class IngestionTaskHelper: def is_ingestion_cancel_task_scheduled( celery_app: Any, collection: str, artifact: str ) -> bool: - for task in find_tasks(celery_app, task_name="delete_ingested_task"): + from private_gpt.celery.tasks.ingestion.delete_tasks import ( + DELETE_INGESTED_TASK_NAME, + ) + + for task in find_tasks(celery_app, task_name=DELETE_INGESTED_TASK_NAME): if not task.args or not isinstance( task.args[0], DeleteIngestedDocumentAsyncBody ): @@ -86,9 +90,12 @@ class IngestionTaskHelper: @staticmethod def revoke_ingestion_task(celery_app: Any, collection: str, artifact: str) -> bool: + from private_gpt.celery.tasks.ingestion.extraction_tasks import ( + VECTOR_INDEX_TASK_NAME, + ) from private_gpt.server.ingest.ingest_router import IngestAsyncBody - for task in find_tasks(celery_app, task_name="vector_index_task"): + for task in find_tasks(celery_app, task_name=VECTOR_INDEX_TASK_NAME): if not task.args or not isinstance(task.args[0], IngestAsyncBody): continue @@ -97,7 +104,6 @@ class IngestionTaskHelper: task_body.ingest_body.collection == collection and task_body.ingest_body.artifact == artifact ): - revoke_task(celery_app, task.task_id) return True @@ -105,7 +111,11 @@ class IngestionTaskHelper: @staticmethod def revoke_deletion_task(celery_app: Any, collection: str, artifact: str) -> bool: - for task in find_tasks(celery_app, task_name="delete_ingested_task"): + from private_gpt.celery.tasks.ingestion.delete_tasks import ( + DELETE_INGESTED_TASK_NAME, + ) + + for task in find_tasks(celery_app, task_name=DELETE_INGESTED_TASK_NAME): if not task.args or not isinstance( task.args[0], DeleteIngestedDocumentAsyncBody ): diff --git a/private_gpt/celery/task_registry.py b/private_gpt/celery/task_registry.py index a33b113d..1d66b024 100644 --- a/private_gpt/celery/task_registry.py +++ b/private_gpt/celery/task_registry.py @@ -1,4 +1,7 @@ -BASE_TASK_PACKAGES = ("private_gpt.celery.tasks",) +BASE_TASK_PACKAGES = ( + "private_gpt.celery.tasks.ingestion", + "private_gpt.celery.tasks.tools", +) _EXTERNAL_TASK_PACKAGES: list[str] = [] @@ -10,14 +13,10 @@ def register_task_packages(*task_packages: str) -> None: def get_task_packages(*task_packages: str) -> tuple[str, ...]: - packages = list(BASE_TASK_PACKAGES) + packages = list(task_packages or BASE_TASK_PACKAGES) for task_package in _EXTERNAL_TASK_PACKAGES: if task_package not in packages: packages.append(task_package) - for task_package in task_packages: - if task_package not in packages: - packages.append(task_package) - return tuple(packages) diff --git a/private_gpt/celery/tasks/__init__.py b/private_gpt/celery/tasks/__init__.py index 1cd1d9c8..c9c2ef67 100644 --- a/private_gpt/celery/tasks/__init__.py +++ b/private_gpt/celery/tasks/__init__.py @@ -1,4 +1 @@ -from private_gpt.celery.tasks.ingestion import __all__ as ingestion_all - -__all__ = [] -__all__.extend(ingestion_all) +__all__: list[str] = [] diff --git a/private_gpt/celery/tasks/chat/__init__.py b/private_gpt/celery/tasks/chat/__init__.py new file mode 100644 index 00000000..c9c2ef67 --- /dev/null +++ b/private_gpt/celery/tasks/chat/__init__.py @@ -0,0 +1 @@ +__all__: list[str] = [] diff --git a/private_gpt/celery/tasks/ingestion/delete_tasks.py b/private_gpt/celery/tasks/ingestion/delete_tasks.py index 899e3a3a..c7fbeed8 100644 --- a/private_gpt/celery/tasks/ingestion/delete_tasks.py +++ b/private_gpt/celery/tasks/ingestion/delete_tasks.py @@ -2,7 +2,7 @@ import logging from typing import TYPE_CHECKING from private_gpt.artifact_index.base_artifact_index import IndexNotReadyException -from private_gpt.celery.base import AsyncBackgroundTask +from private_gpt.celery.base import StatelessBackgroundTask from private_gpt.celery.celery import celery_app from private_gpt.celery.task_helper import IngestionTaskHelper from private_gpt.di import get_global_injector @@ -16,10 +16,13 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) logger.setLevel(logging.DEBUG if settings().server.debug_mode else logging.INFO) +DELETE_INGESTED_TASK_NAME = "private_gpt.ingestion.delete" +DELETE_INGESTED_CALLBACK_TASK_NAME = "delete_ingested_task" + @celery_app.task( - name="delete_ingested_task", - base=AsyncBackgroundTask, + name=DELETE_INGESTED_TASK_NAME, + base=StatelessBackgroundTask, # Retry on ValueError and IndexNotReadyException. # ValueError is thrown when the index is not initialized # and we cannot guarantee that the index will not be ready. @@ -51,3 +54,6 @@ def delete_ingested_task(body: "DeleteIngestedDocumentAsyncBody") -> None: # If the task was not revoked, we follow the normal # flow and raise the exception. raise + + +delete_ingested_task.callback_task_name = DELETE_INGESTED_CALLBACK_TASK_NAME # type: ignore[attr-defined] diff --git a/private_gpt/celery/tasks/ingestion/extraction_tasks.py b/private_gpt/celery/tasks/ingestion/extraction_tasks.py index 4a4acf61..786d7834 100644 --- a/private_gpt/celery/tasks/ingestion/extraction_tasks.py +++ b/private_gpt/celery/tasks/ingestion/extraction_tasks.py @@ -5,7 +5,7 @@ from typing import Any, TypeVar from private_gpt.artifact_index.base_artifact_index import IndexNotReadyException from private_gpt.celery import states as custom_states -from private_gpt.celery.base import AsyncBackgroundTask +from private_gpt.celery.base import StatelessBackgroundTask from private_gpt.celery.celery import celery_app from private_gpt.components.ingest.utils import get_extension, get_file_name from private_gpt.components.storage.s3_helper import S3Helper @@ -18,6 +18,9 @@ from private_gpt.settings.settings import settings logger = logging.getLogger(__name__) logger.setLevel(logging.DEBUG if settings().server.debug_mode else logging.INFO) +VECTOR_INDEX_TASK_NAME = "private_gpt.ingestion.vector_index" +VECTOR_INDEX_CALLBACK_TASK_NAME = "vector_index_task" + T = TypeVar("T") @@ -45,8 +48,8 @@ def cleanup_temporal_files(func: Callable[..., T]) -> Callable[..., T]: @celery_app.task( - name="vector_index_task", - base=AsyncBackgroundTask, + name=VECTOR_INDEX_TASK_NAME, + base=StatelessBackgroundTask, autoretry_for=AUTORETRY_EXCEPTIONS, ) @cleanup_temporal_files @@ -100,7 +103,6 @@ def vector_index_task(body: IngestAsyncBody) -> Any: ) service = get_global_injector().get(IngestService) - content = body.ingest_body.input.to_binary_content( filename=get_file_name(body.ingest_body.metadata) ) @@ -135,6 +137,9 @@ def vector_index_task(body: IngestAsyncBody) -> Any: ) +vector_index_task.callback_task_name = VECTOR_INDEX_CALLBACK_TASK_NAME # type: ignore[attr-defined] + + def ensure_to_remove_temporal_files(body: IngestAsyncBody) -> None: """Remove temporal files from S3 if the input was a URI. diff --git a/private_gpt/celery/tasks/tools/__init__.py b/private_gpt/celery/tasks/tools/__init__.py new file mode 100644 index 00000000..c83ea898 --- /dev/null +++ b/private_gpt/celery/tasks/tools/__init__.py @@ -0,0 +1,3 @@ +from private_gpt.celery.tasks.tools.tool_run_task import tool_run_task + +__all__ = ["tool_run_task"] diff --git a/private_gpt/celery/tasks/tools/tool_run_task.py b/private_gpt/celery/tasks/tools/tool_run_task.py new file mode 100644 index 00000000..36e62b10 --- /dev/null +++ b/private_gpt/celery/tasks/tools/tool_run_task.py @@ -0,0 +1,74 @@ +"""Celery task that executes a single tool call on a dedicated tools worker.""" + +from __future__ import annotations + +import logging +from typing import Any + +from private_gpt.celery.base import StatefulBackgroundTask +from private_gpt.celery.celery import celery_app +from private_gpt.components.tools.remote_execution import ( + ToolExecutionRequest, + ToolExecutionResponse, + execute_tool_request, +) +from private_gpt.components.tools.tool_scheduler import ToolSchedulerFactory +from private_gpt.di import get_global_injector +from private_gpt.events.models import TextBlock + +logger = logging.getLogger(__name__) + + +@celery_app.task( + name="private_gpt.tools.run", + base=StatefulBackgroundTask, +) +async def tool_run_task(*, request_data: dict[str, Any]) -> dict[str, Any]: + request = ToolExecutionRequest.model_validate(request_data) + try: + response = await execute_tool_request(request) + await _notify_completion(request, response) + return response.model_dump(mode="json") + except Exception as exc: + logger.exception("Tool '%s' execution failed", request.tool_name) + response = ToolExecutionResponse( + tool_name=request.tool_name, + tool_id=request.tool_id, + result_content=[TextBlock(text=str(exc))], + is_error=True, + tool_message=request_error_message(request, str(exc)), + ) + await _notify_completion(request, response) + return response.model_dump(mode="json") + + +async def _notify_completion( + request: ToolExecutionRequest, + response: ToolExecutionResponse, +) -> None: + correlation_id = request.context.get("correlation_id") + if not correlation_id or not request.tool_id: + return + try: + scheduler = get_global_injector(True).get(ToolSchedulerFactory).get() + await scheduler.complete(request, response) + except Exception: + logger.exception("Tool completion notify failed for %s", request.tool_id) + + +def request_error_message( + request: ToolExecutionRequest, + message: str, +) -> Any: + from llama_index.core.base.llms.types import ChatMessage + + return ChatMessage( + role="tool", + content=message, + additional_kwargs={ + "tool_call_id": request.tool_id, + "tool_call_name": request.tool_name, + "tool_call_args": request.tool_kwargs, + "raw_output": message, + }, + ) diff --git a/private_gpt/cli/commands/worker.py b/private_gpt/cli/commands/worker.py index 7ee454f8..6df9b07a 100644 --- a/private_gpt/cli/commands/worker.py +++ b/private_gpt/cli/commands/worker.py @@ -1,151 +1,6 @@ -import os -import signal -import subprocess -import sys - -import typer - - -def _build_flower_args(mode: str) -> list[str]: - args: list[str] = [] - url_prefix = os.environ.get("PGPT_FLOWER_URL_PREFIX", "") - port = os.environ.get("PGPT_FLOWER_PORT", "5555") - password = os.environ.get("PGPT_FLOWER_PASSWORD", "flower") - persistent = os.environ.get("PGPT_FLOWER_PERSISTENT", "true") - persistent_db = os.environ.get("PGPT_FLOWER_PERSISTENT_DB", "celery_flower.db") - max_tasks = os.environ.get("PGPT_FLOWER_MAXIMUM_TASKS", "1000") - - if url_prefix: - args.append(f"--url_prefix={url_prefix}") - if port: - args.append(f"--port={port}") - if password: - args.append(f"--basic_auth=flower:{password}") - if mode == "mixed": - args.append("--logging=none") - if persistent.lower() == "true": - args.append("--persistent=True") - if persistent_db: - args.append(f"--db={persistent_db}") - if max_tasks: - args.append(f"--max-tasks={max_tasks}") - return args - - -def _build_worker_args() -> list[str]: - args: list[str] = [] - log_level = os.environ.get("PGPT_CELERY_LOG_LEVEL", "info") - queues = os.environ.get("PGPT_CELERY_QUEUES", "") - hostname = os.environ.get("PGPT_CELERY_HOSTNAME", "default") - pool = os.environ.get("PGPT_CELERY_POOL", "prefork") - concurrency = os.environ.get("PGPT_CELERY_CONCURRENCY", "") - time_limit = os.environ.get("PGPT_CELERY_TIME_LIMIT", "") - - if log_level: - args.append(f"--loglevel={log_level}") - if queues: - args.append(f"--queues={queues}") - hostname = queues - if hostname: - args.append(f"--hostname={hostname}%h") - if pool: - args.append(f"--pool={pool}") - if concurrency: - args.append(f"--concurrency={concurrency}") - if time_limit: - args.append(f"--time-limit={time_limit}") - return args +from private_gpt.worker import run_private_gpt_worker def worker_command() -> None: - """Start a background worker process. - - Behaviour is fully controlled through environment variables. - """ - app_module = os.environ.get("PGPT_WORKER_APP_MODULE", "private_gpt") - celery_app = f"{app_module}.celery" - healthcheck_app = f"{app_module}.celery.healthcheck:app" - mode = os.environ.get("PGPT_WORKER_MODE", "mixed") - - procs: list[subprocess.Popen[bytes]] = [] - - def _cleanup(signum: int = 0, frame: object = None) -> None: - typer.echo("Shutting down services...") - for p in procs: - p.terminate() - for p in procs: - try: - p.wait(timeout=10) - except subprocess.TimeoutExpired: - p.kill() - sys.exit(0) - - signal.signal(signal.SIGTERM, _cleanup) - signal.signal(signal.SIGINT, _cleanup) - - if mode in ("flower", "mixed"): - flower_args = _build_flower_args(mode) - typer.echo( - f"Starting flower on port {os.environ.get('PGPT_FLOWER_PORT', '5555')}" - ) - procs.append( - subprocess.Popen( - [ - sys.executable, - "-m", - "celery", - "--app", - celery_app, - "flower", - *flower_args, - ] - ) - ) - - if mode in ("worker", "mixed"): - worker_args = _build_worker_args() - typer.echo(f"Starting celery worker with args: {' '.join(worker_args)}") - procs.append( - subprocess.Popen( - [ - sys.executable, - "-m", - "celery", - "--app", - celery_app, - "worker", - *worker_args, - ] - ) - ) - - if os.environ.get("API_ENABLED", "true").lower() == "true": - api_port = os.environ.get("API_PORT", "8090") - typer.echo(f"Starting healthcheck server on port {api_port}") - procs.append( - subprocess.Popen( - [ - sys.executable, - "-m", - "uvicorn", - healthcheck_app, - "--host", - "0.0.0.0", - "--port", - api_port, - "--no-access-log", - "--log-level", - "critical", - ] - ) - ) - - if not procs: - typer.echo(f"No processes started. Check PGPT_WORKER_MODE={mode!r}", err=True) - raise SystemExit(1) - - try: - for p in procs: - p.wait() - except KeyboardInterrupt: - _cleanup() + """Start the worker mode selected by ``PGPT_WORKER_MODE``.""" + run_private_gpt_worker() diff --git a/private_gpt/components/broker/broker_component.py b/private_gpt/components/broker/broker_component.py index eb3d3360..052f558f 100644 --- a/private_gpt/components/broker/broker_component.py +++ b/private_gpt/components/broker/broker_component.py @@ -54,7 +54,7 @@ class BrokerInstanceRegistry: provider = _PROVIDERS.get(mode) if provider is None: - raise ValueError(f"Unsupported task result broker mode: {mode}") + return None new_instance = provider(self._settings) self._instances[mode] = new_instance diff --git a/private_gpt/components/chat/models/chat_config_models.py b/private_gpt/components/chat/models/chat_config_models.py index ada5c857..a620baac 100644 --- a/private_gpt/components/chat/models/chat_config_models.py +++ b/private_gpt/components/chat/models/chat_config_models.py @@ -3,19 +3,25 @@ import enum import inspect import re from collections.abc import Awaitable, Callable -from typing import Any, Literal +from typing import Any, ClassVar, Literal from llama_index.core.base.llms.types import ChatMessage, MessageRole, TextBlock from llama_index.core.llms import LLM from llama_index.core.tools import BaseTool, FunctionTool from llama_index.core.tools.function_tool import AsyncCallable, _is_context_param from llama_index.core.tools.utils import create_schema_from_function -from pydantic import BaseModel, ConfigDict, Field +from pydantic import ( + BaseModel, + ConfigDict, + Field, + field_serializer, + field_validator, +) from private_gpt.chat.input_models import BlobVisibilityMode, PromptConfig from private_gpt.chat.schema_models import create_model_from_json_schema from private_gpt.components.engines.citations.types import Citation, Document -from private_gpt.components.llm.llm_helper import TokenizerFn +from private_gpt.components.llm.llm_helper import AsyncTokenizerFn, TokenizerFn from private_gpt.components.sandbox.content_bundle import ContentBundle from private_gpt.components.tools.tool_names import resolve_internal_tool_name from private_gpt.components.tools.types import ToolValidationMode @@ -31,7 +37,7 @@ class LLMInstanceConfig(BaseModel): config: LLMModelConfig = Field( description="The LLM model configuration.", ) - tokenizer: TokenizerFn | None = Field( + tokenizer: TokenizerFn | AsyncTokenizerFn | None = Field( default=None, description="The tokenizer function to use for the LLM.", ) @@ -111,6 +117,53 @@ class ToolRequirements(enum.StrEnum): SANDBOX = "sandbox" +class ToolExecutionMetadata(BaseModel): + rebuild_callable: str = Field( + description="Import path to the callable that rebuilds the server tool." + ) + rebuild_kwargs: dict[str, Any] = Field( + default_factory=dict, + description="JSON-serializable kwargs used to rebuild the server tool.", + ) + + MODEL_TAG_KEY: ClassVar[str] = "__pgpt_model__" + + @field_serializer("rebuild_kwargs", when_used="json") + def _serialize_kwargs(self, kwargs: dict[str, Any]) -> dict[str, Any]: + """Tag BaseModel values so they survive the JSON roundtrip.""" + result: dict[str, Any] = {} + for key, value in kwargs.items(): + if isinstance(value, BaseModel): + result[key] = { + self.MODEL_TAG_KEY: f"{type(value).__module__}:{type(value).__qualname__}", + "data": value.model_dump(mode="json"), + } + else: + result[key] = value + return result + + @field_validator("rebuild_kwargs", mode="before") + @classmethod + def _deserialize_kwargs(cls, kwargs: Any) -> Any: + """Untag BaseModel values after JSON roundtrip.""" + if not isinstance(kwargs, dict): + return kwargs + import importlib + + result: dict[str, Any] = {} + for key, value in kwargs.items(): + if isinstance(value, dict) and cls.MODEL_TAG_KEY in value: + module_path, qualname = value[cls.MODEL_TAG_KEY].rsplit(":", 1) + module = importlib.import_module(module_path) + model_cls: Any = module + for attribute in qualname.split("."): + model_cls = getattr(model_cls, attribute) + result[key] = model_cls.model_validate(value["data"]) + else: + result[key] = value + return result + + class ToolSpec(BaseModel): name: str | None = Field(description="Unique name identifier for the tool") type: str | None = Field( @@ -164,6 +217,22 @@ class ToolSpec(BaseModel): default_factory=list, description="List of requirements for the tool, e.g., SANDBOX", ) + execution_metadata: ToolExecutionMetadata | None = Field( + default=None, + description=( + "Optional metadata used to rebuild this server tool in another process." + ), + ) + + @field_serializer("async_fn", "async_callback", when_used="json") + def _serialize_callable(self, _v: Any) -> None: + return None + + @field_validator("async_fn", mode="before") + @classmethod + def _deserialize_callable(cls, v: Any) -> Any: + """Restore callable after JSON deserialization.""" + return _dummy_tool_async_fn if v is None else v def get_original_tool_name(self) -> str: """Get the original tool name without version suffix.""" @@ -190,6 +259,7 @@ class ToolSpec(BaseModel): partial_params: dict[str, Any] | None = None, instructions: str | None = None, requirements: list[ToolRequirements] | None = None, + execution_metadata: ToolExecutionMetadata | None = None, ) -> "ToolSpec": """Create a ToolSpec from default parameters.""" if not input_schema and not async_fn: @@ -214,6 +284,7 @@ class ToolSpec(BaseModel): partial_params=partial_params, instructions=instructions, requirements=requirements or [], + execution_metadata=execution_metadata, ) @classmethod @@ -246,6 +317,7 @@ class ToolSpec(BaseModel): partial_params=tool.partial_params if hasattr(tool, "partial_params") else None, + execution_metadata=None, ) def to_function_tool(self) -> "FunctionTool": @@ -406,6 +478,20 @@ class ResponseFormatConfig(BaseModel): default=None, description="Output class to use for the response format" ) + @field_serializer("output_cls", when_used="json") + def _serialize_output_cls( + self, + output_cls: type[BaseModel] | None, + ) -> dict[str, Any] | None: + return output_cls.model_json_schema() if output_cls is not None else None + + @field_validator("output_cls", mode="before") + @classmethod + def _deserialize_output_cls(cls, output_cls: Any) -> Any: + if isinstance(output_cls, dict): + return create_model_from_json_schema(output_cls) + return output_cls + class ChatRequest(BaseModel): """Request model for chat-based engines with agent capabilities. diff --git a/private_gpt/components/chat/processors/chat_history/memory/strategies/condenser.py b/private_gpt/components/chat/processors/chat_history/memory/strategies/condenser.py index e231ab21..22a91003 100644 --- a/private_gpt/components/chat/processors/chat_history/memory/strategies/condenser.py +++ b/private_gpt/components/chat/processors/chat_history/memory/strategies/condenser.py @@ -34,6 +34,7 @@ from private_gpt.di import get_global_injector from private_gpt.utils.batches import aiter_batch from private_gpt.utils.tokens import ( MessageInputProtocol, + async_tokenizer, estimate_token_count, ) @@ -109,7 +110,11 @@ class CondenserContextMemoryStrategy(BaseMemoryStrategy): async def stop_condition(text: str) -> bool: if max_tokens is not None: - token_count = tokenizer_fn(text) if tokenizer_fn else None + token_count = ( + await async_tokenizer(text, tokenizer_fn=tokenizer_fn) + if tokenizer_fn + else None + ) if token_count is not None: return len(token_count) <= max_tokens @@ -209,6 +214,7 @@ class CondenserContextMemoryStrategy(BaseMemoryStrategy): direction="left", max_length=max_length, tokenizer_fn=tokenizer_fn, + current_tokens=current_tokens, ) if not chat_history: return [] @@ -393,10 +399,12 @@ class CondenserContextMemoryStrategy(BaseMemoryStrategy): direction: Literal["left", "right"], tokenizer_fn: TokenizerFn | None, max_length: int, + current_tokens: int | None = None, ) -> list[ChatMessage]: - current_tokens = await self._get_messages_tokens( - chat_history, tokenizer_fn=tokenizer_fn - ) + if current_tokens is None: + current_tokens = await self._get_messages_tokens( + chat_history, tokenizer_fn=tokenizer_fn + ) if current_tokens <= max_length: return chat_history @@ -408,14 +416,24 @@ class CondenserContextMemoryStrategy(BaseMemoryStrategy): if direction == "right": blocks.reverse() - # Drop blocks from specified direction until under token limit - remaining_blocks = blocks.copy() - while remaining_blocks and current_tokens > max_length: - dropped_block = remaining_blocks.pop(0) - dropped_tokens = await self._get_messages_tokens( - dropped_block, tokenizer_fn=tokenizer_fn + block_token_counts: list[int] = list( + await asyncio.gather( + *[ + self._get_messages_tokens(block, tokenizer_fn=tokenizer_fn) + for block in blocks + ] ) - current_tokens -= dropped_tokens + ) + + # Drop blocks from the front until we fit within max_length + drop_until = 0 + for block_tokens in block_token_counts: + if current_tokens <= max_length: + break + current_tokens -= block_tokens + drop_until += 1 + + remaining_blocks = blocks[drop_until:] # Flatten remaining blocks back to message list result = [msg for block in remaining_blocks for msg in block] @@ -447,6 +465,9 @@ class CondenserContextMemoryStrategy(BaseMemoryStrategy): tokenizer_fn: TokenizerFn | None, chat_history: list[ChatMessage], max_length: int, + left_tokens: int | None = None, + last_user_tokens: int | None = None, + right_tokens: int | None = None, iteration: int = 0, ) -> list[ChatMessage]: if iteration > MAX_CONDENSE_ITERATIONS: @@ -458,30 +479,64 @@ class CondenserContextMemoryStrategy(BaseMemoryStrategy): right_messages, ) = await asyncio.to_thread(self._split_conversation, chat_history) - # -1. Drop thinking: fully from left, fully from right - left_messages = self._drop_thinking(left_messages) - right_messages = self._drop_thinking(right_messages) - current_history = [*left_messages, last_user_message, *right_messages] - current_tokens = await self._get_messages_tokens( - current_history, tokenizer_fn=tokenizer_fn - ) - if current_tokens <= max_length: - return current_history + # Compute per-component token counts if not provided by caller + if left_tokens is None or last_user_tokens is None or right_tokens is None: + left_tokens, last_user_tokens, right_tokens = await asyncio.gather( + self._get_messages_tokens(left_messages, tokenizer_fn=tokenizer_fn), + self._get_messages_tokens(last_user_message, tokenizer_fn=tokenizer_fn), + self._get_messages_tokens(right_messages, tokenizer_fn=tokenizer_fn), + ) + assert left_tokens is not None + assert last_user_tokens is not None + assert right_tokens is not None - # 0. Drop tools + # -1. Drop thinking from left/right, then re-count changed parts in parallel. + new_left = self._drop_thinking(left_messages) + new_right = self._drop_thinking(right_messages) + has_thinking_left = any( + "thinking" in m.additional_kwargs for m in left_messages + ) + has_thinking_right = any( + "thinking" in m.additional_kwargs for m in right_messages + ) + if has_thinking_left or has_thinking_right: + recount_tasks = [] + if has_thinking_left: + recount_tasks.append( + self._get_messages_tokens(new_left, tokenizer_fn=tokenizer_fn) + ) + if has_thinking_right: + recount_tasks.append( + self._get_messages_tokens(new_right, tokenizer_fn=tokenizer_fn) + ) + recount_results = list(await asyncio.gather(*recount_tasks)) + idx = 0 + if has_thinking_left: + left_tokens = recount_results[idx] + idx += 1 + if has_thinking_right: + right_tokens = recount_results[idx] + left_messages = new_left + right_messages = new_right + + current_tokens = left_tokens + last_user_tokens + right_tokens + if current_tokens <= max_length: + return [*left_messages, last_user_message, *right_messages] + + # 0. Drop tools from left; only re-count if tools were present + has_tools = any(m.role == "tool" for m in left_messages) potential_left_history = await asyncio.to_thread( self._drop_tool_messages, left_messages ) - current_history = [ - *potential_left_history, - last_user_message, - *right_messages, - ] - current_tokens = await self._get_messages_tokens( - current_history, tokenizer_fn=tokenizer_fn - ) - if current_tokens <= max_length: - return current_history + if has_tools: + new_left_tokens = await self._get_messages_tokens( + potential_left_history, tokenizer_fn=tokenizer_fn + ) + current_tokens = new_left_tokens + last_user_tokens + right_tokens + if current_tokens <= max_length: + return [*potential_left_history, last_user_message, *right_messages] + left_messages = potential_left_history + left_tokens = new_left_tokens # 1. Summarize the left part of the conversation current_left = await self._summary_left( @@ -490,13 +545,12 @@ class CondenserContextMemoryStrategy(BaseMemoryStrategy): chat_history=left_messages, max_length=max_length, ) - current_history = [*current_left, last_user_message, *right_messages] - current_tokens = await self._get_messages_tokens( - current_history, tokenizer_fn=tokenizer_fn + left_tokens = await self._get_messages_tokens( + current_left, tokenizer_fn=tokenizer_fn ) - + current_tokens = left_tokens + last_user_tokens + right_tokens if current_tokens <= max_length: - return current_history + return [*current_left, last_user_message, *right_messages] # 2. Summarize the right part of the conversation current_right = await self._summary_right( @@ -506,24 +560,22 @@ class CondenserContextMemoryStrategy(BaseMemoryStrategy): last_user_message=last_user_message, max_length=max_length, ) - current_history = [*current_left, last_user_message, *current_right] - current_tokens = await self._get_messages_tokens( - current_history, tokenizer_fn=tokenizer_fn + right_tokens = await self._get_messages_tokens( + current_right, tokenizer_fn=tokenizer_fn ) - + current_tokens = left_tokens + last_user_tokens + right_tokens if current_tokens <= max_length: - return current_history + return [*current_left, last_user_message, *current_right] # 3. Drop messages from the left side until we reach the max length - available_left_tokens = max_length - await self._get_messages_tokens( - [last_user_message, *current_right], - tokenizer_fn=tokenizer_fn, - ) + # available_left_tokens is arithmetic — no Triton call needed + available_left_tokens = max_length - last_user_tokens - right_tokens current_left = await self._drop_in_direction( chat_history=current_left, direction="left", max_length=available_left_tokens, tokenizer_fn=tokenizer_fn, + current_tokens=left_tokens, ) # 4. Repair the left and right parts of the conversation @@ -533,11 +585,14 @@ class CondenserContextMemoryStrategy(BaseMemoryStrategy): repair_with_tools, [last_user_message, *current_right], strict=False ), ) - - current_history = [*current_left, *current_right] - current_tokens = await self._get_messages_tokens( - current_history, tokenizer_fn=tokenizer_fn + # current_right now includes last_user_message after repair_with_tools; + # count both in parallel + repaired_left_tokens, repaired_right_tokens = await asyncio.gather( + self._get_messages_tokens(current_left, tokenizer_fn=tokenizer_fn), + self._get_messages_tokens(current_right, tokenizer_fn=tokenizer_fn), ) + current_history = [*current_left, *current_right] + current_tokens = repaired_left_tokens + repaired_right_tokens if current_tokens > max_length: return await self._condense_from_left( @@ -556,6 +611,9 @@ class CondenserContextMemoryStrategy(BaseMemoryStrategy): tokenizer_fn: TokenizerFn | None, chat_history: list[ChatMessage], max_length: int, + left_tokens: int | None = None, + last_user_tokens: int | None = None, + right_tokens: int | None = None, iteration: int = 0, ) -> list[ChatMessage]: if iteration > MAX_CONDENSE_ITERATIONS: @@ -567,25 +625,46 @@ class CondenserContextMemoryStrategy(BaseMemoryStrategy): right_messages, ) = await asyncio.to_thread(self._split_conversation, chat_history) - # 0a. Drop thinking from right except the last occurrence - right_messages = self._drop_thinking_except_last(right_messages) - left_messages = self._drop_thinking(left_messages) - current_history = [*left_messages, last_user_message, *right_messages] - current_tokens = await self._get_messages_tokens( - current_history, tokenizer_fn=tokenizer_fn - ) - if current_tokens <= max_length: - return current_history + # Compute per-component token counts if not provided by caller + if left_tokens is None or last_user_tokens is None or right_tokens is None: + left_tokens, last_user_tokens, right_tokens = await asyncio.gather( + self._get_messages_tokens(left_messages, tokenizer_fn=tokenizer_fn), + self._get_messages_tokens(last_user_message, tokenizer_fn=tokenizer_fn), + self._get_messages_tokens(right_messages, tokenizer_fn=tokenizer_fn), + ) + assert left_tokens is not None + assert last_user_tokens is not None + assert right_tokens is not None - # 0b. Drop remaining thinking from right and all from left - right_messages = self._drop_thinking(right_messages) - left_messages = self._drop_thinking(left_messages) - current_history = [*left_messages, last_user_message, *right_messages] - current_tokens = await self._get_messages_tokens( - current_history, tokenizer_fn=tokenizer_fn + # 0a. Drop thinking from right (except last occurrence) and from left; + # re-count both in parallel since either may have changed + new_right = self._drop_thinking_except_last(right_messages) + new_left = self._drop_thinking(left_messages) + new_left_tokens, new_right_tokens = await asyncio.gather( + self._get_messages_tokens(new_left, tokenizer_fn=tokenizer_fn), + self._get_messages_tokens(new_right, tokenizer_fn=tokenizer_fn), ) + left_messages = new_left + right_messages = new_right + left_tokens = new_left_tokens + right_tokens = new_right_tokens + current_tokens = left_tokens + last_user_tokens + right_tokens if current_tokens <= max_length: - return current_history + return [*left_messages, last_user_message, *right_messages] + + # 0b. Drop remaining thinking from right (left already fully stripped) + new_right = self._drop_thinking(right_messages) + has_remaining_thinking = any( + "thinking" in m.additional_kwargs for m in right_messages + ) + if has_remaining_thinking: + right_tokens = await self._get_messages_tokens( + new_right, tokenizer_fn=tokenizer_fn + ) + right_messages = new_right + current_tokens = left_tokens + last_user_tokens + right_tokens + if current_tokens <= max_length: + return [*left_messages, last_user_message, *right_messages] # 1. Summarize the right part of the conversation current_right = await self._summary_right( @@ -596,13 +675,12 @@ class CondenserContextMemoryStrategy(BaseMemoryStrategy): max_length=max_length, aggressive_level=iteration, ) - current_history = [*left_messages, last_user_message, *current_right] - current_tokens = await self._get_messages_tokens( - current_history, tokenizer_fn=tokenizer_fn + right_tokens = await self._get_messages_tokens( + current_right, tokenizer_fn=tokenizer_fn ) - + current_tokens = left_tokens + last_user_tokens + right_tokens if current_tokens <= max_length: - return current_history + return [*left_messages, last_user_message, *current_right] # 3. Summarize the left part of the conversation current_left = await self._summary_left( @@ -611,20 +689,18 @@ class CondenserContextMemoryStrategy(BaseMemoryStrategy): chat_history=left_messages, max_length=max_length, ) - current_history = [*current_left, last_user_message, *current_right] - current_tokens = await self._get_messages_tokens( - current_history, tokenizer_fn=tokenizer_fn + left_tokens = await self._get_messages_tokens( + current_left, tokenizer_fn=tokenizer_fn ) - + current_tokens = left_tokens + last_user_tokens + right_tokens if current_tokens <= max_length: - return current_history + return [*current_left, last_user_message, *current_right] # 4. Drop messages from the left side until we reach the max length - right_side = await self._get_messages_tokens( - [last_user_message, *current_right], - tokenizer_fn=tokenizer_fn, - ) + # right_side is arithmetic — no Triton call needed + right_side = last_user_tokens + right_tokens if right_side >= max_length: + current_history = [*current_left, last_user_message, *current_right] return await self._condense_from_right( llm=llm, tokenizer_fn=tokenizer_fn, @@ -639,20 +715,23 @@ class CondenserContextMemoryStrategy(BaseMemoryStrategy): direction="left", max_length=available_left_tokens, tokenizer_fn=tokenizer_fn, + current_tokens=left_tokens, ) - # 4. Repair the left and right parts of the conversation + # 5. Repair the left and right parts of the conversation current_left, current_right = await asyncio.gather( asyncio.to_thread(repair_without_tools, current_left), asyncio.to_thread( repair_with_tools, [last_user_message, *current_right], strict=False ), ) - - current_history = [*current_left, *current_right] - current_tokens = await self._get_messages_tokens( - current_history, tokenizer_fn=tokenizer_fn + # current_right now includes last_user_message; count both in parallel + repaired_left_tokens, repaired_right_tokens = await asyncio.gather( + self._get_messages_tokens(current_left, tokenizer_fn=tokenizer_fn), + self._get_messages_tokens(current_right, tokenizer_fn=tokenizer_fn), ) + current_history = [*current_left, *current_right] + current_tokens = repaired_left_tokens + repaired_right_tokens if current_tokens > max_length: return await self._condense_from_right( @@ -716,6 +795,9 @@ class CondenserContextMemoryStrategy(BaseMemoryStrategy): tokenizer_fn=tokenizer, chat_history=chat_history, max_length=max_length, + left_tokens=left_tokens, + last_user_tokens=last_user_message_tokens, + right_tokens=right_tokens, ) else: chat_history = await self._condense_from_right( @@ -723,6 +805,9 @@ class CondenserContextMemoryStrategy(BaseMemoryStrategy): tokenizer_fn=tokenizer, chat_history=chat_history, max_length=max_length, + left_tokens=left_tokens, + last_user_tokens=last_user_message_tokens, + right_tokens=right_tokens, ) # After condensation, we ensure that the last user message is preserved diff --git a/private_gpt/components/engines/chat_loop/__init__.py b/private_gpt/components/engines/chat/__init__.py similarity index 100% rename from private_gpt/components/engines/chat_loop/__init__.py rename to private_gpt/components/engines/chat/__init__.py diff --git a/private_gpt/components/engines/chat/async_chat_engine.py b/private_gpt/components/engines/chat/async_chat_engine.py new file mode 100644 index 00000000..6a35cf14 --- /dev/null +++ b/private_gpt/components/engines/chat/async_chat_engine.py @@ -0,0 +1,1622 @@ +"""Job-native async chat loop engine — each method is one job phase. + +Parallel implementation of ChatLoopEngine: shares the same Pydantic state +models but exposes a job-friendly API instead of a while loop. Intended to +replace ChatLoopEngine once regression tests confirm identical event output. + +Job graph:: + + run_start + └── run_iteration(0) + ├── COMPLETED → close_chat_job(stop_reason) + ├── CONTINUE → run_iteration(N+1) + └── WAITING → IterationScheduler.on_waiting + └── (all tools done) → resume("after_tool") + └── run_iteration(N+1) +""" + +import asyncio +import contextlib +import json +import logging +from abc import ABC, abstractmethod +from collections.abc import AsyncGenerator, Awaitable, Callable +from contextlib import suppress +from dataclasses import dataclass, field +from datetime import UTC, datetime, timedelta +from functools import partial +from typing import TYPE_CHECKING, Any, Literal, cast +from uuid import uuid4 + +from llama_index.core.base.llms.types import ChatMessage, ChatResponse +from llama_index.core.llms import MessageRole +from llama_index.core.llms.function_calling import FunctionCallingLLM +from llama_index.core.tools import AsyncBaseTool, ToolSelection, adapt_to_async_tool +from pydantic import BaseModel, Field + +from private_gpt.components.chat.models.chat_config_models import ( + ChatRequest, + ResolvedChatRequest, + ToolSpec, +) +from private_gpt.components.container_registry import ContainerRegistry +from private_gpt.components.context.models.context_stack import ContextStack +from private_gpt.components.engines.chat.chat_engine_interface import ( + ChatEngineExecution, +) +from private_gpt.components.engines.chat.chat_runner import ChatRunner +from private_gpt.components.engines.chat.interceptors.chat_interceptor import ( + ChatRequestLoopInterceptor, + ChatResponseLoopInterceptor, +) +from private_gpt.components.engines.chat.interceptors.ensure_index_is_refreshed_interceptor import ( + EnsureIndexIsRefreshedInterceptor, +) +from private_gpt.components.engines.chat.interceptors.ensure_timestamp_in_content_blocks_interceptors import ( + EnsureTimestampInContentBlocksInterceptor, +) +from private_gpt.components.engines.chat.interceptors.ensure_tools_are_flatten_interceptor import ( + EnsureToolAreFlattenInterceptor, +) +from private_gpt.components.engines.chat.interceptors.restore_stateless_input_interceptor import ( + RestoreStatelessInputInterceptorRequest, +) +from private_gpt.components.engines.chat.models.chat_interceptor_context import ( + ChatInterceptorContext, +) +from private_gpt.components.engines.chat.models.chat_phase import ( + InterceptorPhase, + TimelinePhase, +) +from private_gpt.components.engines.chat.models.chat_state import ( + ChatInputState, + ChatOutputState, + ChatRuntimeState, + ChatState, + ChatStatus, + ChatTimelineEntry, +) +from private_gpt.components.engines.chat.models.execution_hooks import ( + ExecutionHooks, +) +from private_gpt.components.engines.chat.utils.request_builder import ( + build_initial_context_stack, + build_request_from_context_stack, +) +from private_gpt.components.engines.chat.utils.tool_utils import ( + select_tool_names, +) +from private_gpt.components.llm.custom.base import StructuredOutputsParams, ZylonLLM +from private_gpt.components.llm.llm_component import LLMComponent +from private_gpt.components.llm.models import ReasoningEffort +from private_gpt.components.llm.priorities import DefinedPriorities +from private_gpt.components.tools.processors.base import _session_id +from private_gpt.components.tools.remote_execution import ( + ToolExecutionInterceptor, + ToolExecutionRequest, + ToolExecutionResponse, + build_tool_execution_context, +) +from private_gpt.components.tools.tool_scheduler import ( + BaseToolScheduler, + LocalToolScheduler, +) +from private_gpt.events.models import ( + Container, + Event, + InputJSONDelta, + MessageOutputDelta, + RawContentBlockDeltaEvent, + RawContentBlockStartEvent, + RawContentBlockStopEvent, + RawMessageDeltaEvent, + RawMessageStartEvent, + RawMessageStopEvent, + StopReasonEnum, + TextDelta, + ThinkingBlock, + ThinkingDelta, + ToolResultBlock, + ToolUseBlock, + Usage, +) + +if TYPE_CHECKING: + from private_gpt.components.streaming.tasks.chat_scheduler import ( + BaseChatScheduler, + ) +from private_gpt.server.chat.interceptors.schema_coercing_tool_interceptor import ( + _coerce_kwargs, +) + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Internal data classes (mirrors of ChatLoopEngine — no shared import so that +# AsyncChatEngine stays independent once ChatLoopEngine is deleted) +# --------------------------------------------------------------------------- + + +class _ToolExecutionStatus: + EXECUTED = "executed" + NOT_EXECUTED = "not_executed" + PENDING = "pending" + NOT_FOUND = "not_found" + + +class _IterationCheckpoint: + START = "start" + BEFORE_ITERATION = "before_iteration" + TOOLS = "tools" + AFTER_ITERATION = "after_iteration" + CLOSE = "close" + + +class IterationCheckpointPayload(BaseModel): + pending_async_tools: dict[str, str] = Field(default_factory=dict) + tool_responses: list[ToolExecutionResponse] = Field(default_factory=list) + pending_external_tool_calls: list[ToolSelection] = Field(default_factory=list) + total_input_tokens: int = 0 + total_output_tokens: int = 0 + has_input_usage: bool = False + has_output_usage: bool = False + + +@dataclass +class _CheckpointContext: + checkpoint: str + stop_reason: str | None = None + payload: IterationCheckpointPayload = field( + default_factory=IterationCheckpointPayload + ) + + +@dataclass +class _ToolExecutionResult: + status: str + tool_selection: ToolSelection | None = None + async_handle: str | None = None + + +@dataclass +class _Run: + """Hold mutable runtime references for one phase invocation.""" + + state: ChatState + llm: FunctionCallingLLM + total_input_tokens: int = 0 + total_output_tokens: int = 0 + has_input_usage: bool = False + has_output_usage: bool = False + block_count: int = 0 + hooks: ExecutionHooks = field(default_factory=ExecutionHooks) + + +@dataclass +class _ToolDeltaState: + active_tool_block: RawContentBlockStartEvent | None = None + active_tool_raw_id: str | None = None + tool_id_map: dict[str, str] = field(default_factory=dict) + finished_tool_raw_ids: set[str] = field(default_factory=set) + last_serialized: dict[str, str] = field(default_factory=dict) + pending_tasks: list[asyncio.Task[_ToolExecutionResult]] = field( + default_factory=list + ) + tool_semaphore: asyncio.Semaphore | None = None + + +@dataclass +class _StreamDeltaState: + active_block: RawContentBlockStartEvent | None = None + active_block_kind: Literal["text", "thinking"] | None = None + tool_state: _ToolDeltaState = field(default_factory=_ToolDeltaState) + + +class _Done: + """Sentinel: producer finished.""" + + +class EventChannel(ABC): + """Sink for events produced by the engine.""" + + @abstractmethod + def emit(self, event: Event) -> None: + ... + + @abstractmethod + async def close(self) -> None: + ... + + +class LocalEventChannel(EventChannel): + """asyncio.Queue-backed channel for local (in-process) execution.""" + + def __init__(self) -> None: + self._queue: asyncio.Queue[Event | _Done] = asyncio.Queue() + + def emit(self, event: Event) -> None: + self._queue.put_nowait(event) + + async def close(self) -> None: + self._queue.put_nowait(_Done()) + + async def stream( + self, task: asyncio.Task[Any] | None = None + ) -> AsyncGenerator[Event, None]: + try: + while True: + item = await self._queue.get() + if isinstance(item, _Done): + break + yield item + finally: + if task is not None and not task.done(): + task.cancel() + if task is not None: + with suppress(asyncio.CancelledError): + await task + + +@dataclass +class _EventHandler(EventChannel): + queue: asyncio.Queue[Event | _Done] + + def emit(self, event: Event) -> None: + self.queue.put_nowait(event) + + async def close(self) -> None: + self.queue.put_nowait(_Done()) + + async def stream(self, producer: asyncio.Task[Any]) -> AsyncGenerator[Event, None]: + try: + while True: + item = await self.queue.get() + if isinstance(item, _Done): + break + yield item + finally: + if not producer.done(): + producer.cancel() + with suppress(asyncio.CancelledError): + await producer + + +# --------------------------------------------------------------------------- +# AsyncChatEngine +# --------------------------------------------------------------------------- + + +class AsyncChatEngine: + """Job-native chat loop engine. + + Each public method corresponds to exactly one ARQ job phase. The engine + never drives the next phase itself — job code reads ``final_state.output.status`` + and dispatches the next job accordingly. + + status values returned: + CONTINUE — phase complete; job should dispatch the next phase + COMPLETED — no more LLM iterations needed; job should dispatch close_chat_job + WAITING — async tools are pending; job should call IterationScheduler.on_waiting + """ + + def __init__( + self, + llm_component: LLMComponent, + chat_scheduler: "BaseChatScheduler", + request_interceptors: list[ChatRequestLoopInterceptor] | None = None, + response_interceptors: list[ChatResponseLoopInterceptor] | None = None, + max_iterations: int = 40, + container_registry: ContainerRegistry | None = None, + tool_scheduler: BaseToolScheduler | None = None, + tool_interceptors: list[ToolExecutionInterceptor] | None = None, + runner: ChatRunner | None = None, + ) -> None: + self._llm_component = llm_component + self._request_interceptors = [ + RestoreStatelessInputInterceptorRequest(), + *(request_interceptors or []), + EnsureToolAreFlattenInterceptor(), + ] + self._response_interceptors = [ + *(response_interceptors or []), + EnsureIndexIsRefreshedInterceptor(), + EnsureTimestampInContentBlocksInterceptor(), + ] + self._max_iterations = max_iterations + self._container_registry = container_registry + self._tool_scheduler: BaseToolScheduler = tool_scheduler or LocalToolScheduler() + self._tool_interceptors = tool_interceptors or [] + self._chat_scheduler = chat_scheduler + self._runner = runner + self._checkpoint_handlers = { + _IterationCheckpoint.START: self._execute_start_checkpoint, + _IterationCheckpoint.BEFORE_ITERATION: self._execute_before_iteration_checkpoint, + _IterationCheckpoint.TOOLS: self._process_tools_checkpoint, + _IterationCheckpoint.AFTER_ITERATION: self._execute_after_iteration_checkpoint, + _IterationCheckpoint.CLOSE: self._execute_close_checkpoint, + } + + # ------------------------------------------------------------------ + # Public entry points + # ------------------------------------------------------------------ + + async def execute( + self, + request: ChatRequest, + hooks: ExecutionHooks | None = None, + *, + channel: EventChannel, + ) -> ChatState: + """Run from start through the loop until WAITING or COMPLETED.""" + state = await self._execute_start(request, hooks, channel) + return await self._continue_loop( + state.input.request, + state.runtime.iteration, + state.runtime.next_block_count, + IterationCheckpointPayload(), + hooks, + channel, + ) + + async def resume( + self, + resume_type: str, + request: ChatRequest, + iteration: int = 0, + next_block_count: int = 0, + hooks: ExecutionHooks | None = None, + checkpoint_payload: IterationCheckpointPayload | None = None, + *, + channel: EventChannel, + ) -> ChatState: + """Route to the saved checkpoint, then continue the loop.""" + payload = checkpoint_payload or IterationCheckpointPayload() + context = self._build_checkpoint_context( + checkpoint=resume_type, payload=payload + ) + handler = self._resolve_checkpoint_handler(resume_type) + state = await handler( + request, iteration, next_block_count, channel, hooks, context + ) + new_payload = IterationCheckpointPayload( + total_input_tokens=state.runtime.total_input_tokens, + total_output_tokens=state.runtime.total_output_tokens, + has_input_usage=state.runtime.has_input_usage, + has_output_usage=state.runtime.has_output_usage, + ) + if state.output.status == ChatStatus.WAITING: + return state + elif state.output.status == ChatStatus.COMPLETED: + close_state = await self._execute_close( + state.input.request, + state.output.stop_reason, + new_payload, + hooks, + channel, + ) + if state.output.pending_external_tool_calls: + close_state = close_state.model_copy(deep=True) + close_state.output.pending_external_tool_calls = list( + state.output.pending_external_tool_calls + ) + return close_state + return await self._continue_loop( + state.input.request, + state.runtime.iteration, + state.runtime.next_block_count, + new_payload, + hooks, + channel, + ) + + async def run( + self, + request: ChatRequest, + hooks: ExecutionHooks | None = None, + *, + runner: ChatRunner | None = None, + ) -> ChatEngineExecution: + """Submit execution through the transport-independent resumable runner.""" + runner = runner or self._runner + if runner is None: + raise RuntimeError("AsyncChatEngine requires a ResumableChatRunner") + execution_id, broker_events = await runner.submit( + request_data=request.model_dump(mode="json"), + stream_type="chat_completion", + metadata={ + "message_count": len(request.messages), + "thinking_enabled": request.thinking.enabled, + }, + execution_id=getattr(request.context, "correlation_id", None), + ) + + async def _events() -> AsyncGenerator[Event, None]: + completed = False + try: + async for event in broker_events: + yield event + completed = True + finally: + if not completed: + await runner.cancel(execution_id) + + return ChatEngineExecution( + events=_events(), final_state_task=None, execution_id=execution_id + ) + + async def execute_scheduled_start( + self, + *, + execution_id: str, + request_data: dict[str, Any], + stream_type: str, + metadata: dict[str, Any], + ) -> None: + runner = self._require_runner() + await runner.start( + engine=self, + execution_id=execution_id, + request_data=request_data, + stream_type=stream_type, + metadata=metadata, + ) + + async def execute_scheduled_resume(self, *, execution_id: str) -> None: + runner = self._require_runner() + await runner.resume(engine=self, execution_id=execution_id) + + async def record_callback( + self, *, execution_id: str, tool_id: str, result: dict[str, Any] + ) -> None: + runner = self._require_runner() + await runner.callback( + execution_id=execution_id, + tool_id=tool_id, + result=result, + ) + + async def execute_timeout(self, *, execution_id: str, checkpoint_id: str) -> None: + runner = self._require_runner() + await runner.timeout( + execution_id=execution_id, + checkpoint_id=checkpoint_id, + ) + + async def validate(self, request: ResolvedChatRequest) -> None: + await self.run_interceptor_phase( + run=self.initialize_run(request), + phase=InterceptorPhase.VALIDATION, + interceptors=self._request_interceptors, + ) + + async def cancel(self, correlation_id: str) -> bool: + if self._runner is not None: + return bool(await self._runner.cancel(correlation_id)) + return bool(await self._chat_scheduler.cancel(correlation_id)) + + def _require_runner(self) -> ChatRunner: + if self._runner is None: + raise RuntimeError("AsyncChatEngine requires a ResumableChatRunner") + return self._runner + + # ------------------------------------------------------------------ + # Internal loop helpers + # ------------------------------------------------------------------ + + async def _execute_start( + self, + request: ChatRequest, + hooks: ExecutionHooks | None, + channel: EventChannel, + ) -> ChatState: + context = self._build_checkpoint_context(checkpoint=_IterationCheckpoint.START) + return await self._execute_start_checkpoint( + request, 0, 0, channel, hooks, context + ) + + async def _execute_before_iteration( + self, + request: ChatRequest, + iteration: int, + next_block_count: int, + payload: IterationCheckpointPayload, + hooks: ExecutionHooks | None, + channel: EventChannel, + ) -> ChatState: + context = self._build_checkpoint_context( + checkpoint=_IterationCheckpoint.BEFORE_ITERATION, + payload=payload, + ) + return await self._execute_before_iteration_checkpoint( + request, iteration, next_block_count, channel, hooks, context + ) + + async def _execute_close( + self, + request: ChatRequest, + stop_reason: str | None, + payload: IterationCheckpointPayload, + hooks: ExecutionHooks | None, + channel: EventChannel, + ) -> ChatState: + context = self._build_checkpoint_context( + checkpoint=_IterationCheckpoint.CLOSE, + stop_reason=stop_reason, + payload=payload, + ) + return await self._execute_close_checkpoint( + request, 0, 0, channel, hooks, context + ) + + async def _continue_loop( + self, + request: ChatRequest, + iteration: int, + next_block_count: int, + payload: IterationCheckpointPayload, + hooks: ExecutionHooks | None, + channel: EventChannel, + ) -> ChatState: + while True: + if iteration >= self._max_iterations: + return await self._execute_close( + request, StopReasonEnum.MAX_TOKENS.value, payload, hooks, channel + ) + state = await self._execute_before_iteration( + request, iteration, next_block_count, payload, hooks, channel + ) + new_payload = IterationCheckpointPayload( + total_input_tokens=state.runtime.total_input_tokens, + total_output_tokens=state.runtime.total_output_tokens, + has_input_usage=state.runtime.has_input_usage, + has_output_usage=state.runtime.has_output_usage, + ) + if state.output.status == ChatStatus.WAITING: + return state + elif state.output.status == ChatStatus.COMPLETED: + close_state = await self._execute_close( + state.input.request, + state.output.stop_reason, + new_payload, + hooks, + channel, + ) + if state.output.pending_external_tool_calls: + close_state = close_state.model_copy(deep=True) + close_state.output.pending_external_tool_calls = list( + state.output.pending_external_tool_calls + ) + return close_state + iteration = state.runtime.iteration + next_block_count = state.runtime.next_block_count + request = state.input.request + payload = new_payload + + # ------------------------------------------------------------------ + # Phase producers + # ------------------------------------------------------------------ + + def _build_checkpoint_context( + self, + *, + checkpoint: str, + stop_reason: str | None = None, + payload: IterationCheckpointPayload | None = None, + ) -> _CheckpointContext: + return _CheckpointContext( + checkpoint=checkpoint, + stop_reason=stop_reason, + payload=payload or IterationCheckpointPayload(), + ) + + def _resolve_checkpoint_handler( + self, checkpoint: str + ) -> Callable[ + [ + ChatRequest, + int, + int, + EventChannel, + ExecutionHooks | None, + _CheckpointContext, + ], + Awaitable[ChatState], + ]: + if checkpoint not in self._checkpoint_handlers: + raise ValueError(f"Unknown checkpoint: {checkpoint!r}") + return cast( + Callable[ + [ + ChatRequest, + int, + int, + EventChannel, + ExecutionHooks | None, + _CheckpointContext, + ], + Awaitable[ChatState], + ], + self._checkpoint_handlers[checkpoint], + ) + + async def _execute_before_iteration_checkpoint( + self, + request: ChatRequest, + iteration: int, + next_block_count: int, + channel: EventChannel, + hooks: ExecutionHooks | None, + checkpoint_context: _CheckpointContext, + ) -> ChatState: + run = self._initialize_run(request, hooks=hooks) + run.state.runtime.iteration = iteration + run.state.runtime.next_block_count = next_block_count + run.block_count = next_block_count + self._apply_payload_usage(run, checkpoint_context.payload) + await self._run_intercepted_iteration(run, channel) + self._store_runtime_usage(run) + run.state = run.state.model_copy(deep=True) + run.state.runtime.next_block_count = run.block_count + return run.state.model_copy(deep=True) + + async def _execute_start_checkpoint( + self, + request: ChatRequest, + iteration: int, + next_block_count: int, + channel: EventChannel, + hooks: ExecutionHooks | None, + checkpoint_context: _CheckpointContext, + ) -> ChatState: + del iteration, next_block_count, checkpoint_context + run = self._initialize_run(request, hooks=hooks) + channel.emit(RawMessageStartEvent.from_defaults()) + run.state = self._snapshot(run.state, TimelinePhase.START) + run.state = run.state.model_copy(deep=True) + await self.run_interceptor_phase( + run, + InterceptorPhase.VALIDATION, + self._request_interceptors, + channel, + ) + run.state = run.state.model_copy(deep=True) + run.state.output.status = ChatStatus.CONTINUE + return run.state + + async def _process_tools_checkpoint( + self, + request: ChatRequest, + iteration: int, + next_block_count: int, + channel: EventChannel, + hooks: ExecutionHooks | None, + checkpoint_context: _CheckpointContext, + ) -> ChatState: + return await self._continue_tools_checkpoint( + request, + iteration, + next_block_count, + channel, + hooks=hooks, + checkpoint_payload=checkpoint_context.payload, + ) + + async def _continue_tools_checkpoint( + self, + request: ChatRequest, + iteration: int, + next_block_count: int, + channel: EventChannel, + hooks: ExecutionHooks | None = None, + checkpoint_payload: IterationCheckpointPayload | None = None, + ) -> ChatState: + """Continue from the tools checkpoint without re-running the LLM call.""" + try: + run = self._initialize_run(request, hooks=hooks) + run.state.runtime.iteration = iteration + run.state.runtime.next_block_count = next_block_count + run.block_count = next_block_count + payload = checkpoint_payload or IterationCheckpointPayload() + self._apply_payload_usage(run, payload) + pending_external = payload.pending_external_tool_calls + + for response in payload.tool_responses: + result_start = RawContentBlockStartEvent( + index=run.block_count, + block_id=f"block_{uuid4().hex}", + content_block=ToolResultBlock( + tool_use_id=response.tool_id, + content=response.result_content, + is_error=response.is_error, + ), + ) + run.block_count += 1 + channel.emit(result_start) + channel.emit(RawContentBlockStopEvent.from_start(result_start)) + + if pending_external: + run.state = run.state.model_copy(deep=True) + run.state.output.pending_external_tool_calls = pending_external + run.state.output.stop_reason = StopReasonEnum.TOOL_USE.value + run.state.output.status = ChatStatus.COMPLETED + run.state = self._snapshot(run.state, TimelinePhase.STOP) + return run.state + + run.state = self._snapshot(run.state, TimelinePhase.AFTER_TOOLS) + return await self._execute_after_iteration_checkpoint( + request, + iteration, + next_block_count, + channel, + hooks, + self._build_checkpoint_context( + checkpoint=_IterationCheckpoint.AFTER_ITERATION, + ), + run=run, + ) + except asyncio.CancelledError: + logger.debug("tools checkpoint continuation cancelled") + raise + + async def _execute_after_iteration_checkpoint( + self, + request: ChatRequest, + iteration: int, + next_block_count: int, + channel: EventChannel, + hooks: ExecutionHooks | None, + checkpoint_context: _CheckpointContext, + run: _Run | None = None, + ) -> ChatState: + del request, iteration, next_block_count, checkpoint_context + if run is None: + raise ValueError("after_iteration checkpoint requires prepared run state") + await self.run_interceptor_phase( + run, + InterceptorPhase.AFTER_ITERATION, + self._request_interceptors, + channel, + ) + run.state = run.state.model_copy(deep=True) + run.state.runtime.next_block_count = run.block_count + run.state.output.status = ChatStatus.CONTINUE + return run.state + + async def _execute_close_checkpoint( + self, + request: ChatRequest, + iteration: int, + next_block_count: int, + channel: EventChannel, + hooks: ExecutionHooks | None, + checkpoint_context: _CheckpointContext, + ) -> ChatState: + del iteration, next_block_count + run = self._initialize_run(request, hooks=hooks) + self._apply_payload_usage(run, checkpoint_context.payload) + stop_reason = checkpoint_context.stop_reason + run.state = run.state.model_copy(deep=True) + run.state.output.stop_reason = stop_reason + channel.emit( + RawMessageDeltaEvent( + delta=MessageOutputDelta( + stop_reason=stop_reason, + container=self._build_container(run), + ), + usage=self._build_total_usage(run), + ) + ) + channel.emit(RawMessageStopEvent.from_defaults()) + run.state.output.status = ChatStatus.COMPLETED + run.state = self._snapshot(run.state, TimelinePhase.STOP) + return run.state + + # ------------------------------------------------------------------ + # Iteration machinery + # ------------------------------------------------------------------ + + def initialize_run( + self, + request: ChatRequest, + context_stack: ContextStack | None = None, + hooks: ExecutionHooks | None = None, + ) -> _Run: + return self._initialize_run(request, context_stack=context_stack, hooks=hooks) + + async def _run_intercepted_iteration( + self, + run: _Run, + channel: EventChannel, + ) -> None: + iter_handler = _EventHandler(queue=asyncio.Queue()) + + def current_context() -> ChatInterceptorContext: + return ChatInterceptorContext( + state=run.state, + llm=run.llm, + phase=InterceptorPhase.STREAMING, + emit_fn=iter_handler.emit, + ) + + async def _run_and_close() -> None: + try: + for interceptor in self._response_interceptors: + await interceptor.on_iteration_start(current_context()) + await self._run_iteration(run, iter_handler) + finally: + for interceptor in self._response_interceptors: + with suppress(Exception): + await interceptor.on_iteration_end(current_context()) + await iter_handler.close() + + iter_task = asyncio.create_task(_run_and_close()) + + async for raw_event in iter_handler.stream(iter_task): + processed: Event | None = raw_event + if processed is None: + continue + for interceptor in self._response_interceptors: + processed = await interceptor.intercept_event( + processed, current_context() + ) + if processed is None: + break + if processed is not None: + channel.emit(processed) + + async def _run_iteration(self, run: _Run, handler: _EventHandler) -> None: + """One LLM call. Sets status; stop events are emitted by run_close.""" + run.state = run.state.model_copy(deep=True) + + await self.run_interceptor_phase( + run, + InterceptorPhase.BEFORE_ITERATION, + self._request_interceptors, + handler, + ) + + run.state.input.request = build_request_from_context_stack( + run.state.input.request, + run.state.input.context_stack, + ) + run.state.runtime.iteration += 1 + run.state = self._snapshot(run.state, TimelinePhase.BEFORE_LLM) + + llm_tools = self._build_tools(run.state) + tool_specs_by_name = { + tool.name: tool + for tool in run.state.input.context_stack.all_tools() + if tool.name + } + schema_by_name: dict[str, dict[str, Any]] = { + tool.name: tool.input_schema or {} + for tool in run.state.input.context_stack.all_tools() + if tool.name + } + llm_response = ChatResponse(message=ChatMessage(role=MessageRole.ASSISTANT)) + + lock = asyncio.Lock() + stream_delta_state = _StreamDeltaState() + if not run.state.input.request.tool_config.allow_parallel_tool_calls: + stream_delta_state.tool_state.tool_semaphore = asyncio.Semaphore(1) + + llm_kwargs = run.state.input.llm_kwargs.copy() + if isinstance(run.llm, ZylonLLM): + llm_kwargs["priority"] = ( + run.state.input.request.system.priority + or DefinedPriorities.LLM.CHAT_PRIORITY, + ) + + response_stream = await run.llm.astream_chat_with_tools( + llm_tools, + chat_history=run.state.input.request.to_messages(), + allow_parallel_tool_calls=run.state.input.request.tool_config.allow_parallel_tool_calls, + **run.state.input.llm_kwargs, + ) + async for chunk in response_stream: + llm_response = await self._handle_stream_chunk( + run=run, + llm=run.llm, + chunk=chunk, + current_response=llm_response, + stream_delta_state=stream_delta_state, + handler=handler, + tool_specs_by_name=tool_specs_by_name, + schema_by_name=schema_by_name, + lock=lock, + ) + + self._close_active_stream_blocks( + run=run, + stream_delta_state=stream_delta_state, + handler=handler, + tool_specs_by_name=tool_specs_by_name, + schema_by_name=schema_by_name, + lock=lock, + ) + + for key, value in llm_response.additional_kwargs.items(): + llm_response.message.additional_kwargs.setdefault(key, value) + + tool_calls = await asyncio.to_thread( + partial( + run.llm.get_tool_calls_from_response, + response=llm_response, + error_on_no_tool_call=False, + ) + ) + if tool_calls: + llm_response.message.additional_kwargs["tool_calls"] = tool_calls + else: + llm_response.message.additional_kwargs.pop("tool_calls", None) + + assistant_message = llm_response.message + self._ensure_update_tool_ids_in_tool_selection( + run, assistant_message, stream_delta_state.tool_state.tool_id_map + ) + self._accumulate_usage(run, assistant_message) + + run.state = run.state.model_copy(deep=True) + run.state.input.request.messages = [ + *run.state.input.request.messages, + assistant_message, + ] + run.state = self._snapshot(run.state, TimelinePhase.AFTER_LLM) + + if not tool_calls: + stop_reason = assistant_message.additional_kwargs.get("stop_reason") + run.state = run.state.model_copy(deep=True) + run.state.output.stop_reason = stop_reason + # COMPLETED: job dispatches close_chat_job — it emits stop events + run.state.output.status = ChatStatus.COMPLETED + return + + tool_call_results = await asyncio.gather( + *stream_delta_state.tool_state.pending_tasks, return_exceptions=True + ) + + has_external_tool = False + has_pending_tool = False + pending_external = list(run.state.output.pending_external_tool_calls) + pending_async = dict(run.state.output.pending_async_tools) + for result in tool_call_results: + if isinstance(result, Exception): + raise result + elif isinstance(result, _ToolExecutionResult): + if result.status == _ToolExecutionStatus.NOT_EXECUTED: + has_external_tool = True + assert result.tool_selection is not None + pending_external.append(result.tool_selection) + elif result.status == _ToolExecutionStatus.PENDING: + has_pending_tool = True + assert result.tool_selection is not None + pending_async[result.tool_selection.tool_id] = ( + result.async_handle or "" + ) + + if has_pending_tool: + run.state = run.state.model_copy(deep=True) + run.state.output.status = ChatStatus.WAITING + run.state.output.pause_type = _IterationCheckpoint.TOOLS + run.state.output.pending_async_tools = pending_async + run.state.output.pending_external_tool_calls = pending_external + run.state = self._snapshot(run.state, TimelinePhase.STOP) + return + + if has_external_tool: + run.state = run.state.model_copy(deep=True) + run.state.output.pending_external_tool_calls = pending_external + run.state.output.stop_reason = StopReasonEnum.TOOL_USE.value + # COMPLETED: job dispatches close_chat_job + run.state.output.status = ChatStatus.COMPLETED + run.state = self._snapshot(run.state, TimelinePhase.STOP) + return + + # Sync tools ran inline — run AFTER_TOOLS + AFTER_ITERATION, then CONTINUE + run.state = self._snapshot(run.state, TimelinePhase.AFTER_TOOLS) + + await self.run_interceptor_phase( + run, + InterceptorPhase.AFTER_ITERATION, + self._request_interceptors, + handler, + ) + + run.state = run.state.model_copy(deep=True) + run.state.output.status = ChatStatus.CONTINUE + + # ------------------------------------------------------------------ + # Stream chunk processing (identical logic to ChatLoopEngine) + # ------------------------------------------------------------------ + + async def _handle_stream_chunk( + self, + run: _Run, + llm: FunctionCallingLLM, + chunk: ChatResponse, + current_response: ChatResponse, + stream_delta_state: _StreamDeltaState, + handler: _EventHandler, + tool_specs_by_name: dict[str, ToolSpec], + schema_by_name: dict[str, dict[str, Any]], + lock: asyncio.Lock, + ) -> ChatResponse: + assistant_message = current_response.message.model_copy(deep=True) + if chunk.delta: + assistant_message.content = (assistant_message.content or "") + chunk.delta + elif ( + assistant_message.content is None + and isinstance(chunk.message.content, str) + and chunk.message.content + ): + assistant_message.content = chunk.message.content + + assistant_message.role = chunk.message.role or assistant_message.role + for source in (chunk.additional_kwargs, chunk.message.additional_kwargs): + for key, value in source.items(): + if key == "thinking_delta" and isinstance(value, str): + previous = assistant_message.additional_kwargs.get("thinking") + assistant_message.additional_kwargs["thinking"] = ( + previous if isinstance(previous, str) else "" + ) + value + continue + if key == "token_ids_delta" and isinstance(value, list): + previous = assistant_message.additional_kwargs.get( + "token_ids_delta" + ) + existing = previous if isinstance(previous, list) else [] + if not ( + value + and len(existing) >= len(value) + and existing[-len(value) :] == value + ): + assistant_message.additional_kwargs["token_ids_delta"] = ( + existing + value + ) + continue + if key == "tool_calls": + if isinstance(value, list) and value: + assistant_message.additional_kwargs["tool_calls"] = value + continue + assistant_message.additional_kwargs[key] = value + + folded_response = ChatResponse( + message=assistant_message, + additional_kwargs=assistant_message.additional_kwargs, + ) + + tool_calls = await asyncio.to_thread( + partial( + llm.get_tool_calls_from_response, + response=folded_response, + error_on_no_tool_call=False, + ) + ) + if tool_calls: + for tool_call in tool_calls: + if tool_call.tool_id is None: + continue + + raw_id = tool_call.tool_id + tool_state = stream_delta_state.tool_state + + if raw_id not in tool_state.tool_id_map: + tool_state.tool_id_map[raw_id] = f"tool_{uuid4().hex}" + + if raw_id in tool_state.finished_tool_raw_ids: + continue + + if tool_state.active_tool_raw_id != raw_id: + if tool_state.active_tool_block is not None: + self._close_active_tool_block( + run=run, + stream_delta_state=stream_delta_state, + handler=handler, + tool_specs_by_name=tool_specs_by_name, + schema_by_name=schema_by_name, + lock=lock, + ) + + self._close_active_block(stream_delta_state, handler) + + async with lock: + unique_id = tool_state.tool_id_map[raw_id] + use_start = RawContentBlockStartEvent( + index=run.block_count, + block_id=f"block_{uuid4().hex}", + content_block=ToolUseBlock( + id=unique_id, + name=tool_call.tool_name, + input={}, + ), + ) + run.block_count += 1 + handler.emit(use_start) + + tool_state.active_tool_block = use_start + tool_state.active_tool_raw_id = raw_id + tool_state.last_serialized[raw_id] = "" + + if tool_call.tool_kwargs and tool_state.active_tool_block is not None: + tool_schema = schema_by_name.get(tool_call.tool_name or "", {}) + coerced_kwargs = ( + _coerce_kwargs(tool_call.tool_kwargs, tool_schema) + if tool_schema + else tool_call.tool_kwargs + ) + current_json = json.dumps(coerced_kwargs) + tool_state.last_serialized[raw_id] = current_json + handler.emit( + RawContentBlockDeltaEvent.from_content_block_start( + tool_state.active_tool_block, + InputJSONDelta(partial_json_obj=coerced_kwargs), + ) + ) + return folded_response + + reasoning = self._extract_reasoning(chunk.message) + if reasoning: + self._switch_active_block( + run=run, + stream_delta_state=stream_delta_state, + handler=handler, + target_kind="thinking", + ) + if stream_delta_state.active_block is not None: + handler.emit( + RawContentBlockDeltaEvent.from_content_block_start( + stream_delta_state.active_block, + ThinkingDelta.from_text(reasoning), + ) + ) + return folded_response + + if chunk.delta: + self._switch_active_block( + run=run, + stream_delta_state=stream_delta_state, + handler=handler, + target_kind="text", + ) + if stream_delta_state.active_block is not None: + handler.emit( + RawContentBlockDeltaEvent.from_content_block_start( + stream_delta_state.active_block, + TextDelta(text=chunk.delta), + ) + ) + return folded_response + + self._close_active_block(stream_delta_state, handler) + return folded_response + + def _close_active_stream_blocks( + self, + run: _Run, + stream_delta_state: _StreamDeltaState, + handler: _EventHandler, + tool_specs_by_name: dict[str, ToolSpec], + schema_by_name: dict[str, dict[str, Any]], + lock: asyncio.Lock, + ) -> None: + if stream_delta_state.tool_state.active_tool_block is not None: + self._close_active_tool_block( + run, + stream_delta_state, + handler, + tool_specs_by_name, + schema_by_name, + lock, + ) + self._close_active_block(stream_delta_state, handler) + + def _close_active_tool_block( + self, + run: _Run, + stream_delta_state: _StreamDeltaState, + handler: _EventHandler, + tool_specs_by_name: dict[str, ToolSpec], + schema_by_name: dict[str, dict[str, Any]], + lock: asyncio.Lock, + ) -> None: + tool_state = stream_delta_state.tool_state + if tool_state.active_tool_block is None: + return + + prev_raw_id = tool_state.active_tool_raw_id or "" + final_json = tool_state.last_serialized.get(prev_raw_id, "") + final_obj: Any = json.loads(final_json) if final_json else {} + + tool_name = tool_state.active_tool_block.content_block.name # type: ignore[union-attr] + tool_schema = schema_by_name.get(tool_name, {}) + if tool_schema: + final_obj = _coerce_kwargs(final_obj, tool_schema) + + handler.emit( + RawContentBlockDeltaEvent.from_content_block_start( + tool_state.active_tool_block, + InputJSONDelta( + partial_json=final_json, + partial_json_obj=final_obj, + ), + ) + ) + handler.emit(RawContentBlockStopEvent.from_start(tool_state.active_tool_block)) + tool_state.finished_tool_raw_ids.add(prev_raw_id) + + tool_call = ToolSelection( + tool_id=prev_raw_id, + tool_name=tool_name, + tool_kwargs=final_obj, + ) + task = asyncio.create_task( + self._handle_tool_use( + run=run, + tool_call=tool_call, + tool_specs_by_name=tool_specs_by_name, + handler=handler, + tool_id_map=tool_state.tool_id_map, + lock=lock, + semaphore=stream_delta_state.tool_state.tool_semaphore, + ) + ) + tool_state.pending_tasks.append(task) + tool_state.active_tool_block = None + tool_state.active_tool_raw_id = None + + @staticmethod + def _close_active_block( + stream_delta_state: _StreamDeltaState, + handler: _EventHandler, + ) -> None: + if stream_delta_state.active_block is None: + return + handler.emit( + RawContentBlockStopEvent.from_start(stream_delta_state.active_block) + ) + stream_delta_state.active_block = None + stream_delta_state.active_block_kind = None + + @staticmethod + def _switch_active_block( + run: _Run, + stream_delta_state: _StreamDeltaState, + handler: _EventHandler, + target_kind: Literal["text", "thinking"], + ) -> None: + if stream_delta_state.active_block_kind == target_kind: + return + + AsyncChatEngine._close_active_block(stream_delta_state, handler) + if target_kind == "text": + stream_delta_state.active_block = RawContentBlockStartEvent.from_text() + else: + stream_delta_state.active_block = RawContentBlockStartEvent( + block_id=f"block_{uuid4().hex}", + content_block=ThinkingBlock( + thinking="", + signature=f"sig_{uuid4().hex}", + ), + ) + stream_delta_state.active_block.index = run.block_count + run.block_count += 1 + stream_delta_state.active_block_kind = target_kind + handler.emit(stream_delta_state.active_block) + + # ------------------------------------------------------------------ + # Tool execution + # ------------------------------------------------------------------ + + async def _handle_tool_use( + self, + run: _Run, + tool_call: ToolSelection, + tool_specs_by_name: dict[str, ToolSpec], + handler: _EventHandler, + tool_id_map: dict[str, str], + lock: asyncio.Lock, + semaphore: asyncio.Semaphore | None = None, + ) -> _ToolExecutionResult: + raw_id = tool_call.tool_id or "" + call_id = tool_id_map.get(raw_id) or "" + if not call_id: + raise RuntimeError( + f"Missing tool ID mapping for '{tool_call.tool_name}' with raw ID '{raw_id}'" + ) + + tool_spec = tool_specs_by_name.get(tool_call.tool_name or "") + + if tool_spec is None: + error_content = f"Tool '{tool_call.tool_name}' not found." + async with lock: + error_start = RawContentBlockStartEvent( + index=run.block_count, + block_id=f"block_{uuid4().hex}", + content_block=ToolResultBlock( + tool_use_id=call_id, + content=error_content, + is_error=True, + ), + ) + run.block_count += 1 + handler.emit(error_start) + handler.emit(RawContentBlockStopEvent.from_start(error_start)) + + error_message = ChatMessage( + role="tool", + content=error_content, + additional_kwargs={ + "tool_call_id": call_id, + "tool_call_name": tool_call.tool_name, + "tool_call_args": tool_call.tool_kwargs, + "raw_output": error_content, + }, + ) + run.state = run.state.model_copy(deep=True) + run.state.input.request.messages = [ + *run.state.input.request.messages, + error_message, + ] + + return _ToolExecutionResult(status=_ToolExecutionStatus.NOT_FOUND) + + if tool_spec.runtime != "server": + return _ToolExecutionResult( + status=_ToolExecutionStatus.NOT_EXECUTED, + tool_selection=ToolSelection( + tool_id=call_id, + tool_name=tool_call.tool_name, + tool_kwargs=tool_call.tool_kwargs, + ), + ) + + if self._tool_scheduler.is_async: + handle = await self._tool_scheduler.async_execute( + ToolExecutionRequest( + tool_id=call_id, + tool_name=tool_call.tool_name or "", + tool_kwargs=tool_call.tool_kwargs, + tool_spec=tool_spec, + context=build_tool_execution_context(run.state), + hooks=run.hooks, + ), + state_ctx=run.state, + interceptors=self._tool_interceptors, + ) + return _ToolExecutionResult( + status=_ToolExecutionStatus.PENDING, + tool_selection=ToolSelection( + tool_id=call_id, + tool_name=tool_call.tool_name, + tool_kwargs=tool_call.tool_kwargs, + ), + async_handle=handle, + ) + + async with semaphore if semaphore is not None else contextlib.nullcontext(): + response = await self._tool_scheduler.execute( + ToolExecutionRequest( + tool_id=call_id, + tool_name=tool_call.tool_name or "", + tool_kwargs=tool_call.tool_kwargs, + tool_spec=tool_spec, + context=build_tool_execution_context(run.state), + hooks=run.hooks, + ), + state_ctx=run.state, + interceptors=self._tool_interceptors, + ) + + async with lock: + run.state = run.state.model_copy(deep=True) + run.state.input.request.messages = [ + *run.state.input.request.messages, + response.tool_message, + ] + + result_content = response.result_content + + async with lock: + result_start = RawContentBlockStartEvent( + index=run.block_count, + block_id=f"block_{uuid4().hex}", + content_block=ToolResultBlock( + tool_use_id=call_id, + content=result_content, + is_error=response.is_error, + ), + ) + run.block_count += 1 + handler.emit(result_start) + handler.emit(RawContentBlockStopEvent.from_start(result_start)) + + return _ToolExecutionResult(status=_ToolExecutionStatus.EXECUTED) + + # ------------------------------------------------------------------ + # Initialization and interceptor phases + # ------------------------------------------------------------------ + + def _initialize_run( + self, + request: ChatRequest, + context_stack: ContextStack | None = None, + hooks: ExecutionHooks | None = None, + ) -> _Run: + llm = self._llm_component.get_llm(request.system.model) + if not isinstance(llm, FunctionCallingLLM): + raise ValueError("Configured model does not support function calling") + + if not isinstance(request, ResolvedChatRequest) and context_stack is None: + raise ValueError("Configured context stack is required") + + llm_kwargs = dict(request.sampling_params) + if request.thinking.enabled and request.thinking.type: + llm_kwargs["reasoning_effort"] = ReasoningEffort.from_str( + request.thinking.type + ) + if request.response_format and request.response_format.output_cls: + structured = StructuredOutputsParams.from_optional( + output_cls=request.response_format.output_cls, + ) + if structured is not None: + llm_kwargs["structured_outputs"] = structured + + state = ChatState( + input=ChatInputState( + request=request, + context_stack=context_stack or build_initial_context_stack(request), + sampling_params=dict(request.sampling_params), + llm_kwargs=llm_kwargs, + ), + runtime=ChatRuntimeState( + iteration=0, + max_iterations=self._max_iterations, + ), + output=ChatOutputState(), + timeline=[], + ) + state.original_input = state.input.model_copy(deep=True) + run = _Run( + state=state, + llm=llm, + hooks=hooks or ExecutionHooks(), + ) + for message in request.messages: + self._accumulate_usage(run, message) + self._store_runtime_usage(run) + return run + + def _apply_payload_usage( + self, run: _Run, payload: IterationCheckpointPayload + ) -> None: + if payload.has_input_usage: + run.total_input_tokens = payload.total_input_tokens + run.has_input_usage = True + if payload.has_output_usage: + run.total_output_tokens = payload.total_output_tokens + run.has_output_usage = True + self._store_runtime_usage(run) + + def _store_runtime_usage(self, run: _Run) -> None: + run.state.runtime.total_input_tokens = run.total_input_tokens + run.state.runtime.total_output_tokens = run.total_output_tokens + run.state.runtime.has_input_usage = run.has_input_usage + run.state.runtime.has_output_usage = run.has_output_usage + + async def run_interceptor_phase( + self, + run: _Run, + phase: InterceptorPhase, + interceptors: list[ChatRequestLoopInterceptor], + channel: EventChannel | None = None, + ) -> None: + if not interceptors: + return + + phase_marker = ( + TimelinePhase.BEFORE_INTERCEPTORS + if phase == InterceptorPhase.BEFORE_ITERATION + else TimelinePhase.AFTER_INTERCEPTORS + ) + run.state = self._snapshot(run.state, phase_marker) + + for interceptor in interceptors: + context = ChatInterceptorContext( + state=run.state, + llm=run.llm, + phase=phase, + emit_fn=channel.emit if channel is not None else lambda _event: None, + ) + await interceptor.intercept(context) + run.state = context.state + + def _build_tools(self, state: ChatState) -> list[AsyncBaseTool]: + tool_specs = state.input.context_stack.all_tools() + if not tool_specs: + return [] + + allowed_names = set( + select_tool_names( + tool_choices=state.input.request.tool_config.tool_choices, + tool_names=[tool.name or "" for tool in tool_specs], + ) + ) + selected = [tool for tool in tool_specs if (tool.name or "") in allowed_names] + return [adapt_to_async_tool(tool.to_function_tool()) for tool in selected] + + # ------------------------------------------------------------------ + # Utility helpers (identical to ChatLoopEngine) + # ------------------------------------------------------------------ + + def _snapshot(self, state: ChatState, phase: TimelinePhase) -> ChatState: + new_state = state.model_copy(deep=True) + new_state.timeline.append( + ChatTimelineEntry( + iteration=new_state.runtime.iteration, + phase=phase, + conversation_size=len(new_state.input.request.to_messages()), + tool_count=len(new_state.input.context_stack.all_tools()), + stop_reason=new_state.output.stop_reason, + ) + ) + return new_state + + @staticmethod + def _extract_reasoning(message: ChatMessage) -> str | None: + if message.additional_kwargs.get("stop_reason") is not None: + return None + raw_reasoning = message.additional_kwargs.get("thinking_delta") + if isinstance(raw_reasoning, str): + return raw_reasoning + return None + + @staticmethod + def _ensure_update_tool_ids_in_tool_selection( + run: _Run, message: ChatMessage, tool_id_map: dict[str, str] + ) -> None: + tool_calls = message.additional_kwargs.get("tool_calls", []) + if not isinstance(tool_calls, list): + return + for tool_call in tool_calls: + if isinstance(tool_call, ToolSelection): + raw_id = tool_call.tool_id + if raw_id and raw_id in tool_id_map: + tool_call.tool_id = tool_id_map[raw_id] + + @staticmethod + def _accumulate_usage(run: _Run, message: ChatMessage) -> None: + raw_input_tokens = message.additional_kwargs.get("input_tokens") + if isinstance(raw_input_tokens, int): + run.total_input_tokens += raw_input_tokens + run.has_input_usage = True + + raw_output_tokens = message.additional_kwargs.get("output_tokens") + if isinstance(raw_output_tokens, int): + run.total_output_tokens += raw_output_tokens + run.has_output_usage = True + + @staticmethod + def _build_total_usage(run: _Run) -> Usage: + return Usage( + input_tokens=run.total_input_tokens if run.has_input_usage else None, + output_tokens=run.total_output_tokens if run.has_output_usage else None, + ) + + def _build_container(self, run: _Run) -> Container | None: + if self._container_registry is None: + return None + + request = run.state.input.request + if not isinstance(request, ResolvedChatRequest): + return None + + session_id = _session_id(request) + ttl = self._container_registry.get_ttl(session_id) + if ttl is None: + return None + + return Container( + id=session_id, + expires_at=datetime.now(tz=UTC) + timedelta(seconds=ttl), + ) diff --git a/private_gpt/components/engines/chat_loop/chat_loop_engine.py b/private_gpt/components/engines/chat/chat_engine.py similarity index 81% rename from private_gpt/components/engines/chat_loop/chat_loop_engine.py rename to private_gpt/components/engines/chat/chat_engine.py index 33638cd4..6db4a524 100644 --- a/private_gpt/components/engines/chat_loop/chat_loop_engine.py +++ b/private_gpt/components/engines/chat/chat_engine.py @@ -21,45 +21,50 @@ from llama_index.core.tools import AsyncBaseTool, ToolSelection, adapt_to_async_ from private_gpt.components.chat.models.chat_config_models import ( ChatRequest, ResolvedChatRequest, + ToolSpec, ) from private_gpt.components.container_registry import ContainerRegistry from private_gpt.components.context.models.context_stack import ContextStack -from private_gpt.components.engines.chat_loop.interceptors.chat_loop_interceptor import ( +from private_gpt.components.engines.chat.interceptors.chat_interceptor import ( ChatRequestLoopInterceptor, ChatResponseLoopInterceptor, ) -from private_gpt.components.engines.chat_loop.interceptors.ensure_index_is_refreshed_interceptor import ( +from private_gpt.components.engines.chat.interceptors.ensure_index_is_refreshed_interceptor import ( EnsureIndexIsRefreshedInterceptor, ) -from private_gpt.components.engines.chat_loop.interceptors.ensure_timestamp_in_content_blocks_interceptors import ( +from private_gpt.components.engines.chat.interceptors.ensure_timestamp_in_content_blocks_interceptors import ( EnsureTimestampInContentBlocksInterceptor, ) -from private_gpt.components.engines.chat_loop.interceptors.ensure_tools_are_flatten_interceptor import ( +from private_gpt.components.engines.chat.interceptors.ensure_tools_are_flatten_interceptor import ( EnsureToolAreFlattenInterceptor, ) -from private_gpt.components.engines.chat_loop.interceptors.restore_stateless_input_interceptor import ( +from private_gpt.components.engines.chat.interceptors.restore_stateless_input_interceptor import ( RestoreStatelessInputInterceptorRequest, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_interceptor_context import ( - ChatLoopInterceptorContext, +from private_gpt.components.engines.chat.models.chat_interceptor_context import ( + ChatInterceptorContext, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_phase import ( +from private_gpt.components.engines.chat.models.chat_phase import ( InterceptorPhase, TimelinePhase, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_state import ( - ChatLoopInputState, - ChatLoopOutputState, - ChatLoopRuntimeState, - ChatLoopState, - ChatLoopTimelineEntry, +from private_gpt.components.engines.chat.models.chat_state import ( + ChatInputState, + ChatOutputState, + ChatRuntimeState, + ChatState, + ChatStatus, + ChatTimelineEntry, ) -from private_gpt.components.engines.chat_loop.utils.request_builder import ( +from private_gpt.components.engines.chat.models.execution_hooks import ( + ExecutionHooks, + ToolExecutionHook, +) +from private_gpt.components.engines.chat.utils.request_builder import ( build_initial_context_stack, build_request_from_context_stack, ) -from private_gpt.components.engines.chat_loop.utils.tool_utils import ( - execute_tool_call, +from private_gpt.components.engines.chat.utils.tool_utils import ( select_tool_names, ) from private_gpt.components.llm.custom.base import StructuredOutputsParams, ZylonLLM @@ -67,6 +72,15 @@ from private_gpt.components.llm.llm_component import LLMComponent from private_gpt.components.llm.models import ReasoningEffort from private_gpt.components.llm.priorities import DefinedPriorities from private_gpt.components.tools.processors.base import _session_id +from private_gpt.components.tools.remote_execution import ( + ToolExecutionInterceptor, + ToolExecutionRequest, + build_tool_execution_context, +) +from private_gpt.components.tools.tool_scheduler import ( + BaseToolScheduler, + LocalToolScheduler, +) from private_gpt.events.models import ( Container, Event, @@ -85,7 +99,6 @@ from private_gpt.events.models import ( ToolResultBlock, ToolUseBlock, Usage, - from_tool_output, ) from private_gpt.server.chat.interceptors.schema_coercing_tool_interceptor import ( _coerce_kwargs, @@ -97,6 +110,7 @@ logger = logging.getLogger(__name__) class ToolExecutionStatus(StrEnum): EXECUTED = "executed" NOT_EXECUTED = "not_executed" + PENDING = "pending" NOT_FOUND = "not_found" @@ -104,13 +118,20 @@ class ToolExecutionStatus(StrEnum): class ToolExecutionResult: status: str tool_selection: ToolSelection | None = None + async_handle: str | None = None + + +@dataclass +class LoopExecution: + events: AsyncGenerator[Event, None] + final_state_task: asyncio.Task[ChatState] @dataclass class _LoopRun: """Hold mutable runtime references for one run invocation.""" - state: ChatLoopState + state: ChatState llm: FunctionCallingLLM stopped: bool = False total_input_tokens: int = 0 @@ -118,6 +139,7 @@ class _LoopRun: has_input_usage: bool = False has_output_usage: bool = False block_count: int = 0 + hooks: list[ToolExecutionHook] = field(default_factory=list) @dataclass @@ -162,7 +184,7 @@ class _LoopEventHandler: """Close event stream for the current loop.""" self.queue.put_nowait(_LoopDone()) - async def stream(self, producer: asyncio.Task[None]) -> AsyncGenerator[Event, None]: + async def stream(self, producer: asyncio.Task[Any]) -> AsyncGenerator[Event, None]: """Stream queue events and cancel producer when consumer stops.""" try: while True: @@ -187,6 +209,8 @@ class ChatLoopEngine: response_interceptors: list[ChatResponseLoopInterceptor] | None = None, max_iterations: int = 40, container_registry: ContainerRegistry | None = None, + tool_scheduler: BaseToolScheduler | None = None, + tool_interceptors: list[ToolExecutionInterceptor] | None = None, ) -> None: self._llm_component = llm_component self._request_interceptors = [ @@ -201,36 +225,54 @@ class ChatLoopEngine: ] self._max_iterations = max_iterations self._container_registry = container_registry + self._tool_scheduler: BaseToolScheduler = tool_scheduler or LocalToolScheduler() + self._tool_interceptors = tool_interceptors or [] async def run( self, request: ChatRequest, context_stack: ContextStack | None = None, - ) -> AsyncGenerator[Event, None]: + hooks: list[ToolExecutionHook] | None = None, + ) -> LoopExecution: """Execute the loop and stream produced events immediately.""" handler = _LoopEventHandler(queue=asyncio.Queue()) - producer = asyncio.create_task(self._run_loop(request, context_stack, handler)) - try: - async for event in handler.stream(producer): - yield event - except asyncio.CancelledError: - if not producer.done(): - producer.cancel() - with suppress(asyncio.CancelledError): - await producer - raise - except: - raise + producer = asyncio.create_task( + self._run_loop( + request, + context_stack, + handler, + hooks=hooks, + ) + ) + + async def event_stream() -> AsyncGenerator[Event, None]: + try: + async for event in handler.stream(producer): + yield event + except asyncio.CancelledError: + if not producer.done(): + producer.cancel() + with suppress(asyncio.CancelledError): + await producer + raise + + return LoopExecution(events=event_stream(), final_state_task=producer) async def _run_loop( self, request: ChatRequest, context_stack: ContextStack | None, handler: _LoopEventHandler, - ) -> None: + hooks: list[ToolExecutionHook] | None = None, + ) -> ChatState: try: - run = self.initialize_run(request, context_stack) + run = self.initialize_run( + request, + context_stack, + hooks=hooks, + ) await self._run_loop_core(run, handler) + return run.state.model_copy(deep=True) except asyncio.CancelledError: logger.debug("Chat loop producer cancelled") raise @@ -267,6 +309,7 @@ class ChatLoopEngine: ) ) handler.emit(RawMessageStopEvent.from_defaults()) + run.state.output.status = ChatStatus.COMPLETED run.state = self._snapshot(run.state, TimelinePhase.STOP) async def _run_intercepted_iteration( @@ -276,8 +319,8 @@ class ChatLoopEngine: ) -> None: iter_handler = _LoopEventHandler(queue=asyncio.Queue()) - def current_context() -> ChatLoopInterceptorContext: - return ChatLoopInterceptorContext( + def current_context() -> ChatInterceptorContext: + return ChatInterceptorContext( state=run.state, llm=run.llm, phase=InterceptorPhase.STREAMING, @@ -348,7 +391,11 @@ class ChatLoopEngine: run.state = self._snapshot(run.state, TimelinePhase.BEFORE_LLM) llm_tools = self._build_tools(run.state) - tool_by_name = {tool.metadata.name: tool for tool in llm_tools} + tool_specs_by_name = { + tool.name: tool + for tool in run.state.input.context_stack.all_tools() + if tool.name + } schema_by_name: dict[str, dict[str, Any]] = { tool.name: tool.input_schema or {} for tool in run.state.input.context_stack.all_tools() @@ -382,7 +429,7 @@ class ChatLoopEngine: current_response=llm_response, stream_delta_state=stream_delta_state, handler=handler, - tool_by_name=tool_by_name, + tool_specs_by_name=tool_specs_by_name, schema_by_name=schema_by_name, lock=lock, ) @@ -391,11 +438,14 @@ class ChatLoopEngine: run=run, stream_delta_state=stream_delta_state, handler=handler, - tool_by_name=tool_by_name, + tool_specs_by_name=tool_specs_by_name, schema_by_name=schema_by_name, lock=lock, ) + for key, value in llm_response.additional_kwargs.items(): + llm_response.message.additional_kwargs.setdefault(key, value) + # The order is important: # OpenAI returns a custom model instead of ToolSelection # We have to extract, and override before to add the history again @@ -428,6 +478,7 @@ class ChatLoopEngine: stop_reason = assistant_message.additional_kwargs.get("stop_reason") run.state = run.state.model_copy(deep=True) run.state.output.stop_reason = stop_reason + run.state.output.status = ChatStatus.COMPLETED run.stopped = True handler.emit( RawMessageDeltaEvent( @@ -448,7 +499,9 @@ class ChatLoopEngine: ) has_external_tool = False + has_pending_tool = False pending_external = list(run.state.output.pending_external_tool_calls) + pending_async = dict(run.state.output.pending_async_tools) for result in tool_call_results: if isinstance(result, Exception): raise result @@ -457,6 +510,20 @@ class ChatLoopEngine: has_external_tool = True assert result.tool_selection is not None pending_external.append(result.tool_selection) + elif result.status == ToolExecutionStatus.PENDING: + has_pending_tool = True + assert result.tool_selection is not None + pending_async[result.tool_selection.tool_id] = ( + result.async_handle or "" + ) + + if has_pending_tool: + run.state = run.state.model_copy(deep=True) + run.state.output.status = ChatStatus.WAITING + run.state.output.pending_async_tools = pending_async + run.stopped = True + run.state = self._snapshot(run.state, TimelinePhase.STOP) + return if has_external_tool: run.state = run.state.model_copy(deep=True) @@ -470,6 +537,7 @@ class ChatLoopEngine: ) ) handler.emit(RawMessageStopEvent.from_defaults()) + run.state.output.status = ChatStatus.COMPLETED run.state = self._snapshot(run.state, TimelinePhase.STOP) return @@ -490,7 +558,7 @@ class ChatLoopEngine: current_response: ChatResponse, stream_delta_state: _StreamDeltaState, handler: _LoopEventHandler, - tool_by_name: dict[str | None, AsyncBaseTool], + tool_specs_by_name: dict[str, ToolSpec], schema_by_name: dict[str, dict[str, Any]], lock: asyncio.Lock, ) -> ChatResponse: @@ -514,6 +582,20 @@ class ChatLoopEngine: previous if isinstance(previous, str) else "" ) + value continue + if key == "token_ids_delta" and isinstance(value, list): + previous = assistant_message.additional_kwargs.get( + "token_ids_delta" + ) + existing = previous if isinstance(previous, list) else [] + if not ( + value + and len(existing) >= len(value) + and existing[-len(value) :] == value + ): + assistant_message.additional_kwargs["token_ids_delta"] = ( + existing + value + ) + continue if key == "tool_calls": if isinstance(value, list) and value: assistant_message.additional_kwargs["tool_calls"] = value @@ -553,7 +635,7 @@ class ChatLoopEngine: run=run, stream_delta_state=stream_delta_state, handler=handler, - tool_by_name=tool_by_name, + tool_specs_by_name=tool_specs_by_name, schema_by_name=schema_by_name, lock=lock, ) @@ -637,14 +719,19 @@ class ChatLoopEngine: run: _LoopRun, stream_delta_state: _StreamDeltaState, handler: _LoopEventHandler, - tool_by_name: dict[str | None, AsyncBaseTool], + tool_specs_by_name: dict[str, ToolSpec], schema_by_name: dict[str, dict[str, Any]], lock: asyncio.Lock, ) -> None: """Close all active streaming blocks and spawn the last tool task.""" if stream_delta_state.tool_state.active_tool_block is not None: self._close_active_tool_block( - run, stream_delta_state, handler, tool_by_name, schema_by_name, lock + run, + stream_delta_state, + handler, + tool_specs_by_name, + schema_by_name, + lock, ) self._close_active_block(stream_delta_state, handler) @@ -653,7 +740,7 @@ class ChatLoopEngine: run: _LoopRun, stream_delta_state: _StreamDeltaState, handler: _LoopEventHandler, - tool_by_name: dict[str | None, AsyncBaseTool], + tool_specs_by_name: dict[str, ToolSpec], schema_by_name: dict[str, dict[str, Any]], lock: asyncio.Lock, ) -> None: @@ -692,7 +779,7 @@ class ChatLoopEngine: self._handle_tool_use( run=run, tool_call=tool_call, - tool_by_name=tool_by_name, + tool_specs_by_name=tool_specs_by_name, handler=handler, tool_id_map=tool_state.tool_id_map, lock=lock, @@ -749,6 +836,7 @@ class ChatLoopEngine: self, request: ChatRequest, context_stack: ContextStack | None = None, + hooks: list[ToolExecutionHook] | None = None, ) -> _LoopRun: """Build initial llm and state for one run.""" llm = self._llm_component.get_llm(request.system.model) @@ -771,22 +859,26 @@ class ChatLoopEngine: if structured is not None: llm_kwargs["structured_outputs"] = structured - state = ChatLoopState( - input=ChatLoopInputState( + state = ChatState( + input=ChatInputState( request=request, context_stack=context_stack or build_initial_context_stack(request), sampling_params=dict(request.sampling_params), llm_kwargs=llm_kwargs, ), - runtime=ChatLoopRuntimeState( + runtime=ChatRuntimeState( iteration=0, max_iterations=self._max_iterations, ), - output=ChatLoopOutputState(), + output=ChatOutputState(), timeline=[], ) state.original_input = state.input.model_copy(deep=True) - return _LoopRun(state=state, llm=llm) + return _LoopRun( + state=state, + llm=llm, + hooks=list(hooks or []), + ) async def run_interceptor_phase( self, @@ -807,7 +899,7 @@ class ChatLoopEngine: run.state = self._snapshot(run.state, phase_marker) for interceptor in interceptors: - context = ChatLoopInterceptorContext( + context = ChatInterceptorContext( state=run.state, llm=run.llm, phase=phase, @@ -816,7 +908,7 @@ class ChatLoopEngine: await interceptor.intercept(context) run.state = context.state - def _build_tools(self, state: ChatLoopState) -> list[AsyncBaseTool]: + def _build_tools(self, state: ChatState) -> list[AsyncBaseTool]: """Resolve tool specs into async llama-index tools for current state.""" tool_specs = state.input.context_stack.all_tools() if not tool_specs: @@ -835,7 +927,7 @@ class ChatLoopEngine: self, run: _LoopRun, tool_call: ToolSelection, - tool_by_name: dict[str | None, AsyncBaseTool], + tool_specs_by_name: dict[str, ToolSpec], handler: _LoopEventHandler, tool_id_map: dict[str, str], lock: asyncio.Lock, @@ -848,9 +940,9 @@ class ChatLoopEngine: f"Missing tool ID mapping for '{tool_call.tool_name}' with raw ID '{raw_id}'" ) - tool = tool_by_name.get(tool_call.tool_name) + tool_spec = tool_specs_by_name.get(tool_call.tool_name or "") - if tool is None: + if tool_spec is None: error_content = f"Tool '{tool_call.tool_name}' not found." async with lock: error_start = RawContentBlockStartEvent( @@ -884,7 +976,7 @@ class ChatLoopEngine: return ToolExecutionResult(status=ToolExecutionStatus.NOT_FOUND) - if tool.metadata.return_direct: + if tool_spec.runtime != "server": return ToolExecutionResult( status=ToolExecutionStatus.NOT_EXECUTED, tool_selection=ToolSelection( @@ -894,27 +986,51 @@ class ChatLoopEngine: ), ) - async with semaphore if semaphore is not None else contextlib.nullcontext(): - result, tool_message = await execute_tool_call( - tool=tool, - tool_name=tool_call.tool_name, - tool_id=call_id, - tool_kwargs=tool_call.tool_kwargs, + if self._tool_scheduler.is_async: + handle = await self._tool_scheduler.async_execute( + ToolExecutionRequest( + tool_id=call_id, + tool_name=tool_call.tool_name or "", + tool_kwargs=tool_call.tool_kwargs, + tool_spec=tool_spec, + context=build_tool_execution_context(run.state), + hooks=ExecutionHooks(tool_result=run.hooks), + ), state_ctx=run.state, + interceptors=self._tool_interceptors, + ) + return ToolExecutionResult( + status=ToolExecutionStatus.PENDING, + tool_selection=ToolSelection( + tool_id=call_id, + tool_name=tool_call.tool_name, + tool_kwargs=tool_call.tool_kwargs, + ), + async_handle=handle, + ) + + async with semaphore if semaphore is not None else contextlib.nullcontext(): + response = await self._tool_scheduler.execute( + ToolExecutionRequest( + tool_id=call_id, + tool_name=tool_call.tool_name or "", + tool_kwargs=tool_call.tool_kwargs, + tool_spec=tool_spec, + context=build_tool_execution_context(run.state), + hooks=ExecutionHooks(tool_result=run.hooks), + ), + state_ctx=run.state, + interceptors=self._tool_interceptors, ) async with lock: run.state = run.state.model_copy(deep=True) run.state.input.request.messages = [ *run.state.input.request.messages, - tool_message, + response.tool_message, ] - result_content = ( - from_tool_output(result.tool_output.raw_output) - if result.tool_output.raw_output is not None - else result.tool_output.content - ) + result_content = response.result_content async with lock: result_start = RawContentBlockStartEvent( @@ -923,7 +1039,7 @@ class ChatLoopEngine: content_block=ToolResultBlock( tool_use_id=call_id, content=result_content, - is_error=result.tool_output.is_error, + is_error=response.is_error, ), ) @@ -933,11 +1049,11 @@ class ChatLoopEngine: return ToolExecutionResult(status=ToolExecutionStatus.EXECUTED) - def _snapshot(self, state: ChatLoopState, phase: TimelinePhase) -> ChatLoopState: + def _snapshot(self, state: ChatState, phase: TimelinePhase) -> ChatState: """Append one immutable timeline entry.""" new_state = state.model_copy(deep=True) new_state.timeline.append( - ChatLoopTimelineEntry( + ChatTimelineEntry( iteration=new_state.runtime.iteration, phase=phase, conversation_size=len(new_state.input.request.to_messages()), diff --git a/private_gpt/components/engines/chat/chat_engine_interface.py b/private_gpt/components/engines/chat/chat_engine_interface.py new file mode 100644 index 00000000..d7acaacb --- /dev/null +++ b/private_gpt/components/engines/chat/chat_engine_interface.py @@ -0,0 +1,72 @@ +import asyncio +from collections.abc import AsyncGenerator +from dataclasses import dataclass +from typing import Protocol + +from private_gpt.components.chat.models.chat_config_models import ( + ChatRequest, + ResolvedChatRequest, +) +from private_gpt.components.engines.chat.chat_engine import ChatLoopEngine +from private_gpt.components.engines.chat.chat_runner import ChatRunner +from private_gpt.components.engines.chat.models.chat_phase import InterceptorPhase +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.events.models import Event + + +@dataclass +class ChatEngineExecution: + events: AsyncGenerator[Event, None] + final_state_task: asyncio.Task[ChatState] | None + execution_id: str | None = None + + +class ChatEngine(Protocol): + async def run( + self, + request: ChatRequest, + hooks: ExecutionHooks | None = None, + *, + runner: ChatRunner | None = None, + ) -> ChatEngineExecution: + ... + + async def validate(self, request: ResolvedChatRequest) -> None: + ... + + async def cancel(self, correlation_id: str) -> bool: + ... + + +class LoopChatEngineAdapter: + def __init__(self, engine: ChatLoopEngine) -> None: + self._engine = engine + + async def run( + self, + request: ChatRequest, + hooks: ExecutionHooks | None = None, + *, + runner: ChatRunner | None = None, + ) -> ChatEngineExecution: + del runner + execution = await self._engine.run( + request=request, + hooks=hooks.tool_result if hooks is not None else None, + ) + return ChatEngineExecution( + events=execution.events, + final_state_task=execution.final_state_task, + ) + + async def validate(self, request: ResolvedChatRequest) -> None: + await self._engine.run_interceptor_phase( + run=self._engine.initialize_run(request), + phase=InterceptorPhase.VALIDATION, + interceptors=self._engine._request_interceptors, + ) + + async def cancel(self, correlation_id: str) -> bool: + del correlation_id + return True diff --git a/private_gpt/components/engines/chat/chat_runner.py b/private_gpt/components/engines/chat/chat_runner.py new file mode 100644 index 00000000..af230f27 --- /dev/null +++ b/private_gpt/components/engines/chat/chat_runner.py @@ -0,0 +1,69 @@ +from collections.abc import AsyncGenerator +from typing import TYPE_CHECKING, Any, Protocol, cast + +from injector import Injector, inject, singleton + +from private_gpt.events.models import Event + +if TYPE_CHECKING: + from private_gpt.components.engines.chat.async_chat_engine import AsyncChatEngine + + +class ChatRunner(Protocol): + async def submit( + self, + *, + request_data: dict[str, Any], + stream_type: str, + metadata: dict[str, Any], + execution_id: str | None = None, + ) -> tuple[str, AsyncGenerator[Event, None]]: + ... + + async def cancel(self, execution_id: str) -> bool: + ... + + async def start( + self, + *, + engine: "AsyncChatEngine", + execution_id: str, + request_data: dict[str, Any], + stream_type: str, + metadata: dict[str, Any], + ) -> None: + ... + + async def resume( + self, + *, + engine: "AsyncChatEngine", + execution_id: str, + ) -> None: + ... + + async def callback( + self, + *, + execution_id: str, + tool_id: str, + result: dict[str, Any], + ) -> None: + ... + + async def timeout(self, *, execution_id: str, checkpoint_id: str) -> None: + ... + + +@singleton +class ChatRunnerFactory: + @inject + def __init__(self, injector: Injector) -> None: + self._injector = injector + + def get(self) -> ChatRunner: + from private_gpt.components.engines.chat.resumable_runner import ( + ResumableChatRunner, + ) + + return cast(ChatRunner, self._injector.get(ResumableChatRunner)) diff --git a/private_gpt/components/engines/chat/checkpoint_store.py b/private_gpt/components/engines/chat/checkpoint_store.py new file mode 100644 index 00000000..eedfb61a --- /dev/null +++ b/private_gpt/components/engines/chat/checkpoint_store.py @@ -0,0 +1,132 @@ +from __future__ import annotations + +import asyncio +from abc import ABC, abstractmethod +from datetime import datetime +from typing import Any, cast + +from injector import Injector, inject, singleton +from pydantic import BaseModel, Field + +from private_gpt.components.engines.chat.async_chat_engine import ( + IterationCheckpointPayload, +) +from private_gpt.components.tools.remote_execution import ToolExecutionResponse +from private_gpt.settings.settings import Settings + + +class ChatCheckpoint(BaseModel): + correlation_id: str + request_data: dict[str, Any] + stream_type: str + metadata: dict[str, Any] + iteration: int + checkpoint: str = "before_iteration" + checkpoint_payload: IterationCheckpointPayload = Field( + default_factory=IterationCheckpointPayload + ) + next_block_count: int = 0 + checkpoint_id: str = "" + deadline: datetime | None = None + + +class ChatCheckpointStore(ABC): + @abstractmethod + async def save(self, checkpoint: ChatCheckpoint) -> None: + ... + + @abstractmethod + async def load(self, execution_id: str) -> ChatCheckpoint | None: + ... + + @abstractmethod + async def record_result( + self, execution_id: str, tool_id: str, result: dict[str, Any] + ) -> dict[str, ToolExecutionResponse] | None: + ... + + @abstractmethod + async def get_results(self, execution_id: str) -> dict[str, ToolExecutionResponse]: + ... + + @abstractmethod + async def claim_resume(self, execution_id: str) -> bool: + ... + + @abstractmethod + async def cleanup(self, execution_id: str) -> None: + ... + + +@singleton +class InMemoryChatCheckpointStore(ChatCheckpointStore): + """Single-process checkpoint storage for local execution and tests.""" + + def __init__(self) -> None: + self._checkpoints: dict[str, ChatCheckpoint] = {} + self._results: dict[str, dict[str, ToolExecutionResponse]] = {} + self._resumed: set[str] = set() + self._lock = asyncio.Lock() + + async def save(self, checkpoint: ChatCheckpoint) -> None: + async with self._lock: + self._checkpoints[checkpoint.correlation_id] = checkpoint + self._resumed.discard(checkpoint.correlation_id) + + async def load(self, execution_id: str) -> ChatCheckpoint | None: + async with self._lock: + return self._checkpoints.get(execution_id) + + async def record_result( + self, execution_id: str, tool_id: str, result: dict[str, Any] + ) -> dict[str, ToolExecutionResponse] | None: + async with self._lock: + checkpoint = self._checkpoints.get(execution_id) + if checkpoint is None: + results = self._results.setdefault(execution_id, {}) + results[tool_id] = ToolExecutionResponse.model_validate(result) + return None + if tool_id not in checkpoint.checkpoint_payload.pending_async_tools: + return None + results = self._results.setdefault(execution_id, {}) + results[tool_id] = ToolExecutionResponse.model_validate(result) + expected = set(checkpoint.checkpoint_payload.pending_async_tools) + return dict(results) if expected.issubset(results) else None + + async def get_results(self, execution_id: str) -> dict[str, ToolExecutionResponse]: + async with self._lock: + return dict(self._results.get(execution_id, {})) + + async def claim_resume(self, execution_id: str) -> bool: + async with self._lock: + if execution_id in self._resumed: + return False + self._resumed.add(execution_id) + return True + + async def cleanup(self, execution_id: str) -> None: + async with self._lock: + self._checkpoints.pop(execution_id, None) + self._results.pop(execution_id, None) + self._resumed.discard(execution_id) + + +@singleton +class ChatCheckpointStoreFactory: + @inject + def __init__(self, settings: Settings, injector: Injector) -> None: + self._settings = settings + self._injector = injector + + def get(self) -> ChatCheckpointStore: + if self._settings.scheduler.chat.mode == "arq": + from private_gpt.arq.chat.iteration_state import RedisChatCheckpointStore + + return cast( + ChatCheckpointStore, + self._injector.get(RedisChatCheckpointStore), + ) + return cast( + ChatCheckpointStore, + self._injector.get(InMemoryChatCheckpointStore), + ) diff --git a/private_gpt/components/engines/chat/event_broker.py b/private_gpt/components/engines/chat/event_broker.py new file mode 100644 index 00000000..a4a3c3c7 --- /dev/null +++ b/private_gpt/components/engines/chat/event_broker.py @@ -0,0 +1,140 @@ +from __future__ import annotations + +import asyncio +from abc import ABC, abstractmethod +from collections import defaultdict +from collections.abc import Awaitable +from typing import TYPE_CHECKING, Any, cast + +import redis.asyncio as redis # type: ignore[import-untyped] +from injector import Injector, inject, singleton + +from private_gpt.arq.settings import get_redis_settings +from private_gpt.events.event_serializer import StreamingEventHandler +from private_gpt.settings.settings import Settings + +if TYPE_CHECKING: + from collections.abc import AsyncGenerator + + from private_gpt.events.models import Event + + +class EngineEventBroker(ABC): + """Transport used between chat execution and router-side stream processing.""" + + @abstractmethod + async def publish(self, execution_id: str, event: Event) -> None: + ... + + @abstractmethod + async def finish(self, execution_id: str) -> None: + ... + + @abstractmethod + def listen(self, execution_id: str) -> AsyncGenerator[Event, None]: + ... + + @abstractmethod + async def cleanup(self, execution_id: str) -> None: + ... + + +@singleton +class InMemoryEngineEventBroker(EngineEventBroker): + def __init__(self) -> None: + self._queues: dict[str, asyncio.Queue[Event | None]] = defaultdict( + asyncio.Queue + ) + + async def publish(self, execution_id: str, event: Event) -> None: + await self._queues[execution_id].put(event) + + async def finish(self, execution_id: str) -> None: + await self._queues[execution_id].put(None) + + async def listen(self, execution_id: str) -> AsyncGenerator[Event, None]: + queue = self._queues[execution_id] + try: + while True: + event = await queue.get() + if event is None: + return + yield event + finally: + self._queues.pop(execution_id, None) + + async def cleanup(self, execution_id: str) -> None: + self._queues.pop(execution_id, None) + + +@singleton +class RedisEngineEventBroker(EngineEventBroker): + _FINISHED = "__private_gpt_engine_finished__" + + @inject + def __init__(self, settings: Settings) -> None: + redis_settings = get_redis_settings(settings) + self._redis = redis.Redis( + host=cast(str, redis_settings.host), + port=redis_settings.port, + db=redis_settings.database, + username=redis_settings.username, + password=redis_settings.password, + decode_responses=True, + socket_timeout=5, + socket_connect_timeout=5, + ) + self._handler = StreamingEventHandler() + self._prefix = "private_gpt:engine:events" + self._ttl = settings.stream.stream_expiration + + def _key(self, execution_id: str) -> str: + return f"{self._prefix}:{execution_id}" + + async def publish(self, execution_id: str, event: Event) -> None: + key = self._key(execution_id) + await cast( + Awaitable[int], self._redis.rpush(key, self._handler.serialize(event)) + ) + await cast(Awaitable[int], self._redis.expire(key, self._ttl)) + + async def finish(self, execution_id: str) -> None: + key = self._key(execution_id) + await cast(Awaitable[int], self._redis.rpush(key, self._FINISHED)) + await cast(Awaitable[int], self._redis.expire(key, self._ttl)) + + async def listen(self, execution_id: str) -> AsyncGenerator[Event, None]: + key = self._key(execution_id) + try: + while True: + item = cast( + tuple[str, str] | None, + await cast(Awaitable[Any], self._redis.blpop([key], timeout=1)), + ) + if item is None: + continue + payload = item[1] + if payload == self._FINISHED: + return + yield self._handler.deserialize(payload) + finally: + await self.cleanup(execution_id) + + async def cleanup(self, execution_id: str) -> None: + await self._redis.delete(self._key(execution_id)) + + async def close(self) -> None: + await self._redis.aclose() + + +@singleton +class EngineEventBrokerFactory: + @inject + def __init__(self, settings: Settings, injector: Injector) -> None: + self._settings = settings + self._injector = injector + + def get(self) -> EngineEventBroker: + if self._settings.scheduler.chat.mode == "arq": + return self._injector.get(RedisEngineEventBroker) + return self._injector.get(InMemoryEngineEventBroker) diff --git a/private_gpt/components/engines/chat/event_channel.py b/private_gpt/components/engines/chat/event_channel.py new file mode 100644 index 00000000..670875cf --- /dev/null +++ b/private_gpt/components/engines/chat/event_channel.py @@ -0,0 +1,33 @@ +from __future__ import annotations + +import asyncio +from typing import TYPE_CHECKING + +from private_gpt.components.engines.chat.async_chat_engine import EventChannel + +if TYPE_CHECKING: + from private_gpt.components.engines.chat.event_broker import EngineEventBroker + from private_gpt.events.models import Event + + +class BrokerEventChannel(EventChannel): + """Ordered EventChannel adapter backed by an EngineEventBroker.""" + + def __init__(self, broker: EngineEventBroker, execution_id: str) -> None: + self._broker = broker + self._execution_id = execution_id + self._tail: asyncio.Future[None] | None = None + + def emit(self, event: Event) -> None: + previous = self._tail + + async def _publish() -> None: + if previous is not None: + await previous + await self._broker.publish(self._execution_id, event) + + self._tail = asyncio.create_task(_publish()) + + async def close(self) -> None: + if self._tail is not None: + await self._tail diff --git a/private_gpt/components/engines/chat/execution_scheduler.py b/private_gpt/components/engines/chat/execution_scheduler.py new file mode 100644 index 00000000..8b8505f6 --- /dev/null +++ b/private_gpt/components/engines/chat/execution_scheduler.py @@ -0,0 +1,192 @@ +from __future__ import annotations + +import asyncio +from abc import ABC, abstractmethod +from typing import Any + +from injector import Injector, inject, singleton + +from private_gpt.settings.settings import Settings + + +class ChatExecutionScheduler(ABC): + @abstractmethod + async def start( + self, + *, + execution_id: str, + request_data: dict[str, Any], + stream_type: str, + metadata: dict[str, Any], + ) -> None: + ... + + @abstractmethod + async def resume(self, *, execution_id: str) -> None: + ... + + @abstractmethod + async def callback( + self, *, execution_id: str, tool_id: str, result: dict[str, Any] + ) -> None: + ... + + @abstractmethod + async def timeout( + self, *, execution_id: str, checkpoint_id: str, delay_seconds: int + ) -> None: + ... + + @abstractmethod + async def cancel(self, execution_id: str) -> bool: + ... + + +@singleton +class LocalChatExecutionScheduler(ChatExecutionScheduler): + def __init__(self) -> None: + self._tasks: set[asyncio.Task[Any]] = set() + + def _schedule(self, coro: Any, *, name: str) -> None: + task = asyncio.create_task(coro, name=name) + self._tasks.add(task) + task.add_done_callback(self._tasks.discard) + + async def start( + self, + *, + execution_id: str, + request_data: dict[str, Any], + stream_type: str, + metadata: dict[str, Any], + ) -> None: + from private_gpt.di import get_global_injector + from private_gpt.server.chat.chat_service import ChatService + + engine = get_global_injector().get(ChatService).build_async_engine() + self._schedule( + engine.execute_scheduled_start( + execution_id=execution_id, + request_data=request_data, + stream_type=stream_type, + metadata=metadata, + ), + name=f"chat_{execution_id}", + ) + + async def resume(self, *, execution_id: str) -> None: + from private_gpt.di import get_global_injector + from private_gpt.server.chat.chat_service import ChatService + + engine = get_global_injector().get(ChatService).build_async_engine() + self._schedule( + engine.execute_scheduled_resume(execution_id=execution_id), + name=f"chat_{execution_id}", + ) + + async def callback( + self, *, execution_id: str, tool_id: str, result: dict[str, Any] + ) -> None: + from private_gpt.di import get_global_injector + from private_gpt.server.chat.chat_service import ChatService + + chat_service = get_global_injector().get(ChatService) + engine = chat_service.build_async_engine() + await engine.record_callback( + execution_id=execution_id, + tool_id=tool_id, + result=result, + ) + + async def timeout( + self, *, execution_id: str, checkpoint_id: str, delay_seconds: int + ) -> None: + from private_gpt.di import get_global_injector + from private_gpt.server.chat.chat_service import ChatService + + async def _timeout() -> None: + await asyncio.sleep(delay_seconds) + engine = get_global_injector().get(ChatService).build_async_engine() + await engine.execute_timeout( + execution_id=execution_id, checkpoint_id=checkpoint_id + ) + + self._schedule(_timeout(), name=f"chat_timeout_{execution_id}_{checkpoint_id}") + + async def cancel(self, execution_id: str) -> bool: + cancelled = False + for task in asyncio.all_tasks(): + if task.get_name().startswith(f"chat_{execution_id}"): + task.cancel() + cancelled = True + return cancelled + + +@singleton +class ArqChatExecutionScheduler(ChatExecutionScheduler): + async def start( + self, + *, + execution_id: str, + request_data: dict[str, Any], + stream_type: str, + metadata: dict[str, Any], + ) -> None: + from private_gpt.arq.tasks.chat.start import enqueue_start_chat_job + + await enqueue_start_chat_job( + request_data=request_data, + correlation_id=execution_id, + stream_type=stream_type, + metadata=metadata, + job_id=f"{execution_id}:start", + ) + + async def resume(self, *, execution_id: str) -> None: + from private_gpt.arq.tasks.chat.resume import enqueue_resume_iteration_job + + await enqueue_resume_iteration_job( + correlation_id=execution_id, + job_id=f"{execution_id}:resume", + ) + + async def callback( + self, *, execution_id: str, tool_id: str, result: dict[str, Any] + ) -> None: + from private_gpt.arq.tasks.chat.resume import enqueue_tool_resume_job + + await enqueue_tool_resume_job( + correlation_id=execution_id, + tool_id=tool_id, + result=result, + ) + + async def timeout( + self, *, execution_id: str, checkpoint_id: str, delay_seconds: int + ) -> None: + from private_gpt.arq.tasks.chat.resume import enqueue_chat_timeout_job + + await enqueue_chat_timeout_job( + correlation_id=execution_id, + checkpoint_id=checkpoint_id, + delay_seconds=delay_seconds, + job_id=f"{execution_id}:timeout:{checkpoint_id}", + ) + + async def cancel(self, execution_id: str) -> bool: + from private_gpt.arq.tasks.chat import abort_chat_job + + return await abort_chat_job(correlation_id=execution_id) + + +@singleton +class ChatExecutionSchedulerFactory: + @inject + def __init__(self, settings: Settings, injector: Injector) -> None: + self._settings = settings + self._injector = injector + + def get(self) -> ChatExecutionScheduler: + if self._settings.scheduler.chat.mode == "arq": + return self._injector.get(ArqChatExecutionScheduler) + return self._injector.get(LocalChatExecutionScheduler) diff --git a/private_gpt/components/engines/chat_loop/interceptors/__init__.py b/private_gpt/components/engines/chat/interceptors/__init__.py similarity index 100% rename from private_gpt/components/engines/chat_loop/interceptors/__init__.py rename to private_gpt/components/engines/chat/interceptors/__init__.py diff --git a/private_gpt/components/engines/chat_loop/interceptors/chat_loop_interceptor.py b/private_gpt/components/engines/chat/interceptors/chat_interceptor.py similarity index 76% rename from private_gpt/components/engines/chat_loop/interceptors/chat_loop_interceptor.py rename to private_gpt/components/engines/chat/interceptors/chat_interceptor.py index 9287f40b..de5b1b6b 100644 --- a/private_gpt/components/engines/chat_loop/interceptors/chat_loop_interceptor.py +++ b/private_gpt/components/engines/chat/interceptors/chat_interceptor.py @@ -4,8 +4,8 @@ from typing import Any, Self from pydantic import BaseModel -from private_gpt.components.engines.chat_loop.models.chat_loop_interceptor_context import ( - ChatLoopInterceptorContext, +from private_gpt.components.engines.chat.models.chat_interceptor_context import ( + ChatInterceptorContext, ) from private_gpt.events.models import Event @@ -14,7 +14,7 @@ class ChatRequestLoopInterceptor(ABC, BaseModel): """Transform loop context before inference starts.""" @abstractmethod - async def intercept(self, context: ChatLoopInterceptorContext) -> None: + async def intercept(self, context: ChatInterceptorContext) -> None: """Mutate context state and emit events when needed.""" def model_copy( @@ -28,11 +28,11 @@ class ChatRequestLoopInterceptor(ABC, BaseModel): class ChatResponseLoopInterceptor(ABC, BaseModel): """Process events emitted during inference, and mutate loop context as needed.""" - async def on_iteration_start(self, context: ChatLoopInterceptorContext) -> None: + async def on_iteration_start(self, context: ChatInterceptorContext) -> None: """Called once before each iteration — reset per-iteration state here.""" return - async def on_iteration_end(self, context: ChatLoopInterceptorContext) -> None: + async def on_iteration_end(self, context: ChatInterceptorContext) -> None: """Called once after each iteration.""" return @@ -40,7 +40,7 @@ class ChatResponseLoopInterceptor(ABC, BaseModel): async def intercept_event( self, event: Event, - context: ChatLoopInterceptorContext, + context: ChatInterceptorContext, ) -> Event | None: """Process one event. Mutate context.state as needed.""" diff --git a/private_gpt/components/engines/chat_loop/interceptors/chat_loop_interceptor_chain.py b/private_gpt/components/engines/chat/interceptors/chat_interceptor_chain.py similarity index 74% rename from private_gpt/components/engines/chat_loop/interceptors/chat_loop_interceptor_chain.py rename to private_gpt/components/engines/chat/interceptors/chat_interceptor_chain.py index 482e0647..5927024a 100644 --- a/private_gpt/components/engines/chat_loop/interceptors/chat_loop_interceptor_chain.py +++ b/private_gpt/components/engines/chat/interceptors/chat_interceptor_chain.py @@ -1,12 +1,18 @@ -from collections.abc import Callable, Mapping -from typing import Any, Self, cast +from __future__ import annotations + +from collections.abc import Callable +from typing import TYPE_CHECKING, Any, Self, cast + +if TYPE_CHECKING: + from collections.abc import Mapping from pydantic import BaseModel, ConfigDict, Field -from private_gpt.components.engines.chat_loop.interceptors.chat_loop_interceptor import ( +from private_gpt.components.engines.chat.interceptors.chat_interceptor import ( ChatRequestLoopInterceptor, ChatResponseLoopInterceptor, ) +from private_gpt.components.tools.remote_execution import ToolExecutionInterceptor Condition = bool | Callable[[], bool] @@ -16,11 +22,12 @@ def _evaluate(condition: Condition) -> bool: class InterceptorChainEntry(BaseModel): - """One named group in the chain holding N request and/or response interceptors.""" + """One named group holding request, response, and/or tool interceptors.""" name: str requests: list[ChatRequestLoopInterceptor] = Field(default_factory=list) responses: list[ChatResponseLoopInterceptor] = Field(default_factory=list) + tools: list[ToolExecutionInterceptor] = Field(default_factory=list) condition: Condition = True model_config = ConfigDict(arbitrary_types_allowed=True) @@ -39,6 +46,7 @@ class InterceptorChainEntry(BaseModel): name=str(self.name), requests=[r.model_copy(deep=True) for r in self.requests], responses=[r.model_copy(deep=True) for r in self.responses], + tools=list(self.tools), # tool interceptors are stateless singletons condition=self.condition, # shared — callable/bool are stateless ) return cast(Self, deep_copy) @@ -55,7 +63,7 @@ class ChatLoopInterceptorChain(BaseModel): request: ChatRequestLoopInterceptor | None = None, response: ChatResponseLoopInterceptor | None = None, condition: Condition = True, - ) -> "ChatLoopInterceptorChain": + ) -> ChatLoopInterceptorChain: if request is None and response is None: raise ValueError( f"Group '{name}' must provide at least one of request or response interceptor." @@ -76,25 +84,28 @@ class ChatLoopInterceptorChain(BaseModel): *, requests: list[ChatRequestLoopInterceptor] | None = None, responses: list[ChatResponseLoopInterceptor] | None = None, + tools: list[ToolExecutionInterceptor] | None = None, condition: Condition = True, - ) -> "ChatLoopInterceptorChain": + ) -> ChatLoopInterceptorChain: resolved_requests = requests or [] resolved_responses = responses or [] - if not resolved_requests and not resolved_responses: + resolved_tools = tools or [] + if not resolved_requests and not resolved_responses and not resolved_tools: raise ValueError( - f"Group '{name}' must provide at least one request or response interceptor." + f"Group '{name}' must provide at least one request, response, or tool interceptor." ) self.entries.append( InterceptorChainEntry( name=name, requests=resolved_requests, responses=resolved_responses, + tools=resolved_tools, condition=condition, ) ) return self - def clone(self) -> "ChatLoopInterceptorChain": + def clone(self) -> ChatLoopInterceptorChain: return ChatLoopInterceptorChain( entries=[entry.model_copy(deep=True) for entry in self.entries] ) @@ -106,3 +117,7 @@ class ChatLoopInterceptorChain(BaseModel): @property def response_interceptors(self) -> list[ChatResponseLoopInterceptor]: return [r for e in self.entries if e.is_active for r in e.responses] + + @property + def tool_interceptors(self) -> list[ToolExecutionInterceptor]: + return [t for e in self.entries if e.is_active for t in e.tools] diff --git a/private_gpt/components/engines/chat_loop/interceptors/ensure_index_is_refreshed_interceptor.py b/private_gpt/components/engines/chat/interceptors/ensure_index_is_refreshed_interceptor.py similarity index 79% rename from private_gpt/components/engines/chat_loop/interceptors/ensure_index_is_refreshed_interceptor.py rename to private_gpt/components/engines/chat/interceptors/ensure_index_is_refreshed_interceptor.py index 094c9dea..b305a005 100644 --- a/private_gpt/components/engines/chat_loop/interceptors/ensure_index_is_refreshed_interceptor.py +++ b/private_gpt/components/engines/chat/interceptors/ensure_index_is_refreshed_interceptor.py @@ -1,11 +1,11 @@ from collections.abc import Mapping from typing import Any -from private_gpt.components.engines.chat_loop.interceptors.chat_loop_interceptor import ( +from private_gpt.components.engines.chat.interceptors.chat_interceptor import ( ChatResponseLoopInterceptor, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_interceptor_context import ( - ChatLoopInterceptorContext, +from private_gpt.components.engines.chat.models.chat_interceptor_context import ( + ChatInterceptorContext, ) from private_gpt.events.models import ( Event, @@ -20,11 +20,12 @@ class EnsureIndexIsRefreshedInterceptor(ChatResponseLoopInterceptor): _current_index: int = 0 _block_id_map: dict[str, int] | None = None - async def on_iteration_start(self, context: ChatLoopInterceptorContext) -> None: + async def on_iteration_start(self, context: ChatInterceptorContext) -> None: self._block_id_map = {} + self._current_index = context.state.runtime.next_block_count async def intercept_event( - self, event: Event, context: ChatLoopInterceptorContext + self, event: Event, context: ChatInterceptorContext ) -> Event | None: if self._block_id_map is None: raise ValueError("Block ID map is not initialized. This should not happen.") @@ -36,6 +37,7 @@ class EnsureIndexIsRefreshedInterceptor(ChatResponseLoopInterceptor): ) self._block_id_map[block_id] = self._current_index self._current_index += 1 + context.state.runtime.next_block_count = self._current_index return event.model_copy(update={"index": self._block_id_map[block_id]}) case RawContentBlockDeltaEvent( diff --git a/private_gpt/components/engines/chat_loop/interceptors/ensure_timestamp_in_content_blocks_interceptors.py b/private_gpt/components/engines/chat/interceptors/ensure_timestamp_in_content_blocks_interceptors.py similarity index 79% rename from private_gpt/components/engines/chat_loop/interceptors/ensure_timestamp_in_content_blocks_interceptors.py rename to private_gpt/components/engines/chat/interceptors/ensure_timestamp_in_content_blocks_interceptors.py index 4d9cc882..e52fa21d 100644 --- a/private_gpt/components/engines/chat_loop/interceptors/ensure_timestamp_in_content_blocks_interceptors.py +++ b/private_gpt/components/engines/chat/interceptors/ensure_timestamp_in_content_blocks_interceptors.py @@ -2,11 +2,11 @@ import datetime from collections.abc import Mapping from typing import Any -from private_gpt.components.engines.chat_loop.interceptors.chat_loop_interceptor import ( +from private_gpt.components.engines.chat.interceptors.chat_interceptor import ( ChatResponseLoopInterceptor, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_interceptor_context import ( - ChatLoopInterceptorContext, +from private_gpt.components.engines.chat.models.chat_interceptor_context import ( + ChatInterceptorContext, ) from private_gpt.events.models import ( Event, @@ -17,7 +17,7 @@ from private_gpt.events.models import ( class EnsureTimestampInContentBlocksInterceptor(ChatResponseLoopInterceptor): async def intercept_event( - self, event: Event, context: ChatLoopInterceptorContext + self, event: Event, context: ChatInterceptorContext ) -> Event | None: match event: case RawContentBlockStartEvent(): diff --git a/private_gpt/components/engines/chat_loop/interceptors/ensure_tools_are_flatten_interceptor.py b/private_gpt/components/engines/chat/interceptors/ensure_tools_are_flatten_interceptor.py similarity index 89% rename from private_gpt/components/engines/chat_loop/interceptors/ensure_tools_are_flatten_interceptor.py rename to private_gpt/components/engines/chat/interceptors/ensure_tools_are_flatten_interceptor.py index aad8f746..40ff7977 100644 --- a/private_gpt/components/engines/chat_loop/interceptors/ensure_tools_are_flatten_interceptor.py +++ b/private_gpt/components/engines/chat/interceptors/ensure_tools_are_flatten_interceptor.py @@ -1,19 +1,19 @@ from llama_index.core.base.llms.types import ChatMessage, MessageRole from llama_index.core.llms.llm import ToolSelection -from private_gpt.components.engines.chat_loop.interceptors.chat_loop_interceptor import ( +from private_gpt.components.engines.chat.interceptors.chat_interceptor import ( ChatRequestLoopInterceptor, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_interceptor_context import ( - ChatLoopInterceptorContext, +from private_gpt.components.engines.chat.models.chat_interceptor_context import ( + ChatInterceptorContext, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_phase import ( +from private_gpt.components.engines.chat.models.chat_phase import ( InterceptorPhase, ) class EnsureToolAreFlattenInterceptor(ChatRequestLoopInterceptor): - async def intercept(self, context: ChatLoopInterceptorContext) -> None: + async def intercept(self, context: ChatInterceptorContext) -> None: if context.phase != InterceptorPhase.AFTER_ITERATION: return diff --git a/private_gpt/components/engines/chat_loop/interceptors/restore_stateless_input_interceptor.py b/private_gpt/components/engines/chat/interceptors/restore_stateless_input_interceptor.py similarity index 90% rename from private_gpt/components/engines/chat_loop/interceptors/restore_stateless_input_interceptor.py rename to private_gpt/components/engines/chat/interceptors/restore_stateless_input_interceptor.py index f44b709c..2dd9d8fc 100644 --- a/private_gpt/components/engines/chat_loop/interceptors/restore_stateless_input_interceptor.py +++ b/private_gpt/components/engines/chat/interceptors/restore_stateless_input_interceptor.py @@ -4,13 +4,13 @@ from pydantic import Field from private_gpt.components.context.models.context_stack import ContextStack from private_gpt.components.context.models.layer_type import LayerType -from private_gpt.components.engines.chat_loop.interceptors.chat_loop_interceptor import ( +from private_gpt.components.engines.chat.interceptors.chat_interceptor import ( ChatRequestLoopInterceptor, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_interceptor_context import ( - ChatLoopInterceptorContext, +from private_gpt.components.engines.chat.models.chat_interceptor_context import ( + ChatInterceptorContext, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_phase import ( +from private_gpt.components.engines.chat.models.chat_phase import ( InterceptorPhase, ) @@ -72,7 +72,7 @@ class RestoreStatelessInputInterceptorRequest(ChatRequestLoopInterceptor): return ContextStack(layers=merged_layers) - async def intercept(self, context: ChatLoopInterceptorContext) -> None: + async def intercept(self, context: ChatInterceptorContext) -> None: """Copy stateless fields from original_input into the current input state.""" if context.phase != InterceptorPhase.BEFORE_ITERATION: return diff --git a/private_gpt/components/engines/chat_loop/models/__init__.py b/private_gpt/components/engines/chat/models/__init__.py similarity index 100% rename from private_gpt/components/engines/chat_loop/models/__init__.py rename to private_gpt/components/engines/chat/models/__init__.py diff --git a/private_gpt/components/engines/chat/models/chat_interceptor_context.py b/private_gpt/components/engines/chat/models/chat_interceptor_context.py new file mode 100644 index 00000000..bbf4dc7f --- /dev/null +++ b/private_gpt/components/engines/chat/models/chat_interceptor_context.py @@ -0,0 +1,33 @@ +from collections.abc import Callable +from typing import Any + +from llama_index.core.llms.function_calling import FunctionCallingLLM +from pydantic import BaseModel, ConfigDict, Field + +from private_gpt.components.engines.chat.models.chat_phase import ( + InterceptorPhase, +) +from private_gpt.components.engines.chat.models.chat_state import ( + ChatState, +) +from private_gpt.events.models import Event + + +class ChatInterceptorContext(BaseModel): + """Carry state, llm runtime, and helpers across interceptors.""" + + state: ChatState + llm: FunctionCallingLLM + phase: InterceptorPhase + emit_fn: Callable[[Event], None] + metadata: dict[str, Any] = Field(default_factory=dict) + + model_config = ConfigDict(arbitrary_types_allowed=True) + + def emit_event(self, event: Event) -> None: + """Emit one event immediately to the loop stream.""" + self.emit_fn(event) + + def set_state(self, state: ChatState) -> None: + """Replace context state after interceptor transformation.""" + self.state = state diff --git a/private_gpt/components/engines/chat_loop/models/chat_loop_interceptor_context.py b/private_gpt/components/engines/chat/models/chat_loop_interceptor_context.py similarity index 72% rename from private_gpt/components/engines/chat_loop/models/chat_loop_interceptor_context.py rename to private_gpt/components/engines/chat/models/chat_loop_interceptor_context.py index 5f4707ca..bbf4dc7f 100644 --- a/private_gpt/components/engines/chat_loop/models/chat_loop_interceptor_context.py +++ b/private_gpt/components/engines/chat/models/chat_loop_interceptor_context.py @@ -4,19 +4,19 @@ from typing import Any from llama_index.core.llms.function_calling import FunctionCallingLLM from pydantic import BaseModel, ConfigDict, Field -from private_gpt.components.engines.chat_loop.models.chat_loop_phase import ( +from private_gpt.components.engines.chat.models.chat_phase import ( InterceptorPhase, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_state import ( - ChatLoopState, +from private_gpt.components.engines.chat.models.chat_state import ( + ChatState, ) from private_gpt.events.models import Event -class ChatLoopInterceptorContext(BaseModel): +class ChatInterceptorContext(BaseModel): """Carry state, llm runtime, and helpers across interceptors.""" - state: ChatLoopState + state: ChatState llm: FunctionCallingLLM phase: InterceptorPhase emit_fn: Callable[[Event], None] @@ -28,6 +28,6 @@ class ChatLoopInterceptorContext(BaseModel): """Emit one event immediately to the loop stream.""" self.emit_fn(event) - def set_state(self, state: ChatLoopState) -> None: + def set_state(self, state: ChatState) -> None: """Replace context state after interceptor transformation.""" self.state = state diff --git a/private_gpt/components/engines/chat_loop/models/chat_loop_phase.py b/private_gpt/components/engines/chat/models/chat_loop_phase.py similarity index 90% rename from private_gpt/components/engines/chat_loop/models/chat_loop_phase.py rename to private_gpt/components/engines/chat/models/chat_loop_phase.py index bd1c7fcb..71662654 100644 --- a/private_gpt/components/engines/chat_loop/models/chat_loop_phase.py +++ b/private_gpt/components/engines/chat/models/chat_loop_phase.py @@ -6,7 +6,9 @@ class InterceptorPhase(StrEnum): VALIDATION = "validation" BEFORE_ITERATION = "before_iteration" + BEFORE_TOOL = "before_tool" STREAMING = "streaming" + AFTER_TOOL = "after_tool" AFTER_ITERATION = "after_iteration" diff --git a/private_gpt/components/engines/chat_loop/models/chat_loop_state.py b/private_gpt/components/engines/chat/models/chat_loop_state.py similarity index 74% rename from private_gpt/components/engines/chat_loop/models/chat_loop_state.py rename to private_gpt/components/engines/chat/models/chat_loop_state.py index 5a405ab5..f484d44f 100644 --- a/private_gpt/components/engines/chat_loop/models/chat_loop_state.py +++ b/private_gpt/components/engines/chat/models/chat_loop_state.py @@ -1,4 +1,5 @@ from collections.abc import Mapping +from enum import StrEnum from typing import Any, Self from llama_index.core.llms.llm import ToolSelection @@ -6,16 +7,16 @@ from pydantic import BaseModel, ConfigDict, Field from private_gpt.components.chat.models.chat_config_models import ChatRequest from private_gpt.components.context.models.context_stack import ContextStack -from private_gpt.components.engines.chat_loop.models.chat_loop_phase import ( +from private_gpt.components.engines.chat.models.chat_phase import ( TimelinePhase, ) -from private_gpt.components.llm.llm_helper import TokenizerFn +from private_gpt.components.llm.llm_helper import AsyncTokenizerFn, TokenizerFn from private_gpt.components.skills.models.skill_entities import ( SkillVersionWithSkillEntity, ) -class ChatLoopInputState(BaseModel): +class ChatInputState(BaseModel): """Request-scoped inputs for one loop run. ``request`` — the original immutable ChatRequest (restored each iteration). @@ -35,27 +36,40 @@ class ChatLoopInputState(BaseModel): model_config = ConfigDict(arbitrary_types_allowed=True) -class ChatLoopRuntimeState(BaseModel): +class ChatRuntimeState(BaseModel): """Store runtime counters.""" effective_token_limit: int | None = None - tokenizer_fn: TokenizerFn | None = None + tokenizer_fn: TokenizerFn | AsyncTokenizerFn | None = None iteration: int = 0 max_iterations: int = 40 - cache: "ChatLoopRuntimeCache" = Field( - default_factory=lambda: ChatLoopRuntimeCache() - ) + cache: "ChatRuntimeCache" = Field(default_factory=lambda: ChatRuntimeCache()) + next_block_count: int = 0 + total_input_tokens: int = 0 + total_output_tokens: int = 0 + has_input_usage: bool = False + has_output_usage: bool = False -class ChatLoopOutputState(BaseModel): +class ChatStatus(StrEnum): + RUNNING = "running" + WAITING = "waiting" + COMPLETED = "completed" + CONTINUE = "continue" + + +class ChatOutputState(BaseModel): """Store loop outputs and pending external handoffs.""" stop_reason: str | None = None pending_external_tool_calls: list[ToolSelection] = Field(default_factory=list) + status: ChatStatus = ChatStatus.RUNNING + pending_async_tools: dict[str, str] = Field(default_factory=dict) + pause_type: str = "after_tool" -class ChatLoopTimelineEntry(BaseModel): +class ChatTimelineEntry(BaseModel): """Capture one immutable timeline snapshot for debugging.""" iteration: int @@ -76,21 +90,21 @@ class SkillsRuntimeCache(BaseModel): ) -class ChatLoopRuntimeCache(BaseModel): +class ChatRuntimeCache(BaseModel): """Runtime cache buckets for interceptors.""" skill: SkillsRuntimeCache | None = None -class ChatLoopState(BaseModel): +class ChatState(BaseModel): """Aggregate clonable loop state sections and history timeline.""" - input: ChatLoopInputState - runtime: ChatLoopRuntimeState - output: ChatLoopOutputState + input: ChatInputState + runtime: ChatRuntimeState + output: ChatOutputState - original_input: ChatLoopInputState | None = Field(default=None) - timeline: list[ChatLoopTimelineEntry] = Field(default_factory=list) + original_input: ChatInputState | None = Field(default=None) + timeline: list[ChatTimelineEntry] = Field(default_factory=list) model_config = ConfigDict(arbitrary_types_allowed=True) diff --git a/private_gpt/components/engines/chat/models/chat_phase.py b/private_gpt/components/engines/chat/models/chat_phase.py new file mode 100644 index 00000000..71662654 --- /dev/null +++ b/private_gpt/components/engines/chat/models/chat_phase.py @@ -0,0 +1,24 @@ +from enum import StrEnum + + +class InterceptorPhase(StrEnum): + """Represent interceptor execution phase in the loop.""" + + VALIDATION = "validation" + BEFORE_ITERATION = "before_iteration" + BEFORE_TOOL = "before_tool" + STREAMING = "streaming" + AFTER_TOOL = "after_tool" + AFTER_ITERATION = "after_iteration" + + +class TimelinePhase(StrEnum): + """Represent timeline checkpoints captured by the loop state.""" + + START = "start" + BEFORE_INTERCEPTORS = "before_interceptors" + AFTER_INTERCEPTORS = "after_interceptors" + BEFORE_LLM = "before_llm" + AFTER_LLM = "after_llm" + AFTER_TOOLS = "after_tools" + STOP = "stop" diff --git a/private_gpt/components/engines/chat/models/chat_state.py b/private_gpt/components/engines/chat/models/chat_state.py new file mode 100644 index 00000000..f484d44f --- /dev/null +++ b/private_gpt/components/engines/chat/models/chat_state.py @@ -0,0 +1,138 @@ +from collections.abc import Mapping +from enum import StrEnum +from typing import Any, Self + +from llama_index.core.llms.llm import ToolSelection +from pydantic import BaseModel, ConfigDict, Field + +from private_gpt.components.chat.models.chat_config_models import ChatRequest +from private_gpt.components.context.models.context_stack import ContextStack +from private_gpt.components.engines.chat.models.chat_phase import ( + TimelinePhase, +) +from private_gpt.components.llm.llm_helper import AsyncTokenizerFn, TokenizerFn +from private_gpt.components.skills.models.skill_entities import ( + SkillVersionWithSkillEntity, +) + + +class ChatInputState(BaseModel): + """Request-scoped inputs for one loop run. + + ``request`` — the original immutable ChatRequest (restored each iteration). + ``context_stack`` — the working stack for the current iteration, built by + ``build_initial_context_stack`` at loop start then enriched by + interceptors (skills, MCP tools, RAG, …). + + Interceptors append layers to ``context_stack`` via ``stack.append_layer()``. + The engine materializes ``request`` from the stack just before calling the LLM. + """ + + request: ChatRequest + context_stack: ContextStack = Field(default_factory=ContextStack) + sampling_params: dict[str, Any] = Field(default_factory=dict) + llm_kwargs: dict[str, Any] = Field(default_factory=dict) + + model_config = ConfigDict(arbitrary_types_allowed=True) + + +class ChatRuntimeState(BaseModel): + """Store runtime counters.""" + + effective_token_limit: int | None = None + tokenizer_fn: TokenizerFn | AsyncTokenizerFn | None = None + + iteration: int = 0 + max_iterations: int = 40 + cache: "ChatRuntimeCache" = Field(default_factory=lambda: ChatRuntimeCache()) + next_block_count: int = 0 + total_input_tokens: int = 0 + total_output_tokens: int = 0 + has_input_usage: bool = False + has_output_usage: bool = False + + +class ChatStatus(StrEnum): + RUNNING = "running" + WAITING = "waiting" + COMPLETED = "completed" + CONTINUE = "continue" + + +class ChatOutputState(BaseModel): + """Store loop outputs and pending external handoffs.""" + + stop_reason: str | None = None + pending_external_tool_calls: list[ToolSelection] = Field(default_factory=list) + status: ChatStatus = ChatStatus.RUNNING + pending_async_tools: dict[str, str] = Field(default_factory=dict) + pause_type: str = "after_tool" + + +class ChatTimelineEntry(BaseModel): + """Capture one immutable timeline snapshot for debugging.""" + + iteration: int + phase: TimelinePhase + conversation_size: int + tool_count: int + stop_reason: str | None = None + + +class SkillsRuntimeCache(BaseModel): + """Validated/resolved skills cached during runtime.""" + + entries: list[SkillVersionWithSkillEntity] = Field(default_factory=list) + resources: dict[str, list[str]] = Field( + default_factory=dict, + description="skill_id → bundled file paths relative to the skill dir " + "(SKILL.md excluded).", + ) + + +class ChatRuntimeCache(BaseModel): + """Runtime cache buckets for interceptors.""" + + skill: SkillsRuntimeCache | None = None + + +class ChatState(BaseModel): + """Aggregate clonable loop state sections and history timeline.""" + + input: ChatInputState + runtime: ChatRuntimeState + output: ChatOutputState + + original_input: ChatInputState | None = Field(default=None) + timeline: list[ChatTimelineEntry] = Field(default_factory=list) + + model_config = ConfigDict(arbitrary_types_allowed=True) + + def model_copy( + self, *, update: Mapping[str, Any] | None = None, deep: bool = False + ) -> Self: + + # Temporarily detach fields that must not be deep-copied: + # - tokenizer_fn: HF tokenizer is not safely copyable and expensive + # - context_stack: pydantic frozen — all mutations return new instances + # - original_input: set once at loop init, never mutated + tokenizer_fn = self.runtime.tokenizer_fn + context_stack = self.input.context_stack + original_input = self.original_input + + self.runtime.tokenizer_fn = None + self.input.context_stack = ContextStack() + self.original_input = None + + try: + copied = super().model_copy(update=update, deep=deep) + finally: + self.runtime.tokenizer_fn = tokenizer_fn + self.input.context_stack = context_stack + self.original_input = original_input + + copied.runtime.tokenizer_fn = tokenizer_fn + copied.input.context_stack = context_stack + copied.original_input = original_input + + return copied diff --git a/private_gpt/components/engines/chat/models/execution_hooks.py b/private_gpt/components/engines/chat/models/execution_hooks.py new file mode 100644 index 00000000..af924f53 --- /dev/null +++ b/private_gpt/components/engines/chat/models/execution_hooks.py @@ -0,0 +1,12 @@ +from typing import Any + +from pydantic import BaseModel, Field + + +class ToolExecutionHook(BaseModel): + callable_path: str + kwargs: dict[str, Any] = Field(default_factory=dict) + + +class ExecutionHooks(BaseModel): + tool_result: list[ToolExecutionHook] = Field(default_factory=list) diff --git a/private_gpt/components/engines/chat/resumable_runner.py b/private_gpt/components/engines/chat/resumable_runner.py new file mode 100644 index 00000000..11bee659 --- /dev/null +++ b/private_gpt/components/engines/chat/resumable_runner.py @@ -0,0 +1,309 @@ +from __future__ import annotations + +import asyncio +from datetime import UTC, datetime, timedelta +from typing import TYPE_CHECKING, Annotated, Any +from uuid import uuid4 + +from injector import inject, singleton +from llama_index.core.tools import ToolSelection +from pydantic import Field, TypeAdapter, ValidationError + +from private_gpt.components.chat.models.chat_config_models import ResolvedChatRequest +from private_gpt.components.engines.chat.checkpoint_store import ( + ChatCheckpoint, + ChatCheckpointStoreFactory, +) +from private_gpt.components.engines.chat.event_broker import EngineEventBrokerFactory +from private_gpt.components.engines.chat.event_channel import BrokerEventChannel +from private_gpt.components.engines.chat.execution_scheduler import ( + ChatExecutionSchedulerFactory, +) +from private_gpt.components.engines.chat.models.chat_state import ChatStatus +from private_gpt.components.engines.chat.models.execution_hooks import ( + ExecutionHooks, + ToolExecutionHook, +) +from private_gpt.components.tools.tool_scheduler import ToolSchedulerFactory +from private_gpt.events.event_serializer import StreamingEventHandler +from private_gpt.events.models import ContentBlockType +from private_gpt.settings.settings import Settings + +if TYPE_CHECKING: + from collections.abc import AsyncGenerator + + 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 + + +_RESUME_HOOKS = ExecutionHooks( + tool_result=[ + ToolExecutionHook( + callable_path="private_gpt.arq.tasks.chat.callback:resume_chat_callback" + ) + ] +) + +_CONTENT_BLOCK_ADAPTER: TypeAdapter[ContentBlockType] = TypeAdapter( + Annotated[ContentBlockType, Field(discriminator="type")] +) + + +@singleton +class ResumableChatRunner: + """Single start/resume/timeout implementation for local and ARQ execution.""" + + @inject + def __init__( + self, + settings: Settings, + checkpoint_store_factory: ChatCheckpointStoreFactory, + event_broker_factory: EngineEventBrokerFactory, + scheduler_factory: ChatExecutionSchedulerFactory, + tool_scheduler_factory: ToolSchedulerFactory, + ) -> None: + self._settings = settings + self._state = checkpoint_store_factory.get() + self._events = event_broker_factory.get() + self._scheduler = scheduler_factory.get() + self._tool_scheduler = tool_scheduler_factory.get() + + async def submit( + self, + *, + request_data: dict[str, Any], + stream_type: str, + metadata: dict[str, Any], + execution_id: str | None = None, + ) -> tuple[str, AsyncGenerator[Event, None]]: + execution_id = execution_id or str(uuid4()) + events = self._events.listen(execution_id) + await self._scheduler.start( + execution_id=execution_id, + request_data=request_data, + stream_type=stream_type, + metadata=metadata, + ) + return execution_id, events + + async def cancel(self, execution_id: str) -> bool: + checkpoint = await self._state.load(execution_id) + tool_task_ids = ( + list(checkpoint.checkpoint_payload.pending_async_tools.values()) + if checkpoint is not None + else [] + ) + tool_cancellations = await asyncio.gather( + *( + self._tool_scheduler.cancel_task(task_id=task_id) + for task_id in tool_task_ids + ) + ) + chat_cancelled = await self._scheduler.cancel(execution_id) + await self._state.cleanup(execution_id) + await self._events.finish(execution_id) + return chat_cancelled or any(tool_cancellations) + + async def start( + self, + *, + engine: AsyncChatEngine, + execution_id: str, + request_data: dict[str, Any], + stream_type: str, + metadata: dict[str, Any], + ) -> None: + channel = BrokerEventChannel(self._events, execution_id) + try: + request = self._request(request_data) + state = await engine.execute(request, hooks=_RESUME_HOOKS, channel=channel) + await channel.close() + await self._handle_state( + execution_id=execution_id, + state=state, + stream_type=stream_type, + metadata=metadata, + ) + except Exception as exc: + await self._fail(execution_id, exc, channel) + raise + + async def resume(self, *, engine: AsyncChatEngine, execution_id: str) -> None: + saved = await self._state.load(execution_id) + if saved is None: + return + channel = BrokerEventChannel(self._events, execution_id) + try: + 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", [])), + *( + response.tool_message.model_dump(mode="json") + for response in responses + ), + ] + state = await engine.resume( + saved.checkpoint, + self._request(request_data), + iteration=saved.iteration, + next_block_count=saved.next_block_count, + hooks=_RESUME_HOOKS, + checkpoint_payload=saved.checkpoint_payload.model_copy( + update={"tool_responses": responses} + ), + channel=channel, + ) + await channel.close() + await self._state.cleanup(execution_id) + await self._handle_state( + execution_id=execution_id, + state=state, + stream_type=saved.stream_type, + metadata=saved.metadata, + ) + except Exception as exc: + await self._fail(execution_id, exc, channel) + raise + + async def callback( + self, *, execution_id: str, tool_id: str, result: dict[str, Any] + ) -> None: + ready = await self._state.record_result(execution_id, tool_id, result) + if ready is None: + return + await self._resume_if_ready(execution_id) + + async def timeout(self, *, execution_id: str, checkpoint_id: str) -> None: + saved = await self._state.load(execution_id) + if saved is None or saved.checkpoint_id != checkpoint_id: + return + if not await self._state.claim_resume(execution_id): + return + await self._fail( + execution_id, + TimeoutError( + f"Chat callback timed out after {self._settings.scheduler.chat.callback_timeout_seconds} seconds" + ), + ) + + async def _handle_state( + self, + *, + execution_id: str, + state: ChatState, + stream_type: str, + metadata: dict[str, Any], + ) -> None: + if state.output.status != ChatStatus.WAITING: + await self._state.cleanup(execution_id) + await self._events.finish(execution_id) + return + + checkpoint_id = uuid4().hex + timeout_seconds = self._settings.scheduler.chat.callback_timeout_seconds + await self._state.save( + ChatCheckpoint( + correlation_id=execution_id, + request_data=state.input.request.model_dump(mode="json"), + stream_type=stream_type, + metadata=metadata, + iteration=state.runtime.iteration, + checkpoint=state.output.pause_type, + checkpoint_payload=self._checkpoint_payload(state), + next_block_count=state.runtime.next_block_count, + checkpoint_id=checkpoint_id, + deadline=datetime.now(UTC) + timedelta(seconds=timeout_seconds), + ) + ) + await self._resume_if_ready(execution_id) + await self._scheduler.timeout( + execution_id=execution_id, + checkpoint_id=checkpoint_id, + delay_seconds=timeout_seconds, + ) + + async def _resume_if_ready(self, execution_id: str) -> None: + checkpoint = await self._state.load(execution_id) + if checkpoint is None: + return + expected = set(checkpoint.checkpoint_payload.pending_async_tools) + results = await self._state.get_results(execution_id) + if not expected or not expected.issubset(results): + return + if await self._state.claim_resume(execution_id): + await self._scheduler.resume(execution_id=execution_id) + + async def _fail( + self, + execution_id: str, + exc: Exception, + channel: BrokerEventChannel | None = None, + ) -> None: + event = StreamingEventHandler().error_event(execution_id, exc) + if channel is None: + await self._events.publish(execution_id, event) + else: + channel.emit(event) + await channel.close() + 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) + for message in request.messages: + tool_calls = message.additional_kwargs.get("tool_calls") + if isinstance(tool_calls, list): + message.additional_kwargs["tool_calls"] = [ + ToolSelection.model_validate(tool_call) + if isinstance(tool_call, dict) + else tool_call + for tool_call in tool_calls + ] + + message.additional_kwargs = { + key: ResumableChatRunner._restore_content_blocks(value) + for key, value in message.additional_kwargs.items() + } + return request + + @staticmethod + def _restore_content_blocks(value: Any) -> Any: + if isinstance(value, list): + return [ResumableChatRunner._restore_content_blocks(item) for item in value] + if not isinstance(value, dict) or not isinstance(value.get("type"), str): + return value + try: + return _CONTENT_BLOCK_ADAPTER.validate_python(value) + except ValidationError: + return value + + @staticmethod + def _checkpoint_payload(state: ChatState) -> Any: + from private_gpt.components.engines.chat.async_chat_engine import ( + IterationCheckpointPayload, + ) + + return IterationCheckpointPayload( + pending_async_tools=state.output.pending_async_tools, + pending_external_tool_calls=state.output.pending_external_tool_calls, + total_input_tokens=state.runtime.total_input_tokens, + total_output_tokens=state.runtime.total_output_tokens, + has_input_usage=state.runtime.has_input_usage, + has_output_usage=state.runtime.has_output_usage, + ) diff --git a/private_gpt/components/engines/chat_loop/utils/__init__.py b/private_gpt/components/engines/chat/schedulers/__init__.py similarity index 100% rename from private_gpt/components/engines/chat_loop/utils/__init__.py rename to private_gpt/components/engines/chat/schedulers/__init__.py diff --git a/private_gpt/components/engines/chat_loop/types.py b/private_gpt/components/engines/chat/types.py similarity index 100% rename from private_gpt/components/engines/chat_loop/types.py rename to private_gpt/components/engines/chat/types.py diff --git a/private_gpt/components/engines/chat/utils/__init__.py b/private_gpt/components/engines/chat/utils/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/private_gpt/components/engines/chat_loop/utils/message_utils.py b/private_gpt/components/engines/chat/utils/message_utils.py similarity index 100% rename from private_gpt/components/engines/chat_loop/utils/message_utils.py rename to private_gpt/components/engines/chat/utils/message_utils.py diff --git a/private_gpt/components/engines/chat_loop/utils/request_builder.py b/private_gpt/components/engines/chat/utils/request_builder.py similarity index 100% rename from private_gpt/components/engines/chat_loop/utils/request_builder.py rename to private_gpt/components/engines/chat/utils/request_builder.py diff --git a/private_gpt/components/engines/chat_loop/utils/tool_utils.py b/private_gpt/components/engines/chat/utils/tool_utils.py similarity index 100% rename from private_gpt/components/engines/chat_loop/utils/tool_utils.py rename to private_gpt/components/engines/chat/utils/tool_utils.py diff --git a/private_gpt/components/ingestion/__init__.py b/private_gpt/components/ingestion/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/private_gpt/components/ingestion/ingestion_scheduler.py b/private_gpt/components/ingestion/ingestion_scheduler.py new file mode 100644 index 00000000..6ffd2614 --- /dev/null +++ b/private_gpt/components/ingestion/ingestion_scheduler.py @@ -0,0 +1,278 @@ +from __future__ import annotations + +import logging +from abc import ABC, abstractmethod +from collections.abc import Callable +from typing import TYPE_CHECKING + +from injector import Injector, inject, singleton + +from private_gpt.components.ingest.utils import get_extension, get_file_name +from private_gpt.components.storage.s3_helper import S3Helper +from private_gpt.server.ingest.ingest_service import IngestService +from private_gpt.settings.settings import Settings, settings + +if TYPE_CHECKING: + from private_gpt.server.ingest.ingest_router import ( + DeleteIngestedDocumentAsyncBody, + IngestAsyncBody, + IngestBody, + IngestResponse, + ) + +logger = logging.getLogger(__name__) + +IngestionSchedulerProvider = ( + type["BaseIngestionScheduler"] | Callable[[Injector], "BaseIngestionScheduler"] +) +_INGESTION_SCHEDULERS: dict[str, IngestionSchedulerProvider] = {} + + +def register_ingestion_scheduler( + mode: str, provider: IngestionSchedulerProvider +) -> None: + _INGESTION_SCHEDULERS[mode] = provider + + +class BaseIngestionScheduler(ABC): + """Shared ingestion logic between local and Celery modes.""" + + @property + def is_async(self) -> bool: + return False + + def __init__( + self, + ingest_service: IngestService, + s3_helper: S3Helper | None = None, + ) -> None: + self._ingest_service = ingest_service + self._s3_helper = s3_helper + + def _require_s3_helper(self) -> S3Helper: + if self._s3_helper is None: + raise RuntimeError("S3Helper is required for remote ingestion flows") + return self._s3_helper + + @abstractmethod + def ingest(self, ingest_body: IngestBody) -> IngestResponse: + ... + + def ingest_async(self, ingest_body: IngestAsyncBody) -> str: + del ingest_body + raise NotImplementedError + + def delete(self, collection: str, artifact: str) -> None: + self._ingest_service.delete(collection=collection, artifact=artifact) + + def delete_async(self, delete_body: DeleteIngestedDocumentAsyncBody) -> str: + del delete_body + raise NotImplementedError + + +@singleton +class LocalIngestionScheduler(BaseIngestionScheduler): + """Run sync ingestion in-process.""" + + @inject + def __init__(self, ingest_service: IngestService) -> None: + super().__init__(ingest_service) + + def ingest(self, ingest_body: IngestBody) -> IngestResponse: + from private_gpt.server.ingest.ingest_router import ingest_data_sync + + content = ingest_body.input.to_binary_content( + get_file_name(ingest_body.metadata) + ) + with self._ingest_service.temporary_file( + lambda: self._ingest_service.data_path_from_bin_data( + content.data, get_extension(content.filename) + ) + ) as data_path: + return ingest_data_sync( + collection=ingest_body.collection, + artifact=ingest_body.artifact, + data_path=data_path, + service=self._ingest_service, + metadata={ + **(ingest_body.metadata or {}), + "file_name": content.filename, + }, + ) + + +@singleton +class CeleryIngestionScheduler(BaseIngestionScheduler): + """Dispatch sync ingestion to a Celery worker and wait.""" + + @inject + def __init__( + self, settings: Settings, s3_helper: S3Helper, ingest_service: IngestService + ) -> None: + super().__init__(ingest_service, s3_helper) + self._settings = settings + + @property + def is_async(self) -> bool: + return True + + def ingest_async(self, ingest_body: IngestAsyncBody) -> str: + import uuid + + from private_gpt.celery.dispatch import dispatch_task + from private_gpt.celery.tasks.ingestion.extraction_tasks import ( + VECTOR_INDEX_TASK_NAME, + ) + from private_gpt.server.utils.artifact_input import UriArtifact + + config = settings() + self._ingest_service.initialize_artifact_indices( + collection=ingest_body.ingest_body.collection, + artifact=ingest_body.ingest_body.artifact, + ) + if ingest_body.ingest_body.input: + should_upload = ( + not isinstance(ingest_body.ingest_body.input, UriArtifact) + or ingest_body.ingest_body.input.is_base64() + ) + s3_helper = self._require_s3_helper() + if should_upload and s3_helper.is_available(): + content = ingest_body.ingest_body.input.to_binary_content() + object_name = str(uuid.uuid4()) + s3_url = s3_helper.upload_file_to_s3( + filename=content.filename, + bytes_data=content.data.read(), + bucket_name=config.s3.temporary_bucket_name, + object_name=object_name, + ) + if not s3_url: + raise ValueError("Failed to upload file to S3 for async ingestion") + ingest_body.ingest_body.input = UriArtifact(value=s3_url) + + result = dispatch_task( + task_name=VECTOR_INDEX_TASK_NAME, + args=(ingest_body,), + queue=config.scheduler.ingestion.celery_queue, + ) + task_id = result.task_id + if not isinstance(task_id, str): + raise ValueError("Async ingestion task did not return a valid task_id") + return task_id + + def delete_async(self, delete_body: DeleteIngestedDocumentAsyncBody) -> str: + from private_gpt.celery.celery import celery_app + from private_gpt.celery.dispatch import dispatch_task + from private_gpt.celery.task_helper import IngestionTaskHelper + from private_gpt.celery.tasks.ingestion.delete_tasks import ( + DELETE_INGESTED_TASK_NAME, + ) + + revoked = IngestionTaskHelper.revoke_ingestion_task( + celery_app=celery_app, + collection=delete_body.delete_body.collection, + artifact=delete_body.delete_body.artifact, + ) + if revoked: + return "revoked" + + config = settings() + result = dispatch_task( + task_name=DELETE_INGESTED_TASK_NAME, + args=(delete_body,), + queue=config.scheduler.ingestion.celery_queue, + ) + task_id = result.task_id + if not isinstance(task_id, str): + raise ValueError("Delete ingestion task did not return a valid task_id") + return task_id + + def ingest(self, ingest_body: IngestBody) -> IngestResponse: + import uuid + + from private_gpt.celery.dispatch import dispatch_task + from private_gpt.celery.tasks.ingestion.extraction_tasks import ( + VECTOR_INDEX_TASK_NAME, + ) + from private_gpt.server.ingest.ingest_router import ( + IngestAsyncBody, + IngestResponse, + ) + from private_gpt.server.utils.artifact_input import UriArtifact + from private_gpt.server.utils.callback import AMQP, Callback + + config = settings() + self._ingest_service.initialize_artifact_indices( + collection=ingest_body.collection, + artifact=ingest_body.artifact, + ) + async_body = IngestAsyncBody( + ingest_body=ingest_body, + callback=Callback( + amqp=AMQP( + exchange="ingestion_events", + routing_key_done="ingest.completed", + routing_key_error="ingest.failed", + routing_key_progress="ingest.progress", + ) + ), + ) + + if async_body.ingest_body.input: + s3_helper = self._require_s3_helper() + if s3_helper.is_available(): + content = async_body.ingest_body.input.to_binary_content( + async_body.ingest_body.metadata.get("file_name") + if async_body.ingest_body.metadata + else None + ) + object_name = str(uuid.uuid4()) + s3_url = s3_helper.upload_file_to_s3( + filename=content.filename, + bytes_data=content.data.read(), + bucket_name=config.s3.temporary_bucket_name, + object_name=object_name, + ) + if not s3_url: + raise ValueError( + "Failed to upload file to S3 for sync celery ingest" + ) + async_body.ingest_body.input = UriArtifact(value=s3_url) + + result = dispatch_task( + task_name=VECTOR_INDEX_TASK_NAME, + args=(async_body,), + queue=config.scheduler.ingestion.celery_queue, + ) + while not result.ready(): + import time + + time.sleep(0.1) + if result.failed(): + raise result.result + return IngestResponse.model_validate(result.result) + + +register_ingestion_scheduler("local", LocalIngestionScheduler) +register_ingestion_scheduler("celery", CeleryIngestionScheduler) + + +@singleton +class IngestionSchedulerFactory: + @inject + def __init__(self, settings: Settings, injector: Injector) -> None: + self._settings = settings + self._injector = injector + self._scheduler: BaseIngestionScheduler | None = None + + def get(self) -> BaseIngestionScheduler: + if self._scheduler is None: + mode = self._settings.scheduler.ingestion.mode + provider = _INGESTION_SCHEDULERS.get(mode) + if provider is None: + raise ValueError(f"Unknown scheduler.ingestion.mode: {mode}") + self._scheduler = ( + self._injector.get(provider) + if isinstance(provider, type) + else provider(self._injector) + ) + return self._scheduler diff --git a/private_gpt/components/llm/llm_component.py b/private_gpt/components/llm/llm_component.py index bfe41b03..be836677 100644 --- a/private_gpt/components/llm/llm_component.py +++ b/private_gpt/components/llm/llm_component.py @@ -9,9 +9,17 @@ from llama_index.core.llms import LLM from private_gpt.components.llm.custom.base import ZylonLLM from private_gpt.components.llm.discovery import get_models from private_gpt.components.llm.factories.registry import LLMFactoryRegistry -from private_gpt.components.llm.llm_helper import TokenizerFn, get_tokenizer_fn +from private_gpt.components.llm.llm_helper import ( + AsyncTokenizerFn, + TokenizerFn, + get_async_tokenizer_fn, + get_tokenizer_fn, +) from private_gpt.components.llm.registry import LLMInstance, LLMRegistry -from private_gpt.components.llm.tokenizers.tokenizer_base import TokenizerBase +from private_gpt.components.llm.tokenizers.tokenizer_base import ( + AsyncTokenizerBase, + TokenizerBase, +) from private_gpt.components.model_discovery.service import are_distinct_api_bases from private_gpt.settings.settings import LLMModelConfig, Settings @@ -175,7 +183,7 @@ class LLMComponent: def filter( self, predicate: Callable[[LLM, LLMModelConfig], bool] - ) -> Iterator[tuple[LLM, LLMModelConfig, TokenizerFn | None]]: + ) -> Iterator[tuple[LLM, LLMModelConfig, TokenizerFn | AsyncTokenizerFn | None]]: """Filter LLM components based on a predicate function. :param predicate: A function that takes an LLM @@ -192,11 +200,12 @@ class LLMComponent: ): seen.add(id(model_id)) - tokenizer_fn = ( - get_tokenizer_fn(component.tokenizer) - if component.tokenizer - else None - ) + tokenizer_fn: TokenizerFn | AsyncTokenizerFn | None = None + if component.tokenizer: + if isinstance(component.tokenizer, AsyncTokenizerBase): + tokenizer_fn = get_async_tokenizer_fn(component.tokenizer) + else: + tokenizer_fn = get_tokenizer_fn(component.tokenizer) yield component.llm, model_config, tokenizer_fn @property diff --git a/private_gpt/components/llm/llm_helper.py b/private_gpt/components/llm/llm_helper.py index 764ac041..a5d5e44b 100644 --- a/private_gpt/components/llm/llm_helper.py +++ b/private_gpt/components/llm/llm_helper.py @@ -1,10 +1,13 @@ -from collections.abc import Callable -from typing import Any +import asyncio +import concurrent.futures +from collections.abc import Awaitable, Callable +from typing import Any, cast from llama_index.core.llms import LLM from llama_index.core.multi_modal_llms import MultiModalLLMMetadata from private_gpt.components.llm.tokenizers.tokenizer_base import ( + AsyncTokenizerBase, AudioLike, ImageLike, TextLike, @@ -47,6 +50,7 @@ def max_audios_supported( TokenizerFn = Callable[..., TokenizedInput] +AsyncTokenizerFn = Callable[..., Awaitable[TokenizedInput]] def get_tokenizer_fn(tokenizer: TokenizerBase | None) -> TokenizerFn | None: @@ -65,6 +69,41 @@ def get_tokenizer_fn(tokenizer: TokenizerBase | None) -> TokenizerFn | None: return fn +def get_async_tokenizer_fn(tokenizer: AsyncTokenizerBase) -> AsyncTokenizerFn: + async def fn( + texts: TextLike | None = None, + images: ImageLike | None = None, + audios: AudioLike | None = None, + **kwargs: Any, + ) -> TokenizedInput: + return await tokenizer.acall( + texts=texts, images=images, audios=audios, **kwargs + ) + + return fn + + +def as_sync_tokenizer_fn( + tokenizer_fn: TokenizerFn | AsyncTokenizerFn | None, +) -> TokenizerFn | None: + """Wrap an AsyncTokenizerFn so it can be called synchronously. + + Runs the coroutine in a fresh event loop on a dedicated thread, avoiding + conflicts with any running loop in the caller's thread. + """ + if tokenizer_fn is None: + return None + if not asyncio.iscoroutinefunction(tokenizer_fn): + return cast(TokenizerFn, tokenizer_fn) + _async_fn = tokenizer_fn + + def _sync_wrapper(*args: Any, **kwargs: Any) -> Any: + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as ex: + return ex.submit(asyncio.run, _async_fn(*args, **kwargs)).result() + + return _sync_wrapper + + def get_tokenizer() -> TokenizerFn | None: from private_gpt.components.llm.llm_component import LLMComponent diff --git a/private_gpt/components/llm/tokenizers/remote.py b/private_gpt/components/llm/tokenizers/remote.py index bcd6203b..620d6a6a 100644 --- a/private_gpt/components/llm/tokenizers/remote.py +++ b/private_gpt/components/llm/tokenizers/remote.py @@ -4,15 +4,15 @@ from typing import Any import httpx from private_gpt.components.llm.tokenizers.tokenizer_base import ( + AsyncTokenizerBase, AudioLike, ImageLike, TextLike, TokenizedInput, - TokenizerBase, ) -class RemoteTokenizeTokenizer(TokenizerBase): +class RemoteTokenizeTokenizer(AsyncTokenizerBase): """Tokenizer backed by remote `/tokenize` and `/detokenize` endpoints. Supports the current wire formats exposed by: @@ -167,6 +167,63 @@ class RemoteTokenizeTokenizer(TokenizerBase): "RemoteTokenizeTokenizer does not expose piecewise token strings" ) + async def acall( + self, + texts: TextLike | None = None, + images: ImageLike | None = None, + audios: AudioLike | None = None, + add_special_tokens: bool = True, + truncation: bool = False, + max_length: int | None = None, + **kwargs: Any, + ) -> TokenizedInput: + del add_special_tokens, truncation, max_length, kwargs + + if images or audios: + raise NotImplementedError( + "RemoteTokenizeTokenizer only supports text tokenization" + ) + if texts is None: + return TokenizedInput(input_ids=[]) + + if isinstance(texts, str): + return TokenizedInput(input_ids=await self.aencode(texts)) + + if isinstance(texts, Sequence): + input_ids: list[int] = [] + for text in texts: + input_ids.extend(await self.aencode(str(text))) + return TokenizedInput(input_ids=input_ids) + + return TokenizedInput(input_ids=await self.aencode(str(texts))) + + async def aencode( + self, text: str, add_special_tokens: bool | None = None + ) -> list[int]: + del add_special_tokens + response = await self._apost_tokenize(text) + return self._extract_tokens(response) + + async def adecode( + self, ids: list[int] | int, skip_special_tokens: bool = True + ) -> str: + del skip_special_tokens + token_ids = [ids] if isinstance(ids, int) else ids + response = await self._apost_detokenize(token_ids) + return self._extract_text(response) + + async def aapply_chat_template( + self, + conversation: list[dict[str, str | list[dict[str, str]]]], + tools: list[dict[str, Any]] | None = None, + documents: list[dict[str, str]] | None = None, + **kwargs: Any, + ) -> list[int] | str: + del conversation, tools, documents, kwargs + raise NotImplementedError( + "RemoteTokenizeTokenizer cannot render chat templates locally" + ) + def _post_tokenize(self, text: str) -> Any: tokenize_url = self._build_url(self.TOKENIZE_ENDPOINT) last_error: Exception | None = None @@ -215,6 +272,56 @@ class RemoteTokenizeTokenizer(TokenizerBase): raise ValueError("Remote detokenization failed without an HTTP error") raise last_error + async def _apost_tokenize(self, text: str) -> Any: + tokenize_url = self._build_url(self.TOKENIZE_ENDPOINT) + last_error: Exception | None = None + + async with httpx.AsyncClient() as client: + for payload in ( + {"model": self.model_id, "prompt": text}, + {"content": text}, + ): + try: + response = await client.post( + tokenize_url, + json=payload, + headers=self._headers(), + timeout=self.request_timeout, + ) + response.raise_for_status() + return response.json() + except httpx.HTTPError as e: + last_error = e + + if last_error is None: + raise ValueError("Remote tokenization failed without an HTTP error") + raise last_error + + async def _apost_detokenize(self, token_ids: list[int]) -> Any: + detokenize_url = self._build_url(self.DETOKENIZE_ENDPOINT) + last_error: Exception | None = None + + async with httpx.AsyncClient() as client: + for payload in ( + {"model": self.model_id, "tokens": token_ids}, + {"tokens": token_ids}, + ): + try: + response = await client.post( + detokenize_url, + json=payload, + headers=self._headers(), + timeout=self.request_timeout, + ) + response.raise_for_status() + return response.json() + except httpx.HTTPError as e: + last_error = e + + if last_error is None: + raise ValueError("Remote detokenization failed without an HTTP error") + raise last_error + def _build_url(self, path: str) -> str: return f"{self.api_base}/{path.lstrip('/')}" diff --git a/private_gpt/components/llm/tokenizers/tokenizer_base.py b/private_gpt/components/llm/tokenizers/tokenizer_base.py index 304142bc..a6b98281 100644 --- a/private_gpt/components/llm/tokenizers/tokenizer_base.py +++ b/private_gpt/components/llm/tokenizers/tokenizer_base.py @@ -121,3 +121,42 @@ class TokenizerBase(ABC): skip_special_tokens: bool = True, ) -> list[str]: raise NotImplementedError() + + +class AsyncTokenizerBase(TokenizerBase, ABC): + """TokenizerBase extension for natively-async tokenizer backends.""" + + @abstractmethod + async def acall( + self, + texts: TextLike | None = None, + images: ImageLike | None = None, + audios: AudioLike | None = None, + add_special_tokens: bool = True, + truncation: bool = False, + max_length: int | None = None, + **kwargs: Any, + ) -> "TokenizedInput": + raise NotImplementedError() + + @abstractmethod + async def aencode( + self, text: str, add_special_tokens: bool | None = None + ) -> list[int]: + raise NotImplementedError() + + @abstractmethod + async def adecode( + self, ids: list[int] | int, skip_special_tokens: bool = True + ) -> str: + raise NotImplementedError() + + @abstractmethod + async def aapply_chat_template( + self, + conversation: list[dict[str, str | list[dict[str, str]]]], + tools: list[dict[str, Any]] | None = None, + documents: list[dict[str, str]] | None = None, + **kwargs: Any, + ) -> list[int] | str: + raise NotImplementedError() diff --git a/private_gpt/components/memory/trimming_memory.py b/private_gpt/components/memory/trimming_memory.py index bf0f1a1a..d981e010 100644 --- a/private_gpt/components/memory/trimming_memory.py +++ b/private_gpt/components/memory/trimming_memory.py @@ -1,6 +1,6 @@ import enum import json -from collections.abc import Awaitable, Callable +from collections.abc import Callable from typing import Any, Literal from llama_index.core.base.llms.types import ChatMessage, MessageRole @@ -12,6 +12,9 @@ from llama_index.core.memory.types import ( ) from llama_index.core.storage.chat_store import BaseChatStore, SimpleChatStore +from private_gpt.components.llm.llm_helper import TokenizerFn +from private_gpt.utils.tokens import async_tokenizer + DEFAULT_TOKEN_LIMIT_RATIO = 0.75 DEFAULT_TOKEN_LIMIT = 3000 @@ -58,7 +61,7 @@ class TrimmingMemory(BaseChatStoreMemory): default_factory=lambda: _default_text_splitter, exclude=True, ) - tokenizer_fn: Callable[[str], Awaitable[list[int]]] = Field( + tokenizer_fn: TokenizerFn = Field( exclude=True, ) @@ -113,7 +116,7 @@ class TrimmingMemory(BaseChatStoreMemory): allow_partial: bool = False, start_on: MessageRole | list[MessageRole] | None = None, end_on: MessageRole | list[MessageRole] | None = None, - tokenizer_fn: Callable[[str], Awaitable[list[int]]] | None = None, + tokenizer_fn: TokenizerFn | None = None, text_splitter: Callable[[str], list[str]] | None = None, **kwargs: Any, ) -> "TrimmingMemory": @@ -331,7 +334,7 @@ class TrimmingMemory(BaseChatStoreMemory): # Convert messages to string representation for token counting msg_str = " ".join(str(m.content) for m in messages) - return len(await self.tokenizer_fn(msg_str)) + return len(await async_tokenizer(msg_str, tokenizer_fn=self.tokenizer_fn)) def to_string(self) -> str: """Convert memory to string.""" diff --git a/private_gpt/components/storage/s3_helper.py b/private_gpt/components/storage/s3_helper.py index 815e8e59..8bbeaf78 100644 --- a/private_gpt/components/storage/s3_helper.py +++ b/private_gpt/components/storage/s3_helper.py @@ -28,13 +28,16 @@ class S3Helper: self._s3_settings = settings().s3 self._s3_client = self._get_s3_client(self._s3_settings) + def is_available(self) -> bool: + return self._s3_client is not None + @staticmethod def _get_s3_client(s3_settings: S3Settings) -> S3Client | None: if ( not s3_settings - or s3_settings.access_key_id is None - or s3_settings.secret_access_key is None - or s3_settings.endpoint_url is None + or not s3_settings.access_key_id + or not s3_settings.secret_access_key + or not s3_settings.endpoint_url ): return None diff --git a/private_gpt/components/streaming/providers/in_memory_stream_service.py b/private_gpt/components/streaming/providers/in_memory_stream_service.py index 2fd6e7dc..1e277627 100644 --- a/private_gpt/components/streaming/providers/in_memory_stream_service.py +++ b/private_gpt/components/streaming/providers/in_memory_stream_service.py @@ -24,7 +24,10 @@ class InMemoryStreamService(StreamService): str, list[tuple[str, str]] ] = {} # Like Redis stream: [(id, data)] self._event_counters: dict[str, int] = {} + self._cancel_flags: set[str] = set() self._lock = asyncio.Lock() + # Per-stream waiters: readers park here instead of sleeping + self._waiters: dict[str, set[asyncio.Event]] = defaultdict(set) async def create_stream( self, @@ -106,6 +109,8 @@ class InMemoryStreamService(StreamService): message_id = f"{int(datetime.now(UTC).timestamp() * 1000)}-{self._event_counters[correlation_id]}" self._events[correlation_id].append((message_id, event_data)) + for waiter in list(self._waiters.get(correlation_id, set())): + waiter.set() return message_id async def push_event_batch(self, events: list[Event]) -> dict[str, str]: @@ -161,11 +166,27 @@ class InMemoryStreamService(StreamService): return event_data, next_last_id + waiter: asyncio.Event | None = None async with self._lock: event_data, next_last_id = _read_available() + if not event_data and block_ms: + # Register waiter atomically with the "no events" check so we + # cannot miss a push_event that fires between check and wait. + waiter = asyncio.Event() + self._waiters[correlation_id].add(waiter) + + if waiter is not None: + assert block_ms is not None + try: + await asyncio.wait_for(waiter.wait(), timeout=block_ms / 1000) + except TimeoutError: + pass + finally: + async with self._lock: + self._waiters[correlation_id].discard(waiter) + async with self._lock: + event_data, next_last_id = _read_available() - if not event_data: - await asyncio.sleep(0.1) # simulate async context switch return event_data, next_last_id async def stream_exists(self, correlation_id: str) -> bool: @@ -184,6 +205,8 @@ class InMemoryStreamService(StreamService): del self._metadata[correlation_id] del self._events[correlation_id] del self._event_counters[correlation_id] + for waiter in self._waiters.pop(correlation_id, set()): + waiter.set() async def list_streams( self, @@ -213,10 +236,28 @@ class InMemoryStreamService(StreamService): async def clean_up_stream(self, correlation_id: str) -> None: """Clean up a stream by deleting it and its events.""" await self.delete_stream(correlation_id) + await self.clear_cancel_flag(correlation_id) + + async def set_cancel_flag(self, correlation_id: str) -> None: + """Set a cancellation flag for the stream.""" + self._cancel_flags.add(correlation_id) + + async def is_cancelled(self, correlation_id: str) -> bool: + """Check whether the cancellation flag has been set.""" + return correlation_id in self._cancel_flags + + async def clear_cancel_flag(self, correlation_id: str) -> None: + """Remove the cancellation flag.""" + self._cancel_flags.discard(correlation_id) async def close(self) -> None: """Close any open connections and clean up resources.""" async with self._lock: + for waiters in self._waiters.values(): + for waiter in waiters: + waiter.set() + self._waiters.clear() self._metadata.clear() self._events.clear() self._event_counters.clear() + self._cancel_flags.clear() diff --git a/private_gpt/components/streaming/providers/redis_stream_service.py b/private_gpt/components/streaming/providers/redis_stream_service.py index c52f505e..28654991 100644 --- a/private_gpt/components/streaming/providers/redis_stream_service.py +++ b/private_gpt/components/streaming/providers/redis_stream_service.py @@ -1,4 +1,5 @@ import asyncio +import contextlib import json import threading import uuid @@ -42,7 +43,12 @@ class RedisStreamService(StreamService): task = asyncio.create_task( self._pre_create_connections(redis_client, config.minimum_connections) ) - task.add_done_callback(lambda t: t.exception()) + task.add_done_callback(self._consume_background_result) + + @staticmethod + def _consume_background_result(task: asyncio.Task[None]) -> None: + with contextlib.suppress(asyncio.CancelledError, Exception): + task.exception() async def _pre_create_connections(self, client: redis.Redis, n: int) -> None: client.auto_close_connection_pool = False @@ -64,6 +70,10 @@ class RedisStreamService(StreamService): """Generate status key for correlation ID.""" return f"{self._config.status_prefix}:{correlation_id}" + def _get_cancel_key(self, correlation_id: str) -> str: + """Generate cancel-flag key for correlation ID.""" + return f"{self._config.status_prefix}:cancel:{correlation_id}" + async def create_stream( self, stream_type: str, @@ -303,6 +313,22 @@ class RedisStreamService(StreamService): async def clean_up_stream(self, correlation_id: str) -> None: """Clean up a stream by deleting it and its events.""" await self.delete_stream(correlation_id) + await self.clear_cancel_flag(correlation_id) + + async def set_cancel_flag(self, correlation_id: str) -> None: + """Set a cancellation flag in Redis for the chat worker to observe.""" + cancel_key = self._get_cancel_key(correlation_id) + await self._client.set(cancel_key, "1", ex=self._config.expiry_seconds) # type: ignore + + async def is_cancelled(self, correlation_id: str) -> bool: + """Check whether the cancellation flag has been set.""" + cancel_key = self._get_cancel_key(correlation_id) + return bool(await self._client.exists(cancel_key)) # type: ignore + + async def clear_cancel_flag(self, correlation_id: str) -> None: + """Remove the cancellation flag.""" + cancel_key = self._get_cancel_key(correlation_id) + await self._client.delete(cancel_key) # type: ignore async def close(self) -> None: """Close the redis stream service.""" diff --git a/private_gpt/components/streaming/providers/stream_service.py b/private_gpt/components/streaming/providers/stream_service.py index e7013f6f..7c6dadc3 100644 --- a/private_gpt/components/streaming/providers/stream_service.py +++ b/private_gpt/components/streaming/providers/stream_service.py @@ -94,6 +94,26 @@ class StreamService(ABC): """Clean up resources associated with a stream.""" pass + @abstractmethod + async def set_cancel_flag(self, correlation_id: str) -> None: + """Set a cancellation flag for a stream. + + Used by the API process to signal a long-running chat worker that the + stream should be cancelled. The worker polls :meth:`is_cancelled` inside + its event loop. + """ + pass + + @abstractmethod + async def is_cancelled(self, correlation_id: str) -> bool: + """Check whether a cancellation flag has been set for a stream.""" + pass + + @abstractmethod + async def clear_cancel_flag(self, correlation_id: str) -> None: + """Remove the cancellation flag for a stream (e.g. on cleanup).""" + pass + @abstractmethod async def close(self) -> None: """Close any open connections.""" diff --git a/private_gpt/components/streaming/stream/stream_manager.py b/private_gpt/components/streaming/stream/stream_manager.py index b17811cc..67bba093 100644 --- a/private_gpt/components/streaming/stream/stream_manager.py +++ b/private_gpt/components/streaming/stream/stream_manager.py @@ -30,9 +30,8 @@ class StreamManager: ): self.stream_service = stream_component.stream self.processor = stream_processor - - should_multiplexing = bool(settings.chat.multiplexing_threshold) - if should_multiplexing: + self.reader: StreamReader | AdaptiveStreamReader = stream_reader + if settings.chat.multiplexing_threshold: self.reader = AdaptiveStreamReader(settings, stream_reader) async def create_and_start_stream( @@ -60,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/components/streaming/stream/stream_processor.py b/private_gpt/components/streaming/stream/stream_processor.py index 9a742e49..ea66186d 100644 --- a/private_gpt/components/streaming/stream/stream_processor.py +++ b/private_gpt/components/streaming/stream/stream_processor.py @@ -31,6 +31,7 @@ class StreamProcessor: event_generator: AsyncGenerator[BaseModel, None], event_handler: EventHandler, metadata: dict[str, Any] | None = None, + mark_completed: bool = True, ) -> None: """Process a stream of events and push to Redis.""" try: @@ -68,10 +69,11 @@ class StreamProcessor: event_data=event_data, ) - await self.stream_service.update_stream_status( - correlation_id, - StreamStatus.COMPLETED, - ) + if mark_completed: + await self.stream_service.update_stream_status( + correlation_id, + StreamStatus.COMPLETED, + ) except asyncio.CancelledError: await self.stream_service.update_stream_status( @@ -120,6 +122,7 @@ class StreamProcessor: event_generator: AsyncGenerator[BaseModel, None], event_handler: EventHandler, metadata: dict[str, Any] | None = None, + mark_completed: bool = True, ) -> asyncio.Task[Any]: """Start processing a stream in the background.""" task: asyncio.Task[Any] = await self.task_manager.create_task( @@ -130,6 +133,7 @@ class StreamProcessor: event_generator=event_generator, event_handler=event_handler, metadata=metadata, + mark_completed=mark_completed, ), name=f"stream_processor_{correlation_id}", ) @@ -137,7 +141,7 @@ class StreamProcessor: async def cancel_stream_processing(self, correlation_id: str) -> bool: """Cancel stream processing.""" - success = await self.task_manager.cancel_task(correlation_id) + success = await self.task_manager.cancel(correlation_id) if success: await self.stream_service.update_stream_status( correlation_id, diff --git a/private_gpt/components/streaming/stream/stream_reader.py b/private_gpt/components/streaming/stream/stream_reader.py index d4bb0dd4..af9734af 100644 --- a/private_gpt/components/streaming/stream/stream_reader.py +++ b/private_gpt/components/streaming/stream/stream_reader.py @@ -3,7 +3,7 @@ import contextlib import logging from collections import defaultdict from collections.abc import AsyncGenerator -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final from injector import inject, singleton from pydantic import BaseModel @@ -18,9 +18,11 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) - DEFAULT_BLOCK_MS: Final[int] = 1000 BATCH_SIZE: Final[int] = 5 +# How often (seconds) to poll for terminal status when no events arrive. +# Halves Redis metadata round-trips vs. checking every block cycle (1 s). +STATUS_CHECK_INTERVAL: Final[float] = 2.0 TERMINAL_STATUSES: Final[set[StreamStatus]] = { StreamStatus.COMPLETED, StreamStatus.CANCELLED, @@ -28,41 +30,118 @@ TERMINAL_STATUSES: Final[set[StreamStatus]] = { } +class StreamBroadcast: + """Append-only batch log that broadcasts to N concurrent readers. + + A single producer calls ``push`` / ``close`` (both synchronous); each + subscriber holds an integer *cursor* into ``_batches`` and parks on a + ``asyncio.Future`` until new data arrives. One ``push`` resolves *all* + pending futures in a single pass — no per-subscriber queues. + + Safety: all operations are synchronous except ``read_from``, so no lock is + needed (asyncio cooperative threading guarantees no interleaving between + non-await statements). + """ + + def __init__(self) -> None: + self._batches: list[list[BaseModel] | None] = [] + self._waiters: list[asyncio.Future[None]] = [] + + def _notify(self) -> None: + waiters, self._waiters = self._waiters, [] + for fut in waiters: + if not fut.done(): + fut.set_result(None) + + def push(self, batch: list[BaseModel]) -> None: + self._batches.append(batch) + self._notify() + + def close(self) -> None: + """Append the terminal sentinel and wake all parked readers.""" + self._batches.append(None) + self._notify() + + def next_cursor(self) -> int: + """Cursor for a subscriber joining right now (no backfill).""" + return len(self._batches) + + @property + def is_closed(self) -> bool: + return bool(self._batches) and self._batches[-1] is None + + async def read_from( + self, + cursor: int, + stop_event: asyncio.Event | None = None, + ) -> AsyncGenerator[BaseModel, None]: + """Yield events starting at *cursor*, blocking until new ones arrive.""" + while True: + # Drain already-available batches. Safe without a lock: _batches is + # append-only and list operations are atomic between await points. + while cursor < len(self._batches): + batch = self._batches[cursor] + cursor += 1 + if batch is None: + return + for event in batch: + yield event + + if stop_event and stop_event.is_set(): + return + + # Park until the producer pushes more data. + fut: asyncio.Future[None] = asyncio.get_event_loop().create_future() + self._waiters.append(fut) + try: + await asyncio.wait_for(fut, timeout=1.0) + except TimeoutError: + pass # Re-check stop_event and drain any new batches + finally: + # Clean up if push() hasn't already swapped us out. + if fut in self._waiters: + self._waiters.remove(fut) + + class StreamConsumer: - """Consumer that receives events for a specific stream.""" + """A subscriber on a ``StreamBroadcast`` with an individual read cursor.""" def __init__( self, correlation_id: str, event_handler: EventHandler, - ): + broadcast: StreamBroadcast, + cursor: int, + ) -> None: self.correlation_id = correlation_id self.event_handler = event_handler - self.queue: asyncio.Queue[BaseModel | None] = asyncio.Queue() - self.ref_count = 1 - - async def send(self, event: BaseModel) -> None: - self.queue.put_nowait(event) - - async def close(self) -> None: - await self.queue.put(None) + self.broadcast = broadcast + self.cursor = cursor # Index into broadcast._batches class StreamState: - """Tracks state for a multiplexed stream.""" + """Read position and broadcast log for a single multiplexed stream.""" - def __init__(self, last_id: str = "0"): + def __init__(self, last_id: str = "0") -> None: self.last_id = last_id - self.last_flush_time = asyncio.get_event_loop().time() self.cached_events: list[BaseModel] | None = None + # 0.0 → check on the first idle cycle (ancient enough to pass the gate) + self.last_status_check: float = 0.0 + self.broadcast = StreamBroadcast() @singleton class StreamReader: - """Reads events from Redis and deserializes them.""" + """Reads events from the backend (Redis or in-memory) and deserializes them. + + On Redis, ``read_events`` with ``block_ms > 0`` issues ``XREAD BLOCK`` which + parks the connection — the event loop stays free. On the in-memory backend + the provider uses an ``asyncio.Event`` to avoid busy-waiting. + ``block_ms=None`` is always non-blocking (used for catch-up / drain reads). + """ @inject - def __init__(self, stream_component: StreamComponent): + def __init__(self, stream_component: StreamComponent) -> None: self.stream_service = stream_component.stream async def read_events( @@ -73,7 +152,6 @@ class StreamReader: count: int | None = None, block_ms: int | None = None, ) -> tuple[list[BaseModel], str]: - """Read events from Redis and deserialize them.""" raw_events, next_last_id = await self.stream_service.read_events( correlation_id=correlation_id, last_id=last_id, @@ -81,32 +159,27 @@ class StreamReader: block_ms=block_ms, ) - def sync_deserialization() -> list[BaseModel]: - events = [] - for raw_data in raw_events: - try: - event = event_handler.deserialize(raw_data) - events.append(event) - except Exception as e: - logger.error(f"Error deserializing event: {e}") - continue - return events - - final_events = sync_deserialization() - return final_events, next_last_id + events: list[BaseModel] = [] + for raw_data in raw_events: + try: + events.append(event_handler.deserialize(raw_data)) + except Exception as e: + logger.error(f"Error deserializing event: {e}") + return events, next_last_id async def check_terminal_status( self, correlation_id: str, event_handler: EventHandler, cached_events: list[BaseModel] | None, - last_flush_time: float, - cache_flush_interval: int, ) -> bool: - """Check if stream has reached terminal status.""" - current_time = asyncio.get_event_loop().time() + """Return True when the stream has reached a terminal status. - if cached_events and (current_time - last_flush_time) >= cache_flush_interval: + Tries the last cached event first; only hits Redis when that check is + inconclusive, avoiding a round-trip when the cached event already + carries terminal state. + """ + if cached_events: try: status = await event_handler.get_current_status(cached_events[-1]) if status in TERMINAL_STATUSES: @@ -133,16 +206,22 @@ class StreamReader: block_ms: int | None = None, stop_event: asyncio.Event | None = None, ) -> AsyncGenerator[BaseModel, None]: - """Stream events as an async generator.""" + """Stream events as an async generator (direct, non-multiplexed path). + + A producer task reads from the backend in blocking mode and feeds an + internal queue; the generator drains the queue and yields individual + events. After the terminal status is detected a final non-blocking + drain read is performed so no event published between the last read + and the terminal signal is lost. + """ stop_event = stop_event or asyncio.Event() block_ms = block_ms if block_ms is not None else DEFAULT_BLOCK_MS - cache_flush_interval: Final[int] = 5 * block_ms event_queue: asyncio.Queue[list[BaseModel] | None] = asyncio.Queue() async def produce_events(current_last_id: str) -> None: try: - cached_events: list[Any] | None = None - last_flush = asyncio.get_event_loop().time() + cached_events: list[BaseModel] | None = None + last_status_check: float = 0.0 # 0 → check on first idle cycle while not stop_event.is_set(): events, current_last_id = await self.read_events( @@ -154,23 +233,23 @@ class StreamReader: if events: cached_events = events - last_flush = asyncio.get_event_loop().time() await event_queue.put(events) - elif await self.check_terminal_status( - correlation_id, - event_handler, - cached_events, - last_flush, - cache_flush_interval, - ): - break + else: + current_time = asyncio.get_event_loop().time() + if (current_time - last_status_check) >= STATUS_CHECK_INTERVAL: + if await self.check_terminal_status( + correlation_id, + event_handler, + cached_events, + ): + break + last_status_check = current_time events, _ = await self.read_events( event_handler=event_handler, correlation_id=correlation_id, last_id=current_last_id, ) - if events: await event_queue.put(events) @@ -189,7 +268,6 @@ class StreamReader: break for event in events: yield event - finally: producer_task.cancel() with contextlib.suppress(asyncio.CancelledError): @@ -206,16 +284,30 @@ class StreamReader: class StreamMultiplexer: - """Multiplexes multiple stream consumers onto a single worker.""" + """Multiplexes many stream consumers onto a single worker task. - def __init__(self, stream_reader: StreamReader): + Each correlation id has one ``StreamState`` containing a ``StreamBroadcast`` + shared by all subscribers. The worker calls ``_process_stream`` once per + active stream per cycle: one Redis read, one ``broadcast.push`` — no + per-consumer queues. + + Key invariants + --------------- + * ``broadcast.push`` is synchronous and resolves all parked futures in one + pass, so all N subscribers wake atomically after a single call. + * State is updated under the lock; ``broadcast.push`` happens outside it so + the lock is held for as short a time as possible. + * ``close`` on the broadcast terminates all ``read_from`` generators, making + ``stop`` and terminal-status cleanup self-contained. + """ + + def __init__(self, stream_reader: StreamReader) -> None: self.stream_reader = stream_reader self.consumers: dict[str, list[StreamConsumer]] = defaultdict(list) self.states: dict[str, StreamState] = {} self.worker_task: asyncio.Task[None] | None = None self.shutdown_event = asyncio.Event() self._lock = asyncio.Lock() - self._cache_flush_interval = 5 * DEFAULT_BLOCK_MS async def start(self) -> None: if self.worker_task is None or self.worker_task.done(): @@ -224,47 +316,68 @@ class StreamMultiplexer: logger.info("Stream multiplexer started") async def stop(self) -> None: + """Cancel the worker, clear all state, and close every broadcast.""" self.shutdown_event.set() - if self.worker_task: - await self.worker_task - logger.info("Stream multiplexer stopped") + task = self.worker_task + self.worker_task = None + if task is not None and not task.done(): + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task + + async with self._lock: + states_to_close = list(self.states.values()) + self.consumers.clear() + self.states.clear() + + for state in states_to_close: + state.broadcast.close() async def add_consumer( self, correlation_id: str, event_handler: EventHandler, + last_id: str = "0", ) -> StreamConsumer: + """Register a new subscriber for *correlation_id*. + + Each call returns a distinct ``StreamConsumer`` with its own cursor so + all concurrent subscribers receive every event independently (true + fan-out). When this is the first consumer for the stream the read + position is initialised to *last_id*. + """ async with self._lock: - existing_consumers = self.consumers.get(correlation_id, []) - - for consumer in existing_consumers: - if isinstance(consumer.event_handler, type(event_handler)): - consumer.ref_count += 1 - return consumer - - consumer = StreamConsumer(correlation_id, event_handler) - self.consumers[correlation_id].append(consumer) - if correlation_id not in self.states: - self.states[correlation_id] = StreamState() + self.states[correlation_id] = StreamState(last_id=last_id) + state = self.states[correlation_id] + # last_id="0" means "from the start" — replay every batch already in + # the broadcast so a late joiner never misses history (e.g. the + # content_block_start that arrived before it connected). + # A specific last_id means the consumer is already caught up to that + # position, so start at the current end and receive only new batches. + cursor = 0 if last_id == "0" else state.broadcast.next_cursor() + consumer = StreamConsumer( + correlation_id, event_handler, state.broadcast, cursor + ) + self.consumers[correlation_id].append(consumer) return consumer async def remove_consumer(self, consumer: StreamConsumer) -> None: + broadcast_to_close: StreamBroadcast | None = None async with self._lock: correlation_id = consumer.correlation_id - consumer.ref_count -= 1 + cl = self.consumers.get(correlation_id) + if cl and consumer in cl: + cl.remove(consumer) + if correlation_id in self.consumers and not self.consumers[correlation_id]: + del self.consumers[correlation_id] + state = self.states.pop(correlation_id, None) + if state: + broadcast_to_close = state.broadcast - if consumer.ref_count <= 0: - if correlation_id in self.consumers: - self.consumers[correlation_id].remove(consumer) - - if not self.consumers[correlation_id]: - del self.consumers[correlation_id] - if correlation_id in self.states: - del self.states[correlation_id] - - await consumer.close() + if broadcast_to_close is not None: + broadcast_to_close.close() async def _run_worker(self) -> None: try: @@ -273,23 +386,21 @@ class StreamMultiplexer: correlation_ids = list(self.states.keys()) if not correlation_ids: - await asyncio.sleep(1) + await asyncio.sleep(0.5) continue - tasks = [ - self._process_stream(correlation_id) - for correlation_id in correlation_ids - ] + tasks = [self._process_stream(cid) for cid in correlation_ids] await asyncio.gather(*tasks, return_exceptions=True) await asyncio.sleep(0) - except Exception as e: logger.error(f"Worker error: {e}", exc_info=True) finally: async with self._lock: - for consumer_list in self.consumers.values(): - for consumer in consumer_list: - await consumer.close() + states_to_close = list(self.states.values()) + self.consumers.clear() + self.states.clear() + for state in states_to_close: + state.broadcast.close() async def _process_stream(self, correlation_id: str) -> None: try: @@ -297,75 +408,122 @@ class StreamMultiplexer: if correlation_id not in self.states: return state = self.states[correlation_id] - consumer_list = self.consumers.get(correlation_id, []) + consumer_list = list(self.consumers.get(correlation_id, [])) if not consumer_list: return + event_handler = consumer_list[0].event_handler events, new_last_id = await self.stream_reader.read_events( - event_handler=consumer_list[0].event_handler, + event_handler=event_handler, correlation_id=correlation_id, last_id=state.last_id, + block_ms=DEFAULT_BLOCK_MS, ) current_time = asyncio.get_event_loop().time() if events: - async with self._lock: - state.last_id = new_last_id - state.last_flush_time = current_time - state.cached_events = events + await self._dispatch( + correlation_id, state, events, new_last_id, current_time + ) - send_tasks = [] - for consumer in consumer_list: - for event in events: - send_tasks.append(consumer.send(event)) - - await asyncio.gather(*send_tasks) - - elif await self.stream_reader.check_terminal_status( - correlation_id, - consumer_list[0].event_handler, - state.cached_events, - state.last_flush_time, - self._cache_flush_interval, - ): - await self._close_all_consumers(correlation_id) - logger.info(f"Stream {correlation_id} reached terminal status") - else: - async with self._lock: - state.last_flush_time = current_time + elif (current_time - state.last_status_check) >= STATUS_CHECK_INTERVAL: + state.last_status_check = current_time + if await self.stream_reader.check_terminal_status( + correlation_id, + event_handler, + state.cached_events, + ): + drain_events, drain_last_id = await self.stream_reader.read_events( + event_handler=event_handler, + correlation_id=correlation_id, + last_id=state.last_id, + block_ms=None, + ) + if drain_events: + await self._dispatch( + correlation_id, + state, + drain_events, + drain_last_id, + current_time, + ) + await self._close_all_consumers(correlation_id) + logger.info(f"Stream {correlation_id} reached terminal status") except Exception as e: logger.error(f"Error processing {correlation_id}: {e}", exc_info=True) + async def _dispatch( + self, + correlation_id: str, + state: StreamState, + events: list[BaseModel], + new_last_id: str, + current_time: float, + ) -> None: + """Advance stream state then broadcast *events* to all subscribers. + + State is updated under the lock; ``broadcast.push`` is synchronous and + happens outside so the lock is held for the minimum time. + """ + async with self._lock: + if correlation_id not in self.states: + return + if not self.consumers.get(correlation_id): + return + state.last_id = new_last_id + state.cached_events = events + state.last_status_check = current_time + + state.broadcast.push(events) + async def _close_all_consumers(self, correlation_id: str) -> None: async with self._lock: - consumer_list = self.consumers.get(correlation_id, []) - for consumer in consumer_list: - await consumer.close() - - if correlation_id in self.consumers: - del self.consumers[correlation_id] - if correlation_id in self.states: - del self.states[correlation_id] + state = self.states.pop(correlation_id, None) + self.consumers.pop(correlation_id, None) + if state: + state.broadcast.close() @singleton class AdaptiveStreamReader: - """Adaptively switches between direct and multiplexed streaming based on load.""" + """Routes stream consumers between direct and multiplexed paths. + + Direct path: one XREAD BLOCK per consumer → one Redis connection held. + Multiplexed path: one shared worker for all active streams on this process. + + When concurrent direct-path readers reach ``multiplexing_threshold`` new + consumers are routed to the multiplexer instead of opening another Redis + connection. When the multiplexer drains to zero active streams it is + stopped, so load-shedding is fully reversible — consumers that arrive when + load is back below the threshold go direct again. + + ``multiplexing_threshold=None`` disables multiplexing entirely. + ``enable_multiplexing=False`` is a hard kill-switch (e.g. for tests). + """ @inject def __init__( self, settings: Settings, stream_reader: StreamReader, - ): + ) -> None: self.stream_reader = stream_reader self.multiplexer: StreamMultiplexer | None = None - self._active_count = 0 self._lock = asyncio.Lock() self.enable_multiplexing = True - self.multiplexing_threshold = settings.chat.multiplexing_threshold or None + self.multiplexing_threshold: int | None = settings.chat.multiplexing_threshold + # Number of consumers currently on the direct path. + # Each one holds an XREAD BLOCK connection; this is what we throttle. + self._direct_count = 0 + + def _should_multiplex(self) -> bool: + return ( + self.enable_multiplexing + and self.multiplexing_threshold is not None + and self._direct_count >= self.multiplexing_threshold + ) async def read_events( self, @@ -392,16 +550,14 @@ class AdaptiveStreamReader: stop_event: asyncio.Event | None = None, ) -> AsyncGenerator[BaseModel, None]: async def gen() -> AsyncGenerator[BaseModel, None]: + # Decide path under the lock so _direct_count is consistent. async with self._lock: - self._active_count += 1 - should_multiplex = ( - self.enable_multiplexing - and self.multiplexing_threshold is not None - and self._active_count >= self.multiplexing_threshold - ) + use_mux = self._should_multiplex() + if not use_mux: + self._direct_count += 1 try: - if should_multiplex: + if use_mux: async for event in await self._stream_multiplexed( event_handler, correlation_id, last_id, stop_event ): @@ -412,8 +568,9 @@ class AdaptiveStreamReader: ): yield event finally: - async with self._lock: - self._active_count -= 1 + if not use_mux: + async with self._lock: + self._direct_count -= 1 return gen() @@ -425,49 +582,38 @@ class AdaptiveStreamReader: stop_event: asyncio.Event | None, ) -> AsyncGenerator[BaseModel, None]: async def gen() -> AsyncGenerator[BaseModel, None]: - if self.multiplexer is None: - self.multiplexer = StreamMultiplexer(self.stream_reader) - await self.multiplexer.start() - logger.info("Multiplexed streaming activated") + # Start the multiplexer on first use; capture the reference so the + # finally block targets the right instance even if it is replaced. + async with self._lock: + if self.multiplexer is None: + self.multiplexer = StreamMultiplexer(self.stream_reader) + await self.multiplexer.start() + logger.info("Multiplexer started (threshold reached)") + mux = self.multiplexer - consumer = await self.multiplexer.add_consumer( - correlation_id, event_handler + consumer = await mux.add_consumer( + correlation_id, event_handler, last_id=last_id ) - try: - if last_id != "0": - initial_events, _ = await self.read_events( - event_handler, correlation_id, "0", None, 0 - ) - for initial_event in initial_events: - yield initial_event - - events_processed = 0 - while True: - if stop_event and stop_event.is_set(): - break - - try: - event: Any | None = await asyncio.wait_for( - consumer.queue.get(), timeout=1.0 - ) - if event is None: - break - yield event - - events_processed += 1 - if events_processed % BATCH_SIZE == 0: - await asyncio.sleep(0) - - except TimeoutError: - continue - + async for event in consumer.broadcast.read_from( + cursor=consumer.cursor, + stop_event=stop_event, + ): + yield event finally: - await self.multiplexer.remove_consumer(consumer) + await mux.remove_consumer(consumer) + # Stop the multiplexer when it has no more active streams so the + # next consumer below the threshold can go direct again. + async with self._lock: + if self.multiplexer is mux and not mux.states: + await mux.stop() + self.multiplexer = None + logger.info("Multiplexer stopped (no active streams)") return gen() async def shutdown(self) -> None: - if self.multiplexer: - await self.multiplexer.stop() - self.multiplexer = None + async with self._lock: + mux, self.multiplexer = self.multiplexer, None + if mux: + await mux.stop() diff --git a/private_gpt/components/streaming/tasks/chat_scheduler.py b/private_gpt/components/streaming/tasks/chat_scheduler.py new file mode 100644 index 00000000..c7bcc2f7 --- /dev/null +++ b/private_gpt/components/streaming/tasks/chat_scheduler.py @@ -0,0 +1,76 @@ +"""Chat schedulers — execution lifecycle management.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from collections.abc import Callable + +from injector import Injector, inject, singleton + +from private_gpt.arq.tasks.chat import abort_chat_job +from private_gpt.settings.settings import Settings + +ChatSchedulerProvider = ( + type["BaseChatScheduler"] | Callable[[Injector], "BaseChatScheduler"] +) +_CHAT_SCHEDULERS: dict[str, ChatSchedulerProvider] = {} + + +def register_chat_scheduler(mode: str, provider: ChatSchedulerProvider) -> None: + _CHAT_SCHEDULERS[mode] = provider + + +class BaseChatScheduler(ABC): + """Provides lifecycle operations for chat executions.""" + + @abstractmethod + async def cancel(self, correlation_id: str) -> bool: + """Cancel a running execution.""" + + +class LocalChatScheduler(BaseChatScheduler): + """Cancels local asyncio tasks by name.""" + + async def cancel(self, correlation_id: str) -> bool: + import asyncio + + for task in asyncio.all_tasks(): + if task.get_name() == f"chat_{correlation_id}": + task.cancel() + return True + return False + + +class ArqChatScheduler(BaseChatScheduler): + """Aborts ARQ jobs.""" + + async def cancel(self, correlation_id: str) -> bool: + return await abort_chat_job(correlation_id=correlation_id) + + +register_chat_scheduler("local", LocalChatScheduler) +register_chat_scheduler("arq", ArqChatScheduler) + + +@singleton +class ChatSchedulerFactory: + @inject + def __init__(self, settings: Settings, injector: Injector) -> None: + self._settings = settings + self._injector = injector + self._scheduler: BaseChatScheduler | None = None + + mode = self._settings.scheduler.chat.mode + if mode not in _CHAT_SCHEDULERS: + raise ValueError(f"Unknown scheduler.chat.mode: {mode}") + + def get(self) -> BaseChatScheduler: + if self._scheduler is None: + mode = self._settings.scheduler.chat.mode + provider = _CHAT_SCHEDULERS[mode] + self._scheduler = ( + self._injector.get(provider) + if isinstance(provider, type) + else provider(self._injector) + ) + return self._scheduler diff --git a/private_gpt/components/streaming/tasks/task_manager.py b/private_gpt/components/streaming/tasks/task_manager.py index eab3c1f8..15945fa3 100644 --- a/private_gpt/components/streaming/tasks/task_manager.py +++ b/private_gpt/components/streaming/tasks/task_manager.py @@ -5,25 +5,18 @@ from collections.abc import Coroutine from contextvars import copy_context from typing import Any -from injector import inject, singleton +from injector import singleton -from private_gpt.settings.settings import Settings +from private_gpt.di import get_global_injector logger = logging.getLogger(__name__) @singleton class TaskManager: - """Manages asyncio tasks with concurrency control and cancellation support.""" + """Manages asyncio tasks with cancellation support.""" - @inject - def __init__(self, settings: Settings) -> None: - max_concurrent_tasks = settings.chat.maximum_concurrent_requests or None - self._semaphore = ( - asyncio.Semaphore(max_concurrent_tasks) - if max_concurrent_tasks is not None - else None - ) + def __init__(self) -> None: self._active_tasks: dict[str, asyncio.Task[Any | None]] = {} self._cancellation_tokens: dict[str, asyncio.Event] = {} @@ -33,31 +26,16 @@ class TaskManager: coro: Coroutine[Any, Any, None], name: str | None = None, ) -> asyncio.Task[Any]: - """Create and register a new task with concurrency control.""" + """Create and register a new task.""" if correlation_id in self._active_tasks: logger.warning(f"Task {correlation_id} already exists") return self._active_tasks[correlation_id] - async def wrapped_coro() -> None: - try: - if self._semaphore is not None: - await self._semaphore.acquire() - logger.debug( - f"Acquired slot for {correlation_id} " - f"(active: {len(self._active_tasks)})" - ) - - await coro - finally: - if self._semaphore is not None: - self._semaphore.release() - logger.debug(f"Released slot for {correlation_id}") - cancellation_token = asyncio.Event() self._cancellation_tokens[correlation_id] = cancellation_token ctx = copy_context() - task = asyncio.create_task(wrapped_coro(), name=name, context=ctx) + task = asyncio.create_task(coro, name=name, context=ctx) self._active_tasks[correlation_id] = task task.add_done_callback(lambda _: self._cleanup_task(correlation_id)) @@ -76,7 +54,7 @@ class TaskManager: token = self._cancellation_tokens.get(correlation_id) return token is not None and token.is_set() - async def cancel_task(self, correlation_id: str, timeout: float = 2.0) -> bool: + async def cancel_task(self, correlation_id: str) -> bool: """Cancel a task and wait for completion.""" if correlation_id in self._cancellation_tokens: self._cancellation_tokens[correlation_id].set() @@ -87,11 +65,19 @@ class TaskManager: task.cancel() - with contextlib.suppress(TimeoutError, asyncio.CancelledError): - await asyncio.wait_for(task, timeout=timeout) + with contextlib.suppress(asyncio.CancelledError): + await task return True + async def cancel(self, correlation_id: str) -> bool: + from private_gpt.server.chat.chat_service import ChatService + + chat_service = get_global_injector().get(ChatService) + success = await self.cancel_task(correlation_id) + scheduled_cancel = await chat_service.cancel(correlation_id) + return success or scheduled_cancel + def _cleanup_task(self, correlation_id: str) -> None: """Clean up task references.""" self._active_tasks.pop(correlation_id, None) diff --git a/private_gpt/components/tools/builders/bash_tool_builder.py b/private_gpt/components/tools/builders/bash_tool_builder.py index 6f888c92..c1d56daa 100644 --- a/private_gpt/components/tools/builders/bash_tool_builder.py +++ b/private_gpt/components/tools/builders/bash_tool_builder.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any, cast from injector import inject, singleton @@ -11,9 +11,11 @@ from private_gpt.components.chat.models.chat_config_models import ( from private_gpt.components.code_execution.code_execution_component import ( CodeExecutionComponent, ) +from private_gpt.components.tools.remote_execution import build_rebuild_metadata from private_gpt.components.tools.tool_names import BASH_TOOL_NAME from private_gpt.components.tools.tool_placeholders import BASH_TOOL_FN from private_gpt.components.tools.utils import truncate_output +from private_gpt.di import get_global_injector from private_gpt.events.models import TextBlock from private_gpt.settings.settings import Settings @@ -75,4 +77,19 @@ class BashToolBuilder: description=description, async_fn=run_bash, requirements=[ToolRequirements.SANDBOX], + execution_metadata=build_rebuild_metadata( + rebuild_bash_tool, + { + "session_id": session_id, + "bundles": bundles, + "name": name, + "type": type, + "description": description, + }, + ), ) + + +async def rebuild_bash_tool(**kwargs: Any) -> ToolSpec: + builder = get_global_injector().get(BashToolBuilder) + return await builder.build_tool(**cast(Any, kwargs)) diff --git a/private_gpt/components/tools/builders/database_query_builder.py b/private_gpt/components/tools/builders/database_query_builder.py index 43fcf5ba..aa68c9bd 100644 --- a/private_gpt/components/tools/builders/database_query_builder.py +++ b/private_gpt/components/tools/builders/database_query_builder.py @@ -1,5 +1,5 @@ import asyncio -from typing import TYPE_CHECKING, Literal +from typing import TYPE_CHECKING, Any, Literal, cast from injector import inject, singleton from llama_index.core.base.llms.types import ( @@ -16,9 +16,11 @@ from private_gpt.components.llm.llm_component import LLMComponent from private_gpt.components.tools.binary_block_decorators import ( auto_resolve_media_blocks, ) +from private_gpt.components.tools.remote_execution import build_rebuild_metadata from private_gpt.components.tools.tool_names import DATABASE_QUERY_TOOL_NAME from private_gpt.components.tools.tool_placeholders import DATABASE_QUERY_TOOL_FN from private_gpt.components.tools.types import ToolValidationMode +from private_gpt.di import get_global_injector from private_gpt.events.models import ( BinaryBlock, ResultContentBlockType, @@ -387,4 +389,22 @@ class DatabaseQueryToolBuilder: runtime=runtime, description=description, async_fn=run_tool, + execution_metadata=build_rebuild_metadata( + rebuild_database_query_tool, + { + "sql_artifacts": sql_artifacts, + "chat_history": chat_history, + "name": name, + "type": type, + "description": description, + "validate": validate, + "runtime": runtime, + "blob_visibility": blob_visibility, + }, + ), ) + + +async def rebuild_database_query_tool(**kwargs: Any) -> ToolSpec: + builder = get_global_injector().get(DatabaseQueryToolBuilder) + return await builder.build_tool(**cast(Any, kwargs)) diff --git a/private_gpt/components/tools/builders/present_files_tool_builder.py b/private_gpt/components/tools/builders/present_files_tool_builder.py index 1dde14ea..b2382908 100644 --- a/private_gpt/components/tools/builders/present_files_tool_builder.py +++ b/private_gpt/components/tools/builders/present_files_tool_builder.py @@ -3,7 +3,7 @@ from __future__ import annotations import base64 import mimetypes from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any, cast from injector import inject, singleton @@ -14,8 +14,10 @@ from private_gpt.components.chat.models.chat_config_models import ( from private_gpt.components.code_execution.code_execution_component import ( CodeExecutionComponent, ) +from private_gpt.components.tools.remote_execution import build_rebuild_metadata from private_gpt.components.tools.tool_names import PRESENT_FILES_TOOL_NAME from private_gpt.components.tools.tool_placeholders import PRESENT_FILES_TOOL_FN +from private_gpt.di import get_global_injector from private_gpt.events.models import LocalResourceBlock, TextBlock if TYPE_CHECKING: @@ -95,4 +97,19 @@ class PresentFilesToolBuilder: description=description, async_fn=present_files, requirements=[ToolRequirements.SANDBOX], + execution_metadata=build_rebuild_metadata( + rebuild_present_files_tool, + { + "session_id": session_id, + "bundles": bundles, + "name": name, + "type": type, + "description": description, + }, + ), ) + + +async def rebuild_present_files_tool(**kwargs: Any) -> ToolSpec: + builder = get_global_injector().get(PresentFilesToolBuilder) + return await builder.build_tool(**cast(Any, kwargs)) diff --git a/private_gpt/components/tools/builders/semantic_search_builder.py b/private_gpt/components/tools/builders/semantic_search_builder.py index 17f07b96..4425b45d 100644 --- a/private_gpt/components/tools/builders/semantic_search_builder.py +++ b/private_gpt/components/tools/builders/semantic_search_builder.py @@ -24,6 +24,7 @@ from private_gpt.components.postprocessor.tree_expansion.tree_expansion_replacem TreeExpansionReplacementPostProcessor, ) from private_gpt.components.prompts.prompt_builder import PromptBuilderService +from private_gpt.components.tools.remote_execution import build_rebuild_metadata from private_gpt.components.tools.tool_names import SEMANTIC_SEARCH_TOOL_NAME from private_gpt.components.tools.tool_placeholders import SEMANTIC_SEARCH_TOOL_FN from private_gpt.components.tools.types import ToolValidationMode @@ -34,6 +35,7 @@ from private_gpt.components.workflows.retrieval.semantic_search import ( SemanticSearchInputEvent, SemanticSearchWorkflow, ) +from private_gpt.di import get_global_injector from private_gpt.events.models import ( ResultContentBlockType, SourceBlock, @@ -185,7 +187,9 @@ class SemanticSearchToolBuilder: tokenizer: TokenizerBase | None = None, ) -> SemanticSearchWorkflow: context_filter = await self._validate_context(context_filter) - retriever = self._create_vector_index_retriever(context_filter, embed_model_id) + retriever = await asyncio.to_thread( + self._create_vector_index_retriever, context_filter, embed_model_id + ) if not llm: llm = self.llm_component.get_llm(model_id) @@ -326,4 +330,25 @@ class SemanticSearchToolBuilder: description=description, async_fn=semantic_search, async_callback=format_result, + execution_metadata=build_rebuild_metadata( + rebuild_semantic_search_tool, + { + "context_filter": context_filter, + "model_id": model_id, + "embed_model_id": embed_model_id, + "name": name, + "type": type, + "description": description, + "validate": validate, + "runtime": runtime, + **kwargs, + "generate_citations": generate_citations, + "token_limit": token_limit, + }, + ), ) + + +async def rebuild_semantic_search_tool(**kwargs: Any) -> ToolSpec: + builder = get_global_injector().get(SemanticSearchToolBuilder) + return await builder.build_tool(**kwargs) diff --git a/private_gpt/components/tools/builders/skill_management_builder.py b/private_gpt/components/tools/builders/skill_management_builder.py index 7bcd8397..b25aebb1 100644 --- a/private_gpt/components/tools/builders/skill_management_builder.py +++ b/private_gpt/components/tools/builders/skill_management_builder.py @@ -10,11 +10,13 @@ from private_gpt.components.skills.models.skill_entities import ( SkillVersionEntity, ) from private_gpt.components.skills.services.skill_service import SkillService +from private_gpt.components.tools.remote_execution import build_rebuild_metadata from private_gpt.components.tools.tool_names import ( SKILL_LIST_TOOL_NAME, SKILL_LOAD_TOOL_NAME, SKILL_UNLOAD_TOOL_NAME, ) +from private_gpt.di import get_global_injector from private_gpt.events.models import ResultContentBlockType, TextBlock @@ -72,6 +74,15 @@ class SkillManagementToolBuilder: description="Mark one available skill as loaded for this conversation.", async_fn=load_skill, requirements=[ToolRequirements.SANDBOX], + execution_metadata=build_rebuild_metadata( + rebuild_load_skill_tool, + { + "skill_filter": self._skill_filter, + "skill_injection_mode": self._skill_injection_mode, + "name": name, + "type": type, + }, + ), ) def build_unload_skill( @@ -89,6 +100,15 @@ class SkillManagementToolBuilder: description="Mark one loaded skill as unloaded for this conversation.", async_fn=unload_skill, requirements=[ToolRequirements.SANDBOX], + execution_metadata=build_rebuild_metadata( + rebuild_unload_skill_tool, + { + "skill_filter": self._skill_filter, + "skill_injection_mode": self._skill_injection_mode, + "name": name, + "type": type, + }, + ), ) def build_list_skills( @@ -128,4 +148,61 @@ class SkillManagementToolBuilder: description="Browse the skill catalog (paginated). Use page/page_size to navigate large catalogs.", async_fn=list_skills, requirements=[ToolRequirements.SANDBOX], + execution_metadata=build_rebuild_metadata( + rebuild_list_skills_tool, + { + "skill_filter": self._skill_filter, + "skill_injection_mode": self._skill_injection_mode, + "name": name, + "type": type, + }, + ), ) + + +def _builder( + skill_filter: SkillFilter, + skill_injection_mode: str, +) -> "SkillManagementToolBuilder": + injector = get_global_injector() + return SkillManagementToolBuilder( + skill_service=injector.get(SkillService), + skill_filter=skill_filter, + skill_injection_mode=skill_injection_mode, + ) + + +def rebuild_load_skill_tool( + skill_filter: SkillFilter, + skill_injection_mode: str, + name: str = SKILL_LOAD_TOOL_NAME, + type: str = SKILL_LOAD_TOOL_NAME + "_v1", +) -> ToolSpec: + return _builder(skill_filter, skill_injection_mode).build_load_skill( + name=name, + type=type, + ) + + +def rebuild_unload_skill_tool( + skill_filter: SkillFilter, + skill_injection_mode: str, + name: str = SKILL_UNLOAD_TOOL_NAME, + type: str = SKILL_UNLOAD_TOOL_NAME + "_v1", +) -> ToolSpec: + return _builder(skill_filter, skill_injection_mode).build_unload_skill( + name=name, + type=type, + ) + + +def rebuild_list_skills_tool( + skill_filter: SkillFilter, + skill_injection_mode: str, + name: str = SKILL_LIST_TOOL_NAME, + type: str = SKILL_LIST_TOOL_NAME + "_v1", +) -> ToolSpec: + return _builder(skill_filter, skill_injection_mode).build_list_skills( + name=name, + type=type, + ) diff --git a/private_gpt/components/tools/builders/tabular_data_builder.py b/private_gpt/components/tools/builders/tabular_data_builder.py index 4000ffcf..790cf14a 100644 --- a/private_gpt/components/tools/builders/tabular_data_builder.py +++ b/private_gpt/components/tools/builders/tabular_data_builder.py @@ -28,12 +28,14 @@ from private_gpt.components.sandbox import SandboxComponent from private_gpt.components.tools.binary_block_decorators import ( auto_resolve_media_blocks, ) +from private_gpt.components.tools.remote_execution import build_rebuild_metadata from private_gpt.components.tools.tool_names import TABULAR_DATA_ANALYSIS from private_gpt.components.tools.tool_placeholders import TABULAR_DATA_TOOL_FN from private_gpt.components.tools.types import ToolValidationMode from private_gpt.components.vector_store.vector_store_component import ( VectorStoreComponent, ) +from private_gpt.di import get_global_injector from private_gpt.events.models import ( ResultContentBlockType, from_tool_output, @@ -222,7 +224,9 @@ class TabularDataToolBuilder: _, tabular_data_analysis_workflow_cls = _load_tabular_workflow_dependencies() llm = llm or self.llm_component.get_llm(model_id) context_filter = await self._validate_context(context_filter) - retriever = self._create_vector_index_retriever(context_filter, embed_model_id) + retriever = await asyncio.to_thread( + self._create_vector_index_retriever, context_filter, embed_model_id + ) pandas_ai = self._create_pandas_ai_service() def node_postprocessors_fn( @@ -331,4 +335,24 @@ class TabularDataToolBuilder: description=description, async_fn=tabular_data_analysis, requirements=[ToolRequirements.SANDBOX], + execution_metadata=build_rebuild_metadata( + rebuild_tabular_data_tool, + { + "context_filter": context_filter, + "model_id": model_id, + "embed_model_id": embed_model_id, + "name": name, + "type": type, + "description": description, + "validate": validate, + "runtime": runtime, + "blob_visibility": blob_visibility, + **kwargs, + }, + ), ) + + +async def rebuild_tabular_data_tool(**kwargs: Any) -> ToolSpec: + builder = get_global_injector().get(TabularDataToolBuilder) + return await builder.build_tool(**kwargs) diff --git a/private_gpt/components/tools/builders/text_editor_tool_builder.py b/private_gpt/components/tools/builders/text_editor_tool_builder.py index 580897c2..ef8ff688 100644 --- a/private_gpt/components/tools/builders/text_editor_tool_builder.py +++ b/private_gpt/components/tools/builders/text_editor_tool_builder.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any, cast from injector import inject, singleton @@ -11,6 +11,7 @@ from private_gpt.components.chat.models.chat_config_models import ( from private_gpt.components.code_execution.code_execution_component import ( CodeExecutionComponent, ) +from private_gpt.components.tools.remote_execution import build_rebuild_metadata from private_gpt.components.tools.tool_names import ( TEXT_EDITOR_CREATE_TOOL_NAME, TEXT_EDITOR_INSERT_TOOL_NAME, @@ -24,6 +25,7 @@ from private_gpt.components.tools.tool_placeholders import ( TEXT_EDITOR_VIEW_TOOL_FN, ) from private_gpt.components.tools.utils import truncate_output +from private_gpt.di import get_global_injector from private_gpt.events.models import TextBlock from private_gpt.settings.settings import Settings @@ -99,6 +101,16 @@ class TextEditorToolBuilder: description=description, async_fn=view, requirements=[ToolRequirements.SANDBOX], + execution_metadata=build_rebuild_metadata( + rebuild_text_editor_view_tool, + { + "session_id": session_id, + "bundles": bundles, + "name": name, + "type": type, + "description": description, + }, + ), ) async def build_str_replace_tool( @@ -129,6 +141,16 @@ class TextEditorToolBuilder: description=description, async_fn=str_replace, requirements=[ToolRequirements.SANDBOX], + execution_metadata=build_rebuild_metadata( + rebuild_text_editor_str_replace_tool, + { + "session_id": session_id, + "bundles": bundles, + "name": name, + "type": type, + "description": description, + }, + ), ) async def build_create_tool( @@ -158,6 +180,16 @@ class TextEditorToolBuilder: description=description, async_fn=create, requirements=[ToolRequirements.SANDBOX], + execution_metadata=build_rebuild_metadata( + rebuild_text_editor_create_tool, + { + "session_id": session_id, + "bundles": bundles, + "name": name, + "type": type, + "description": description, + }, + ), ) async def build_insert_tool( @@ -188,4 +220,34 @@ class TextEditorToolBuilder: description=description, async_fn=insert, requirements=[ToolRequirements.SANDBOX], + execution_metadata=build_rebuild_metadata( + rebuild_text_editor_insert_tool, + { + "session_id": session_id, + "bundles": bundles, + "name": name, + "type": type, + "description": description, + }, + ), ) + + +async def rebuild_text_editor_view_tool(**kwargs: Any) -> ToolSpec: + builder = get_global_injector().get(TextEditorToolBuilder) + return await builder.build_view_tool(**cast(Any, kwargs)) + + +async def rebuild_text_editor_str_replace_tool(**kwargs: Any) -> ToolSpec: + builder = get_global_injector().get(TextEditorToolBuilder) + return await builder.build_str_replace_tool(**cast(Any, kwargs)) + + +async def rebuild_text_editor_create_tool(**kwargs: Any) -> ToolSpec: + builder = get_global_injector().get(TextEditorToolBuilder) + return await builder.build_create_tool(**cast(Any, kwargs)) + + +async def rebuild_text_editor_insert_tool(**kwargs: Any) -> ToolSpec: + builder = get_global_injector().get(TextEditorToolBuilder) + return await builder.build_insert_tool(**cast(Any, kwargs)) diff --git a/private_gpt/components/tools/builders/web_fetch_builder.py b/private_gpt/components/tools/builders/web_fetch_builder.py index 4df54803..ef84e9c5 100644 --- a/private_gpt/components/tools/builders/web_fetch_builder.py +++ b/private_gpt/components/tools/builders/web_fetch_builder.py @@ -1,12 +1,14 @@ -from typing import Literal +from typing import Any, Literal, cast from injector import inject, singleton from private_gpt.components.chat.models.chat_config_models import ToolSpec from private_gpt.components.llm.llm_component import LLMComponent +from private_gpt.components.tools.remote_execution import build_rebuild_metadata from private_gpt.components.tools.tool_names import WEB_FETCH_TOOL_NAME from private_gpt.components.tools.tool_placeholders import WEB_FETCH_TOOL_FN from private_gpt.components.web.web_scraper_service import WebScraperService +from private_gpt.di import get_global_injector from private_gpt.events.models import ( ResultContentBlockType, TextBlock, @@ -63,4 +65,18 @@ class WebFetchToolBuilder: runtime=runtime, description=description, async_fn=run_tool, + execution_metadata=build_rebuild_metadata( + rebuild_web_fetch_tool, + { + "name": name, + "type": type, + "description": description, + "runtime": runtime, + }, + ), ) + + +def rebuild_web_fetch_tool(**kwargs: Any) -> ToolSpec: + builder = get_global_injector().get(WebFetchToolBuilder) + return builder.build_tool(**cast(Any, kwargs)) diff --git a/private_gpt/components/tools/builders/web_search_builder.py b/private_gpt/components/tools/builders/web_search_builder.py index 0e546611..52f7f634 100644 --- a/private_gpt/components/tools/builders/web_search_builder.py +++ b/private_gpt/components/tools/builders/web_search_builder.py @@ -1,16 +1,18 @@ import asyncio -from typing import Literal +from typing import Any, Literal, cast from injector import inject, singleton from private_gpt.components.chat.models.chat_config_models import ToolSpec from private_gpt.components.chunk.models import Website from private_gpt.components.llm.llm_component import LLMComponent +from private_gpt.components.tools.remote_execution import build_rebuild_metadata from private_gpt.components.tools.tool_names import WEB_SEARCH_TOOL_NAME from private_gpt.components.tools.tool_placeholders import WEB_SEARCH_TOOL_FN from private_gpt.components.tools.types import ToolValidationMode from private_gpt.components.web.web_search.models import WebSearchResult from private_gpt.components.web.web_search.web_search_service import WebSearchService +from private_gpt.di import get_global_injector from private_gpt.events.models import ( ResultContentBlockType, TextBlock, @@ -87,4 +89,20 @@ class WebSearchToolBuilder: runtime=runtime, description=description, async_fn=run_tool, + execution_metadata=build_rebuild_metadata( + rebuild_web_search_tool, + { + "model_id": model_id, + "name": name, + "type": type, + "description": description, + "validate": validate, + "runtime": runtime, + }, + ), ) + + +async def rebuild_web_search_tool(**kwargs: Any) -> ToolSpec: + builder = get_global_injector().get(WebSearchToolBuilder) + return await builder.build_tool(**cast(Any, kwargs)) diff --git a/private_gpt/components/tools/remote_execution.py b/private_gpt/components/tools/remote_execution.py new file mode 100644 index 00000000..85904e89 --- /dev/null +++ b/private_gpt/components/tools/remote_execution.py @@ -0,0 +1,206 @@ +from __future__ import annotations + +import importlib +import inspect +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Any + +from llama_index.core.base.llms.types import ChatMessage +from llama_index.core.tools import adapt_to_async_tool +from pydantic import BaseModel, Field + +from private_gpt.components.chat.models.chat_config_models import ( + ToolExecutionMetadata, + ToolSpec, +) +from private_gpt.components.engines.chat.models.chat_phase import ( + InterceptorPhase, +) +from private_gpt.components.engines.chat.models.execution_hooks import ( + ExecutionHooks, +) +from private_gpt.components.engines.chat.utils.tool_utils import execute_tool_call +from private_gpt.events.models import ( + ResultContentBlockType, + TextBlock, + from_tool_output, +) + +if TYPE_CHECKING: + from llama_index.core.tools import AsyncBaseTool + + from private_gpt.components.engines.chat.models.chat_state import ( + ChatState, + ) + from private_gpt.components.engines.chat.models.execution_hooks import ( + ToolExecutionHook, + ) + + +class ToolExecutionRequest(BaseModel): + tool_id: str + tool_name: str + tool_kwargs: dict[str, Any] = Field(default_factory=dict) + tool_spec: ToolSpec + context: dict[str, Any] = Field(default_factory=dict) + hooks: ExecutionHooks = Field(default_factory=ExecutionHooks) + + +async def invoke_execution_hook( + hook: ToolExecutionHook, + request: ToolExecutionRequest, + response: ToolExecutionResponse, +) -> None: + callback_callable = _import_callable(hook.callable_path) + result = callback_callable(request=request, response=response, **hook.kwargs) + if inspect.isawaitable(result): + await result + + +class ToolExecutionResponse(BaseModel): + tool_name: str + tool_id: str + result_content: list[ResultContentBlockType] = Field(default_factory=list) + is_error: bool = False + tool_message: ChatMessage + + +class ToolExecutionInterceptorContext(BaseModel): + phase: InterceptorPhase + request: ToolExecutionRequest + tool_kwargs: dict[str, Any] + response: ToolExecutionResponse | None = None + + def set_tool_kwargs(self, tool_kwargs: dict[str, Any]) -> None: + self.tool_kwargs = tool_kwargs + + def set_response(self, response: ToolExecutionResponse) -> None: + self.response = response + + +class ToolExecutionInterceptor(ABC): + @abstractmethod + async def intercept(self, context: ToolExecutionInterceptorContext) -> None: + """Mutate tool execution context before/after tool invocation.""" + + +class ToolExecutor: + def __init__( + self, + interceptors: list[ToolExecutionInterceptor] | None = None, + ) -> None: + self._interceptors = interceptors or [] + + async def execute( + self, + request: ToolExecutionRequest, + state_ctx: ChatState | None = None, + ) -> ToolExecutionResponse: + tool = await rebuild_tool_from_spec(request.tool_spec) + + before_context = ToolExecutionInterceptorContext( + phase=InterceptorPhase.BEFORE_TOOL, + request=request, + tool_kwargs=dict(request.tool_kwargs), + ) + for interceptor in self._interceptors: + await interceptor.intercept(before_context) + + result, tool_message = await execute_tool_call( + tool=tool, + tool_name=request.tool_name, + tool_id=request.tool_id, + tool_kwargs=before_context.tool_kwargs, + state_ctx=state_ctx, + ) + response = ToolExecutionResponse( + tool_name=request.tool_name, + tool_id=request.tool_id, + result_content=( + from_tool_output(result.tool_output.raw_output) + if result.tool_output.raw_output is not None + else [TextBlock(text=result.tool_output.content or "")] + ), + is_error=result.tool_output.is_error, + tool_message=tool_message, + ) + + after_context = ToolExecutionInterceptorContext( + phase=InterceptorPhase.AFTER_TOOL, + request=request, + tool_kwargs=before_context.tool_kwargs, + response=response, + ) + for interceptor in self._interceptors: + await interceptor.intercept(after_context) + + assert after_context.response is not None + return after_context.response + + +def build_rebuild_metadata( + rebuild_callable: Any, + rebuild_kwargs: dict[str, Any] | None = None, +) -> ToolExecutionMetadata: + return ToolExecutionMetadata( + rebuild_callable=_callable_path(rebuild_callable), + rebuild_kwargs=rebuild_kwargs or {}, + ) + + +async def rebuild_tool_from_spec(tool_spec: ToolSpec) -> AsyncBaseTool: + metadata = tool_spec.execution_metadata + if metadata is None: + return adapt_to_async_tool(tool_spec.to_function_tool()) + + rebuilt = await _invoke_rebuild(metadata) + return adapt_to_async_tool(rebuilt.to_function_tool()) + + +async def execute_tool_request( + request: ToolExecutionRequest, + state_ctx: ChatState | None = None, + interceptors: list[ToolExecutionInterceptor] | None = None, +) -> ToolExecutionResponse: + executor = ToolExecutor(interceptors=interceptors) + return await executor.execute(request, state_ctx=state_ctx) + + +def build_tool_execution_context(state: ChatState) -> dict[str, Any]: + return { + "correlation_id": state.input.request.context.correlation_id, + "messages": [ + msg.model_dump(mode="json", exclude_none=True) + for msg in state.input.request.messages + ], + } + + +def restore_chat_history_from_context(context: dict[str, Any]) -> list[ChatMessage]: + return [ + ChatMessage.model_validate(message_data) + for message_data in context.get("messages", []) + ] + + +async def _invoke_rebuild(metadata: ToolExecutionMetadata) -> ToolSpec: + rebuild_callable = _import_callable(metadata.rebuild_callable) + rebuilt = rebuild_callable(**metadata.rebuild_kwargs) + if inspect.isawaitable(rebuilt): + rebuilt = await rebuilt + if not isinstance(rebuilt, ToolSpec): + raise TypeError("Tool rebuild callable must return a ToolSpec instance.") + return rebuilt + + +def _callable_path(rebuild_callable: Any) -> str: + return f"{rebuild_callable.__module__}:{rebuild_callable.__qualname__}" + + +def _import_callable(path: str) -> Any: + module_name, attr_path = path.split(":", maxsplit=1) + module = importlib.import_module(module_name) + target = module + for attr in attr_path.split("."): + target = getattr(target, attr) + return target diff --git a/private_gpt/components/tools/tool_scheduler.py b/private_gpt/components/tools/tool_scheduler.py new file mode 100644 index 00000000..7416c1b3 --- /dev/null +++ b/private_gpt/components/tools/tool_scheduler.py @@ -0,0 +1,191 @@ +from __future__ import annotations + +import logging +from abc import ABC, abstractmethod +from collections.abc import Callable +from typing import TYPE_CHECKING + +from injector import Injector, inject, singleton + +from private_gpt.celery.dispatch import dispatch_task +from private_gpt.components.tools.remote_execution import ( + execute_tool_request, + invoke_execution_hook, +) +from private_gpt.settings.settings import Settings + +TOOL_TASK_NAME = "private_gpt.tools.run" + +if TYPE_CHECKING: + from private_gpt.components.engines.chat.models.chat_state import ChatState + from private_gpt.components.tools.remote_execution import ( + ToolExecutionInterceptor, + ToolExecutionRequest, + ToolExecutionResponse, + ) + +logger = logging.getLogger(__name__) + +ToolSchedulerProvider = ( + type["BaseToolScheduler"] | Callable[[Injector], "BaseToolScheduler"] +) +_TOOL_SCHEDULERS: dict[str, ToolSchedulerProvider] = {} + + +def register_tool_scheduler(mode: str, provider: ToolSchedulerProvider) -> None: + _TOOL_SCHEDULERS[mode] = provider + + +class BaseToolScheduler(ABC): + @property + def is_async(self) -> bool: + return False + + @abstractmethod + async def execute( + self, + request: ToolExecutionRequest, + state_ctx: ChatState | None = None, + interceptors: list[ToolExecutionInterceptor] | None = None, + ) -> ToolExecutionResponse: + ... + + async def async_execute( + self, + request: ToolExecutionRequest, + state_ctx: ChatState | None = None, + interceptors: list[ToolExecutionInterceptor] | None = None, + ) -> str: + del request, state_ctx, interceptors + raise NotImplementedError + + @abstractmethod + async def cancel( + self, + request: ToolExecutionRequest, + task_id: str | None = None, + ) -> bool: + ... + + async def cancel_task(self, task_id: str) -> bool: + del task_id + return False + + async def complete( + self, + request: ToolExecutionRequest, + response: ToolExecutionResponse, + ) -> None: + for hook in request.hooks.tool_result: + await invoke_execution_hook(hook, request, response) + + +@singleton +class LocalToolScheduler(BaseToolScheduler): + """Execute tools in-process (no worker dispatch).""" + + async def execute( + self, + request: ToolExecutionRequest, + state_ctx: ChatState | None = None, + interceptors: list[ToolExecutionInterceptor] | None = None, + ) -> ToolExecutionResponse: + return await execute_tool_request( + request, state_ctx=state_ctx, interceptors=interceptors + ) + + async def cancel( + self, + request: ToolExecutionRequest, + task_id: str | None = None, + ) -> bool: + del request, task_id + return False + + async def cancel_task(self, task_id: str) -> bool: + del task_id + return False + + +register_tool_scheduler("local", LocalToolScheduler) + + +@singleton +class CeleryToolScheduler(BaseToolScheduler): + """Dispatch tool calls to a dedicated Celery tools worker.""" + + @inject + def __init__(self, settings: Settings) -> None: + self._settings = settings + + @property + def is_async(self) -> bool: + return True + + async def execute( + self, + request: ToolExecutionRequest, + state_ctx: ChatState | None = None, + interceptors: list[ToolExecutionInterceptor] | None = None, + ) -> ToolExecutionResponse: + del request, state_ctx, interceptors + raise NotImplementedError( + "CeleryToolScheduler only supports async_execute() in resumable chat mode." + ) + + async def cancel( + self, + request: ToolExecutionRequest, + task_id: str | None = None, + ) -> bool: + del request + return await self.cancel_task(task_id) if task_id else False + + async def cancel_task(self, task_id: str) -> bool: + from private_gpt.celery.celery import celery_app + + celery_app.control.revoke(task_id, terminate=True) + return True + + async def async_execute( + self, + request: ToolExecutionRequest, + state_ctx: ChatState | None = None, + interceptors: list[ToolExecutionInterceptor] | None = None, + ) -> str: + del state_ctx, interceptors + + correlation_id = request.context.get("correlation_id") + task_id = f"{correlation_id}:{request.tool_id}" if correlation_id else None + result = dispatch_task( + task_name=TOOL_TASK_NAME, + kwargs={"request_data": request.model_dump(mode="json")}, + queue=self._settings.scheduler.tools.celery_queue, + task_id=task_id, + ) + return str(result.id) + + +register_tool_scheduler("celery", CeleryToolScheduler) + + +@singleton +class ToolSchedulerFactory: + @inject + def __init__(self, settings: Settings, injector: Injector) -> None: + self._settings = settings + self._injector = injector + self._scheduler: BaseToolScheduler | None = None + + def get(self) -> BaseToolScheduler: + if self._scheduler is None: + mode = self._settings.scheduler.tools.mode + provider = _TOOL_SCHEDULERS.get(mode) + if provider is None: + raise ValueError(f"Unknown scheduler.tools.mode: {mode}") + self._scheduler = ( + self._injector.get(provider) + if isinstance(provider, type) + else provider(self._injector) + ) + return self._scheduler diff --git a/private_gpt/components/vector_store/qdrant_client_builder.py b/private_gpt/components/vector_store/qdrant_client_builder.py index eab5d304..2d90fb5c 100644 --- a/private_gpt/components/vector_store/qdrant_client_builder.py +++ b/private_gpt/components/vector_store/qdrant_client_builder.py @@ -11,7 +11,6 @@ from qdrant_client import ( # type: ignore[import-not-found] from private_gpt.settings.settings import Settings logger = logging.getLogger(__name__) -logger.setLevel(logging.DEBUG) class QdrantClients: @@ -74,15 +73,15 @@ class QdrantClientBuilder: config = dict(config) config["path"] = str(resolve_data_path(config["path"])) - # This is a workaround to allow to execute tests/local qdrant - # when we want to support sync/async client. - # To allow, remove .lock from the db_dir + db_dir = config["path"] + # Local Qdrant uses one lock file for sync and async clients. + # Remove it between client construction using the resolved path. client = QdrantClient(**config) - QdrantClientBuilder.clean_lock(settings) + QdrantClientBuilder.clean_lock(db_dir=db_dir) aclient = AsyncQdrantClient(**config) - QdrantClientBuilder.clean_lock(settings) + QdrantClientBuilder.clean_lock(db_dir=db_dir) else: client = QdrantClient(**config) @@ -98,11 +97,7 @@ class QdrantClientBuilder: return config.get("path") is not None @staticmethod - def clean_lock(settings: Settings) -> None: - db_dir = settings.qdrant.path - if db_dir is None: - return - + def clean_lock(db_dir: str) -> None: lock_file = os.path.join(db_dir, ".lock") if os.path.exists(lock_file): os.remove(lock_file) diff --git a/private_gpt/di.py b/private_gpt/di.py index 715918d2..f2ed5a13 100644 --- a/private_gpt/di.py +++ b/private_gpt/di.py @@ -1,4 +1,6 @@ import asyncio +import contextlib +import inspect import logging import threading from asyncio import AbstractEventLoop @@ -70,7 +72,6 @@ def get_injector( ) return get_injector() except RuntimeError: - # Not in an asyncio loop, use global injector if _global_injector is None: with _global_injector_lock: global_injector = create_application_injector() @@ -90,29 +91,32 @@ def set_injector(injector: Injector) -> None: loop = asyncio.get_running_loop() setattr(loop, _INJECTOR_KEY, injector) except RuntimeError: - # Not in an asyncio loop, set global injector with _global_injector_lock: global _global_injector _global_injector = injector -def clean_global_injector(loop: AbstractEventLoop | None = None) -> None: +async def clean_global_injector(loop: AbstractEventLoop | None = None) -> None: try: loop = loop or asyncio.get_running_loop() - if hasattr(loop, _INJECTOR_KEY): - logger.debug("Closing loop injector resources...") - _injector = cast(Injector, getattr(loop, _INJECTOR_KEY)) - for interface, _ in _injector.binder._bindings.items(): - impl: Any = _injector.get(interface) + if not hasattr(loop, _INJECTOR_KEY): + return + logger.debug("Closing loop injector resources...") + injector = cast(Injector | None, getattr(loop, _INJECTOR_KEY, None)) + if injector is None: + return + bindings = getattr(injector.binder, "_bindings", {}) + for interface in list(bindings.keys()): + with contextlib.suppress(Exception): + impl: Any = injector.get(interface) if hasattr(impl, "close"): - impl.close() - - del impl - - with _loop_injector_lock: + res = impl.close() + if inspect.isawaitable(res): + await res + with _loop_injector_lock: + if hasattr(loop, _INJECTOR_KEY): delattr(loop, _INJECTOR_KEY) except RuntimeError: - # Not in an asyncio loop, do nothing pass diff --git a/private_gpt/eager_loading.py b/private_gpt/eager_loading.py new file mode 100644 index 00000000..c324f614 --- /dev/null +++ b/private_gpt/eager_loading.py @@ -0,0 +1,102 @@ +"""Profile-based eager-loading of injector-bound components. + +A single entry point ``warm(injector, profile)`` dispatches to the right +component groups based on the role. New worker types only need a new +profile entry — no changes anywhere else. +""" +import logging +from collections.abc import Callable, Sequence + +from injector import Injector + +from private_gpt.components.code_execution.code_execution_component import ( + CodeExecutionComponent, +) +from private_gpt.components.embedding.embedding_component import EmbeddingComponent +from private_gpt.components.llm.llm_component import LLMComponent +from private_gpt.components.node_store.node_store_component import NodeStoreComponent +from private_gpt.components.prompts.prompt_builder import PromptBuilderService +from private_gpt.components.streaming.stream.stream_manager import StreamManager +from private_gpt.components.streaming.stream_component import StreamComponent +from private_gpt.components.vector_store.vector_store_component import ( + VectorStoreComponent, +) +from private_gpt.server.tools.tool_service import ToolService +from private_gpt.settings.settings import Settings + +logger = logging.getLogger(__name__) + +_Warmer = Callable[[Injector], None] + + +def _warm_base(injector: Injector) -> None: + logger.debug("Warming base (settings + models)") + injector.get(Settings) + injector.get(LLMComponent) + injector.get(EmbeddingComponent) + + +def _warm_stores(injector: Injector) -> None: + logger.debug("Warming stores") + injector.get(NodeStoreComponent) + injector.get(VectorStoreComponent) + + +def _warm_streaming(injector: Injector) -> None: + logger.debug("Warming streaming components") + injector.get(StreamComponent) + injector.get(StreamManager) + + +def _warm_tools(injector: Injector) -> None: + logger.debug("Warming tools") + injector.get(PromptBuilderService) + injector.get(ToolService) + injector.get(CodeExecutionComponent) + + +def _warm_chat(injector: Injector) -> None: + from private_gpt.components.streaming.stream.stream_processor import ( + StreamProcessor, + ) + from private_gpt.server.chat.chat_service import ChatService + + logger.debug("Warming chat-path components") + injector.get(ChatService) + injector.get(StreamProcessor) + + +_GROUPS: dict[str, _Warmer] = { + "base": _warm_base, + "stores": _warm_stores, + "streaming": _warm_streaming, + "tools": _warm_tools, + "chat": _warm_chat, +} + +_PROFILES: dict[str, Sequence[str]] = { + "chat": ("base", "stores", "streaming", "tools", "chat"), + "tools": ("base", "tools"), + "full": ("base", "stores", "streaming", "tools"), +} + + +def warm(injector: Injector, profile: str = "full") -> None: + """Eager-resolve all DI singletons for the given worker profile. + + ``profile`` is one of ``chat``, ``tools``, or ``full`` (API server). + """ + groups = _PROFILES.get(profile) + if groups is None: + raise ValueError( + f"Unknown warm-up profile {profile!r}. " + f"Available: {', '.join(sorted(_PROFILES))}" + ) + logger.info("Warming profile=%s (groups: %s)", profile, ", ".join(groups)) + for group in groups: + _GROUPS[group](injector) + + +def eager_loading(injector: Injector) -> None: + """Full warm-up: everything. Kept for API server backward compat.""" + warm(injector, "full") diff --git a/private_gpt/events/utils.py b/private_gpt/events/utils.py index acbc588b..8154e7f8 100644 --- a/private_gpt/events/utils.py +++ b/private_gpt/events/utils.py @@ -1,4 +1,5 @@ from collections.abc import AsyncGenerator, AsyncIterator +from contextlib import aclosing from typing import Any from private_gpt.events.models import ( @@ -39,5 +40,6 @@ def to_message( async def to_sse_stream( event_generator: AsyncGenerator[Event, None], ) -> AsyncIterator[str]: - async for event in event_generator: - yield SSEFormatter.to_sse_event(event) + async with aclosing(event_generator): + async for event in event_generator: + yield SSEFormatter.to_sse_event(event) diff --git a/private_gpt/launcher.py b/private_gpt/launcher.py index a9c35350..2c0ffc4d 100644 --- a/private_gpt/launcher.py +++ b/private_gpt/launcher.py @@ -17,24 +17,13 @@ from injector import Injector from llama_index.core.embeddings import MockEmbedding from llama_index.core.settings import Settings as LlamaIndexSettings -from private_gpt.components.code_execution.code_execution_component import ( - CodeExecutionComponent, -) -from private_gpt.components.embedding.embedding_component import EmbeddingComponent -from private_gpt.components.llm.llm_component import LLMComponent -from private_gpt.components.node_store.node_store_component import NodeStoreComponent from private_gpt.components.persistence.persistence_component import ( PersistenceComponent, ) -from private_gpt.components.prompts.prompt_builder import PromptBuilderService -from private_gpt.components.streaming.stream.stream_manager import StreamManager -from private_gpt.components.streaming.stream_component import StreamComponent -from private_gpt.components.vector_store.vector_store_component import ( - VectorStoreComponent, -) from private_gpt.constants import PROJECT_ROOT_PATH from private_gpt.di import set_global_injector from private_gpt.docs import DESCRIPTION, TITLE, configure_openapi +from private_gpt.eager_loading import eager_loading from private_gpt.global_handler import ( ExceptionMiddleware, request_validation_exception_adapter, @@ -55,7 +44,6 @@ from private_gpt.server.models.models_router import models_router from private_gpt.server.primitives.primitives_router import primitives_router from private_gpt.server.skills.skill_router import skill_router from private_gpt.server.tools.tool_router import tool_router -from private_gpt.server.tools.tool_service import ToolService from private_gpt.settings.settings import Settings from private_gpt.utils.runner import get_version @@ -63,33 +51,6 @@ logger = logging.getLogger(__name__) UI_DIRECTORY = PROJECT_ROOT_PATH / "ui" -def eager_loading(injector: Injector) -> None: - """Eagerly load modules to avoid race conditions in multi-threaded environments.""" - logger.debug("Initializing mandatory dependencies") - injector.get(Settings) - - # Models - logger.debug("Initializing models") - injector.get(LLMComponent) - injector.get(EmbeddingComponent) - - # Stores - logger.debug("Initializing stores") - injector.get(NodeStoreComponent) - injector.get(VectorStoreComponent) - - # Streaming components - logger.debug("Initializing streaming components") - injector.get(StreamComponent) - injector.get(StreamManager) - - # Auxiliar - logger.debug("Initializing auxiliar services") - injector.get(PromptBuilderService) - injector.get(ToolService) - injector.get(CodeExecutionComponent) - - def apply_migrations(injector: Injector) -> None: """Ensure that all migrations are applied.""" logger.debug("Ensuring migrations are applied") @@ -115,12 +76,17 @@ def create_app(root_injector: Injector) -> FastAPI: # Set nested loop nest_asyncio.apply() - # Set default thread pool limit + # Set default thread pool limit. This executor now only serves genuine + # blocking-I/O offloads (broker waits, sync HTTP, sync file reads); all + # CPU-bound work is routed to dedicated workers, and chat can be routed + # to a long-lived external worker when ``scheduler.chat.mode`` is enabled, + # so a small I/O-only pool is enough + # and stops the GIL from being contended with the event loop. cpu_count = os.cpu_count() or 1 executor = concurrent.futures.ThreadPoolExecutor( max_workers=min(500, cpu_count * 50), thread_name_prefix="Stream-Pool" ) - asyncio.get_event_loop().set_default_executor(executor) + asyncio.get_running_loop().set_default_executor(executor) # Set the global injector as loop injector. set_global_injector(root_injector) diff --git a/private_gpt/server/chat/chat_facade.py b/private_gpt/server/chat/chat_facade.py index de96aa56..dd83ec4d 100644 --- a/private_gpt/server/chat/chat_facade.py +++ b/private_gpt/server/chat/chat_facade.py @@ -31,10 +31,15 @@ class ChatFacadeService: """Create the chat event generator.""" async def coro() -> AsyncGenerator[Event, None]: - # Open the chat service and start streaming completion_gen = await self._chat_service.stream_chat(request) - event_generator: AsyncGenerator[Event, None] = completion_gen.events # type: ignore - async for event in event_generator: - yield event + event_generator = completion_gen.events + try: + async for event in event_generator: + yield event + finally: + await event_generator.aclose() return coro() + + async def cancel(self, correlation_id: str) -> bool: + return await self._chat_service.cancel(correlation_id) diff --git a/private_gpt/server/chat/chat_router.py b/private_gpt/server/chat/chat_router.py index c304cd8e..bea05187 100644 --- a/private_gpt/server/chat/chat_router.py +++ b/private_gpt/server/chat/chat_router.py @@ -258,10 +258,10 @@ async def chat_messages( chat_facade_service: ChatAsyncFacadeService = request.state.injector.get( ChatAsyncFacadeService ) - request_mapper: ChatRequestMapper = request.state.injector.get(ChatRequestMapper) return await chat_facade_service.chat( - request, await request_mapper.create_request_from_body(body) + request, + body, ) diff --git a/private_gpt/server/chat/chat_service.py b/private_gpt/server/chat/chat_service.py index 01d86313..6eb5ae05 100644 --- a/private_gpt/server/chat/chat_service.py +++ b/private_gpt/server/chat/chat_service.py @@ -14,15 +14,28 @@ from private_gpt.components.chat.models.chat_config_models import ( from private_gpt.components.chunk.models import SourceType from private_gpt.components.container_registry import ContainerRegistry from private_gpt.components.embedding.embedding_component import EmbeddingComponent -from private_gpt.components.engines.chat_loop.chat_loop_engine import ChatLoopEngine -from private_gpt.components.engines.chat_loop.models.chat_loop_phase import ( - InterceptorPhase, +from private_gpt.components.engines.chat.async_chat_engine import ( + AsyncChatEngine, +) +from private_gpt.components.engines.chat.chat_engine import ChatLoopEngine +from private_gpt.components.engines.chat.chat_engine_interface import ( + ChatEngine, + LoopChatEngineAdapter, +) +from private_gpt.components.engines.chat.chat_runner import ( + ChatRunner, + ChatRunnerFactory, +) +from private_gpt.components.engines.chat.models.execution_hooks import ( + ExecutionHooks, ) from private_gpt.components.ingest.ingest_component import IngestComponent from private_gpt.components.llm.custom.base import ZylonLLM from private_gpt.components.llm.llm_component import LLMComponent from private_gpt.components.llm.models import ReasoningEffort from private_gpt.components.node_store.node_store_component import NodeStoreComponent +from private_gpt.components.streaming.tasks.chat_scheduler import ChatSchedulerFactory +from private_gpt.components.tools.tool_scheduler import ToolSchedulerFactory from private_gpt.components.vector_store.vector_store_component import ( VectorStoreComponent, ) @@ -114,6 +127,7 @@ class Completion(BaseModel): class CompletionGen(BaseModel): events: AsyncGenerator[Event, None] + final_state_task: Any | None = None class Config: arbitrary_types_allowed = True @@ -150,6 +164,9 @@ class ChatService: chat_interceptor_service: ChatInterceptorService, models_service: ModelsService, container_registry: ContainerRegistry, + scheduler_factory: ToolSchedulerFactory, + chat_scheduler_factory: ChatSchedulerFactory, + runner_factory: ChatRunnerFactory, ) -> None: self.settings = settings self.llm_component = llm_component @@ -160,30 +177,70 @@ class ChatService: self.chat_interceptor_service = chat_interceptor_service self.models_service = models_service self.container_registry = container_registry + self._tool_scheduler = scheduler_factory.get() + self._chat_scheduler = chat_scheduler_factory.get() + self._runner: ChatRunner = runner_factory.get() - def _build_loop_engine(self) -> ChatLoopEngine: + def build_async_engine(self) -> AsyncChatEngine: # Don't build a singleton since the interceptors # and state can be mutated during the loop execution # for several threads handling different requests in parallel chain = self.chat_interceptor_service.get_chain() - return ChatLoopEngine( + return AsyncChatEngine( llm_component=self.llm_component, request_interceptors=chain.request_interceptors, response_interceptors=chain.response_interceptors, + tool_interceptors=chain.tool_interceptors, max_iterations=40, container_registry=self.container_registry, + tool_scheduler=self._tool_scheduler, + chat_scheduler=self._chat_scheduler, + runner=self._runner, + ) + + def build_loop_engine(self) -> LoopChatEngineAdapter: + chain = self.chat_interceptor_service.get_chain() + return LoopChatEngineAdapter( + engine=ChatLoopEngine( + llm_component=self.llm_component, + request_interceptors=chain.request_interceptors, + response_interceptors=chain.response_interceptors, + tool_interceptors=chain.tool_interceptors, + max_iterations=40, + container_registry=self.container_registry, + tool_scheduler=self._tool_scheduler, + ) + ) + + def build_engine(self) -> ChatEngine: + return ( + self.build_async_engine() + if self.settings.chat.engine_mode == "async" + else self.build_loop_engine() ) async def stream_chat( self, request: ChatRequest, + hooks: ExecutionHooks | None = None, ) -> CompletionGen: - loop_engine = await asyncio.to_thread(self._build_loop_engine) - event_generator: AsyncGenerator[Event, None] = loop_engine.run(request) - return CompletionGen(events=event_generator) + engine = await asyncio.to_thread(self.build_engine) + execution = await engine.run( + request=request, + hooks=hooks, + runner=self._runner, + ) + return CompletionGen( + events=execution.events, + final_state_task=execution.final_state_task, + ) - async def chat(self, request: ChatRequest) -> Completion: - completion_gen = await self.stream_chat(request) + async def chat( + self, + request: ChatRequest, + hooks: ExecutionHooks | None = None, + ) -> Completion: + completion_gen = await self.stream_chat(request, hooks=hooks) wrapped_response = await fold(completion_gen.events) usage = Usage.model_validate(wrapped_response.usage or {}) return Completion( @@ -193,16 +250,16 @@ class ChatService: usage=usage, ) + async def cancel(self, correlation_id: str) -> bool: + engine = await asyncio.to_thread(self.build_engine) + return await engine.cancel(correlation_id=correlation_id) + async def validate(self, request: ResolvedChatRequest) -> ChatValidationResult: errors: list[str] = [] try: - loop_engine = await asyncio.to_thread(self._build_loop_engine) - await loop_engine.run_interceptor_phase( - run=loop_engine.initialize_run(request), - phase=InterceptorPhase.VALIDATION, - interceptors=loop_engine._request_interceptors, - ) + engine = await asyncio.to_thread(self.build_engine) + await engine.validate(request=request) except ValidationError as e: for err in e.errors(): field = " -> ".join(str(loc) for loc in err["loc"]) diff --git a/private_gpt/server/chat/interceptors/chat_interceptor_service.py b/private_gpt/server/chat/interceptors/chat_interceptor_service.py index 9f8b810f..f11c8e16 100644 --- a/private_gpt/server/chat/interceptors/chat_interceptor_service.py +++ b/private_gpt/server/chat/interceptors/chat_interceptor_service.py @@ -1,9 +1,9 @@ from injector import inject, singleton -from private_gpt.components.engines.chat_loop.interceptors.chat_loop_interceptor_chain import ( +from private_gpt.components.engines.chat.interceptors.chat_interceptor_chain import ( ChatLoopInterceptorChain, ) -from private_gpt.components.engines.chat_loop.interceptors.restore_stateless_input_interceptor import ( +from private_gpt.components.engines.chat.interceptors.restore_stateless_input_interceptor import ( RestoreStatelessInputInterceptorRequest, ) from private_gpt.components.prompts.prompt_builder import PromptBuilderService @@ -13,6 +13,9 @@ from private_gpt.server.chat.interceptors.citation_interceptor import ( from private_gpt.server.chat.interceptors.condensation_interceptor import ( CondensationRequestInterceptor, ) +from private_gpt.server.chat.interceptors.configure_tool_execution_interceptor import ( + ConfigureToolExecutionInterceptor, +) from private_gpt.server.chat.interceptors.configure_tool_interceptor import ( ConfigureToolRequestInterceptor, ) @@ -88,6 +91,8 @@ class ChatInterceptorService: skill_tool_visibility_interceptor: SkillToolVisibilityInterceptor, tool_choice_interceptor: ToolChoiceRequestInterceptor, configure_tool_interceptor: ConfigureToolRequestInterceptor, + # --- tool execution interceptors --- + configure_tool_execution_interceptor: ConfigureToolExecutionInterceptor, # --- loop interceptors (run each iteration, order matters) --- document_file_interceptor: DocumentFilePreprocessingInterceptor, multimodal_interceptor: MultimodalRequestInterceptor, @@ -122,6 +127,7 @@ class ChatInterceptorService: configure_tool_interceptor, platform_guidelines_interceptor, ], + tools=[configure_tool_execution_interceptor], ) # Preprocess the chat history .add_range( diff --git a/private_gpt/server/chat/interceptors/citation_interceptor.py b/private_gpt/server/chat/interceptors/citation_interceptor.py index 37be678b..f4f81e0e 100644 --- a/private_gpt/server/chat/interceptors/citation_interceptor.py +++ b/private_gpt/server/chat/interceptors/citation_interceptor.py @@ -4,13 +4,13 @@ from injector import singleton from private_gpt.components.context.models.context_layer import DocumentLayer from private_gpt.components.context.models.layer_type import LayerType -from private_gpt.components.engines.chat_loop.interceptors.chat_loop_interceptor import ( +from private_gpt.components.engines.chat.interceptors.chat_interceptor import ( ChatRequestLoopInterceptor, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_interceptor_context import ( - ChatLoopInterceptorContext, +from private_gpt.components.engines.chat.models.chat_interceptor_context import ( + ChatInterceptorContext, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_phase import ( +from private_gpt.components.engines.chat.models.chat_phase import ( InterceptorPhase, ) from private_gpt.components.engines.citations.utils import ( @@ -26,7 +26,7 @@ if TYPE_CHECKING: class CitationRequestInterceptor(ChatRequestLoopInterceptor): """Populate documents and citations from chat history.""" - async def intercept(self, context: ChatLoopInterceptorContext) -> None: + async def intercept(self, context: ChatInterceptorContext) -> None: """Extract citations and source documents from conversation history.""" if context.phase != InterceptorPhase.BEFORE_ITERATION: return diff --git a/private_gpt/server/chat/interceptors/condensation_interceptor.py b/private_gpt/server/chat/interceptors/condensation_interceptor.py index e7e614a4..32739a40 100644 --- a/private_gpt/server/chat/interceptors/condensation_interceptor.py +++ b/private_gpt/server/chat/interceptors/condensation_interceptor.py @@ -12,13 +12,13 @@ from private_gpt.components.chat.processors.chat_history.memory.tldr_processor i CondenseResponse, condense_chat_history, ) -from private_gpt.components.engines.chat_loop.interceptors.chat_loop_interceptor import ( +from private_gpt.components.engines.chat.interceptors.chat_interceptor import ( ChatRequestLoopInterceptor, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_interceptor_context import ( - ChatLoopInterceptorContext, +from private_gpt.components.engines.chat.models.chat_interceptor_context import ( + ChatInterceptorContext, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_phase import ( +from private_gpt.components.engines.chat.models.chat_phase import ( InterceptorPhase, ) from private_gpt.events.models import ( @@ -156,7 +156,7 @@ class CondensationRequestInterceptor(ChatRequestLoopInterceptor): self._condensation_timeout = settings.chat.tldr_timeout self._min_duration = settings.chat.tldr_minimum_threshold_seconds - async def intercept(self, context: ChatLoopInterceptorContext) -> None: + async def intercept(self, context: ChatInterceptorContext) -> None: if context.phase != InterceptorPhase.BEFORE_ITERATION: return diff --git a/private_gpt/server/chat/interceptors/configure_tool_execution_interceptor.py b/private_gpt/server/chat/interceptors/configure_tool_execution_interceptor.py new file mode 100644 index 00000000..be3d72fb --- /dev/null +++ b/private_gpt/server/chat/interceptors/configure_tool_execution_interceptor.py @@ -0,0 +1,32 @@ +from injector import inject, singleton + +from private_gpt.components.tools.remote_execution import ( + ToolExecutionInterceptor, + ToolExecutionInterceptorContext, +) +from private_gpt.server.chat.interceptors.null_tool_values_interceptor import ( + NullToolValuesRequestInterceptor, +) +from private_gpt.server.chat.interceptors.schema_coercing_tool_interceptor import ( + SchemaCoercingToolInterceptor, +) + + +@singleton +class ConfigureToolExecutionInterceptor(ToolExecutionInterceptor): + """Aggregate tool-execution sub-interceptors into a single step.""" + + @inject + def __init__( + self, + null_tool_values_interceptor: NullToolValuesRequestInterceptor, + schema_coercing_interceptor: SchemaCoercingToolInterceptor, + ) -> None: + self._interceptors: list[ToolExecutionInterceptor] = [ + null_tool_values_interceptor, + schema_coercing_interceptor, + ] + + async def intercept(self, context: ToolExecutionInterceptorContext) -> None: + for interceptor in self._interceptors: + await interceptor.intercept(context) diff --git a/private_gpt/server/chat/interceptors/configure_tool_interceptor.py b/private_gpt/server/chat/interceptors/configure_tool_interceptor.py index 8933a2fe..9cf4cec3 100644 --- a/private_gpt/server/chat/interceptors/configure_tool_interceptor.py +++ b/private_gpt/server/chat/interceptors/configure_tool_interceptor.py @@ -1,12 +1,12 @@ from injector import inject, singleton -from private_gpt.components.engines.chat_loop.interceptors.chat_loop_interceptor import ( +from private_gpt.components.engines.chat.interceptors.chat_interceptor import ( ChatRequestLoopInterceptor, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_interceptor_context import ( - ChatLoopInterceptorContext, +from private_gpt.components.engines.chat.models.chat_interceptor_context import ( + ChatInterceptorContext, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_phase import ( +from private_gpt.components.engines.chat.models.chat_phase import ( InterceptorPhase, ) from private_gpt.server.chat.interceptors.null_tool_values_interceptor import ( @@ -32,7 +32,7 @@ class ConfigureToolRequestInterceptor(ChatRequestLoopInterceptor): schema_coercing_interceptor, ] - async def intercept(self, context: ChatLoopInterceptorContext) -> None: + async def intercept(self, context: ChatInterceptorContext) -> None: """Run all tool-configuration interceptors in order.""" if context.phase != InterceptorPhase.BEFORE_ITERATION: return diff --git a/private_gpt/server/chat/interceptors/default_values_interceptor.py b/private_gpt/server/chat/interceptors/default_values_interceptor.py index 4e9395a7..f59bf923 100644 --- a/private_gpt/server/chat/interceptors/default_values_interceptor.py +++ b/private_gpt/server/chat/interceptors/default_values_interceptor.py @@ -3,13 +3,13 @@ import logging from injector import inject, singleton from private_gpt.chat.input_models import PromptConfig -from private_gpt.components.engines.chat_loop.interceptors.chat_loop_interceptor import ( +from private_gpt.components.engines.chat.interceptors.chat_interceptor import ( ChatRequestLoopInterceptor, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_interceptor_context import ( - ChatLoopInterceptorContext, +from private_gpt.components.engines.chat.models.chat_interceptor_context import ( + ChatInterceptorContext, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_phase import ( +from private_gpt.components.engines.chat.models.chat_phase import ( InterceptorPhase, ) from private_gpt.settings.settings import Settings @@ -23,7 +23,7 @@ class DefaultValuesRequestInterceptor(ChatRequestLoopInterceptor): def __init__(self, settings: Settings) -> None: self._settings = settings - async def intercept(self, context: ChatLoopInterceptorContext) -> None: + async def intercept(self, context: ChatInterceptorContext) -> None: if context.phase != InterceptorPhase.VALIDATION: return diff --git a/private_gpt/server/chat/interceptors/document_file_interceptor.py b/private_gpt/server/chat/interceptors/document_file_interceptor.py index 2ba1cf6b..26e1c59b 100644 --- a/private_gpt/server/chat/interceptors/document_file_interceptor.py +++ b/private_gpt/server/chat/interceptors/document_file_interceptor.py @@ -8,13 +8,13 @@ from llama_index.core.llms.llm import ToolSelection from private_gpt.components.chat.processors.chat_history.documents.document_preprocessor import ( preprocess_document_history, ) -from private_gpt.components.engines.chat_loop.interceptors.chat_loop_interceptor import ( +from private_gpt.components.engines.chat.interceptors.chat_interceptor import ( ChatRequestLoopInterceptor, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_interceptor_context import ( - ChatLoopInterceptorContext, +from private_gpt.components.engines.chat.models.chat_interceptor_context import ( + ChatInterceptorContext, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_phase import ( +from private_gpt.components.engines.chat.models.chat_phase import ( InterceptorPhase, ) from private_gpt.events.models import ( @@ -43,7 +43,7 @@ class DocumentFilePreprocessingInterceptor(ChatRequestLoopInterceptor): self._tool_name = DOCUMENT_PROCESSING_TOOL_NAME self._preprocess_settings = settings.chat.preprocess.documents - async def intercept(self, context: ChatLoopInterceptorContext) -> None: + async def intercept(self, context: ChatInterceptorContext) -> None: """Convert document blocks to text before inference.""" if context.phase != InterceptorPhase.BEFORE_ITERATION: return diff --git a/private_gpt/server/chat/interceptors/document_processing_interceptor.py b/private_gpt/server/chat/interceptors/document_processing_interceptor.py index 829023e6..93cb0956 100644 --- a/private_gpt/server/chat/interceptors/document_processing_interceptor.py +++ b/private_gpt/server/chat/interceptors/document_processing_interceptor.py @@ -3,13 +3,13 @@ from injector import inject, singleton from private_gpt.components.chat.processors.chat_history.documents.citations import ( process_chat_history_with_documents, ) -from private_gpt.components.engines.chat_loop.interceptors.chat_loop_interceptor import ( +from private_gpt.components.engines.chat.interceptors.chat_interceptor import ( ChatRequestLoopInterceptor, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_interceptor_context import ( - ChatLoopInterceptorContext, +from private_gpt.components.engines.chat.models.chat_interceptor_context import ( + ChatInterceptorContext, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_phase import ( +from private_gpt.components.engines.chat.models.chat_phase import ( InterceptorPhase, ) @@ -24,7 +24,7 @@ class DocumentProcessingRequestInterceptor(ChatRequestLoopInterceptor): def __init__(self, add_context_to_system_prompt: bool | None = None) -> None: self._add_context_to_system_prompt = add_context_to_system_prompt - async def intercept(self, context: ChatLoopInterceptorContext) -> None: + async def intercept(self, context: ChatInterceptorContext) -> None: """Render document context into conversation history when configured.""" if context.phase != InterceptorPhase.BEFORE_ITERATION: return diff --git a/private_gpt/server/chat/interceptors/extract_citation_interceptor.py b/private_gpt/server/chat/interceptors/extract_citation_interceptor.py index be2f4c12..6fff568f 100644 --- a/private_gpt/server/chat/interceptors/extract_citation_interceptor.py +++ b/private_gpt/server/chat/interceptors/extract_citation_interceptor.py @@ -2,11 +2,11 @@ import asyncio from collections.abc import Mapping from typing import TYPE_CHECKING, Any -from private_gpt.components.engines.chat_loop.interceptors.chat_loop_interceptor import ( +from private_gpt.components.engines.chat.interceptors.chat_interceptor import ( ChatResponseLoopInterceptor, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_interceptor_context import ( - ChatLoopInterceptorContext, +from private_gpt.components.engines.chat.models.chat_interceptor_context import ( + ChatInterceptorContext, ) from private_gpt.components.engines.citations.utils import ( extract_citations_by_original_text, @@ -31,12 +31,12 @@ class ExtractCitationInterceptor(ChatResponseLoopInterceptor): self._documents: list[Document] = [] self._citation_indices: dict[str, int] = {} - async def on_iteration_start(self, context: ChatLoopInterceptorContext) -> None: + async def on_iteration_start(self, context: ChatInterceptorContext) -> None: self._send_text = "" self._send_citations = [] self._current_text = "" - async def on_iteration_end(self, context: ChatLoopInterceptorContext) -> None: + async def on_iteration_end(self, context: ChatInterceptorContext) -> None: # Mutate the context state to include the final citations for this iteration if self._send_citations: new_state = context.state.model_copy(deep=True) @@ -52,7 +52,7 @@ class ExtractCitationInterceptor(ChatResponseLoopInterceptor): async def intercept_event( self, event: Event, - context: ChatLoopInterceptorContext, + context: ChatInterceptorContext, ) -> Event | None: citations_is_enabled = context.state.input.request.citation.enabled documents = self._documents or context.state.input.context_stack.all_documents() diff --git a/private_gpt/server/chat/interceptors/filter_event_by_type_interceptor.py b/private_gpt/server/chat/interceptors/filter_event_by_type_interceptor.py index 199ce8b9..3910a736 100644 --- a/private_gpt/server/chat/interceptors/filter_event_by_type_interceptor.py +++ b/private_gpt/server/chat/interceptors/filter_event_by_type_interceptor.py @@ -1,11 +1,11 @@ from collections.abc import Mapping from typing import Any, Literal -from private_gpt.components.engines.chat_loop.interceptors.chat_loop_interceptor import ( +from private_gpt.components.engines.chat.interceptors.chat_interceptor import ( ChatResponseLoopInterceptor, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_interceptor_context import ( - ChatLoopInterceptorContext, +from private_gpt.components.engines.chat.models.chat_interceptor_context import ( + ChatInterceptorContext, ) from private_gpt.events.models import ( Event, @@ -18,16 +18,16 @@ class FilterZylonInterceptor(ChatResponseLoopInterceptor): def __init__(self) -> None: self._active_blocks: set[str] = set() - async def on_iteration_start(self, context: ChatLoopInterceptorContext) -> None: + async def on_iteration_start(self, context: ChatInterceptorContext) -> None: self._active_blocks.clear() - async def on_iteration_end(self, context: ChatLoopInterceptorContext) -> None: + async def on_iteration_end(self, context: ChatInterceptorContext) -> None: self._active_blocks.clear() async def intercept_event( self, event: Event, - context: ChatLoopInterceptorContext, + context: ChatInterceptorContext, ) -> Event | None: response_format: Literal["zylon", "anthropic"] = ( "zylon" diff --git a/private_gpt/server/chat/interceptors/internal_tools_interceptor.py b/private_gpt/server/chat/interceptors/internal_tools_interceptor.py index e92920e4..352978c6 100644 --- a/private_gpt/server/chat/interceptors/internal_tools_interceptor.py +++ b/private_gpt/server/chat/interceptors/internal_tools_interceptor.py @@ -4,16 +4,16 @@ from private_gpt.components.context.models.context_layer import ( ToolDefinitionsLayer, ) from private_gpt.components.context.models.layer_type import LayerType -from private_gpt.components.engines.chat_loop.interceptors.chat_loop_interceptor import ( +from private_gpt.components.engines.chat.interceptors.chat_interceptor import ( ChatRequestLoopInterceptor, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_interceptor_context import ( - ChatLoopInterceptorContext, +from private_gpt.components.engines.chat.models.chat_interceptor_context import ( + ChatInterceptorContext, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_phase import ( +from private_gpt.components.engines.chat.models.chat_phase import ( InterceptorPhase, ) -from private_gpt.components.engines.chat_loop.utils.request_builder import ( +from private_gpt.components.engines.chat.utils.request_builder import ( build_request_from_context_stack, ) from private_gpt.components.prompts.prompt_builder import PromptBuilderService @@ -31,7 +31,7 @@ class InternalToolRequestInterceptor(ChatRequestLoopInterceptor): self._tool_pipeline = tool_pipeline self._prompt_builder = prompt_builder - async def intercept(self, context: ChatLoopInterceptorContext) -> None: + async def intercept(self, context: ChatInterceptorContext) -> None: if ( context.phase != InterceptorPhase.VALIDATION and context.phase != InterceptorPhase.BEFORE_ITERATION diff --git a/private_gpt/server/chat/interceptors/mcp_interceptor.py b/private_gpt/server/chat/interceptors/mcp_interceptor.py index ee07b6a2..211e195c 100644 --- a/private_gpt/server/chat/interceptors/mcp_interceptor.py +++ b/private_gpt/server/chat/interceptors/mcp_interceptor.py @@ -8,13 +8,13 @@ from private_gpt.components.chat.models.chat_config_models import ( ToolSpec, ) from private_gpt.components.context.models.context_layer import ToolDefinitionsLayer -from private_gpt.components.engines.chat_loop.interceptors.chat_loop_interceptor import ( +from private_gpt.components.engines.chat.interceptors.chat_interceptor import ( ChatRequestLoopInterceptor, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_interceptor_context import ( - ChatLoopInterceptorContext, +from private_gpt.components.engines.chat.models.chat_interceptor_context import ( + ChatInterceptorContext, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_phase import ( +from private_gpt.components.engines.chat.models.chat_phase import ( InterceptorPhase, ) from private_gpt.events.event_errors import Errors @@ -98,7 +98,7 @@ class McpRequestInterceptor(ChatRequestLoopInterceptor): raise e return [] - async def intercept(self, context: ChatLoopInterceptorContext) -> None: + async def intercept(self, context: ChatInterceptorContext) -> None: if ( context.phase != InterceptorPhase.VALIDATION and context.phase != InterceptorPhase.BEFORE_ITERATION diff --git a/private_gpt/server/chat/interceptors/multimodal_interceptor.py b/private_gpt/server/chat/interceptors/multimodal_interceptor.py index 84b47f2e..fd9b57a7 100644 --- a/private_gpt/server/chat/interceptors/multimodal_interceptor.py +++ b/private_gpt/server/chat/interceptors/multimodal_interceptor.py @@ -9,17 +9,17 @@ from llama_index.core.llms.llm import ToolSelection from private_gpt.components.chat.processors.chat_history.multimodality.multimodality_preprocessor import ( preprocess_multimodal_history, ) -from private_gpt.components.engines.chat_loop.interceptors.chat_loop_interceptor import ( +from private_gpt.components.engines.chat.interceptors.chat_interceptor import ( ChatRequestLoopInterceptor, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_interceptor_context import ( - ChatLoopInterceptorContext, +from private_gpt.components.engines.chat.models.chat_interceptor_context import ( + ChatInterceptorContext, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_phase import ( +from private_gpt.components.engines.chat.models.chat_phase import ( InterceptorPhase, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_state import ( - ChatLoopState, +from private_gpt.components.engines.chat.models.chat_state import ( + ChatState, ) from private_gpt.components.llm.llm_component import LLMComponent from private_gpt.components.llm.llm_helper import supports_audio, supports_images @@ -48,7 +48,7 @@ class MultimodalRequestInterceptor(ChatRequestLoopInterceptor): self._tool_name = MULTIMODAL_TOOL_NAME self._preprocess_settings = settings.chat.preprocess.multimodal - async def intercept(self, context: ChatLoopInterceptorContext) -> None: + async def intercept(self, context: ChatInterceptorContext) -> None: """Apply multimodal preprocessing to the current chat history.""" if context.phase != InterceptorPhase.BEFORE_ITERATION: return @@ -157,7 +157,7 @@ class MultimodalRequestInterceptor(ChatRequestLoopInterceptor): def resolve_multimodal_models( self, - state: ChatLoopState, + state: ChatState, main_llm: LLM, ) -> tuple[LLM | None, LLM | None]: """Resolve optional multimodal models using configured LLM registry.""" diff --git a/private_gpt/server/chat/interceptors/null_tool_values_interceptor.py b/private_gpt/server/chat/interceptors/null_tool_values_interceptor.py index 30ce1600..6065756d 100644 --- a/private_gpt/server/chat/interceptors/null_tool_values_interceptor.py +++ b/private_gpt/server/chat/interceptors/null_tool_values_interceptor.py @@ -1,50 +1,46 @@ -from typing import Any +from __future__ import annotations + +from typing import TYPE_CHECKING from injector import singleton -from private_gpt.components.chat.models.chat_config_models import ToolSpec -from private_gpt.components.context.models.context_layer import ToolDefinitionsLayer -from private_gpt.components.context.models.layer_type import LayerType -from private_gpt.components.engines.chat_loop.interceptors.chat_loop_interceptor import ( +from private_gpt.components.engines.chat.interceptors.chat_interceptor import ( ChatRequestLoopInterceptor, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_interceptor_context import ( - ChatLoopInterceptorContext, -) -from private_gpt.components.engines.chat_loop.models.chat_loop_phase import ( +from private_gpt.components.engines.chat.models.chat_phase import ( InterceptorPhase, ) +from private_gpt.components.tools.remote_execution import ( + ToolExecutionInterceptor, + ToolExecutionInterceptorContext, +) + +if TYPE_CHECKING: + from private_gpt.components.engines.chat.models.chat_interceptor_context import ( + ChatInterceptorContext, + ) @singleton -class NullToolValuesRequestInterceptor(ChatRequestLoopInterceptor): - """Patch tool specs to strip None kwargs before async invocation.""" +class NullToolValuesRequestInterceptor( + ChatRequestLoopInterceptor, + ToolExecutionInterceptor, +): + """Strip ``None``-valued kwargs before tool execution.""" - async def intercept(self, context: ChatLoopInterceptorContext) -> None: - if context.phase != InterceptorPhase.BEFORE_ITERATION: + async def intercept( + self, + context: ChatInterceptorContext | ToolExecutionInterceptorContext, + ) -> None: + if not isinstance(context, ToolExecutionInterceptorContext): + return + if context.phase != InterceptorPhase.BEFORE_TOOL: return - state = context.state - tools = [ - self._patch_tool(tool) for tool in state.input.context_stack.all_tools() - ] - - stack = state.input.context_stack - stack = stack.remove_layers_of_type(LayerType.TOOL_DEFINITIONS) - if tools: - stack = stack.append_layer( - ToolDefinitionsLayer(tools=tools, source="tool_patch") - ) - state.input.context_stack = stack - context.set_state(state) - - def _patch_tool(self, tool: ToolSpec) -> ToolSpec: - """Patch the tool to avoid to call using optional None params.""" - - async def async_patched_tool(*args: Any, **kwargs: Any) -> Any: - new_kwargs = {k: v for k, v in kwargs.items() if v is not None} - return await tool.async_fn(*args, **new_kwargs) - - tool_copy = tool.model_copy(deep=True) - tool_copy.async_fn = async_patched_tool - return tool_copy + context.set_tool_kwargs( + { + key: value + for key, value in context.tool_kwargs.items() + if value is not None + } + ) diff --git a/private_gpt/server/chat/interceptors/ping_loop_interceptor.py b/private_gpt/server/chat/interceptors/ping_loop_interceptor.py index 44ee26dd..d5540f13 100644 --- a/private_gpt/server/chat/interceptors/ping_loop_interceptor.py +++ b/private_gpt/server/chat/interceptors/ping_loop_interceptor.py @@ -3,11 +3,11 @@ import contextlib from collections.abc import Mapping from typing import Any -from private_gpt.components.engines.chat_loop.interceptors.chat_loop_interceptor import ( +from private_gpt.components.engines.chat.interceptors.chat_interceptor import ( ChatResponseLoopInterceptor, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_interceptor_context import ( - ChatLoopInterceptorContext, +from private_gpt.components.engines.chat.models.chat_interceptor_context import ( + ChatInterceptorContext, ) from private_gpt.events.models import Event, PingEvent @@ -19,7 +19,7 @@ class PingInterceptor(ChatResponseLoopInterceptor): self._ping_interval = ping_interval self._ping_task: asyncio.Task[None] | None = None - async def on_iteration_start(self, context: ChatLoopInterceptorContext) -> None: + async def on_iteration_start(self, context: ChatInterceptorContext) -> None: emit_fn = context.emit_fn async def _ping_loop() -> None: @@ -30,7 +30,7 @@ class PingInterceptor(ChatResponseLoopInterceptor): self._ping_task = asyncio.create_task(_ping_loop()) - async def on_iteration_end(self, context: ChatLoopInterceptorContext) -> None: + async def on_iteration_end(self, context: ChatInterceptorContext) -> None: if self._ping_task and not self._ping_task.done(): self._ping_task.cancel() with contextlib.suppress(asyncio.CancelledError): @@ -40,7 +40,7 @@ class PingInterceptor(ChatResponseLoopInterceptor): async def intercept_event( self, event: Event, - context: ChatLoopInterceptorContext, + context: ChatInterceptorContext, ) -> Event | None: return event diff --git a/private_gpt/server/chat/interceptors/platform_guidelines_interceptor.py b/private_gpt/server/chat/interceptors/platform_guidelines_interceptor.py index f586039c..677362b7 100644 --- a/private_gpt/server/chat/interceptors/platform_guidelines_interceptor.py +++ b/private_gpt/server/chat/interceptors/platform_guidelines_interceptor.py @@ -8,13 +8,13 @@ from private_gpt.components.context.models.context_layer import ( ToolInstructionsLayer, ) from private_gpt.components.context.models.context_stack import ContextStack -from private_gpt.components.engines.chat_loop.interceptors.chat_loop_interceptor import ( +from private_gpt.components.engines.chat.interceptors.chat_interceptor import ( ChatRequestLoopInterceptor, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_interceptor_context import ( - ChatLoopInterceptorContext, +from private_gpt.components.engines.chat.models.chat_interceptor_context import ( + ChatInterceptorContext, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_phase import ( +from private_gpt.components.engines.chat.models.chat_phase import ( InterceptorPhase, ) from private_gpt.components.engines.citations.types import Document @@ -61,7 +61,7 @@ class PlatformGuidelinesInterceptor(ChatRequestLoopInterceptor): self._thinking_content: str | None = None self._citation_guidelines_content: str | None = None - async def intercept(self, context: ChatLoopInterceptorContext) -> None: + async def intercept(self, context: ChatInterceptorContext) -> None: if context.phase != InterceptorPhase.BEFORE_ITERATION: return diff --git a/private_gpt/server/chat/interceptors/schema_coercing_tool_interceptor.py b/private_gpt/server/chat/interceptors/schema_coercing_tool_interceptor.py index f51994f7..e09baa8f 100644 --- a/private_gpt/server/chat/interceptors/schema_coercing_tool_interceptor.py +++ b/private_gpt/server/chat/interceptors/schema_coercing_tool_interceptor.py @@ -1,24 +1,29 @@ +from __future__ import annotations + import ast import contextlib import json import logging import math -from typing import Any +from typing import TYPE_CHECKING, Any from injector import singleton -from private_gpt.components.chat.models.chat_config_models import ToolSpec -from private_gpt.components.context.models.context_layer import ToolDefinitionsLayer -from private_gpt.components.context.models.layer_type import LayerType -from private_gpt.components.engines.chat_loop.interceptors.chat_loop_interceptor import ( +from private_gpt.components.engines.chat.interceptors.chat_interceptor import ( ChatRequestLoopInterceptor, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_interceptor_context import ( - ChatLoopInterceptorContext, -) -from private_gpt.components.engines.chat_loop.models.chat_loop_phase import ( +from private_gpt.components.engines.chat.models.chat_phase import ( InterceptorPhase, ) +from private_gpt.components.tools.remote_execution import ( + ToolExecutionInterceptor, + ToolExecutionInterceptorContext, +) + +if TYPE_CHECKING: + from private_gpt.components.engines.chat.models.chat_interceptor_context import ( + ChatInterceptorContext, + ) logger = logging.getLogger(__name__) @@ -86,10 +91,6 @@ def _parse_literal_string( raw: str, expected: type | tuple[type, ...], ) -> Any: - """Parse a string as JSON, falling back to ast.literal_eval for Python repr. - - Returns the parsed value if it is an instance of `expected`, else None. - """ with contextlib.suppress(json.JSONDecodeError, ValueError): parsed = json.loads(raw) if isinstance(parsed, expected): @@ -290,46 +291,31 @@ def _coerce_kwargs( @singleton -class SchemaCoercingToolInterceptor(ChatRequestLoopInterceptor): - """Coerce tool kwargs to match the declared input_schema before invocation.""" +class SchemaCoercingToolInterceptor( + ChatRequestLoopInterceptor, + ToolExecutionInterceptor, +): + """Coerce tool kwargs to the declared schema before execution.""" - async def intercept(self, context: ChatLoopInterceptorContext) -> None: - if context.phase != InterceptorPhase.BEFORE_ITERATION: + async def intercept( + self, + context: ChatInterceptorContext | ToolExecutionInterceptorContext, + ) -> None: + if not isinstance(context, ToolExecutionInterceptorContext): + return + if context.phase != InterceptorPhase.BEFORE_TOOL: return - state = context.state - tools = [ - self._patch_tool(tool) for tool in state.input.context_stack.all_tools() - ] - - stack = state.input.context_stack.remove_layers_of_type( - LayerType.TOOL_DEFINITIONS - ) - if tools: - stack = stack.append_layer( - ToolDefinitionsLayer(tools=tools, source="schema_coercion_patch") + schema = context.request.tool_spec.input_schema or {} + try: + context.set_tool_kwargs( + _coerce_kwargs(context.tool_kwargs, input_schema=schema) + ) + except SchemaCoercionError: + raise + except Exception as e: + logger.exception( + "Schema coercion failed for tool '%s', invoking with original kwargs", + context.request.tool_spec.name, + exc_info=e, ) - state.input.context_stack = stack - context.set_state(state) - - def _patch_tool(self, tool: ToolSpec) -> ToolSpec: - schema = tool.input_schema or {} - original_fn = tool.async_fn - - async def _coerced_fn(*args: Any, **kwargs: Any) -> Any: - try: - fixed_kwargs = _coerce_kwargs(kwargs, input_schema=schema) - except SchemaCoercionError: - raise - except Exception as e: - logger.exception( - "Schema coercion failed for tool '%s', invoking with original kwargs", - tool.name, - exc_info=e, - ) - fixed_kwargs = kwargs - return await original_fn(*args, **fixed_kwargs) - - patched = tool.model_copy(deep=True) - patched.async_fn = _coerced_fn - return patched diff --git a/private_gpt/server/chat/interceptors/skill_tool_visibility_interceptor.py b/private_gpt/server/chat/interceptors/skill_tool_visibility_interceptor.py index d40f7fc2..343fcec2 100644 --- a/private_gpt/server/chat/interceptors/skill_tool_visibility_interceptor.py +++ b/private_gpt/server/chat/interceptors/skill_tool_visibility_interceptor.py @@ -8,13 +8,13 @@ from private_gpt.components.context.models.context_layer import ( ToolDefinitionsLayer, ) from private_gpt.components.context.models.layer_type import LayerType -from private_gpt.components.engines.chat_loop.interceptors.chat_loop_interceptor import ( +from private_gpt.components.engines.chat.interceptors.chat_interceptor import ( ChatRequestLoopInterceptor, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_interceptor_context import ( - ChatLoopInterceptorContext, +from private_gpt.components.engines.chat.models.chat_interceptor_context import ( + ChatInterceptorContext, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_phase import ( +from private_gpt.components.engines.chat.models.chat_phase import ( InterceptorPhase, ) from private_gpt.components.skills.models.skill_entities import SkillFilter @@ -43,7 +43,7 @@ class SkillToolVisibilityInterceptor(ChatRequestLoopInterceptor): def __init__(self) -> None: pass - async def intercept(self, context: ChatLoopInterceptorContext) -> None: + async def intercept(self, context: ChatInterceptorContext) -> None: if context.phase != InterceptorPhase.BEFORE_ITERATION: return diff --git a/private_gpt/server/chat/interceptors/skills_loop_interceptor.py b/private_gpt/server/chat/interceptors/skills_loop_interceptor.py index 88c6be52..22688ab8 100644 --- a/private_gpt/server/chat/interceptors/skills_loop_interceptor.py +++ b/private_gpt/server/chat/interceptors/skills_loop_interceptor.py @@ -13,13 +13,13 @@ from private_gpt.components.context.models.context_layer import ( SkillCatalogLayer, ) from private_gpt.components.context.models.layer_type import LayerType -from private_gpt.components.engines.chat_loop.interceptors.chat_loop_interceptor import ( +from private_gpt.components.engines.chat.interceptors.chat_interceptor import ( ChatRequestLoopInterceptor, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_interceptor_context import ( - ChatLoopInterceptorContext, +from private_gpt.components.engines.chat.models.chat_interceptor_context import ( + ChatInterceptorContext, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_phase import ( +from private_gpt.components.engines.chat.models.chat_phase import ( InterceptorPhase, ) from private_gpt.components.skills.models.skill_entities import SkillFilter @@ -92,7 +92,7 @@ class SkillsInterceptor(ChatRequestLoopInterceptor): self._skill_loader = skill_loader self._skill_injection_mode = settings.skills.skill_injection_mode - async def intercept(self, context: ChatLoopInterceptorContext) -> None: + async def intercept(self, context: ChatInterceptorContext) -> None: if context.phase != InterceptorPhase.BEFORE_ITERATION: return diff --git a/private_gpt/server/chat/interceptors/skills_validation_interceptor.py b/private_gpt/server/chat/interceptors/skills_validation_interceptor.py index 77d529ed..51bfe779 100644 --- a/private_gpt/server/chat/interceptors/skills_validation_interceptor.py +++ b/private_gpt/server/chat/interceptors/skills_validation_interceptor.py @@ -2,16 +2,16 @@ from collections.abc import Sequence from injector import inject, singleton -from private_gpt.components.engines.chat_loop.interceptors.chat_loop_interceptor import ( +from private_gpt.components.engines.chat.interceptors.chat_interceptor import ( ChatRequestLoopInterceptor, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_interceptor_context import ( - ChatLoopInterceptorContext, +from private_gpt.components.engines.chat.models.chat_interceptor_context import ( + ChatInterceptorContext, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_phase import ( +from private_gpt.components.engines.chat.models.chat_phase import ( InterceptorPhase, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_state import ( +from private_gpt.components.engines.chat.models.chat_state import ( SkillsRuntimeCache, ) from private_gpt.components.skills.models.skill_entities import SkillFilter @@ -27,7 +27,7 @@ class SkillsValidationInterceptor(ChatRequestLoopInterceptor): def __init__(self, skill_service: SkillService) -> None: self._skill_service = skill_service - async def intercept(self, context: ChatLoopInterceptorContext) -> None: + async def intercept(self, context: ChatInterceptorContext) -> None: if context.phase not in { InterceptorPhase.VALIDATION, InterceptorPhase.BEFORE_ITERATION, diff --git a/private_gpt/server/chat/interceptors/system_prompt_interceptor.py b/private_gpt/server/chat/interceptors/system_prompt_interceptor.py index 56878907..801e46d0 100644 --- a/private_gpt/server/chat/interceptors/system_prompt_interceptor.py +++ b/private_gpt/server/chat/interceptors/system_prompt_interceptor.py @@ -9,18 +9,19 @@ from private_gpt.components.context.models.context_layer import ( UserInstructionsLayer, ) from private_gpt.components.context.models.layer_type import LayerType -from private_gpt.components.engines.chat_loop.interceptors.chat_loop_interceptor import ( +from private_gpt.components.engines.chat.interceptors.chat_interceptor import ( ChatRequestLoopInterceptor, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_interceptor_context import ( - ChatLoopInterceptorContext, +from private_gpt.components.engines.chat.models.chat_interceptor_context import ( + ChatInterceptorContext, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_phase import ( +from private_gpt.components.engines.chat.models.chat_phase import ( InterceptorPhase, ) -from private_gpt.components.engines.chat_loop.utils.request_builder import ( +from private_gpt.components.engines.chat.utils.request_builder import ( build_request_from_context_stack, ) +from private_gpt.components.llm.llm_helper import as_sync_tokenizer_fn from private_gpt.components.prompts.prompt_builder import PromptBuilderService @@ -37,7 +38,7 @@ class SystemPromptRequestInterceptor(ChatRequestLoopInterceptor): self._prompt_builder_service = prompt_builder_service self._add_context_to_system_prompt = add_context_to_system_prompt - async def intercept(self, context: ChatLoopInterceptorContext) -> None: + async def intercept(self, context: ChatInterceptorContext) -> None: if context.phase != InterceptorPhase.BEFORE_ITERATION: return @@ -62,7 +63,7 @@ class SystemPromptRequestInterceptor(ChatRequestLoopInterceptor): async def _build_generated_layers( self, - context: ChatLoopInterceptorContext, + context: ChatInterceptorContext, request: ChatRequest, documents: list[Any], ) -> list[UserInstructionsLayer | ContextPromptLayer]: @@ -91,7 +92,9 @@ class SystemPromptRequestInterceptor(ChatRequestLoopInterceptor): generate_citations=request.citation.enabled, guidelines_prompt=None, token_limit=context.state.runtime.effective_token_limit, - tokenizer_fn=context.state.runtime.tokenizer_fn, + tokenizer_fn=as_sync_tokenizer_fn( + context.state.runtime.tokenizer_fn + ), ) if context_prompt: context_text = context_prompt.format().strip() diff --git a/private_gpt/server/chat/interceptors/tool_choice_interceptor.py b/private_gpt/server/chat/interceptors/tool_choice_interceptor.py index 072767b6..b1f3c70c 100644 --- a/private_gpt/server/chat/interceptors/tool_choice_interceptor.py +++ b/private_gpt/server/chat/interceptors/tool_choice_interceptor.py @@ -6,13 +6,13 @@ from private_gpt.components.chat.processors.chat_history.tools.tool_choices impo ) from private_gpt.components.context.models.context_layer import ToolDefinitionsLayer from private_gpt.components.context.models.layer_type import LayerType -from private_gpt.components.engines.chat_loop.interceptors.chat_loop_interceptor import ( +from private_gpt.components.engines.chat.interceptors.chat_interceptor import ( ChatRequestLoopInterceptor, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_interceptor_context import ( - ChatLoopInterceptorContext, +from private_gpt.components.engines.chat.models.chat_interceptor_context import ( + ChatInterceptorContext, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_phase import ( +from private_gpt.components.engines.chat.models.chat_phase import ( InterceptorPhase, ) @@ -21,14 +21,14 @@ from private_gpt.components.engines.chat_loop.models.chat_loop_phase import ( class ToolChoiceRequestInterceptor(ChatRequestLoopInterceptor): """Filter tools and update user message hint for forced tool choice.""" - async def intercept(self, context: ChatLoopInterceptorContext) -> None: + async def intercept(self, context: ChatInterceptorContext) -> None: """Apply tool choice policy to history and available tools.""" if context.phase == InterceptorPhase.BEFORE_ITERATION: await self.intercept_before(context) elif context.phase == InterceptorPhase.AFTER_ITERATION: await self.intercept_after(context) - async def intercept_before(self, context: ChatLoopInterceptorContext) -> None: + async def intercept_before(self, context: ChatInterceptorContext) -> None: state = context.state history, tools, new_tool_choice = await process_tool_choices( chat_history=state.input.request.messages, @@ -50,7 +50,7 @@ class ToolChoiceRequestInterceptor(ChatRequestLoopInterceptor): context.set_state(new_state) - async def intercept_after(self, context: ChatLoopInterceptorContext) -> None: + async def intercept_after(self, context: ChatInterceptorContext) -> None: state = context.state tool_choices = state.input.request.tool_config.tool_choices diff --git a/private_gpt/server/chat/interceptors/validator_request_interceptor.py b/private_gpt/server/chat/interceptors/validator_request_interceptor.py index bffd6c53..ccedf730 100644 --- a/private_gpt/server/chat/interceptors/validator_request_interceptor.py +++ b/private_gpt/server/chat/interceptors/validator_request_interceptor.py @@ -7,13 +7,13 @@ from llama_index.core.base.llms.types import ( TextBlock, ) -from private_gpt.components.engines.chat_loop.interceptors.chat_loop_interceptor import ( +from private_gpt.components.engines.chat.interceptors.chat_interceptor import ( ChatRequestLoopInterceptor, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_interceptor_context import ( - ChatLoopInterceptorContext, +from private_gpt.components.engines.chat.models.chat_interceptor_context import ( + ChatInterceptorContext, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_phase import ( +from private_gpt.components.engines.chat.models.chat_phase import ( InterceptorPhase, ) from private_gpt.components.llm.custom.base import ZylonLLM @@ -33,7 +33,7 @@ class ValidatorRequestInterceptor(ChatRequestLoopInterceptor): def __init__(self, llm_component: LLMComponent) -> None: self._llm_component = llm_component - async def intercept(self, context: ChatLoopInterceptorContext) -> None: + async def intercept(self, context: ChatInterceptorContext) -> None: if context.phase != InterceptorPhase.VALIDATION: return diff --git a/private_gpt/server/chat_async/chat_async_facade.py b/private_gpt/server/chat_async/chat_async_facade.py index 67709564..1aea58fd 100644 --- a/private_gpt/server/chat_async/chat_async_facade.py +++ b/private_gpt/server/chat_async/chat_async_facade.py @@ -6,10 +6,11 @@ from injector import inject, singleton from starlette.requests import Request from starlette.responses import StreamingResponse -from private_gpt.components.chat.models.chat_config_models import ChatRequest from private_gpt.events.event_folding import fold from private_gpt.events.models import Event, FatalError, Message from private_gpt.events.utils import to_message, to_sse_stream +from private_gpt.server.chat.chat_models import ChatBody +from private_gpt.server.chat.chat_request_mapper import ChatRequestMapper from private_gpt.server.chat_async.chat_async_service import ChatAsyncService logger = logging.getLogger(__name__) @@ -24,11 +25,13 @@ class ChatAsyncFacadeService: def __init__( self, chat_async_service: ChatAsyncService, + chat_request_mapper: ChatRequestMapper, ) -> None: self._chat_async_service = chat_async_service + self._chat_request_mapper = chat_request_mapper async def chat( - self, http_request: Request, request: ChatRequest, message_id: str | None = None + self, http_request: Request, body: ChatBody, message_id: str | None = None ) -> Message | FatalError | StreamingResponse: """Handle chat with proper cancellation support for FastAPI. @@ -37,8 +40,9 @@ class ChatAsyncFacadeService: 2. Resources are cleaned up 3. CancelledError is re-raised to maintain asyncio semantics """ + chat_request = await self._chat_request_mapper.create_request_from_body(body) message_id = await self._chat_async_service.initiate_chat_stream( - request=request, message_id=message_id + request=chat_request, message_id=message_id ) event_generator = await self._chat_async_service.get_stream_events( message_id=message_id, @@ -49,7 +53,7 @@ class ChatAsyncFacadeService: cancellable_generator = self._cancellable_stream_generator( http_request, event_generator, message_id ) - if request.stream: + if body.stream: sse_stream = to_sse_stream(cancellable_generator) return StreamingResponse( sse_stream, diff --git a/private_gpt/server/chat_async/chat_async_router.py b/private_gpt/server/chat_async/chat_async_router.py index 9b6f160c..1dadc76a 100644 --- a/private_gpt/server/chat_async/chat_async_router.py +++ b/private_gpt/server/chat_async/chat_async_router.py @@ -274,11 +274,11 @@ async def chat_messages( * Stream status can be monitored via status endpoint """ chat_service: ChatAsyncService = request.state.injector.get(ChatAsyncService) - chat_request_mapper: ChatRequestMapper = request.state.injector.get( - ChatRequestMapper - ) + request_mapper: ChatRequestMapper = request.state.injector.get(ChatRequestMapper) + chat_request = await request_mapper.create_request_from_body(body) message_id = await chat_service.initiate_chat_stream( - await chat_request_mapper.create_request_from_body(body), message_id=message_id + request=chat_request, + message_id=message_id, ) return ChatResponse( diff --git a/private_gpt/server/chat_async/chat_async_service.py b/private_gpt/server/chat_async/chat_async_service.py index af66e92f..94ebb30f 100644 --- a/private_gpt/server/chat_async/chat_async_service.py +++ b/private_gpt/server/chat_async/chat_async_service.py @@ -1,4 +1,5 @@ from collections.abc import AsyncGenerator +from uuid import uuid4 from injector import inject, singleton @@ -13,42 +14,43 @@ from private_gpt.server.chat.chat_facade import ChatFacadeService @singleton class ChatAsyncService: @inject - def __init__(self, chat_facade: ChatFacadeService, stream_manager: StreamManager): - self._chat_facade = chat_facade + def __init__( + self, + stream_manager: StreamManager, + chat_facade: ChatFacadeService, + ): self.stream_manager = stream_manager + self._chat_facade = chat_facade async def initiate_chat_stream( self, request: ChatRequest, message_id: str | None = None ) -> str: - """Initiate a chat completion stream. - - Returns: - Tuple of (message_id, message) - - Raises: - HTTPException: If message_id already exists - """ + """Initiate a chat completion stream.""" + message_id = message_id or str(uuid4()) if message_id and await self.stream_manager.stream_exists(message_id): raise ValueError("Stream with this message_id already exists") + request = request.model_copy( + update={ + "context": request.context.model_copy( + update={"correlation_id": message_id} + ) + } + ) event_generator = await self._chat_facade.create_chat_event_generator( request=request ) - message_id = await self.stream_manager.create_and_start_stream( + return await self.stream_manager.create_and_start_stream( event_handler=StreamingEventHandler(), 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, - "citations_enabled": request.citation.enabled, }, ) - return message_id - async def get_stream_events( self, message_id: str ) -> AsyncGenerator[Event, None] | None: diff --git a/private_gpt/server/ingest/ingest_router.py b/private_gpt/server/ingest/ingest_router.py index b724c7a8..a53ba842 100644 --- a/private_gpt/server/ingest/ingest_router.py +++ b/private_gpt/server/ingest/ingest_router.py @@ -1,27 +1,23 @@ import logging -import uuid from pathlib import Path -from typing import TYPE_CHECKING, Annotated, Any, Literal +from typing import Annotated, Any, Literal from fastapi import APIRouter, Body, Depends, Request from pydantic import BaseModel, Field -from private_gpt.components.ingest.utils import get_extension, get_file_name -from private_gpt.components.storage.s3_helper import S3Helper +from private_gpt.components.ingestion.ingestion_scheduler import ( + IngestionSchedulerFactory, +) from private_gpt.server.ingest.ingest_service import IngestService from private_gpt.server.ingest.model import IngestedDoc -from private_gpt.server.utils.artifact_input import IngestableArtifactType, UriArtifact +from private_gpt.server.utils.artifact_input import IngestableArtifactType from private_gpt.server.utils.auth import authenticated from private_gpt.server.utils.callback import BaseCallbackInput -from private_gpt.settings.settings import settings try: from private_gpt.celery.celery import celery_app from private_gpt.celery.model import Task, TaskStatus - if TYPE_CHECKING: - from celery.result import AsyncResult - _CELERY_AVAILABLE = True except ImportError: _CELERY_AVAILABLE = False @@ -429,249 +425,206 @@ def ingest_content( The ingested content becomes immediately available for use in other API endpoints like /messages, /artifacts/search, and /artifacts/content. + + The configured ingestion scheduler decides whether the work runs + in-process or is dispatched to a worker. """ - service: IngestService = request.state.injector.get(IngestService) - - content = body.input.to_binary_content(get_file_name(body.metadata)) - with service.temporary_file( - lambda: service.data_path_from_bin_data( - content.data, get_extension(content.filename) - ) - ) as data_path: - return ingest_data_sync( - collection=body.collection, - artifact=body.artifact, - data_path=data_path, - service=service, - metadata={**(body.metadata or {}), "file_name": content.filename}, - ) + scheduler = request.state.injector.get(IngestionSchedulerFactory).get() + result: IngestResponse = scheduler.ingest(body) + return result -if _CELERY_AVAILABLE: - - @ingest_router.post( - "/ingest/async", - response_model=Task, - summary="Ingest Content Asynchronously", - description="Initiate asynchronous ingestion of content from multiple sources", - responses={ - 200: { - "description": "Successfully initiated ingestion task", - "content": { - "application/json": { - "example": {"task_id": "123e4567-e89b-12d3-a456-426614174000"} - } - }, - }, - 422: { - "description": "Invalid input format or request parameters", - "content": { - "application/json": { - "example": { - "detail": "Input format is invalid or parameters are missing" - } - } - }, - }, - 404: { - "description": "Resource not accessible (for URI inputs)", - "content": { - "application/json": { - "example": {"detail": "The URI resource could not be accessed"} - } - }, +@ingest_router.post( + "/ingest/async", + response_model=Task, + summary="Ingest Content Asynchronously", + description="Initiate asynchronous ingestion of content from multiple sources", + responses={ + 200: { + "description": "Successfully initiated ingestion task", + "content": { + "application/json": { + "example": {"task_id": "123e4567-e89b-12d3-a456-426614174000"} + } }, }, - tags=["Artifacts"], - openapi_extra={ - "requestBody": { - "description": "JSON request body containing ingestion parameters and callback configuration for asynchronous processing notifications", - "content": { - "application/json": { - "examples": { - "file_async": { - "summary": "Async file ingestion", - "value": { - "callback": { - "amqp": { - "exchange": "ingestion_events", - "routing_key_done": "ingest.completed", - "routing_key_error": "ingest.failed", - "routing_key_progress": "ingest.progress", - } - }, - "ingest_body": { - "input": { - "type": "file", - "value": "JVBERi0xLjQKJaqrrK0KMS...", - }, - "artifact": "quarterly_report", - "collection": "financial_docs", - "metadata": {"file_name": "Q3_Report.pdf"}, - }, - }, - }, - "uri_async": { - "summary": "Async URI ingestion", - "value": { - "callback": { - "amqp": { - "exchange": "ingestion_events", - "routing_key_done": "ingest.completed", - "routing_key_error": "ingest.failed", - "routing_key_progress": "ingest.progress", - } - }, - "ingest_body": { - "input": { - "type": "uri", - "value": "s3://company-docs/annual-2023.pdf", - }, - "artifact": "annual_report_2023", - "collection": "financial_reports", - "metadata": {"year": "2023"}, - }, - }, - }, - "text_async": { - "summary": "Async text ingestion", - "value": { - "callback": { - "amqp": { - "exchange": "ingestion_events", - "routing_key_done": "ingest.completed", - "routing_key_error": "ingest.failed", - "routing_key_progress": "ingest.progress", - } - }, - "ingest_body": { - "input": { - "type": "text", - "value": "Our company was founded in 2020...", - }, - "artifact": "company_profile", - "collection": "corporate_docs", - "metadata": {"author": "Marketing Team"}, - }, - }, - }, - } + 422: { + "description": "Invalid input format or request parameters", + "content": { + "application/json": { + "example": { + "detail": "Input format is invalid or parameters are missing" } - }, + } + }, + }, + 404: { + "description": "Resource not accessible (for URI inputs)", + "content": { + "application/json": { + "example": {"detail": "The URI resource could not be accessed"} + } + }, + }, + }, + tags=["Artifacts"], + openapi_extra={ + "requestBody": { + "description": "JSON request body containing ingestion parameters and callback configuration for asynchronous processing notifications", + "content": { + "application/json": { + "examples": { + "file_async": { + "summary": "Async file ingestion", + "value": { + "callback": { + "amqp": { + "exchange": "ingestion_events", + "routing_key_done": "ingest.completed", + "routing_key_error": "ingest.failed", + "routing_key_progress": "ingest.progress", + } + }, + "ingest_body": { + "input": { + "type": "file", + "value": "JVBERi0xLjQKJaqrrK0KMS...", + }, + "artifact": "quarterly_report", + "collection": "financial_docs", + "metadata": {"file_name": "Q3_Report.pdf"}, + }, + }, + }, + "uri_async": { + "summary": "Async URI ingestion", + "value": { + "callback": { + "amqp": { + "exchange": "ingestion_events", + "routing_key_done": "ingest.completed", + "routing_key_error": "ingest.failed", + "routing_key_progress": "ingest.progress", + } + }, + "ingest_body": { + "input": { + "type": "uri", + "value": "s3://company-docs/annual-2023.pdf", + }, + "artifact": "annual_report_2023", + "collection": "financial_reports", + "metadata": {"year": "2023"}, + }, + }, + }, + "text_async": { + "summary": "Async text ingestion", + "value": { + "callback": { + "amqp": { + "exchange": "ingestion_events", + "routing_key_done": "ingest.completed", + "routing_key_error": "ingest.failed", + "routing_key_progress": "ingest.progress", + } + }, + "ingest_body": { + "input": { + "type": "text", + "value": "Our company was founded in 2020...", + }, + "artifact": "company_profile", + "collection": "corporate_docs", + "metadata": {"author": "Marketing Team"}, + }, + }, + }, + } + } + }, + } + }, +) +def ingest_content_async(request: Request, body: IngestAsyncBody) -> Task: + """Asynchronously process and ingest content from multiple sources. + + This endpoint initiates an asynchronous task to process and ingest content + from various sources including files, URIs, and text. It's particularly + useful for large files or when you need non-blocking operations. + + Supported input types: + - Base64 encoded files + - Remote URIs (HTTP/HTTPS, S3, etc.) + - Plain text content + - Already processed content (ContextFilter objects) + + The context obtained from ingested content is later used in + `/chat/completions`, `/completions`, and `/artifacts/search` APIs. + + A file can generate different Documents + (for example a PDF generates one Document + per page). All Documents are returned in the response, together with the + extracted Metadata and artifact id, + which is later used to improve context retrieval + and can be used to filter the context used to create responses in + `/chat/completions`, `/completions`, + `/artifacts/search` and `/artifact/content`. + + Example request: + ```json + { + "callback": { + "amqp": { + "exchange": "ingestion", + "routing_key_done": "ingest.done", + "routing_key_error": "ingest.error", + "routing_key_progress": "ingest.progress" } }, - ) - def ingest_content_async(request: Request, body: IngestAsyncBody) -> Task: - """Asynchronously process and ingest content from multiple sources. - - This endpoint initiates an asynchronous task to process and ingest content - from various sources including files, URIs, and text. It's particularly - useful for large files or when you need non-blocking operations. - - Supported input types: - - Base64 encoded files - - Remote URIs (HTTP/HTTPS, S3, etc.) - - Plain text content - - Already processed content (ContextFilter objects) - - The context obtained from ingested content is later used in - `/chat/completions`, `/completions`, and `/artifacts/search` APIs. - - A file can generate different Documents - (for example a PDF generates one Document - per page). All Documents are returned in the response, together with the - extracted Metadata and artifact id, - which is later used to improve context retrieval - and can be used to filter the context used to create responses in - `/chat/completions`, `/completions`, - `/artifacts/search` and `/artifact/content`. - - Example request: - ```json - { - "callback": { - "amqp": { - "exchange": "ingestion", - "routing_key_done": "ingest.done", - "routing_key_error": "ingest.error", - "routing_key_progress": "ingest.progress" - } + "ingest_body": { + "input": { + "type": "uri", + "value": "s3://company-docs/annual-2023.pdf" }, - "ingest_body": { - "input": { - "type": "uri", - "value": "s3://company-docs/annual-2023.pdf" - }, - "artifact": "annual_report_2023", - "collection": "financial_reports", - "metadata": { - "file_name": "annual_report_2023.pdf", - "department": "finance", - "year": "2023" - } + "artifact": "annual_report_2023", + "collection": "financial_reports", + "metadata": { + "file_name": "annual_report_2023.pdf", + "department": "finance", + "year": "2023" } } - ``` + } + ``` - Notes: - * URIs must be accessible from the server - * Base64 content should be properly encoded - * File name with extension is recommended in metadata - * PDFs generate one document per page - * Large files are automatically split into chunks - * Progress can be monitored via /artifacts/ingest/async/{task_id} - * Resulting context is available in chat/completion APIs - * Ingested content can be filtered using metadata + Notes: + * URIs must be accessible from the server + * Base64 content should be properly encoded + * File name with extension is recommended in metadata + * PDFs generate one document per page + * Large files are automatically split into chunks + * Progress can be monitored via /artifacts/ingest/async/{task_id} + * Resulting context is available in chat/completion APIs + * Ingested content can be filtered using metadata - Important to know: - - Since binary data cannot be passed directly to Celery tasks, files are - first uploaded to a temporary S3 bucket. The Celery task then retrieves - the file from S3 for processing. This can incur additional time. + Important to know: + - Since binary data cannot be passed directly to Celery tasks, files are + first uploaded to a temporary S3 bucket. The Celery task then retrieves + the file from S3 for processing. This can incur additional time. - """ - service = request.state.injector.get(IngestService) - config = settings() - s3_helper = request.state.injector.get(S3Helper) + """ + scheduler = request.state.injector.get(IngestionSchedulerFactory).get() + try: + task_id = scheduler.ingest_async(body) + except NotImplementedError as exc: + from fastapi import HTTPException - # Initialize ingestion - service.initialize_artifact_indices( - collection=body.ingest_body.collection, - artifact=body.ingest_body.artifact, - ) - - # Since we cannot store in the Celery task args the actual binary data, - # we upload the file to a temporary S3 bucket and pass the S3 URI instead. - if body.ingest_body.input: - should_upload = ( - # If it's not a URI, we need to upload - not isinstance(body.ingest_body.input, UriArtifact) - # It's a URI in Base64, we need to upload - or body.ingest_body.input.is_base64() - ) - - if should_upload: - content = body.ingest_body.input.to_binary_content() - object_name = str(uuid.uuid4()) - - s3_url = s3_helper.upload_file_to_s3( - filename=content.filename, - bytes_data=content.data.read(), - bucket_name=config.s3.temporary_bucket_name, - object_name=object_name, - ) - - if not s3_url: - raise ValueError("Failed to upload file to S3 for async ingestion") - - body.ingest_body.input = UriArtifact(value=s3_url) - - # Enqueue ingestion tasks (index population) - async_result = celery_app.send_task( - "vector_index_task", - args=(body,), - ) - - return Task(task_id=async_result.task_id) + raise HTTPException( + status_code=501, + detail="Async ingestion is not supported by the configured scheduler", + ) from exc + return Task(task_id=task_id) def ingest_data_sync( @@ -900,77 +853,66 @@ def delete_ingested(request: Request, body: DeleteIngestedDocumentBody) -> None: and all related chunks from the specified collection. The deletion is immediate and cannot be undone. """ - service = request.state.injector.get(IngestService) - service.delete( - collection=body.collection, - artifact=body.artifact, - ) + scheduler = request.state.injector.get(IngestionSchedulerFactory).get() + scheduler.delete(collection=body.collection, artifact=body.artifact) -if _CELERY_AVAILABLE: - - @ingest_router.post( - "/delete/async", - response_model=Task, - summary="Delete Document Asynchronously", - responses={ - 200: {"model": Task, "description": "Successfully initiated deletion task"}, - 422: {"description": "Invalid request parameters"}, - }, - tags=["Artifacts"], - openapi_extra={ - "requestBody": { - "description": "JSON request body containing deletion parameters and optional callback configuration for asynchronous processing notifications", - "content": { - "application/json": { - "example": { - "callback": { - "amqp": { - "exchange": "deletion_events", - "routing_key_done": "delete.completed", - "routing_key_error": "delete.failed", - "routing_key_progress": "delete.progress", - } - }, - "delete_body": { - "collection": "financial_reports", - "artifact": "obsolete_report_2022", - }, - } +@ingest_router.post( + "/delete/async", + response_model=Task, + summary="Delete Document Asynchronously", + responses={ + 200: {"model": Task, "description": "Successfully initiated deletion task"}, + 422: {"description": "Invalid request parameters"}, + }, + tags=["Artifacts"], + openapi_extra={ + "requestBody": { + "description": "JSON request body containing deletion parameters and optional callback configuration for asynchronous processing notifications", + "content": { + "application/json": { + "example": { + "callback": { + "amqp": { + "exchange": "deletion_events", + "routing_key_done": "delete.completed", + "routing_key_error": "delete.failed", + "routing_key_progress": "delete.progress", + } + }, + "delete_body": { + "collection": "financial_reports", + "artifact": "obsolete_report_2022", + }, } - }, - } - }, - ) - def delete_ingested_async(body: DeleteIngestedDocumentAsyncBody) -> Task: - """Initiates asynchronous deletion of a document and all associated data. + } + }, + } + }, +) +def delete_ingested_async( + request: Request, body: DeleteIngestedDocumentAsyncBody +) -> Task: + """Initiates asynchronous deletion of a document and all associated data. - This endpoint queues a deletion task for background processing, making it - suitable for large documents or when non-blocking operation is required. - The task can be monitored using the returned task ID. + This endpoint queues a deletion task for background processing, making it + suitable for large documents or when non-blocking operation is required. + The task can be monitored using the returned task ID. - If an ingestion task is currently running for the same document, it will - be automatically revoked before initiating the deletion. - """ - # First we try to revoke the running task if it's already running. - # In case that yes, we return a task_id with status "revoked" - from private_gpt.celery.task_helper import IngestionTaskHelper + If an ingestion task is currently running for the same document, it will + be automatically revoked before initiating the deletion. + """ + scheduler = request.state.injector.get(IngestionSchedulerFactory).get() + try: + task_id = scheduler.delete_async(body) + except NotImplementedError as exc: + from fastapi import HTTPException - revoked = IngestionTaskHelper.revoke_ingestion_task( - celery_app=celery_app, - collection=body.delete_body.collection, - artifact=body.delete_body.artifact, - ) - if revoked: - logger.info(f"Revoked ingestion task for {body.delete_body.artifact}") - return Task(task_id="revoked") - - # Enqueue deletion task - async_result: AsyncResult[None] = celery_app.send_task( - "delete_ingested_task", - args=(body,), - ) - return Task(task_id=async_result.task_id) + raise HTTPException( + status_code=501, + detail="Async deletion is not supported by the configured scheduler", + ) from exc + return Task(task_id=task_id) if _CELERY_AVAILABLE: @@ -1033,7 +975,7 @@ if _CELERY_AVAILABLE: description="Unique identifier of the deletion task to check status for", examples=["123e4567-e89b-12d3-a456-426614174000"], ), - ] + ], ) -> TaskStatus[None]: """Retrieves the current status of an asynchronous deletion task. diff --git a/private_gpt/settings/settings.py b/private_gpt/settings/settings.py index fac61a73..d26da16f 100644 --- a/private_gpt/settings/settings.py +++ b/private_gpt/settings/settings.py @@ -2,7 +2,14 @@ import inspect import json from typing import Annotated, Any, Literal -from pydantic import AnyUrl, BaseModel, ConfigDict, Field, field_validator +from pydantic import ( + AnyUrl, + BaseModel, + ConfigDict, + Field, + field_validator, + model_validator, +) from private_gpt.settings.settings_loader import load_active_settings @@ -454,7 +461,45 @@ class PreprocessSettings(BaseModel): ) +class SchedulerSettings(BaseModel): + mode: str = Field( + default="local", + description=( + "Scheduler mode name. Built-ins include ``local``, ``arq``, and ``celery``. " + "Tests and extensions may register additional provider names." + ), + ) + celery_queue: str = Field( + default="", + description="Celery queue name when mode is 'celery'.", + ) + callback_timeout_seconds: int = Field( + default=300, + gt=0, + description="Maximum time to wait for resumable chat callbacks.", + ) + + +class SchedulerConfig(BaseModel): + ingestion: SchedulerSettings = Field( + default_factory=lambda: SchedulerSettings(celery_queue="ingestion"), + description="Ingestion worker scheduler configuration.", + ) + chat: SchedulerSettings = Field( + default_factory=lambda: SchedulerSettings(celery_queue="chat"), + description="Chat worker scheduler configuration.", + ) + tools: SchedulerSettings = Field( + default_factory=lambda: SchedulerSettings(celery_queue="tools"), + description="Tool worker scheduler configuration.", + ) + + class ChatSettings(BaseModel): + engine_mode: Literal["loop", "async"] = Field( + default="async", + description="Chat engine selected behind the runtime feature flag.", + ) allow_use_default_prompt: bool = Field( True, description="Flag indicating if the chat engine should use default prompts or not.", @@ -528,10 +573,6 @@ class ChatSettings(BaseModel): None, description="The threshold for the number of context items to switch to multiplexing mode.", ) - maximum_concurrent_requests: int | None = Field( - None, - description="The maximum number of concurrent requests that can be handled by the chat engine.", - ) maximum_blob_size: int = Field( 25 * 1024 * 1024, description="The maximum size in bytes of a blob that can be processed by the chat engine.", @@ -1055,6 +1096,11 @@ class CelerySettings(BaseModel): description="The visibility timeout for tasks in seconds", default=None, ) + max_tasks_per_child: int = Field( + description="Maximum tasks handled by a stateful Celery child before recycling", + default=1000, + gt=0, + ) def __init__(self, **data: Any) -> None: if "soft_time_limit" in data: @@ -1674,6 +1720,42 @@ class Settings(BaseModel): skills: SkillSettings transformation: TransformationSettings semaphore: SemaphoreSettings + scheduler: SchedulerConfig = Field( + default_factory=SchedulerConfig, + description="Scheduler configuration for chat and tool workers.", + ) + + @model_validator(mode="after") + def validate_chat_scheduler_configuration(self) -> "Settings": + if self.scheduler.chat.mode not in {"local", "arq"}: + raise ValueError( + f"Unsupported scheduler.chat.mode={self.scheduler.chat.mode!r}. " + "Supported chat scheduler modes are 'local' and 'arq'." + ) + + if self.scheduler.tools.mode not in {"local", "celery"}: + raise ValueError( + f"Unsupported scheduler.tools.mode={self.scheduler.tools.mode!r}. " + "Supported tool scheduler modes are 'local' and 'celery'." + ) + + if self.scheduler.tools.mode == "celery" and self.scheduler.chat.mode != "arq": + raise ValueError( + f"scheduler.tools.mode={self.scheduler.tools.mode!r} requires " + "scheduler.chat.mode='arq' so tool callbacks can resume shared " + "chat state." + ) + + if self.scheduler.chat.mode == "local": + return self + + if self.stream.broker != "redis": + raise ValueError( + f"scheduler.chat.mode={self.scheduler.chat.mode!r} requires stream.broker=redis " + "because API and chat worker processes must share stream state." + ) + + return self """ diff --git a/private_gpt/worker/__init__.py b/private_gpt/worker/__init__.py new file mode 100644 index 00000000..c483949f --- /dev/null +++ b/private_gpt/worker/__init__.py @@ -0,0 +1,12 @@ +from collections.abc import Sequence + +from private_gpt.worker.modes import register_private_gpt_worker_modes +from private_gpt.worker.registry import register_worker_mode, run_worker + + +def run_private_gpt_worker(args: Sequence[str] = ()) -> None: + register_private_gpt_worker_modes() + run_worker(args) + + +__all__ = ["register_worker_mode", "run_private_gpt_worker", "run_worker"] diff --git a/private_gpt/worker/__main__.py b/private_gpt/worker/__main__.py new file mode 100644 index 00000000..ea3422a2 --- /dev/null +++ b/private_gpt/worker/__main__.py @@ -0,0 +1,6 @@ +import sys + +from private_gpt.worker import run_private_gpt_worker + +if __name__ == "__main__": + run_private_gpt_worker(sys.argv[1:]) diff --git a/private_gpt/worker/modes.py b/private_gpt/worker/modes.py new file mode 100644 index 00000000..37c73a24 --- /dev/null +++ b/private_gpt/worker/modes.py @@ -0,0 +1,177 @@ +import importlib +import os +import signal +import subprocess +import sys +from collections.abc import Callable, Sequence +from functools import partial + +import typer + +from private_gpt.settings.settings import settings +from private_gpt.worker.registry import register_worker_mode + + +def _app_module() -> str: + return os.environ.get("PGPT_WORKER_APP_MODULE", "private_gpt") + + +def _build_flower_args() -> list[str]: + args: list[str] = [] + url_prefix = os.environ.get("PGPT_FLOWER_URL_PREFIX", "") + port = os.environ.get("PGPT_FLOWER_PORT", "5555") + password = os.environ.get("PGPT_FLOWER_PASSWORD", "flower") + persistent = os.environ.get("PGPT_FLOWER_PERSISTENT", "true") + persistent_db = os.environ.get("PGPT_FLOWER_PERSISTENT_DB", "celery_flower.db") + max_tasks = os.environ.get("PGPT_FLOWER_MAXIMUM_TASKS", "1000") + + if url_prefix: + args.append(f"--url_prefix={url_prefix}") + if port: + args.append(f"--port={port}") + if password: + args.append(f"--basic_auth=flower:{password}") + if persistent.lower() == "true": + args.append("--persistent=True") + if persistent_db: + args.append(f"--db={persistent_db}") + if max_tasks: + args.append(f"--max-tasks={max_tasks}") + return args + + +def _build_celery_args(max_tasks_per_child: int) -> list[str]: + args: list[str] = [] + log_level = os.environ.get("PGPT_CELERY_LOG_LEVEL", "info") + queues = os.environ.get("PGPT_CELERY_QUEUES", "") + hostname = os.environ.get("PGPT_CELERY_HOSTNAME", "default") + pool = os.environ.get("PGPT_CELERY_POOL", "prefork") + concurrency = os.environ.get("PGPT_CELERY_CONCURRENCY", "") + time_limit = os.environ.get("PGPT_CELERY_TIME_LIMIT", "") + stateful_type = os.environ.get("PGPT_STATEFUL_WORKER_TYPE", "").strip() + + if log_level: + args.append(f"--loglevel={log_level}") + if stateful_type: + queues = queues or stateful_type + hostname = hostname or stateful_type + args.append(f"--max-tasks-per-child={max_tasks_per_child}") + if queues: + args.append(f"--queues={queues}") + if hostname: + args.append(f"--hostname={hostname}%h") + if pool: + args.append(f"--pool={pool}") + if concurrency: + args.append(f"--concurrency={concurrency}") + if time_limit: + args.append(f"--time-limit={time_limit}") + return args + + +def _run_processes(commands: Sequence[Sequence[str]]) -> None: + processes = [subprocess.Popen(command) for command in commands] + + def cleanup(signum: int = 0, frame: object = None) -> None: + del signum, frame + typer.echo("Shutting down services...") + for process in processes: + process.terminate() + for process in processes: + try: + process.wait(timeout=10) + except subprocess.TimeoutExpired: + process.kill() + + signal.signal(signal.SIGTERM, cleanup) + signal.signal(signal.SIGINT, cleanup) + + try: + for process in processes: + process.wait() + except KeyboardInterrupt: + cleanup() + + +def _healthcheck_command() -> list[str]: + return [ + sys.executable, + "-m", + "uvicorn", + f"{_app_module()}.celery.healthcheck:app", + "--host", + "0.0.0.0", + "--port", + os.environ.get("API_PORT", "8090"), + "--no-access-log", + "--log-level", + "critical", + ] + + +def _with_healthcheck(command: list[str]) -> list[list[str]]: + commands = [command] + if os.environ.get("API_ENABLED", "true").lower() == "true": + commands.append(_healthcheck_command()) + return commands + + +def run_arq(args: Sequence[str]) -> None: + if args: + raise ValueError("The arq worker mode does not support arguments yet") + arq_app = importlib.import_module(f"{_app_module()}.arq") + arq_app.run_arq_worker() + + +def run_celery( + args: Sequence[str], + *, + max_tasks_per_child_resolver: Callable[[], int] = ( + lambda: settings().celery.max_tasks_per_child + ), +) -> None: + celery_args = _build_celery_args(max_tasks_per_child_resolver()) + typer.echo(f"Starting celery worker with args: {' '.join(celery_args)}") + command = [ + sys.executable, + "-m", + "celery", + "--app", + f"{_app_module()}.celery", + "worker", + *celery_args, + *args, + ] + _run_processes(_with_healthcheck(command)) + + +def run_flower(args: Sequence[str]) -> None: + flower_args = _build_flower_args() + typer.echo(f"Starting flower on port {os.environ.get('PGPT_FLOWER_PORT', '5555')}") + command = [ + sys.executable, + "-m", + "celery", + "--app", + f"{_app_module()}.celery", + "flower", + *flower_args, + *args, + ] + _run_processes(_with_healthcheck(command)) + + +def register_private_gpt_worker_modes( + max_tasks_per_child_resolver: Callable[[], int] = ( + lambda: settings().celery.max_tasks_per_child + ), +) -> None: + register_worker_mode("arq", run_arq) + register_worker_mode( + "celery", + partial( + run_celery, + max_tasks_per_child_resolver=max_tasks_per_child_resolver, + ), + ) + register_worker_mode("flower", run_flower) diff --git a/private_gpt/worker/registry.py b/private_gpt/worker/registry.py new file mode 100644 index 00000000..4405479a --- /dev/null +++ b/private_gpt/worker/registry.py @@ -0,0 +1,31 @@ +import os +from collections.abc import Callable, Sequence + +WorkerModeHandler = Callable[[Sequence[str]], None] + +_worker_modes: dict[str, WorkerModeHandler] = {} + + +def register_worker_mode(name: str, handler: WorkerModeHandler) -> None: + normalized_name = name.strip().lower() + if not normalized_name: + raise ValueError("Worker mode name cannot be empty") + _worker_modes[normalized_name] = handler + + +def get_worker_mode(name: str) -> WorkerModeHandler: + normalized_name = name.strip().lower() + try: + return _worker_modes[normalized_name] + except KeyError as exc: + supported_modes = ", ".join(sorted(_worker_modes)) + raise ValueError( + f"Unsupported PGPT_WORKER_MODE={name!r}. Registered modes: {supported_modes}" + ) from exc + + +def run_worker(args: Sequence[str] = ()) -> None: + mode = os.environ.get("PGPT_WORKER_MODE", "").strip() + if not mode: + raise ValueError("PGPT_WORKER_MODE is required") + get_worker_mode(mode)(args) diff --git a/pyproject.toml b/pyproject.toml index b1437e8a..ad6d9856 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,6 +21,7 @@ dependencies = [ "retry-async>=0.1.4,<0.2", "nest-asyncio>=1.6.0", "nltk>=3.9.1,<4", + "arq>=0.26,<0.27", ] [project.scripts] diff --git a/scripts/worker_entrypoint b/scripts/worker_entrypoint index 67ceb5f4..c4d91ffc 100755 --- a/scripts/worker_entrypoint +++ b/scripts/worker_entrypoint @@ -1,104 +1,7 @@ #!/bin/bash +set -euo pipefail -: "${PGPT_WORKER_MODE:=mixed}" # mixed, worker, flower -: "${PGPT_FLOWER_PORT:=5555}" -: "${PGPT_FLOWER_PASSWORD:=flower}" -: "${PGPT_FLOWER_URL_PREFIX:=}" -: "${PGPT_FLOWER_PERSISTENT:=true}" -: "${PGPT_FLOWER_PERSISTENT_DB:=celery_flower.db}" -: "${PGPT_FLOWER_MAXIMUM_TASKS:=1000}" -: "${PGPT_CELERY_LOG_LEVEL:=info}" -: "${PGPT_CELERY_QUEUES:=}" -: "${PGPT_CELERY_HOSTNAME:=default}" -: "${PGPT_CELERY_POOL:=prefork}" -: "${PGPT_CELERY_CONCURRENCY:=}" -: "${PGPT_CELERY_TIME_LIMIT:=}" +: "${PGPT_WORKER_MODE:?PGPT_WORKER_MODE is required}" : "${PGPT_WORKER_APP_MODULE:=private_gpt}" -: "${API_ENABLED:=true}" -: "${API_PORT:=8090}" -CELERY_APP_MODULE="${PGPT_WORKER_APP_MODULE}.celery" -HEALTHCHECK_APP_MODULE="${PGPT_WORKER_APP_MODULE}.celery.healthcheck:app" - -# Create directory for PID files -mkdir -p /tmp/private_gpt_pids - -# Add signal handling -cleanup() { - echo "Shutting down services..." - jobs -p | xargs -r kill - exit 0 -} -trap cleanup EXIT INT TERM - -# Run flower on the specified port -if [[ "$PGPT_WORKER_MODE" == "flower" || "$PGPT_WORKER_MODE" == "mixed" ]]; then - echo "Starting flower on port $PGPT_FLOWER_PORT" - FLOWER_ARGS=() - if [[ -n "$PGPT_FLOWER_URL_PREFIX" ]]; then - FLOWER_ARGS+=("--url_prefix=$PGPT_FLOWER_URL_PREFIX") - fi - if [[ -n "$PGPT_FLOWER_PORT" ]]; then - FLOWER_ARGS+=("--port=$PGPT_FLOWER_PORT") - fi - if [[ -n "$PGPT_FLOWER_PASSWORD" ]]; then - FLOWER_ARGS+=("--basic_auth=flower:$PGPT_FLOWER_PASSWORD") - fi - if [[ "$PGPT_WORKER_MODE" == "mixed" ]]; then - FLOWER_ARGS+=("--logging=none") - fi - if [[ "$PGPT_FLOWER_PERSISTENT" == "true" ]]; then - FLOWER_ARGS+=("--persistent=True") - if [[ -n "$PGPT_FLOWER_PERSISTENT_DB" ]]; then - FLOWER_ARGS+=("--db=$PGPT_FLOWER_PERSISTENT_DB") - fi - if [[ -n "$PGPT_FLOWER_MAXIMUM_TASKS" ]]; then - FLOWER_ARGS+=("--max-tasks=$PGPT_FLOWER_MAXIMUM_TASKS") - fi - fi - - # print all args except password - args=("${FLOWER_ARGS[@]}") - args=("${args[@]//:$PGPT_FLOWER_PASSWORD/}") - - echo "Starting flower with args: ${args[*]}" - python -m celery --app "$CELERY_APP_MODULE" flower "${FLOWER_ARGS[@]}" & -fi - -# Run celery worker using the config -if [[ "$PGPT_WORKER_MODE" == "worker" || "$PGPT_WORKER_MODE" == "mixed" ]]; then - WORKER_ARGS=() - if [[ -n "$PGPT_CELERY_LOG_LEVEL" ]]; then - WORKER_ARGS+=("--loglevel=$PGPT_CELERY_LOG_LEVEL") - fi - if [[ -n "$PGPT_CELERY_QUEUES" ]]; then - WORKER_ARGS+=("--queues=$PGPT_CELERY_QUEUES") - if [[ -n "$PGPT_CELERY_HOSTNAME" ]]; then - PGPT_CELERY_HOSTNAME="${PGPT_CELERY_QUEUES}" - fi - fi - if [[ -n "$PGPT_CELERY_HOSTNAME" ]]; then - WORKER_ARGS+=("--hostname=${PGPT_CELERY_HOSTNAME}%h") - fi - if [[ -n "$PGPT_CELERY_POOL" ]]; then - WORKER_ARGS+=("--pool=$PGPT_CELERY_POOL") - fi - if [[ -n "$PGPT_CELERY_CONCURRENCY" ]]; then - WORKER_ARGS+=("--concurrency=$PGPT_CELERY_CONCURRENCY") - fi - if [[ -n "$PGPT_CELERY_TIME_LIMIT" ]]; then - WORKER_ARGS+=("--time-limit=$PGPT_CELERY_TIME_LIMIT") - fi - - echo "Starting celery worker with args: ${WORKER_ARGS[*]}" - python -m celery --app "$CELERY_APP_MODULE" worker "${WORKER_ARGS[@]}" & -fi - -# Start the Healthcheck server -if [[ "$API_ENABLED" == "true" ]]; then - echo "Starting healthcheck server on port $API_PORT" - python -m uvicorn "$HEALTHCHECK_APP_MODULE" --host 0.0.0.0 --port $API_PORT --no-access-log --log-level critical & -fi - -# Wait for all background processes -wait +exec python -m "${PGPT_WORKER_APP_MODULE}.worker" "$@" diff --git a/settings-test.yaml b/settings-test.yaml index 589ec749..60cb063b 100644 --- a/settings-test.yaml +++ b/settings-test.yaml @@ -15,6 +15,13 @@ qdrant: stream: broker: memory +scheduler: + ingestion: + # Tests run local ingestion and local retrieval in the same process. This + # mode reuses the already-open local Qdrant handle instead of opening a + # second sync/async client pair against the same on-disk database. + mode: local + docling: mode: api use_ocr: true @@ -85,4 +92,3 @@ skills: storage_provider: local skill_injection_mode: system_prompt maximum_loaded_skills: 1 - diff --git a/settings.yaml b/settings.yaml index 890f1d9b..6b1d0ab2 100644 --- a/settings.yaml +++ b/settings.yaml @@ -62,6 +62,7 @@ retrieval: max_merging_recalculations: ${PGPT_MAX_MERGING_RECALCULATIONS:3} chat: + engine_mode: ${PGPT_CHAT_ENGINE_MODE:async} allow_use_default_prompt: ${PGPT_ALLOW_USE_DEFAULT_PROMPT:true} allow_use_context: ${PGPT_ALLOW_USE_CONTEXT:true} allow_generate_citations: ${PGPT_ALLOW_GENERATE_CITATIONS:true} @@ -80,7 +81,6 @@ chat: tldr_timeout: ${PGPT_TLDR_TIMEOUT:300} tldr_minimum_threshold_seconds: ${PGPT_TLDR_MINIMUM_THRESHOLD_SECONDS:1.5} multiplexing_threshold: ${PGPT_MULTIPLEXING_THRESHOLD:10} - maximum_concurrent_requests: ${PGPT_MAXIMUM_CONCURRENT_REQUESTS:10} maximum_blob_size: ${PGPT_MAXIMUM_BLOB_SIZE:26214400} preprocess: documents: @@ -163,6 +163,7 @@ celery: soft_time_limit: ${PGPT_CELERY_SOFT_TIME_LIMIT:} hard_time_limit: ${PGPT_CELERY_HARD_TIME_LIMIT:} visibility_timeout: ${PGPT_CELERY_VISIBILITY_TIMEOUT:21600} + max_tasks_per_child: ${PGPT_CELERY_MAX_TASKS_PER_CHILD:1000} # Broker to which workers will notify the result of tasks tasks_results_broker: @@ -290,6 +291,18 @@ skills: semaphore: mode: ${PGPT_SEMAPHORE_MODE:memory} +scheduler: + ingestion: + mode: ${PGPT_INGESTION_SCHEDULER_MODE:local} + celery_queue: ${PGPT_INGESTION_CELERY_QUEUE:ingestion} + chat: + mode: ${PGPT_CHAT_SCHEDULER_MODE:local} + celery_queue: ${PGPT_CHAT_CELERY_QUEUE:chat} + callback_timeout_seconds: ${PGPT_CHAT_CALLBACK_TIMEOUT_SECONDS:300} + tools: + mode: ${PGPT_TOOLS_SCHEDULER_MODE:local} + celery_queue: ${PGPT_TOOLS_CELERY_QUEUE:tools} + transformation: pptx: vision: @@ -320,4 +333,3 @@ transformation: max_iterations: ${PGPT_TRANSFORMATION_VISION_DOCUMENTS_VISION_MAX_ITERATIONS:2} enable_semaphore: ${PGPT_TRANSFORMATION_VISION_DOCUMENTS_ENABLE_SEMAPHORE:false} mode: ${PGPT_TRANSFORMATION_VISION_DOCUMENTS_VISION_MODE:none} - diff --git a/tests/arq/tasks/chat/test_start.py b/tests/arq/tasks/chat/test_start.py new file mode 100644 index 00000000..1f6cca46 --- /dev/null +++ b/tests/arq/tasks/chat/test_start.py @@ -0,0 +1,53 @@ +from unittest.mock import AsyncMock + +import pytest + +from private_gpt.arq.tasks.chat import enqueue_start_chat_job +from private_gpt.arq.tasks.chat.settings import ( + START_CHAT_TASK_NAME, + get_queue_name, +) +from private_gpt.settings.settings import settings + + +@pytest.mark.parametrize( + ("job_id", "expected_job_id"), + [ + (None, "execution-id:start"), + ("custom-job-id", "custom-job-id"), + ], +) +async def test_enqueue_start_chat_job_dispatches_generic_arq_job( + monkeypatch: pytest.MonkeyPatch, + job_id: str | None, + expected_job_id: str, +) -> None: + enqueue_job = AsyncMock() + monkeypatch.setattr( + "private_gpt.arq.tasks.chat.start.enqueue_job", + enqueue_job, + ) + + request_data = {"messages": [{"role": "user", "content": "Hello"}]} + metadata = {"conversation_id": "conversation-id"} + + await enqueue_start_chat_job( + request_data=request_data, + correlation_id="execution-id", + stream_type="text/event-stream", + metadata=metadata, + job_id=job_id, + ) + + enqueue_job.assert_awaited_once_with( + task_name=START_CHAT_TASK_NAME, + queue_name=get_queue_name(settings()), + args=( + request_data, + "execution-id", + "text/event-stream", + metadata, + ), + job_id=expected_job_id, + correlation_id="execution-id", + ) diff --git a/tests/arq/test_healthcheck.py b/tests/arq/test_healthcheck.py new file mode 100644 index 00000000..35638246 --- /dev/null +++ b/tests/arq/test_healthcheck.py @@ -0,0 +1,9 @@ +from private_gpt.arq.healthcheck import health + + +async def test_healthcheck_does_not_require_redis() -> None: + assert await health() == { + "status": "healthy", + "mode": "arq-worker", + "services": {"worker": "healthy"}, + } diff --git a/tests/arq/test_runner.py b/tests/arq/test_runner.py new file mode 100644 index 00000000..33bb1589 --- /dev/null +++ b/tests/arq/test_runner.py @@ -0,0 +1,37 @@ +import pytest + +from private_gpt.arq.runner import _queue_name, _task_packages + + +def test_task_packages_are_loaded_from_environment( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv( + "PGPT_ARQ_TASK_PACKAGES", + "private_gpt.arq.tasks.chat, custom.tasks ", + ) + + assert _task_packages() == ( + "private_gpt.arq.tasks.chat", + "custom.tasks", + ) + + +def test_task_packages_are_required(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("PGPT_ARQ_TASK_PACKAGES", raising=False) + + with pytest.raises(ValueError, match="PGPT_ARQ_TASK_PACKAGES"): + _task_packages() + + +def test_queue_is_loaded_from_environment(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("PGPT_ARQ_QUEUE", "chat") + + assert _queue_name() == "private_gpt:arq:queue:chat" + + +def test_queue_is_required(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("PGPT_ARQ_QUEUE", raising=False) + + with pytest.raises(ValueError, match="PGPT_ARQ_QUEUE"): + _queue_name() diff --git a/tests/arq/test_task_registry.py b/tests/arq/test_task_registry.py new file mode 100644 index 00000000..2207135f --- /dev/null +++ b/tests/arq/test_task_registry.py @@ -0,0 +1,20 @@ +from private_gpt.arq.task_registry import get_task_packages +from private_gpt.arq.tasks import autodiscover_registered_tasks + + +def test_explicit_task_packages_replace_defaults() -> None: + assert get_task_packages("private_gpt.arq.tasks.chat") == ( + "private_gpt.arq.tasks.chat", + ) + + +def test_chat_package_discovers_only_chat_tasks() -> None: + assert { + task.name + for task in autodiscover_registered_tasks("private_gpt.arq.tasks.chat") + } == { + "private_gpt.chat.resume_iteration", + "private_gpt.chat.start", + "private_gpt.chat.timeout", + "private_gpt.tool.resume", + } diff --git a/tests/celery/test_chat_worker.py b/tests/celery/test_chat_worker.py new file mode 100644 index 00000000..f4f57e69 --- /dev/null +++ b/tests/celery/test_chat_worker.py @@ -0,0 +1,57 @@ +import asyncio +from typing import Any +from unittest.mock import Mock + +import pytest + +from private_gpt.celery.base import StatefulBackgroundTask + + +def test_stateful_task_reuses_one_event_loop() -> None: + class LoopIdTask(StatefulBackgroundTask): + name = "test_chat_loop_id_task" + + @classmethod + def warm_up(cls) -> None: + cls._ensure_runtime() + cls._warmed = True + + async def run(self, *args: Any, **kwargs: Any) -> int: + return id(asyncio.get_running_loop()) + + try: + task = LoopIdTask() + + first_loop_id = task() + second_loop_id = task() + + assert first_loop_id == second_loop_id + finally: + LoopIdTask.shutdown_runtime() + + +async def test_stateful_worker_uses_explicit_warm_profile( + monkeypatch: pytest.MonkeyPatch, +) -> None: + injector = Mock() + warm = Mock() + monkeypatch.setenv("PGPT_STATEFUL_WORKER_TYPE", "stateful-type") + monkeypatch.setenv("PGPT_WORKER_WARM_PROFILE", "tools") + monkeypatch.setattr( + "private_gpt.celery.base.get_global_injector", + Mock(return_value=injector), + ) + monkeypatch.setattr("private_gpt.eager_loading.warm", warm) + + await StatefulBackgroundTask._warm_async() + + warm.assert_called_once_with(injector, profile="tools") + + +async def test_stateful_worker_requires_explicit_warm_profile( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("PGPT_WORKER_WARM_PROFILE", raising=False) + + with pytest.raises(ValueError, match="PGPT_WORKER_WARM_PROFILE"): + await StatefulBackgroundTask._warm_async() diff --git a/tests/celery/test_healthcheck.py b/tests/celery/test_healthcheck.py new file mode 100644 index 00000000..286c7c4d --- /dev/null +++ b/tests/celery/test_healthcheck.py @@ -0,0 +1,32 @@ +import pytest + +from private_gpt.celery import healthcheck + + +@pytest.mark.parametrize("mode", ["", "mixed", "worker", "unknown"]) +async def test_healthcheck_rejects_unregistered_modes( + monkeypatch: pytest.MonkeyPatch, + mode: str, +) -> None: + monkeypatch.setenv("PGPT_WORKER_MODE", mode) + + status = await healthcheck.health_check() + + assert status["status"] == "unhealthy" + assert status["services"] == {} + + +async def test_healthcheck_checks_only_celery_worker( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("PGPT_WORKER_MODE", "celery") + monkeypatch.setattr(healthcheck, "check_worker", lambda: _async_result(True)) + + status = await healthcheck.health_check() + + assert status["status"] == "healthy" + assert status["services"] == {"worker": "healthy"} + + +async def _async_result(value: bool) -> bool: + return value diff --git a/tests/celery/test_task_registry.py b/tests/celery/test_task_registry.py new file mode 100644 index 00000000..3c5b03e6 --- /dev/null +++ b/tests/celery/test_task_registry.py @@ -0,0 +1,52 @@ +import ast +import os +import subprocess +import sys + +import pytest + +from private_gpt.celery.task_registry import get_task_packages + + +def test_explicit_task_packages_replace_defaults() -> None: + assert get_task_packages("private_gpt.celery.tasks.tools") == ( + "private_gpt.celery.tasks.tools", + ) + + +@pytest.mark.parametrize( + ("task_package", "expected_tasks"), + [ + ( + "private_gpt.celery.tasks.ingestion", + { + "private_gpt.ingestion.delete", + "private_gpt.ingestion.vector_index", + }, + ), + ( + "private_gpt.celery.tasks.tools", + {"private_gpt.tools.run"}, + ), + ], +) +def test_worker_registers_only_configured_task_package( + task_package: str, + expected_tasks: set[str], +) -> None: + env = os.environ.copy() + env["PGPT_CELERY_TASK_PACKAGES"] = task_package + output = subprocess.check_output( + [ + sys.executable, + "-c", + "from private_gpt.celery.celery import celery_app; " + "celery_app.loader.import_default_modules(); " + "print(sorted(name for name in celery_app.tasks " + "if name.startswith('private_gpt.')))", + ], + env=env, + text=True, + ) + + assert set(ast.literal_eval(output.strip())) == expected_tasks diff --git a/tests/components/streaming/__init__.py b/tests/components/streaming/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/components/streaming/test_chat_scheduler.py b/tests/components/streaming/test_chat_scheduler.py new file mode 100644 index 00000000..9941a78d --- /dev/null +++ b/tests/components/streaming/test_chat_scheduler.py @@ -0,0 +1,87 @@ +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +from injector import Injector + +from private_gpt.components.streaming.tasks.chat_scheduler import ( + ArqChatScheduler, + ChatSchedulerFactory, + LocalChatScheduler, +) + + +@pytest.fixture +def injector() -> Injector: + return Injector() + + +@pytest.mark.anyio +async def test_arq_chat_scheduler_cancel_aborts_job( + monkeypatch: pytest.MonkeyPatch, +) -> None: + abort_chat_job = AsyncMock(return_value=True) + monkeypatch.setattr( + "private_gpt.components.streaming.tasks.chat_scheduler.abort_chat_job", + abort_chat_job, + ) + + scheduler = ArqChatScheduler() + cancelled = await scheduler.cancel("msg-arq-3") + + assert cancelled is True + abort_chat_job.assert_awaited_once_with(correlation_id="msg-arq-3") + + +@pytest.mark.anyio +async def test_local_chat_scheduler_cancel_cancels_task() -> None: + async def _work() -> None: + await asyncio.sleep(100) + + task = asyncio.create_task(_work(), name="chat_msg-local-2") + + scheduler = LocalChatScheduler() + cancelled = await scheduler.cancel("msg-local-2") + + assert cancelled is True + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + +@pytest.mark.anyio +async def test_local_chat_scheduler_cancel_returns_false_when_no_task() -> None: + scheduler = LocalChatScheduler() + cancelled = await scheduler.cancel("nonexistent-msg") + assert cancelled is False + + +def test_chat_scheduler_factory_selects_local_mode(injector: Injector) -> None: + factory = ChatSchedulerFactory( + settings=SimpleNamespace( + scheduler=SimpleNamespace(chat=SimpleNamespace(mode="local")) + ), + injector=injector, + ) + assert isinstance(factory.get(), LocalChatScheduler) + + +def test_chat_scheduler_factory_selects_arq_mode(injector: Injector) -> None: + factory = ChatSchedulerFactory( + settings=SimpleNamespace( + scheduler=SimpleNamespace(chat=SimpleNamespace(mode="arq")) + ), + injector=injector, + ) + assert isinstance(factory.get(), ArqChatScheduler) + + +def test_chat_scheduler_factory_raises_on_unknown_mode(injector: Injector) -> None: + with pytest.raises(ValueError, match=r"Unknown scheduler\.chat\.mode"): + ChatSchedulerFactory( + settings=SimpleNamespace( + scheduler=SimpleNamespace(chat=SimpleNamespace(mode="missing")) + ), + injector=injector, + ) diff --git a/tests/components/streaming/test_stream_multiplexer.py b/tests/components/streaming/test_stream_multiplexer.py new file mode 100644 index 00000000..e83bea3c --- /dev/null +++ b/tests/components/streaming/test_stream_multiplexer.py @@ -0,0 +1,703 @@ +"""Tests for StreamMultiplexer and the streaming readers. + +The core strategy is **differential testing**: each scenario is run through the +direct path (StreamReader.stream_events) AND the multiplexed path +(StreamMultiplexer), and the two results must be identical. If they ever +diverge, the multiplexer is broken. + +Section layout +-------------- +1. Differential tests — direct == multiplexed for every scenario. +2. Multiplexer concurrency — late joiners, concurrent remove, multi-consumer. +3. Terminal detection — events in the terminal gap are drained. +4. Lifecycle — stop/restart leave no stale state. +5. Performance — worker does not busy-poll. +6. Backend block semantics — block_ms=None is non-blocking. +""" + +import asyncio +import json +import random +from typing import Any + +import pytest +from pydantic import BaseModel + +from private_gpt.components.streaming.providers.in_memory_stream_service import ( + InMemoryStreamService, +) +from private_gpt.components.streaming.providers.models import StreamStatus +from private_gpt.components.streaming.providers.stream_service import StreamService +from private_gpt.components.streaming.stream.stream_reader import ( + DEFAULT_BLOCK_MS, + StreamConsumer, + StreamMultiplexer, + StreamReader, +) + +STRESS_N = 500 + + +# ──────────────────────────── helpers / models ──────────────────────────────── + + +class SimpleEvent(BaseModel): + n: int + + +class SimpleEventHandler: + def serialize(self, event: BaseModel) -> str: + return event.model_dump_json() + + def deserialize(self, data: str) -> BaseModel: + return SimpleEvent.model_validate_json(data) + + async def get_current_status(self, event: BaseModel) -> StreamStatus | None: + return None + + def error_event(self, correlation_id: str, error: Exception) -> BaseModel: + return SimpleEvent(n=-1) + + +class SimpleEventHandler2(SimpleEventHandler): + """Distinct type so add_consumer creates a separate consumer.""" + + +class MockStreamComponent: + def __init__(self, service: StreamService) -> None: + self.stream = service + + +async def collect_from_consumer( + consumer: StreamConsumer, + timeout: float = 10.0, +) -> list[BaseModel]: + """Collect all events from a consumer's broadcast until it is closed.""" + events: list[BaseModel] = [] + + async def drain() -> None: + async for event in consumer.broadcast.read_from(cursor=consumer.cursor): + events.append(event) + + await asyncio.wait_for(drain(), timeout=timeout) + return events + + +async def run_direct( + stream_reader: StreamReader, + handler: SimpleEventHandler, + cid: str, + last_id: str = "0", +) -> list[BaseModel]: + received: list[BaseModel] = [] + gen = await stream_reader.stream_events( + event_handler=handler, + correlation_id=cid, + last_id=last_id, + ) + async for event in gen: + received.append(event) + return received + + +async def run_multiplexed( + multiplexer: StreamMultiplexer, + handler: SimpleEventHandler, + cid: str, + last_id: str = "0", +) -> list[BaseModel]: + consumer = await multiplexer.add_consumer(cid, handler, last_id=last_id) + return await collect_from_consumer(consumer, timeout=30.0) + + +async def push_n_and_complete( + service: StreamService, + cid: str, + n: int, + *, + burst: int = 10, + max_jitter: float = 0.0, +) -> None: + for i in range(n): + await service.push_event(cid, json.dumps({"n": i})) + if (i + 1) % burst == 0: + await asyncio.sleep(random.uniform(0, max_jitter)) + await asyncio.sleep(random.uniform(0, max_jitter * 2)) + await service.update_stream_status(cid, StreamStatus.COMPLETED) + + +def assert_ordered_0_to_n(events: list[BaseModel], n: int) -> None: + ns = [e.n for e in events] # type: ignore[attr-defined] + missing = sorted(set(range(n)) - set(ns)) + duplicates = sorted({x for x in ns if ns.count(x) > 1}) + assert not missing, f"missing {len(missing)} events: {missing[:10]}" + assert not duplicates, f"duplicates: {duplicates[:10]}" + assert ns == sorted(ns), "events arrived out of order" + assert len(ns) == n + + +# ──────────────────────────── fixtures ─────────────────────────────────────── + + +@pytest.fixture +def service() -> StreamService: + return InMemoryStreamService() + + +@pytest.fixture +def stream_reader(service: StreamService) -> StreamReader: + return StreamReader(MockStreamComponent(service)) # type: ignore[arg-type] + + +@pytest.fixture +def multiplexer(stream_reader: StreamReader) -> StreamMultiplexer: + return StreamMultiplexer(stream_reader) + + +@pytest.fixture +def handler() -> SimpleEventHandler: + return SimpleEventHandler() + + +@pytest.fixture +def handler2() -> SimpleEventHandler2: + return SimpleEventHandler2() + + +# ═══════════════════════════════════════════════════════════════════════════════ +# 1. Differential tests — direct path must equal multiplexed path +# ═══════════════════════════════════════════════════════════════════════════════ + + +async def test_both_paths_prepopulated( + service: StreamService, + stream_reader: StreamReader, + multiplexer: StreamMultiplexer, + handler: SimpleEventHandler, +) -> None: + cid = await service.create_stream("test") + await service.push_event(cid, json.dumps({"n": 1})) + await service.push_event(cid, json.dumps({"n": 2})) + await service.push_event(cid, json.dumps({"n": 3})) + await service.update_stream_status(cid, StreamStatus.COMPLETED) + + await multiplexer.start() + try: + mux = await run_multiplexed(multiplexer, handler, cid) + finally: + await multiplexer.stop() + direct = await run_direct(stream_reader, handler, cid) + + assert [e.n for e in mux] == [e.n for e in direct] == [1, 2, 3] # type: ignore[attr-defined] + + +async def test_both_paths_empty_stream( + service: StreamService, + stream_reader: StreamReader, + multiplexer: StreamMultiplexer, + handler: SimpleEventHandler, +) -> None: + cid = await service.create_stream("test") + await service.update_stream_status(cid, StreamStatus.COMPLETED) + + await multiplexer.start() + try: + mux = await run_multiplexed(multiplexer, handler, cid) + finally: + await multiplexer.stop() + direct = await run_direct(stream_reader, handler, cid) + + assert mux == direct == [] + + +async def test_both_paths_reconnect_from_mid( + service: StreamService, + stream_reader: StreamReader, + multiplexer: StreamMultiplexer, + handler: SimpleEventHandler, +) -> None: + """A consumer that already has event n=1 (it knows last_id=e1_id) must + receive only events after that point — never a replay from "0".""" + cid = await service.create_stream("test") + e1_id = await service.push_event(cid, json.dumps({"n": 1})) + await service.push_event(cid, json.dumps({"n": 2})) + await service.push_event(cid, json.dumps({"n": 3})) + await service.update_stream_status(cid, StreamStatus.COMPLETED) + + await multiplexer.start() + try: + mux = await run_multiplexed(multiplexer, handler, cid, last_id=e1_id) + finally: + await multiplexer.stop() + direct = await run_direct(stream_reader, handler, cid, last_id=e1_id) + + assert [e.n for e in mux] == [e.n for e in direct] == [2, 3] # type: ignore[attr-defined] + + +async def test_both_paths_reconnect_from_tail( + service: StreamService, + stream_reader: StreamReader, + multiplexer: StreamMultiplexer, + handler: SimpleEventHandler, +) -> None: + """Consumer is caught up to the last event; both paths return nothing.""" + cid = await service.create_stream("test") + await service.push_event(cid, json.dumps({"n": 1})) + await service.push_event(cid, json.dumps({"n": 2})) + e3_id = await service.push_event(cid, json.dumps({"n": 3})) + await service.update_stream_status(cid, StreamStatus.COMPLETED) + + await multiplexer.start() + try: + mux = await run_multiplexed(multiplexer, handler, cid, last_id=e3_id) + finally: + await multiplexer.stop() + direct = await run_direct(stream_reader, handler, cid, last_id=e3_id) + + assert mux == direct == [] + + +async def test_both_paths_live_producer( + service: StreamService, + stream_reader: StreamReader, + multiplexer: StreamMultiplexer, + handler: SimpleEventHandler, +) -> None: + """Events arrive while both consumers are already blocking on the stream. + Both paths must receive every event in order.""" + cid = await service.create_stream("test") + + await multiplexer.start() + try: + mux_task = asyncio.create_task(run_multiplexed(multiplexer, handler, cid)) + direct_task = asyncio.create_task(run_direct(stream_reader, handler, cid)) + await asyncio.sleep(0.15) + await push_n_and_complete(service, cid, 100, burst=10, max_jitter=0.003) + mux = await asyncio.wait_for(mux_task, timeout=30.0) + direct = await asyncio.wait_for(direct_task, timeout=30.0) + finally: + await multiplexer.stop() + + mux_ns = [e.n for e in mux] # type: ignore[attr-defined] + direct_ns = [e.n for e in direct] # type: ignore[attr-defined] + assert mux_ns == direct_ns + assert_ordered_0_to_n(mux, 100) + assert_ordered_0_to_n(direct, 100) + + +@pytest.mark.parametrize( + ("burst", "max_jitter"), + [ + (1, 0.0), + (10, 0.003), + (STRESS_N, 0.0), + ], + ids=["one-at-a-time", "small-bursts", "single-burst"], +) +async def test_both_paths_stress( + service: StreamService, + stream_reader: StreamReader, + multiplexer: StreamMultiplexer, + handler: SimpleEventHandler, + burst: int, + max_jitter: float, +) -> None: + cid = await service.create_stream("test") + await push_n_and_complete( + service, cid, STRESS_N, burst=burst, max_jitter=max_jitter + ) + + await multiplexer.start() + try: + mux = await run_multiplexed(multiplexer, handler, cid) + finally: + await multiplexer.stop() + direct = await run_direct(stream_reader, handler, cid) + + assert_ordered_0_to_n(mux, STRESS_N) + assert_ordered_0_to_n(direct, STRESS_N) + assert [e.n for e in mux] == [e.n for e in direct] # type: ignore[attr-defined] + + +# ═══════════════════════════════════════════════════════════════════════════════ +# 2. Multiplexer concurrency +# ═══════════════════════════════════════════════════════════════════════════════ + + +async def test_late_joiner_receives_inflight_batch( + service: StreamService, + stream_reader: StreamReader, + multiplexer: StreamMultiplexer, + handler: SimpleEventHandler, + handler2: SimpleEventHandler2, +) -> None: + """A consumer that joins while the worker is blocked inside read_events() + must still receive the batch that the worker is about to dispatch.""" + cid = await service.create_stream("test") + + original_read = stream_reader.read_events + late_box: dict[str, Any] = {} + + async def read_then_join( + event_handler: Any, + correlation_id: str, + last_id: str = "0", + count: int | None = None, + block_ms: int | None = None, + ) -> tuple[list[BaseModel], str]: + result = await original_read( + event_handler=event_handler, + correlation_id=correlation_id, + last_id=last_id, + count=count, + block_ms=block_ms, + ) + if result[0] and "c2" not in late_box: + late_box["c2"] = await multiplexer.add_consumer(cid, handler2) + return result + + stream_reader.read_events = read_then_join # type: ignore[method-assign] + + await multiplexer.start() + c1 = await multiplexer.add_consumer(cid, handler) + await service.push_event(cid, json.dumps({"n": 1})) + await service.update_stream_status(cid, StreamStatus.COMPLETED) + + try: + recv1 = await collect_from_consumer(c1, timeout=10.0) + c2 = late_box["c2"] + recv2 = await collect_from_consumer(c2, timeout=10.0) + finally: + await multiplexer.stop() + stream_reader.read_events = original_read # type: ignore[method-assign] + + assert 1 in [e.n for e in recv1] # type: ignore[attr-defined] + assert 1 in [e.n for e in recv2], "late joiner missed the in-flight batch" # type: ignore[attr-defined] + + +async def test_concurrent_remove_preserves_delivery( + service: StreamService, + stream_reader: StreamReader, + multiplexer: StreamMultiplexer, + handler: SimpleEventHandler, + handler2: SimpleEventHandler2, +) -> None: + """If one consumer is removed while the worker is reading, the remaining + consumers still receive the batch.""" + cid = await service.create_stream("test") + + original_read = stream_reader.read_events + removed = False + + async def read_then_remove( + event_handler: Any, + correlation_id: str, + last_id: str = "0", + count: int | None = None, + block_ms: int | None = None, + ) -> tuple[list[BaseModel], str]: + result = await original_read( + event_handler=event_handler, + correlation_id=correlation_id, + last_id=last_id, + count=count, + block_ms=block_ms, + ) + nonlocal removed + if result[0] and not removed: + await multiplexer.remove_consumer(c2) + removed = True + return result + + stream_reader.read_events = read_then_remove # type: ignore[method-assign] + + await multiplexer.start() + c1 = await multiplexer.add_consumer(cid, handler) + c2 = await multiplexer.add_consumer(cid, handler2) + await service.push_event(cid, json.dumps({"n": 1})) + await service.update_stream_status(cid, StreamStatus.COMPLETED) + + try: + recv1 = await collect_from_consumer(c1, timeout=5.0) + finally: + await multiplexer.stop() + stream_reader.read_events = original_read # type: ignore[method-assign] + + assert [e.n for e in recv1] == [1] # type: ignore[attr-defined] + + +async def test_multiple_consumers_all_receive( + service: StreamService, + stream_reader: StreamReader, + multiplexer: StreamMultiplexer, + handler: SimpleEventHandler, + handler2: SimpleEventHandler2, +) -> None: + """Two distinct consumers on the same stream both receive every event.""" + cid = await service.create_stream("test") + await service.push_event(cid, json.dumps({"n": 1})) + await service.push_event(cid, json.dumps({"n": 2})) + await service.update_stream_status(cid, StreamStatus.COMPLETED) + + await multiplexer.start() + c1 = await multiplexer.add_consumer(cid, handler) + c2 = await multiplexer.add_consumer(cid, handler2) + try: + recv1 = await collect_from_consumer(c1, timeout=10.0) + recv2 = await collect_from_consumer(c2, timeout=10.0) + finally: + await multiplexer.stop() + + assert [e.n for e in recv1] == [1, 2] # type: ignore[attr-defined] + assert [e.n for e in recv2] == [1, 2] # type: ignore[attr-defined] + + +async def test_late_joiner_gets_full_history( + service: StreamService, + stream_reader: StreamReader, + multiplexer: StreamMultiplexer, + handler: SimpleEventHandler, + handler2: SimpleEventHandler2, +) -> None: + """A consumer that joins AFTER the first batch has already been dispatched + must still receive every event from the beginning (last_id='0'). + + This is the 'Delta without start' bug: if a content_block_start event is + in batch 0 and a second consumer only sees batch 1 onwards, the client + receives content_block_delta with no preceding start. + """ + cid = await service.create_stream("test") + + # Patch read_events so that after the first successful dispatch we can + # confirm the broadcast already has batch 0, then add a second consumer. + original_read = stream_reader.read_events + second_consumer_box: dict[str, Any] = {} + dispatch_count = 0 + + async def read_after_first_dispatch( + event_handler: Any, + correlation_id: str, + last_id: str = "0", + count: int | None = None, + block_ms: int | None = None, + ) -> tuple[list[BaseModel], str]: + nonlocal dispatch_count + result = await original_read( + event_handler=event_handler, + correlation_id=correlation_id, + last_id=last_id, + count=count, + block_ms=block_ms, + ) + # After the first batch is returned (and will be dispatched), add c2. + # At this point broadcast._batches is still empty (dispatch hasn't run), + # but after dispatch it will have 1 batch — c2 must still see it. + if result[0] and dispatch_count == 0 and "c2" not in second_consumer_box: + dispatch_count += 1 + # Wait one event-loop tick so _dispatch runs and pushes batch 0. + await asyncio.sleep(0) + second_consumer_box["c2"] = await multiplexer.add_consumer( + cid, handler2, last_id="0" + ) + return result + + stream_reader.read_events = read_after_first_dispatch # type: ignore[method-assign] + + await multiplexer.start() + c1 = await multiplexer.add_consumer(cid, handler, last_id="0") + + await service.push_event(cid, json.dumps({"n": 1})) + await service.push_event(cid, json.dumps({"n": 2})) + await service.update_stream_status(cid, StreamStatus.COMPLETED) + + try: + recv1 = await collect_from_consumer(c1) + c2 = second_consumer_box["c2"] + recv2 = await collect_from_consumer(c2) + finally: + await multiplexer.stop() + stream_reader.read_events = original_read # type: ignore[method-assign] + + assert [e.n for e in recv1] == [1, 2], "consumer 1 must receive all events" # type: ignore[attr-defined] + assert [e.n for e in recv2] == [1, 2], "late joiner must not miss batch 0" # type: ignore[attr-defined] + + +# ═══════════════════════════════════════════════════════════════════════════════ +# 3. Terminal detection — events in the gap are drained +# ═══════════════════════════════════════════════════════════════════════════════ + + +async def test_events_in_terminal_gap_drained( + service: StreamService, + stream_reader: StreamReader, + multiplexer: StreamMultiplexer, + handler: SimpleEventHandler, +) -> None: + """Events published between the worker's empty read and the terminal-status + check must be drained and delivered, not dropped.""" + cid = await service.create_stream("test") + + original_check = stream_reader.check_terminal_status + + async def injecting_check( + correlation_id: str, + event_handler: Any, + cached_events: Any, + ) -> bool: + await service.push_event(correlation_id, json.dumps({"n": 1})) + await service.push_event(correlation_id, json.dumps({"n": 2})) + await service.update_stream_status(correlation_id, StreamStatus.COMPLETED) + return True + + stream_reader.check_terminal_status = injecting_check # type: ignore[method-assign] + + await multiplexer.start() + consumer = await multiplexer.add_consumer(cid, handler) + try: + received = await collect_from_consumer(consumer, timeout=10.0) + finally: + await multiplexer.stop() + stream_reader.check_terminal_status = original_check # type: ignore[method-assign] + + assert [e.n for e in received] == [1, 2] # type: ignore[attr-defined] + + +# ═══════════════════════════════════════════════════════════════════════════════ +# 4. Lifecycle — stop / restart leave no stale state +# ═══════════════════════════════════════════════════════════════════════════════ + + +async def test_stop_clears_state_and_closes_consumers( + service: StreamService, + stream_reader: StreamReader, + multiplexer: StreamMultiplexer, + handler: SimpleEventHandler, +) -> None: + cid = await service.create_stream("test") + await service.push_event(cid, json.dumps({"n": 1})) + + await multiplexer.start() + consumer = await multiplexer.add_consumer(cid, handler) + + # Wait until the worker has pushed at least one batch to the broadcast. + async def get_one() -> None: + async for _ in consumer.broadcast.read_from(cursor=consumer.cursor): + break + + await asyncio.wait_for(get_one(), timeout=5.0) + + await multiplexer.stop() + + assert not multiplexer.consumers, "consumers dict must be empty after stop()" + assert not multiplexer.states, "states dict must be empty after stop()" + assert consumer.broadcast.is_closed, "broadcast must be closed after stop()" + + +async def test_restart_after_stop( + service: StreamService, + stream_reader: StreamReader, + multiplexer: StreamMultiplexer, + handler: SimpleEventHandler, +) -> None: + cid = await service.create_stream("test") + await service.push_event(cid, json.dumps({"n": 1})) + await service.update_stream_status(cid, StreamStatus.COMPLETED) + + await multiplexer.start() + await multiplexer.add_consumer(cid, handler) + await multiplexer.stop() + + # Second stream on a fresh cid after restart. + cid2 = await service.create_stream("test") + await service.push_event(cid2, json.dumps({"n": 42})) + await service.update_stream_status(cid2, StreamStatus.COMPLETED) + + await multiplexer.start() + try: + consumer = await multiplexer.add_consumer(cid2, handler) + received = await collect_from_consumer(consumer, timeout=10.0) + finally: + await multiplexer.stop() + + assert [e.n for e in received] == [42] # type: ignore[attr-defined] + + +async def test_stop_idempotent( + multiplexer: StreamMultiplexer, +) -> None: + await multiplexer.start() + await multiplexer.stop() + await multiplexer.stop() # must not raise + + +# ═══════════════════════════════════════════════════════════════════════════════ +# 5. Performance — worker does not busy-poll +# ═══════════════════════════════════════════════════════════════════════════════ + + +async def test_no_busy_poll( + service: StreamService, + stream_reader: StreamReader, + multiplexer: StreamMultiplexer, + handler: SimpleEventHandler, +) -> None: + """With block_ms=DEFAULT_BLOCK_MS the worker parks on the backend; it must + not spin. In 300 ms we expect at most one read call.""" + cid = await service.create_stream("test") + + read_count = 0 + original_read = stream_reader.read_events + + async def counting_read( + event_handler: Any, + correlation_id: str, + last_id: str = "0", + count: int | None = None, + block_ms: int | None = None, + ) -> tuple[list[BaseModel], str]: + nonlocal read_count + read_count += 1 + return await original_read( + event_handler=event_handler, + correlation_id=correlation_id, + last_id=last_id, + count=count, + block_ms=block_ms, + ) + + stream_reader.read_events = counting_read # type: ignore[method-assign] + + await multiplexer.start() + await multiplexer.add_consumer(cid, handler) + await asyncio.sleep(0.3) + await multiplexer.stop() + stream_reader.read_events = original_read # type: ignore[method-assign] + + assert read_count <= 1, ( + f"Worker called read_events {read_count} times in 300 ms " + f"(expected ≤1 with block_ms={DEFAULT_BLOCK_MS})." + ) + + +# ═══════════════════════════════════════════════════════════════════════════════ +# 6. Backend block semantics — block_ms=None is non-blocking +# ═══════════════════════════════════════════════════════════════════════════════ + + +async def test_block_ms_none_is_non_blocking( + service: StreamService, + stream_reader: StreamReader, + handler: SimpleEventHandler, +) -> None: + """block_ms=None means non-blocking on both backends. This is what the + direct path's final drain relies on, and it must never hang.""" + cid = await service.create_stream("test") + await service.push_event(cid, json.dumps({"n": 1})) + e1_id = await service.push_event(cid, json.dumps({"n": 2})) + + # Read past e1; no events remain. + events, _ = await asyncio.wait_for( + stream_reader.read_events(handler, cid, e1_id, None, None), + timeout=2.0, + ) + assert events == [] diff --git a/tests/components/tools/test_tool_scheduler.py b/tests/components/tools/test_tool_scheduler.py new file mode 100644 index 00000000..8492f844 --- /dev/null +++ b/tests/components/tools/test_tool_scheduler.py @@ -0,0 +1,59 @@ +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +from private_gpt.components.chat.models.chat_config_models import ToolSpec +from private_gpt.components.tools.remote_execution import ToolExecutionRequest +from private_gpt.components.tools.tool_scheduler import CeleryToolScheduler + + +@pytest.mark.anyio +async def test_celery_tool_scheduler_execute_raises_not_implemented() -> None: + scheduler = CeleryToolScheduler( + settings=SimpleNamespace( + scheduler=SimpleNamespace(tools=SimpleNamespace(celery_queue="tools")) + ), + ) + + request = ToolExecutionRequest.model_validate( + { + "tool_id": "tool-1", + "tool_name": "bash", + "tool_kwargs": {}, + "tool_spec": ToolSpec(name="bash", runtime="server", input_schema={}), + "context": {"correlation_id": "msg-1", "messages": []}, + } + ) + + with pytest.raises(NotImplementedError): + await scheduler.execute(request) + + +@pytest.mark.anyio +async def test_celery_tool_scheduler_cancel_revokes_task( + monkeypatch: pytest.MonkeyPatch, +) -> None: + celery_app = MagicMock() + monkeypatch.setattr("private_gpt.celery.celery.celery_app", celery_app) + + scheduler = CeleryToolScheduler( + settings=SimpleNamespace( + scheduler=SimpleNamespace(tools=SimpleNamespace(celery_queue="tools")) + ), + ) + + request = ToolExecutionRequest.model_validate( + { + "tool_id": "tool-1", + "tool_name": "bash", + "tool_kwargs": {}, + "tool_spec": ToolSpec(name="bash", runtime="server", input_schema={}), + "context": {"correlation_id": "msg-1", "messages": []}, + } + ) + + cancelled = await scheduler.cancel(request, task_id="task-abc") + + assert cancelled is True + celery_app.control.revoke.assert_called_once_with("task-abc", terminate=True) diff --git a/tests/engines/test_async_chat_engine.py b/tests/engines/test_async_chat_engine.py new file mode 100644 index 00000000..2671c0c4 --- /dev/null +++ b/tests/engines/test_async_chat_engine.py @@ -0,0 +1,539 @@ +import asyncio +from collections.abc import AsyncGenerator +from dataclasses import dataclass +from typing import Any +from unittest.mock import MagicMock + +import pytest +from llama_index.core.base.llms.types import ChatMessage, MessageRole +from llama_index.core.llms.llm import ToolSelection + +from private_gpt.components.chat.models.chat_config_models import ( + ResolvedChatRequest, + ResolvedSystemConfig, + ResolvedToolConfig, + ToolSpec, +) +from private_gpt.components.engines.chat.async_chat_engine import ( + AsyncChatEngine, + IterationCheckpointPayload, + LocalEventChannel, +) +from private_gpt.components.engines.chat.chat_engine import ChatLoopEngine +from private_gpt.components.engines.chat.models.chat_state import ( + ChatState, + ChatStatus, +) +from private_gpt.components.llm.llm_component import LLMComponent +from private_gpt.components.tools.remote_execution import ( + ToolExecutionRequest, + ToolExecutionResponse, + build_rebuild_metadata, + execute_tool_request, +) +from private_gpt.components.tools.tool_scheduler import ( + BaseToolScheduler, + LocalToolScheduler, +) +from private_gpt.events.models import ( + RawContentBlockStartEvent, + TextBlock, + ToolResultBlock, +) +from tests.fixtures.mock_function_llm import get_mock_function_calling_llm + + +async def _noop_tool(value: str) -> str: + await asyncio.sleep(0.01) + return f"ok:{value}" + + +def _client_tool(name: str) -> ToolSpec: + return ToolSpec.from_defaults( + name=name, + type=name, + runtime="client", + input_schema={"type": "object", "properties": {"value": {"type": "string"}}}, + ) + + +def _rebuild_server_tool(name: str) -> ToolSpec: + return ToolSpec.from_defaults( + name=name, + type=name, + runtime="server", + async_fn=_noop_tool, + ) + + +def _server_tool(name: str) -> ToolSpec: + return ToolSpec.from_defaults( + name=name, + type=name, + runtime="server", + async_fn=_noop_tool, + execution_metadata=build_rebuild_metadata( + _rebuild_server_tool, + {"name": name}, + ), + ) + + +@pytest.fixture +def base_request() -> ResolvedChatRequest: + return ResolvedChatRequest( + messages=[ChatMessage(role=MessageRole.USER, content="hello")], + system=ResolvedSystemConfig(prompt="test"), + ) + + +class _FakeAsyncToolScheduler(BaseToolScheduler): + def __init__(self) -> None: + self.pending: dict[str, tuple[ToolExecutionRequest, str]] = {} + self.cancelled: list[str] = [] + self._next = 0 + + @property + def is_async(self) -> bool: + return True + + async def execute( + self, + request: ToolExecutionRequest, + state_ctx=None, + interceptors=None, + ) -> ToolExecutionResponse: + del request, state_ctx, interceptors + raise NotImplementedError + + async def async_execute( + self, + request: ToolExecutionRequest, + state_ctx=None, + interceptors=None, + ) -> str: + del state_ctx, interceptors + self._next += 1 + handle = f"handle-{self._next}" + self.pending[request.tool_id] = (request, handle) + return handle + + async def cancel( + self, + request: ToolExecutionRequest, + task_id: str | None = None, + ) -> bool: + del request + if task_id is None: + return False + self.cancelled.append(task_id) + return True + + async def complete_pending(self) -> list[ToolExecutionResponse]: + responses: list[ToolExecutionResponse] = [] + for tool_id in list(self.pending): + request, _ = self.pending.pop(tool_id) + responses.append(await execute_tool_request(request)) + return responses + + +class _FakeChatScheduler: + def __init__(self) -> None: + self.cancelled: list[str] = [] + + async def cancel(self, correlation_id: str) -> bool: + self.cancelled.append(correlation_id) + return True + + +@dataclass +class _AsyncRunResult: + events: list[Any] + states: list[ChatState] + + +def _make_llm_component(mock_llm: Any) -> MagicMock: + llm_component = MagicMock(spec=LLMComponent) + llm_component.get_llm.return_value = mock_llm + return llm_component + + +async def _collect_events(events: AsyncGenerator[Any, None]) -> list[Any]: + return [event async for event in events] + + +def _normalize(value: Any) -> Any: + if hasattr(value, "model_dump"): + value = value.model_dump(mode="json", exclude_none=True) + if isinstance(value, dict): + return { + key: _normalize(item) + for key, item in value.items() + if key + not in { + "block_id", + "id", + "start_timestamp", + "stop_timestamp", + "expires_at", + "tool_use_id", + } + } + if isinstance(value, list): + return [_normalize(item) for item in value] + return value + + +def _normalize_events(events: list[Any]) -> list[Any]: + return [ + { + "type": event.__class__.__name__, + "payload": _normalize(event), + } + for event in events + ] + + +def _tool_result_texts(events: list[Any]) -> list[str]: + texts: list[str] = [] + for event in events: + if isinstance(event, RawContentBlockStartEvent) and isinstance( + event.content_block, ToolResultBlock + ): + for block in event.content_block.content: + if isinstance(block, TextBlock): + texts.append(block.text) + return texts + + +async def _run_sync_engine( + request: ResolvedChatRequest, + mock_llm: Any, + tool_scheduler: BaseToolScheduler | None = None, +) -> list[Any]: + engine = ChatLoopEngine( + llm_component=_make_llm_component(mock_llm), + request_interceptors=[], + response_interceptors=[], + max_iterations=6, + tool_scheduler=tool_scheduler or LocalToolScheduler(), + ) + execution = await engine.run(request) + events = await _collect_events(execution.events) + await execution.final_state_task + return events + + +async def _drain(channel: LocalEventChannel) -> list[Any]: + """Drain all events from a closed LocalEventChannel.""" + return [e async for e in channel.stream()] + + +async def _run_async_engine( + request: ResolvedChatRequest, + mock_llm: Any, + tool_scheduler: BaseToolScheduler, +) -> _AsyncRunResult: + engine = AsyncChatEngine( + llm_component=_make_llm_component(mock_llm), + request_interceptors=[], + response_interceptors=[], + max_iterations=6, + tool_scheduler=tool_scheduler, + chat_scheduler=_FakeChatScheduler(), + ) + + all_events: list[Any] = [] + states: list[ChatState] = [] + + channel = LocalEventChannel() + state = await engine.execute(request, channel=channel) + await channel.close() + all_events.extend(await _drain(channel)) + states.append(state) + + while state.output.status == ChatStatus.WAITING: + assert isinstance(tool_scheduler, _FakeAsyncToolScheduler) + responses = await tool_scheduler.complete_pending() + resumed_request = state.input.request.model_copy(deep=True) + resumed_request.messages = [ + *resumed_request.messages, + *(response.tool_message for response in responses), + ] + channel2 = LocalEventChannel() + state = await engine.resume( + state.output.pause_type, + resumed_request, + iteration=state.runtime.iteration, + next_block_count=state.runtime.next_block_count, + checkpoint_payload=IterationCheckpointPayload( + pending_async_tools=state.output.pending_async_tools, + tool_responses=responses, + pending_external_tool_calls=state.output.pending_external_tool_calls, + total_input_tokens=state.runtime.total_input_tokens, + total_output_tokens=state.runtime.total_output_tokens, + has_input_usage=state.runtime.has_input_usage, + has_output_usage=state.runtime.has_output_usage, + ), + channel=channel2, + ) + await channel2.close() + all_events.extend(await _drain(channel2)) + states.append(state) + + return _AsyncRunResult(events=all_events, states=states) + + +@pytest.mark.asyncio +async def test_async_engine_matches_sync_simple_message( + base_request: ResolvedChatRequest, +) -> None: + sync_events = await _run_sync_engine( + base_request.model_copy(deep=True), + get_mock_function_calling_llm(["hello", " world"]), + ) + async_result = await _run_async_engine( + base_request.model_copy(deep=True), + get_mock_function_calling_llm(["hello", " world"]), + tool_scheduler=_FakeAsyncToolScheduler(), + ) + + assert _normalize_events(async_result.events) == _normalize_events(sync_events) + assert [state.output.status for state in async_result.states] == [ + ChatStatus.COMPLETED, + ] + + +@pytest.mark.asyncio +async def test_async_engine_matches_sync_one_client_tool_and_stops_first_iteration( + base_request: ResolvedChatRequest, +) -> None: + request = base_request.model_copy(deep=True) + request.tool_config = ResolvedToolConfig(tools=[_client_tool("browser")]) + sync_events = await _run_sync_engine( + request.model_copy(deep=True), + get_mock_function_calling_llm( + [ + [ + ToolSelection( + tool_id="tool_1", + tool_name="browser", + tool_kwargs={"value": "x"}, + ) + ] + ] + ), + ) + async_result = await _run_async_engine( + request.model_copy(deep=True), + get_mock_function_calling_llm( + [ + [ + ToolSelection( + tool_id="tool_1", + tool_name="browser", + tool_kwargs={"value": "x"}, + ) + ] + ] + ), + tool_scheduler=_FakeAsyncToolScheduler(), + ) + + assert _normalize_events(async_result.events) == _normalize_events(sync_events) + assert async_result.states[0].output.status == ChatStatus.COMPLETED + assert async_result.states[0].output.stop_reason == "tool_use" + assert len(async_result.states[0].output.pending_external_tool_calls) == 1 + + +@pytest.mark.asyncio +async def test_async_engine_matches_sync_one_server_tool_and_resumes_same_point( + base_request: ResolvedChatRequest, +) -> None: + request = base_request.model_copy(deep=True) + request.tool_config = ResolvedToolConfig(tools=[_server_tool("echo")]) + sync_events = await _run_sync_engine( + request.model_copy(deep=True), + get_mock_function_calling_llm( + [ + [ + ToolSelection( + tool_id="tool_1", tool_name="echo", tool_kwargs={"value": "x"} + ) + ], + ["done"], + ] + ), + ) + async_result = await _run_async_engine( + request.model_copy(deep=True), + get_mock_function_calling_llm( + [ + [ + ToolSelection( + tool_id="tool_1", tool_name="echo", tool_kwargs={"value": "x"} + ) + ], + ["done"], + ] + ), + tool_scheduler=_FakeAsyncToolScheduler(), + ) + + assert [state.output.status for state in async_result.states] == [ + ChatStatus.WAITING, + ChatStatus.COMPLETED, + ] + assert async_result.states[0].output.pause_type == "tools" + assert _tool_result_texts(async_result.events) == ["ok:x"] + assert _tool_result_texts(sync_events) == ["ok:x"] + assert _normalize_events(async_result.events) == _normalize_events(sync_events) + + +@pytest.mark.asyncio +async def test_async_engine_matches_sync_two_server_tools_plus_one_client_tool( + base_request: ResolvedChatRequest, +) -> None: + request = base_request.model_copy(deep=True) + request.tool_config = ResolvedToolConfig( + tools=[ + _server_tool("echo"), + _server_tool("echo2"), + _client_tool("browser"), + ] + ) + sync_events = await _run_sync_engine( + request.model_copy(deep=True), + get_mock_function_calling_llm( + [ + [ + ToolSelection( + tool_id="tool_1", tool_name="echo", tool_kwargs={"value": "a"} + ), + ToolSelection( + tool_id="tool_2", tool_name="echo2", tool_kwargs={"value": "b"} + ), + ToolSelection( + tool_id="tool_3", + tool_name="browser", + tool_kwargs={"value": "c"}, + ), + ] + ] + ), + ) + async_result = await _run_async_engine( + request.model_copy(deep=True), + get_mock_function_calling_llm( + [ + [ + ToolSelection( + tool_id="tool_1", tool_name="echo", tool_kwargs={"value": "a"} + ), + ToolSelection( + tool_id="tool_2", tool_name="echo2", tool_kwargs={"value": "b"} + ), + ToolSelection( + tool_id="tool_3", + tool_name="browser", + tool_kwargs={"value": "c"}, + ), + ] + ] + ), + tool_scheduler=_FakeAsyncToolScheduler(), + ) + + assert async_result.states[0].output.status == ChatStatus.WAITING + assert async_result.states[1].output.status == ChatStatus.COMPLETED + assert async_result.states[1].output.stop_reason == "tool_use" + assert len(async_result.states[1].output.pending_external_tool_calls) == 1 + assert sorted(_tool_result_texts(async_result.events)) == ["ok:a", "ok:b"] + assert sorted(_tool_result_texts(sync_events)) == ["ok:a", "ok:b"] + assert _normalize_events(async_result.events) == _normalize_events(sync_events) + + +@pytest.mark.asyncio +async def test_async_engine_matches_sync_across_multiple_server_tool_iterations( + base_request: ResolvedChatRequest, +) -> None: + request = base_request.model_copy(deep=True) + request.tool_config = ResolvedToolConfig(tools=[_server_tool("echo")]) + deltas = [ + [ToolSelection(tool_id="tool_1", tool_name="echo", tool_kwargs={"value": "a"})], + [ToolSelection(tool_id="tool_2", tool_name="echo", tool_kwargs={"value": "b"})], + ["all", " done"], + ] + sync_events = await _run_sync_engine( + request.model_copy(deep=True), + get_mock_function_calling_llm(deltas), + ) + async_result = await _run_async_engine( + request.model_copy(deep=True), + get_mock_function_calling_llm(deltas), + tool_scheduler=_FakeAsyncToolScheduler(), + ) + + assert [state.output.status for state in async_result.states] == [ + ChatStatus.WAITING, + ChatStatus.WAITING, + ChatStatus.COMPLETED, + ] + assert _tool_result_texts(async_result.events) == ["ok:a", "ok:b"] + assert _normalize_events(async_result.events) == _normalize_events(sync_events) + + +@pytest.mark.asyncio +async def test_async_engine_cancel_schedules_chat_cancellation( + base_request: ResolvedChatRequest, +) -> None: + request = base_request.model_copy(deep=True) + request.context.correlation_id = "msg-cancel-1" + request.tool_config = ResolvedToolConfig(tools=[_server_tool("echo")]) + mock_llm = get_mock_function_calling_llm( + [ + [ + ToolSelection( + tool_id="tool_1", tool_name="echo", tool_kwargs={"value": "x"} + ) + ] + ] + ) + tool_scheduler = _FakeAsyncToolScheduler() + chat_scheduler = _FakeChatScheduler() + engine = AsyncChatEngine( + llm_component=_make_llm_component(mock_llm), + request_interceptors=[], + response_interceptors=[], + max_iterations=4, + tool_scheduler=tool_scheduler, + chat_scheduler=chat_scheduler, + ) + + channel = LocalEventChannel() + state = await engine.execute(request, channel=channel) + await channel.close() + await _collect_events(channel.stream()) + + assert state.output.status == ChatStatus.WAITING + assert list(state.output.pending_async_tools.values()) == ["handle-1"] + + cancelled = await engine.cancel("msg-cancel-1") + + assert cancelled is True + assert chat_scheduler.cancelled == ["msg-cancel-1"] + + +@pytest.mark.asyncio +async def test_async_engine_cancel_without_scheduler_returns_false() -> None: + engine = AsyncChatEngine( + llm_component=_make_llm_component(get_mock_function_calling_llm(["ok"])), + request_interceptors=[], + response_interceptors=[], + max_iterations=2, + tool_scheduler=_FakeAsyncToolScheduler(), + chat_scheduler=_FakeChatScheduler(), + ) + + # _FakeChatScheduler returns True for any correlation_id + assert await engine.cancel("msg-no-scheduler") is True diff --git a/tests/engines/test_chat_agent_engine.py b/tests/engines/test_chat_agent_engine.py index 97336b50..84ff7d45 100644 --- a/tests/engines/test_chat_agent_engine.py +++ b/tests/engines/test_chat_agent_engine.py @@ -1,3 +1,5 @@ +import asyncio +from collections.abc import AsyncGenerator from typing import Any from unittest.mock import MagicMock @@ -16,8 +18,21 @@ from private_gpt.components.chat.models.chat_config_models import ( ResolvedToolConfig, ToolSpec, ) -from private_gpt.components.engines.chat_loop.chat_loop_engine import ChatLoopEngine +from private_gpt.components.engines.chat.async_chat_engine import ( + AsyncChatEngine, + LocalEventChannel, + _EventHandler, + _StreamDeltaState, +) +from private_gpt.components.engines.chat.chat_engine import ChatLoopEngine +from private_gpt.components.engines.chat.chat_engine_interface import ( + ChatEngine, + LoopChatEngineAdapter, +) +from private_gpt.components.engines.chat.chat_runner import ChatRunner from private_gpt.components.llm.llm_component import LLMComponent +from private_gpt.components.streaming.tasks.chat_scheduler import LocalChatScheduler +from private_gpt.components.tools.tool_scheduler import LocalToolScheduler from private_gpt.events.models import ( RawContentBlockDeltaEvent, RawContentBlockStartEvent, @@ -34,6 +49,90 @@ async def _noop_tool(value: str) -> str: return f"ok:{value}" +async def _collect_events(events: AsyncGenerator[Any, None]) -> list[Any]: + return [event async for event in events] + + +class _LocalTestRunner: + def __init__(self, engine: AsyncChatEngine) -> None: + self._engine = engine + self._tasks: dict[str, asyncio.Task[Any]] = {} + + async def submit( + self, + *, + request_data: dict[str, Any], + stream_type: str, + metadata: dict[str, Any], + execution_id: str | None = None, + ) -> tuple[str, AsyncGenerator[Any, None]]: + del stream_type, metadata + correlation_id = execution_id or "test-execution" + channel = LocalEventChannel() + + async def execute() -> None: + try: + request = ResolvedChatRequest.model_validate(request_data) + await self._engine.execute(request=request, channel=channel) + finally: + await channel.close() + + task = asyncio.create_task(execute()) + self._tasks[correlation_id] = task + return correlation_id, channel.stream(task) + + async def cancel(self, execution_id: str) -> bool: + task = self._tasks.get(execution_id) + if task is None: + return False + task.cancel() + return True + + +async def _run_engine( + engine: ChatEngine, + request: ResolvedChatRequest, + runner: ChatRunner | None, +) -> list[Any]: + execution = await engine.run(request=request, runner=runner) + events = await _collect_events(execution.events) + if execution.final_state_task is not None: + await execution.final_state_task + return events + + +def _build_engine( + engine_cls: Any, + engine_kwargs: dict[str, Any], + llm_component: LLMComponent, + max_iterations: int, +) -> tuple[ChatEngine, ChatRunner | None]: + engine = engine_cls( + llm_component=llm_component, + request_interceptors=[], + response_interceptors=[], + max_iterations=max_iterations, + **engine_kwargs, + ) + if isinstance(engine, AsyncChatEngine): + runner = _LocalTestRunner(engine) + return engine, runner + return LoopChatEngineAdapter(engine=engine), None + + +ENGINE_CONFIGS = [ + pytest.param( + AsyncChatEngine, + { + "tool_scheduler": LocalToolScheduler(), + "chat_scheduler": LocalChatScheduler(), + }, + id="async", + ), + pytest.param(ChatLoopEngine, {}, id="sync"), +] + + @pytest.fixture def base_request() -> ResolvedChatRequest: return ResolvedChatRequest( @@ -43,26 +142,36 @@ def base_request() -> ResolvedChatRequest: @pytest.mark.asyncio -async def test_loop_emits_text_and_stop(base_request: ResolvedChatRequest) -> None: +@pytest.mark.parametrize(("engine_cls", "engine_kwargs"), ENGINE_CONFIGS) +async def test_loop_emits_text_and_stop( + base_request: ResolvedChatRequest, engine_cls: Any, engine_kwargs: dict +) -> None: mock_llm = get_mock_function_calling_llm(["hello", " world"]) llm_component = MagicMock(spec=LLMComponent) llm_component.get_llm.return_value = mock_llm - engine = ChatLoopEngine( + engine, runner = _build_engine( + engine_cls=engine_cls, + engine_kwargs=engine_kwargs, llm_component=llm_component, - request_interceptors=[], - response_interceptors=[], max_iterations=2, ) - events = [event async for event in engine.run(base_request)] + events = await _run_engine( + engine=engine, + request=base_request, + runner=runner, + ) assert any(isinstance(event, RawContentBlockDeltaEvent) for event in events) assert any(isinstance(event, RawMessageStopEvent) for event in events) @pytest.mark.asyncio +@pytest.mark.parametrize(("engine_cls", "engine_kwargs"), ENGINE_CONFIGS) async def test_loop_streams_tool_use_and_tool_result( base_request: ResolvedChatRequest, + engine_cls: Any, + engine_kwargs: dict, ) -> None: request = base_request.model_copy(deep=True) request.tool_config = ResolvedToolConfig( @@ -91,14 +200,18 @@ async def test_loop_streams_tool_use_and_tool_result( llm_component = MagicMock(spec=LLMComponent) llm_component.get_llm.return_value = mock_llm - engine = ChatLoopEngine( + engine, runner = _build_engine( + engine_cls=engine_cls, + engine_kwargs=engine_kwargs, llm_component=llm_component, - request_interceptors=[], - response_interceptors=[], max_iterations=4, ) - events = [event async for event in engine.run(request)] + events = await _run_engine( + engine=engine, + request=request, + runner=runner, + ) assert any( isinstance(event, RawContentBlockStartEvent) and isinstance(event.content_block, ToolUseBlock) @@ -112,7 +225,10 @@ async def test_loop_streams_tool_use_and_tool_result( @pytest.mark.asyncio -async def test_loop_streams_reasoning_blocks(base_request: ResolvedChatRequest) -> None: +@pytest.mark.parametrize(("engine_cls", "engine_kwargs"), ENGINE_CONFIGS) +async def test_loop_streams_reasoning_blocks( + base_request: ResolvedChatRequest, engine_cls: Any, engine_kwargs: dict +) -> None: mock_llm = MagicMock(spec=FunctionCallingLLM) mock_llm.metadata.context_window = 4096 mock_llm.metadata.num_output = 1024 @@ -165,14 +281,18 @@ async def test_loop_streams_reasoning_blocks(base_request: ResolvedChatRequest) llm_component = MagicMock(spec=LLMComponent) llm_component.get_llm.return_value = mock_llm - engine = ChatLoopEngine( + engine, runner = _build_engine( + engine_cls=engine_cls, + engine_kwargs=engine_kwargs, llm_component=llm_component, - request_interceptors=[], - response_interceptors=[], max_iterations=2, ) - events = [event async for event in engine.run(base_request)] + events = await _run_engine( + engine=engine, + request=base_request, + runner=runner, + ) assert any( isinstance(event, RawContentBlockStartEvent) and isinstance(event.content_block, ThinkingBlock) @@ -181,8 +301,11 @@ async def test_loop_streams_reasoning_blocks(base_request: ResolvedChatRequest) @pytest.mark.asyncio +@pytest.mark.parametrize(("engine_cls", "engine_kwargs"), ENGINE_CONFIGS) async def test_loop_accumulates_usage_across_iterations( base_request: ResolvedChatRequest, + engine_cls: Any, + engine_kwargs: dict, ) -> None: request = base_request.model_copy(deep=True) request.tool_config = ResolvedToolConfig( @@ -269,14 +392,18 @@ async def test_loop_accumulates_usage_across_iterations( llm_component = MagicMock(spec=LLMComponent) llm_component.get_llm.return_value = mock_llm - engine = ChatLoopEngine( + engine, runner = _build_engine( + engine_cls=engine_cls, + engine_kwargs=engine_kwargs, llm_component=llm_component, - request_interceptors=[], - response_interceptors=[], max_iterations=4, ) - events = [event async for event in engine.run(request)] + events = await _run_engine( + engine=engine, + request=request, + runner=runner, + ) message_deltas = [ event for event in events if isinstance(event, RawMessageDeltaEvent) ] @@ -287,8 +414,11 @@ async def test_loop_accumulates_usage_across_iterations( @pytest.mark.asyncio +@pytest.mark.parametrize(("engine_cls", "engine_kwargs"), ENGINE_CONFIGS) async def test_loop_preserves_tool_calls_when_last_chunk_has_empty_tool_calls( base_request: ResolvedChatRequest, + engine_cls: Any, + engine_kwargs: dict, ) -> None: request = base_request.model_copy(deep=True) request.tool_config = ResolvedToolConfig( @@ -362,16 +492,107 @@ async def test_loop_preserves_tool_calls_when_last_chunk_has_empty_tool_calls( llm_component = MagicMock(spec=LLMComponent) llm_component.get_llm.return_value = mock_llm - engine = ChatLoopEngine( + engine, runner = _build_engine( + engine_cls=engine_cls, + engine_kwargs=engine_kwargs, llm_component=llm_component, - request_interceptors=[], - response_interceptors=[], max_iterations=2, ) - events = [event async for event in engine.run(request)] + events = await _run_engine( + engine=engine, + request=request, + runner=runner, + ) assert any( isinstance(event, RawContentBlockStartEvent) and isinstance(event.content_block, ToolUseBlock) for event in events ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("engine_cls", "engine_kwargs"), ENGINE_CONFIGS) +async def test_handle_stream_chunk_accumulates_token_ids_delta( + base_request: ResolvedChatRequest, + engine_cls: Any, + engine_kwargs: dict, +) -> None: + mock_llm = MagicMock(spec=FunctionCallingLLM) + mock_llm.metadata.context_window = 4096 + mock_llm.metadata.num_output = 1024 + mock_llm.metadata.is_function_calling_model = True + mock_llm.callback_manager = MagicMock() + mock_llm.completion_to_prompt = lambda prompt, **kwargs: prompt + mock_llm.messages_to_prompt = lambda messages, **kwargs: "\n".join( + [message.content for message in messages or [] if message and message.content] + ) + mock_llm.get_tool_calls_from_response = lambda *args, **kwargs: [] + + llm_component = MagicMock(spec=LLMComponent) + llm_component.get_llm.return_value = mock_llm + + engine = engine_cls( + llm_component=llm_component, + request_interceptors=[], + response_interceptors=[], + max_iterations=2, + **engine_kwargs, + ) + + run = engine.initialize_run(base_request) + current_response = ChatResponse( + message=ChatMessage( + role=MessageRole.ASSISTANT, + content=None, + additional_kwargs={}, + ), + additional_kwargs={}, + ) + handler = _EventHandler(queue=asyncio.Queue()) + stream_delta_state = _StreamDeltaState() + lock = asyncio.Lock() + + first = ChatResponse( + message=ChatMessage( + role=MessageRole.ASSISTANT, + content="he", + additional_kwargs={"token_ids_delta": [11, 12]}, + ), + delta="he", + additional_kwargs={"token_ids_delta": [11, 12]}, + ) + second = ChatResponse( + message=ChatMessage( + role=MessageRole.ASSISTANT, + content="llo", + additional_kwargs={"token_ids_delta": [13]}, + ), + delta="llo", + additional_kwargs={"token_ids_delta": [13]}, + ) + + current_response = await engine._handle_stream_chunk( + run=run, + llm=mock_llm, + chunk=first, + current_response=current_response, + stream_delta_state=stream_delta_state, + handler=handler, + tool_specs_by_name={}, + schema_by_name={}, + lock=lock, + ) + current_response = await engine._handle_stream_chunk( + run=run, + llm=mock_llm, + chunk=second, + current_response=current_response, + stream_delta_state=stream_delta_state, + handler=handler, + tool_specs_by_name={}, + schema_by_name={}, + lock=lock, + ) + + assert current_response.message.additional_kwargs["token_ids_delta"] == [11, 12, 13] diff --git a/tests/engines/test_resumable_runtime.py b/tests/engines/test_resumable_runtime.py new file mode 100644 index 00000000..e9b95e37 --- /dev/null +++ b/tests/engines/test_resumable_runtime.py @@ -0,0 +1,410 @@ +import asyncio + +import pytest + +from private_gpt.components.engines.chat.async_chat_engine import ( + IterationCheckpointPayload, +) +from private_gpt.components.engines.chat.checkpoint_store import ( + ChatCheckpoint, + InMemoryChatCheckpointStore, +) +from private_gpt.components.engines.chat.event_broker import ( + InMemoryEngineEventBroker, +) +from private_gpt.components.engines.chat.event_channel import BrokerEventChannel +from private_gpt.components.tools.remote_execution import ToolExecutionResponse +from private_gpt.events.models import PingEvent, TextBlock + + +@pytest.mark.asyncio +async def test_memory_event_broker_preserves_publish_order_and_finishes() -> None: + broker = InMemoryEngineEventBroker() + channel = BrokerEventChannel(broker, "execution-1") + channel.emit(PingEvent()) + channel.emit(PingEvent()) + await channel.close() + await broker.finish("execution-1") + + events = [event async for event in broker.listen("execution-1")] + + assert [event.type for event in events] == ["ping", "ping"] + + +@pytest.mark.asyncio +async def test_memory_event_broker_waits_for_late_publication() -> None: + broker = InMemoryEngineEventBroker() + + async def collect(): + return [event async for event in broker.listen("execution-2")] + + collector = asyncio.create_task(collect()) + await asyncio.sleep(0) + await broker.publish("execution-2", PingEvent()) + await broker.finish("execution-2") + + assert [event.type for event in await collector] == ["ping"] + + +@pytest.mark.asyncio +async def test_memory_iteration_state_aggregates_once_and_cleans_up() -> None: + service = InMemoryChatCheckpointStore() + response = ToolExecutionResponse( + tool_name="echo", + tool_id="tool-1", + result_content=[TextBlock(text="ok")], + tool_message={ + "role": "tool", + "content": "ok", + "additional_kwargs": {"tool_call_id": "tool-1"}, + }, + ) + await service.save( + ChatCheckpoint( + correlation_id="execution-3", + request_data={}, + stream_type="chat_completion", + metadata={}, + iteration=0, + checkpoint="tools", + checkpoint_payload=IterationCheckpointPayload( + pending_async_tools={"tool-1": "job-1"} + ), + checkpoint_id="checkpoint-1", + ) + ) + + ready = await service.record_result( + "execution-3", "tool-1", response.model_dump(mode="json") + ) + + assert ready == {"tool-1": response} + assert await service.claim_resume("execution-3") is True + assert await service.claim_resume("execution-3") is False + await service.cleanup("execution-3") + assert await service.load("execution-3") is None + + +@pytest.mark.asyncio +async def test_memory_iteration_state_preserves_result_received_before_checkpoint() -> None: + service = InMemoryChatCheckpointStore() + response = ToolExecutionResponse( + tool_name="echo", + tool_id="tool-early", + result_content=[TextBlock(text="ok")], + tool_message={ + "role": "tool", + "content": "ok", + "additional_kwargs": {"tool_call_id": "tool-early"}, + }, + ) + + assert ( + await service.record_result( + "execution-early", "tool-early", response.model_dump(mode="json") + ) + is None + ) + await service.save( + ChatCheckpoint( + correlation_id="execution-early", + request_data={}, + stream_type="chat_completion", + metadata={}, + iteration=0, + checkpoint="tools", + checkpoint_payload=IterationCheckpointPayload( + pending_async_tools={"tool-early": "job-early"} + ), + ) + ) + + assert await service.get_results("execution-early") == {"tool-early": response} + + +@pytest.mark.asyncio +async def test_memory_iteration_state_ignores_unknown_tool_result() -> None: + service = InMemoryChatCheckpointStore() + await service.save( + ChatCheckpoint( + correlation_id="execution-unknown", + request_data={}, + stream_type="chat_completion", + metadata={}, + iteration=0, + checkpoint="tools", + checkpoint_payload=IterationCheckpointPayload( + pending_async_tools={"tool-expected": "job-expected"} + ), + ) + ) + response = ToolExecutionResponse( + tool_name="echo", + tool_id="tool-unknown", + result_content=[TextBlock(text="ok")], + tool_message={ + "role": "tool", + "content": "ok", + "additional_kwargs": {"tool_call_id": "tool-unknown"}, + }, + ) + + assert ( + await service.record_result( + "execution-unknown", "tool-unknown", response.model_dump(mode="json") + ) + is None + ) + assert await service.get_results("execution-unknown") == {} + + +@pytest.mark.asyncio +async def test_runner_cancel_revokes_pending_tool_tasks() -> None: + from unittest.mock import AsyncMock, MagicMock + + from private_gpt.components.engines.chat.resumable_runner import ResumableChatRunner + + checkpoint_store = InMemoryChatCheckpointStore() + await checkpoint_store.save( + ChatCheckpoint( + correlation_id="execution-cancel", + request_data={}, + stream_type="chat_completion", + metadata={}, + iteration=1, + checkpoint="tools", + checkpoint_payload=IterationCheckpointPayload( + pending_async_tools={ + "tool-1": "celery-task-1", + "tool-2": "celery-task-2", + } + ), + ) + ) + event_broker = InMemoryEngineEventBroker() + chat_scheduler = MagicMock() + chat_scheduler.cancel = AsyncMock(return_value=False) + tool_scheduler = MagicMock() + tool_scheduler.cancel_task = AsyncMock(return_value=True) + + checkpoint_factory = MagicMock() + checkpoint_factory.get.return_value = checkpoint_store + event_factory = MagicMock() + event_factory.get.return_value = event_broker + scheduler_factory = MagicMock() + scheduler_factory.get.return_value = chat_scheduler + tool_scheduler_factory = MagicMock() + tool_scheduler_factory.get.return_value = tool_scheduler + + runner = ResumableChatRunner( + settings=MagicMock(), + checkpoint_store_factory=checkpoint_factory, + event_broker_factory=event_factory, + scheduler_factory=scheduler_factory, + tool_scheduler_factory=tool_scheduler_factory, + ) + + cancelled = await runner.cancel(execution_id="execution-cancel") + + assert cancelled is True + assert tool_scheduler.cancel_task.await_count == 2 + tool_scheduler.cancel_task.assert_any_await(task_id="celery-task-1") + tool_scheduler.cancel_task.assert_any_await(task_id="celery-task-2") + chat_scheduler.cancel.assert_awaited_once_with("execution-cancel") + assert await checkpoint_store.load("execution-cancel") is None + + +def test_runner_restores_additional_kwargs_after_serialization() -> None: + from llama_index.core.base.llms.types import ChatMessage, MessageRole + from llama_index.core.tools import ToolSelection + + from private_gpt.components.chat.models.chat_config_models import ( + ResolvedChatRequest, + ) + from private_gpt.components.engines.chat.resumable_runner import ( + ResumableChatRunner, + ) + from private_gpt.events.models import DocumentBlock, ThinkingBlock + + document = DocumentBlock( + source=DocumentBlock.PlainTextSource( + data="Document text", + media_type="text/plain", + ), + title="Test document", + ) + thinking = ThinkingBlock(thinking="Reasoning", signature="signature") + tool_call = ToolSelection( + tool_id="tool-1", + tool_name="lookup", + tool_kwargs={"query": "test"}, + ) + request = ResolvedChatRequest( + messages=[ + ChatMessage( + role=MessageRole.USER, + content="Summarize this document", + additional_kwargs={ + "document": [document], + "thinking": [thinking], + "tool_calls": [tool_call], + }, + ) + ] + ) + + restored = ResumableChatRunner._request(request.model_dump(mode="json")) + additional_kwargs = restored.messages[0].additional_kwargs + + assert additional_kwargs["document"] == [document] + assert additional_kwargs["thinking"] == [thinking] + assert additional_kwargs["tool_calls"] == [tool_call] + assert isinstance(additional_kwargs["document"][0], DocumentBlock) + assert isinstance(additional_kwargs["thinking"][0], ThinkingBlock) + 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 + + from llama_index.core.base.llms.types import ChatMessage, MessageRole + from llama_index.core.tools import ToolSelection + + from private_gpt.chat.schema_models import create_model_from_json_schema + from private_gpt.components.chat.models.chat_config_models import ( + CitationConfig, + CondensationConfig, + ResolvedChatRequest, + ResolvedContextConfig, + ResolvedSystemConfig, + ResolvedToolConfig, + ResponseFormatConfig, + ThinkingConfig, + ToolSpec, + ) + from private_gpt.components.engines.chat.resumable_runner import ( + ResumableChatRunner, + ) + from private_gpt.events.models import DocumentBlock, ThinkingBlock + + document = DocumentBlock( + source=DocumentBlock.PlainTextSource( + data="Document text", + media_type="text/plain", + ), + title="Test document", + ) + thinking = ThinkingBlock(thinking="Reasoning", signature="signature") + tool_call = ToolSelection( + tool_id="tool-1", + tool_name="lookup", + tool_kwargs={"query": "test"}, + ) + output_schema = { + "title": "Profile", + "type": "object", + "properties": {"name": {"type": "string"}}, + "required": ["name"], + } + output_cls = create_model_from_json_schema(output_schema) + request = ResolvedChatRequest( + stream=True, + messages=[ + ChatMessage( + role=MessageRole.USER, + content="Create a profile", + additional_kwargs={ + "document": [document], + "thinking": [thinking], + "tool_calls": [tool_call], + }, + ) + ], + system=ResolvedSystemConfig(prompt="System prompt", model="default"), + tool_config=ResolvedToolConfig( + tools=[ + ToolSpec.from_defaults( + name="lookup", + type="lookup", + runtime="server", + input_schema={ + "type": "object", + "properties": {"query": {"type": "string"}}, + }, + ) + ], + tool_choices=["lookup"], + allow_parallel_tool_calls=False, + ), + context=ResolvedContextConfig( + add_context_to_system_prompt=True, + deduplicate_context_in_history=True, + maximum_context_length=2048, + correlation_id="correlation-1", + user_id="user-1", + container="container-1", + maximum_loaded_skills=2, + ), + condensation=CondensationConfig(enabled=False), + citation=CitationConfig( + enabled=True, + force_to_return_citations=True, + return_missing_citations=True, + ), + thinking=ThinkingConfig(enabled=True, type="high"), + response_format=ResponseFormatConfig(output_cls=output_cls), + sampling_params={ + "temperature": 0.2, + "max_tokens": 512, + "stop": ["done"], + }, + ) + + serialized = request.model_dump(mode="json") + encoded = json.dumps(serialized) + restored = ResumableChatRunner._request(json.loads(encoded)) + + assert set(serialized) == set(ResolvedChatRequest.model_fields) + assert restored.model_dump(mode="json") == serialized diff --git a/tests/fixtures/mock_injector.py b/tests/fixtures/mock_injector.py index c73d2b0d..eb916966 100644 --- a/tests/fixtures/mock_injector.py +++ b/tests/fixtures/mock_injector.py @@ -3,7 +3,7 @@ from typing import Any from unittest.mock import MagicMock import pytest -from injector import Provider, ScopeDecorator, singleton +from injector import Provider, ScopeDecorator, SingletonScope, singleton from private_gpt.di import create_application_injector from private_gpt.settings.settings import Settings, unsafe_settings @@ -25,6 +25,7 @@ class MockInjector: if mock is None: mock = MagicMock() self.test_injector.binder.bind(interface, to=mock, scope=scope) + self.test_injector.get(SingletonScope)._context.pop(interface, None) return mock # type: ignore def bind_settings(self, settings: dict[str, Any]) -> Settings: diff --git a/tests/server/chat/anthropic/test_anthropic_client.py b/tests/server/chat/anthropic/test_anthropic_client.py index e1170a41..5e59a52b 100644 --- a/tests/server/chat/anthropic/test_anthropic_client.py +++ b/tests/server/chat/anthropic/test_anthropic_client.py @@ -119,8 +119,8 @@ def setup_mock_llm( ) llm_component = injector.get(LLMComponent) - llm_component.llm = mock_llm_instance - injector.bind_mock(LLMComponent, mock_llm_instance) + llm_component.get_llm = Mock(return_value=mock_llm_instance) + injector.bind_mock(LLMComponent, llm_component) def create_mock_http_client( @@ -692,6 +692,7 @@ def test_messages_count_tokens_endpoint_forwards_message_input_and_options( test_client: TestClient, httpx_mock: HTTPXMock, ) -> None: + injector.get(ChatService) chat_service = Mock() chat_service.count_tokens = AsyncMock( return_value=CountTokensOutput(input_tokens=11) diff --git a/tests/server/chat/anthropic/test_langchain_anthropic.py b/tests/server/chat/anthropic/test_langchain_anthropic.py index 94a2930d..0117992a 100644 --- a/tests/server/chat/anthropic/test_langchain_anthropic.py +++ b/tests/server/chat/anthropic/test_langchain_anthropic.py @@ -1,6 +1,7 @@ import uuid from collections.abc import Iterator from typing import Any, get_args +from unittest.mock import Mock import anthropic import httpx @@ -139,8 +140,8 @@ def setup_mock_llm( ) llm_component = injector.get(LLMComponent) - llm_component.llm = mock_llm_instance - injector.bind_mock(LLMComponent, mock_llm_instance) + llm_component.get_llm = Mock(return_value=mock_llm_instance) + injector.bind_mock(LLMComponent, llm_component) def create_mock_http_client( diff --git a/tests/server/chat/interceptors/test_validation_request_interceptor.py b/tests/server/chat/interceptors/test_validation_request_interceptor.py index dc2c73ec..6a4a64ae 100644 --- a/tests/server/chat/interceptors/test_validation_request_interceptor.py +++ b/tests/server/chat/interceptors/test_validation_request_interceptor.py @@ -18,19 +18,19 @@ from private_gpt.components.chat.models.chat_config_models import ( ResolvedSystemConfig, ToolSpec, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_interceptor_context import ( - ChatLoopInterceptorContext, +from private_gpt.components.engines.chat.models.chat_interceptor_context import ( + ChatInterceptorContext, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_phase import ( +from private_gpt.components.engines.chat.models.chat_phase import ( InterceptorPhase, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_state import ( - ChatLoopInputState, - ChatLoopOutputState, - ChatLoopRuntimeState, - ChatLoopState, +from private_gpt.components.engines.chat.models.chat_state import ( + ChatInputState, + ChatOutputState, + ChatRuntimeState, + ChatState, ) -from private_gpt.components.engines.chat_loop.utils.request_builder import ( +from private_gpt.components.engines.chat.utils.request_builder import ( build_initial_context_stack, ) from private_gpt.events.event_errors import Errors @@ -124,18 +124,18 @@ async def _run_interceptor( interceptor: ValidatorRequestInterceptor, request: ChatRequest, ) -> None: - state = ChatLoopState( - input=ChatLoopInputState( + state = ChatState( + input=ChatInputState( request=request, context_stack=build_initial_context_stack(request), sampling_params=dict(request.sampling_params), llm_kwargs=dict(request.sampling_params), ), - runtime=ChatLoopRuntimeState(), - output=ChatLoopOutputState(), + runtime=ChatRuntimeState(), + output=ChatOutputState(), timeline=[], ) - context = ChatLoopInterceptorContext( + context = ChatInterceptorContext( state=state, llm=get_mock_function_calling_llm(["ok"]), phase=InterceptorPhase.VALIDATION, diff --git a/tests/server/chat/test_chat_async_worker.py b/tests/server/chat/test_chat_async_worker.py new file mode 100644 index 00000000..6f2c3c6b --- /dev/null +++ b/tests/server/chat/test_chat_async_worker.py @@ -0,0 +1,90 @@ +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from private_gpt.server.chat_async.chat_async_service import ChatAsyncService + + +def _service( + stream_manager: MagicMock | None = None, + chat_facade: MagicMock | None = None, +) -> ChatAsyncService: + if stream_manager is None: + stream_manager = MagicMock() + stream_manager.stream_exists = AsyncMock(return_value=False) + stream_manager.cancel_stream = AsyncMock(return_value=False) + if chat_facade is None: + chat_facade = MagicMock() + + return ChatAsyncService(stream_manager=stream_manager, chat_facade=chat_facade) + + +@pytest.mark.anyio +async def test_initiate_chat_stream_starts_engine_event_processing() -> None: + stream_manager = MagicMock() + stream_manager.stream_exists = AsyncMock(return_value=False) + stream_manager.create_and_start_stream = AsyncMock(return_value="msg-1") + chat_facade = MagicMock() + + async def events(): + if False: + yield None + + event_generator = events() + chat_facade.create_chat_event_generator = AsyncMock(return_value=event_generator) + + service = _service(stream_manager=stream_manager, chat_facade=chat_facade) + request = MagicMock() + request.messages = [] + request.thinking.enabled = False + updated_request = MagicMock() + updated_request.messages = [] + updated_request.thinking.enabled = False + request.model_copy.return_value = updated_request + correlation_id = await service.initiate_chat_stream( + request=request, + message_id="msg-1", + ) + + assert correlation_id == "msg-1" + request.context.model_copy.assert_called_once_with( + update={"correlation_id": "msg-1"} + ) + chat_facade.create_chat_event_generator.assert_awaited_once_with( + request=updated_request + ) + stream_manager.create_and_start_stream.assert_awaited_once_with( + event_handler=stream_manager.create_and_start_stream.await_args.kwargs[ + "event_handler" + ], + stream_type="chat_completion", + event_generator=event_generator, + correlation_id="msg-1", + metadata={"message_count": 0, "thinking_enabled": False}, + ) + + +@pytest.mark.anyio +async def test_initiate_chat_stream_rejects_duplicate_message_id() -> None: + stream_manager = MagicMock() + stream_manager.stream_exists = AsyncMock(return_value=True) + + service = _service(stream_manager=stream_manager) + + with pytest.raises(ValueError, match="already exists"): + await service.initiate_chat_stream(request=MagicMock(), message_id="msg-2") + + +@pytest.mark.anyio +async def test_cancel_only_delegates_to_stream_manager() -> None: + stream_manager = MagicMock() + stream_manager.cancel_stream = AsyncMock(return_value=False) + chat_facade = MagicMock() + chat_facade.cancel = AsyncMock(return_value=True) + + service = _service(stream_manager=stream_manager, chat_facade=chat_facade) + + await service.cancel_stream("msg-3") + + stream_manager.cancel_stream.assert_awaited_once_with("msg-3") + chat_facade.cancel.assert_not_awaited() diff --git a/tests/server/chat/test_chat_facade.py b/tests/server/chat/test_chat_facade.py new file mode 100644 index 00000000..de2fc6bf --- /dev/null +++ b/tests/server/chat/test_chat_facade.py @@ -0,0 +1,35 @@ +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from private_gpt.events.models import PingEvent +from private_gpt.server.chat.chat_facade import ChatFacadeService + + +@pytest.mark.asyncio +async def test_event_generator_close_propagates_to_engine_generator() -> None: + closed = asyncio.Event() + + async def engine_events(): + try: + yield PingEvent() + await asyncio.Event().wait() + finally: + closed.set() + + chat_service = MagicMock() + chat_service.stream_chat = AsyncMock( + return_value=SimpleNamespace(events=engine_events()) + ) + facade = ChatFacadeService( + chat_service=chat_service, + stream_manager=MagicMock(), + ) + event_generator = await facade.create_chat_event_generator(request=MagicMock()) + + await anext(event_generator) + await event_generator.aclose() + + assert closed.is_set() diff --git a/tests/server/chat/test_chat_knowledge_revamp.py b/tests/server/chat/test_chat_knowledge_revamp.py index b66c7510..5d90b099 100644 --- a/tests/server/chat/test_chat_knowledge_revamp.py +++ b/tests/server/chat/test_chat_knowledge_revamp.py @@ -1,5 +1,6 @@ import uuid from typing import Any +from unittest.mock import Mock import pytest from httpx import AsyncClient @@ -72,8 +73,8 @@ async def mock_llm_with_capture( mock_llm_instance.astream_chat_with_tools = coro llm_component = injector.get(LLMComponent) - llm_component.llm = mock_llm_instance - injector.bind_mock(LLMComponent, mock_llm_instance) + llm_component.get_llm = Mock(return_value=mock_llm_instance) + injector.bind_mock(LLMComponent, llm_component) def create_tool_definition( diff --git a/tests/server/chat/test_chat_routes.py b/tests/server/chat/test_chat_routes.py index 3ec64ef6..b72e5a58 100644 --- a/tests/server/chat/test_chat_routes.py +++ b/tests/server/chat/test_chat_routes.py @@ -3,7 +3,7 @@ import json import uuid from collections.abc import AsyncGenerator from typing import Any -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, Mock import pytest from httpx import AsyncClient @@ -107,8 +107,8 @@ async def mock_llm( mock_llm = get_mock_function_calling_llm(deltas) llm_component = injector.get(LLMComponent) - llm_component.llm = mock_llm - injector.bind_mock(LLMComponent, mock_llm) + llm_component.get_llm = Mock(return_value=mock_llm) + injector.bind_mock(LLMComponent, llm_component) @pytest.mark.anyio @@ -2181,67 +2181,6 @@ async def test_chat_accepts_empty_system_list( assert isinstance(output.content[0], TextBlock) -@pytest.mark.anyio -async def test_chat_cancels_llm_astream_on_client_disconnection( - async_test_client: AsyncClient, injector: MockInjector -) -> None: - llm_started = asyncio.Event() - llm_generator_closed = asyncio.Event() - - class SlowStreamingLLM(FunctionCallingLLMMock): - async def astream_chat_with_tools( - self, *args: Any, **kwargs: Any - ) -> AsyncGenerator[ChatResponse, None]: - async def _gen() -> AsyncGenerator[ChatResponse, None]: - try: - llm_started.set() - yield ChatResponse( - message=ChatMessage(role=MessageRole.ASSISTANT, content=""), - delta="Hi", - ) - await asyncio.sleep(30) - yield ChatResponse( - message=ChatMessage(role=MessageRole.ASSISTANT, content=""), - delta=" never", - ) - finally: - llm_generator_closed.set() - - return _gen() - - llm_component = injector.get(LLMComponent) - llm_component.llm = SlowStreamingLLM() - - body = ChatBody( - messages=[MessageInput(content="What is Python?", role="user")], - stream=False, - ) - - request_task = asyncio.create_task( - async_test_client.post("/v1/messages", json=body.model_dump()) - ) - - try: - await asyncio.wait_for(llm_started.wait(), timeout=5.0) - except TimeoutError: - request_task.cancel() - pytest.fail("LLM astream_chat_with_tools never started") - - request_task.cancel() - - with pytest.raises(asyncio.CancelledError): - await request_task - - try: - await asyncio.wait_for(llm_generator_closed.wait(), timeout=2.0) - except TimeoutError: - pytest.fail( - "LLM astream_chat_with_tools generator was not closed after client disconnection" - ) - - assert llm_generator_closed.is_set() - - @pytest.mark.anyio async def test_concurrent_requests_dont_share_state( async_test_client: AsyncClient, injector: MockInjector diff --git a/tests/server/chat/test_chat_routes_skills_integration.py b/tests/server/chat/test_chat_routes_skills_integration.py index 4e58d3ce..488272d2 100644 --- a/tests/server/chat/test_chat_routes_skills_integration.py +++ b/tests/server/chat/test_chat_routes_skills_integration.py @@ -1,6 +1,7 @@ import json import uuid from typing import Any +from unittest.mock import Mock import pytest from httpx import AsyncClient @@ -72,8 +73,8 @@ async def mock_llm_with_capture( mock_llm.astream_chat_with_tools = coro llm_component = injector.get(LLMComponent) - llm_component.llm = mock_llm - injector.bind_mock(LLMComponent, mock_llm) + llm_component.get_llm = Mock(return_value=mock_llm) + injector.bind_mock(LLMComponent, llm_component) def _skill_md(name: str, description: str, body: str) -> bytes: diff --git a/tests/server/chat/test_chat_service.py b/tests/server/chat/test_chat_service.py index e714f8a3..3efb4ee5 100644 --- a/tests/server/chat/test_chat_service.py +++ b/tests/server/chat/test_chat_service.py @@ -1,5 +1,5 @@ from collections.abc import AsyncGenerator -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock import pytest from llama_index.core.base.llms.types import MessageRole @@ -9,11 +9,14 @@ from private_gpt.components.chat.models.chat_config_models import ( ResolvedChatRequest, ResolvedSystemConfig, ) -from private_gpt.components.engines.chat_loop.interceptors.chat_loop_interceptor import ( +from private_gpt.components.engines.chat.chat_engine_interface import ( + ChatEngineExecution, +) +from private_gpt.components.engines.chat.interceptors.chat_interceptor import ( ChatRequestLoopInterceptor, ) -from private_gpt.components.engines.chat_loop.models.chat_loop_interceptor_context import ( - ChatLoopInterceptorContext, +from private_gpt.components.engines.chat.models.chat_interceptor_context import ( + ChatInterceptorContext, ) from private_gpt.events.models import ( Event, @@ -36,7 +39,7 @@ class _PromptInterceptorRequest(ChatRequestLoopInterceptor): def __init__(self, label: str) -> None: self._label = label - async def intercept(self, context: ChatLoopInterceptorContext) -> None: + async def intercept(self, context: ChatInterceptorContext) -> None: state = context.state.model_copy(deep=True) current = state.input.request.system.prompt or "" state.input.request.system.prompt = f"{current}{self._label}" @@ -66,7 +69,11 @@ async def _fatal_event_stream() -> AsyncGenerator[Event, None]: def _mock_engine_for(stream: AsyncGenerator[Event, None]) -> MagicMock: mock_engine = MagicMock() - mock_engine.run.return_value = stream + execution = ChatEngineExecution( + events=stream, + final_state_task=MagicMock(), + ) + mock_engine.run = AsyncMock(return_value=execution) return mock_engine @@ -74,7 +81,7 @@ def _mock_engine_for(stream: AsyncGenerator[Event, None]) -> MagicMock: async def test_chat_folds_loop_events(injector: MockInjector) -> None: service: ChatService = injector.get(ChatService) mock_engine = _mock_engine_for(_basic_event_stream()) - service._build_loop_engine = MagicMock(return_value=mock_engine) + service.build_engine = MagicMock(return_value=mock_engine) request = ResolvedChatRequest( messages=[ChatMessage(content="hi", role=MessageRole.USER)], @@ -87,8 +94,8 @@ async def test_chat_folds_loop_events(injector: MockInjector) -> None: assert result.response == "hello" assert result.stop_reason == "end_turn" assert result.usage is not None - assert result.usage["input_tokens"] == 10 - assert result.usage["output_tokens"] == 20 + assert result.usage.input_tokens == 10 + assert result.usage.output_tokens == 20 mock_engine.run.assert_called_once() @@ -96,7 +103,7 @@ async def test_chat_folds_loop_events(injector: MockInjector) -> None: async def test_stream_chat_returns_loop_generator(injector: MockInjector) -> None: service: ChatService = injector.get(ChatService) mock_engine = _mock_engine_for(_basic_event_stream()) - service._build_loop_engine = MagicMock(return_value=mock_engine) + service.build_engine = MagicMock(return_value=mock_engine) request = ResolvedChatRequest( messages=[ChatMessage(content="hi", role=MessageRole.USER)], @@ -138,7 +145,7 @@ async def test_chat_propagates_fatal_error_to_completion( ) -> None: service: ChatService = injector.get(ChatService) mock_engine = _mock_engine_for(_fatal_event_stream()) - service._build_loop_engine = MagicMock(return_value=mock_engine) + service.build_engine = MagicMock(return_value=mock_engine) request = ResolvedChatRequest( messages=[ChatMessage(content="hi", role=MessageRole.USER)], diff --git a/tests/server/ingest/test_ingest_routes.py b/tests/server/ingest/test_ingest_routes.py index d94d5758..711890c4 100644 --- a/tests/server/ingest/test_ingest_routes.py +++ b/tests/server/ingest/test_ingest_routes.py @@ -7,6 +7,9 @@ from unittest.mock import Mock from fastapi.testclient import TestClient from private_gpt.components.broker.broker_component import BrokerComponent +from private_gpt.components.ingestion.ingestion_scheduler import ( + IngestionSchedulerFactory, +) from private_gpt.components.storage.s3_helper import S3Helper from private_gpt.di import set_global_injector from private_gpt.server.ingest.ingest_router import ( @@ -22,6 +25,15 @@ from tests.fixtures.ingest_helper import IngestHelper from tests.fixtures.mock_injector import MockInjector +def _use_celery_ingestion(injector: MockInjector) -> None: + settings = injector.bind_settings({"scheduler": {"ingestion": {"mode": "celery"}}}) + factory = IngestionSchedulerFactory( + settings=settings, + injector=injector.test_injector, + ) + injector.bind_mock(IngestionSchedulerFactory, factory) + + def test_ingest_accepts_txt_files( test_client: TestClient, ingest_helper: IngestHelper ) -> None: @@ -103,6 +115,7 @@ def test_ingest_uri_async( test_client: TestClient, injector: MockInjector, ingest_helper: IngestHelper ) -> None: collection = str(uuid.uuid4()) + _use_celery_ingestion(injector=injector) # Mock broker to receive callback broker_mock = Mock(BrokerComponent) @@ -285,6 +298,7 @@ def test_delete_async( test_client: TestClient, injector: MockInjector, ingest_helper: IngestHelper ) -> None: collection = str(uuid.uuid4()) + _use_celery_ingestion(injector=injector) # Mock broker to receive callback broker_mock = Mock(BrokerComponent) injector.bind_mock(BrokerComponent, broker_mock) @@ -303,7 +317,6 @@ def test_delete_async( "collection": collection, }, ) - body = DeleteIngestedDocumentAsyncBody( delete_body=DeleteIngestedDocumentBody( collection=collection, diff --git a/tests/server/ingest/test_ingestion_async.py b/tests/server/ingest/test_ingestion_async.py index 1cc31d2d..3c6b67bb 100644 --- a/tests/server/ingest/test_ingestion_async.py +++ b/tests/server/ingest/test_ingestion_async.py @@ -79,7 +79,7 @@ def test_delete_during_ingestion(mock_setup, test_bodies): "worker1": [ { "id": "task1", - "name": "vector_index_task", + "name": "private_gpt.ingestion.vector_index", "args": [ingestion_body], } ] @@ -146,7 +146,7 @@ def test_delete_different_artifact_ingesting(mock_setup, test_bodies): "worker1": [ { "id": "task1", - "name": "vector_index_task", + "name": "private_gpt.ingestion.vector_index", "args": [different_ingestion_body], } ] @@ -195,17 +195,17 @@ def test_delete_with_multiple_ingestion_tasks(mock_setup, test_bodies): "worker1": [ { "id": "task1", - "name": "vector_index_task", + "name": "private_gpt.ingestion.vector_index", "args": [different_artifact], }, { "id": "task2", - "name": "vector_index_task", + "name": "private_gpt.ingestion.vector_index", "args": [same_artifact_diff_collection], }, { "id": "task3", - "name": "vector_index_task", + "name": "private_gpt.ingestion.vector_index", "args": [ingestion_body], }, {"id": "task4", "name": "different_task", "args": [{}]}, @@ -234,7 +234,7 @@ def test_delete_terminates_pending_tasks(mock_setup, test_bodies): pending_task = { "id": "task1", - "name": "vector_index_task", + "name": "private_gpt.ingestion.vector_index", "args": [ingestion_body], "status": "PENDING", } @@ -270,7 +270,7 @@ def test_delete_scheduled_when_ingestion_will_run(mock_setup, test_bodies): "worker1": [ { "id": "task1", - "name": "delete_ingested_task", + "name": "private_gpt.ingestion.delete", "args": [delete_body], } ] diff --git a/tests/settings/test_settings.py b/tests/settings/test_settings.py index f10bdd81..15b8f5a8 100644 --- a/tests/settings/test_settings.py +++ b/tests/settings/test_settings.py @@ -1,5 +1,8 @@ +import pytest + from private_gpt.settings.settings import Settings, settings -from tests.fixtures.mock_injector import MockInjector +from private_gpt.settings.settings_loader import merge_settings +from tests.fixtures.mock_injector import MockInjector, unsafe_settings def test_settings_are_loaded_and_merged() -> None: @@ -10,3 +13,63 @@ def test_settings_can_be_overriden(injector: MockInjector) -> None: injector.bind_settings({"server": {"env_name": "overriden"}}) mocked_settings = injector.get(Settings) assert mocked_settings.server.env_name == "overriden" + + +def test_chat_scheduler_rejects_unsupported_celery_mode() -> None: + merged = merge_settings( + [ + unsafe_settings, + { + "scheduler": {"chat": {"mode": "celery"}}, + "stream": {"broker": "redis"}, + "celery": {"use_workers": True}, + }, + ] + ) + + with pytest.raises(ValueError, match=r"Unsupported scheduler\.chat\.mode='celery'"): + Settings(**merged) + + +def test_arq_chat_scheduler_requires_redis_stream_backend() -> None: + merged = merge_settings( + [ + unsafe_settings, + { + "scheduler": {"chat": {"mode": "arq"}}, + "stream": {"broker": "memory"}, + }, + ] + ) + + with pytest.raises(ValueError, match=r"stream\.broker=redis"): + Settings(**merged) + + +def test_celery_tool_scheduler_requires_arq_chat() -> None: + merged = merge_settings( + [ + unsafe_settings, + { + "scheduler": { + "chat": {"mode": "local"}, + "tools": {"mode": "celery"}, + }, + }, + ] + ) + + with pytest.raises(ValueError, match=r"scheduler\.chat\.mode='arq'"): + Settings(**merged) + + +def test_tool_scheduler_rejects_unsupported_arq_mode() -> None: + merged = merge_settings( + [ + unsafe_settings, + {"scheduler": {"tools": {"mode": "arq"}}}, + ] + ) + + with pytest.raises(ValueError, match=r"Unsupported scheduler\.tools\.mode='arq'"): + Settings(**merged) diff --git a/tests/sse/test_sse.py b/tests/sse/test_sse.py index c0288fbd..0a3fdb7a 100644 --- a/tests/sse/test_sse.py +++ b/tests/sse/test_sse.py @@ -28,9 +28,9 @@ from private_gpt.components.chat.models.chat_config_models import ( ResolvedToolConfig, ToolSpec, ) -from private_gpt.components.engines.chat_loop.chat_loop_engine import ChatLoopEngine -from private_gpt.components.engines.chat_loop.models.chat_loop_interceptor_context import ( - ChatLoopInterceptorContext, +from private_gpt.components.engines.chat.chat_engine import ChatLoopEngine +from private_gpt.components.engines.chat.models.chat_interceptor_context import ( + ChatInterceptorContext, ) from private_gpt.components.llm.llm_component import LLMComponent from private_gpt.events.models import ( @@ -59,7 +59,7 @@ from private_gpt.server.chat.interceptors.ping_loop_interceptor import PingInter from tests.fixtures.mock_function_llm import get_mock_function_calling_llm if TYPE_CHECKING: - from private_gpt.components.engines.chat_loop.interceptors.chat_loop_interceptor import ( + from private_gpt.components.engines.chat.interceptors.chat_interceptor import ( ChatResponseLoopInterceptor, ) @@ -87,7 +87,7 @@ async def _collect_from_gen( queue: asyncio.Queue = asyncio.Queue() sentinel = object() - ctx = MagicMock(spec=ChatLoopInterceptorContext) + ctx = MagicMock(spec=ChatInterceptorContext) ctx.state = MagicMock() ctx.state.input = MagicMock() ctx.state.input.request = MagicMock() @@ -163,7 +163,8 @@ async def _noop_tool(value: str) -> str: async def _collect_engine(engine: ChatLoopEngine, request: ResolvedChatRequest) -> list: result = [] - async for event in engine.run(request): + execution = await engine.run(request) + async for event in execution.events: result.append(event) return result diff --git a/tests/test_di.py b/tests/test_di.py index ff617074..0fd6456f 100644 --- a/tests/test_di.py +++ b/tests/test_di.py @@ -159,7 +159,7 @@ def test_clean_global_injector() -> None: old_global_injector.get(NodeStoreComponent) running_loop = asyncio.get_running_loop() - clean_global_injector(running_loop) + await clean_global_injector(running_loop) new_injector = get_global_injector() assert new_injector is not None diff --git a/tests/utils/test_tokens.py b/tests/utils/test_tokens.py new file mode 100644 index 00000000..9e1765ec --- /dev/null +++ b/tests/utils/test_tokens.py @@ -0,0 +1,57 @@ +from unittest.mock import MagicMock + +import pytest +from llama_index.core.base.llms.types import ChatMessage, MessageRole + +from private_gpt.components.llm.llm_helper import get_async_tokenizer_fn +from private_gpt.components.llm.tokenizers.tokenizer_base import ( + TokenizedInput, + TokenizerBase, +) +from private_gpt.utils.tokens import async_tokenizer, estimate_token_count + + +class AsyncCapableTokenizer(MagicMock): + def __init__(self): + super().__init__(spec=TokenizerBase) + self.sync_calls = 0 + self.async_calls = 0 + + def __call__(self, texts=None, images=None, audios=None, **kwargs): + del images, audios, kwargs + self.sync_calls += 1 + return TokenizedInput(input_ids=[99]) + + async def acall(self, texts=None, images=None, audios=None, **kwargs): + del images, audios, kwargs + self.async_calls += 1 + text = texts or "" + return TokenizedInput(input_ids=list(range(len(str(text).split())))) + + +@pytest.mark.asyncio +async def test_async_tokenizer_prefers_underlying_acall(): + tokenizer = AsyncCapableTokenizer() + tokenizer_fn = get_async_tokenizer_fn(tokenizer) + + tokens = await async_tokenizer("one two three", tokenizer_fn=tokenizer_fn) + + assert tokens == [0, 1, 2] + assert tokenizer.async_calls == 1 + assert tokenizer.sync_calls == 0 + + +@pytest.mark.asyncio +async def test_estimate_token_count_uses_async_tokenizer_wrapper(): + tokenizer = AsyncCapableTokenizer() + tokenizer_fn = get_async_tokenizer_fn(tokenizer) + + count = await estimate_token_count( + chat_history=[ChatMessage(role=MessageRole.USER, content="one two three")], + tokenizer_fn=tokenizer_fn, + message_to_input=lambda messages: str(messages[0].content), + ) + + assert count == 3 + assert tokenizer.async_calls == 1 + assert tokenizer.sync_calls == 0 diff --git a/tests/worker/test_modes.py b/tests/worker/test_modes.py new file mode 100644 index 00000000..f42dcabc --- /dev/null +++ b/tests/worker/test_modes.py @@ -0,0 +1,18 @@ +import pytest + +from private_gpt.worker import modes + + +def test_celery_mode_forwards_worker_arguments( + monkeypatch: pytest.MonkeyPatch, +) -> None: + commands: list[list[list[str]]] = [] + monkeypatch.setenv("API_ENABLED", "false") + monkeypatch.setattr(modes, "_run_processes", lambda value: commands.append(value)) + + modes.run_celery( + ["--without-gossip"], + max_tasks_per_child_resolver=lambda: 1000, + ) + + assert commands[0][0][-1] == "--without-gossip" diff --git a/tests/worker/test_registry.py b/tests/worker/test_registry.py new file mode 100644 index 00000000..d5c66146 --- /dev/null +++ b/tests/worker/test_registry.py @@ -0,0 +1,33 @@ +import pytest + +from private_gpt.worker.registry import ( + get_worker_mode, + register_worker_mode, + run_worker, +) + + +def test_registered_application_mode_can_be_dispatched( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls: list[str] = [] + register_worker_mode("custom", lambda args: calls.extend(args)) + monkeypatch.setenv("PGPT_WORKER_MODE", "CUSTOM") + + run_worker(["custom"]) + + assert calls == ["custom"] + + +def test_worker_mode_is_required(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("PGPT_WORKER_MODE", raising=False) + + with pytest.raises(ValueError, match="PGPT_WORKER_MODE is required"): + run_worker() + + +def test_unknown_worker_mode_lists_registered_modes() -> None: + register_worker_mode("known", lambda args: None) + + with pytest.raises(ValueError, match=r"Registered modes: .*known"): + get_worker_mode("unknown") diff --git a/uv.lock b/uv.lock index 26b888e5..3a74aa75 100644 --- a/uv.lock +++ b/uv.lock @@ -162,6 +162,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/38/0e/27be9fdef66e72d64c0cdc3cc2823101b80585f8119b5c112c2e8f5f7dab/anyio-4.12.1-py3-none-any.whl", hash = "sha256:d405828884fc140aa80a3c667b8beed277f1dfedec42ba031bd6ac3db606ab6c", size = 113592, upload-time = "2026-01-06T11:45:19.497Z" }, ] +[[package]] +name = "arq" +version = "0.26.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "click" }, + { name = "redis", extra = ["hiredis"] }, +] +sdist = { url = "https://files.pythonhosted.org/packages/4f/65/5add7049297a449d1453e26a8d5924f0d5440b3876edc9e80d5dc621f16d/arq-0.26.3.tar.gz", hash = "sha256:362063ea3c726562fb69c723d5b8ee80827fdefda782a8547da5be3d380ac4b1", size = 291111, upload-time = "2025-01-06T22:44:49.771Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/85/b3/a24a183c628da633b7cafd1759b14aaf47958de82ba6bcae9f1c2898781d/arq-0.26.3-py3-none-any.whl", hash = "sha256:9f4b78149a58c9dc4b88454861a254b7c4e7a159f2c973c89b548288b77e9005", size = 25968, upload-time = "2025-01-06T22:44:45.771Z" }, +] + [[package]] name = "astor" version = "0.8.1" @@ -3181,6 +3194,7 @@ name = "private-gpt" version = "1.0.1" source = { editable = "." } dependencies = [ + { name = "arq" }, { name = "cachetools" }, { name = "fastapi", extra = ["all"] }, { name = "injector" }, @@ -3586,6 +3600,7 @@ vectorstore-qdrant = [ requires-dist = [ { name = "aiobotocore", marker = "extra == 'storage-s3'", specifier = ">=3.6.0,<3.7" }, { name = "anthropic", marker = "extra == 'llm-anthropic'", specifier = ">=0.116.0" }, + { name = "arq", specifier = ">=0.26,<0.27" }, { name = "asyncpg", marker = "extra == 'database-postgres'", specifier = ">=0.31.0,<0.32" }, { name = "beautifulsoup4", marker = "extra == 'ingest-markup'", specifier = ">=4.14.3,<5" }, { name = "black", extras = ["jupyter"], marker = "extra == 'lint'", specifier = ">=22,<23" },