diff --git a/src/kb/query/engine.py b/src/kb/query/engine.py index 9bc13b7..f2ee395 100644 --- a/src/kb/query/engine.py +++ b/src/kb/query/engine.py @@ -10,7 +10,8 @@ import asyncio import re import time -from typing import Any +from dataclasses import dataclass +from typing import Any, Literal import structlog @@ -41,6 +42,94 @@ _CITE_RE = re.compile(r"\[(\d+)\]") +@dataclass(frozen=True) +class _CitationSource: + via: Literal["graph_route", "retrieval"] + file_id: str + filename: str + page_start: int + page_end: int + excerpt: str + hit: dict[str, Any] | None = None + + +def _build_graph_sources( + mentions: list[dict[str, Any]], filenames: dict[str, str], excerpt_chars: int +) -> list[_CitationSource]: + """Normalize graph-route mentions into citation sources. + + Dedupes by file/page/excerpt so one graph mention does not get duplicated + into the source list and shift downstream numbering. + """ + out: list[_CitationSource] = [] + seen: set[tuple[str, int, int, str]] = set() + for m in mentions: + fid = str(m.get("file_id") or "") + if not fid: + continue + ps = int(m.get("page_start") or 0) + pe = int(m.get("page_end") or ps) + excerpt = (m.get("excerpt") or "")[:excerpt_chars] + key = (fid, ps, pe, excerpt) + if key in seen: + continue + seen.add(key) + out.append( + _CitationSource( + via="graph_route", + file_id=fid, + filename=filenames.get(fid, "unknown"), + page_start=ps, + page_end=pe, + excerpt=excerpt, + ) + ) + return out + + +def _build_retrieval_sources(hits: list[dict[str, Any]]) -> list[_CitationSource]: + return [ + _CitationSource( + via="retrieval", + file_id=str(h["metadata"].get("file_id", "")), + filename="", + page_start=int(h["metadata"].get("page_start") or 0), + page_end=int(h["metadata"].get("page_end") or h["metadata"].get("page_start") or 0), + excerpt="", + hit=h, + ) + for h in hits + ] + + +def _format_numbered_sources( + graph_sources: list[_CitationSource], retrieval_sources: list[_CitationSource] +) -> str: + out: list[str] = [] + if graph_sources: + out.append("Graph sources:") + for i, src in enumerate(graph_sources, start=1): + page = src.page_start + if src.page_end and src.page_end != page: + page = f"{src.page_start}-{src.page_end}" + out.append(f"[{i}] (file={src.file_id[:8]} page={page})\n{src.excerpt}") + if retrieval_sources: + if out: + out.append("") + out.append("Retrieval sources:") + offset = len(graph_sources) + for i, src in enumerate(retrieval_sources, start=1): + h = src.hit or {} + md = h.get("metadata", {}) + page = md.get("page_start", "?") + if md.get("page_end") and md.get("page_end") != page: + page = f"{md['page_start']}-{md['page_end']}" + out.append( + f"[{offset + i}] (file={md.get('file_id', '?')[:8]} page={page})\n{h.get('text', '')}" + ) + return "\n\n".join(out) + + def _extract_cited_indices(answer: str) -> list[int]: return sorted({int(m.group(1)) for m in _CITE_RE.finditer(answer)}) @@ -568,11 +657,22 @@ def _section_score(h: Any) -> float: ) ) - # Map chunk_id -> [primary file_id, ...also_in_files] for multi-source citations. - sources_by_chunk = consolidate_sources(hits) - # ── Stage 4: synthesis ─────────────────────────────────────────────────── sys_prompt = pipeline.get(cfg, "prompts.synthesize_system", "") + graph_sources: list[_CitationSource] = [] + if graph_result and graph_result.get("mentions"): + graph_filenames = await _resolve_filenames( + list({m["file_id"] for m in graph_result["mentions"] if m.get("file_id")}) + ) + graph_sources = _build_graph_sources( + graph_result["mentions"], + graph_filenames, + int(pipeline.get(cfg, "synthesize.excerpt_chars", 400)), + ) + retrieval_sources = _build_retrieval_sources(serializable_hits) + # Keep the existing chunk/source mapping for retrieval-only citations. + sources_by_chunk = consolidate_sources(hits) + combined_sources = graph_sources + retrieval_sources structured_block = "" if structured: structured_block += ( @@ -593,7 +693,7 @@ def _section_score(h: Any) -> float: ) user_prompt = ( f"Question: {body.question}{history_block}{structured_block}\n\n" - f"Sources:\n{_format_sources(serializable_hits)}\n\n" + f"Sources:\n{_format_numbered_sources(graph_sources, retrieval_sources)}\n\n" "Answer the question using ONLY the sources above. Cite using inline [n] markers " "tied to the source numbers. If the sources don't support an answer, say so " "explicitly and report low confidence. " @@ -648,7 +748,7 @@ def _section_score(h: Any) -> float: confidence_reason = f"CRAG retrieval score {crag_score:.2f}: {crag_reason}" cited_indices = _extract_cited_indices(answer_text) - if not cited_indices and refuse_no_cite and serializable_hits: + if not cited_indices and refuse_no_cite and combined_sources: answer_text = ( "I cannot answer with citations from the provided sources. " "The retrieved excerpts do not directly support a confident answer to this question." @@ -730,39 +830,58 @@ def _section_score(h: Any) -> float: } ) - # ── Stage 5: span-level citations (multi-source aware) ─────────────────── - # Resolve filenames for the cited chunks + every file in their also_in_files. - file_ids_to_resolve: set[str] = set() + # ── Stage 5: span-level citations (global numbering across graph + retrieval) ─── + started = time.time() + retrieval_file_ids_to_resolve: set[str] = set() for i in cited_indices: - if not (1 <= i <= len(serializable_hits)): - continue - h = serializable_hits[i - 1] - for fid in sources_by_chunk.get(h["id"], []): - if fid: - file_ids_to_resolve.add(fid) - filenames = await _resolve_filenames(list(file_ids_to_resolve)) + if len(graph_sources) < i <= len(combined_sources): + src = combined_sources[i - 1] + if src.via == "retrieval": + h = src.hit or {} + md = h.get("metadata", {}) + primary = md.get("file_id", "") + if primary: + retrieval_file_ids_to_resolve.add(primary) + for fid in sources_by_chunk.get(h.get("id", ""), []): + if fid: + retrieval_file_ids_to_resolve.add(fid) + retrieval_filenames = await _resolve_filenames(list(retrieval_file_ids_to_resolve)) excerpt_chars = int(pipeline.get(cfg, "synthesize.excerpt_chars", 400)) - started = time.time() citations: list[Citation] = [] for i in cited_indices: - if not (1 <= i <= len(serializable_hits)): + if not (1 <= i <= len(combined_sources)): continue - h = serializable_hits[i - 1] - md = h["metadata"] + src = combined_sources[i - 1] + if src.via == "graph_route": + citations.append( + Citation( + file_id=src.file_id, + filename=src.filename, + page_start=src.page_start, + page_end=src.page_end, + excerpt=src.excerpt, + also_in=[], + bbox=None, + via="graph_route", + ) + ) + continue + h = src.hit or {} + md = h.get("metadata", {}) excerpt = await pick_best_span( - query=body.question, chunk_text=h["text"], max_chars=excerpt_chars + query=body.question, chunk_text=h.get("text", ""), max_chars=excerpt_chars ) primary = md.get("file_id", "") also_files = [ - CitationSource(file_id=f, filename=filenames.get(f, "unknown")) - for f in sources_by_chunk.get(h["id"], []) + CitationSource(file_id=f, filename=retrieval_filenames.get(f, "unknown")) + for f in sources_by_chunk.get(h.get("id", ""), []) if f and f != primary ] citations.append( Citation( file_id=primary, - filename=filenames.get(primary, "unknown"), + filename=retrieval_filenames.get(primary, "unknown"), page_start=int(md.get("page_start") or 0), page_end=int(md.get("page_end") or md.get("page_start") or 0), excerpt=excerpt, @@ -772,41 +891,6 @@ def _section_score(h: Any) -> float: ) ) - # ── Graph-route citation backfill ────────────────────────────────────── - # When the GraphRAG-shaped theme route fired, its narrative summary - # influenced the answer's structure (themes, groupings) but its - # underlying entity_mentions weren't surfaced as citations above — - # because they came from the entity graph, not from retrieval. Backfill - # them here, deduped against retrieval-sourced citations by - # (file_id, page_start). Marked via="graph_route" so consumers can tell - # which evidence shaped the prose vs which directly grounds a [n] marker. - if graph_result and graph_result.get("mentions"): - seen_keys = {(c.file_id, c.page_start) for c in citations} - graph_filenames = await _resolve_filenames( - list({m["file_id"] for m in graph_result["mentions"] if m.get("file_id")}) - ) - for m in graph_result["mentions"]: - fid = m.get("file_id") - if not fid: - continue - ps = int(m.get("page_start") or 0) - key = (fid, ps) - if key in seen_keys: - continue - seen_keys.add(key) - citations.append( - Citation( - file_id=fid, - filename=graph_filenames.get(fid, "unknown"), - page_start=ps, - page_end=int(m.get("page_end") or ps), - excerpt=(m.get("excerpt") or "")[:excerpt_chars], - also_in=[], - bbox=None, - via="graph_route", - ) - ) - stages.append(_stage("span_cite", started, citations=len(citations))) retrieved_nodes = [ diff --git a/src/kb/query/structured.py b/src/kb/query/structured.py index e197b48..2607f97 100644 --- a/src/kb/query/structured.py +++ b/src/kb/query/structured.py @@ -15,6 +15,7 @@ import json import operator +import re from typing import Any import structlog @@ -28,6 +29,14 @@ # Map of simple comparison operators we accept in NL-parsed numeric filters. _OPS = {">": operator.gt, ">=": operator.ge, "<": operator.lt, "<=": operator.le, "=": operator.eq} +_FIELD_KEY_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") + + +def _safe_field_key(key: str) -> str | None: + """Return a SQL-safe JSON field key or None if the key is not valid.""" + if isinstance(key, str) and _FIELD_KEY_RE.fullmatch(key): + return key + return None async def list_entities_matching( @@ -51,16 +60,20 @@ async def list_entities_matching( params["t"] = entity_type for i, (k, v) in enumerate(filters.items()): + safe_key = _safe_field_key(k) + if not safe_key: + logger.warning("dropping unsafe structured filter key: %r", k) + continue key = f"f{i}" if isinstance(v, tuple) and len(v) == 2 and v[0] in _OPS: op, val = v - conds.append(f"(fields->>'{k}')::numeric {op} :{key}") + conds.append(f"(fields->>'{safe_key}')::numeric {op} :{key}") params[key] = val else: # Fuzzy-string match by default: "Q2 2024" matches "Q2 FY2024", etc. # Intent extraction often produces normalized strings that don't # exactly equal the stored field; ILIKE with %v% catches both. - conds.append(f"(fields->>'{k}') ILIKE :{key}") + conds.append(f"(fields->>'{safe_key}') ILIKE :{key}") params[key] = f"%{v}%" sql = ( @@ -134,7 +147,7 @@ async def maybe_structured_answer( Otherwise returns a dict with `entities`, `mentions`, and a natural-language summary the synthesis step can cite from. """ - if intent.kind != "aggregate": + if intent.kind not in ("aggregate", "compare"): return None threshold = parse_numeric_threshold(question) filters: dict[str, Any] = {} diff --git a/tests/test_query_synthesis.py b/tests/test_query_synthesis.py index c8d0cfc..72ba858 100644 --- a/tests/test_query_synthesis.py +++ b/tests/test_query_synthesis.py @@ -2,9 +2,18 @@ from __future__ import annotations +import asyncio import re -from kb.query.engine import _extract_cited_indices, _format_sources +from kb.query.engine import ( + _build_graph_sources, + _build_retrieval_sources, + _extract_cited_indices, + _format_numbered_sources, + _format_sources, +) +from kb.query.intent import QueryIntent +from kb.query.structured import _safe_field_key, maybe_structured_answer def test_cite_re_finds_numbers() -> None: @@ -26,6 +35,81 @@ def test_format_sources_includes_index_and_page() -> None: assert "page=8-10" in out +def test_numbered_sources_put_graph_before_retrieval() -> None: + graph_sources = _build_graph_sources( + [ + { + "file_id": "graphfile-1", + "page_start": 1, + "page_end": 1, + "excerpt": "graph evidence", + } + ], + {"graphfile-1": "graph.txt"}, + 200, + ) + retrieval_sources = _build_retrieval_sources( + [ + { + "text": "retrieval evidence", + "metadata": {"file_id": "retrfile-1", "page_start": 2, "page_end": 2}, + } + ] + ) + + out = _format_numbered_sources(graph_sources, retrieval_sources) + assert out.index("[1]") < out.index("[2]") + assert "graph evidence" in out + assert "retrieval evidence" in out + + +def test_safe_field_key_rejects_injection_shape() -> None: + assert _safe_field_key("ticker") == "ticker" + assert _safe_field_key("x') IS NULL OR 1=1 --") is None + + +def test_compare_questions_can_use_structured_path(monkeypatch) -> None: + async def fake_list_entities_matching( + *, domain: str, entity_type: str | None, filters: dict, limit: int, project: str + ): + assert domain == "sec" + assert entity_type == "FinancialMetric" + return [ + {"id": "e1", "display_name": "Revenue", "identity_key": "rev", "fields": {"value": 1}} + ] + + async def fake_mentions_for(entity_ids: list[str], project: str = "default"): + assert entity_ids == ["e1"] + return [ + { + "entity_id": "e1", + "file_id": "f1", + "filename": "file.txt", + "page_start": 1, + "page_end": 1, + "excerpt": "x", + } + ] + + import kb.query.structured as structured_mod + + monkeypatch.setattr(structured_mod, "list_entities_matching", fake_list_entities_matching) + monkeypatch.setattr(structured_mod, "mentions_for", fake_mentions_for) + out = asyncio.run( + maybe_structured_answer( + intent=QueryIntent( + kind="compare", entity_type="FinancialMetric", filters={"ticker": "AAPL"} + ), + domain="sec", + question="Compare NVIDIA and Apple revenue", + project="default", + ) + ) + + assert out is not None + assert out["entities"][0]["id"] == "e1" + + def test_trailing_confidence_json_pattern() -> None: sample = 'The answer is X. [1] {"confidence": 0.83, "confidence_reason": "well-supported"}' m = re.search(r"\{[^{}]*confidence[^{}]*\}\s*$", sample)