mirror of
https://github.com/arsvendg/Stirling-PDF.git
synced 2026-07-14 10:34:06 +02:00
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:
+23
-7
@@ -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
|
||||
|
||||
@@ -26,6 +26,7 @@ COPY .taskfiles/ ./.taskfiles/
|
||||
|
||||
ENV PATH="/app/engine/.venv/bin:$PATH"
|
||||
ENV PYTHONUNBUFFERED=1
|
||||
ENV STIRLING_ENGINE_WORKERS=4
|
||||
|
||||
EXPOSE 5001
|
||||
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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",
|
||||
)
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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),
|
||||
)
|
||||
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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
|
||||
@@ -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",
|
||||
|
||||
@@ -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",
|
||||
|
||||
Generated
+17
-2
@@ -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"
|
||||
|
||||
Reference in New Issue
Block a user