mirror of
https://github.com/VectifyAI/PageIndex.git
synced 2026-07-15 21:11:05 +02:00
fix: prompt-injection delimiter escape, legacy config coercion, gather resilience, true cross-thread concurrency bound
Addresses items 4-8 from the max-effort review of PR #272 (VectifyAI/PageIndex#272). - agent.py: wrap_with_doc_context() strips '<'/'>' from doc_name/doc_description (untrusted: unsanitized filename / LLM-generated from document content) before inserting them into the <docs>...</docs> block, so embedded content can never form a literal </docs> that closes the delimiter early and escapes the untrusted-data boundary SCOPED_SYSTEM_PROMPT relies on. Deterministic per-field transform, doesn't touch the (cacheable) static system prompt. - ConfigLoader.load() (legacy 0.2.x compat) now routes merged overrides through IndexConfig before returning, so a legacy 'no' string gets pydantic's bool coercion instead of surviving as a truthy non-empty string — page_index_main's bare `if opt.if_add_node_summary:` checks (changed from `== 'yes'` elsewhere in this PR) were silently inverting caller intent and firing unwanted billed LLM calls. - verify_toc, process_large_node_recursively, tree_parser, generate_summaries_for_structure, generate_summaries_for_structure_md: added return_exceptions=True to their asyncio.gather calls (llm_completion/ llm_acompletion raise RuntimeError on retry exhaustion, added earlier in this PR), each with a degrade path matching the pattern already used by sibling hardened gathers in the same files. One transient LLM failure no longer aborts the whole document's indexing. - _llm_semaphore is now a true process-wide ceiling (threading.Semaphore, shared across every thread/event loop) instead of one asyncio.Semaphore per event loop -- concurrently indexing N documents on N threads no longer multiplies the effective cap by N. A max_concurrency_scope() override is layered as a second, nested, per-loop restriction that can only tighten the effective cap within the ceiling, never widen past it. - set_llm_params() mutated a bare process-wide dict with no per-call isolation, unlike max_concurrency which already had ContextVar scoping. Added llm_params_scope() (mirrors max_concurrency_scope) + IndexConfig.llm_params, wired into build_index() the same way max_concurrency already was, so concurrent indexing jobs with different llm kwargs don't leak into each other. Adds regression tests for all five. Full suite: 221 passed, 2 skipped (one pre-existing, unrelated flaky cloud-streaming test intermittently fails on rerun; confirmed independent of this change). Claude-Session: https://claude.ai/code/session_01Kx5DgKbhK1N8autqXH8SmS
This commit is contained in:
parent
623ce927f1
commit
b9d021916f
11 changed files with 408 additions and 55 deletions
|
|
@ -29,6 +29,11 @@ class IndexConfig(BaseModel):
|
|||
# default (get_max_concurrency(), overridable via PAGEINDEX_MAX_CONCURRENCY).
|
||||
# An explicit value here wins for this client.
|
||||
max_concurrency: int | None = None
|
||||
# Per-call litellm completion kwargs for this client's indexing calls only
|
||||
# (e.g. {"temperature": 1}). None = use the process-wide defaults
|
||||
# (get_llm_params(), overridable via set_llm_params()). Scoped via
|
||||
# llm_params_scope so it doesn't leak into other concurrent indexing calls.
|
||||
llm_params: dict | None = None
|
||||
|
||||
@field_validator("max_concurrency", mode="before")
|
||||
@classmethod
|
||||
|
|
@ -57,6 +62,16 @@ def _env_drop_params_default() -> bool:
|
|||
# PAGEINDEX_DROP_PARAMS env shortcut.
|
||||
_LLM_PARAMS: dict = {"temperature": 0, "drop_params": _env_drop_params_default()}
|
||||
|
||||
# Per-call override, isolated per thread / async context — mirrors
|
||||
# _MAX_CONCURRENCY_OVERRIDE below. Without this, set_llm_params() is the only
|
||||
# way to change llm params and it mutates the process-wide dict directly, so
|
||||
# concurrently indexing two documents with different llm_params_scope() would
|
||||
# otherwise leak one caller's settings (e.g. temperature) into the other's
|
||||
# in-flight calls. None = no override -> fall back to the process-wide _LLM_PARAMS.
|
||||
_LLM_PARAMS_OVERRIDE: ContextVar[dict | None] = ContextVar(
|
||||
"pageindex_llm_params_override", default=None
|
||||
)
|
||||
|
||||
# Structural kwargs PageIndex always supplies itself — not overridable here.
|
||||
_RESERVED_LLM_PARAMS = ("model", "messages")
|
||||
|
||||
|
|
@ -121,6 +136,16 @@ def get_max_concurrency() -> int:
|
|||
return override if override is not None else _MAX_CONCURRENCY
|
||||
|
||||
|
||||
def _process_wide_max_concurrency() -> int:
|
||||
"""The process-wide default cap, ignoring any active max_concurrency_scope
|
||||
override. This is the TRUE ceiling shared across every thread/event loop in
|
||||
the process (see index/utils.py's _llm_semaphore) — a per-call override may
|
||||
only narrow the effective cap within that ceiling, never widen it, so the
|
||||
ceiling itself must not vary with a context-local override.
|
||||
"""
|
||||
return _MAX_CONCURRENCY
|
||||
|
||||
|
||||
def set_max_concurrency(value: int) -> None:
|
||||
"""Set the process-wide default cap on concurrent in-flight LLM calls."""
|
||||
global _MAX_CONCURRENCY
|
||||
|
|
@ -147,19 +172,53 @@ def max_concurrency_scope(value: int | None):
|
|||
|
||||
|
||||
def get_llm_params() -> dict:
|
||||
"""Return a copy of the per-call kwargs PageIndex passes to litellm."""
|
||||
return dict(_LLM_PARAMS)
|
||||
"""Return a copy of the effective per-call kwargs PageIndex passes to litellm.
|
||||
|
||||
A per-index override (llm_params_scope) is merged over the process-wide
|
||||
defaults for the current context; otherwise just the process-wide defaults
|
||||
apply.
|
||||
"""
|
||||
params = dict(_LLM_PARAMS)
|
||||
override = _LLM_PARAMS_OVERRIDE.get()
|
||||
if override:
|
||||
params.update(override)
|
||||
return params
|
||||
|
||||
|
||||
def set_llm_params(**kwargs) -> None:
|
||||
"""Override or extend the litellm completion kwargs PageIndex sends per call.
|
||||
"""Override or extend the process-wide default litellm completion kwargs.
|
||||
|
||||
e.g. ``set_llm_params(drop_params=False, temperature=1, num_retries=5)``.
|
||||
Applied per call; never writes litellm's global state, so it can't leak into
|
||||
other litellm users in the same process. ``model`` / ``messages`` are
|
||||
reserved (PageIndex supplies them) and rejected.
|
||||
Never writes litellm's global state, so it can't leak into other litellm
|
||||
users in the same process — but it DOES mutate PageIndex's own process-wide
|
||||
default, so it affects every concurrent caller in this process. For a
|
||||
one-off override scoped to a single indexing call, use ``llm_params_scope``
|
||||
instead. ``model`` / ``messages`` are reserved (PageIndex supplies them) and
|
||||
rejected.
|
||||
"""
|
||||
reserved = [k for k in kwargs if k in _RESERVED_LLM_PARAMS]
|
||||
if reserved:
|
||||
raise ValueError(f"cannot override reserved litellm kwargs: {reserved}")
|
||||
_LLM_PARAMS.update(kwargs)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def llm_params_scope(overrides: dict | None):
|
||||
"""Scope a per-index override of the litellm completion kwargs to the
|
||||
current context.
|
||||
|
||||
``overrides=None`` (or ``{}``) means "no override" (fall back to the
|
||||
process-wide defaults). Isolated per thread / async context and reset on
|
||||
exit, so concurrent indexing doesn't leak one call's kwargs into another's
|
||||
and a one-off override never becomes the sticky new process default —
|
||||
mirrors ``max_concurrency_scope``.
|
||||
"""
|
||||
if overrides:
|
||||
reserved = [k for k in overrides if k in _RESERVED_LLM_PARAMS]
|
||||
if reserved:
|
||||
raise ValueError(f"cannot override reserved litellm kwargs: {reserved}")
|
||||
token = _LLM_PARAMS_OVERRIDE.set(overrides or None)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
_LLM_PARAMS_OVERRIDE.reset(token)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue