Studio: read Deep Research synthesis context from the inference orchestrator

Make the adaptive synthesis-evidence budget actually engage in the normal Studio
architecture. _loaded_context_length read core.inference.inference, the low-level backend that
lives in the model subprocess and stays unpopulated in the main web process where the research
supervisor runs, so it returned None and the budget silently fell back to the 32000 character
cap (leaving the report exposed to the truncation this was meant to fix). Read the inference
orchestrator instead, and the llama.cpp backend for GGUF, mirroring
routes.inference._monitor_context_length so the budget sizes to the context the API layer
serves. Verified on a running server: at a 12288 token load the probe now reports 12288 and the
budget adapts to 24576 characters instead of the 32000 fallback.

Also:
- Reserve context for the generated report as well as the prompt scaffolding (raise the reserve
  to 4096 tokens) so evidence does not crowd out the output on a small window.
- Honor a numeric UNSLOTH_RESEARCH_AUTO_SCRAPE by passing the per-run maxAutoScrape as the page
  cap to the scraper, instead of always reading the maximum.
- Guard the web-RAG connection acquisition so a get_connection failure returns the documented
  empty result rather than propagating.
- Add a synthesis-context test that patches the real backend accessor (not the probe itself) so
  the production wiring is exercised, plus a scrape page-cap test.
This commit is contained in:
danielhanchen 2026-07-21 11:05:19 +00:00
commit a92f462df3
3 changed files with 99 additions and 12 deletions

View file

