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
This commit is contained in:
James Brunton
2026-06-11 16:31:35 +01:00
committed by GitHub
parent 606964ee52
commit d52c7ced7c
17 changed files with 315 additions and 90 deletions
+1 -1
View File
@@ -43,7 +43,7 @@ tasks:
env: env:
PYTHONUNBUFFERED: "1" PYTHONUNBUFFERED: "1"
cmds: 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: dev:
desc: "Start engine dev server with hot reload" desc: "Start engine dev server with hot reload"
+23 -7
View File
@@ -14,19 +14,30 @@ STIRLING_FAST_MODEL=anthropic:claude-haiku-4-5
STIRLING_SMART_MODEL_MAX_TOKENS=8192 STIRLING_SMART_MODEL_MAX_TOKENS=8192
STIRLING_FAST_MODEL_MAX_TOKENS=2048 STIRLING_FAST_MODEL_MAX_TOKENS=2048
# RAG Configuration — retrieval-augmented generation is always on. # Process-wide cap on concurrent model API calls, shared by both model tiers.
# Embedding provider credentials are handled natively (e.g. VOYAGE_API_KEY for VoyageAI). # Per-request fan-outs (chunked reasoner workers, contradiction detection) are
STIRLING_RAG_EMBEDDING_MODEL=voyageai:voyage-4 # bounded per request; this bounds their product across concurrent requests.
STIRLING_MODEL_MAX_CONCURRENCY=32
# Vector store backend: "sqlite" (embedded) or "pgvector" (external Postgres). # Document store: the one database holding vector chunks, ordered page text,
STIRLING_RAG_BACKEND=sqlite # 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). # 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. # Postgres DSN for pgvector (used when backend=pgvector). Leave empty when backend=sqlite.
# Example: postgresql://user:password@host:5432/dbname # 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_SIZE=512
STIRLING_RAG_CHUNK_OVERLAP=64 STIRLING_RAG_CHUNK_OVERLAP=64
@@ -57,6 +68,11 @@ STIRLING_CHUNKED_REASONER_NOTES_CHAR_BUDGET=250000
STIRLING_MAX_PAGES=200 STIRLING_MAX_PAGES=200
STIRLING_MAX_CHARACTERS=200000 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. # PostHog analytics. Set STIRLING_POSTHOG_ENABLED=true and provide an API key to enable.
STIRLING_POSTHOG_ENABLED=false STIRLING_POSTHOG_ENABLED=false
STIRLING_POSTHOG_API_KEY=phc_VOdeYnlevc2T63m3myFGjeBlRcIusRgmhfx6XL5a1iz STIRLING_POSTHOG_API_KEY=phc_VOdeYnlevc2T63m3myFGjeBlRcIusRgmhfx6XL5a1iz
+1
View File
@@ -26,6 +26,7 @@ COPY .taskfiles/ ./.taskfiles/
ENV PATH="/app/engine/.venv/bin:$PATH" ENV PATH="/app/engine/.venv/bin:$PATH"
ENV PYTHONUNBUFFERED=1 ENV PYTHONUNBUFFERED=1
ENV STIRLING_ENGINE_WORKERS=4
EXPOSE 5001 EXPOSE 5001
+1 -1
View File
@@ -6,7 +6,7 @@ requires-python = ">=3.13"
dependencies = [ dependencies = [
"fastapi>=0.116.0", "fastapi>=0.116.0",
"pgvector>=0.3.6", "pgvector>=0.3.6",
"psycopg[binary]>=3.2", "psycopg[binary,pool]>=3.2",
"pydantic>=2.0.0", "pydantic>=2.0.0",
"pydantic-ai>=1.67.0", "pydantic-ai>=1.67.0",
"pydantic-ai-slim[voyageai]>=1.67.0", "pydantic-ai-slim[voyageai]>=1.67.0",
+13 -9
View File
@@ -18,6 +18,7 @@ from stirling.agents import (
) )
from stirling.agents.ledger import MathAuditorAgent from stirling.agents.ledger import MathAuditorAgent
from stirling.agents.pdf_comment import PdfCommentAgent 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.engine_auth import EngineSharedSecretMiddleware
from stirling.api.middleware import UserIdMiddleware from stirling.api.middleware import UserIdMiddleware
from stirling.api.routes import ( 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 = FastAPI(title="Stirling AI Engine", lifespan=lifespan, version="0.1.0")
app.add_middleware(UserIdMiddleware) app.add_middleware(UserIdMiddleware)
app.add_middleware(EngineSharedSecretMiddleware) app.add_middleware(EngineSharedSecretMiddleware)
app.include_router(orchestrator_router) # Every router gets the same configurable identity gate; /health stays open
app.include_router(pdf_edit_router) # for liveness probes. See enforce_required_user_id for the policy.
app.include_router(pdf_question_router) _user_gate = [Depends(enforce_required_user_id)]
app.include_router(agent_draft_router) app.include_router(orchestrator_router, dependencies=_user_gate)
app.include_router(execution_router) app.include_router(pdf_edit_router, dependencies=_user_gate)
app.include_router(document_router) app.include_router(pdf_question_router, dependencies=_user_gate)
app.include_router(ledger_router) app.include_router(agent_draft_router, dependencies=_user_gate)
app.include_router(pdf_comments_router) app.include_router(execution_router, dependencies=_user_gate)
app.include_router(agent_capabilities_router) 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) @app.get("/health", response_model=HealthResponse)
+17 -1
View File
@@ -1,6 +1,8 @@
from __future__ import annotations 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 ( from stirling.agents import (
ExecutionPlanningAgent, ExecutionPlanningAgent,
@@ -11,6 +13,7 @@ from stirling.agents import (
) )
from stirling.agents.ledger import MathAuditorAgent from stirling.agents.ledger import MathAuditorAgent
from stirling.agents.pdf_comment import PdfCommentAgent from stirling.agents.pdf_comment import PdfCommentAgent
from stirling.config import AppSettings, load_settings
from stirling.documents import DocumentService from stirling.documents import DocumentService
from stirling.models import UserId from stirling.models import UserId
from stirling.services import AppRuntime, current_user_id from stirling.services import AppRuntime, current_user_id
@@ -67,3 +70,16 @@ def require_user_id() -> UserId:
detail="X-User-Id header is required", detail="X-User-Id header is required",
) )
return user_id 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",
)
+2 -2
View File
@@ -1,10 +1,10 @@
"""Configuration models and loaders for the Stirling AI service.""" """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__ = [ __all__ = [
"ENGINE_ROOT", "ENGINE_ROOT",
"AppSettings", "AppSettings",
"RagBackend", "DocumentsBackend",
"load_settings", "load_settings",
] ]
+21 -5
View File
@@ -15,7 +15,7 @@ ENV_FILE = ENGINE_ROOT / ".env"
ENV_LOCAL_FILE = ENGINE_ROOT / ".env.local" ENV_LOCAL_FILE = ENGINE_ROOT / ".env.local"
class RagBackend(StrEnum): class DocumentsBackend(StrEnum):
SQLITE = "sqlite" SQLITE = "sqlite"
PGVECTOR = "pgvector" PGVECTOR = "pgvector"
@@ -27,12 +27,22 @@ class AppSettings(BaseSettings):
fast_model_name: str = Field(validation_alias="STIRLING_FAST_MODEL") fast_model_name: str = Field(validation_alias="STIRLING_FAST_MODEL")
smart_model_max_tokens: int = Field(validation_alias="STIRLING_SMART_MODEL_MAX_TOKENS") 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") 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. # Document store: the one database holding vector chunks, ordered page
rag_backend: RagBackend = Field(validation_alias="STIRLING_RAG_BACKEND") # 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_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_size: int = Field(validation_alias="STIRLING_RAG_CHUNK_SIZE")
rag_chunk_overlap: int = Field(validation_alias="STIRLING_RAG_CHUNK_OVERLAP") rag_chunk_overlap: int = Field(validation_alias="STIRLING_RAG_CHUNK_OVERLAP")
rag_default_top_k: int = Field(validation_alias="STIRLING_RAG_TOP_K") 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_pages: int = Field(validation_alias="STIRLING_MAX_PAGES")
max_characters: int = Field(validation_alias="STIRLING_MAX_CHARACTERS") 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_level: str = Field(default="INFO", validation_alias="STIRLING_LOG_LEVEL")
log_file: str = Field(default="", validation_alias="STIRLING_LOG_FILE") log_file: str = Field(default="", validation_alias="STIRLING_LOG_FILE")
# When true, raises httpx + httpcore logger levels so every outgoing # When true, raises httpx + httpcore logger levels so every outgoing
+4 -4
View File
@@ -54,10 +54,10 @@ everything = RagCapability(runtime.documents)
Non-secret defaults live in the committed `engine/.env`: 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_EMBEDDING_MODEL=voyageai:voyage-4
STIRLING_RAG_STORE_PATH=data/rag.db # used when backend=sqlite STIRLING_DOCUMENTS_SQLITE_PATH=data/rag.db # used when backend=sqlite
STIRLING_RAG_PGVECTOR_DSN= # used when backend=pgvector STIRLING_DOCUMENTS_PGVECTOR_DSN= # used when backend=pgvector
STIRLING_RAG_CHUNK_SIZE=512 STIRLING_RAG_CHUNK_SIZE=512
STIRLING_RAG_CHUNK_OVERLAP=64 STIRLING_RAG_CHUNK_OVERLAP=64
STIRLING_RAG_TOP_K=5 STIRLING_RAG_TOP_K=5
@@ -76,7 +76,7 @@ VOYAGE_API_KEY=your-key
and self-hosted deployments. and self-hosted deployments.
**`pgvector`** - External PostgreSQL with the `vector` extension. Point **`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 Both backends implement the same `DocumentStore` interface, so agents and the
service work identically regardless of which you pick. service work identically regardless of which you pick.
+59 -36
View File
@@ -1,10 +1,12 @@
from __future__ import annotations from __future__ import annotations
import asyncio
import json import json
from datetime import datetime from datetime import datetime
import psycopg import psycopg
from pgvector.psycopg import register_vector_async from pgvector.psycopg import register_vector_async
from psycopg_pool import AsyncConnectionPool
from stirling.contracts.documents import Page, PageRange from stirling.contracts.documents import Page, PageRange
from stirling.documents.store import Document, DocumentStore, SearchResult, StoredPage from stirling.documents.store import Document, DocumentStore, SearchResult, StoredPage
@@ -13,12 +15,20 @@ from stirling.models import OwnerId, PrincipalId
_READ_PERMISSION = "read" _READ_PERMISSION = "read"
async def _register_vector(conn: psycopg.AsyncConnection) -> None:
await register_vector_async(conn)
class PgVectorStore(DocumentStore): class PgVectorStore(DocumentStore):
"""PostgreSQL + pgvector backed store, scoped by owner with ACL-gated reads. """PostgreSQL + pgvector backed store, scoped by owner with ACL-gated reads.
Connects to an external Postgres instance (DSN provided via config) and uses the 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. `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: Owned tables:
* ``documents_meta`` - parent row, one per ``(collection, owner_id)``. * ``documents_meta`` - parent row, one per ``(collection, owner_id)``.
@@ -31,21 +41,36 @@ class PgVectorStore(DocumentStore):
Writes are owner-scoped; reads are ACL-scoped. 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: 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._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._initialized = False
self._init_lock = asyncio.Lock()
async def _connect(self) -> psycopg.AsyncConnection: async def _ensure_ready(self) -> None:
conn = await psycopg.AsyncConnection.connect(self._dsn) """Create the schema and open the pool, exactly once across concurrent callers."""
await register_vector_async(conn)
return conn
async def _ensure_schema(self) -> None:
if self._initialized: if self._initialized:
return 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: async with conn.cursor() as cur:
await cur.execute("CREATE EXTENSION IF NOT EXISTS vector") await cur.execute("CREATE EXTENSION IF NOT EXISTS vector")
await cur.execute( 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)" "CREATE INDEX IF NOT EXISTS idx_acl_principal_permission ON document_acl(principal_id, permission)"
) )
await conn.commit() await conn.commit()
self._initialized = True
# ── lifecycle of the (collection, owner_id) row ──────────────────────── # ── lifecycle of the (collection, owner_id) row ────────────────────────
@@ -129,8 +153,8 @@ class PgVectorStore(DocumentStore):
owner_id: OwnerId, owner_id: OwnerId,
expires_at: datetime | None, expires_at: datetime | None,
) -> None: ) -> None:
await self._ensure_schema() await self._ensure_ready()
async with await self._connect() as conn: async with self._pool.connection() as conn:
async with conn.cursor() as cur: async with conn.cursor() as cur:
await cur.execute( await cur.execute(
""" """
@@ -145,8 +169,8 @@ class PgVectorStore(DocumentStore):
await conn.commit() await conn.commit()
async def purge_owner(self, owner_id: OwnerId) -> int: async def purge_owner(self, owner_id: OwnerId) -> int:
await self._ensure_schema() await self._ensure_ready()
async with await self._connect() as conn: async with self._pool.connection() as conn:
async with conn.cursor() as cur: async with conn.cursor() as cur:
await cur.execute( await cur.execute(
"DELETE FROM documents_meta WHERE owner_id = %s", "DELETE FROM documents_meta WHERE owner_id = %s",
@@ -157,8 +181,8 @@ class PgVectorStore(DocumentStore):
return deleted return deleted
async def reap_expired(self) -> int: async def reap_expired(self) -> int:
await self._ensure_schema() await self._ensure_ready()
async with await self._connect() as conn: async with self._pool.connection() as conn:
async with conn.cursor() as cur: async with conn.cursor() as cur:
await cur.execute("DELETE FROM documents_meta WHERE expires_at IS NOT NULL AND expires_at < NOW()") await cur.execute("DELETE FROM documents_meta WHERE expires_at IS NOT NULL AND expires_at < NOW()")
deleted = cur.rowcount deleted = cur.rowcount
@@ -166,8 +190,8 @@ class PgVectorStore(DocumentStore):
return deleted return deleted
async def delete_collection(self, collection: str, owner_id: OwnerId) -> bool: async def delete_collection(self, collection: str, owner_id: OwnerId) -> bool:
await self._ensure_schema() await self._ensure_ready()
async with await self._connect() as conn: async with self._pool.connection() as conn:
async with conn.cursor() as cur: async with conn.cursor() as cur:
# Cascade FKs handle rag_documents, document_pages, document_acl. # Cascade FKs handle rag_documents, document_pages, document_acl.
await cur.execute( await cur.execute(
@@ -192,8 +216,8 @@ class PgVectorStore(DocumentStore):
if not documents: if not documents:
return return
await self._ensure_schema() await self._ensure_ready()
async with await self._connect() as conn: async with self._pool.connection() as conn:
async with conn.cursor() as cur: async with conn.cursor() as cur:
for doc, emb in zip(documents, embeddings): for doc, emb in zip(documents, embeddings):
await cur.execute( await cur.execute(
@@ -211,8 +235,8 @@ class PgVectorStore(DocumentStore):
await conn.commit() await conn.commit()
async def add_pages(self, collection: str, pages: list[StoredPage], owner_id: OwnerId) -> None: async def add_pages(self, collection: str, pages: list[StoredPage], owner_id: OwnerId) -> None:
await self._ensure_schema() await self._ensure_ready()
async with await self._connect() as conn: async with self._pool.connection() as conn:
async with conn.cursor() as cur: async with conn.cursor() as cur:
await cur.execute( await cur.execute(
"DELETE FROM document_pages WHERE collection = %s AND owner_id = %s", "DELETE FROM document_pages WHERE collection = %s AND owner_id = %s",
@@ -238,8 +262,8 @@ class PgVectorStore(DocumentStore):
) -> None: ) -> None:
if not principals: if not principals:
return return
await self._ensure_schema() await self._ensure_ready()
async with await self._connect() as conn: async with self._pool.connection() as conn:
async with conn.cursor() as cur: async with conn.cursor() as cur:
await cur.executemany( await cur.executemany(
""" """
@@ -257,8 +281,8 @@ class PgVectorStore(DocumentStore):
owner_id: OwnerId, owner_id: OwnerId,
principal: PrincipalId, principal: PrincipalId,
) -> None: ) -> None:
await self._ensure_schema() await self._ensure_ready()
async with await self._connect() as conn: async with self._pool.connection() as conn:
async with conn.cursor() as cur: async with conn.cursor() as cur:
await cur.execute( await cur.execute(
"DELETE FROM document_acl WHERE collection = %s AND owner_id = %s AND principal_id = %s", "DELETE FROM document_acl WHERE collection = %s AND owner_id = %s AND principal_id = %s",
@@ -277,8 +301,8 @@ class PgVectorStore(DocumentStore):
) -> list[SearchResult]: ) -> list[SearchResult]:
if not principals: if not principals:
return [] return []
await self._ensure_schema() await self._ensure_ready()
async with await self._connect() as conn: async with self._pool.connection() as conn:
async with conn.cursor() as cur: async with conn.cursor() as cur:
owner_id = await self._readable_owner_for(cur, collection, principals) owner_id = await self._readable_owner_for(cur, collection, principals)
if owner_id is None: if owner_id is None:
@@ -332,8 +356,8 @@ class PgVectorStore(DocumentStore):
) -> list[Page]: ) -> list[Page]:
if not principals: if not principals:
return [] return []
await self._ensure_schema() await self._ensure_ready()
async with await self._connect() as conn: async with self._pool.connection() as conn:
async with conn.cursor() as cur: async with conn.cursor() as cur:
owner_id = await self._readable_owner_for(cur, collection, principals) owner_id = await self._readable_owner_for(cur, collection, principals)
if owner_id is None: if owner_id is None:
@@ -358,8 +382,8 @@ class PgVectorStore(DocumentStore):
async def has_collection(self, collection: str, principals: list[PrincipalId]) -> bool: async def has_collection(self, collection: str, principals: list[PrincipalId]) -> bool:
if not principals: if not principals:
return False return False
await self._ensure_schema() await self._ensure_ready()
async with await self._connect() as conn: async with self._pool.connection() as conn:
async with conn.cursor() as cur: async with conn.cursor() as cur:
owner_id = await self._readable_owner_for(cur, collection, principals) owner_id = await self._readable_owner_for(cur, collection, principals)
return owner_id is not None return owner_id is not None
@@ -367,8 +391,8 @@ class PgVectorStore(DocumentStore):
async def list_collections(self, principals: list[PrincipalId]) -> list[str]: async def list_collections(self, principals: list[PrincipalId]) -> list[str]:
if not principals: if not principals:
return [] return []
await self._ensure_schema() await self._ensure_ready()
async with await self._connect() as conn: async with self._pool.connection() as conn:
async with conn.cursor() as cur: async with conn.cursor() as cur:
await cur.execute( await cur.execute(
""" """
@@ -383,5 +407,4 @@ class PgVectorStore(DocumentStore):
return [r[0] for r in rows] return [r[0] for r in rows]
async def close(self) -> None: async def close(self) -> None:
# Connections are opened and closed per call, so nothing persistent to release. await self._pool.close()
return None
+54 -10
View File
@@ -1,16 +1,22 @@
from __future__ import annotations from __future__ import annotations
import asyncio
import logging import logging
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from dataclasses import dataclass from dataclasses import dataclass
from typing import assert_never from typing import Any, assert_never
import httpx 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.anthropic import AnthropicModel
from pydantic_ai.models.wrapper import WrapperModel
from pydantic_ai.providers.anthropic import AnthropicProvider from pydantic_ai.providers.anthropic import AnthropicProvider
from pydantic_ai.settings import ModelSettings 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 ( from stirling.documents import (
DocumentService, DocumentService,
DocumentStore, DocumentStore,
@@ -40,6 +46,37 @@ def _build_anthropic_http_client() -> httpx.AsyncClient:
return httpx.AsyncClient(limits=httpx.Limits(max_keepalive_connections=0)) 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) @dataclass(frozen=True)
class AppRuntime: class AppRuntime:
settings: AppSettings settings: AppSettings
@@ -74,17 +111,21 @@ def validate_structured_output_support(model: Model, model_name: str) -> None:
def _build_document_store(settings: AppSettings) -> DocumentStore: def _build_document_store(settings: AppSettings) -> DocumentStore:
"""Build the configured document store backend.""" """Build the configured document store backend."""
if settings.rag_backend == RagBackend.SQLITE: if settings.documents_backend == DocumentsBackend.SQLITE:
store_path = settings.rag_store_path store_path = settings.documents_sqlite_path
# Treat ":memory:" as a special in-process token; otherwise resolve against the engine root. # 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(): if str(store_path) != ":memory:" and not store_path.is_absolute():
store_path = ENGINE_ROOT / store_path store_path = ENGINE_ROOT / store_path
logger.info("Document store backend=sqlite, db_path=%s", store_path) logger.info("Document store backend=sqlite, db_path=%s", store_path)
return SqliteVecStore(db_path=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>") logger.info("Document store backend=pgvector, dsn=<configured>")
return PgVectorStore(dsn=settings.rag_pgvector_dsn) return PgVectorStore(
assert_never(settings.rag_backend) 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: 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(fast_model, settings.fast_model_name)
validate_structured_output_support(smart_model, settings.smart_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( return AppRuntime(
settings=settings, settings=settings,
fast_model=fast_model, fast_model=ConcurrencyLimitedModel(fast_model, model_semaphore),
smart_model=smart_model, smart_model=ConcurrencyLimitedModel(smart_model, model_semaphore),
documents=_build_documents(settings), documents=_build_documents(settings),
) )
+8 -4
View File
@@ -5,7 +5,7 @@ from pathlib import Path
import pytest 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 import build_runtime
from stirling.services.runtime import AppRuntime from stirling.services.runtime import AppRuntime
@@ -23,10 +23,13 @@ def build_app_settings() -> AppSettings:
fast_model_name="test", fast_model_name="test",
smart_model_max_tokens=8192, smart_model_max_tokens=8192,
fast_model_max_tokens=2048, fast_model_max_tokens=2048,
rag_backend=RagBackend.SQLITE, model_max_concurrency=32,
documents_backend=DocumentsBackend.SQLITE,
rag_embedding_model="voyageai:voyage-4", rag_embedding_model="voyageai:voyage-4",
rag_store_path=Path(":memory:"), documents_sqlite_path=Path(":memory:"),
rag_pgvector_dsn="", documents_pgvector_dsn="",
documents_pgvector_pool_min_size=1,
documents_pgvector_pool_max_size=10,
rag_chunk_size=512, rag_chunk_size=512,
rag_chunk_overlap=64, rag_chunk_overlap=64,
rag_default_top_k=5, rag_default_top_k=5,
@@ -41,6 +44,7 @@ def build_app_settings() -> AppSettings:
contradiction_canonicaliser_batch_size=500, contradiction_canonicaliser_batch_size=500,
max_pages=200, max_pages=200,
max_characters=200_000, max_characters=200_000,
require_user_id=False,
posthog_enabled=False, posthog_enabled=False,
posthog_api_key="", posthog_api_key="",
posthog_host="https://eu.i.posthog.com", posthog_host="https://eu.i.posthog.com",
+8 -4
View File
@@ -16,7 +16,7 @@ from fastapi.testclient import TestClient
from stirling.api import app from stirling.api import app
from stirling.api.dependencies import get_pdf_comment_agent 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 ( from stirling.contracts.pdf_comments import (
PdfCommentInstruction, PdfCommentInstruction,
PdfCommentRequest, PdfCommentRequest,
@@ -35,10 +35,13 @@ class StubSettingsProvider:
fast_model_name="test", fast_model_name="test",
smart_model_max_tokens=8192, smart_model_max_tokens=8192,
fast_model_max_tokens=2048, fast_model_max_tokens=2048,
rag_backend=RagBackend.SQLITE, model_max_concurrency=32,
documents_backend=DocumentsBackend.SQLITE,
rag_embedding_model="test-embed", rag_embedding_model="test-embed",
rag_store_path=Path(":memory:"), documents_sqlite_path=Path(":memory:"),
rag_pgvector_dsn="", documents_pgvector_dsn="",
documents_pgvector_pool_min_size=1,
documents_pgvector_pool_max_size=10,
rag_chunk_size=512, rag_chunk_size=512,
rag_chunk_overlap=64, rag_chunk_overlap=64,
rag_default_top_k=5, rag_default_top_k=5,
@@ -49,6 +52,7 @@ class StubSettingsProvider:
chunked_reasoner_notes_char_budget=250_000, chunked_reasoner_notes_char_budget=250_000,
max_pages=100, max_pages=100,
max_characters=100_000, max_characters=100_000,
require_user_id=False,
posthog_enabled=False, posthog_enabled=False,
posthog_api_key="", posthog_api_key="",
posthog_host="https://eu.i.posthog.com", posthog_host="https://eu.i.posthog.com",
+58
View File
@@ -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
+20
View File
@@ -214,6 +214,26 @@ def test_pdf_edit_route() -> None:
assert response.json()["outcome"] == "cannot_do" 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: def test_pdf_questions_route() -> None:
response = client.post( response = client.post(
"/api/v1/pdf/questions", "/api/v1/pdf/questions",
+8 -4
View File
@@ -77,17 +77,20 @@ def test_pdf_question_answer_defaults_evidence_list() -> None:
def test_app_settings_accepts_model_configuration() -> None: def test_app_settings_accepts_model_configuration() -> None:
from pathlib import Path from pathlib import Path
from stirling.config import RagBackend from stirling.config import DocumentsBackend
settings = AppSettings( settings = AppSettings(
smart_model_name="claude-sonnet-4-5-20250929", smart_model_name="claude-sonnet-4-5-20250929",
fast_model_name="claude-haiku-4-5-20251001", fast_model_name="claude-haiku-4-5-20251001",
smart_model_max_tokens=8192, smart_model_max_tokens=8192,
fast_model_max_tokens=2048, fast_model_max_tokens=2048,
rag_backend=RagBackend.SQLITE, model_max_concurrency=32,
documents_backend=DocumentsBackend.SQLITE,
rag_embedding_model="voyageai:voyage-4", rag_embedding_model="voyageai:voyage-4",
rag_store_path=Path(":memory:"), documents_sqlite_path=Path(":memory:"),
rag_pgvector_dsn="", documents_pgvector_dsn="",
documents_pgvector_pool_min_size=1,
documents_pgvector_pool_max_size=10,
rag_chunk_size=512, rag_chunk_size=512,
rag_chunk_overlap=64, rag_chunk_overlap=64,
rag_default_top_k=5, rag_default_top_k=5,
@@ -98,6 +101,7 @@ def test_app_settings_accepts_model_configuration() -> None:
chunked_reasoner_notes_char_budget=250_000, chunked_reasoner_notes_char_budget=250_000,
max_pages=200, max_pages=200,
max_characters=200_000, max_characters=200_000,
require_user_id=False,
posthog_enabled=False, posthog_enabled=False,
posthog_api_key="", posthog_api_key="",
posthog_host="https://eu.i.posthog.com", posthog_host="https://eu.i.posthog.com",
+17 -2
View File
@@ -607,7 +607,7 @@ dependencies = [
{ name = "opentelemetry-sdk" }, { name = "opentelemetry-sdk" },
{ name = "pgvector" }, { name = "pgvector" },
{ name = "posthog" }, { name = "posthog" },
{ name = "psycopg", extra = ["binary"] }, { name = "psycopg", extra = ["binary", "pool"] },
{ name = "pydantic" }, { name = "pydantic" },
{ name = "pydantic-ai" }, { name = "pydantic-ai" },
{ name = "pydantic-ai-slim", extra = ["voyageai"] }, { name = "pydantic-ai-slim", extra = ["voyageai"] },
@@ -633,7 +633,7 @@ requires-dist = [
{ name = "opentelemetry-sdk", specifier = ">=1.39.0" }, { name = "opentelemetry-sdk", specifier = ">=1.39.0" },
{ name = "pgvector", specifier = ">=0.3.6" }, { name = "pgvector", specifier = ">=0.3.6" },
{ name = "posthog", specifier = ">=3.0.0" }, { 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", specifier = ">=2.0.0" },
{ name = "pydantic-ai", specifier = ">=1.67.0" }, { name = "pydantic-ai", specifier = ">=1.67.0" },
{ name = "pydantic-ai-slim", extras = ["voyageai"], specifier = ">=1.67.0" }, { name = "pydantic-ai-slim", extras = ["voyageai"], specifier = ">=1.67.0" },
@@ -2165,6 +2165,9 @@ wheels = [
binary = [ binary = [
{ name = "psycopg-binary", marker = "implementation_name != 'pypy'" }, { name = "psycopg-binary", marker = "implementation_name != 'pypy'" },
] ]
pool = [
{ name = "psycopg-pool" },
]
[[package]] [[package]]
name = "psycopg-binary" 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" }, { 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]] [[package]]
name = "py-key-value-aio" name = "py-key-value-aio"
version = "0.4.4" version = "0.4.4"