feat: add adaptative

private_gpt/components/tools/tool_scheduler.py

@@ -1,14 +1,17 @@
 from __future__ import annotations
 
 import asyncio
+import dataclasses
 import logging
+import time
 import uuid
 from abc import ABC, abstractmethod
-from collections.abc import Awaitable, Callable
-from typing import Any
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from collections.abc import Awaitable, Callable
 
 from injector import inject, singleton
-from pydantic import BaseModel, ConfigDict
 
 from private_gpt.components.tools.tool_names import (
     BASH_TOOL_NAME,
@@ -64,9 +67,8 @@ class ImmediateToolScheduler(BaseToolScheduler):
         return await func()
 
 
-class _PendingCall(BaseModel):
-    model_config = ConfigDict(arbitrary_types_allowed=True)
-
+@dataclasses.dataclass
+class _PendingCall:
     score: float
     counter: int
     entry_id: str
@@ -187,3 +189,69 @@ class QueuedToolScheduler(BaseToolScheduler):
             if not future.done():
                 future.cancel()
             raise
+
+
+@singleton
+class AdaptiveToolScheduler(BaseToolScheduler):
+    @inject
+    def __init__(
+        self,
+        immediate: ImmediateToolScheduler,
+        queued: QueuedToolScheduler,
+        settings: Settings,
+    ) -> None:
+        cfg = settings.tool_scheduler.adaptive
+        self._immediate = immediate
+        self._queued = queued
+        self._threshold_ms: float = cfg.threshold_ms
+        self._recovery_ms: float = cfg.threshold_ms * cfg.recovery_ratio
+        self._alpha: float = cfg.ewma_alpha
+        self._ewma_ms: float | None = None
+        self._use_queue: bool = False
+
+    def _update(self, elapsed_ms: float) -> None:
+        self._ewma_ms = (
+            elapsed_ms
+            if self._ewma_ms is None
+            else self._alpha * elapsed_ms + (1 - self._alpha) * self._ewma_ms
+        )
+        if not self._use_queue and self._ewma_ms > self._threshold_ms:
+            self._use_queue = True
+            logger.info("AdaptiveToolScheduler: -> queued (ewma=%.0fms)", self._ewma_ms)
+        elif self._use_queue and self._ewma_ms < self._recovery_ms:
+            self._use_queue = False
+            logger.info(
+                "AdaptiveToolScheduler: -> immediate (ewma=%.0fms)", self._ewma_ms
+            )
+
+    async def execute(
+        self,
+        tool_name: str,
+        chat_priority: int | None,
+        func: Callable[[], Awaitable[Any]],
+    ) -> Any:
+        scheduler = self._queued if self._use_queue else self._immediate
+        start = time.monotonic()
+        try:
+            return await scheduler.execute(tool_name, chat_priority, func)
+        finally:
+            self._update((time.monotonic() - start) * 1000)
+
+
+@singleton
+class ToolSchedulerFactory:
+    @inject
+    def __init__(
+        self,
+        settings: Settings,
+        queued: QueuedToolScheduler,
+        adaptive: AdaptiveToolScheduler,
+    ) -> None:
+        self._scheduler: BaseToolScheduler = {
+            "immediate": ImmediateToolScheduler(),
+            "queued": queued,
+            "adaptive": adaptive,
+        }[settings.tool_scheduler.mode]
+
+    def get(self) -> BaseToolScheduler:
+        return self._scheduler

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:

private_gpt/server/chat/chat_service.py

@@ -23,11 +23,7 @@ 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.tools.tool_scheduler import (
-    BaseToolScheduler,
-    ImmediateToolScheduler,
-    QueuedToolScheduler,
-)
+from private_gpt.components.tools.tool_scheduler import ToolSchedulerFactory
 from private_gpt.components.vector_store.vector_store_component import (
     VectorStoreComponent,
 )
@@ -155,7 +151,7 @@ class ChatService:
         chat_interceptor_service: ChatInterceptorService,
         models_service: ModelsService,
         container_registry: ContainerRegistry,
-        tool_scheduler: QueuedToolScheduler,
+        scheduler_factory: ToolSchedulerFactory,
     ) -> None:
         self.settings = settings
         self.llm_component = llm_component
@@ -166,11 +162,7 @@ class ChatService:
         self.chat_interceptor_service = chat_interceptor_service
         self.models_service = models_service
         self.container_registry = container_registry
