diff --git a/studio/backend/core/research_runs.py b/studio/backend/core/research_runs.py index 2e7437e970..959fbc0c19 100644 --- a/studio/backend/core/research_runs.py +++ b/studio/backend/core/research_runs.py @@ -90,10 +90,10 @@ _MAX_CONTEXT_CHARS = 12_000 _MAX_CONTEXT_MESSAGE_CHARS = 4_000 _MAX_SYNTHESIS_EVIDENCE_CHARS = 32_000 # The synthesis prompt must fit the loaded context or it is silently truncated and the report -# degenerates (echoes the evidence tail). Studio defaults context to 2048 tokens, far below the -# cap above, so the evidence budget adapts to the loaded context: reserve tokens for the prompt -# scaffolding (system prompt, plan, source catalogs) AND the generated report, then convert the -# remainder to chars. Unknown context keeps the full cap. +# degenerates (echoes the evidence tail). GGUF auto-fit floors at 4096 and transformers models +# default to 4096, but the context box accepts anything from 128 up, so the budget adapts: the +# reserve covers the generated report, and every trimmable section is measured against what the +# untrimmable scaffolding leaves. Unknown context keeps the full cap. _MIN_SYNTHESIS_EVIDENCE_CHARS = 1_500 _SYNTHESIS_EVIDENCE_CHARS_PER_TOKEN = 3.0 _SYNTHESIS_CONTEXT_RESERVE_TOKENS = 4_096 @@ -537,15 +537,33 @@ def _local_model_ready() -> bool: return not probed -def _synthesis_evidence_budget() -> int: - """Char budget for synthesis evidence, sized to fit the loaded context (falls back to the - full cap when the context is unknown).""" +def _prompt_char_budget(reserve_tokens: int) -> int | None: + """Chars the whole prompt may occupy on the loaded context, or None when it is unknown.""" ctx = _loaded_context_length() if not ctx: - return _MAX_SYNTHESIS_EVIDENCE_CHARS - usable_tokens = max(0, ctx - _SYNTHESIS_CONTEXT_RESERVE_TOKENS) - budget = int(usable_tokens * _SYNTHESIS_EVIDENCE_CHARS_PER_TOKEN) - return max(_MIN_SYNTHESIS_EVIDENCE_CHARS, min(budget, _MAX_SYNTHESIS_EVIDENCE_CHARS)) + return None + return int(max(0, ctx - reserve_tokens) * _SYNTHESIS_EVIDENCE_CHARS_PER_TOKEN) + + +def _trimmable_budget(total: int | None, fixed_chars: int, hard_cap: int) -> int: + """Chars left for a trimmable section once the rest of the prompt is counted. + + Budgeting one section against the context while the others are unbounded does not stop an + overflow: at a 2048-token context the untrimmable scaffolding alone is several times the + window. Returns 0 rather than a floor, since a short report beats a failed run. + """ + if total is None: + return hard_cap + return max(0, min(hard_cap, total - fixed_chars)) + + +def _synthesis_evidence_budget(fixed_chars: int = 0) -> int: + """Char budget for synthesis evidence (full cap when the context is unknown).""" + return _trimmable_budget( + _prompt_char_budget(_SYNTHESIS_CONTEXT_RESERVE_TOKENS), + fixed_chars, + _MAX_SYNTHESIS_EVIDENCE_CHARS, + ) def _bounded_synthesis_evidence( @@ -1787,30 +1805,44 @@ class ResearchSupervisor: for source in sources ) evidence = "\n\n".join(decision_notes) + decision_system = _system_prompt_with_instructions( + _AGENT_SYSTEM_PROMPT + (f"\n\n{policy_prompt}" if policy_prompt else ""), + run["config"], + ) + # Same whole-prompt budget as synthesis: a fixed 60k evidence tail is many times a + # small loaded context, and this runs on every step, so an overflow here kills the + # run long before it can synthesize what it already gathered. + decision_plan_json = json.dumps(run["plan"], ensure_ascii = False) + decision_scaffold = ( + len(decision_system) + len(question) + len(decision_plan_json) + len(source_catalog) + ) + decision_total = _prompt_char_budget(_SYNTHESIS_CONTEXT_RESERVE_TOKENS) + evidence_chars = _trimmable_budget( + decision_total, decision_scaffold, _MAX_SYNTHESIS_EVIDENCE_CHARS + ) + decision_context = conversation_context[ + : _trimmable_budget( + decision_total, decision_scaffold + evidence_chars, _MAX_CONTEXT_CHARS + ) + ] decision, decision_reasoning, _finish_reason = await self._stream_completion( run, [ { "role": "system", - "content": ( - _system_prompt_with_instructions( - _AGENT_SYSTEM_PROMPT - + (f"\n\n{policy_prompt}" if policy_prompt else ""), - run["config"], - ) - ), + "content": decision_system, }, { "role": "user", "content": ( - f"Conversation context JSON:\n{_shield_untrusted(conversation_context)}\n\n" + f"Conversation context JSON:\n{_shield_untrusted(decision_context)}\n\n" f"Question:\n{_shield_untrusted(question)}\n\n" f"Approved plan (guidance only):\n" - f"{_shield_untrusted(json.dumps(run['plan'], ensure_ascii = False))}\n\n" + f"{_shield_untrusted(decision_plan_json)}\n\n" f"Actions remaining after this one: {max_steps - position - 1}\n" f"\n" f"Gathered sources:\n{_shield_untrusted(source_catalog) or '(none)'}\n\n" - f"{_shield_untrusted(evidence[-60000:]) or '(none)'}\n" + f"{_shield_untrusted(evidence[-evidence_chars:] if evidence_chars else '') or '(none)'}\n" f"" ), }, @@ -2070,16 +2102,32 @@ class ResearchSupervisor: f" Chunk ID: {source.get('chunkId') or '(unknown)'}" for index, source in enumerate(document_sources, 1) ) - evidence_text = _bounded_synthesis_evidence(notes, _synthesis_evidence_budget()) + # Budget the whole prompt, not just the evidence: trim the conversation context first, + # then the evidence, so the untrimmable scaffolding cannot push the request past the + # loaded context and turn a finished run into a failure. + report_system = _system_prompt_with_instructions(_REPORT_SYSTEM_PROMPT, run["config"]) + plan_json = json.dumps(run["plan"], ensure_ascii = False) + scaffold_chars = ( + len(report_system) + + len(question) + + len(plan_json) + + len(source_catalog) + + len(document_source_catalog) + ) + # Evidence is the report, so it is budgeted first and the chat history takes what is left. + total_budget = _prompt_char_budget(_SYNTHESIS_CONTEXT_RESERVE_TOKENS) + evidence_text = _bounded_synthesis_evidence( + notes, _synthesis_evidence_budget(scaffold_chars) + ) + conversation_context = conversation_context[ + : _trimmable_budget(total_budget, scaffold_chars + len(evidence_text), _MAX_CONTEXT_CHARS) + ] report, synthesis_reasoning, synthesis_finish_reason = await self._stream_completion( run, [ { "role": "system", - "content": _system_prompt_with_instructions( - _REPORT_SYSTEM_PROMPT, - run["config"], - ), + "content": report_system, }, { "role": "user", diff --git a/studio/backend/tests/test_research_runs_hardening.py b/studio/backend/tests/test_research_runs_hardening.py index 3af7be1453..d598e34dea 100644 --- a/studio/backend/tests/test_research_runs_hardening.py +++ b/studio/backend/tests/test_research_runs_hardening.py @@ -161,6 +161,23 @@ def test_citation_title_strips_brackets_for_catalog_and_citation(): assert _citation_title({}, "https://x/a") == "https://x/a" +def test_prompt_budget_counts_the_whole_prompt(monkeypatch): + # Budgeting only the evidence cannot prevent an overflow: at a small context the + # untrimmable scaffolding (system prompt, plan, source catalogs) is already several times + # the window, and the old floor added 1500 chars on top of that. + monkeypatch.setattr(research_runs, "_loaded_context_length", lambda: None) + assert research_runs._prompt_char_budget(4096) is None + assert research_runs._trimmable_budget(None, 99_999, 500) == 500 + + monkeypatch.setattr(research_runs, "_loaded_context_length", lambda: 16384) + total = research_runs._prompt_char_budget(4096) + assert total == int((16384 - 4096) * research_runs._SYNTHESIS_EVIDENCE_CHARS_PER_TOKEN) + # A trimmable section never exceeds what is left, and never goes negative. + assert research_runs._trimmable_budget(total, 0, 1_000) == 1_000 + assert research_runs._trimmable_budget(total, total - 10, 1_000) == 10 + assert research_runs._trimmable_budget(total, total + 5_000, 1_000) == 0 + + def _make_payload(**overrides) -> CreateResearchRun: payload = {"threadId": "t1", "userMessageId": "u1", "inferenceRequest": {"model": "m"}} payload.update(overrides) diff --git a/studio/backend/tests/test_research_runs_storage.py b/studio/backend/tests/test_research_runs_storage.py index c87ef586b8..161064c6ae 100644 --- a/studio/backend/tests/test_research_runs_storage.py +++ b/studio/backend/tests/test_research_runs_storage.py @@ -187,10 +187,16 @@ def test_synthesis_evidence_budget_tracks_loaded_context(monkeypatch): monkeypatch.setattr(worker, "_loaded_context_length", lambda: None) assert worker._synthesis_evidence_budget() == worker._MAX_SYNTHESIS_EVIDENCE_CHARS - # A small context (Studio's 2048 default) shrinks the budget so evidence fits. + # A small context (Studio's 2048 default) shrinks the budget so evidence fits. It must reach + # 0 rather than a floor: the old floor handed back 1500 chars even when the context left room + # for none, so the request still overflowed and failed the run after every search had run. monkeypatch.setattr(worker, "_loaded_context_length", lambda: 2048) - small = worker._synthesis_evidence_budget() - assert worker._MIN_SYNTHESIS_EVIDENCE_CHARS <= small < worker._MAX_SYNTHESIS_EVIDENCE_CHARS + assert worker._synthesis_evidence_budget() == 0 + + # The rest of the prompt counts against the same budget, not just the evidence. + monkeypatch.setattr(worker, "_loaded_context_length", lambda: 16384) + roomy = worker._synthesis_evidence_budget() + assert 0 < worker._synthesis_evidence_budget(8_000) < roomy # A large context uses (and clamps to) the full cap. monkeypatch.setattr(worker, "_loaded_context_length", lambda: 32768)