diff --git a/arborist/qa/corpus_query.py b/arborist/qa/corpus_query.py index da7713e..9c18eb2 100644 --- a/arborist/qa/corpus_query.py +++ b/arborist/qa/corpus_query.py @@ -138,11 +138,54 @@ def run_query( } # 2. Pull chunk content per hit, build evidence rows. + # + # Two budget shapes: + # - policy=None: flat split `max_context_chars / len(hits)`, role + # is rank-based ("primary_answer_source" for rank 1, else + # "background_source"). Byte-identical to pre-step-6b behavior. + # - policy provided: classify each hit's role via + # arborist.qa.source_roles.classify_source_role; split the + # budget by SOURCE_ROLE_BUDGET_WEIGHTS[role] proportionally. + # Primary gets 2× the slice; noisy/sequel get 0.5×. Matches + # the legacy query() pipeline so Phase 2's cache wrapper can + # swap in. ts = _time.time() - per_doc_budget = max(1000, max_context_chars // max(1, len(hits))) + if policy is not None: + from arborist.qa.source_roles import ( + classify_source_role, + SOURCE_ROLE_BUDGET_WEIGHTS, + ) + from arborist.qa._text_norm import stem_for_match + from arborist.qa.query import _title_query_tokens + qtokens_stem = { + stem_for_match(t) for t in _title_query_tokens(question) + } + hit_roles = [ + classify_source_role( + h.title, qtokens_stem, document_uri=h.document_uri, + ) + for h in hits + ] + weights = [ + SOURCE_ROLE_BUDGET_WEIGHTS.get(r, 1.0) for r in hit_roles + ] + weight_sum = sum(weights) or 1.0 + per_hit_budgets = [ + max(1000, int(max_context_chars * w / weight_sum)) for w in weights + ] + else: + hit_roles = [ + "primary_answer_source" if i == 0 else "background_source" + for i in range(len(hits)) + ] + per_hit_budgets = [ + max(1000, max_context_chars // max(1, len(hits))) + ] * len(hits) chunks_for_evidence: list[dict] = [] total_chars = 0 - for rank, h in enumerate(hits, 1): + for rank, (h, role, per_doc_budget) in enumerate( + zip(hits, hit_roles, per_hit_budgets), 1, + ): rows = corpus.chunks_for_doc(h.document_root, limit=1) if not rows: continue @@ -157,9 +200,7 @@ def run_query( "chunk_idx": r.idx, "chunk_root": r.leaf_hash or ("0" * 64), "span": span, - "source_role": ( - "primary_answer_source" if rank == 1 else "background_source" - ), + "source_role": role, }) total_chars += len(span) if total_chars >= max_context_chars: diff --git a/tests/test_qa_corpus_query.py b/tests/test_qa_corpus_query.py index b13a5f4..7c93fea 100644 --- a/tests/test_qa_corpus_query.py +++ b/tests/test_qa_corpus_query.py @@ -188,6 +188,23 @@ def test_run_query_policy_forwards_verifier_kwarg(corpus): ), f"expected TOO_MANY_CLAIMS violation in capped run; got {violations!r}" +def test_run_query_policy_classifies_source_role(corpus): + """With policy provided, hit roles come from the title-based + classifier (arborist.qa.source_roles.classify_source_role), not + a flat 'rank 1 = primary, else = background' default. For the + 'anarchism' query and the Anarchism fixture title, the top hit + must classify as primary_answer_source either way (matches + the legacy behavior for single-token clean-title queries).""" + result = run_query( + corpus, "anarchism", + StubClient(answer="Anarchism is a philosophy. [E1]"), + model_id="stub", top_k=2, policy={}, + ) + src = result["sources"][0] + # Classifier identifies clean Anarchism title as primary. + assert src["source_role"] == "primary_answer_source" + + def test_run_query_policy_ignores_unknown_keys(corpus): """policy with verifier-irrelevant keys (e.g. base_version that Phase 1 step 6a doesn't honor yet) must not blow up — unknown keys