Improvements to Stirling Engine to prepare for SaaS release (#6603)

# Description of Changes
- Use pool for postgres connections
- Add ability to require user ID to be set on API calls to the engine
- Add process-wide concurrency cap on AI access (in addition to existing
user caps)
- Allow number of workers (threads) to be specified for stirling engine
- Update env var names to reflect that the DB is not just for RAG

engine/src/stirling/services/runtime.py

@@ -1,16 +1,22 @@
 from __future__ import annotations
 
+import asyncio
 import logging
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
 from dataclasses import dataclass
-from typing import assert_never
+from typing import Any, assert_never
 
 import httpx
-from pydantic_ai.models import Model, infer_model
+from pydantic_ai import RunContext
+from pydantic_ai.messages import ModelMessage, ModelResponse
+from pydantic_ai.models import Model, ModelRequestParameters, StreamedResponse, infer_model
 from pydantic_ai.models.anthropic import AnthropicModel
+from pydantic_ai.models.wrapper import WrapperModel
 from pydantic_ai.providers.anthropic import AnthropicProvider
 from pydantic_ai.settings import ModelSettings
 
-from stirling.config import ENGINE_ROOT, AppSettings, RagBackend
+from stirling.config import ENGINE_ROOT, AppSettings, DocumentsBackend
 from stirling.documents import (
     DocumentService,
     DocumentStore,
@@ -40,6 +46,37 @@ def _build_anthropic_http_client() -> httpx.AsyncClient:
     return httpx.AsyncClient(limits=httpx.Limits(max_keepalive_connections=0))
 
 
+class ConcurrencyLimitedModel(WrapperModel):
+    """Caps in-flight model API calls with a semaphore shared across the process."""
+
+    def __init__(self, wrapped: Model, semaphore: asyncio.Semaphore) -> None:
+        super().__init__(wrapped)
+        self._semaphore = semaphore
+
+    async def request(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> ModelResponse:
+        async with self._semaphore:
+            return await super().request(messages, model_settings, model_request_parameters)
+
+    @asynccontextmanager
+    async def request_stream(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+        run_context: RunContext[Any] | None = None,
+    ) -> AsyncIterator[StreamedResponse]:
+        async with self._semaphore:
+            async with super().request_stream(
+                messages, model_settings, model_request_parameters, run_context
+            ) as response_stream:
+                yield response_stream
+
+
 @dataclass(frozen=True)
 class AppRuntime:
     settings: AppSettings
@@ -74,17 +111,21 @@ def validate_structured_output_support(model: Model, model_name: str) -> None:
 
 def _build_document_store(settings: AppSettings) -> DocumentStore:
     """Build the configured document store backend."""
-    if settings.rag_backend == RagBackend.SQLITE:
-        store_path = settings.rag_store_path
+    if settings.documents_backend == DocumentsBackend.SQLITE:
+        store_path = settings.documents_sqlite_path
         # Treat ":memory:" as a special in-process token; otherwise resolve against the engine root.
         if str(store_path) != ":memory:" and not store_path.is_absolute():
             store_path = ENGINE_ROOT / store_path
         logger.info("Document store backend=sqlite, db_path=%s", store_path)
         return SqliteVecStore(db_path=store_path)
-    if settings.rag_backend == RagBackend.PGVECTOR:
+    if settings.documents_backend == DocumentsBackend.PGVECTOR:
         logger.info("Document store backend=pgvector, dsn=<configured>")
-        return PgVectorStore(dsn=settings.rag_pgvector_dsn)
-    assert_never(settings.rag_backend)
+        return PgVectorStore(
+            dsn=settings.documents_pgvector_dsn,
+            pool_min_size=settings.documents_pgvector_pool_min_size,
+            pool_max_size=settings.documents_pgvector_pool_max_size,
+        )
+    assert_never(settings.documents_backend)
 
 
 def _build_documents(settings: AppSettings) -> DocumentService:
@@ -105,10 +146,13 @@ def build_runtime(settings: AppSettings) -> AppRuntime:
     validate_structured_output_support(fast_model, settings.fast_model_name)
     validate_structured_output_support(smart_model, settings.smart_model_name)
 
+    # One semaphore across both tiers: the cap protects the provider account
+    # and process resources, which the tiers share.
+    model_semaphore = asyncio.Semaphore(settings.model_max_concurrency)
     return AppRuntime(
         settings=settings,
-        fast_model=fast_model,
-        smart_model=smart_model,
+        fast_model=ConcurrencyLimitedModel(fast_model, model_semaphore),
+        smart_model=ConcurrencyLimitedModel(smart_model, model_semaphore),
         documents=_build_documents(settings),
     )