Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 20 additions & 2 deletions backend/src/agents/main_agent/session/compaction_summary.py
Original file line number Diff line number Diff line change
Expand Up @@ -84,6 +84,12 @@
# cap (half the budget) is 4k, so this is what binds; it is the figure the
# harness screen ran with.
_EXTRACTION_MAX_OUTPUT_TOKENS = 3_000
# A generation is used only on a positive completion signal. ``max_tokens``
# is the one exception, because it is salvaged (see the module docstring).
# Anything else (a guardrail stop, a content-filter refusal, an unknown
# reason) could be the refusal text itself, and a summary persists into the
# cacheable history until the next cut, so it falls back instead.
_ACCEPTED_STOP_REASONS = frozenset({"end_turn", "max_tokens"})

_COMPRESSION_SYSTEM_PROMPT = """You maintain the running summary of a long conversation between a user and an AI assistant. You are given the existing summary notes (oldest first). Rewrite them into ONE compact summary the assistant can continue the conversation from.

Expand Down Expand Up @@ -240,8 +246,14 @@ async def _compress(
"maxTokens": min(_max_output_tokens(model_id), max(256, int(budget_tokens))),
},
)
stop_reason = response.get("stopReason")
if stop_reason not in _ACCEPTED_STOP_REASONS:
logger.info(
"compaction_summary_model_refused: stopReason=%s; falling back to truncation", stop_reason,
)
return None, False
out = response["output"]["message"]["content"][0]["text"].strip()
if response.get("stopReason") == "max_tokens":
if stop_reason == "max_tokens":
salvaged = _salvage(out, budget_tokens)
logger.info(
"compaction_summary_model_truncated: generation hit the token ceiling; kept %d of %d chars",
Expand Down Expand Up @@ -305,8 +317,14 @@ async def extract_with_model(
messages=[{"role": "user", "content": [{"text": "Summary notes, oldest first:\n\n" + text}]}],
inferenceConfig={"temperature": 0.0, "maxTokens": max(256, int(max_tokens))},
)
stop_reason = response.get("stopReason")
if stop_reason not in _ACCEPTED_STOP_REASONS:
logger.info(
"compaction_summary_extract_refused: stopReason=%s; falling back to plain compression", stop_reason,
)
return None
out = response["output"]["message"]["content"][0]["text"].strip()
if response.get("stopReason") == "max_tokens":
if stop_reason == "max_tokens":
logger.info("compaction_summary_extract_truncated: keeping the complete lines")
cut = out.rfind("\n")
out = out[:cut].rstrip() if cut > 0 else ""
Expand Down
8 changes: 8 additions & 0 deletions backend/src/apis/inference_api/chat/service.py
Original file line number Diff line number Diff line change
Expand Up @@ -664,6 +664,14 @@ async def generate_conversation_title(
inferenceConfig=request_body["inferenceConfig"],
)

# A guardrail or content-filter stop may carry the refusal as its
# text, which must not become the sidebar title. `max_tokens` is
# still accepted: the 50-char clip below already handles an overrun.
stop_reason = response.get("stopReason")
if stop_reason not in ("end_turn", "max_tokens"):
logger.warning("Title generation stopped with stopReason=%s; keeping the placeholder", stop_reason)
return "New Conversation"

# Extract generated title from response
title = response["output"]["message"]["content"][0]["text"].strip()

Expand Down
8 changes: 7 additions & 1 deletion backend/src/apis/shared/files/document_digest.py
Original file line number Diff line number Diff line change
Expand Up @@ -277,9 +277,15 @@ async def generate_abstract(outline: DocumentDigest, sample: str, model_id: str
# silently drop every abstract.
inferenceConfig={"temperature": 0.2, "maxTokens": _ABSTRACT_MAX_OUTPUT_TOKENS},
)
if response.get("stopReason") == "max_tokens":
stop_reason = response.get("stopReason")
if stop_reason == "max_tokens":
logger.debug("Document abstract hit the token ceiling; discarding")
return None
# A guardrail or content-filter stop may carry the refusal as its
# text; only a finished generation is an abstract.
if stop_reason != "end_turn":
logger.debug("Document abstract stopped with stopReason=%s; discarding", stop_reason)
return None
text = response["output"]["message"]["content"][0]["text"].strip()
return re.sub(r"\s+", " ", text) or None
except Exception: # noqa: BLE001 - an abstract is never worth a failed upload
Expand Down
9 changes: 8 additions & 1 deletion backend/src/apis/shared/tool_summaries/summarizer.py
Original file line number Diff line number Diff line change
Expand Up @@ -187,9 +187,16 @@ async def summarize_tool_batch(calls: List[Dict[str, Any]]) -> Optional[str]:
# deterministic formatter speak. (Defensive: the dangling quotes seen
# on dev turned out to be `_unwrap_quotes`, not truncation, but
# nothing guarded this boundary and a fragment must never persist.)
if response.get("stopReason") == "max_tokens":
# Any other non-`end_turn` stop (a guardrail or content-filter
# refusal, an unknown reason) is dropped too: its text may be the
# refusal itself, which is no summary of the batch.
stop_reason = response.get("stopReason")
if stop_reason == "max_tokens":
logger.debug("Tool-batch summary hit the token ceiling; discarding")
return None
if stop_reason != "end_turn":
logger.debug("Tool-batch summary stopped with stopReason=%s; discarding", stop_reason)
return None

summary = _clean(response["output"]["message"]["content"][0]["text"])
if not summary:
Expand Down
36 changes: 36 additions & 0 deletions backend/tests/agents/main_agent/session/test_compaction_summary.py
Original file line number Diff line number Diff line change
Expand Up @@ -90,6 +90,21 @@ async def test_model_ceiling_hit_with_no_complete_line_falls_back(self, bedrock)
assert result.outcome == "truncated_after_model"
assert approx_tokens(result.text) <= BUDGET

@pytest.mark.asyncio
@pytest.mark.parametrize("stop", ["guardrail_intervened", "content_filtered", "tool_use", "something_new", None])
async def test_a_refused_generation_falls_back_and_is_never_the_summary(self, bedrock, stop):
"""Only a completed (or salvageable) generation becomes the checkpoint.