@ -75,7 +75,11 @@ def retrieve_web_chunks(
overlap = config.CHUNK_OVERLAP if overlap is None else overlap
count = embeddings.token_counter(model)
conn = rag_db.get_connection()
try:
conn = rag_db.get_connection()
except Exception:
logger.warning("research.web_rank_failed", exc_info = True)
return "", []
scope = f"research_scrape_{uuid.uuid4().hex}"
doc_ids: list[str] = []
try:

View file

@ -85,10 +85,11 @@ _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, then convert the remainder to chars. Unknown context keeps the full cap.
# scaffolding (system prompt, plan, source catalogs) AND the generated report, then convert the
# remainder to chars. Unknown context keeps the full cap.
_MIN_SYNTHESIS_EVIDENCE_CHARS = 1_500
_SYNTHESIS_EVIDENCE_CHARS_PER_TOKEN = 3.0
_SYNTHESIS_CONTEXT_RESERVE_TOKENS = 2_048
_SYNTHESIS_CONTEXT_RESERVE_TOKENS = 4_096
# Below this loaded context the prompt scaffolding alone fills the window and the grounded
# report degenerates, so grounding is skipped (snippet-only) for smaller loads.
_AUTO_SCRAPE_MIN_CONTEXT_TOKENS = 8_192
@ -443,16 +444,44 @@ def _research_question_context(thread_id: str, user_message_id: str) -> tuple[st
return question, json.dumps(turns, ensure_ascii = False)
def _positive_int_or_none(value: object) -> int | None:
return value if isinstance(value, int) and not isinstance(value, bool) and value > 0 else None
def _loaded_context_length() -> int | None:
"""Best-effort read of the active model's context window in tokens, or None if unknown."""
"""Best-effort read of the active model's context window in tokens, or None if unknown.
Mirrors routes.inference._monitor_context_length (llama.cpp backend, else the inference
orchestrator) so grounding sizes evidence to the same context the API layer serves. The ML
backends live in a worker subprocess, so the low-level core.inference.inference singleton is
unpopulated in this (main) process and importing it pulls in the ML stack; read the
orchestrator the routes use instead."""
# GGUF / llama.cpp keeps context on its own backend (checked first, like the API layer).
try:
from core.inference.inference import get_inference_backend
from routes.inference import get_llama_cpp_backend
llama = get_llama_cpp_backend()
if getattr(llama, "is_loaded", False):
ctx = _positive_int_or_none(getattr(llama, "context_length", None))
if ctx is not None:
return ctx
except Exception:
logger.debug("research.context_probe_llama_failed", exc_info = True)
# Native / transformers: the orchestrator the API layer reads (not the subprocess singleton).
try:
from core.inference import get_inference_backend
backend = get_inference_backend()
name = getattr(backend, "active_model_name", None)
if name:
ctx = (getattr(backend, "models", {}).get(name) or {}).get("context_length")
if isinstance(ctx, int) and not isinstance(ctx, bool) and ctx > 0:
models = getattr(backend, "models", {}) or {}
info = models.get(name) if (name and isinstance(models, dict)) else None
for candidate in (
(info or {}).get("context_length"),
getattr(backend, "context_length", None),
getattr(backend, "max_seq_length", None),
):
ctx = _positive_int_or_none(candidate)
if ctx is not None:
return ctx
except Exception:
logger.debug("research.context_probe_failed", exc_info = True)
@ -858,21 +887,25 @@ class ResearchSupervisor:
step_sources: list[dict],
fetched_urls: set[str],
*,
limit: int,
tool_timeout: int,
website_policy: dict | None,
) -> tuple[str, list[str]]:
"""Concurrently read the top few of this step's accepted source URLs, rank their
"""Concurrently read up to ``limit`` of this step's accepted source URLs, rank their
content against the research question with the knowledge-base embedding model, and
return the most relevant chunks as ``<chunk>`` evidence plus the URLs actually read.
URLs are already access checked and deduplicated by the caller, so no new sources are
created. Failures, timeouts, unreadable pages, and low-relevance chunks are dropped;
the caller enforces cancellation."""
cap = max(0, min(limit, _AUTO_SCRAPE_TOP_K))
if cap <= 0:
return "", []
targets = []
for source in step_sources:
url = str(source.get("url") or "")
if url and url not in fetched_urls:
targets.append(source)
if len(targets) >= _AUTO_SCRAPE_TOP_K:
if len(targets) >= cap:
break
if not targets:
return "", []
@ -1782,6 +1815,7 @@ class ResearchSupervisor:
question,
step_sources,
fetched_urls,
limit = max_auto_scrape,
tool_timeout = tool_timeout,
website_policy = website_policy,
)

View file

@ -196,6 +196,31 @@ def test_synthesis_evidence_budget_tracks_loaded_context(monkeypatch):
assert worker._synthesis_evidence_budget() == worker._MAX_SYNTHESIS_EVIDENCE_CHARS
def test_loaded_context_length_reads_orchestrator(monkeypatch):
# The probe must read the inference ORCHESTRATOR (what the API layer serves), not the
# low-level in-subprocess singleton that stays unpopulated in the main process. Patch the
# real accessor (not _loaded_context_length) so this exercises the production wiring; a probe
# that read the wrong backend would return None here and the adaptive budget would not engage.
import core.inference as core_inference
from core import research_runs as worker
class _Orchestrator:
active_model_name = "Qwen2.5-14B-Instruct"
models = {"Qwen2.5-14B-Instruct": {"context_length": 8192}}
monkeypatch.setattr(core_inference, "get_inference_backend", lambda: _Orchestrator(), raising = False)
assert worker._loaded_context_length() == 8192
assert worker._synthesis_evidence_budget() < worker._MAX_SYNTHESIS_EVIDENCE_CHARS
class _NoModel:
active_model_name = None
models: dict = {}
monkeypatch.setattr(core_inference, "get_inference_backend", lambda: _NoModel(), raising = False)
assert worker._loaded_context_length() is None
assert worker._synthesis_evidence_budget() == worker._MAX_SYNTHESIS_EVIDENCE_CHARS
def test_bounded_synthesis_evidence_respects_small_budget():
from core import research_runs as worker
@ -1544,7 +1569,7 @@ def test_auto_scrape_respects_char_budgets(research_home, monkeypatch):
section, fetched = asyncio.run(
supervisor._auto_scrape_sources(
{"id": "run-x"}, "question", step_sources, set(),
tool_timeout = 10, website_policy = None,
limit = worker._AUTO_SCRAPE_TOP_K, tool_timeout = 10, website_policy = None,
)
)
# the folded evidence is bounded chunks, not the 150k of raw page bodies
@ -1567,7 +1592,7 @@ def test_auto_scrape_falls_back_when_no_relevant_chunks(research_home, monkeypat
section, fetched = asyncio.run(
supervisor._auto_scrape_sources(
{"id": "run-x"}, "find the special token", step_sources, set(),
tool_timeout = 10, website_policy = None,
limit = worker._AUTO_SCRAPE_TOP_K, tool_timeout = 10, website_policy = None,
)
)
assert section == ""
@ -1621,6 +1646,7 @@ def test_auto_scrape_skips_already_fetched_urls(research_home, monkeypatch):
"question",
step_sources,
{"https://x.example.com"},
limit = worker._AUTO_SCRAPE_TOP_K,
tool_timeout = 10,
website_policy = None,
)
@ -1630,6 +1656,29 @@ def test_auto_scrape_skips_already_fetched_urls(research_home, monkeypatch):
assert "https://x.example.com" not in section
def test_auto_scrape_honors_numeric_limit(research_home, monkeypatch):
# A numeric UNSLOTH_RESEARCH_AUTO_SCRAPE (persisted as maxAutoScrape=N) caps the pages read,
# rather than always scraping _AUTO_SCRAPE_TOP_K.
worker, supervisor = _bare_supervisor(monkeypatch)
_patch_web_rank(monkeypatch)
called = []
def fake_tool(name, arguments, *args, **kwargs):
called.append(arguments["url"])
return "body for " + arguments["url"]
monkeypatch.setattr(worker, "execute_tool", fake_tool)
step_sources = [{"url": f"https://s{i}.example.com", "title": f"S{i}"} for i in range(3)]
_section, fetched = asyncio.run(
supervisor._auto_scrape_sources(
{"id": "run-x"}, "question", step_sources, set(),
limit = 1, tool_timeout = 10, website_policy = None,
)
)
assert len(called) == 1
assert len(fetched) == 1
def test_recovered_running_research_resumes_durable_progress(research_home, monkeypatch):
from core import research_runs as worker