From 999e30fda33e95e7537664fdad5fd779af51786a Mon Sep 17 00:00:00 2001 From: Adilbek Bazarkulov Date: Mon, 8 Jun 2026 12:31:36 -0400 Subject: [PATCH 1/6] Ignore local working documents --- .gitignore | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.gitignore b/.gitignore index 0e7a27a..321b48c 100644 --- a/.gitignore +++ b/.gitignore @@ -3,7 +3,9 @@ # local-only working docs (kept on disk, not published) AGENTS.md +AUDIT.md CLAUDE.md CONTEXT.md DEPLOY.md PIPELINE.md +PLAN.md From cca410fd336df07e05933dd0a1a1123a7bccc856 Mon Sep 17 00:00:00 2001 From: Adilbek Bazarkulov Date: Mon, 8 Jun 2026 12:31:36 -0400 Subject: [PATCH 2/6] Demote text-mined edges and add response chunking primitive --- pipeline/code/reduction.py | 441 +++++++++++++----- pipeline/code/tests/test_chunking.py | 150 ++++++ .../code/tests/test_response_reduction.py | 102 +++- 3 files changed, 548 insertions(+), 145 deletions(-) create mode 100644 pipeline/code/tests/test_chunking.py diff --git a/pipeline/code/reduction.py b/pipeline/code/reduction.py index 361f85d..e37463a 100644 --- a/pipeline/code/reduction.py +++ b/pipeline/code/reduction.py @@ -16,6 +16,9 @@ # faithfulness evaluator) can rely on field names not silently # changing. from dataclasses import dataclass +# json: stdlib. used only to estimate a chunk's token budget by measuring +# the serialised character count of its sub-body. +import json # logging: stdlib. one INFO line at exit summarising the reduction, # one WARNING per malformed result. import logging @@ -54,6 +57,18 @@ _AT_RANK_UNKNOWN = 5 _AT_RANK_ABSENT = 4 +# text-mining provenance. SemMedDB and similar text-mined sources are the +# most abundant edges in KG2c but the noisiest; we DEMOTE every text-mined +# edge below every curated edge (source_tier is the LEADING sort key) so a +# well-cited text-mined edge can no longer dominate a curated one. an edge +# is text-mined when its agent_type is text_mining_agent OR its primary +# knowledge source is a known text-mining infores. +_TEXT_MINING_AGENT = "text_mining_agent" +_TEXT_MINING_SOURCES: frozenset[str] = frozenset({ + "infores:semmeddb", + "infores:text-mining-provider-targeted", +}) + # the strategy label that lands in the artifact. constant in v1; would # change here if we ever introduced a Strategy-C or hybrid variant. _STRATEGY_NAME = "B" @@ -109,7 +124,7 @@ class _Enriched: bound_edge_ids: tuple[str, ...] representative_edge_id: str primary_predicate: str - sort_key: tuple[int, int, int, str] + sort_key: tuple[int, int, int, int, str] original: dict[str, Any] node_ids_in_bindings: tuple[str, ...] @@ -149,6 +164,29 @@ def _n_publications(edge: dict[str, Any]) -> int: return 0 +def _primary_knowledge_source(edge: dict[str, Any]) -> str | None: + # TRAPI edges declare provenance via `sources` (a list of + # {resource_id, resource_role}); the primary source is the entry whose + # role is "primary_knowledge_source". returns None if absent. + sources = edge.get("sources") + if not isinstance(sources, list): + return None + for source in sources: + if isinstance(source, dict) and source.get("resource_role") == "primary_knowledge_source": + resource_id = source.get("resource_id") + return resource_id if isinstance(resource_id, str) else None + return None + + +def _source_tier(edge: dict[str, Any]) -> int: + # 0 = curated, 1 = text-mined (demoted). leading key in the sort so + # text-mined edges sink below ALL curated edges within a predicate + # group, regardless of publication count. + if _read_attr(edge, _ATTR_AGENT_TYPE) == _TEXT_MINING_AGENT: + return 1 + return 1 if _primary_knowledge_source(edge) in _TEXT_MINING_SOURCES else 0 + + def _all_edge_ids_in_result(result: dict[str, Any]) -> tuple[str, ...]: # TRAPI 1.5: edge_bindings live under analyses[i].edge_bindings as # {qg_edge_id: [{"id": kg_edge_id, ...}, ...]}. a single result can @@ -248,20 +286,19 @@ def _enrich_results( ) continue # sort tuple ordering, ascending (smaller = stronger): - # 1. knowledge_level rank (smaller = stronger) - # 2. -n_publications (more pubs = first; continuous evidence signal) - # 3. agent_type rank (smaller = stronger; binary provenance label) - # 4. edge_id alphabetical (full determinism tie-break) + # 1. source_tier (0 curated, 1 text-mined → text-mined last) + # 2. knowledge_level rank (smaller = stronger) + # 3. -n_publications (more pubs = first; continuous evidence signal) + # 4. agent_type rank (smaller = stronger; binary provenance label) + # 5. edge_id alphabetical (full determinism tie-break) # - # n_pubs is promoted ahead of agent_type because publication - # count is a continuous corroboration signal across sources, - # whereas agent_type is a binary label about extraction - # provenance. when knowledge_level ties across a whole - # predicate group (as it does for KG2c gene-disease edges - # where every edge is kl=prediction), letting agent_type - # outrank n_pubs drops well-cited text_mining_agent edges - # below single-PMID automated_agent edges. see the spec at - # docs/specs/response-reduction-strategy-b.md §4.5. + # source_tier is the LEADING key: SemMedDB-style text-mined edges + # are abundant in KG2c but noisy, so every text-mined edge sinks + # below every curated edge regardless of pub count — a well-cited + # text-mined edge can no longer dominate a curated one. within a + # tier, n_pubs is promoted ahead of agent_type because publication + # count is a continuous corroboration signal whereas agent_type is + # a binary provenance label. # # multi-binding handling: a single result can bind multiple # kg edges (one per source asserting the same fact). we score @@ -270,9 +307,10 @@ def _enrich_results( # high-quality source elevates the whole result, even if the # other bindings are weak. all bound edges are retained in # the reduced kg_edges if the result survives the top-N. - def edge_sort_key(eid: str) -> tuple[int, int, int, str]: + def edge_sort_key(eid: str) -> tuple[int, int, int, int, str]: e = kg_edges[eid] return ( + _source_tier(e), _rank_knowledge_level(e), -_n_publications(e), _rank_agent_type(e), @@ -299,29 +337,24 @@ def edge_sort_key(eid: str) -> tuple[int, int, int, str]: return enriched -def reduce_plover_response( - plover_body: dict[str, Any], - *, - top_n_per_predicate: int, - logger: logging.Logger, -) -> ReductionResult: - # entry point. see docs/specs/response-reduction-strategy-b.md. - # this is a pure function: it does not mutate plover_body, does - # not write to disk, and does not call any network service. - - # boundary guards: programmer-error cases. callers MUST pass a - # dict body and a positive top-N; anything else is a bug worth - # surfacing at the call site, not a runtime degradation we silently - # paper over. - if not isinstance(plover_body, dict): - raise TypeError( - f"plover_body must be a dict (got {type(plover_body).__name__})" - ) - if top_n_per_predicate < 1: - raise ValueError( - f"top_n_per_predicate must be >= 1 (got {top_n_per_predicate!r})" - ) +@dataclass(frozen=True) +class _ParsedBody: + # the TRAPI sub-blocks pulled out of a PloverDB body once, so the + # single-shot reducer and the chunker parse identically. + message: dict[str, Any] + raw_results: list[Any] + kg_nodes: dict[str, Any] + kg_edges: dict[str, dict[str, Any]] + aux_graphs: dict[str, Any] + query_graph: dict[str, Any] + original_result_count: int + original_edge_count: int + original_node_count: int + +def _parse_body(plover_body: dict[str, Any]) -> _ParsedBody: + # defensive extraction: any block may be missing or the wrong type in a + # malformed message, so each falls back to an empty container. message = plover_body.get("message") if not isinstance(message, dict): message = {} @@ -341,123 +374,171 @@ def reduce_plover_response( if not isinstance(aux_graphs_in, dict): aux_graphs_in = {} query_graph = message.get("query_graph") or {"nodes": {}, "edges": {}} + return _ParsedBody( + message=message, + raw_results=raw_results, + kg_nodes=kg_nodes_in, + kg_edges=kg_edges_in, + aux_graphs=aux_graphs_in, + query_graph=query_graph, + original_result_count=len(raw_results), + original_edge_count=len(kg_edges_in), + original_node_count=len(kg_nodes_in), + ) - original_result_count = len(raw_results) - original_edge_count = len(kg_edges_in) - original_node_count = len(kg_nodes_in) - - # no-op fast paths. either case means "nothing to score, return - # the body unchanged and record matching counts". strategy is - # still recorded as "B" so the artifact shape stays uniform across - # runs (some runs being labelled "B" and others "no-op" would - # complicate the eval harness for no benefit). - if not raw_results or not kg_edges_in: - if raw_results and not kg_edges_in: - # results exist but reference no kg.edges — TRAPI-malformed - # at the source. one WARNING is enough; per-result warnings - # would be noisy. - logger.warning( - "reduction results present but knowledge_graph.edges " - "is empty; returning body unchanged" - ) - metadata = ReductionMetadata( - strategy_applied=_STRATEGY_NAME, - top_n_per_predicate=top_n_per_predicate, - original_result_count=original_result_count, - original_edge_count=original_edge_count, - original_node_count=original_node_count, - reduced_result_count=0 if not raw_results else 0, - reduced_edge_count=0, - reduced_node_count=0, - predicate_groups=[], - edges_kept_per_group={}, - edges_dropped_per_group={}, - ) - # echo the body verbatim. callers MUST treat the returned dict - # as read-only; we don't deepcopy here because the function is - # documented as not mutating its input. - return ReductionResult(reduced_body=plover_body, metadata=metadata) - # enrich + group - enriched = _enrich_results(raw_results, kg_edges_in, logger) +def _sorted_groups( + raw_results: list[Any], + kg_edges: dict[str, dict[str, Any]], + logger: logging.Logger, +) -> tuple[dict[str, list[_Enriched]], list[str], list[_Enriched]]: + # enrich + group by predicate + sort each group by sort_key (strongest + # first), WITHOUT truncating. this is the un-fused "sort" half: the + # single-shot reducer then takes top-N per group, while the chunker + # partitions the flat ordered list by token budget. predicate groups + # are iterated alphabetically so the byte ordering is stable. + enriched = _enrich_results(raw_results, kg_edges, logger) groups: dict[str, list[_Enriched]] = {} for row in enriched: groups.setdefault(row.primary_predicate, []).append(row) - - # sort each group by the 4-key tuple ascending (stronger first), - # then take top-N. predicate_groups is iterated in alphabetical - # order so the reduced_body's byte ordering is stable across runs - # — important for OpenRouter prompt-cache hit rates. - kept_rows: list[_Enriched] = [] - edges_kept_per_group: dict[str, int] = {} - edges_dropped_per_group: dict[str, int] = {} - predicate_groups_sorted: list[str] = sorted(groups.keys()) + predicate_groups_sorted = sorted(groups.keys()) + ordered_rows: list[_Enriched] = [] for predicate in predicate_groups_sorted: - group = sorted(groups[predicate], key=lambda r: r.sort_key) - retained = group[: top_n_per_predicate] - kept_rows.extend(retained) - edges_kept_per_group[predicate] = len(retained) - edges_dropped_per_group[predicate] = max(0, len(group) - len(retained)) + groups[predicate] = sorted(groups[predicate], key=lambda r: r.sort_key) + ordered_rows.extend(groups[predicate]) + return groups, predicate_groups_sorted, ordered_rows - # rebuild knowledge_graph + auxiliary_graphs from the surviving - # rows. all bound edges of each surviving result are retained (so a - # result with multi-source corroboration keeps every source in the - # reduced kg_edges, not just the one we scored by). edge ids are - # deduped across results via set. + +def _rebuild_body( + rows: list[_Enriched], + parsed: _ParsedBody, + plover_body: dict[str, Any], +) -> dict[str, Any]: + # rebuild a valid TRAPI body containing ONLY the nodes/edges/aux the + # given result rows reference. ALL bound edges of each row are retained + # (multi-source corroboration kept). keys emitted in sorted order so the + # body is byte-stable. shared by the single-shot reducer (top-N rows) + # and the chunker (per-chunk rows) — so a chunk reads exactly like a + # full reduced body. retained_edge_ids: set[str] = set() - for row in kept_rows: + for row in rows: retained_edge_ids.update(row.bound_edge_ids) retained_node_ids: set[str] = set() - for row in kept_rows: + for row in rows: retained_node_ids.update(row.node_ids_in_bindings) retained_aux_ids: set[str] = set() for eid in retained_edge_ids: - retained_aux_ids.update(_support_graph_ids(kg_edges_in[eid])) - - # keys are emitted in sorted order so the reduced body is byte- - # stable across runs even when Python's dict insertion order would - # have differed (e.g. across interpreter restarts with different - # hash seeds — Python 3.7+ preserves order but the SOURCE order - # came from PloverDB and is itself nondeterministic). + retained_aux_ids.update(_support_graph_ids(parsed.kg_edges[eid])) reduced_kg_nodes = { - nid: kg_nodes_in[nid] + nid: parsed.kg_nodes[nid] for nid in sorted(retained_node_ids) - if nid in kg_nodes_in + if nid in parsed.kg_nodes } reduced_kg_edges = { - eid: kg_edges_in[eid] for eid in sorted(retained_edge_ids) + eid: parsed.kg_edges[eid] for eid in sorted(retained_edge_ids) } reduced_aux_graphs = { - gid: aux_graphs_in[gid] + gid: parsed.aux_graphs[gid] for gid in sorted(retained_aux_ids) - if gid in aux_graphs_in + if gid in parsed.aux_graphs } - # results keep their group-ordered insertion order (sorted groups, - # each group sorted by sort_key) so the LLM sees the strongest - # evidence FIRST inside each predicate block. - reduced_results = [row.original for row in kept_rows] - reduced_message: dict[str, Any] = { - **message, - "query_graph": query_graph, + **parsed.message, + "query_graph": parsed.query_graph, "knowledge_graph": { "nodes": reduced_kg_nodes, "edges": reduced_kg_edges, }, - "results": reduced_results, + "results": [row.original for row in rows], "auxiliary_graphs": reduced_aux_graphs, } - reduced_body: dict[str, Any] = {**plover_body, "message": reduced_message} + return {**plover_body, "message": reduced_message} + + +def reduce_plover_response( + plover_body: dict[str, Any], + *, + top_n_per_predicate: int, + logger: logging.Logger, +) -> ReductionResult: + # single-shot entry point. see docs/specs/response-reduction-strategy-b.md. + # pure function: does not mutate plover_body, write to disk, or call any + # network service. sort + truncate are now un-fused (_sorted_groups does + # the sort; this function does the top-N truncation) so the chunker can + # consume the same sorted ranking without truncation. + + # boundary guards: programmer-error cases. callers MUST pass a dict body + # and a positive top-N; anything else is a bug worth surfacing at the + # call site, not a runtime degradation we silently paper over. + if not isinstance(plover_body, dict): + raise TypeError( + f"plover_body must be a dict (got {type(plover_body).__name__})" + ) + if top_n_per_predicate < 1: + raise ValueError( + f"top_n_per_predicate must be >= 1 (got {top_n_per_predicate!r})" + ) + + parsed = _parse_body(plover_body) + + # no-op fast paths. either case means "nothing to score, return the body + # unchanged and record matching counts". strategy is still recorded as + # "B" so the artifact shape stays uniform across runs. + if not parsed.raw_results or not parsed.kg_edges: + if parsed.raw_results and not parsed.kg_edges: + # results exist but reference no kg.edges — TRAPI-malformed at + # the source. one WARNING is enough; per-result would be noisy. + logger.warning( + "reduction results present but knowledge_graph.edges " + "is empty; returning body unchanged" + ) + metadata = ReductionMetadata( + strategy_applied=_STRATEGY_NAME, + top_n_per_predicate=top_n_per_predicate, + original_result_count=parsed.original_result_count, + original_edge_count=parsed.original_edge_count, + original_node_count=parsed.original_node_count, + reduced_result_count=0, + reduced_edge_count=0, + reduced_node_count=0, + predicate_groups=[], + edges_kept_per_group={}, + edges_dropped_per_group={}, + ) + # echo the body verbatim. callers MUST treat the returned dict as + # read-only; we don't deepcopy because the function does not mutate. + return ReductionResult(reduced_body=plover_body, metadata=metadata) + + groups, predicate_groups_sorted, _ = _sorted_groups( + parsed.raw_results, parsed.kg_edges, logger, + ) + + # truncate top-N per (already-sorted) predicate group. + kept_rows: list[_Enriched] = [] + edges_kept_per_group: dict[str, int] = {} + edges_dropped_per_group: dict[str, int] = {} + for predicate in predicate_groups_sorted: + group = groups[predicate] + retained = group[: top_n_per_predicate] + kept_rows.extend(retained) + edges_kept_per_group[predicate] = len(retained) + edges_dropped_per_group[predicate] = max(0, len(group) - len(retained)) + + reduced_body = _rebuild_body(kept_rows, parsed, plover_body) + reduced_kg = reduced_body["message"]["knowledge_graph"] + reduced_result_count = len(reduced_body["message"]["results"]) + reduced_edge_count = len(reduced_kg["edges"]) + reduced_node_count = len(reduced_kg["nodes"]) metadata = ReductionMetadata( strategy_applied=_STRATEGY_NAME, top_n_per_predicate=top_n_per_predicate, - original_result_count=original_result_count, - original_edge_count=original_edge_count, - original_node_count=original_node_count, - reduced_result_count=len(reduced_results), - reduced_edge_count=len(reduced_kg_edges), - reduced_node_count=len(reduced_kg_nodes), + original_result_count=parsed.original_result_count, + original_edge_count=parsed.original_edge_count, + original_node_count=parsed.original_node_count, + reduced_result_count=reduced_result_count, + reduced_edge_count=reduced_edge_count, + reduced_node_count=reduced_node_count, predicate_groups=predicate_groups_sorted, edges_kept_per_group=edges_kept_per_group, edges_dropped_per_group=edges_dropped_per_group, @@ -466,10 +547,120 @@ def reduce_plover_response( logger.info( f"[bold cyan]→ reduction[/] strategy={_STRATEGY_NAME} " f"predicates={len(predicate_groups_sorted)} " - f"results={original_result_count}→{len(reduced_results)} " - f"edges={original_edge_count}→{len(reduced_kg_edges)} " - f"nodes={original_node_count}→{len(reduced_kg_nodes)} " + f"results={parsed.original_result_count}→{reduced_result_count} " + f"edges={parsed.original_edge_count}→{reduced_edge_count} " + f"nodes={parsed.original_node_count}→{reduced_node_count} " f"top_n={top_n_per_predicate}" ) return ReductionResult(reduced_body=reduced_body, metadata=metadata) + + +_CHARS_PER_TOKEN = 4 # rough JSON+English heuristic for the chunk-budget estimate + + +@dataclass(frozen=True) +class ChunkSet: + # the result of partitioning the full sorted (untruncated) ranking into + # token-budgeted chunks. each chunk is a valid TRAPI sub-body the LLM + # reads exactly like a full reduced body. chunk 0 holds the strongest + # evidence (curated first, text-mined last) so an early-stopping reader + # sees the best answers first. + chunks: list[dict[str, Any]] + n_chunks: int + total_rows: int # sorted result rows available + chunked_rows: int # rows actually placed into chunks + truncated_at_max_chunks: bool # True if max_chunks cut off the tail + + +def _row_weight(row: _Enriched, parsed: _ParsedBody) -> int: + # rough char-count proxy for a row's contribution to a chunk body: its + # result plus its bound edges plus its bound nodes. shared nodes across + # rows in the same chunk are double-counted, which makes the estimate an + # UPPER bound, so chunks land at or under budget (safe). + weight = len(json.dumps(row.original, ensure_ascii=False)) + for eid in row.bound_edge_ids: + edge = parsed.kg_edges.get(eid) + if edge is not None: + weight += len(json.dumps(edge, ensure_ascii=False)) + for nid in row.node_ids_in_bindings: + node = parsed.kg_nodes.get(nid) + if node is not None: + weight += len(json.dumps(node, ensure_ascii=False)) + return weight + + +def chunk_plover_response( + plover_body: dict[str, Any], + *, + chunk_token_budget: int, + max_chunks: int, + logger: logging.Logger, +) -> ChunkSet: + # partition the full sorted-but-UNTRUNCATED ranking into chunks bounded + # by chunk_token_budget. each chunk is a valid TRAPI sub-body (built by + # _rebuild_body) containing only the nodes/edges/aux its results + # reference. the iterative Stage-11 loop reads chunks in order until it + # is confident, so nothing is truncated up front — only max_chunks caps + # the total. pure function: no IO, no network. + if not isinstance(plover_body, dict): + raise TypeError( + f"plover_body must be a dict (got {type(plover_body).__name__})" + ) + if chunk_token_budget < 1: + raise ValueError( + f"chunk_token_budget must be >= 1 (got {chunk_token_budget!r})" + ) + if max_chunks < 1: + raise ValueError(f"max_chunks must be >= 1 (got {max_chunks!r})") + + parsed = _parse_body(plover_body) + if not parsed.raw_results or not parsed.kg_edges: + # nothing to chunk — hand back the body as a single chunk so the + # caller's loop runs exactly once and terminates cleanly. + return ChunkSet( + chunks=[plover_body], n_chunks=1, total_rows=0, + chunked_rows=0, truncated_at_max_chunks=False, + ) + + _, _, ordered_rows = _sorted_groups(parsed.raw_results, parsed.kg_edges, logger) + + char_budget = chunk_token_budget * _CHARS_PER_TOKEN + chunk_rows: list[list[_Enriched]] = [] + current: list[_Enriched] = [] + current_weight = 0 + for row in ordered_rows: + weight = _row_weight(row, parsed) + # close the current chunk before it would exceed the budget, but + # never emit an empty chunk (a single oversized row gets its own). + if current and current_weight + weight > char_budget: + chunk_rows.append(current) + current = [] + current_weight = 0 + if len(chunk_rows) >= max_chunks: + break + current.append(row) + current_weight += weight + if current and len(chunk_rows) < max_chunks: + chunk_rows.append(current) + + chunked = sum(len(rows) for rows in chunk_rows) + truncated = chunked < len(ordered_rows) + if truncated: + logger.warning( + f"chunking max_chunks={max_chunks} reached; " + f"{len(ordered_rows) - chunked} of {len(ordered_rows)} sorted " + f"results were not chunked" + ) + chunks = [_rebuild_body(rows, parsed, plover_body) for rows in chunk_rows] + logger.info( + f"[bold cyan]→ chunking[/] results={len(ordered_rows)} " + f"chunks={len(chunks)} budget={chunk_token_budget}tok" + ) + return ChunkSet( + chunks=chunks, + n_chunks=len(chunks), + total_rows=len(ordered_rows), + chunked_rows=chunked, + truncated_at_max_chunks=truncated, + ) diff --git a/pipeline/code/tests/test_chunking.py b/pipeline/code/tests/test_chunking.py new file mode 100644 index 0000000..2176f75 --- /dev/null +++ b/pipeline/code/tests/test_chunking.py @@ -0,0 +1,150 @@ +# chunk_plover_response: the iterative-mode chunking primitive. +# +# partitions the full sorted-but-untruncated ranking into token-budgeted +# valid TRAPI sub-bodies. pins: each chunk is a closed sub-body (every edge +# its results reference is present), the union of chunks covers every +# result when not truncated, the strongest evidence (curated before +# text-mined) lands in chunk 0, and max_chunks caps + flags the tail. pure +# function, no network. + +import logging +from typing import Any + +from code.reduction import chunk_plover_response + +_LOG = logging.getLogger("test_chunking") + + +def _attr(type_id: str, value: object) -> dict[str, object]: + return {"attribute_type_id": type_id, "value": value} + + +def _edge( + subject: str, + object_: str, + *, + predicate: str = "biolink:treats", + agent_type: str = "manual_agent", +) -> dict[str, object]: + return { + "subject": subject, + "object": object_, + "predicate": predicate, + "attributes": [ + _attr("biolink:knowledge_level", "knowledge_assertion"), + _attr("biolink:agent_type", agent_type), + ], + } + + +def _result(edge_id: str, subject: str, object_: str) -> dict[str, object]: + return { + "node_bindings": {"n0": [{"id": subject}], "n1": [{"id": object_}]}, + "analyses": [{"edge_bindings": {"e0": [{"id": edge_id}]}}], + } + + +def _body( + edges: dict[str, object], results: list[dict[str, object]], nodes: dict[str, object] +) -> dict[str, object]: + return { + "message": { + "query_graph": {"nodes": {}, "edges": {}}, + "knowledge_graph": {"nodes": nodes, "edges": edges}, + "results": results, + "auxiliary_graphs": {}, + } + } + + +def _n_drug_body(n: int) -> dict[str, object]: + # n results: drug_i --treats--> the same pinned disease. + edges: dict[str, object] = {} + results: list[dict[str, object]] = [] + nodes: dict[str, object] = {"MONDO:1": {"name": "disease"}} + for i in range(n): + drug = f"CHEBI:{i}" + edges[f"e{i}"] = _edge(drug, "MONDO:1") + results.append(_result(f"e{i}", drug, "MONDO:1")) + nodes[drug] = {"name": f"drug{i}"} + return _body(edges, results, nodes) + + +def _chunk_edge_ids(chunk: dict[str, Any]) -> set[str]: + return set(chunk["message"]["knowledge_graph"]["edges"].keys()) + + +def test_huge_budget_yields_single_chunk(): + out = chunk_plover_response( + _n_drug_body(5), chunk_token_budget=100_000, max_chunks=50, logger=_LOG, + ) + assert out.n_chunks == 1 + assert out.total_rows == 5 + assert out.chunked_rows == 5 + assert out.truncated_at_max_chunks is False + assert len(_chunk_edge_ids(out.chunks[0])) == 5 + + +def test_tiny_budget_splits_one_row_per_chunk_and_covers_all(): + # budget=10 tokens (40 chars) is below any single row's weight, so each + # row gets its own chunk; the union must still cover every result edge. + out = chunk_plover_response( + _n_drug_body(5), chunk_token_budget=10, max_chunks=50, logger=_LOG, + ) + assert out.n_chunks == 5 + covered: set[str] = set() + for chunk in out.chunks: + covered |= _chunk_edge_ids(chunk) + assert covered == {f"e{i}" for i in range(5)} + + +def test_strongest_evidence_lands_in_first_chunk(): + # a curated edge and a text-mined edge under the same predicate; with a + # tiny budget each is its own chunk. source_tier demotion means the + # curated edge sorts first, so it must be in chunk 0. + edges = { + "curated": _edge("CHEBI:1", "MONDO:1", agent_type="manual_agent"), + "textmined": _edge("CHEBI:2", "MONDO:1", agent_type="text_mining_agent"), + } + results = [_result("curated", "CHEBI:1", "MONDO:1"), + _result("textmined", "CHEBI:2", "MONDO:1")] + nodes = {"MONDO:1": {}, "CHEBI:1": {}, "CHEBI:2": {}} + out = chunk_plover_response( + _body(edges, results, nodes), chunk_token_budget=10, max_chunks=50, logger=_LOG, + ) + assert out.n_chunks == 2 + assert "curated" in _chunk_edge_ids(out.chunks[0]) + assert "textmined" in _chunk_edge_ids(out.chunks[1]) + + +def test_max_chunks_caps_and_flags_truncation(): + out = chunk_plover_response( + _n_drug_body(5), chunk_token_budget=10, max_chunks=2, logger=_LOG, + ) + assert out.n_chunks == 2 + assert out.total_rows == 5 + assert out.chunked_rows == 2 + assert out.truncated_at_max_chunks is True + + +def test_each_chunk_is_a_closed_subbody(): + # every edge a chunk's results bind must exist in that chunk's + # knowledge_graph.edges (the LLM can read it standalone). + out = chunk_plover_response( + _n_drug_body(6), chunk_token_budget=10, max_chunks=50, logger=_LOG, + ) + for chunk in out.chunks: + msg = chunk["message"] + edge_ids = set(msg["knowledge_graph"]["edges"].keys()) + for result in msg["results"]: + for binding_list in result["analyses"][0]["edge_bindings"].values(): + for binding in binding_list: + assert binding["id"] in edge_ids + + +def test_noop_body_returns_single_chunk(): + out = chunk_plover_response( + _body({}, [], {}), chunk_token_budget=1000, max_chunks=50, logger=_LOG, + ) + assert out.n_chunks == 1 + assert out.total_rows == 0 diff --git a/pipeline/code/tests/test_response_reduction.py b/pipeline/code/tests/test_response_reduction.py index 08c94c9..63dc054 100644 --- a/pipeline/code/tests/test_response_reduction.py +++ b/pipeline/code/tests/test_response_reduction.py @@ -46,6 +46,7 @@ def _edge( agent_type: str | None = "manual_agent", pubs: list[str] | None = None, support_graphs: list[str] | None = None, + primary_knowledge_source: str | None = None, ) -> dict[str, object]: # None for kl/at means "attribute deliberately absent" (tests # the missing-attribute degradation path). empty-list pubs means @@ -59,12 +60,20 @@ def _edge( attrs.append(_attr("biolink:publications", pubs)) if support_graphs is not None: attrs.append(_attr("biolink:support_graphs", support_graphs)) - return { + edge: dict[str, object] = { "subject": subject, "object": object_, "predicate": predicate, "attributes": attrs, } + # primary_knowledge_source lives in the TRAPI `sources` block, not in + # attributes — used by the source-tier (text-mined) demotion. + if primary_knowledge_source is not None: + edge["sources"] = [{ + "resource_id": primary_knowledge_source, + "resource_role": "primary_knowledge_source", + }] + return edge def _result( @@ -341,22 +350,19 @@ def test_more_publications_beats_fewer_within_same_kl_and_at(): assert kept == ["well_cited"] -def test_more_pubs_beats_better_agent_type_when_kl_ties(): - # the CFTR ↔ cystic fibrosis case in real data: every edge in a - # predicate group shares the same kl (e.g. all kl=prediction for - # gene-disease edges in KG2c). under the *prior* sort order - # (kl, at, -pubs, edge_id), a 1-PMID automated_agent edge would - # beat a 20-PMID text_mining_agent edge because at was the second - # key. under the corrected order (kl, -pubs, at, edge_id), the - # 20-PMID text-mining edge wins because n_pubs is now ahead of - # at. this test pins the new ordering. +def test_text_mined_demoted_below_curated_regardless_of_pubs(): + # text-mined edges are demoted below curated edges (source_tier is the + # leading sort key), so a 20-PMID text_mining_agent edge now LOSES to a + # 1-PMID automated_agent (curated) edge with the same knowledge_level. + # this is the deliberate reversal of the old "well-cited text-mining + # wins" behaviour — text-mined evidence can no longer dominate. body = _body( results=[ - _result("cftr_cf_text_mined", "NCBIGene:1080", "MONDO:0009061"), - _result("cftr_other_automated", "NCBIGene:1080", "MONDO:0004975"), + _result("cf_text_mined", "NCBIGene:1080", "MONDO:0009061"), + _result("cf_automated", "NCBIGene:1080", "MONDO:0004975"), ], edges={ - "cftr_cf_text_mined": _edge( + "cf_text_mined": _edge( subject="NCBIGene:1080", object_="MONDO:0009061", predicate="biolink:gene_associated_with_condition", @@ -364,7 +370,7 @@ def test_more_pubs_beats_better_agent_type_when_kl_ties(): agent_type="text_mining_agent", pubs=[f"PMID:{i}" for i in range(20)], ), - "cftr_other_automated": _edge( + "cf_automated": _edge( subject="NCBIGene:1080", object_="MONDO:0004975", predicate="biolink:gene_associated_with_condition", @@ -377,20 +383,76 @@ def test_more_pubs_beats_better_agent_type_when_kl_ties(): ) out = _reduce(body, top_n=1) kept = list(out.reduced_body["message"]["knowledge_graph"]["edges"].keys()) - assert kept == ["cftr_cf_text_mined"] + assert kept == ["cf_automated"] + + +def test_text_mined_demoted_even_with_stronger_kl_and_more_pubs(): + # source_tier dominates knowledge_level AND n_pubs: a text_mining edge + # with the strongest kl (knowledge_assertion) and 50 pubs still loses + # to a curated edge with a weaker kl (prediction) and a single pub. + body = _body( + results=[ + _result("tm_strong", "S", "O"), + _result("curated_weak", "S", "O2"), + ], + edges={ + "tm_strong": _edge( + object_="O", knowledge_level="knowledge_assertion", + agent_type="text_mining_agent", + pubs=[f"PMID:{i}" for i in range(50)], + ), + "curated_weak": _edge( + object_="O2", knowledge_level="prediction", + agent_type="automated_agent", pubs=["PMID:1"], + ), + }, + nodes={"S": {}, "O": {}, "O2": {}}, + ) + out = _reduce(body, top_n=1) + kept = list(out.reduced_body["message"]["knowledge_graph"]["edges"].keys()) + assert kept == ["curated_weak"] + + +def test_semmeddb_source_demoted_even_without_text_mining_agent_type(): + # demotion also triggers on the primary knowledge source: an edge + # sourced from infores:semmeddb is text-mined even when its agent_type + # is absent, so it sinks below a curated (drugcentral) edge. + body = _body( + results=[ + _result("semmed", "S", "O"), + _result("curated", "S", "O2"), + ], + edges={ + "semmed": _edge( + object_="O", knowledge_level="knowledge_assertion", + agent_type=None, pubs=[f"PMID:{i}" for i in range(30)], + primary_knowledge_source="infores:semmeddb", + ), + "curated": _edge( + object_="O2", knowledge_level="prediction", + agent_type="automated_agent", pubs=["PMID:1"], + primary_knowledge_source="infores:drugcentral", + ), + }, + nodes={"S": {}, "O": {}, "O2": {}}, + ) + out = _reduce(body, top_n=1) + kept = list(out.reduced_body["message"]["knowledge_graph"]["edges"].keys()) + assert kept == ["curated"] def test_better_agent_type_breaks_tie_when_pubs_equal(): - # when kl AND n_pubs both tie, agent_type is the next tiebreaker - # (still earns its keep — just no longer outranks evidence count). + # within the SAME source tier (both curated), when kl AND n_pubs both + # tie, agent_type is the next tiebreaker: manual_agent beats + # automated_agent. (both non-text-mined, so source_tier ties at 0.) body = _body( results=[ _result("manual_pick", "S", "O"), - _result("nlp_pick", "S", "O"), + _result("auto_pick", "S", "O"), ], edges={ "manual_pick": _edge(agent_type="manual_agent", pubs=["PMID:1"]), - "nlp_pick": _edge(agent_type="text_mining_agent", pubs=["PMID:1"]), + "auto_pick": _edge(agent_type="automated_agent", pubs=["PMID:1"]), }, nodes={"S": {}, "O": {}}, ) @@ -700,7 +762,7 @@ def test_multi_binding_result_scored_by_strongest_bound_edge(): subject="NCBIGene:1080", object_="MONDO:0009061", predicate="biolink:gene_associated_with_condition", knowledge_level="prediction", - agent_type="text_mining_agent", + agent_type="automated_agent", pubs=[f"PMID:{i}" for i in range(20)], ), "other": _edge( From 0a2a2e0b7335158747e1e1e892d333aa266a6693 Mon Sep 17 00:00:00 2001 From: Adilbek Bazarkulov Date: Mon, 8 Jun 2026 12:31:36 -0400 Subject: [PATCH 3/6] Add iterative and context-window config; key gold questions by id --- pipeline/code/config.py | 49 +++++++++++++++++++++- pipeline/code/runner.py | 1 + pipeline/code/tests/test_load_questions.py | 14 +++++++ pipeline/config.yaml | 41 ++++++++++++++---- 4 files changed, 96 insertions(+), 9 deletions(-) create mode 100644 pipeline/code/tests/test_load_questions.py diff --git a/pipeline/code/config.py b/pipeline/code/config.py index 1dbec04..0242e43 100644 --- a/pipeline/code/config.py +++ b/pipeline/code/config.py @@ -53,6 +53,10 @@ class ModelSpec: provider: str # anthropic / google / openai / xai / deepseek price_in: float # USD per 1M input tokens price_out: float # USD per 1M output tokens + # max input context window in tokens, set MANUALLY in config.yaml. drives + # the iterative chunk budget (context_window * context_fraction). adjust + # by hand when a provider changes a model's window. + context_window: int = 128000 @dataclass(frozen=True) @@ -86,6 +90,24 @@ class ReductionConfig: top_n_per_predicate: int +@dataclass(frozen=True) +class Stage11IterativeConfig: + # iterative chunked answer-picking. default OFF: the single-shot + # Strategy B path runs unless this is enabled, so the baseline stays + # runnable for the A/B comparison. when enabled, the full sorted + # PloverDB response is split into chunks the LLM reads strongest-first + # until it has good answers, instead of truncating to top_n_per_predicate. + enabled: bool + # the chunk budget is dynamic per model: context_window * context_fraction + # tokens. 0.8 leaves headroom for the system prompt and the completion. + context_fraction: float + # soft target for how many answers to return. the loop stops once it has + # this many good picks (or the LLM is satisfied with fewer); the final + # set is clamped to this many. NOT a hard truncation of the response. + answer_target: int + max_chunks: int + + @dataclass(frozen=True) class Paths: # absolute filesystem locations derived from the YAML's relative paths. @@ -101,6 +123,11 @@ class Config: generation: Generation paths: Paths reduction: ReductionConfig + stage11_iterative: Stage11IterativeConfig + # Stage 1 scope-check guardrail. default OFF: it is an extra OpenRouter + # call per query (and a rate-limit source). when off, every input is + # treated as in-scope; scope is checked manually instead. + scope_check_enabled: bool def model(self, model_id: str) -> ModelSpec: # tiny lookup helper. raising explicitly here means a typo in @@ -138,12 +165,23 @@ def load_config() -> Config: reduction = ReductionConfig( top_n_per_predicate=int(reduction_raw.get("top_n_per_predicate", 10)), ) + # stage11.iterative is optional and defaults to OFF, so existing + # configs (and the single-shot baseline) load unchanged. + iterative_raw = (raw.get("stage11") or {}).get("iterative") or {} + stage11_iterative = Stage11IterativeConfig( + enabled=bool(iterative_raw.get("enabled", False)), + context_fraction=float(iterative_raw.get("context_fraction", 0.8)), + answer_target=int(iterative_raw.get("answer_target", 5)), + max_chunks=int(iterative_raw.get("max_chunks", 16)), + ) return Config( endpoints=eps, models=models, generation=gen, paths=paths, reduction=reduction, + stage11_iterative=stage11_iterative, + scope_check_enabled=bool((raw.get("scope_check") or {}).get("enabled", False)), ) @@ -163,4 +201,13 @@ def load_questions(cfg: Config) -> list[dict[str, Any]]: qdir.glob("q*.json"), key=lambda p: int(p.stem[1:]) if p.stem[1:].isdigit() else 1_000_000, ) - return [json.loads(p.read_text()) for p in paths] + # gold files key the id as `question_id`, but the pipeline and runner + # read `q["id"]` everywhere (the ad-hoc path also uses id="adhoc"), so + # normalise to `id` at load time. without this, `runner --questions q1` + # KeyErrors on q["id"] and the gold set cannot run at all. + out: list[dict[str, Any]] = [] + for p in paths: + record = json.loads(p.read_text()) + record.setdefault("id", record.get("question_id")) + out.append(record) + return out diff --git a/pipeline/code/runner.py b/pipeline/code/runner.py index 1f94893..74fc505 100644 --- a/pipeline/code/runner.py +++ b/pipeline/code/runner.py @@ -381,6 +381,7 @@ def main() -> int: "slug": m.slug, "tier": m.tier, "provider": m.provider, + "context_window": m.context_window, } for m in chosen_models ], diff --git a/pipeline/code/tests/test_load_questions.py b/pipeline/code/tests/test_load_questions.py new file mode 100644 index 0000000..03d6547 --- /dev/null +++ b/pipeline/code/tests/test_load_questions.py @@ -0,0 +1,14 @@ +# load_questions must normalise the gold files' `question_id` field to the +# `id` key the pipeline + runner read everywhere (q["id"]). the gold files +# key it as `question_id`; without normalisation the runner KeyErrors on +# `--questions q1` and the gold set cannot run. pure file read, no network. + +from code.config import load_config, load_questions + + +def test_every_gold_question_exposes_id_matching_question_id(): + qs = load_questions(load_config()) + assert qs, "no gold questions loaded" + for q in qs: + assert q.get("id"), f"gold question missing id (question_id={q.get('question_id')})" + assert q["id"] == q.get("question_id") diff --git a/pipeline/config.yaml b/pipeline/config.yaml index 8f7a94d..7ac1678 100644 --- a/pipeline/config.yaml +++ b/pipeline/config.yaml @@ -23,15 +23,18 @@ endpoints: # 4 western providers (anthropic, google, openai, xai), so the # frontier-vs-budget comparison is held constant on provider identity. # DeepSeek (Chinese) is intentionally excluded. +# context_window is the max input tokens per model, set MANUALLY here. it +# drives the iterative chunk budget (context_window * stage11.context_fraction). +# adjust by hand when a provider changes a model's window. models: - - { id: m1, tier: frontier, slug: "anthropic/claude-opus-4.6", provider: anthropic, price_in: 5.00, price_out: 25.00 } - - { id: m2, tier: frontier, slug: "google/gemini-3.1-pro-preview", provider: google, price_in: 2.00, price_out: 12.00 } - - { id: m3, tier: frontier, slug: "openai/gpt-5.5", provider: openai, price_in: 5.00, price_out: 30.00 } - - { id: m4, tier: frontier, slug: "x-ai/grok-4", provider: xai, price_in: 3.00, price_out: 15.00 } - - { id: m5, tier: budget, slug: "anthropic/claude-haiku-4.5", provider: anthropic, price_in: 1.00, price_out: 5.00 } - - { id: m6, tier: budget, slug: "google/gemini-3-flash-preview", provider: google, price_in: 0.50, price_out: 3.00 } - - { id: m7, tier: budget, slug: "openai/gpt-5", provider: openai, price_in: 1.25, price_out: 10.00 } - - { id: m8, tier: budget, slug: "x-ai/grok-4.3", provider: xai, price_in: 1.25, price_out: 2.50 } + - { id: m1, tier: frontier, slug: "anthropic/claude-opus-4.6", provider: anthropic, price_in: 5.00, price_out: 25.00, context_window: 1000000 } + - { id: m2, tier: frontier, slug: "google/gemini-3.1-pro-preview", provider: google, price_in: 2.00, price_out: 12.00, context_window: 1048576 } + - { id: m3, tier: frontier, slug: "openai/gpt-5.5", provider: openai, price_in: 5.00, price_out: 30.00, context_window: 1050000 } + - { id: m4, tier: frontier, slug: "x-ai/grok-4", provider: xai, price_in: 3.00, price_out: 15.00, context_window: 256000 } + - { id: m5, tier: budget, slug: "anthropic/claude-haiku-4.5", provider: anthropic, price_in: 1.00, price_out: 5.00, context_window: 200000 } + - { id: m6, tier: budget, slug: "google/gemini-3-flash-preview", provider: google, price_in: 0.50, price_out: 3.00, context_window: 1048576 } + - { id: m7, tier: budget, slug: "openai/gpt-5", provider: openai, price_in: 1.25, price_out: 10.00, context_window: 400000 } + - { id: m8, tier: budget, slug: "x-ai/grok-4.3", provider: xai, price_in: 1.25, price_out: 2.50, context_window: 1000000 } # LLM sampling parameters. shared across every stage of every model. # temperature=0 where the provider supports it; max_tokens caps any @@ -65,5 +68,27 @@ paths: # top_n_per_predicate is the v1 default. the benchmark sweeps it # across {5, 10, 15, 20} to characterise the answer-quality vs # token-cost curve before locking the published default. +# Stage 1 scope-check guardrail (an LLM call that refuses non-biomedical +# input). default OFF: it is an extra OpenRouter call per query and a +# rate-limit source. when off, every input is treated as in-scope and flows +# to entity extraction; scope is checked manually. re-enable for the +# irrelevant-probe benchmark or to refuse chit-chat in the live UI. +scope_check: + enabled: true + reduction: top_n_per_predicate: 10 + +# iterative chunked answer-picking (Stage 11). default OFF — the +# single-shot Strategy B path above runs unless this is enabled, so the +# baseline stays runnable for the A/B comparison. when enabled, the full +# sorted PloverDB response is split into token-budgeted chunks the LLM +# reads in order until confident, instead of truncating to +# top_n_per_predicate (which can drop correct answers before the LLM +# ever sees them). +stage11: + iterative: + enabled: true + context_fraction: 0.8 # chunk budget = model.context_window * this + answer_target: 5 # stop once this many good answers are found + max_chunks: 16 From 59de28299141808fa6613bdd71c23365a63a7ca5 Mon Sep 17 00:00:00 2001 From: Adilbek Bazarkulov Date: Mon, 8 Jun 2026 12:31:49 -0400 Subject: [PATCH 4/6] Rank answer selection by relevance and tidy explanation output --- pipeline/code/prompts.py | 143 +++++++++++++++++++++++++++------------ 1 file changed, 101 insertions(+), 42 deletions(-) diff --git a/pipeline/code/prompts.py b/pipeline/code/prompts.py index 30031a4..22630cf 100644 --- a/pipeline/code/prompts.py +++ b/pipeline/code/prompts.py @@ -277,8 +277,9 @@ Each entry carries `curie`, `label`, `types` (Biolink categories), `bm25_score` (raw NameRes score; kept for traceability — the LIST ORDER is the rerank tier, NOT the BM25 score), and (when available) - `kg2c_edges_to_` (live count of edges in KG2c from - that CURIE to the answer category, used for fix (e) below). + `kg_edges_to_` (live count of edges in the hosted + knowledge graph from that CURIE to the answer category, used for + fix (e) below). Decide: @@ -313,21 +314,21 @@ return chosen_curie=null with a reason so the pipeline can fail loudly rather than silently grounding to garbage. - (e) KG2c COVERAGE: each candidate may carry a `kg2c_edges_to_=N` + (e) GRAPH COVERAGE: each candidate may carry a `kg_edges_to_=N` count, measured live from PloverDB at query time. This is the number - of edges in KG2c connecting that CURIE to any node of the answer - category (either direction). When two candidates are semantically - close, PREFER the one with non-zero KG2c edges — picking a CURIE - with `kg2c_edges_to_*=0` will return no results downstream and the - run will fail with `outcome=no_results`. A slightly lower-ranked - candidate that actually has KG2c coverage beats a perfect-label one - with no data. Concrete case: for "cholesterol biosynthesis", the - top BM25 result is often PANTHER.PATHWAY:P00014 (Pathway type, - perfect label) but with `kg2c_edges_to_biolink:Gene=0`, while the - lower-ranked GO:0006695 or REACT:R-HSA-191273 entry has dozens of - edges and is what KG2c actually populates. If the counts are not - provided (older runs, probe disabled), fall back to label+score - picking as before. + of edges in the hosted knowledge graph connecting that CURIE to any + node of the answer category (either direction). When two candidates + are semantically close, PREFER the one with non-zero edges — picking + a CURIE with `kg_edges_to_*=0` will return no results downstream and + the run will fail with `outcome=no_results`. The candidate ORDER + already prefers non-zero coverage among same-type candidates, so + trust the order; a perfect-label candidate with no edges has been + deprioritised for you. Concrete case: for "cholesterol biosynthesis", + the top BM25 result is often PANTHER.PATHWAY:P00014 (Pathway type, + perfect label) but with `kg_edges_to_biolink:Gene=0`, while the + GO:0006695 or REACT:R-HSA-191273 entry has dozens of edges and is + what the graph actually populates. If the counts are not provided + (older runs, probe disabled), fall back to label+score picking. 3. Return exactly one of these JSON objects on a single line: @@ -483,6 +484,69 @@ """ +# stage 11 (iterative mode): the PloverDB response is split into ordered +# chunks the LLM reads until confident, instead of truncating to top-N. +# the model accumulates picks across chunks (may overturn), declares the +# expected answer count (variable N, no fixed cap), and signals when it has +# enough. selection is re-validated in code against each chunk's edges. +SYS_ANSWER_PICK_ITER = """You are PloverAI's relevance-ranking answer selector. + +The PloverDB response may be too large for one message, so you read it in +ordered CHUNKS — strongest evidence first, weakest (text-mined) last. Each +turn you receive: +- The user's question. +- A target number of answers to return (an upper bound). +- shortlist_so_far: your current best-ranked answers from earlier chunks, + each with its one-line relevance reason (empty on the first chunk). This is + your running ranking — merge this chunk's candidates into it. +- ONE chunk: a TRAPI sub-response with its own knowledge_graph.nodes/edges. + +RANK BY RELEVANCE FIRST. Your primary job is to choose the answers most +RELEVANT and REPRESENTATIVE of what the question actually asks. Evidence +strength (knowledge_level, n_publications, source) is ONLY a TIE-BREAKER: +use it to choose between candidates that are equally relevant. NEVER drop a +clearly more relevant answer in favour of a better-documented but less +relevant one. (Example: for "what cells are in the brain", a neuron is more +relevant than a brain-vasculature smooth-muscle cell even if the latter has +a better-cited edge.) + +For a "list / which / what X" question, prefer a REPRESENTATIVE spread across +the distinct answer types over several near-duplicates of one type. + +Hard constraints: +- Every answer's curie must appear in THIS chunk's knowledge_graph.nodes OR + already in shortlist_so_far. Never invent a CURIE. +- Merge this chunk's candidates into shortlist_so_far: a more relevant new + candidate may displace a less relevant one. Return the FULL updated, ranked + shortlist each turn (most relevant first), UP TO the target. Returning + fewer is fine; do not pad to the target. +- Text-mined edges are ordered last and are LOW-CONFIDENCE. You MAY pick one + if it is the most relevant answer and nothing better exists, but do not + prefer it over an equally relevant non-text-mined edge. + +When to stop vs. read another chunk: +- Set confidence_sufficient = true when you are confident you have seen the + candidates needed to rank the most relevant answers — e.g. you have read + the whole response, or the remaining (weaker-evidence) chunks are unlikely + to hold a MORE RELEVANT answer than your current shortlist. +- Set confidence_sufficient = false when a relevant answer might still be in + a later chunk and you want to see it before finalising the ranking. + +Output a single JSON object: +{ + "answers": [ + { "curie": "...", "label": "...", "why": "", + "supporting_edge_ids": ["..."] } + ], + "evidence_tier": "", + "confidence_sufficient": , + "rationale": "" +} + +Return ONLY the JSON object. No markdown fences, no commentary. +""" + + # stage 15: write the user-facing explanation as structured Markdown. # the four-section template (Answer / Evidence / Confidence / Limitations) # maps 1:1 to the cards in the research-grade result UI so the LLM's @@ -554,38 +618,33 @@ supporting edge is rather than curator-attested; treat as a research lead, not an established fact". -If ALL picked entities are MODERATE or WEAK, lead the **Answer** section -with a GROUNDED caveat that uses the pipeline-context numbers and the -specific provenance profile you observed. Template: - - "PloverDB returned [plover_total_results] edges for the query - ([pinned_entity] --[predicate_used]--> [answer_category]). Strategy B - reduction kept [reduction_results_kept]. The [picked_edges_count] - edges selected to answer share the same provenance profile: - knowledge_level=[X], agent_type=[Y], primary_knowledge_source=[Z] - (null if none recorded), supporting_publications=[N] (zero if - empty). This means [one short sentence explaining what the absence - of pks / PMIDs implies — e.g. 'KG2c does not record which source - database originated these edges or which papers support them, so - independent verification of these claims is not possible from the - pipeline output alone']. Treat each entry below as a research lead, - not an established fact." - -Fill in the bracketed values from the pipeline-context block and the -actual edge attributes. Do not omit numbers. Do not paraphrase -"PloverDB returned N edges" into "the knowledge graph returned -candidates" — name the source and quote the count. +If ALL picked entities are MODERATE or WEAK, HEDGE the phrasing in the +**Answer** section (per the language gates above) — e.g. "the knowledge +graph links X to Y" rather than "X treats Y". Do NOT put pipeline metrics, +edge counts, provenance profiles, or verifiability caveats in the Answer — +those belong ONLY in the Confidence and Limitations sections, where you +cite the pipeline-context numbers. Keep the Answer about the answer. ## Template (use these exact `##` headings, in this order) ## Answer A direct 2-4 sentence answer to the question, in plain prose. State -the headline finding clearly. **Name each top entity with both its -human-readable label AND its CURIE in parentheses**, e.g. -"metformin (CHEBI:6801) and insulin (CHEBI:5931) are the most -common treatments..." — the CURIE makes the answer unambiguously -identifiable across knowledge graphs. No bullet list here. +the headline finding clearly. **Bold each entity name** and put its +CURIE in parentheses right after, e.g. "**metformin** (CHEBI:6801) and +**sitagliptin** (CHEBI:40237) are among the treatments..." — the bold +makes the answers scannable and the CURIE makes them unambiguous. +No bullet list here. + +Attribute the finding to the knowledge graph consistently — EVERY sentence, +including the first, must read as "the knowledge graph identifies / lists..." +rather than stating a graph-derived result as a bare biological fact. + +If the graph returned MORE results than you list (pipeline-context +`plover_total_results` exceeds your number of answers), add ONE short clause +saying these are the most relevant of the larger set — e.g. "the most +relevant of [N] the knowledge graph returned." One clause only; keep the +detailed provenance numbers for Confidence/Limitations. ## Evidence From da77d1ca5f17db4a58c9d4dbe0baa60a8849fa4c Mon Sep 17 00:00:00 2001 From: Adilbek Bazarkulov Date: Mon, 8 Jun 2026 12:31:49 -0400 Subject: [PATCH 5/6] Simplify entity resolution and add iterative chunked answer-picking --- pipeline/code/pipeline.py | 680 +++++++++++------- pipeline/code/tests/test_answer_graph_view.py | 54 ++ .../code/tests/test_iterative_answer_pick.py | 30 + pipeline/code/tests/test_query_arity.py | 68 ++ .../code/tests/test_rerank_coverage_tier.py | 74 ++ pipeline/code/tests/test_smoke.py | 2 +- 6 files changed, 666 insertions(+), 242 deletions(-) create mode 100644 pipeline/code/tests/test_iterative_answer_pick.py create mode 100644 pipeline/code/tests/test_query_arity.py create mode 100644 pipeline/code/tests/test_rerank_coverage_tier.py diff --git a/pipeline/code/pipeline.py b/pipeline/code/pipeline.py index 92c8eba..ba68176 100644 --- a/pipeline/code/pipeline.py +++ b/pipeline/code/pipeline.py @@ -74,7 +74,7 @@ # reduction: Strategy B reduces the PloverDB body to top-N results per # predicate before Stage 11 sees it, ranked by knowledge_level then # agent_type. full spec: docs/specs/response-reduction-strategy-b.md. -from .reduction import reduce_plover_response +from .reduction import chunk_plover_response, reduce_plover_response # prompts: SYS_ENTITY_EXTRACT, SYS_TRAPI_BUILD, SYS_ANSWER_PICK, # SYS_EXPLAIN. flat strings re-exported by name so diffs for tuning @@ -97,6 +97,9 @@ # plover_executed_rate, etc. keep them stable across runs. STATUS_OK = "ok" STATUS_INVALID_QUERY = "invalid_query" # validator rejected the LLM's TRAPI +STATUS_INVALID_QUERY_ARITY = "invalid_query_arity" # query graph is not one-hop + # (not exactly 2 nodes + 1 edge); + # caught in code before the validator STATUS_LLM_BAD_JSON = "llm_bad_json" # LLM returned non-JSON in a JSON-only stage STATUS_PLOVER_ERROR = "plover_error" # PloverDB returned non-200 / network error STATUS_LLM_ERROR = "llm_error" # OpenRouter returned non-200 / network error @@ -106,10 +109,6 @@ STATUS_NO_CANDIDATE_MATCH = "no_candidate_match" # Stage 4 LLM declared none of the NameRes # candidates fits the user's intent (typo # survived Stage 2, or all candidates wrong) -STATUS_LOW_CONFIDENCE_RESOLUTION = "low_confidence_resolution" # Stage 7: resolved label - # has low textual similarity to the user - # mention; pipeline refuses to query KG - # against a probably-wrong entity # Stage 1 (scope check) decided the input is outside PloverAI's # biomedical-KG scope (politics, weather, chit-chat, ...). this is a # deliberate refusal, not a failure — pipeline exits cleanly with a @@ -254,17 +253,22 @@ def _build_answer_graph_view( pinned_category: str | None, picked_answer_curies: list[str], plover_response: dict[str, Any], + supporting_edge_ids: set[str] | None = None, ) -> dict[str, Any]: # Stage 13: reshape (pinned entity + picked answers + PloverDB KG) # into a node-link graph view with full provenance per edge. pure # function, no IO. consumed by the frontend to render a research- - # grade graph card with hover-able evidence on each edge. + # grade graph card with hover-able evidence on each edge, and fed to + # the Stage 15 explainer — so its edge set IS the anti-hallucination + # boundary: the explainer can only cite edges that appear here. # # contract (see test_answer_graph_view.py for the strict spec): # - never drop a picked answer, even if it's missing from the KG # (label/category fall back to None) - # - keep only edges that touch the pinned node AND a picked answer - # (edges between two non-relevant nodes are noise — drop them) + # - when supporting_edge_ids is given, keep ONLY those edges (the + # answer-picker's cited support) — the faithfulness invariant. when + # None (legacy), keep edges that structurally touch the pinned node + # AND a picked answer. # - preserve TRAPI subject/object orientation verbatim (don't flip) # - degrade gracefully on missing/empty attributes blocks nodes_block: dict[str, Any] = ( @@ -309,15 +313,25 @@ def _build_answer_graph_view( for edge_id, e in edges_block.items(): subj = e.get("subject") obj = e.get("object") - if not (subj in relevant_pair and obj in relevant_pair): - continue - if pinned_curie not in (subj, obj): - continue - # at least one endpoint must be a picked answer (otherwise it's a - # pinned↔pinned self-edge, which TRAPI shouldn't produce but - # we guard against) - if not ((subj in picked_set) or (obj in picked_set)): - continue + if supporting_edge_ids is not None: + # faithfulness invariant: keep ONLY edges the answer-picker + # cited as support. the edge id is authoritative, so we do NOT + # re-apply the structural node-pair filter here — the explainer + # sees exactly the picked edges, nothing the picker didn't select. + if edge_id not in supporting_edge_ids: + continue + else: + # legacy structural filter (no cited edges supplied): keep edges + # that touch the pinned node AND a picked answer; drop the rest. + if not (subj in relevant_pair and obj in relevant_pair): + continue + if pinned_curie not in (subj, obj): + continue + # at least one endpoint must be a picked answer (otherwise it's a + # pinned↔pinned self-edge, which TRAPI shouldn't produce but + # we guard against) + if not ((subj in picked_set) or (obj in picked_set)): + continue attrs = e.get("attributes") or [] # walk the attributes list once, picking up each provenance field # by its attribute_type_id. None / empty defaults for absent ones. @@ -395,20 +409,24 @@ def _rerank_nameres_candidates( candidates: list[dict[str, Any]], mention: str, expected_category: str | None, + coverage_by_curie: dict[str, int] | None = None, ) -> list[dict[str, Any]]: # local re-rank that the BM25 score alone won't deliver. tiers go # MOST-discriminating first, BM25 last. lexicographic sort over the # tier tuple means a Tier-1 hit (exact label match) wins over a # Tier-4 BM25 spike no matter how big the BM25 difference is. # - # tiers (each is 0 or 1, summed in a tuple and sorted DESC): + # tiers (sorted DESC): # T1 exact_label — candidate.label == mention (case-folded) # T2 exact_synonym — mention appears verbatim in candidate.synonyms # T3 token_match — mention appears as a whole token (split on # whitespace) in label OR any synonym; catches # "seizures" inside "Febrile Seizure" / etc. # T4 type_match — Stage 2's expected_category is in candidate.types - # T5 bm25_score — fall-through; NameRes's original ranking + # T5 has_edges — candidate has >=1 edge to the answer category in + # the hosted graph (from the coverage probe) + # T6 edge_count — more edges first, within has_edges + # T7 bm25_score — fall-through; NameRes's original ranking # # rationale per tier: # - T1/T2 fix the canonical-short-label problem: "Seizure" (HP) gets @@ -420,11 +438,17 @@ def _rerank_nameres_candidates( # biolink_type filter lets adjacent-type entries in. for "seizures" # that means HP (PhenotypicFeature) candidates beat MONDO (Disease) # candidates even though MONDO scores ~165 points higher in BM25. - # - T5 is the safety net so candidates with NO discriminating signal + # - T5/T6 sit BELOW type_match on purpose: graph coverage breaks ties + # among same-type candidates (so we stop pinning a real-looking but + # data-less entity) WITHOUT overriding the type-match fix — a + # type-matched candidate with zero edges still beats a wrong-type one + # with more edges. inert when coverage_by_curie is None (the first, + # pre-probe rerank used only to choose which candidates to probe). + # - T7 is the safety net so candidates with NO discriminating signal # are still ordered deterministically by their original rank. m = mention.lower().strip() - def key(c: dict[str, Any]) -> tuple[int, int, int, int, float]: + def key(c: dict[str, Any]) -> tuple[int, int, int, int, int, int, float]: label = str(c.get("label") or "").lower().strip() synonyms = [ str(s).lower().strip() @@ -444,7 +468,16 @@ def key(c: dict[str, Any]) -> tuple[int, int, int, int, float]: token_match = 1 if (m in label_tokens or m in syn_tokens) else 0 type_match = 1 if expected_category and expected_category in types else 0 - return (exact_label, exact_syn, token_match, type_match, bm25) + edge_count = 0 + if coverage_by_curie is not None: + count = coverage_by_curie.get(str(c.get("curie"))) + edge_count = count if isinstance(count, int) and count > 0 else 0 + has_edges = 1 if edge_count > 0 else 0 + + return ( + exact_label, exact_syn, token_match, type_match, + has_edges, edge_count, bm25, + ) return sorted(candidates, key=key, reverse=True) @@ -510,6 +543,17 @@ def _has_any_kg_coverage(probes: dict[str, Any]) -> bool: ) +def _coverage_from_probes(probes: dict[str, Any]) -> dict[str, int]: + # extract per-candidate edge counts from the probe results so they can + # be folded into the rerank as the coverage tier. failed probes (error + # set) are skipped — a probe failure is "unknown", not "zero edges". + return { + curie: int(p.get("total_edges") or 0) + for curie, p in probes.items() + if not p.get("error") + } + + def _check_label_consistency(mention: str, label: str) -> tuple[float, dict[str, Any]]: # Stage 7's similarity check, factored out for unit-testability. # returns (similarity_score, debug_info_dict). pure function — no @@ -549,33 +593,25 @@ def _check_label_consistency(mention: str, label: str) -> tuple[float, dict[str, } -def _format_low_confidence_explanation( - question: str, mention: str, canonical: str, label: str, similarity: float, -) -> str: - # Markdown body for the explanation.md artifact when Stage 7 refuses - # to query PloverDB because the resolved entity has low textual - # similarity to what the user typed. mirrors the section structure of - # the normal explainer + the out_of_scope refusal so the UI renders - # the same component, just with a low_confidence_resolution badge. - return ( - "## Answer\n\n" - f"PloverAI could not confidently resolve **{mention}** to a " - f"known biomedical entity in RTX-KG2.10.2c.\n\n" - "## Reason\n\n" - f"The closest match found was **{canonical}** " - f"(*{label}*), but the resolved label is too different from what " - f"you typed (similarity {similarity:.2f}, threshold 0.50). " - f"Querying PloverDB with this entity would likely return results " - f"for a different concept than you intended, so the pipeline " - f"stopped here rather than silently grounding the wrong entity.\n\n" - "## What to try\n\n" - "- Check the spelling of the entity in your question.\n" - "- Use the canonical name (e.g., *type 2 diabetes mellitus* " - "instead of *type 2 diabites*).\n" - "- Expand abbreviations if the entity is well-known by its full " - "name (e.g., *T2DM* → *type 2 diabetes mellitus*).\n" - "- If you intended a different entity, please rephrase the " - "question with the entity's full name.\n" +def _check_query_graph_arity(trapi_msg: dict[str, Any]) -> tuple[bool, str]: + # one-hop arity gate: PloverDB accepts exactly 2 query-graph nodes and + # 1 edge. Stage 8 has an LLM assemble the query graph, and Stage 9's + # reasoner-validator only checks Biolink term compliance, NOT shape, so + # this is the only code that structurally enforces the one-hop + # invariant before the query is sent. returns (ok, reason); reason is + # "" when ok and names the offending counts otherwise. degrades + # gracefully (False, reason) on a malformed message rather than raising. + message = trapi_msg.get("message") + query_graph = message.get("query_graph") if isinstance(message, dict) else None + nodes = query_graph.get("nodes") if isinstance(query_graph, dict) else None + edges = query_graph.get("edges") if isinstance(query_graph, dict) else None + n_nodes = len(nodes) if isinstance(nodes, dict) else 0 + n_edges = len(edges) if isinstance(edges, dict) else 0 + if n_nodes == 2 and n_edges == 1: + return True, "" + return False, ( + f"query graph has {n_nodes} nodes and {n_edges} edges; " + f"one-hop requires exactly 2 nodes and 1 edge" ) @@ -722,6 +758,27 @@ def _extract_json(text: str) -> dict[str, Any]: return loaded2 +def _valid_iterative_picks( + answers: list[Any], chunk_node_ids: set[str], prior_curies: set[str], +) -> list[dict[str, Any]]: + # keep only answers whose curie is in THIS chunk's nodes or was already a + # valid prior pick (carry-forward). drops invented/hallucinated curies so + # selection is enforced in code, not trusted from the model. dedupes by + # curie, first occurrence wins. + out: list[dict[str, Any]] = [] + seen: set[str] = set() + for answer in answers: + if not isinstance(answer, dict): + continue + curie = answer.get("curie") + if not isinstance(curie, str) or curie in seen: + continue + if curie in chunk_node_ids or curie in prior_curies: + out.append(answer) + seen.add(curie) + return out + + def run_grounded( *, cfg: Config, @@ -768,63 +825,65 @@ def run_grounded( # by referencing it. _ = cfg - # ---------- Stage 1: scope-check guardrail ---------- - # decides whether the question is biomedical at all. refusing - # here costs ~1 cheap LLM call and saves the full pipeline cost - # (~6 LLM calls + 4 RENCI calls + PloverDB) on out-of-scope - # input. failures fall through to STAGE_0 (fail-open) because - # blocking a real biomedical question on a guardrail glitch is - # worse than letting the rare chit-chat through. - user_msg_scope = f"Question: {nl_question}" - prompt_log["stage_1_scope_check"] = { - "system": prompts.SYS_SCOPE_CHECK, - "user": user_msg_scope, - } - write_json(qp.prompt, prompt_log) - try: - rep_scope = llm.chat( - model=model, - system=prompts.SYS_SCOPE_CHECK, - user=user_msg_scope, - stage="scope_check", - ) - except OpenRouterError as e: - return _finish_failure(qp, ledger, t_start, q["id"], STATUS_LLM_ERROR, str(e)) - prompt_log["stage_1_scope_check"]["response"] = _llm_response_meta(rep_scope) - write_json(qp.prompt, prompt_log) - ledger.add( - "scope_check", model.id, model.slug, - rep_scope.input_tokens, rep_scope.output_tokens, - rep_scope.cost.input_usd, rep_scope.cost.output_usd, rep_scope.cost.total_usd, - rep_scope.latency_s, - ) - scope = _parse_scope_check_output(rep_scope.content) - logger.info( - f"{tag} scope_check in_scope={scope.in_scope} " - f"reason=[white]{scope.reason or '(none)'!r}[/]" - ) - if not scope.in_scope: - # write a clean markdown explanation so the UI has something - # to render — no TRAPI query, no graph hit, no further cost. - explanation = _format_out_of_scope_explanation(nl_question, scope.reason) - write_text(qp.explanation, explanation) - _flush_meta_and_cost( - qp, ledger, t_start, q["id"], - STATUS_OUT_OF_SCOPE, scope.reason or "question is outside biomedical-KG scope", - outcome="out_of_scope", - outcome_reason=scope.reason, - plover_n_results=-1, answers_n_picked=-1, + # ---------- Stage 1: scope-check guardrail (optional) ---------- + # an LLM call that decides whether the question is biomedical at all, + # refusing chit-chat / out-of-scope input before the full pipeline cost. + # DISABLED by default (cfg.scope_check_enabled): it is an extra OpenRouter + # call per query and a rate-limit source. when off, every input is treated + # as in-scope and flows to entity extraction; scope is checked manually. + if cfg.scope_check_enabled: + user_msg_scope = f"Question: {nl_question}" + prompt_log["stage_1_scope_check"] = { + "system": prompts.SYS_SCOPE_CHECK, + "user": user_msg_scope, + } + write_json(qp.prompt, prompt_log) + try: + rep_scope = llm.chat( + model=model, + system=prompts.SYS_SCOPE_CHECK, + user=user_msg_scope, + stage="scope_check", + ) + except OpenRouterError as e: + return _finish_failure(qp, ledger, t_start, q["id"], STATUS_LLM_ERROR, str(e)) + prompt_log["stage_1_scope_check"]["response"] = _llm_response_meta(rep_scope) + write_json(qp.prompt, prompt_log) + ledger.add( + "scope_check", model.id, model.slug, + rep_scope.input_tokens, rep_scope.output_tokens, + rep_scope.cost.input_usd, rep_scope.cost.output_usd, rep_scope.cost.total_usd, + rep_scope.latency_s, ) - return QuestionResult( - q_id=q["id"], - status=STATUS_OUT_OF_SCOPE, - cost_total_usd=ledger.total_usd(), - cost_total_tokens=ledger.total_tokens(), - elapsed_s=round(time.perf_counter() - t_start, 3), - error=None, - outcome="out_of_scope", - outcome_reason=scope.reason, + scope = _parse_scope_check_output(rep_scope.content) + logger.info( + f"{tag} scope_check in_scope={scope.in_scope} " + f"reason=[white]{scope.reason or '(none)'!r}[/]" ) + if not scope.in_scope: + # write a clean markdown explanation so the UI has something + # to render — no TRAPI query, no graph hit, no further cost. + explanation = _format_out_of_scope_explanation(nl_question, scope.reason) + write_text(qp.explanation, explanation) + _flush_meta_and_cost( + qp, ledger, t_start, q["id"], + STATUS_OUT_OF_SCOPE, scope.reason or "question is outside biomedical-KG scope", + outcome="out_of_scope", + outcome_reason=scope.reason, + plover_n_results=-1, answers_n_picked=-1, + ) + return QuestionResult( + q_id=q["id"], + status=STATUS_OUT_OF_SCOPE, + cost_total_usd=ledger.total_usd(), + cost_total_tokens=ledger.total_tokens(), + elapsed_s=round(time.perf_counter() - t_start, 3), + error=None, + outcome="out_of_scope", + outcome_reason=scope.reason, + ) + else: + logger.info(f"{tag} scope_check skipped (disabled in config)") # ---------- Stage 2: extract focal entity mention from NL ---------- # we inject the list of Biolink categories PloverDB ACTUALLY has @@ -954,6 +1013,17 @@ def run_grounded( logger=logger, tag=f"{tag} strict", ) + # fold the probe's edge counts back into the rerank so the order the + # LLM sees already prefers candidates with graph coverage (a tier below + # type_match, so the type-match fix still dominates). + nr_candidates_full = _rerank_nameres_candidates( + nr_candidates_full, mention, expected_category, + coverage_by_curie=_coverage_from_probes(candidate_probes_by_curie), + ) + rerank_top1 = nr_candidates_full[0].get("curie") if nr_candidates_full else None + rerank_top1_label = ( + nr_candidates_full[0].get("label") if nr_candidates_full else None + ) nameres_filter_used: list[str] | None = strict_filter strict_has_coverage = _has_any_kg_coverage(candidate_probes_by_curie) fallback_to_loose = False @@ -986,6 +1056,10 @@ def run_grounded( logger=logger, tag=f"{tag} loose", ) + loose_reranked = _rerank_nameres_candidates( + loose_reranked, mention, expected_category, + coverage_by_curie=_coverage_from_probes(loose_probes), + ) # adopt the loose pass as the working set. strict results stay # available for the audit trail in `strict_attempt` below. strict_attempt = { @@ -1105,9 +1179,9 @@ def run_grounded( c_score = c.get("score") # if we ran the candidate-density probe, attach the per-candidate # edge count to the answer category. this lets the LLM avoid - # the "lexical match but zero KG2c coverage" trap (e.g. picking + # the "lexical match but zero graph coverage" trap (e.g. picking # PANTHER.PATHWAY:P00014 for cholesterol biosynthesis when - # PANTHER pathways have 0 edges to Gene nodes in KG2c — the + # PANTHER pathways have 0 edges to Gene nodes in the graph — the # Reactome / GO equivalents in the same candidate set are the # ones that actually have edges). probe_suffix = "" @@ -1116,9 +1190,9 @@ def run_grounded( n = cp.get("total_edges", 0) err = cp.get("error") if err: - probe_suffix = f" kg2c_edges_to_{answer_category}=probe_failed({err})" + probe_suffix = f" kg_edges_to_{answer_category}=probe_failed({err})" else: - probe_suffix = f" kg2c_edges_to_{answer_category}={n}" + probe_suffix = f" kg_edges_to_{answer_category}={n}" candidates_block_lines.append( f" {i}. curie={c_curie!r} label={c_label!r} " f"types={c_types} bm25_score={c_score}{probe_suffix}" @@ -1130,13 +1204,12 @@ def run_grounded( density_preamble = "" if candidate_probes_by_curie: density_preamble = ( - f"Each candidate carries a `kg2c_edges_to_{answer_category}` " - f"count: the number of edges in KG2c from that CURIE to any " - f"node of the answer category (in either direction). A " - f"high-BM25 candidate with ZERO KG2c edges will produce no " - f"results downstream — prefer a slightly lower-ranked " - f"candidate that has actual edges over a perfect-label one " - f"with no data.\n\n" + f"Each candidate carries a `kg_edges_to_{answer_category}` " + f"count: the number of edges in the hosted knowledge graph " + f"from that CURIE to any node of the answer category (in " + f"either direction). The candidate ORDER already prefers " + f"non-zero coverage; a candidate with ZERO edges will produce " + f"no results downstream.\n\n" ) user_msg_pick = ( f"User question: {nl_question}\n" @@ -1262,54 +1335,6 @@ def run_grounded( ) candidate_pick_fell_back = True - # ---------- Stage 5: re-rank NameRes candidates by information_content ---------- - # NameRes ranks by Solr BM25 — biased toward longer labels that - # contain the query string ("Hypoglycemic seizures" beats "Seizure" - # for the query "seizures"). when the question wants a broad concept - # (granularity=general), we override that by sorting candidates by - # information_content ASCENDING (lower IC = more general concept). - # for granularity=specific we leave the Stage-4 pick alone. - # - # information_content comes from NodeNorm, so this costs one extra - # NodeNorm batch call on the top-K (K=5). NodeNorm is fast and - # batches all 5 in a single POST. - reranked = False - if granularity == "general" and len(candidate_curies) > 1: - try: - nn_candidates = nodenorm.normalize(candidate_curies) - # smaller IC = more general → lowest IC first. NodeNorm - # returns None for unresolvable CURIEs; we treat those as - # large (sort them last) so unresolvable candidates can't - # accidentally win the broad-concept lottery. - def _ic_key(curie: str) -> float: - ic = nn_candidates.information_content.get(curie) - return ic if ic is not None else float("inf") - sorted_by_ic = sorted(candidate_curies, key=_ic_key) - if sorted_by_ic[0] != chosen_curie: - prev_curie = chosen_curie - chosen_curie = sorted_by_ic[0] - reranked = True - logger.info( - f"{tag} ic_rerank SWAP " - f"{prev_curie} (IC={nn_candidates.information_content.get(prev_curie)}) " - f"→ {chosen_curie} (IC={nn_candidates.information_content.get(chosen_curie)}) " - f"granularity=general" - ) - else: - # log "ran but didn't swap" so the progress stream shows - # the IC rerank was considered — silent inaction looks - # the same as "stage didn't run" without this line. - logger.info( - f"{tag} ic_rerank no_swap " - f"{chosen_curie} is already lowest-IC " - f"(IC={nn_candidates.information_content.get(chosen_curie)}) " - f"granularity=general" - ) - except NodeNormError as e: - # don't fail the pipeline over a re-rank optimisation — - # fall back to the existing chosen_curie. - logger.warning(f"Stage 5 IC re-rank skipped (NodeNorm error): {e}") - # ---------- Stage 6: NodeNorm canonicalises the chosen CURIE ---------- try: nn_pinned = nodenorm.normalize([chosen_curie]) @@ -1343,7 +1368,6 @@ def _ic_key(curie: str) -> float: "fell_back_to_top1": candidate_pick_fell_back, }, "chosen_curie": chosen_curie, - "reranked_by_ic": reranked, "latency_s": round(nr.latency_s, 3), }) @@ -1382,8 +1406,17 @@ def _ic_key(curie: str) -> float: f"substring={sim_debug['substring_match']})" ) if similarity < LOW_CONFIDENCE_THRESHOLD: - # write the failure-resolution info to nameres.json so the - # artifact captures why the pipeline stopped here + # low textual similarity between the user's mention and the + # resolved label. this used to hard-abort, but the 0.50 cutoff is + # only anecdotally calibrated, so we now WARN, record the flag in + # nameres.json, and continue rather than refuse a possibly-valid + # question. the flag is available downstream for an answer caveat. + logger.warning( + f"Stage 7 low-confidence resolution: mention={mention!r} " + f"resolved to {canonical_pinned!r} ({pinned_label!r}) at " + f"similarity {similarity:.2f} < {LOW_CONFIDENCE_THRESHOLD}; " + f"continuing (flagged low_confidence_resolution)" + ) write_json(qp.nameres, { "mention": nr.mention, "expected_category": expected_category, @@ -1400,7 +1433,7 @@ def _ic_key(curie: str) -> float: "fell_back_to_top1": candidate_pick_fell_back, }, "chosen_curie": chosen_curie, - "reranked_by_ic": reranked, + "low_confidence_resolution": True, "latency_s": round(nr.latency_s, 3), "consistency_check": { "mention": sim_debug["mention_normalized"], @@ -1412,30 +1445,6 @@ def _ic_key(curie: str) -> float: "passed": False, }, }) - # surface a "did you mean?" style failure to the user via the - # standard explanation.md path (same Markdown structure as the - # OUT_OF_SCOPE refusal so the UI renders the same component). - # the message names both the user's mention and what we resolved - # it to, so the user can tell whether to rephrase. - did_you_mean_md = _format_low_confidence_explanation( - question=nl_question, - mention=mention, - canonical=canonical_pinned, - label=pinned_label or "(no label)", - similarity=similarity, - ) - write_text(qp.explanation, did_you_mean_md) - short_err = ( - f"Resolved {canonical_pinned!r} ({pinned_label!r}) has low " - f"similarity ({similarity:.2f} < {LOW_CONFIDENCE_THRESHOLD}) to " - f"user mention {mention!r}; pipeline refused to query KG against " - f"a probably-wrong entity." - ) - return _finish_failure( - qp, ledger, t_start, q["id"], - STATUS_LOW_CONFIDENCE_RESOLUTION, - short_err, - ) # ---------- Stage 8: LLM builds TRAPI query graph ---------- # the predicate list (fix B) is the meta_KG slice for the @@ -1627,6 +1636,18 @@ def _ic_key(curie: str) -> float: write_json(qp.trapi_query, trapi_msg) + # one-hop arity gate (code-enforced, BEFORE the Biolink-compliance + # validator): a hallucinated multi-hop / wrong-node-count query is + # stopped here rather than sent to PloverDB. reasoner-validator checks + # Biolink terms, not shape, so this is the only structural one-hop check. + arity_ok, arity_reason = _check_query_graph_arity(trapi_msg) + if not arity_ok: + logger.warning(f"Stage 9 arity gate rejected: {arity_reason}") + return _finish_failure( + qp, ledger, t_start, q["id"], + STATUS_INVALID_QUERY_ARITY, arity_reason, + ) + # ---------- Stage 9: validate ---------- val = validate_query(trapi_msg, logger=logger) write_json(qp.validation, { @@ -1662,11 +1683,11 @@ def _ic_key(curie: str) -> float: plover_n_results = len(prep.body.get("message", {}).get("results") or []) # ---------- Stage 11: LLM picks answer ---------- - # Strategy B reduction (predicate-grouped knowledge_level ranking) - # shrinks the PloverDB body BEFORE the LLM sees it. the top-N is a - # config knob; the reduced body + per-group stats are written out - # as artifacts so the faithfulness evaluator can grade answers - # against what the LLM saw, not the full PloverDB response. + # reduction is computed in both modes: the single-shot path feeds its + # top-N reduced body to the LLM, and Stage 15's pipeline-context cites + # its metadata. when iterative mode is on, the LLM instead reads the + # FULL sorted response in token-budgeted chunks (no top-N truncation), + # so a correct answer can't be dropped before the LLM ever sees it. reduction = reduce_plover_response( prep.body, top_n_per_predicate=cfg.reduction.top_n_per_predicate, @@ -1674,45 +1695,195 @@ def _ic_key(curie: str) -> float: ) write_json(qp.reduced_data, reduction.reduced_body) write_json(qp.reduction_metadata, asdict(reduction.metadata)) - user_msg_4 = ( - f"User question: {nl_question}\n\n" - f"TRAPI response (reduced — top " - f"{reduction.metadata.top_n_per_predicate} results per predicate, " - f"ranked by knowledge_level then agent_type):\n" - f"{json.dumps(reduction.reduced_body, ensure_ascii=False)}\n" - ) - prompt_log["stage_11_answer_pick"] = { - "system": prompts.SYS_ANSWER_PICK, - "user_truncated": user_msg_4[:2000], - } - write_json(qp.prompt, prompt_log) - try: - rep4 = llm.chat( - model=model, - system=prompts.SYS_ANSWER_PICK, - user=user_msg_4, - stage="answer_pick", + answer_obj: dict[str, Any] + # selection_mode + iter_summary feed the Stage 15 pipeline-context so its + # caveats cite what ACTUALLY happened (single-shot top-N truncation vs + # iterative chunked reading), not the unused Strategy B numbers. + selection_mode = "single_shot" + iter_summary: dict[str, Any] = {} + if cfg.stage11_iterative.enabled: + # iterative chunked answer-picking: read chunks strongest-first and + # stop as soon as the model has `answer_target` good picks (or is + # satisfied with fewer). chunk 0 carries the strongest evidence, so + # most questions finish in one call; we only read further chunks when + # a chunk was a dud. the chunk budget is sized to THIS model's context + # window (context_window * context_fraction), so a big-context model + # sends the whole response in one chunk while a small one splits it. + # every returned curie is re-validated in code against the chunk's + # nodes or the carry-forward set, so invented curies are dropped. + question_dir = qp.answer.parent + chunk_budget = max( + 1, int(model.context_window * cfg.stage11_iterative.context_fraction), ) - except OpenRouterError as e: - return _finish_failure(qp, ledger, t_start, q["id"], STATUS_LLM_ERROR, str(e)) - prompt_log["stage_11_answer_pick"]["response"] = _llm_response_meta(rep4) - write_json(qp.prompt, prompt_log) - ledger.add( - "answer_pick", model.id, model.slug, - rep4.input_tokens, rep4.output_tokens, - rep4.cost.input_usd, rep4.cost.output_usd, rep4.cost.total_usd, - rep4.latency_s, - ) + chunkset = chunk_plover_response( + prep.body, + chunk_token_budget=chunk_budget, + max_chunks=cfg.stage11_iterative.max_chunks, + logger=logger, + ) + picks: list[dict[str, Any]] = [] + prior_curies: set[str] = set() + iter_tier: Any = None + iter_rationale = "" + iterations: list[dict[str, Any]] = [] + termination = "chunks_exhausted" + for i, chunk in enumerate(chunkset.chunks): + chunk_nodes = chunk.get("message", {}).get("knowledge_graph", {}).get("nodes", {}) + chunk_node_ids: set[str] = set(chunk_nodes) if isinstance(chunk_nodes, dict) else set() + # the compact running shortlist (curie + label + relevance + # reason) is carried forward so the LLM can rank-merge across + # chunks without re-reading full edge JSON. + shortlist_so_far = json.dumps( + [ + {"curie": p.get("curie"), "label": p.get("label"), "why": p.get("why")} + for p in picks + ], + ensure_ascii=False, + ) + user_msg_iter = ( + f"User question: {nl_question}\n\n" + f"Return up to {cfg.stage11_iterative.answer_target} answers, " + f"ranked by relevance.\n\n" + f"shortlist_so_far: {shortlist_so_far}\n\n" + f"Chunk {i + 1} of {chunkset.n_chunks} (strongest evidence " + f"first; text-mined last):\n" + f"{json.dumps(chunk, ensure_ascii=False)}\n" + ) + try: + rep_iter = llm.chat( + model=model, + system=prompts.SYS_ANSWER_PICK_ITER, + user=user_msg_iter, + stage="answer_pick", + ) + except OpenRouterError as e: + return _finish_failure( + qp, ledger, t_start, q["id"], STATUS_LLM_ERROR, str(e), + ) + ledger.add( + "answer_pick", model.id, model.slug, + rep_iter.input_tokens, rep_iter.output_tokens, + rep_iter.cost.input_usd, rep_iter.cost.output_usd, + rep_iter.cost.total_usd, rep_iter.latency_s, + ) + try: + chunk_obj = _extract_json(rep_iter.content) + except json.JSONDecodeError: + # a single malformed chunk response must not kill the run: + # keep picks_so_far and move to the next chunk. + logger.warning(f"Stage 11 chunk {i} returned non-JSON; skipping it") + write_json(question_dir / f"stage11_chunk_{i}.json", { + "chunk_index": i, + "error": "non_json_response", + "response": _llm_response_meta(rep_iter), + }) + continue + raw_answers = chunk_obj.get("answers") + picks = _valid_iterative_picks( + raw_answers if isinstance(raw_answers, list) else [], + chunk_node_ids, prior_curies, + ) + prior_curies = {str(p["curie"]) for p in picks} + iter_tier = chunk_obj.get("evidence_tier", iter_tier) + iter_rationale = str(chunk_obj.get("rationale") or iter_rationale) + confident = bool(chunk_obj.get("confidence_sufficient")) + write_json(question_dir / f"stage11_chunk_{i}.json", { + "chunk_index": i, + "n_picks_after": len(picks), + "confidence_sufficient": confident, + "picks": picks, + "rationale": iter_rationale, + "response": _llm_response_meta(rep_iter), + "user_truncated": user_msg_iter[:2000], + }) + iterations.append({ + "chunk_index": i, + "n_picks": len(picks), + "confidence_sufficient": confident, + }) + # relevance-ranking stop: only the LLM's confidence ends the loop + # early. we do NOT stop just because we hold answer_target picks — + # a more relevant answer can still be in a later (weaker-evidence) + # chunk, so the model decides when it has seen enough to rank. + if confident: + termination = "confidence_sufficient" + break + else: + termination = ( + "max_chunks" if chunkset.truncated_at_max_chunks else "chunks_exhausted" + ) + # clamp to the soft target: picks are the model's ranked set, so the + # first answer_target are the strongest. enforced in code, not trusted. + picks = picks[: cfg.stage11_iterative.answer_target] + write_json(question_dir / "stage11_loop_summary.json", { + "n_chunks_available": chunkset.n_chunks, + "total_rows": chunkset.total_rows, + "chunked_rows": chunkset.chunked_rows, + "truncated_at_max_chunks": chunkset.truncated_at_max_chunks, + "iterations": iterations, + "termination": termination, + "n_picks_final": len(picks), + }) + selection_mode = "iterative" + iter_summary = { + "chunks_read": len(iterations), + "chunks_available": chunkset.n_chunks, + "total_results_available": chunkset.total_rows, + "termination": termination, + } + answer_obj = { + "answers": picks, + "evidence_tier": iter_tier, + "rationale": iter_rationale, + } + prompt_log["stage_11_answer_pick"] = { + "system": prompts.SYS_ANSWER_PICK_ITER, + "mode": "iterative", + "n_chunks": chunkset.n_chunks, + "termination": termination, + } + write_json(qp.prompt, prompt_log) + else: + user_msg_4 = ( + f"User question: {nl_question}\n\n" + f"TRAPI response (reduced — top " + f"{reduction.metadata.top_n_per_predicate} results per predicate, " + f"ranked by source tier then knowledge_level then n_publications):\n" + f"{json.dumps(reduction.reduced_body, ensure_ascii=False)}\n" + ) + prompt_log["stage_11_answer_pick"] = { + "system": prompts.SYS_ANSWER_PICK, + "user_truncated": user_msg_4[:2000], + } + write_json(qp.prompt, prompt_log) - try: - answer_obj = _extract_json(rep4.content) - except json.JSONDecodeError: - return _finish_failure( - qp, ledger, t_start, q["id"], - STATUS_LLM_BAD_JSON, "Stage 11 output was not valid JSON", + try: + rep4 = llm.chat( + model=model, + system=prompts.SYS_ANSWER_PICK, + user=user_msg_4, + stage="answer_pick", + ) + except OpenRouterError as e: + return _finish_failure(qp, ledger, t_start, q["id"], STATUS_LLM_ERROR, str(e)) + prompt_log["stage_11_answer_pick"]["response"] = _llm_response_meta(rep4) + write_json(qp.prompt, prompt_log) + ledger.add( + "answer_pick", model.id, model.slug, + rep4.input_tokens, rep4.output_tokens, + rep4.cost.input_usd, rep4.cost.output_usd, rep4.cost.total_usd, + rep4.latency_s, ) + try: + answer_obj = _extract_json(rep4.content) + except json.JSONDecodeError: + return _finish_failure( + qp, ledger, t_start, q["id"], + STATUS_LLM_BAD_JSON, "Stage 11 output was not valid JSON", + ) + # ---------- Stage 12: NodeNorm canonicalises every answer CURIE ---------- raw_answer_curies: list[str] = [ a["curie"] for a in (answer_obj.get("answers") or []) if a.get("curie") @@ -1764,12 +1935,31 @@ def _ic_key(curie: str) -> float: # rendering as a hoverable graph card in the UI. pure function, # unit-tested in tests/test_answer_graph_view.py. canonical_curies_for_view: list[str] = answer_obj.get("canonical_curies") or [] + # faithfulness invariant: feed the explainer ONLY the edges the + # answer-picker cited (supporting_edge_ids), not the full body filtered + # by node-pair. so the explainer cannot cite an edge the picker didn't + # select. if the picker cited nothing, fall back to the legacy node-pair + # filter (None) so the explanation isn't empty — and warn loudly. + picked_edge_ids: set[str] = { + eid + for answer in (answer_obj.get("answers") or []) + if isinstance(answer, dict) + for eid in (answer.get("supporting_edge_ids") or []) + if isinstance(eid, str) + } + if not picked_edge_ids: + logger.warning( + "Stage 13 answer-picker cited no supporting_edge_ids; graph view " + "falls back to node-pair edges (explainer may see edges beyond " + "the picked set)" + ) answer_graph_view = _build_answer_graph_view( pinned_curie=canonical_pinned, pinned_label=pinned_label, pinned_category=primary_pinned_cat, picked_answer_curies=canonical_curies_for_view, plover_response=prep.body, + supporting_edge_ids=picked_edge_ids or None, ) # ---------- Stage 13 enrichment: PubTator co-mention verification ---------- @@ -1926,21 +2116,29 @@ def _ic_key(curie: str) -> float: qg_edges = trapi_msg.get("message", {}).get("query_graph", {}).get("edges", {}) qg_e0 = qg_edges.get("e0") or {} qg_predicates = qg_e0.get("predicates") or [] - pipeline_context = { + pipeline_context: dict[str, Any] = { "predicate_used": qg_predicates[0] if qg_predicates else None, "plover_total_results": plover_n_results, - "reduction_strategy": reduction.metadata.strategy_applied, - "reduction_top_n_per_predicate": reduction.metadata.top_n_per_predicate, - "reduction_results_kept": reduction.metadata.reduced_result_count, - "reduction_results_dropped": ( - reduction.metadata.original_result_count - - reduction.metadata.reduced_result_count - ), - "reduction_predicate_groups": reduction.metadata.predicate_groups, - "reduction_edges_kept_per_group": reduction.metadata.edges_kept_per_group, - "reduction_edges_dropped_per_group": reduction.metadata.edges_dropped_per_group, + "selection_mode": selection_mode, "picked_edges_count": len(answer_graph_view.get("edges") or []), } + if selection_mode == "iterative": + # iterative read the FULL response in chunks (no top-N truncation), + # so cite the chunking numbers, not the unused Strategy B reduction. + pipeline_context.update(iter_summary) + else: + pipeline_context.update({ + "reduction_strategy": reduction.metadata.strategy_applied, + "reduction_top_n_per_predicate": reduction.metadata.top_n_per_predicate, + "reduction_results_kept": reduction.metadata.reduced_result_count, + "reduction_results_dropped": ( + reduction.metadata.original_result_count + - reduction.metadata.reduced_result_count + ), + "reduction_predicate_groups": reduction.metadata.predicate_groups, + "reduction_edges_kept_per_group": reduction.metadata.edges_kept_per_group, + "reduction_edges_dropped_per_group": reduction.metadata.edges_dropped_per_group, + }) user_msg_5 = ( f"User question: {nl_question}\n\n" f"Pipeline context (cite these numbers in your Confidence and " diff --git a/pipeline/code/tests/test_answer_graph_view.py b/pipeline/code/tests/test_answer_graph_view.py index 0d840e5..ec4e00c 100644 --- a/pipeline/code/tests/test_answer_graph_view.py +++ b/pipeline/code/tests/test_answer_graph_view.py @@ -323,3 +323,57 @@ def test_multiple_edges_between_same_pair_all_kept(): assert len(view["edges"]) == 2 levels = sorted([e["knowledge_level"] for e in view["edges"]]) assert levels == ["knowledge_assertion", "prediction"] + + +# ---- anti-hallucination invariant: supporting_edge_ids restricts the view ---- + +def _two_edge_kg(): + # two edges BOTH connecting metformin -> T2DM (both pass the legacy + # node-pair filter); used to show supporting_edge_ids keeps only cited. + base_attrs = [{"attribute_type_id": "biolink:knowledge_level", + "value": "knowledge_assertion"}] + return { + "message": { + "knowledge_graph": { + "nodes": { + "MONDO:0005148": {"name": "type 2 diabetes mellitus", + "categories": ["biolink:Disease"]}, + "CHEBI:6801": {"name": "metformin", "categories": ["biolink:Drug"]}, + }, + "edges": { + "edge_cited": {"subject": "CHEBI:6801", "object": "MONDO:0005148", + "predicate": "biolink:treats", "attributes": base_attrs}, + "edge_uncited": {"subject": "CHEBI:6801", "object": "MONDO:0005148", + "predicate": "biolink:treats", "attributes": base_attrs}, + }, + } + } + } + + +def test_supporting_edge_ids_restricts_to_cited_edges_only(): + # both edges connect metformin -> T2DM, so the legacy node-pair filter + # would keep both. citing only edge_cited must yield ONLY edge_cited — + # the explainer cannot see an edge the answer-picker did not select. + view = _build_answer_graph_view( + pinned_curie="MONDO:0005148", + pinned_label="type 2 diabetes mellitus", + pinned_category="biolink:Disease", + picked_answer_curies=["CHEBI:6801"], + plover_response=_two_edge_kg(), + supporting_edge_ids={"edge_cited"}, + ) + assert {e["id"] for e in view["edges"]} == {"edge_cited"} + + +def test_legacy_node_pair_filter_keeps_both_without_supporting_ids(): + # contrast: with no cited edges, the legacy filter keeps both edges + # connecting the pinned node to the picked answer (the old behaviour). + view = _build_answer_graph_view( + pinned_curie="MONDO:0005148", + pinned_label="type 2 diabetes mellitus", + pinned_category="biolink:Disease", + picked_answer_curies=["CHEBI:6801"], + plover_response=_two_edge_kg(), + ) + assert {e["id"] for e in view["edges"]} == {"edge_cited", "edge_uncited"} diff --git a/pipeline/code/tests/test_iterative_answer_pick.py b/pipeline/code/tests/test_iterative_answer_pick.py new file mode 100644 index 0000000..efdd1b0 --- /dev/null +++ b/pipeline/code/tests/test_iterative_answer_pick.py @@ -0,0 +1,30 @@ +# iterative Stage-11 loop control: the pure helper that validates which +# picks survive (the relevance-stop is now just the LLM's confidence flag, +# checked inline). the LLM call and loop wiring are integration-level; here +# we pin the deterministic carry-forward validation. no network. + +from code.pipeline import _valid_iterative_picks + + +def test_valid_picks_keeps_chunk_and_carry_forward_drops_invented(): + answers = [ + {"curie": "CHEBI:1", "label": "in chunk"}, + {"curie": "CHEBI:2", "label": "carried from earlier chunk"}, + {"curie": "CHEBI:999", "label": "invented, not anywhere"}, + {"label": "no curie"}, + "not a dict", + ] + out = _valid_iterative_picks( + answers, chunk_node_ids={"CHEBI:1"}, prior_curies={"CHEBI:2"}, + ) + assert [a["curie"] for a in out] == ["CHEBI:1", "CHEBI:2"] + + +def test_valid_picks_dedupes_by_curie_first_wins(): + answers = [ + {"curie": "CHEBI:1", "label": "first"}, + {"curie": "CHEBI:1", "label": "dup"}, + ] + out = _valid_iterative_picks(answers, chunk_node_ids={"CHEBI:1"}, prior_curies=set()) + assert len(out) == 1 + assert out[0]["label"] == "first" diff --git a/pipeline/code/tests/test_query_arity.py b/pipeline/code/tests/test_query_arity.py new file mode 100644 index 0000000..78eed8d --- /dev/null +++ b/pipeline/code/tests/test_query_arity.py @@ -0,0 +1,68 @@ +# Stage 9 one-hop arity gate. +# +# PloverDB accepts exactly one hop: 2 query-graph nodes (n0, n1) and 1 +# edge (e0). Stage 8 has an LLM assemble that query graph, so a malformed +# shape (1 node, 3 nodes, a second edge for an illegal multi-hop) is a +# real hallucination risk. reasoner-validator checks Biolink term +# compliance, NOT arity — so this pure check is the only thing that +# structurally enforces the one-hop invariant before the query is sent. +# +# pure function, zero network. asserts the gate accepts the one legal +# shape and rejects every malformed arity, naming the actual counts so a +# failure is diagnosable. + +from code.pipeline import _check_query_graph_arity + + +def _msg(n_nodes: int, n_edges: int) -> dict[str, object]: + nodes = { + f"n{i}": {"categories": ["biolink:NamedThing"]} for i in range(n_nodes) + } + edges = { + f"e{i}": {"subject": "n0", "object": "n1", "predicates": ["biolink:related_to"]} + for i in range(n_edges) + } + return {"message": {"query_graph": {"nodes": nodes, "edges": edges}}} + + +def test_one_hop_passes(): + ok, reason = _check_query_graph_arity(_msg(2, 1)) + assert ok is True + assert reason == "" + + +def test_single_node_rejected(): + # an LLM that emits only n0 (no answer node) — would query nothing + ok, reason = _check_query_graph_arity(_msg(1, 1)) + assert ok is False + assert "1 node" in reason + + +def test_three_nodes_rejected(): + ok, reason = _check_query_graph_arity(_msg(3, 1)) + assert ok is False + assert "3 node" in reason + + +def test_two_edges_is_an_illegal_multi_hop(): + # the core invariant: a second edge means a 2-hop traversal, which + # PloverDB does not accept and the project forbids + ok, reason = _check_query_graph_arity(_msg(2, 2)) + assert ok is False + assert "2 edge" in reason + + +def test_zero_edges_rejected(): + ok, reason = _check_query_graph_arity(_msg(2, 0)) + assert ok is False + assert "0 edge" in reason + + +def test_missing_query_graph_is_rejected_not_crashed(): + # a malformed message must fail the gate gracefully, never raise + ok, reason = _check_query_graph_arity({"message": {}}) + assert ok is False + assert reason + ok_empty, reason_empty = _check_query_graph_arity({}) + assert ok_empty is False + assert reason_empty diff --git a/pipeline/code/tests/test_rerank_coverage_tier.py b/pipeline/code/tests/test_rerank_coverage_tier.py new file mode 100644 index 0000000..ff83fa7 --- /dev/null +++ b/pipeline/code/tests/test_rerank_coverage_tier.py @@ -0,0 +1,74 @@ +# Stage 3 candidate rerank: graph-coverage tier. +# +# _rerank_nameres_candidates orders NameRes candidates by lexical/type +# tiers that BM25 alone won't deliver. this file pins the coverage tier: +# an optional per-candidate edge count (how many edges the candidate has +# to the answer category in the hosted graph) is folded in as the LOWEST +# tier — below type_match, above BM25. that ordering means: +# - a candidate that actually has edges beats an equal one that doesn't +# (so we stop pinning real-looking but data-less entities), but +# - type_match still dominates coverage, so the seizure-case fix (an +# HP PhenotypicFeature candidate beating a higher-BM25 MONDO Disease +# for "seizures") survives even when the MONDO term has more edges. +# +# pure function, zero network. coverage is passed in (the probe that +# produces it is exercised elsewhere); here we pin only the ordering. + +from code.pipeline import _rerank_nameres_candidates + + +def _cand( + curie: str, + label: str, + types: list[str], + score: float, + synonyms: list[str] | None = None, +) -> dict[str, object]: + return { + "curie": curie, + "label": label, + "synonyms": synonyms or [], + "types": types, + "score": score, + } + + +def test_coverage_breaks_ties_within_same_lexical_tier(): + # two candidates identical on every lexical/type tier and BM25; the + # only difference is graph coverage. the one with edges must win. + a = _cand("CURIE:A", "foo", ["biolink:Disease"], 10.0) + b = _cand("CURIE:B", "foo", ["biolink:Disease"], 10.0) + ordered = _rerank_nameres_candidates( + [a, b], "foo", "biolink:Disease", + coverage_by_curie={"CURIE:A": 0, "CURIE:B": 5}, + ) + assert [c["curie"] for c in ordered] == ["CURIE:B", "CURIE:A"] + + +def test_type_match_outranks_coverage(): + # seizure-case preservation: the HP candidate type-matches the + # expected category but has ZERO edges; the MONDO candidate has a much + # higher BM25 and 500 edges but the WRONG type. type_match is a higher + # tier than coverage, so HP must still win despite having no edges. + hp = _cand("HP:1", "phenotype label", ["biolink:PhenotypicFeature"], 10.0) + mondo = _cand("MONDO:1", "disease label", ["biolink:Disease"], 200.0) + ordered = _rerank_nameres_candidates( + [mondo, hp], "seizures", "biolink:PhenotypicFeature", + coverage_by_curie={"HP:1": 0, "MONDO:1": 500}, + ) + assert ordered[0]["curie"] == "HP:1" + + +def test_coverage_none_orders_by_bm25_not_edges(): + # when no coverage is supplied the tier is inert: candidates identical + # on lexical/type tiers fall back to BM25 order, exactly as before the + # coverage tier existed. (contrast: with coverage favouring B, B wins.) + a = _cand("CURIE:A", "foo", ["biolink:Disease"], 20.0) + b = _cand("CURIE:B", "foo", ["biolink:Disease"], 10.0) + without = _rerank_nameres_candidates([a, b], "foo", "biolink:Disease") + assert [c["curie"] for c in without] == ["CURIE:A", "CURIE:B"] + with_cov = _rerank_nameres_candidates( + [a, b], "foo", "biolink:Disease", + coverage_by_curie={"CURIE:A": 0, "CURIE:B": 5}, + ) + assert [c["curie"] for c in with_cov] == ["CURIE:B", "CURIE:A"] diff --git a/pipeline/code/tests/test_smoke.py b/pipeline/code/tests/test_smoke.py index 1be4f1b..511a926 100644 --- a/pipeline/code/tests/test_smoke.py +++ b/pipeline/code/tests/test_smoke.py @@ -18,6 +18,7 @@ EXPECTED_STATUSES = { "ok", "invalid_query", + "invalid_query_arity", "llm_bad_json", "plover_error", "llm_error", @@ -26,7 +27,6 @@ "entity_empty", "out_of_scope", "no_candidate_match", - "low_confidence_resolution", } From a66174ba677a5effd7ba22075a64340f3a54b7ce Mon Sep 17 00:00:00 2001 From: Adilbek Bazarkulov Date: Mon, 8 Jun 2026 12:31:49 -0400 Subject: [PATCH 6/6] Add golden-question evaluation harness --- pipeline/code/golden_eval.py | 401 +++++++++++++++++++++++++++++++++++ 1 file changed, 401 insertions(+) create mode 100644 pipeline/code/golden_eval.py diff --git a/pipeline/code/golden_eval.py b/pipeline/code/golden_eval.py new file mode 100644 index 0000000..34b9120 --- /dev/null +++ b/pipeline/code/golden_eval.py @@ -0,0 +1,401 @@ +# golden_eval.py — run gold questions through the full grounded pipeline +# with a chosen model and score each result against the gold +# verified_answers. drives a cost-conscious one-at-a-time loop: with +# --stop-on-fail (default) it halts at the first question that does not +# clear the objective floor, so we never pay to run all 10 to discover #3 +# broke. +# +# scoring philosophy: the gold `verified_answers` is the FLOOR (manually +# verified known-correct answers), not the ceiling. recall of that set is +# the objective number; EXTRA answers the pipeline returns are surfaced for +# human review, NOT counted as errors (they may be additional correct +# answers). entity + predicate correctness are the Q1 (NL->TRAPI) signals; +# explanation.md faithfulness is a human read. +# +# this hits the network and spends real OpenRouter credits (~6 LLM calls +# per question), so it is a CLI tool, not a CI test. invoke from pipeline/: +# python -m code.golden_eval --model m6 --questions q1 +# python -m code.golden_eval --model m6 --no-stop-on-fail # final all-10 + +from __future__ import annotations + +import argparse +import json +import logging +from dataclasses import dataclass, replace +from pathlib import Path +from typing import Any + +from dotenv import load_dotenv + +from .biolink_helper import ( + build_neighborhood_map as build_biolink_neighborhood_map, + make_toolkit as make_biolink_toolkit, +) +from .config import Config, load_config +from .logging_setup import console, setup_logger, utc_stamp +from .nameres_client import NameResClient +from .nodenorm_client import NodeNormClient, NodeNormError +from .openrouter_client import OpenRouterClient +from .pipeline import QuestionResult, run_grounded +from .plover_client import PloverClient, PloverError +from .trace import QuestionPaths, make_run_dir, make_run_root, write_json + +ENV_FILE = Path(__file__).resolve().parent.parent / ".env" + +# the single-shot answer step still caps at 5 picks (the iterative redesign +# lifts this). when the cap is hit AND gold is still missing, we flag the +# run as "capped" rather than failed — the miss may be the cap, not a bug. +SINGLE_SHOT_ANSWER_CAP = 5 + + +@dataclass +class Runtime: + llm: OpenRouterClient + plover: PloverClient + nameres: NameResClient + nodenorm: NodeNormClient + predicate_index: dict[tuple[str, str], list[str]] + biolink_neighborhoods: dict[str, list[str]] + + def close(self) -> None: + self.llm.close() + self.plover.close() + self.nameres.close() + self.nodenorm.close() + + +@dataclass +class Score: + qid: str + nl_question: str + status: str + outcome: str | None + cost_usd: float + n_picked: int + entity_expected: str | None + entity_pinned: str | None + entity_ok: bool + predicate_expected: str | None + predicate_built: list[str] + predicate_ok: bool + gold_curies: list[str] + matched: list[str] + missing: list[str] + extras: list[str] + recall: float + capped: bool + passes_floor: bool + explanation_path: Path + + +def build_runtime(cfg: Config, logger: logging.Logger) -> Runtime: + # mirrors the per-run setup in runner.main(): one client of each kind, + # the cached meta_KG predicate index (Stage 8 constraint), and the BMT + # loose-neighborhood map (Stage 3 filter). kept here so the eval tool + # does not reach into the benchmark runner's CLI internals. + llm = OpenRouterClient(cfg, logger) + plover = PloverClient(cfg, logger) + nameres = NameResClient(cfg, logger) + nodenorm = NodeNormClient(cfg, logger) + + predicate_index: dict[tuple[str, str], list[str]] = {} + try: + meta_kg = plover.fetch_meta_kg() + for edge in meta_kg.get("edges") or []: + subject, obj, predicate = edge.get("subject"), edge.get("object"), edge.get("predicate") + if subject and obj and predicate: + predicate_index.setdefault((subject, obj), []).append(predicate) + for pair in list(predicate_index.keys()): + predicate_index[pair] = sorted(set(predicate_index[pair])) + except PloverError as e: + logger.warning(f"could not fetch meta_KG: {e} (Stage 8 runs unconstrained)") + + available_categories = sorted({c for pair in predicate_index for c in pair}) + biolink_neighborhoods: dict[str, list[str]] = {} + try: + toolkit = make_biolink_toolkit() + biolink_neighborhoods = build_biolink_neighborhood_map( + available_categories, toolkit, logger, + ) + except Exception as e: # non-fatal startup cache (mirrors runner); Stage 3 degrades to strict + logger.warning(f"could not build biolink neighborhoods: {e}") + + return Runtime(llm, plover, nameres, nodenorm, predicate_index, biolink_neighborhoods) + + +def _load_json(path: Path) -> dict[str, Any]: + # artifacts may be absent when the run failed before that stage; treat + # a missing/unreadable file as empty rather than crashing the scorer. + if not path.is_file(): + return {} + try: + loaded = json.loads(path.read_text()) + except (json.JSONDecodeError, OSError): + return {} + if isinstance(loaded, dict): + return {str(k): v for k, v in loaded.items()} + return {} + + +def _load_gold(cfg: Config, qid: str) -> dict[str, Any]: + loaded: Any = json.loads((cfg.paths.questions / f"{qid}.json").read_text()) + if not isinstance(loaded, dict): + raise TypeError(f"gold file {qid}.json is not a JSON object") + return {str(k): v for k, v in loaded.items()} + + +def _predicates_in_query(trapi_msg: dict[str, Any]) -> list[str]: + query_graph = trapi_msg.get("message", {}).get("query_graph", {}) + edges = query_graph.get("edges", {}) if isinstance(query_graph, dict) else {} + out: list[str] = [] + if isinstance(edges, dict): + for edge in edges.values(): + if not isinstance(edge, dict): + continue + predicates = edge.get("predicates") + if isinstance(predicates, list): + out.extend(str(p) for p in predicates) + elif isinstance(edge.get("predicate"), str): + out.append(str(edge["predicate"])) + return out + + +def _picked_curies(answer_doc: dict[str, Any]) -> list[str]: + out: list[str] = [] + for answer in answer_doc.get("answers", []) or []: + if isinstance(answer, dict) and isinstance(answer.get("curie"), str): + out.append(answer["curie"]) + return out + + +def _canonical_map( + nodenorm: NodeNormClient, curies: list[str], logger: logging.Logger, +) -> dict[str, str]: + # canonicalise every curie so gold and picked are compared on the same + # footing (the gold curies are already canonical, but normalising both + # is robust to equivalent-identifier drift). falls back to identity on + # a NodeNorm error rather than silently dropping curies. + unique = [c for c in dict.fromkeys(curies) if c] + if not unique: + return {} + try: + result = nodenorm.normalize(unique) + except NodeNormError as e: + logger.warning(f"NodeNorm failed during scoring; comparing raw curies: {e}") + return {c: c for c in unique} + return {c: (result.canonical.get(c) or c) for c in unique} + + +def _score( + qid: str, + gold: dict[str, Any], + qp: QuestionPaths, + result: QuestionResult, + nodenorm: NodeNormClient, + logger: logging.Logger, +) -> Score: + entity_expected = (gold.get("pinned_entity") or {}).get("curie") + predicate_expected = gold.get("predicate") + gold_curies = [ + va["curie"] + for va in (gold.get("verified_answers") or []) + if isinstance(va, dict) and va.get("curie") + ] + + pinned = _load_json(qp.nodenorm).get("pinned") or {} + entity_pinned = pinned.get("canonical_curie") if isinstance(pinned, dict) else None + predicate_built = _predicates_in_query(_load_json(qp.trapi_query)) + picked = _picked_curies(_load_json(qp.answer)) + + canon = _canonical_map( + nodenorm, + [*gold_curies, *picked, entity_expected or "", entity_pinned or ""], + logger, + ) + + def canon_set(items: list[str]) -> set[str]: + return {canon.get(x, x) for x in items} + + def canon_one(curie: str | None) -> str | None: + return canon.get(curie, curie) if curie else None + + gold_set = canon_set(gold_curies) + picked_set = canon_set(picked) + matched = sorted(gold_set & picked_set) + missing = sorted(gold_set - picked_set) + extras = sorted(picked_set - gold_set) + recall = len(matched) / len(gold_set) if gold_set else 0.0 + + entity_ok = bool(entity_pinned) and canon_one(entity_pinned) == canon_one(entity_expected) + predicate_ok = bool(predicate_expected) and predicate_expected in predicate_built + capped = len(picked) >= SINGLE_SHOT_ANSWER_CAP and bool(missing) + # objective floor: the run succeeded, pinned the right entity, built the + # right predicate, and found at least one gold answer. the final + # hold/advance call layers human judgement (faithfulness, extras) on top. + passes_floor = ( + result.status == "ok" and entity_ok and predicate_ok and len(matched) > 0 + ) + + return Score( + qid=qid, + nl_question=gold.get("nl_question", ""), + status=result.status, + outcome=result.outcome, + cost_usd=result.cost_total_usd, + n_picked=len(picked), + entity_expected=entity_expected, + entity_pinned=entity_pinned, + entity_ok=entity_ok, + predicate_expected=predicate_expected, + predicate_built=predicate_built, + predicate_ok=predicate_ok, + gold_curies=gold_curies, + matched=matched, + missing=missing, + extras=extras, + recall=recall, + capped=capped, + passes_floor=passes_floor, + explanation_path=qp.explanation, + ) + + +def _print_verdict(score: Score) -> None: + ok = "[green]OK[/]" + bad = "[red]WRONG[/]" + console.rule(f"[bold]{score.qid}[/] {score.nl_question}") + console.print( + f"status=[bold]{score.status}[/] outcome={score.outcome} " + f"picks={score.n_picked} cost=${score.cost_usd:.4f}" + ) + console.print( + f"Q1 entity {ok if score.entity_ok else bad} " + f"expected={score.entity_expected} pinned={score.entity_pinned}" + ) + console.print( + f"Q1 predicate {ok if score.predicate_ok else bad} " + f"expected={score.predicate_expected} built={score.predicate_built}" + ) + console.print( + f"Q2 recall [bold]{len(score.matched)}/{len(score.gold_curies)}[/] " + f"gold found ({score.recall:.0%})" + ) + console.print(f" matched: {score.matched or '—'}") + console.print(f" [yellow]MISSING[/]: {score.missing or '—'}") + console.print(f" extras (review, not errors): {score.extras or '—'}") + if score.capped: + console.print( + " [yellow]capped[/]: 5-pick limit hit with gold still missing — " + "likely the single-shot answer cap (lifted by the iterative redesign)" + ) + console.print(f" explanation: {score.explanation_path}") + console.print( + f"floor: {'[green]PASS[/]' if score.passes_floor else '[red]HOLD[/]'} " + "(advance decision adds a human read of faithfulness + extras)" + ) + + +def _parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser( + prog="ploverai-golden-eval", + description="run gold questions through the full pipeline and score vs gold.", + ) + p.add_argument("--model", default="m6", help="model id from config.yaml (default: m6, Gemini Flash).") + p.add_argument( + "--questions", nargs="+", default=None, + help="gold question ids in run order (e.g. --questions q1 q2). default: all 10.", + ) + p.add_argument( + "--stop-on-fail", action=argparse.BooleanOptionalAction, default=True, + help="stop at the first question that fails the objective floor (default: on).", + ) + p.add_argument( + "--iterative", action="store_true", + help="enable iterative chunked answer-picking (overrides config default-off).", + ) + return p.parse_args() + + +def main() -> int: + load_dotenv(dotenv_path=ENV_FILE) + args = _parse_args() + cfg = load_config() + # --iterative flips the (frozen) config's default-off iterative flag so a + # single command can A/B single-shot vs iterative without editing yaml. + if args.iterative: + cfg = replace( + cfg, stage11_iterative=replace(cfg.stage11_iterative, enabled=True), + ) + question_ids = args.questions or [f"q{i}" for i in range(1, 11)] + + run_id = utc_stamp() + logger, log_path = setup_logger(cfg.paths.logs, run_id) + runtime = build_runtime(cfg, logger) + model = cfg.model(args.model) + # chunk budget is sized from the MANUAL context_window in config.yaml. + budget = int(model.context_window * cfg.stage11_iterative.context_fraction) + logger.info( + f"golden_eval started run_id={run_id} model={model.id} " + f"context_window={model.context_window} chunk_budget={budget}" + ) + console.print( + f"[bold]golden_eval[/] model=[cyan]{model.id}[/] ({model.slug}) " + f"context_window={model.context_window} iterative={cfg.stage11_iterative.enabled} " + f"questions={question_ids} stop_on_fail={args.stop_on_fail} log={log_path}" + ) + + run_root = make_run_root(make_run_dir(cfg.paths.results, run_id), model.id, model.slug) + # record the resolved context window + iterative config for reproducibility. + write_json(run_root.run_meta, { + "run_id": run_id, + "model": {"id": model.id, "slug": model.slug, "context_window": model.context_window}, + "iterative": { + "enabled": cfg.stage11_iterative.enabled, + "context_fraction": cfg.stage11_iterative.context_fraction, + "answer_target": cfg.stage11_iterative.answer_target, + "max_chunks": cfg.stage11_iterative.max_chunks, + "chunk_budget": budget, + }, + }) + + scores: list[Score] = [] + try: + for qid in question_ids: + gold = _load_gold(cfg, qid) + qp = QuestionPaths.under(run_root.root, qid) + result = run_grounded( + cfg=cfg, model=model, q={"id": qid, "nl_question": gold["nl_question"]}, + qp=qp, llm=runtime.llm, nameres=runtime.nameres, nodenorm=runtime.nodenorm, + plover=runtime.plover, logger=logger, + predicate_index=runtime.predicate_index, + biolink_neighborhoods=runtime.biolink_neighborhoods, + ) + score = _score(qid, gold, qp, result, runtime.nodenorm, logger) + scores.append(score) + _print_verdict(score) + if args.stop_on_fail and not score.passes_floor: + console.print( + f"[red]stopping[/]: {qid} did not clear the floor. " + "fix, then re-run this question before advancing." + ) + break + finally: + runtime.close() + + n_pass = sum(1 for s in scores if s.passes_floor) + console.rule("[bold]summary[/]") + for s in scores: + mark = "[green]PASS[/]" if s.passes_floor else "[red]HOLD[/]" + console.print( + f" {s.qid:>4} {mark} recall={len(s.matched)}/{len(s.gold_curies)} " + f"entity={'ok' if s.entity_ok else 'WRONG'} " + f"predicate={'ok' if s.predicate_ok else 'WRONG'} status={s.status}" + ) + total_cost = sum(s.cost_usd for s in scores) + console.print(f"[bold]{n_pass}/{len(scores)} cleared the floor[/] total_cost=${total_cost:.4f}") + return 0 if n_pass == len(scores) and scores else 1 + + +if __name__ == "__main__": + raise SystemExit(main())