A refusal persisted as the summary would sit in the cacheable history
until the next cut.
"""
bedrock.return_value = _model_reply("Sorry, the model cannot answer this.", stop=stop)
records = ["old " * 100, "new " * 50]
result = await bound_summary(records, BUDGET, model_enabled=True, model_id="m")
assert result.outcome == "truncated_after_model"
assert "Sorry" not in result.text and result.text.startswith("new")
assert await compress_with_model(records, BUDGET, model_id="m") is None

@pytest.mark.asyncio
async def test_model_overshoot_is_tail_trimmed(self, bedrock):
bedrock.return_value = _model_reply("y" * 2000 + "END")
Expand Down Expand Up @@ -390,6 +405,27 @@ async def test_extraction_ceiling_hit_keeps_the_complete_lines(self, bedrock):
assert result.outcome == "extract_then_compress"
assert "- PRJ-4417" in result.text and "- PRJ-44\n" not in result.text

@pytest.mark.asyncio
async def test_refused_extraction_keeps_the_narrative_without_a_third_call(self, bedrock):
bedrock.side_effect = _route(
_model_reply("Sorry, the model cannot answer this.", stop="guardrail_intervened"),
_model_reply("plain compressed summary"),
)
result = await _extract(["r" * 2000])
assert result.outcome == "model"
assert result.text == "plain compressed summary"
assert bedrock.call_count == 2

@pytest.mark.asyncio
async def test_refused_narrative_keeps_the_pinned_block_and_truncates(self, bedrock):
bedrock.side_effect = _route(
_model_reply(PINNED),
_model_reply("Sorry, the model cannot answer this.", stop="guardrail_intervened"),
)
result = await _extract(["old " * 300, "newest record"])
assert result.outcome == "extract_then_truncate"
assert "Sorry" not in result.text and result.text.endswith("newest record")

@pytest.mark.asyncio
async def test_narrative_ceiling_hit_keeps_its_complete_lines(self, bedrock):
"""A cut-off narrative is salvaged, not swapped for raw records (extract_then_truncate)."""
Expand Down
21 changes: 21 additions & 0 deletions backend/tests/apis/shared/tool_summaries/test_summarizer.py
Original file line number Diff line number Diff line change
Expand Up @@ -92,6 +92,27 @@ async def test_truncation_is_judged_by_stop_reason_not_by_length(bedrock):
assert await summarize_tool_batch(_calls()) == long_but_finished




@pytest.mark.asyncio
@pytest.mark.parametrize("stop_reason", ["guardrail_intervened", "content_filtered", "tool_use", "something_new"])
async def test_any_stop_other_than_end_turn_is_discarded(bedrock, stop_reason):
"""A guardrail or content-filter stop can carry the refusal as its text.

