From d52c7ced7cb7b2078a4312a7671451522708ba6f Mon Sep 17 00:00:00 2001 From: James Brunton Date: Thu, 11 Jun 2026 16:31:35 +0100 Subject: [PATCH] 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 --- .taskfiles/engine.yml | 2 +- engine/.env | 30 ++++-- engine/Dockerfile | 1 + engine/pyproject.toml | 2 +- engine/src/stirling/api/app.py | 22 +++-- engine/src/stirling/api/dependencies.py | 18 +++- engine/src/stirling/config/__init__.py | 4 +- engine/src/stirling/config/settings.py | 26 ++++- engine/src/stirling/documents/README.md | 8 +- .../src/stirling/documents/pgvector_store.py | 95 ++++++++++++------- engine/src/stirling/services/runtime.py | 64 +++++++++++-- engine/tests/conftest.py | 12 ++- engine/tests/pdf_comment/test_routes.py | 12 ++- engine/tests/test_model_concurrency.py | 58 +++++++++++ engine/tests/test_stirling_api.py | 20 ++++ engine/tests/test_stirling_contracts.py | 12 ++- engine/uv.lock | 19 +++- 17 files changed, 315 insertions(+), 90 deletions(-) create mode 100644 engine/tests/test_model_concurrency.py diff --git a/.taskfiles/engine.yml b/.taskfiles/engine.yml index 1834495b1..244b48353 100644 --- a/.taskfiles/engine.yml +++ b/.taskfiles/engine.yml @@ -43,7 +43,7 @@ tasks: env: PYTHONUNBUFFERED: "1" cmds: - - uv run uvicorn stirling.api.app:app --host 0.0.0.0 --port {{.PORT}} + - uv run uvicorn stirling.api.app:app --host 0.0.0.0 --port {{.PORT}} --workers "${STIRLING_ENGINE_WORKERS:-4}" dev: desc: "Start engine dev server with hot reload" diff --git a/engine/.env b/engine/.env index e334a9bf9..a99abf0cc 100644 --- a/engine/.env +++ b/engine/.env @@ -14,19 +14,30 @@ STIRLING_FAST_MODEL=anthropic:claude-haiku-4-5 STIRLING_SMART_MODEL_MAX_TOKENS=8192 STIRLING_FAST_MODEL_MAX_TOKENS=2048 -# RAG Configuration — retrieval-augmented generation is always on. -# Embedding provider credentials are handled natively (e.g. VOYAGE_API_KEY for VoyageAI). -STIRLING_RAG_EMBEDDING_MODEL=voyageai:voyage-4 +# Process-wide cap on concurrent model API calls, shared by both model tiers. +# Per-request fan-outs (chunked reasoner workers, contradiction detection) are +# bounded per request; this bounds their product across concurrent requests. +STIRLING_MODEL_MAX_CONCURRENCY=32 -# Vector store backend: "sqlite" (embedded) or "pgvector" (external Postgres). -STIRLING_RAG_BACKEND=sqlite +# Document store: the one database holding vector chunks, ordered page text, +# and ACL rows. Backend is "sqlite" (embedded sqlite-vec) or "pgvector" +# (external Postgres). +STIRLING_DOCUMENTS_BACKEND=sqlite # Path to the sqlite-vec database file (used when backend=sqlite). -STIRLING_RAG_STORE_PATH=data/rag.db +STIRLING_DOCUMENTS_SQLITE_PATH=data/rag.db # Postgres DSN for pgvector (used when backend=pgvector). Leave empty when backend=sqlite. # Example: postgresql://user:password@host:5432/dbname -STIRLING_RAG_PGVECTOR_DSN= +STIRLING_DOCUMENTS_PGVECTOR_DSN= + +# Connection pool bounds for the pgvector backend. +STIRLING_DOCUMENTS_PGVECTOR_POOL_MIN_SIZE=1 +STIRLING_DOCUMENTS_PGVECTOR_POOL_MAX_SIZE=10 + +# RAG Configuration - retrieval-augmented generation is always on. +# Embedding provider credentials are handled natively (e.g. VOYAGE_API_KEY for VoyageAI). +STIRLING_RAG_EMBEDDING_MODEL=voyageai:voyage-4 STIRLING_RAG_CHUNK_SIZE=512 STIRLING_RAG_CHUNK_OVERLAP=64 @@ -57,6 +68,11 @@ STIRLING_CHUNKED_REASONER_NOTES_CHAR_BUDGET=250000 STIRLING_MAX_PAGES=200 STIRLING_MAX_CHARACTERS=200000 +# Reject API requests that lack an X-User-Id header. Self-hosted deployments +# with security disabled have no user identity, so this is off by default. +# Multi-tenant (SaaS) deployments must set it to true. +STIRLING_REQUIRE_USER_ID=false + # PostHog analytics. Set STIRLING_POSTHOG_ENABLED=true and provide an API key to enable. STIRLING_POSTHOG_ENABLED=false STIRLING_POSTHOG_API_KEY=phc_VOdeYnlevc2T63m3myFGjeBlRcIusRgmhfx6XL5a1iz diff --git a/engine/Dockerfile b/engine/Dockerfile index 5a0a835fa..e3f548e8d 100644 --- a/engine/Dockerfile +++ b/engine/Dockerfile @@ -26,6 +26,7 @@ COPY .taskfiles/ ./.taskfiles/ ENV PATH="/app/engine/.venv/bin:$PATH" ENV PYTHONUNBUFFERED=1 +ENV STIRLING_ENGINE_WORKERS=4 EXPOSE 5001 diff --git a/engine/pyproject.toml b/engine/pyproject.toml index 4074c21ee..7bd2fa7be 100644 --- a/engine/pyproject.toml +++ b/engine/pyproject.toml @@ -6,7 +6,7 @@ requires-python = ">=3.13" dependencies = [ "fastapi>=0.116.0", "pgvector>=0.3.6", - "psycopg[binary]>=3.2", + "psycopg[binary,pool]>=3.2", "pydantic>=2.0.0", "pydantic-ai>=1.67.0", "pydantic-ai-slim[voyageai]>=1.67.0", diff --git a/engine/src/stirling/api/app.py b/engine/src/stirling/api/app.py index 699cede06..7e6272c28 100644 --- a/engine/src/stirling/api/app.py +++ b/engine/src/stirling/api/app.py @@ -18,6 +18,7 @@ from stirling.agents import ( ) from stirling.agents.ledger import MathAuditorAgent from stirling.agents.pdf_comment import PdfCommentAgent +from stirling.api.dependencies import enforce_required_user_id from stirling.api.engine_auth import EngineSharedSecretMiddleware from stirling.api.middleware import UserIdMiddleware from stirling.api.routes import ( @@ -118,15 +119,18 @@ async def lifespan(fast_api: FastAPI): app = FastAPI(title="Stirling AI Engine", lifespan=lifespan, version="0.1.0") app.add_middleware(UserIdMiddleware) app.add_middleware(EngineSharedSecretMiddleware) -app.include_router(orchestrator_router) -app.include_router(pdf_edit_router) -app.include_router(pdf_question_router) -app.include_router(agent_draft_router) -app.include_router(execution_router) -app.include_router(document_router) -app.include_router(ledger_router) -app.include_router(pdf_comments_router) -app.include_router(agent_capabilities_router) +# Every router gets the same configurable identity gate; /health stays open +# for liveness probes. See enforce_required_user_id for the policy. +_user_gate = [Depends(enforce_required_user_id)] +app.include_router(orchestrator_router, dependencies=_user_gate) +app.include_router(pdf_edit_router, dependencies=_user_gate) +app.include_router(pdf_question_router, dependencies=_user_gate) +app.include_router(agent_draft_router, dependencies=_user_gate) +app.include_router(execution_router, dependencies=_user_gate) +app.include_router(document_router, dependencies=_user_gate) +app.include_router(ledger_router, dependencies=_user_gate) +app.include_router(pdf_comments_router, dependencies=_user_gate) +app.include_router(agent_capabilities_router, dependencies=_user_gate) @app.get("/health", response_model=HealthResponse) diff --git a/engine/src/stirling/api/dependencies.py b/engine/src/stirling/api/dependencies.py index cae20a8ae..780e97c9d 100644 --- a/engine/src/stirling/api/dependencies.py +++ b/engine/src/stirling/api/dependencies.py @@ -1,6 +1,8 @@ from __future__ import annotations -from fastapi import HTTPException, Request, status +from typing import Annotated + +from fastapi import Depends, HTTPException, Request, status from stirling.agents import ( ExecutionPlanningAgent, @@ -11,6 +13,7 @@ from stirling.agents import ( ) from stirling.agents.ledger import MathAuditorAgent from stirling.agents.pdf_comment import PdfCommentAgent +from stirling.config import AppSettings, load_settings from stirling.documents import DocumentService from stirling.models import UserId from stirling.services import AppRuntime, current_user_id @@ -67,3 +70,16 @@ def require_user_id() -> UserId: detail="X-User-Id header is required", ) return user_id + + +def enforce_required_user_id( + settings: Annotated[AppSettings, Depends(load_settings)], +) -> None: + """Router-level boundary gate, applied uniformly to every router.""" + if not settings.require_user_id: + return + if current_user_id.get() is None: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="X-User-Id header is required", + ) diff --git a/engine/src/stirling/config/__init__.py b/engine/src/stirling/config/__init__.py index c5044a214..f038d02e5 100644 --- a/engine/src/stirling/config/__init__.py +++ b/engine/src/stirling/config/__init__.py @@ -1,10 +1,10 @@ """Configuration models and loaders for the Stirling AI service.""" -from .settings import ENGINE_ROOT, AppSettings, RagBackend, load_settings +from .settings import ENGINE_ROOT, AppSettings, DocumentsBackend, load_settings __all__ = [ "ENGINE_ROOT", "AppSettings", - "RagBackend", + "DocumentsBackend", "load_settings", ] diff --git a/engine/src/stirling/config/settings.py b/engine/src/stirling/config/settings.py index 4cd2bbc9e..1293cc371 100644 --- a/engine/src/stirling/config/settings.py +++ b/engine/src/stirling/config/settings.py @@ -15,7 +15,7 @@ ENV_FILE = ENGINE_ROOT / ".env" ENV_LOCAL_FILE = ENGINE_ROOT / ".env.local" -class RagBackend(StrEnum): +class DocumentsBackend(StrEnum): SQLITE = "sqlite" PGVECTOR = "pgvector" @@ -27,12 +27,22 @@ class AppSettings(BaseSettings): fast_model_name: str = Field(validation_alias="STIRLING_FAST_MODEL") smart_model_max_tokens: int = Field(validation_alias="STIRLING_SMART_MODEL_MAX_TOKENS") fast_model_max_tokens: int = Field(validation_alias="STIRLING_FAST_MODEL_MAX_TOKENS") + # Process-wide ceiling on concurrent model API calls, shared by both model + # tiers. Per-request fan-outs (chunked reasoner, contradiction detection) + # carry their own per-request caps, but those multiply under concurrent + # traffic; this is the global backstop. + model_max_concurrency: int = Field(validation_alias="STIRLING_MODEL_MAX_CONCURRENCY") - # RAG settings — always on; the backend picks between embedded sqlite-vec and external pgvector. - rag_backend: RagBackend = Field(validation_alias="STIRLING_RAG_BACKEND") + # Document store: the one database holding vector chunks, ordered page + # text, and ACL rows - embedded sqlite-vec or external pgvector. + documents_backend: DocumentsBackend = Field(validation_alias="STIRLING_DOCUMENTS_BACKEND") + documents_sqlite_path: Path = Field(validation_alias="STIRLING_DOCUMENTS_SQLITE_PATH") + documents_pgvector_dsn: str = Field(validation_alias="STIRLING_DOCUMENTS_PGVECTOR_DSN") + documents_pgvector_pool_min_size: int = Field(validation_alias="STIRLING_DOCUMENTS_PGVECTOR_POOL_MIN_SIZE") + documents_pgvector_pool_max_size: int = Field(validation_alias="STIRLING_DOCUMENTS_PGVECTOR_POOL_MAX_SIZE") + + # RAG settings - always on. rag_embedding_model: str = Field(validation_alias="STIRLING_RAG_EMBEDDING_MODEL") - rag_store_path: Path = Field(validation_alias="STIRLING_RAG_STORE_PATH") - rag_pgvector_dsn: str = Field(validation_alias="STIRLING_RAG_PGVECTOR_DSN") rag_chunk_size: int = Field(validation_alias="STIRLING_RAG_CHUNK_SIZE") rag_chunk_overlap: int = Field(validation_alias="STIRLING_RAG_CHUNK_OVERLAP") rag_default_top_k: int = Field(validation_alias="STIRLING_RAG_TOP_K") @@ -88,6 +98,12 @@ class AppSettings(BaseSettings): max_pages: int = Field(validation_alias="STIRLING_MAX_PAGES") max_characters: int = Field(validation_alias="STIRLING_MAX_CHARACTERS") + # When true, API routes reject requests that lack an X-User-Id header at + # the boundary. Self-hosted deployments with security disabled have no + # user identity and leave this off; multi-tenant deployments turn it on so + # user-scoped work is never processed without a tenant attached. + require_user_id: bool = Field(validation_alias="STIRLING_REQUIRE_USER_ID") + log_level: str = Field(default="INFO", validation_alias="STIRLING_LOG_LEVEL") log_file: str = Field(default="", validation_alias="STIRLING_LOG_FILE") # When true, raises httpx + httpcore logger levels so every outgoing diff --git a/engine/src/stirling/documents/README.md b/engine/src/stirling/documents/README.md index e1c759044..6848d8a10 100644 --- a/engine/src/stirling/documents/README.md +++ b/engine/src/stirling/documents/README.md @@ -54,10 +54,10 @@ everything = RagCapability(runtime.documents) Non-secret defaults live in the committed `engine/.env`: ``` -STIRLING_RAG_BACKEND=sqlite # or "pgvector" +STIRLING_DOCUMENTS_BACKEND=sqlite # or "pgvector" STIRLING_RAG_EMBEDDING_MODEL=voyageai:voyage-4 -STIRLING_RAG_STORE_PATH=data/rag.db # used when backend=sqlite -STIRLING_RAG_PGVECTOR_DSN= # used when backend=pgvector +STIRLING_DOCUMENTS_SQLITE_PATH=data/rag.db # used when backend=sqlite +STIRLING_DOCUMENTS_PGVECTOR_DSN= # used when backend=pgvector STIRLING_RAG_CHUNK_SIZE=512 STIRLING_RAG_CHUNK_OVERLAP=64 STIRLING_RAG_TOP_K=5 @@ -76,7 +76,7 @@ VOYAGE_API_KEY=your-key and self-hosted deployments. **`pgvector`** - External PostgreSQL with the `vector` extension. Point -`STIRLING_RAG_PGVECTOR_DSN` at your Postgres instance. +`STIRLING_DOCUMENTS_PGVECTOR_DSN` at your Postgres instance. Both backends implement the same `DocumentStore` interface, so agents and the service work identically regardless of which you pick. diff --git a/engine/src/stirling/documents/pgvector_store.py b/engine/src/stirling/documents/pgvector_store.py index 8aff5ad3a..1599ae047 100644 --- a/engine/src/stirling/documents/pgvector_store.py +++ b/engine/src/stirling/documents/pgvector_store.py @@ -1,10 +1,12 @@ from __future__ import annotations +import asyncio import json from datetime import datetime import psycopg from pgvector.psycopg import register_vector_async +from psycopg_pool import AsyncConnectionPool from stirling.contracts.documents import Page, PageRange from stirling.documents.store import Document, DocumentStore, SearchResult, StoredPage @@ -13,12 +15,20 @@ from stirling.models import OwnerId, PrincipalId _READ_PERMISSION = "read" +async def _register_vector(conn: psycopg.AsyncConnection) -> None: + await register_vector_async(conn) + + class PgVectorStore(DocumentStore): """PostgreSQL + pgvector backed store, scoped by owner with ACL-gated reads. Connects to an external Postgres instance (DSN provided via config) and uses the `vector` extension for similarity search. The schema is created on first use. + Queries run on a shared async connection pool. The pool is opened lazily on + first use, after the schema bootstrap, because each pooled connection's + configure hook registers the ``vector`` type, which must already exist. + Owned tables: * ``documents_meta`` - parent row, one per ``(collection, owner_id)``. @@ -31,21 +41,36 @@ class PgVectorStore(DocumentStore): Writes are owner-scoped; reads are ACL-scoped. """ - def __init__(self, dsn: str) -> None: + def __init__(self, dsn: str, pool_min_size: int, pool_max_size: int) -> None: if not dsn: - raise ValueError("pgvector backend requires a non-empty DSN (STIRLING_RAG_PGVECTOR_DSN)") + raise ValueError("pgvector backend requires a non-empty DSN (STIRLING_DOCUMENTS_PGVECTOR_DSN)") self._dsn = dsn + self._pool = AsyncConnectionPool( + dsn, + min_size=pool_min_size, + max_size=pool_max_size, + open=False, + configure=_register_vector, + check=AsyncConnectionPool.check_connection, + ) self._initialized = False + self._init_lock = asyncio.Lock() - async def _connect(self) -> psycopg.AsyncConnection: - conn = await psycopg.AsyncConnection.connect(self._dsn) - await register_vector_async(conn) - return conn - - async def _ensure_schema(self) -> None: + async def _ensure_ready(self) -> None: + """Create the schema and open the pool, exactly once across concurrent callers.""" if self._initialized: return - async with await self._connect() as conn: + async with self._init_lock: + if self._initialized: + return + await self._bootstrap_schema() + await self._pool.open() + self._initialized = True + + async def _bootstrap_schema(self) -> None: + # Uses a dedicated non-pool connection: pooled connections register the + # vector type on checkout, which fails until the extension exists. + async with await psycopg.AsyncConnection.connect(self._dsn) as conn: async with conn.cursor() as cur: await cur.execute("CREATE EXTENSION IF NOT EXISTS vector") await cur.execute( @@ -118,7 +143,6 @@ class PgVectorStore(DocumentStore): "CREATE INDEX IF NOT EXISTS idx_acl_principal_permission ON document_acl(principal_id, permission)" ) await conn.commit() - self._initialized = True # ── lifecycle of the (collection, owner_id) row ──────────────────────── @@ -129,8 +153,8 @@ class PgVectorStore(DocumentStore): owner_id: OwnerId, expires_at: datetime | None, ) -> None: - await self._ensure_schema() - async with await self._connect() as conn: + await self._ensure_ready() + async with self._pool.connection() as conn: async with conn.cursor() as cur: await cur.execute( """ @@ -145,8 +169,8 @@ class PgVectorStore(DocumentStore): await conn.commit() async def purge_owner(self, owner_id: OwnerId) -> int: - await self._ensure_schema() - async with await self._connect() as conn: + await self._ensure_ready() + async with self._pool.connection() as conn: async with conn.cursor() as cur: await cur.execute( "DELETE FROM documents_meta WHERE owner_id = %s", @@ -157,8 +181,8 @@ class PgVectorStore(DocumentStore): return deleted async def reap_expired(self) -> int: - await self._ensure_schema() - async with await self._connect() as conn: + await self._ensure_ready() + async with self._pool.connection() as conn: async with conn.cursor() as cur: await cur.execute("DELETE FROM documents_meta WHERE expires_at IS NOT NULL AND expires_at < NOW()") deleted = cur.rowcount @@ -166,8 +190,8 @@ class PgVectorStore(DocumentStore): return deleted async def delete_collection(self, collection: str, owner_id: OwnerId) -> bool: - await self._ensure_schema() - async with await self._connect() as conn: + await self._ensure_ready() + async with self._pool.connection() as conn: async with conn.cursor() as cur: # Cascade FKs handle rag_documents, document_pages, document_acl. await cur.execute( @@ -192,8 +216,8 @@ class PgVectorStore(DocumentStore): if not documents: return - await self._ensure_schema() - async with await self._connect() as conn: + await self._ensure_ready() + async with self._pool.connection() as conn: async with conn.cursor() as cur: for doc, emb in zip(documents, embeddings): await cur.execute( @@ -211,8 +235,8 @@ class PgVectorStore(DocumentStore): await conn.commit() async def add_pages(self, collection: str, pages: list[StoredPage], owner_id: OwnerId) -> None: - await self._ensure_schema() - async with await self._connect() as conn: + await self._ensure_ready() + async with self._pool.connection() as conn: async with conn.cursor() as cur: await cur.execute( "DELETE FROM document_pages WHERE collection = %s AND owner_id = %s", @@ -238,8 +262,8 @@ class PgVectorStore(DocumentStore): ) -> None: if not principals: return - await self._ensure_schema() - async with await self._connect() as conn: + await self._ensure_ready() + async with self._pool.connection() as conn: async with conn.cursor() as cur: await cur.executemany( """ @@ -257,8 +281,8 @@ class PgVectorStore(DocumentStore): owner_id: OwnerId, principal: PrincipalId, ) -> None: - await self._ensure_schema() - async with await self._connect() as conn: + await self._ensure_ready() + async with self._pool.connection() as conn: async with conn.cursor() as cur: await cur.execute( "DELETE FROM document_acl WHERE collection = %s AND owner_id = %s AND principal_id = %s", @@ -277,8 +301,8 @@ class PgVectorStore(DocumentStore): ) -> list[SearchResult]: if not principals: return [] - await self._ensure_schema() - async with await self._connect() as conn: + await self._ensure_ready() + async with self._pool.connection() as conn: async with conn.cursor() as cur: owner_id = await self._readable_owner_for(cur, collection, principals) if owner_id is None: @@ -332,8 +356,8 @@ class PgVectorStore(DocumentStore): ) -> list[Page]: if not principals: return [] - await self._ensure_schema() - async with await self._connect() as conn: + await self._ensure_ready() + async with self._pool.connection() as conn: async with conn.cursor() as cur: owner_id = await self._readable_owner_for(cur, collection, principals) if owner_id is None: @@ -358,8 +382,8 @@ class PgVectorStore(DocumentStore): async def has_collection(self, collection: str, principals: list[PrincipalId]) -> bool: if not principals: return False - await self._ensure_schema() - async with await self._connect() as conn: + await self._ensure_ready() + async with self._pool.connection() as conn: async with conn.cursor() as cur: owner_id = await self._readable_owner_for(cur, collection, principals) return owner_id is not None @@ -367,8 +391,8 @@ class PgVectorStore(DocumentStore): async def list_collections(self, principals: list[PrincipalId]) -> list[str]: if not principals: return [] - await self._ensure_schema() - async with await self._connect() as conn: + await self._ensure_ready() + async with self._pool.connection() as conn: async with conn.cursor() as cur: await cur.execute( """ @@ -383,5 +407,4 @@ class PgVectorStore(DocumentStore): return [r[0] for r in rows] async def close(self) -> None: - # Connections are opened and closed per call, so nothing persistent to release. - return None + await self._pool.close() diff --git a/engine/src/stirling/services/runtime.py b/engine/src/stirling/services/runtime.py index 685509cfe..06629347c 100644 --- a/engine/src/stirling/services/runtime.py +++ b/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=") - 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), ) diff --git a/engine/tests/conftest.py b/engine/tests/conftest.py index 1898b9129..d4c854313 100644 --- a/engine/tests/conftest.py +++ b/engine/tests/conftest.py @@ -5,7 +5,7 @@ from pathlib import Path import pytest -from stirling.config import AppSettings, RagBackend, load_settings +from stirling.config import AppSettings, DocumentsBackend, load_settings from stirling.services import build_runtime from stirling.services.runtime import AppRuntime @@ -23,10 +23,13 @@ def build_app_settings() -> AppSettings: fast_model_name="test", smart_model_max_tokens=8192, fast_model_max_tokens=2048, - rag_backend=RagBackend.SQLITE, + model_max_concurrency=32, + documents_backend=DocumentsBackend.SQLITE, rag_embedding_model="voyageai:voyage-4", - rag_store_path=Path(":memory:"), - rag_pgvector_dsn="", + documents_sqlite_path=Path(":memory:"), + documents_pgvector_dsn="", + documents_pgvector_pool_min_size=1, + documents_pgvector_pool_max_size=10, rag_chunk_size=512, rag_chunk_overlap=64, rag_default_top_k=5, @@ -41,6 +44,7 @@ def build_app_settings() -> AppSettings: contradiction_canonicaliser_batch_size=500, max_pages=200, max_characters=200_000, + require_user_id=False, posthog_enabled=False, posthog_api_key="", posthog_host="https://eu.i.posthog.com", diff --git a/engine/tests/pdf_comment/test_routes.py b/engine/tests/pdf_comment/test_routes.py index 6fce78f25..0056e10fe 100644 --- a/engine/tests/pdf_comment/test_routes.py +++ b/engine/tests/pdf_comment/test_routes.py @@ -16,7 +16,7 @@ from fastapi.testclient import TestClient from stirling.api import app from stirling.api.dependencies import get_pdf_comment_agent -from stirling.config import AppSettings, RagBackend, load_settings +from stirling.config import AppSettings, DocumentsBackend, load_settings from stirling.contracts.pdf_comments import ( PdfCommentInstruction, PdfCommentRequest, @@ -35,10 +35,13 @@ class StubSettingsProvider: fast_model_name="test", smart_model_max_tokens=8192, fast_model_max_tokens=2048, - rag_backend=RagBackend.SQLITE, + model_max_concurrency=32, + documents_backend=DocumentsBackend.SQLITE, rag_embedding_model="test-embed", - rag_store_path=Path(":memory:"), - rag_pgvector_dsn="", + documents_sqlite_path=Path(":memory:"), + documents_pgvector_dsn="", + documents_pgvector_pool_min_size=1, + documents_pgvector_pool_max_size=10, rag_chunk_size=512, rag_chunk_overlap=64, rag_default_top_k=5, @@ -49,6 +52,7 @@ class StubSettingsProvider: chunked_reasoner_notes_char_budget=250_000, max_pages=100, max_characters=100_000, + require_user_id=False, posthog_enabled=False, posthog_api_key="", posthog_host="https://eu.i.posthog.com", diff --git a/engine/tests/test_model_concurrency.py b/engine/tests/test_model_concurrency.py new file mode 100644 index 000000000..ada8b5727 --- /dev/null +++ b/engine/tests/test_model_concurrency.py @@ -0,0 +1,58 @@ +from __future__ import annotations + +import asyncio + +import pytest +from pydantic_ai.messages import ModelMessage, ModelResponse, TextPart +from pydantic_ai.models import Model, ModelRequestParameters +from pydantic_ai.settings import ModelSettings + +from stirling.services.runtime import ConcurrencyLimitedModel + + +class TrackingModel(Model): + """Records the high-water mark of concurrent in-flight requests.""" + + def __init__(self) -> None: + super().__init__() + self.active = 0 + self.max_active = 0 + + @property + def model_name(self) -> str: + return "tracking" + + @property + def system(self) -> str: + return "test" + + async def request( + self, + messages: list[ModelMessage], + model_settings: ModelSettings | None, + model_request_parameters: ModelRequestParameters, + ) -> ModelResponse: + self.active += 1 + self.max_active = max(self.max_active, self.active) + # Yield twice so every gathered task gets a chance to be in flight + # together before any of them completes. + await asyncio.sleep(0) + await asyncio.sleep(0) + self.active -= 1 + return ModelResponse(parts=[TextPart(content="ok")]) + + +@pytest.mark.anyio +async def test_shared_semaphore_caps_concurrency_across_models() -> None: + inner = TrackingModel() + semaphore = asyncio.Semaphore(2) + fast = ConcurrencyLimitedModel(inner, semaphore) + smart = ConcurrencyLimitedModel(inner, semaphore) + params = ModelRequestParameters() + + await asyncio.gather( + *(fast.request([], None, params) for _ in range(5)), + *(smart.request([], None, params) for _ in range(5)), + ) + + assert inner.max_active == 2 diff --git a/engine/tests/test_stirling_api.py b/engine/tests/test_stirling_api.py index 4eaecf90f..3a6497174 100644 --- a/engine/tests/test_stirling_api.py +++ b/engine/tests/test_stirling_api.py @@ -214,6 +214,26 @@ def test_pdf_edit_route() -> None: assert response.json()["outcome"] == "cannot_do" +def test_routes_require_user_id_when_enforced() -> None: + """With STIRLING_REQUIRE_USER_ID on, an identity-less request is rejected + at the boundary before any handler runs; supplying X-User-Id is accepted. + The other route tests in this module run with the flag off and cover the + identity-less self-hosted path.""" + app.dependency_overrides[load_settings] = lambda: build_app_settings().model_copy(update={"require_user_id": True}) + try: + anonymous = client.post("/api/v1/pdf/edit", json={"userMessage": "rotate this"}) + identified = client.post( + "/api/v1/pdf/edit", + json={"userMessage": "rotate this"}, + headers={"X-User-Id": "alice"}, + ) + finally: + app.dependency_overrides[load_settings] = build_app_settings + + assert anonymous.status_code == 401 + assert identified.status_code == 200 + + def test_pdf_questions_route() -> None: response = client.post( "/api/v1/pdf/questions", diff --git a/engine/tests/test_stirling_contracts.py b/engine/tests/test_stirling_contracts.py index ccc057adb..7270955c2 100644 --- a/engine/tests/test_stirling_contracts.py +++ b/engine/tests/test_stirling_contracts.py @@ -77,17 +77,20 @@ def test_pdf_question_answer_defaults_evidence_list() -> None: def test_app_settings_accepts_model_configuration() -> None: from pathlib import Path - from stirling.config import RagBackend + from stirling.config import DocumentsBackend settings = AppSettings( smart_model_name="claude-sonnet-4-5-20250929", fast_model_name="claude-haiku-4-5-20251001", smart_model_max_tokens=8192, fast_model_max_tokens=2048, - rag_backend=RagBackend.SQLITE, + model_max_concurrency=32, + documents_backend=DocumentsBackend.SQLITE, rag_embedding_model="voyageai:voyage-4", - rag_store_path=Path(":memory:"), - rag_pgvector_dsn="", + documents_sqlite_path=Path(":memory:"), + documents_pgvector_dsn="", + documents_pgvector_pool_min_size=1, + documents_pgvector_pool_max_size=10, rag_chunk_size=512, rag_chunk_overlap=64, rag_default_top_k=5, @@ -98,6 +101,7 @@ def test_app_settings_accepts_model_configuration() -> None: chunked_reasoner_notes_char_budget=250_000, max_pages=200, max_characters=200_000, + require_user_id=False, posthog_enabled=False, posthog_api_key="", posthog_host="https://eu.i.posthog.com", diff --git a/engine/uv.lock b/engine/uv.lock index c8f79e637..00a8345c1 100644 --- a/engine/uv.lock +++ b/engine/uv.lock @@ -607,7 +607,7 @@ dependencies = [ { name = "opentelemetry-sdk" }, { name = "pgvector" }, { name = "posthog" }, - { name = "psycopg", extra = ["binary"] }, + { name = "psycopg", extra = ["binary", "pool"] }, { name = "pydantic" }, { name = "pydantic-ai" }, { name = "pydantic-ai-slim", extra = ["voyageai"] }, @@ -633,7 +633,7 @@ requires-dist = [ { name = "opentelemetry-sdk", specifier = ">=1.39.0" }, { name = "pgvector", specifier = ">=0.3.6" }, { name = "posthog", specifier = ">=3.0.0" }, - { name = "psycopg", extras = ["binary"], specifier = ">=3.2" }, + { name = "psycopg", extras = ["binary", "pool"], specifier = ">=3.2" }, { name = "pydantic", specifier = ">=2.0.0" }, { name = "pydantic-ai", specifier = ">=1.67.0" }, { name = "pydantic-ai-slim", extras = ["voyageai"], specifier = ">=1.67.0" }, @@ -2165,6 +2165,9 @@ wheels = [ binary = [ { name = "psycopg-binary", marker = "implementation_name != 'pypy'" }, ] +pool = [ + { name = "psycopg-pool" }, +] [[package]] name = "psycopg-binary" @@ -2195,6 +2198,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/98/5a/291d89f44d3820fffb7a04ebc8f3ef5dda4f542f44a5daea0c55a84abf45/psycopg_binary-3.3.3-cp314-cp314-win_amd64.whl", hash = "sha256:165f22ab5a9513a3d7425ffb7fcc7955ed8ccaeef6d37e369d6cc1dff1582383", size = 3652796, upload-time = "2026-02-18T16:52:14.02Z" }, ] +[[package]] +name = "psycopg-pool" +version = "3.3.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/90/82/7a23d26039827ecd4ebe93905651029ddd307c5182ad59296dfb6f67b528/psycopg_pool-3.3.1.tar.gz", hash = "sha256:b10b10b7a175d5cc1592147dc5b7eec8a9e0834eb3ed2c4a92c858e2f51eb63c", size = 31661, upload-time = "2026-05-01T23:31:59.809Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/37/ed/89c2c620af0e1660354cd8aabf9f5b21f911597ce22acb37c805d6c86bc8/psycopg_pool-3.3.1-py3-none-any.whl", hash = "sha256:2af5b432941c4c9ad5c87b3fa410aec910ec8f7c122855897983a06c45f2e4b5", size = 40023, upload-time = "2026-05-01T23:31:53.136Z" }, +] + [[package]] name = "py-key-value-aio" version = "0.4.4"