diff --git a/studio/backend/core/rag/retrieval.py b/studio/backend/core/rag/retrieval.py index e6eb45c276..d963737213 100644 --- a/studio/backend/core/rag/retrieval.py +++ b/studio/backend/core/rag/retrieval.py @@ -5,6 +5,7 @@ from __future__ import annotations +import re from dataclasses import dataclass from utils.rag.config import ( @@ -16,6 +17,31 @@ from utils.rag.config import ( from . import bm25, embeddings, vector_store +# Match "Figure 1", "Figure 1.2", "Figure B.1", "Table 4", "Fig. 5" anywhere +# in the query. Used to inject a third retrieval source that directly looks +# up chunks anchored by these references — dense vectors don't preserve +# figure numbers, so without this an exact-numbered query gets out-ranked +# by chunks describing other figures that share more vocabulary with the +# question. +_FIGURE_REF_RE = re.compile( + r"\b(Figure|Fig\.|Table|Tab\.)\s+([A-Z]?\.?\d+(?:\.\d+)?)\b", + re.IGNORECASE, +) + + +def _extract_figure_refs(query: str) -> list[str]: + """Return normalized 'Figure N' / 'Table N' references found in query.""" + refs: list[str] = [] + seen: set[str] = set() + for m in _FIGURE_REF_RE.finditer(query): + head = m.group(1).lower() + label = "Figure" if head.startswith("fig") else "Table" + ref = f"{label} {m.group(2)}" + if ref not in seen: + seen.add(ref) + refs.append(ref) + return refs + @dataclass(frozen = True) class Hit: @@ -33,6 +59,67 @@ def retrieve_bm25(scope: str, query: str, k: int | None = None) -> list[Hit]: return [Hit(chunk_id = cid, score = s) for cid, s in bm25.search(scope, query, limit)] +def retrieve_figure_refs( + scope: str, + query: str, + *, + k: int = 5, + document_ids: list[str] | None = None, +) -> list[Hit]: + """Look up chunks anchored at a 'Figure N:' / 'Table N:' caption that + the query references. Returns at most ``k`` hits — usually 0 or 1. + + Chunks produced by the figure-boundary chunker start with the literal + caption, so a SQL prefix match is enough; we don't need full-text + search here. + """ + refs = _extract_figure_refs(query) + if not refs: + return [] + + from .db import get_rag_connection + + placeholders_docs = "" + params: list = [scope] + if document_ids: + placeholders_docs = ( + f" AND document_id IN ({','.join('?' * len(document_ids))})" + ) + params.extend(document_ids) + + like_clauses: list[str] = [] + for ref in refs: + # Match "Figure 1:" and "Figure 1." (period-terminated captions). + like_clauses.append( + "json_extract(payload_json, '$.text') LIKE ?" + " OR json_extract(payload_json, '$.text') LIKE ?" + ) + params.extend([f"{ref}:%", f"{ref}.%"]) + + sql = ( + "SELECT chunk_id, document_id, chunk_index, kind" + " FROM rag_vectors" + f" WHERE scope = ?{placeholders_docs}" + " AND kind = 'text'" + f" AND ({' OR '.join(like_clauses)})" + f" LIMIT {int(k)}" + ) + + out: list[Hit] = [] + with get_rag_connection() as conn: + for row in conn.execute(sql, params): + out.append( + Hit( + chunk_id = row[0], + score = 1.0, + document_id = row[1], + chunk_index = row[2], + kind = row[3] or "text", + ) + ) + return out + + def retrieve_dense( scope: str, query: str, @@ -121,8 +208,12 @@ def retrieve_hybrid( document_ids = document_ids, embedder_model = embedder_model, ) + rankings = [bm25_hits, dense_hits] + fig_hits = retrieve_figure_refs(scope, query, document_ids = document_ids) + if fig_hits: + rankings.append(fig_hits) return _rrf_fuse( - [bm25_hits, dense_hits], + rankings, rrf_k = RAG_RRF_K, top_k = k or RAG_TOP_K_HYBRID, )