Only `end_turn` says the model finished a summary; anything else must
leave the deterministic line in place.
"""
bedrock.converse.return_value = _response("Sorry, the model cannot answer this.", stop_reason=stop_reason)

assert await summarize_tool_batch(_calls()) is None


@pytest.mark.asyncio
async def test_a_missing_stop_reason_is_discarded(bedrock):
bedrock.converse.return_value = {"output": {"message": {"content": [{"text": "Found 3 active courses"}]}}}

assert await summarize_tool_batch(_calls()) is None

# -- every other failure is also a None -----------------------------------


Expand Down
7 changes: 7 additions & 0 deletions backend/tests/shared/test_document_digest.py
Original file line number Diff line number Diff line change
Expand Up @@ -146,6 +146,13 @@ async def test_truncated_or_failed_generations_are_none(self, monkeypatch):
assert await dd.generate_abstract(dd.DocumentDigest(), "text") is None
assert await dd.generate_abstract(dd.DocumentDigest(), " ") is None

@pytest.mark.asyncio
@pytest.mark.parametrize("stop", ["guardrail_intervened", "content_filtered", "something_new", None])
async def test_a_refused_generation_is_none(self, monkeypatch, stop):
"""Only a finished generation is an abstract; a refusal's text is not."""
_bedrock(monkeypatch, text="Sorry, the model cannot answer this.", stop=stop)
assert await dd.generate_abstract(dd.DocumentDigest(), "text") is None


class TestBuild:
@pytest.mark.asyncio
Expand Down
26 changes: 24 additions & 2 deletions backend/tests/shared/test_side_channel_inference_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,12 +17,12 @@
from apis.shared.tool_summaries.summarizer import summarize_tool_batch


def _client(text: str) -> MagicMock:
def _client(text: str, stop_reason: str = "end_turn") -> MagicMock:
def converse(**kwargs):
config = kwargs.get("inferenceConfig", {})
if "temperature" in config and "topP" in config:
raise RuntimeError("ValidationException: `temperature` and `top_p` cannot both be specified")
return {"stopReason": "end_turn", "output": {"message": {"content": [{"text": text}]}}}
return {"stopReason": stop_reason, "output": {"message": {"content": [{"text": text}]}}}

client = MagicMock()
client.converse.side_effect = converse
Expand Down Expand Up @@ -64,3 +64,25 @@ async def test_conversation_title(monkeypatch):
title = await chat_service.generate_conversation_title(session_id="s", user_id="u", user_input="hi")
assert title == "Planning a biology syllabus"
assert "topP" not in client.converse.call_args.kwargs["inferenceConfig"]


@pytest.mark.asyncio
@pytest.mark.parametrize("stop_reason", ["guardrail_intervened", "content_filtered", None])
async def test_conversation_title_refusal_keeps_the_placeholder(monkeypatch, stop_reason):
"""A guardrail stop's text is the refusal; it must never become the title."""
client = _client("Sorry, the model cannot answer this.", stop_reason=stop_reason)
monkeypatch.setattr(chat_service.boto3, "client", MagicMock(return_value=client))
update = AsyncMock()
monkeypatch.setattr(chat_service, "update_session_title", update)
title = await chat_service.generate_conversation_title(session_id="s", user_id="u", user_input="hi")
assert title == "New Conversation"
update.assert_not_awaited()


@pytest.mark.asyncio
async def test_conversation_title_at_the_token_ceiling_is_still_clipped(monkeypatch):
client = _client("A" * 80, stop_reason="max_tokens")
monkeypatch.setattr(chat_service.boto3, "client", MagicMock(return_value=client))
monkeypatch.setattr(chat_service, "update_session_title", AsyncMock())
title = await chat_service.generate_conversation_title(session_id="s", user_id="u", user_input="hi")
assert title == "A" * 47 + "..."