github/privateGPT
1
0
Fork 1
mirror of https://github.com/imartinez/privateGPT.git synced 2026-08-08 23:35:29 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: context stack (#2323)

Browse Source 
* fix: revert all changes

* fix: avoid to duplicate layers

* fix: store original input

* fixx: ty
This commit is contained in:
Javier Martinez
2026-07-30 17:45:51 +02:00
committed by GitHub
parent e534667c98
commit 0216db9f4a
8 changed files with 214 additions and 512 deletions
Download Patch File Download Diff File Expand all files Collapse all files

79
private_gpt/components/context/models/context_stack.py
View File

@@ -47,60 +47,18 @@ class ContextStack(BaseModel):
# ------------------------------------------------------------------
def to_system_prompt(self) -> list[TextBlock]:
"""Render prompt layers by priority (then insertion order).
Layers stay isolated in the stack; deduplication happens only here, at
render time. Two cases must be covered so the LLM never receives the
same text twice:
- *Stale duplicate*: a layer's text is reproduced verbatim by another
kept layer (e.g. the chat header rendered both as a
``RuntimeInstructionsLayer`` and baked back into a re-ingested
``UserInstructionsLayer``). The later candidate is dropped because a
kept block already contains it.
- *Snowballed aggregate*: a re-ingested ``UserInstructionsLayer`` may
embed the header, platform guidelines and even prior skill bodies
that the interceptors regenerate, this iteration, as isolated
layers. The freshly-generated isolated layers carry the latest
state (e.g. the ``ContextPromptLayer`` rebuilt from the latest
documents accumulated across iterations), so they must survive.
The aggregate layer is the one discarded: its rendered text *is a
superset of* (contains) one or more kept isolated blocks.
Rendering therefore walks layers in *descending* priority order — the
freshly-generated isolated layers (highest priority numbers) are kept
first — and a candidate is dropped when its text contains, or is
contained in, the text of any already-kept block. That preserves the
latest isolated layers across iterations and discards the duplicate /
stale aggregate layer at render, restoring the prompt to its original
isolated shape.
"""
"""Render prompt layers by priority (then insertion order)."""
ordered_layers = sorted(
enumerate(self.layers),
key=lambda item: (-item[1].priority, item[0]),
key=lambda item: (item[1].priority, item[0]),
)
blocks: list[TextBlock] = []
kept: list[str] = []
for _, layer in ordered_layers:
rendered = layer.render()
if not rendered or not rendered.strip():
continue
text = rendered.strip()
# Drop the candidate if its text already appears inside a kept
# block (stale duplicate) OR if it reproduces, as a snowballed
# aggregate, any kept isolated block. The freshly-generated
# isolated layers are kept first thanks to the descending-priority
# order, so the aggregate (e.g. a re-ingested UserInstructions
# layer built from a previous response) is the one discarded,
# while the latest isolated layers — including the ContextPrompt
# rebuilt with the latest documents — survive.
is_stale_duplicate = any(text in candidate for candidate in kept)
is_snowballed_aggregate = any(candidate in text for candidate in kept)
if is_stale_duplicate or is_snowballed_aggregate:
continue
blocks.append(TextBlock(text=text))
kept.append(text)
return blocks
chunks = [
rendered
for _, layer in ordered_layers
for rendered in [layer.render()]
if rendered.strip()
]
return [TextBlock(text=chunk) for chunk in chunks]
def all_tools(self) -> list[ToolSpec]:
"""Return deduplicated ToolSpec list from all TOOL_DEFINITIONS layers."""
@@ -150,30 +108,17 @@ class ContextStack(BaseModel):
# ------------------------------------------------------------------
# Immutable mutation helpers
# ------------------------------------------------------------------
def _remove_duplicate(self, layer: AnyContextLayer) -> "list[AnyContextLayer]":
"""Return current layers with any existing (type, source) duplicate removed."""
return [
existing
for existing in self.layers
if not (existing.type == layer.type and existing.source == layer.source)
]
def insert_layer(self, layer: AnyContextLayer, index: int) -> "ContextStack":
"""Return a new stack with *layer* inserted at *index* (default 0)."""
base = self._remove_duplicate(layer)
index = min(index, len(base))
return ContextStack(layers=[*base[:index], layer, *base[index:]])
return ContextStack(layers=[*self.layers[:index], layer, *self.layers[index:]])
def append_layer(self, layer: AnyContextLayer) -> "ContextStack":
"""Return a new stack with *layer* appended."""
return ContextStack(layers=[*self._remove_duplicate(layer), layer])
return ContextStack(layers=[*self.layers, layer])
def append_layers(self, layers: list[AnyContextLayer]) -> "ContextStack":
"""Return a new stack with *layers* appended."""
stack = self
for layer in layers:
stack = stack.append_layer(layer)
return stack
return ContextStack(layers=[*self.layers, *layers])
def remove_layers_of_type(self, layer_type: LayerType) -> "ContextStack":
"""Return a new stack with all layers of *layer_type* removed."""

