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
|
|
@ -8,8 +8,11 @@ import pytest
|
|||
from pageindex.config import (
|
||||
IndexConfig,
|
||||
_env_max_concurrency_default,
|
||||
get_llm_params,
|
||||
get_max_concurrency,
|
||||
llm_params_scope,
|
||||
max_concurrency_scope,
|
||||
set_llm_params,
|
||||
set_max_concurrency,
|
||||
)
|
||||
from pageindex.index.utils import _llm_semaphore, llm_acompletion
|
||||
|
|
@ -23,6 +26,14 @@ def _restore_max_concurrency():
|
|||
set_max_concurrency(prev)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _restore_llm_params():
|
||||
"""Keep tests isolated — llm params are a module global too."""
|
||||
prev = get_llm_params()
|
||||
yield
|
||||
set_llm_params(**prev)
|
||||
|
||||
|
||||
async def _nested_llm_load(state, *, branches=5, leaves=5):
|
||||
"""Drive branches*leaves leaf calls, nested two levels deep, each holding
|
||||
the shared per-loop LLM semaphore — the exact shape of the indexing pipeline
|
||||
|
|
@ -52,6 +63,34 @@ def test_llm_semaphore_bounds_concurrency_even_when_nested():
|
|||
assert state["peak"] == 3
|
||||
|
||||
|
||||
def test_llm_semaphore_is_a_true_process_wide_ceiling_across_threads():
|
||||
# The bug this fixes: each asyncio.run() (its own event loop) used to get
|
||||
# an independent full-size semaphore, so N concurrently-indexing threads
|
||||
# multiplied the effective cap by N. 2 threads, cap=3 -> combined peak
|
||||
# must stay at 3, not 6.
|
||||
set_max_concurrency(3)
|
||||
state = {"in_flight": 0, "peak": 0}
|
||||
lock = threading.Lock()
|
||||
|
||||
async def leaf():
|
||||
async with _llm_semaphore():
|
||||
with lock:
|
||||
state["in_flight"] += 1
|
||||
state["peak"] = max(state["peak"], state["in_flight"])
|
||||
await asyncio.sleep(0.05)
|
||||
with lock:
|
||||
state["in_flight"] -= 1
|
||||
|
||||
async def load():
|
||||
await asyncio.gather(*(leaf() for _ in range(5)))
|
||||
|
||||
threads = [threading.Thread(target=lambda: asyncio.run(load())) for _ in range(2)]
|
||||
[t.start() for t in threads]
|
||||
[t.join() for t in threads]
|
||||
|
||||
assert state["peak"] == 3
|
||||
|
||||
|
||||
def test_llm_semaphore_uses_scoped_override():
|
||||
# A per-index max_concurrency_scope active when the loop's semaphore is first
|
||||
# created must set its size, and must not mutate the process default.
|
||||
|
|
@ -193,6 +232,72 @@ def test_max_concurrency_scope_is_isolated_across_threads():
|
|||
assert seen["main"] == 10 # main is unaffected by the worker's scope
|
||||
|
||||
|
||||
def test_llm_params_scope_overrides_then_restores():
|
||||
set_llm_params(temperature=0)
|
||||
with llm_params_scope({"temperature": 1}):
|
||||
assert get_llm_params()["temperature"] == 1
|
||||
assert get_llm_params()["temperature"] == 0
|
||||
|
||||
|
||||
def test_llm_params_scope_none_is_a_no_op():
|
||||
set_llm_params(temperature=0)
|
||||
with llm_params_scope(None):
|
||||
assert get_llm_params()["temperature"] == 0
|
||||
assert get_llm_params()["temperature"] == 0
|
||||
|
||||
|
||||
def test_llm_params_scope_rejects_reserved_keys():
|
||||
with pytest.raises(ValueError):
|
||||
with llm_params_scope({"model": "x"}):
|
||||
pass
|
||||
|
||||
|
||||
def test_llm_params_scope_is_isolated_across_threads():
|
||||
set_llm_params(temperature=0)
|
||||
seen = {}
|
||||
barrier = threading.Barrier(2)
|
||||
|
||||
def worker():
|
||||
with llm_params_scope({"temperature": 1}):
|
||||
barrier.wait()
|
||||
seen["worker"] = get_llm_params()["temperature"]
|
||||
barrier.wait()
|
||||
|
||||
t = threading.Thread(target=worker)
|
||||
t.start()
|
||||
barrier.wait()
|
||||
seen["main"] = get_llm_params()["temperature"]
|
||||
barrier.wait()
|
||||
t.join()
|
||||
|
||||
assert seen["worker"] == 1
|
||||
assert seen["main"] == 0
|
||||
|
||||
|
||||
def test_llm_params_scope_does_not_leak_across_concurrent_indexing():
|
||||
# The bug this fixes: set_llm_params() mutates a bare process-wide dict, so
|
||||
# two documents indexed concurrently with different llm_params_scope()
|
||||
# overrides must not see each other's temperature.
|
||||
set_llm_params(temperature=0)
|
||||
seen = {"a": None, "b": None}
|
||||
|
||||
async def job(name, temperature, delay_before, delay_after):
|
||||
with llm_params_scope({"temperature": temperature}):
|
||||
await asyncio.sleep(delay_before)
|
||||
seen[name] = get_llm_params()["temperature"]
|
||||
await asyncio.sleep(delay_after)
|
||||
|
||||
async def run():
|
||||
await asyncio.gather(
|
||||
job("a", 1, 0.0, 0.05),
|
||||
job("b", 2, 0.02, 0.0),
|
||||
)
|
||||
|
||||
asyncio.run(run())
|
||||
assert seen["a"] == 1
|
||||
assert seen["b"] == 2
|
||||
|
||||
|
||||
def test_utils_star_import_does_not_leak_config_name():
|
||||
# `from .utils import *` (used by the page_index modules) must not export a
|
||||
# name `config` that would shadow the real pageindex.config submodule for
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue