diff --git a/studio/backend/core/rag/web_rank.py b/studio/backend/core/rag/web_rank.py index 1808e2b5e9..c86e9d9ec6 100644 --- a/studio/backend/core/rag/web_rank.py +++ b/studio/backend/core/rag/web_rank.py @@ -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: diff --git a/studio/backend/core/research_runs.py b/studio/backend/core/research_runs.py index 3b369ab6c5..1a56ad226d 100644 --- a/studio/backend/core/research_runs.py +++ b/studio/backend/core/research_runs.py @@ -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 ```` 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, ) diff --git a/studio/backend/tests/test_research_runs_storage.py b/studio/backend/tests/test_research_runs_storage.py index 2cf9bc2e91..91b4d22093 100644 --- a/studio/backend/tests/test_research_runs_storage.py +++ b/studio/backend/tests/test_research_runs_storage.py @@ -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