73
private_gpt/components/engines/chat/async_chat_engine.py
View File

@@ -172,6 +172,7 @@ class AsyncChatCheckpoint(BaseModel):
payload: IterationCheckpointPayload = Field(
default_factory=IterationCheckpointPayload
)
original_input: ChatInputState | None = None
@dataclass
@@ -182,6 +183,7 @@ class _CheckpointContext:
payload: IterationCheckpointPayload = field(
default_factory=IterationCheckpointPayload
)
original_input: ChatInputState | None = None
@dataclass
@@ -371,6 +373,7 @@ class AsyncChatEngine:
IterationCheckpointPayload(),
hooks,
channel,
original_input=state.original_input,
)
async def resume(
@@ -385,6 +388,7 @@ class AsyncChatEngine:
checkpoint=checkpoint.checkpoint,
payload=checkpoint.payload,
context_stack=checkpoint.input.context_stack,
original_input=checkpoint.original_input,
)
handler = self._resolve_checkpoint_handler(checkpoint.checkpoint)
state = await handler(
@@ -411,6 +415,7 @@ class AsyncChatEngine:
new_payload,
hooks,
channel,
original_input=state.original_input,
)
if state.output.pending_external_tool_calls:
close_state = close_state.model_copy(deep=True)
@@ -425,6 +430,7 @@ class AsyncChatEngine:
new_payload,
hooks,
channel,
original_input=state.original_input,
)
async def run(
@@ -539,10 +545,12 @@ class AsyncChatEngine:
payload: IterationCheckpointPayload,
hooks: ExecutionHooks | None,
channel: EventChannel,
original_input: ChatInputState | None = None,
) -> ChatState:
context = self._build_checkpoint_context(
checkpoint=_IterationCheckpoint.BEFORE_ITERATION,
payload=payload,
original_input=original_input,
)
return await self._execute_before_iteration_checkpoint(
request, iteration, next_block_count, channel, hooks, context
@@ -555,11 +563,13 @@ class AsyncChatEngine:
payload: IterationCheckpointPayload,
hooks: ExecutionHooks | None,
channel: EventChannel,
original_input: ChatInputState | None = None,
) -> ChatState:
context = self._build_checkpoint_context(
checkpoint=_IterationCheckpoint.CLOSE,
stop_reason=stop_reason,
payload=payload,
original_input=original_input,
)
return await self._execute_close_checkpoint(
request, 0, 0, channel, hooks, context
@@ -573,15 +583,31 @@ class AsyncChatEngine:
payload: IterationCheckpointPayload,
hooks: ExecutionHooks | None,
channel: EventChannel,
original_input: ChatInputState | None = None,
) -> ChatState:
while True:
if iteration >= self._max_iterations:
return await self._execute_close(
request, StopReasonEnum.MAX_TOKENS.value, payload, hooks, channel
request,
StopReasonEnum.MAX_TOKENS.value,
payload,
hooks,
channel,
original_input=original_input,
)
state = await self._execute_before_iteration(
request, iteration, next_block_count, payload, hooks, channel
request,
iteration,
next_block_count,
payload,
hooks,
channel,
original_input=original_input,
)
# Keep the first-iteration snapshot for the whole loop. Later
# checkpoints rebuild run state from the materialized request, so
# re-snapshotting here would poison restore/system-prompt rebuilds.
original_input = state.original_input or original_input
new_payload = IterationCheckpointPayload(
model_id=state.runtime.model_id,
total_input_tokens=state.runtime.total_input_tokens,
@@ -598,6 +624,7 @@ class AsyncChatEngine:
new_payload,
hooks,
channel,
original_input=original_input,
)
if state.output.pending_external_tool_calls:
close_state = close_state.model_copy(deep=True)
@@ -621,12 +648,14 @@ class AsyncChatEngine:
stop_reason: str | None = None,
payload: IterationCheckpointPayload | None = None,
context_stack: ContextStack | None = None,
original_input: ChatInputState | None = None,
) -> _CheckpointContext:
return _CheckpointContext(
checkpoint=checkpoint,
stop_reason=stop_reason,
context_stack=context_stack,
payload=payload or IterationCheckpointPayload(),
original_input=original_input,
)
def _resolve_checkpoint_handler(
@@ -669,7 +698,10 @@ class AsyncChatEngine:
checkpoint_context: _CheckpointContext,
) -> ChatState:
run = self._initialize_run(
request, context_stack=checkpoint_context.context_stack, hooks=hooks
request,
context_stack=checkpoint_context.context_stack,
hooks=hooks,
original_input=checkpoint_context.original_input,
)
run.state.runtime.iteration = iteration
run.state.runtime.next_block_count = next_block_count
@@ -692,7 +724,10 @@ class AsyncChatEngine:
) -> ChatState:
del iteration, next_block_count
run = self._initialize_run(
request, context_stack=checkpoint_context.context_stack, hooks=hooks
request,
context_stack=checkpoint_context.context_stack,
hooks=hooks,
original_input=checkpoint_context.original_input,
)
channel.emit(RawMessageStartEvent.from_defaults())
await channel.flush()
@@ -725,6 +760,7 @@ class AsyncChatEngine:
context_stack=checkpoint_context.context_stack,
hooks=hooks,
checkpoint_payload=checkpoint_context.payload,
original_input=checkpoint_context.original_input,
)
async def _continue_tools_checkpoint(
@@ -736,11 +772,15 @@ class AsyncChatEngine:
context_stack: ContextStack | None = None,
hooks: ExecutionHooks | None = None,
checkpoint_payload: IterationCheckpointPayload | None = None,
original_input: ChatInputState | None = None,
) -> ChatState:
"""Continue from the tools checkpoint without re-running the LLM call."""
try:
run = self._initialize_run(
request, context_stack=context_stack, hooks=hooks
request,
context_stack=context_stack,
hooks=hooks,
original_input=original_input,
)
run.state.runtime.iteration = iteration
run.state.runtime.next_block_count = next_block_count
@@ -843,7 +883,10 @@ class AsyncChatEngine:
) -> ChatState:
del iteration, next_block_count
run = self._initialize_run(
request, context_stack=checkpoint_context.context_stack, hooks=hooks
request,
context_stack=checkpoint_context.context_stack,
hooks=hooks,
original_input=checkpoint_context.original_input,
)
self._apply_payload_usage(run, checkpoint_context.payload)
stop_reason = checkpoint_context.stop_reason
@@ -872,8 +915,14 @@ class AsyncChatEngine:
request: ChatRequest,
context_stack: ContextStack | None = None,
hooks: ExecutionHooks | None = None,
original_input: ChatInputState | None = None,
) -> _Run:
return self._initialize_run(request, context_stack=context_stack, hooks=hooks)
return self._initialize_run(
request,
context_stack=context_stack,
hooks=hooks,
original_input=original_input,
)
async def _pipe_events_through_interceptors(
self,
@@ -1547,6 +1596,7 @@ class AsyncChatEngine:
request: ChatRequest,
context_stack: ContextStack | None = None,
hooks: ExecutionHooks | None = None,
original_input: ChatInputState | None = None,
) -> _Run:
llm = self._llm_component.get_llm(request.system.model)
if not isinstance(llm, FunctionCallingLLM):
@@ -1581,7 +1631,14 @@ class AsyncChatEngine:
output=ChatOutputState(),
timeline=[],
)
state.original_input = state.input.model_copy(deep=True)
# original_input is a first-iteration snapshot. Rebuilding it from later
# checkpoint request materializations would make restore/system-prompt
# interceptors treat intermediate prompt state as the user original.
state.original_input = (
original_input
if original_input is not None
else state.input.model_copy(deep=True)
)
run = _Run(
state=state,
llm=llm,

1
private_gpt/components/engines/chat/checkpoint_store.py
View File

@@ -32,6 +32,7 @@ class ChatCheckpoint(BaseModel):
correlation_id: str
request_data: dict[str, Any]
context_stack_data: dict[str, Any] = Field(default_factory=dict)
original_input_data: dict[str, Any] | None = None
stream_type: str
metadata: dict[str, Any]
iteration: int

23
private_gpt/components/engines/chat/resumable_runner.py
View File

@@ -192,6 +192,7 @@ class ResumableChatRunner:
payload=saved.checkpoint_payload.model_copy(
update={"tool_responses": responses}
),
original_input=self._original_input(saved),
),
hooks=_RESUME_HOOKS,
channel=channel,
@@ -249,6 +250,7 @@ class ResumableChatRunner:
correlation_id=execution_id,
request_data=state.input.request.model_dump(mode="json"),
context_stack_data=state.input.context_stack.checkpoint_dump(),
original_input_data=self._dump_original_input(state.original_input),
stream_type=stream_type,
metadata=metadata,
iteration=state.runtime.iteration,
@@ -386,6 +388,27 @@ class ResumableChatRunner:
}
return request
@staticmethod
def _dump_original_input(
original_input: ChatInputState | None,
) -> dict[str, Any] | None:
if original_input is None or not isinstance(original_input, ChatInputState):
return None
return original_input.model_dump(mode="json")
@staticmethod
def _original_input(checkpoint: ChatCheckpoint) -> ChatInputState | None:
if not checkpoint.original_input_data:
return None
data = dict(checkpoint.original_input_data)
request_data = data.get("request")
if isinstance(request_data, dict):
data["request"] = ResumableChatRunner._request(request_data)
context_stack_data = data.get("context_stack")
if isinstance(context_stack_data, dict):
data["context_stack"] = ContextStack.model_validate(context_stack_data)
return ChatInputState.model_validate(data)
@staticmethod
def _context_stack(
checkpoint: ChatCheckpoint, request_data: dict[str, Any]

4
private_gpt/components/engines/chat/utils/request_builder.py
View File

@@ -10,6 +10,7 @@ from private_gpt.components.context.models.context_layer import (
UserInstructionsLayer,
)
from private_gpt.components.context.models.context_stack import ContextStack
from private_gpt.components.context.models.layer_type import LayerType
def build_initial_context_stack(
@@ -23,11 +24,13 @@ def build_initial_context_stack(
# in the context stack if they are present in the request.
if request.system.prompt:
stack = stack.remove_layers_of_type(LayerType.USER_INSTRUCTIONS)
stack = stack.append_layer(
UserInstructionsLayer(text=request.system.prompt, source=source)
)
if request.tool_config.tools:
stack = stack.remove_layers_of_type(LayerType.TOOL_DEFINITIONS)
stack = stack.append_layer(
ToolDefinitionsLayer(
tools=list(request.tool_config.tools),
@@ -36,6 +39,7 @@ def build_initial_context_stack(
)
if request.context.documents:
stack = stack.remove_layers_of_type(LayerType.DOCUMENT)
for document in request.context.documents:
stack = stack.append_layer(
DocumentLayer(document=document, source=source)

137
tests/components/context/test_context_stack_render.py
View File

@@ -1,137 +0,0 @@
"""Render-time deduplication of ``ContextStack.to_system_prompt``.
Layers stay isolated in the stack; deduplication happens only at render. The
contract enforced here:
- The same text never reaches the LLM twice (stale duplicate drop).
- A re-ingested ``UserInstructionsLayer`` that aggregates a previous response
(header + guidelines) is discarded once the freshly-generated isolated
layers that the interceptors rebuilt for this iteration already reproduce
its parts (snowballed-aggregate drop).
- The freshly-generated isolated layers, including the ``ContextPromptLayer``
rebuilt with the latest documents accumulated across iterations, must
survive over any stale aggregate that embeds an older version of them.
"""
from private_gpt.components.context.models.context_layer import (
ContextPromptLayer,
RuntimeInstructionsLayer,
SkillBodyLayer,
ToolInstructionsLayer,
UserInstructionsLayer,
)
from private_gpt.components.context.models.context_stack import ContextStack
def _render(stack: ContextStack) -> list[str]:
return [b.text for b in stack.to_system_prompt() if b.text]
class TestRenderTimeDeduplication:
def test_identical_layers_collapse_to_one_block(self) -> None:
text = "You are Zylon, an AI assistant.\nCurrent date: 2026-07-28."
stack = ContextStack(
layers=[
UserInstructionsLayer(text=text, source="request"),
RuntimeInstructionsLayer(text=text, source="platform_header"),
]
)
rendered = _render(stack)
assert rendered == [text]
def test_stale_duplicate_lower_priority_is_dropped(self) -> None:
"""Runtime header is kept; aggregate UserInstructions is dropped."""
header = "You are Zylon, an AI assistant."
guideline = "<response_formatting>\nWrite clearly.\n</response_formatting>"
bloated = f"{header}\n\n{guideline}"
stack = ContextStack(
layers=[
UserInstructionsLayer(text=bloated, source="request"),
RuntimeInstructionsLayer(text=header, source="platform_header"),
ToolInstructionsLayer(
tool_name="response_formatting",
instructions=guideline,
source="platform:tool_instructions",
),
]
)
rendered = _render(stack)
assert rendered.count(header) == 1
assert rendered.count(guideline) == 1
# The bloated aggregate must NOT survive: only the isolated layers do
assert bloated not in rendered
assert header in rendered
assert guideline in rendered
def test_latest_context_prompt_survives_stale_aggregate(self) -> None:
"""A re-ingested UserInstructions embedding an *older* rendered
context prompt is discarded; the fresh ContextPromptLayer (rebuilt
from the latest documents accumulated across iterations) survives.
"""
stale_ctx = "<context_doc ids=[LVSE]>\nold content\n</context_doc>"
fresh_ctx = "<context_doc ids=[LVSE, WR2J]>\nlatest content\n</context_doc>"
# Aggregate layer reproduces the *stale* version of the context.
bloated = f"You are Zylon.\n\n{stale_ctx}\n\n<response_formatting>...</response_formatting>"
stack = ContextStack(
layers=[
UserInstructionsLayer(text=bloated, source="request"),
RuntimeInstructionsLayer(
text="You are Zylon.", source="platform_header"
),
ContextPromptLayer(text=fresh_ctx, source="system_prompt"),
]
)
rendered = _render(stack)
assert fresh_ctx in rendered, "Latest context prompt must survive"
assert stale_ctx not in rendered, "Stale context prompt must be dropped"
assert bloated not in rendered, "Snowballed aggregate must be dropped"
def test_distinct_isolated_layers_kept(self) -> None:
"""No false positives: layers with non-overlapping content are kept."""
header = "You are Zylon."
guideline = "<response_formatting>\nWrite clearly.\n</response_formatting>"
# SkillBodyLayer wraps the instructions in <skill_content name="...">
skill_body = '<skill_content name="x">\nbody\n</skill_content>'
stack = ContextStack(
layers=[
RuntimeInstructionsLayer(text=header, source="platform_header"),
ToolInstructionsLayer(
tool_name="response_formatting",
instructions=guideline,
source="platform:tool_instructions",
),
SkillBodyLayer(
skill_id="x",
name="x",
version="1",
instructions="body",
source="skill:x",
),
]
)
rendered = _render(stack)
assert set(rendered) == {header, guideline, skill_body}
def test_render_is_idempotent_across_iterations(self) -> None:
"""Simulate two iterations where the same starter stack is rendered
twice — output must not grow with repeated calls.
"""
header = "You are Zylon."
stack = ContextStack(
layers=[
UserInstructionsLayer(text=f"{header}\n<old>", source="request"),
RuntimeInstructionsLayer(text=header, source="platform_header"),
]
)
first = _render(stack)
second = _render(ContextStack(layers=list(stack.layers)))
assert first == second
assert first.count(header) == 1

109
tests/engines/test_async_chat_engine.py
View File

@@ -329,6 +329,7 @@ async def _run_async_engine(
has_input_usage=state.runtime.has_input_usage,
has_output_usage=state.runtime.has_output_usage,
),
original_input=state.original_input,
),
channel=channel2,
)
@@ -869,6 +870,7 @@ async def test_extract_citation_interceptor_converts_bracket_refs_on_resume(
has_input_usage=state.runtime.has_input_usage,
has_output_usage=state.runtime.has_output_usage,
),
original_input=state.original_input,
),
channel=channel2,
)
@@ -894,3 +896,110 @@ async def test_extract_citation_interceptor_converts_bracket_refs_on_resume(
assert "<citation" in full_text, (
f"Expected <citation> XML tag in output, got: {full_text!r}"
)
@pytest.mark.asyncio
async def test_initialize_run_reuses_provided_original_input() -> None:
"""_initialize_run must not resnapshot original_input on later checkpoints."""
from llama_index.core.base.llms.types import TextBlock
from llama_index.core.llms.function_calling import FunctionCallingLLM
from private_gpt.components.context.models.context_layer import (
UserInstructionsLayer,
)
from private_gpt.components.context.models.context_stack import ContextStack
from private_gpt.components.context.models.layer_type import LayerType
class _FakeFunctionLLM(FunctionCallingLLM):
@property
def metadata(self):
return MagicMock(is_function_calling_model=True, context_window=8192)
def _prepare_chat_with_tools(self, *a, **k):
return {}
async def achat(self, *a, **k):
raise NotImplementedError
def chat(self, *a, **k):
raise NotImplementedError
def stream_chat(self, *a, **k):
raise NotImplementedError
async def astream_chat(self, *a, **k):
raise NotImplementedError
def complete(self, *a, **k):
raise NotImplementedError
async def acomplete(self, *a, **k):
raise NotImplementedError
def stream_complete(self, *a, **k):
raise NotImplementedError
async def astream_complete(self, *a, **k):
raise NotImplementedError
def chat_with_tools(self, *a, **k):
raise NotImplementedError
async def achat_with_tools(self, *a, **k):
raise NotImplementedError
def stream_chat_with_tools(self, *a, **k):
raise NotImplementedError
async def astream_chat_with_tools(self, *a, **k):
raise NotImplementedError
def get_tool_calls_from_response(self, *a, **k):
return []
first_request = ResolvedChatRequest(
messages=[ChatMessage(role=MessageRole.USER, content="hello")],
system=ResolvedSystemConfig(
model="default",
prompt=[TextBlock(text="USER PROMPT")],
),
)
later_request = first_request.model_copy(deep=True)
later_request.system.prompt = [TextBlock(text="FULL RENDERED PROMPT")]
later_request.messages = [
*later_request.messages,
ChatMessage(role=MessageRole.ASSISTANT, content="tool-turn"),
]
llm_component = MagicMock(spec=LLMComponent)
llm_component.get_llm.return_value = _FakeFunctionLLM()
engine = AsyncChatEngine(
llm_component=llm_component,
chat_scheduler=MagicMock(),
)
first_run = engine.initialize_run(first_request)
original = first_run.state.original_input
assert original is not None
first_layers = original.context_stack.layers_of_type(LayerType.USER_INSTRUCTIONS)
assert first_layers
assert first_layers[0].text == [TextBlock(text="USER PROMPT")]
second_run = engine.initialize_run(
later_request,
context_stack=ContextStack(
layers=[
UserInstructionsLayer(
text=[TextBlock(text="FULL RENDERED PROMPT")],
source="request",
)
]
),
original_input=original,
)
assert second_run.state.original_input is original
second_layers = second_run.state.original_input.context_stack.layers_of_type(
LayerType.USER_INSTRUCTIONS
)
assert second_layers
assert second_layers[0].text == [TextBlock(text="USER PROMPT")]

300
tests/server/chat/interceptors/test_system_prompt_interceptor.py
View File

@@ -1,300 +0,0 @@
"""Tests for SystemPromptRequestInterceptor layer deduplication.
Verifies that running the interceptor N times (simulating tool-call loops or
the recalculate branch) never accumulates duplicate layers in the context
stack or duplicates text in the rendered system prompt.
"""
from unittest.mock import MagicMock
import pytest
from llama_index.core.base.llms.types import ChatMessage, MessageRole, TextBlock
from private_gpt.components.chat.models.chat_config_models import (
ResolvedChatRequest,
ResolvedSystemConfig,
)
from private_gpt.components.context.models.context_stack import ContextStack
from private_gpt.components.context.models.layer_type import LayerType
from private_gpt.components.engines.chat.models.chat_interceptor_context import (
ChatInterceptorContext,
)
from private_gpt.components.engines.chat.models.chat_phase import InterceptorPhase
from private_gpt.components.engines.chat.models.chat_state import (
ChatInputState,
ChatOutputState,
ChatRuntimeState,
ChatState,
)
from private_gpt.components.engines.chat.utils.request_builder import (
build_initial_context_stack,
)
from private_gpt.server.chat.interceptors.system_prompt_interceptor import (
SystemPromptRequestInterceptor,
)
from tests.fixtures.mock_function_llm import get_mock_function_calling_llm
_SYSTEM_PROMPT = "You are Zylon, an AI assistant.\nCurrent date: 2026-07-27."
def _make_request(
system_prompt: str | list[TextBlock] | None = _SYSTEM_PROMPT,
) -> ResolvedChatRequest:
return ResolvedChatRequest(
messages=[ChatMessage(role=MessageRole.USER, content="hello")],
system=ResolvedSystemConfig(prompt=system_prompt),
)
def _make_context(
request: ResolvedChatRequest,
context_stack: ContextStack | None = None,
phase: InterceptorPhase = InterceptorPhase.BEFORE_ITERATION,
) -> ChatInterceptorContext:
stack = (
context_stack
if context_stack is not None
else build_initial_context_stack(request)
)
state = ChatState(
input=ChatInputState(
request=request,
context_stack=stack,
),
runtime=ChatRuntimeState(),
output=ChatOutputState(),
timeline=[],
)
return ChatInterceptorContext(
state=state,
llm=get_mock_function_calling_llm(["ok"]),
phase=phase,
emit_fn=lambda _: None,
)
def _make_interceptor(
add_context_to_system_prompt: bool = False,
) -> SystemPromptRequestInterceptor:
"""Build a SystemPromptRequestInterceptor with a minimal PromptBuilderService."""
prompt_builder = MagicMock()
prompt_template = MagicMock()
prompt_template.format.return_value = _SYSTEM_PROMPT
prompt_builder.create_chat_header_prompt.return_value = prompt_template
return SystemPromptRequestInterceptor(
prompt_builder_service=prompt_builder,
add_context_to_system_prompt=add_context_to_system_prompt,
)
class TestSystemPromptInterceptorIdempotency:
"""The interceptor must be idempotent across repeated calls."""
@pytest.mark.asyncio
async def test_single_run_produces_one_platform_header_layer(self) -> None:
interceptor = _make_interceptor()
request = _make_request()
context = _make_context(request)
await interceptor.intercept(context)
platform_layers = [
layer
for layer in context.state.input.context_stack.layers
if layer.source == "platform_header"
]
assert len(platform_layers) == 1, "Expected exactly 1 platform_header layer"
@pytest.mark.asyncio
async def test_running_n_times_does_not_accumulate_layers(self) -> None:
"""Simulates multiple BEFORE_ITERATION passes (tool-call loop)."""
interceptor = _make_interceptor()
request = _make_request()
context = _make_context(request)
for _ in range(5):
await interceptor.intercept(context)
platform_layers = [
layer
for layer in context.state.input.context_stack.layers
if layer.source == "platform_header"
]
assert len(platform_layers) == 1, (
f"After 5 iterations got {len(platform_layers)} platform_header layers — "
"interceptor is accumulating duplicates!"
)
@pytest.mark.asyncio
async def test_running_n_times_no_user_instruction_duplication(self) -> None:
"""Multiple BEFORE_ITERATION passes: user instructions appear once."""
interceptor = _make_interceptor()
request = _make_request()
context = _make_context(request)
for _ in range(5):
await interceptor.intercept(context)
user_layers = [
layer
for layer in context.state.input.context_stack.layers
if layer.type == LayerType.USER_INSTRUCTIONS
]
sources = [layer.source for layer in user_layers]
assert sources.count("request") <= 1, (
f"'request' USER_INSTRUCTIONS layer duplicated: {sources}"
)
assert sources.count("platform_header") <= 1, (
f"'platform_header' USER_INSTRUCTIONS layer duplicated: {sources}"
)
@pytest.mark.asyncio
async def test_system_prompt_text_not_duplicated_after_n_iterations(self) -> None:
"""The rendered system prompt text must not repeat after N runs."""
interceptor = _make_interceptor()
request = _make_request()
context = _make_context(request)
for _ in range(5):
await interceptor.intercept(context)
prompt = context.state.input.request.system.prompt
# Normalise to list of text strings
if isinstance(prompt, str):
texts = [prompt]
elif isinstance(prompt, list):
texts = [b.text for b in prompt if isinstance(b, TextBlock) and b.text]
else:
texts = []
full_text = "\n".join(texts)
occurrences = full_text.count(_SYSTEM_PROMPT)
assert (
occurrences <= 2
), ( # at most 2: once in user layer, once in platform_header
f"System prompt text appears {occurrences} times after 5 iterations. "
"Likely a duplication bug!"
)
@pytest.mark.asyncio
async def test_checkpoint_saves_original_not_mutated_prompt(self) -> None:
"""Verify checkpoint round-trip with original request avoids duplication.
The ResumableChatRunner now saves ``state.original_input.request``.
This test simulates the resume path where the original (clean) request
is used together with ``build_initial_context_stack``.
"""
interceptor = _make_interceptor()
original_request = _make_request()
# --- Simulate first request execution ---
context = _make_context(original_request)
await interceptor.intercept(context)
# Verify _render_system_prompt_text returns a single TextBlock
mutated_prompt = context.state.input.request.system.prompt
if isinstance(mutated_prompt, list):
assert len(mutated_prompt) == 1, (
"_render_system_prompt_text should return exactly 1 TextBlock"
)
# --- Simulate resume: build fresh stack from ORIGINAL (clean) request ---
restored_stack = build_initial_context_stack(original_request)
restored_context = _make_context(original_request, context_stack=restored_stack)
await interceptor.intercept(restored_context)
prompt = restored_context.state.input.request.system.prompt
if isinstance(prompt, str):
full_text = prompt
elif isinstance(prompt, list):
full_text = "\n".join(
b.text for b in prompt if isinstance(b, TextBlock) and b.text
)
else:
full_text = ""
occurrences = full_text.count(_SYSTEM_PROMPT)
assert occurrences <= 2, (
f"After resume the system prompt appears {occurrences} times."
)
@pytest.mark.asyncio
async def test_round_tripped_prompt_does_not_duplicate_rendered_blocks(
self,
) -> None:
"""A request whose ``system.prompt`` already carries the rendered
stack (client echoes back a previous response) must not produce a
system message whose content blocks repeat the header or guidelines
that the interceptors regenerate as isolated layers.
Layers stay isolated in the stack; at render time duplicate rendered
text is discarded so the snowball never reaches the LLM.
"""
header = "You are Zylon, an AI assistant.\nCurrent date: 2026-07-28."
guideline = "<response_formatting>\nWrite clearly.\n</response_formatting>"
# Client re-sends the fully rendered prompt (header + guideline baked in)
bloated_prompt = f"{header}\n\n{guideline}"
interceptor = _make_interceptor()
# Header template renders the same header the client already embedded
interceptor._prompt_builder_service.create_chat_header_prompt.return_value.format.return_value = header
request = _make_request(system_prompt=bloated_prompt)
context = _make_context(request)
await interceptor.intercept(context)
prompt = context.state.input.request.system.prompt
if isinstance(prompt, str):
texts = [prompt]
elif isinstance(prompt, list):
texts = [b.text for b in prompt if isinstance(b, TextBlock) and b.text]
else:
texts = []
full = "\n".join(texts)
# Header must appear at most once across all rendered blocks
assert full.count(header) <= 1, (
f"Header rendered {full.count(header)} times after round-trip: {texts!r}"
)
# Guideline must appear at most once
assert full.count(guideline) <= 1, (
f"Guideline rendered {full.count(guideline)} times: {texts!r}"
)
# No rendered block may fully duplicate another kept block
stripped = [t.strip() for t in texts if t.strip()]
for i, block in enumerate(stripped):
for j, other in enumerate(stripped):
if i != j and block and block in other:
raise AssertionError(
f"Rendered block #{i} is contained in block #{j} — "
"duplicate content reached the rendered prompt."
)
@pytest.mark.asyncio
async def test_fallback_build_with_mutated_prompt_is_safe(self) -> None:
"""Defensive: even if build_initial is called on a mutated request,
the system prompt should not explode (snowball test)."""
interceptor = _make_interceptor()
request = _make_request()
context = _make_context(request)
await interceptor.intercept(context)
# Simulate mutated request being re-ingested
mutated_request = context.state.input.request
restored_stack = build_initial_context_stack(mutated_request)
restored_context = _make_context(mutated_request, context_stack=restored_stack)
await interceptor.intercept(restored_context)
prompt = restored_context.state.input.request.system.prompt
if isinstance(prompt, str):
full_text = prompt
elif isinstance(prompt, list):
full_text = "\n".join(
b.text for b in prompt if isinstance(b, TextBlock) and b.text
)
else:
full_text = ""
occurrences = full_text.count(_SYSTEM_PROMPT)
assert occurrences <= 3, (
f"Even on fallback path, the system prompt explodes to {occurrences} occurrences."
)