-        self._tool_scheduler: BaseToolScheduler = (
-            tool_scheduler
-            if settings.tool_scheduler.enabled
-            else ImmediateToolScheduler()
-        )
+        self._tool_scheduler = scheduler_factory.get()
 
     def _build_loop_engine(self) -> ChatLoopEngine:
         # Don't build a singleton since the interceptors

private_gpt/settings/settings.py

@@ -1639,41 +1639,56 @@ class ToolSchedulerWeightsSettings(BaseModel):
         default=0.4,
         description="Weight given to the chat request priority signal (0-1).",
     )
-    age: float = Field(
-        default=0.4,
-        description=(
-            "Weight given to wait time in the queue. "
-            "Higher values reduce starvation of long-waiting calls."
-        ),
-    )
     complexity: float = Field(
         default=0.2,
         description="Weight given to the estimated tool complexity (0-1).",
     )
 
 
-class ToolSchedulerSettings(BaseModel):
-    enabled: bool = Field(
-        default=False,
+class AdaptiveToolSchedulerSettings(BaseModel):
+    threshold_ms: float = Field(
+        default=2000.0,
         description=(
-            "When True, all tool executions are routed through a shared priority queue "
-            "that limits global concurrency and orders calls by urgency."
+            "EWMA of tool execution time (ms) above which the adaptive scheduler "
+            "switches from immediate to queued mode."
+        ),
+    )
+    recovery_ratio: float = Field(
+        default=0.7,
+        description=(
+            "Fraction of threshold_ms at which the adaptive scheduler switches back "
+            "to immediate mode. Must be < 1 to provide hysteresis."
+        ),
+    )
+    ewma_alpha: float = Field(
+        default=0.3,
+        description=(
+            "Smoothing factor for the exponentially weighted moving average. "
+            "Higher values react faster to load changes."
+        ),
+    )
+
+
+class ToolSchedulerSettings(BaseModel):
+    mode: Literal["immediate", "queued", "adaptive"] = Field(
+        default="immediate",
+        description=(
+            "'immediate' runs tools without queuing; "
+            "'queued' routes all calls through a shared priority queue; "
+            "'adaptive' switches between the two based on measured execution time."
         ),
     )
     max_concurrent_tools: int = Field(
         default=5,
-        description="Maximum number of tool calls executed simultaneously across all chats.",
-    )
-    max_age_seconds: float = Field(
-        default=60.0,
-        description=(
-            "Time (seconds) after which a queued tool call reaches maximum age urgency. "
-            "Prevents starvation by boosting priority of long-waiting calls."
-        ),
+        description="Maximum simultaneous tool calls (used by 'queued' and 'adaptive' modes).",
     )
     weights: ToolSchedulerWeightsSettings = Field(
         default_factory=ToolSchedulerWeightsSettings,
-        description="Relative weights for the three priority signals.",
+        description="Priority weights for the queued scheduler.",
+    )
+    adaptive: AdaptiveToolSchedulerSettings = Field(
+        default_factory=AdaptiveToolSchedulerSettings,
+        description="Settings for the adaptive scheduler.",
     )
 
 

settings.yaml

@@ -290,13 +290,15 @@ semaphore:
   mode: ${PGPT_SEMAPHORE_MODE:memory}
 
 tool_scheduler:
-  enabled: ${PGPT_TOOL_SCHEDULER_ENABLED:false}
-  max_concurrent_tools: ${PGPT_TOOL_SCHEDULER_MAX_CONCURRENT_TOOLS:5}
-  max_age_seconds: ${PGPT_TOOL_SCHEDULER_MAX_AGE_SECONDS:60}
+  mode: ${PGPT_TOOL_SCHEDULER_MODE:adaptive}
+  max_concurrent_tools: ${PGPT_TOOL_SCHEDULER_MAX_CONCURRENT_TOOLS:4}
   weights:
     chat_priority: ${PGPT_TOOL_SCHEDULER_WEIGHT_PRIORITY:0.4}
-    age: ${PGPT_TOOL_SCHEDULER_WEIGHT_AGE:0.4}
     complexity: ${PGPT_TOOL_SCHEDULER_WEIGHT_COMPLEXITY:0.2}
+  adaptive:
+    threshold_ms: ${PGPT_TOOL_SCHEDULER_ADAPTIVE_THRESHOLD_MS:800}
+    recovery_ratio: ${PGPT_TOOL_SCHEDULER_ADAPTIVE_RECOVERY_RATIO:0.7}
+    ewma_alpha: ${PGPT_TOOL_SCHEDULER_ADAPTIVE_EWMA_ALPHA:0.5}
 
 transformation:
   pptx: