"""Retrieve MASTER or REFERENCE paragraphs for mapping messy field notes.""" from __future__ import annotations import re from typing import Literal from backend.config import settings from backend.core.rag_store import SearchHit, get_rag_store from backend.models.schema import TemplateSchema InterferenceLevel = Literal["minimum", "medium", "maximum"] RetrievalLevel = Literal["paragraph", "section", "document"] _STOP = frozenset({ "about", "with", "from", "this", "that", "level", "section", "inside", "outside", "other", "property", "report", "survey", "your", "the", "and", "were", "was", "has", "have", "been", "noted", "inspected", "ground", }) def build_retrieval_query( section_label: str, observations: list[str], *, section_id: str = "", ) -> str: """Semantic query anchored to the target section ID + observations.""" sid = (section_id or "").strip().upper() label = section_label.strip() prefix = f"{sid} {label}".strip() if sid else label parts = [o.strip() for o in observations if o and o.strip()] if not parts: return prefix ranked = sorted(parts, key=len, reverse=True) return f"{prefix}: " + " ".join(ranked[:5]) def _content_tokens(text: str) -> set[str]: return { w for w in re.findall(r"[a-z]{4,}", text.lower()) if w not in _STOP } def rerank_hits_by_observations( hits: list[SearchHit], observations: list[str], ) -> list[SearchHit]: if not hits or not observations: return hits obs_tokens = _content_tokens(" ".join(observations)) if not obs_tokens: return hits boost = settings.retrieval_lexical_boost def rank_key(h: SearchHit) -> float: para_tokens = _content_tokens(h.text) overlap = len(obs_tokens & para_tokens) return h.score + overlap * boost return sorted(hits, key=rank_key, reverse=True) def rerank_hits_for_comprehensive_baseline(hits: list[SearchHit]) -> list[SearchHit]: """Among similarly scored REFERENCE hits, prefer the longest / fullest text block.""" if len(hits) < 2: return hits top_score = hits[0].score score_floor = top_score - max(0.06, top_score * 0.08) def sort_key(h: SearchHit) -> tuple[float, int, int]: return ( h.score, len(h.text.split()), -(h.paragraph_index or 0), ) close = [h for h in hits if h.score >= score_floor] tail = [h for h in hits if h.score < score_floor] close.sort(key=sort_key, reverse=True) return close + tail def _filter_hits_to_section( hits: list[SearchHit], paragraph_section_id: str, ) -> list[SearchHit]: """Keep only chunks tagged with the target section (case-insensitive).""" want = (paragraph_section_id or "").strip().upper() if not want: return hits return [ h for h in hits if (h.section_id or "").strip().upper() == want ] def _complete_section_group( tenant_id: str, source_key: str | None, section_id: str, *, tier: str | None, allowed_doc_keys: frozenset[str] | None, property_context: dict | None, ) -> list[SearchHit]: """Pull EVERY stored chunk for a section (optionally pinned to one source). Mirrors the property guard applied to similarity hits so a section-complete baseline cannot reintroduce wrong-property content the ranked path would drop. """ from backend.core.rag_store import TIER_REFERENCE, get_rag_store use_tier = tier or TIER_REFERENCE chunks = get_rag_store().fetch_section_chunks( tenant_id, tier=use_tier, section_id=section_id, source_key=source_key, allowed_doc_keys=allowed_doc_keys, ) if property_context and chunks: from backend.core.property_blocklist import filter_hits_by_property chunks = filter_hits_by_property(chunks, property_context) return chunks def assemble_reference_baseline( hits: list[SearchHit], *, paragraph_section_id: str = "", tenant_id: str | None = None, tier: str | None = None, allowed_doc_keys: frozenset[str] | None = None, property_context: dict | None = None, ) -> tuple[str, list[SearchHit]]: """Build the fullest REFERENCE baseline from ranked hits (same doc + section). When ``tenant_id`` is supplied and section-complete assembly is enabled, the best source is chosen by similarity (as before) but the baseline is then built from EVERY chunk that source holds for the section — in document order — rather than only the top-K chunks similarity surfaced. This is what makes the whole past-report section get mapped instead of a partial fragment. A character budget bounds the assembled length so the mapping prompt stays sane. """ from backend.core.paragraph_merge import combine_reference_blocks if not hits: return "", [] by_source: dict[str, list[SearchHit]] = {} for hit in hits: key = hit.source_filename or hit.doc_id or "unknown" by_source.setdefault(key, []).append(hit) sid = (paragraph_section_id or "").upper() def source_rank(key: str) -> tuple[float, int]: group = by_source[key] max_score = max(h.score for h in group) matched = [ h for h in group if sid and (h.section_id or "").upper() == sid ] or group word_count = sum(len(h.text.split()) for h in matched) return max_score + min(word_count / 400.0, 0.2), word_count best_source = max(by_source, key=source_rank) group = by_source[best_source] # Section-complete expansion: replace the top-K group with the chosen source's # ENTIRE section so no paragraph the surveyor expects is left unmapped. section_complete = False if ( tenant_id and sid and settings.reference_section_complete_enabled ): complete = _complete_section_group( tenant_id, best_source, sid, tier=tier, allowed_doc_keys=allowed_doc_keys, property_context=property_context, ) if complete: group = complete section_complete = True section_hits = [ h for h in group if sid and (h.section_id or "").upper() == sid ] if not section_hits: return "", [] section_hits.sort( key=lambda h: (h.paragraph_index or 0, -len(h.text.split()), -h.score), ) # Only the expanded path is char-budgeted; the bounded top-K path is untouched. budget = ( settings.reference_section_complete_max_chars if section_complete else None ) texts: list[str] = [] contributing: list[SearchHit] = [] seen_para: set[int] = set() used = 0 for hit in section_hits: para_idx = hit.paragraph_index or 0 if para_idx and para_idx in seen_para: continue body = hit.text.strip() if not body: continue if budget is not None and texts and used + len(body) > budget: break if para_idx: seen_para.add(para_idx) if body not in texts: texts.append(body) contributing.append(hit) used += len(body) if not texts: best = max(section_hits, key=lambda h: (h.score, len(h.text.split()))) return best.text.strip(), [best] baseline = combine_reference_blocks(texts[0], texts[1:]) return baseline, contributing def fetch_complete_section_baseline( tenant_id: str, *, paragraph_section_id: str, tier: str | None = None, allowed_doc_keys: frozenset[str] | None = None, property_context: dict | None = None, ) -> tuple[str, list[SearchHit]]: """Metadata-only baseline for a section (no similarity query). Fallback for when similarity search surfaced nothing for a section that nonetheless exists in the index (weak query match, alias drift). Returns the fullest single-source section, or ``("", [])`` when the section is genuinely absent — in which case the caller correctly degrades to a notes-only section. """ if not settings.reference_section_complete_enabled: return "", [] chunks = _complete_section_group( tenant_id, None, paragraph_section_id, tier=tier, allowed_doc_keys=allowed_doc_keys, property_context=property_context, ) if not chunks: return "", [] return assemble_reference_baseline( chunks, paragraph_section_id=paragraph_section_id, tenant_id=tenant_id, tier=tier, allowed_doc_keys=allowed_doc_keys, property_context=property_context, ) def _uses_reference_tier(interference_level: InterferenceLevel) -> bool: return interference_level in ("minimum", "medium", "maximum") def _retrieval_params( retrieval_level: RetrievalLevel, *, paragraph_section_id: str, k: int, ) -> tuple[bool, int]: level = (retrieval_level or "paragraph").lower() if level == "section": return bool(paragraph_section_id), max(k, 8) if level == "document": return bool(paragraph_section_id), max(k, 15) return bool(paragraph_section_id), k def _focus_document_hits(hits: list[SearchHit], k: int) -> list[SearchHit]: """Prefer chunks from the single best-matching past report.""" if not hits: return hits by_source: dict[str, list[SearchHit]] = {} for hit in hits: key = hit.source_filename or hit.doc_id or "unknown" by_source.setdefault(key, []).append(hit) best_key = max(by_source, key=lambda name: max(h.score for h in by_source[name])) focused = sorted(by_source[best_key], key=lambda h: h.score, reverse=True) return focused[:k] def _search_tier( tenant_id: str, query: str, *, paragraph_section_id: str, k: int, interference_level: InterferenceLevel, retrieval_level: RetrievalLevel = "paragraph", allowed_doc_keys: frozenset[str] | None = None, ) -> list[SearchHit]: store = get_rag_store() if _uses_reference_tier(interference_level): search = store.search_for_reference_mapping else: search = store.search_for_generation strict, fetch_k = _retrieval_params( retrieval_level, paragraph_section_id=paragraph_section_id, k=k, ) search_kwargs = { "section_id": paragraph_section_id or None, "allowed_doc_keys": allowed_doc_keys, } # Never widen to cross-section retrieval — empty is safer than wrong-section bleed. use_strict = strict or bool(paragraph_section_id) hits = search( tenant_id, query, top_k=max(fetch_k * 3, fetch_k), section_strict=use_strict, **search_kwargs, ) hits = _filter_hits_to_section(hits, paragraph_section_id) if (retrieval_level or "paragraph").lower() == "document": hits = _focus_document_hits(hits, fetch_k) return hits def retrieve_paragraphs_for_mapping( tenant_id: str, *, section_label: str, paragraph_section_id: str, observations: list[str], interference_level: InterferenceLevel = "minimum", retrieval_level: RetrievalLevel = "paragraph", top_k: int | None = None, allowed_doc_keys: frozenset[str] | None = None, property_context: dict | None = None, ) -> list[SearchHit]: """Ranked paragraphs from uploaded past reports for all interference levels. When ``property_context`` is supplied, paragraphs whose terminology is incompatible with the confirmed property type (e.g. flat-only content in a house report) are dropped before ranking — the next compatible candidate is used rather than leaving a wrong-property baseline. """ k = top_k or settings.reference_baseline_top_k if interference_level == "medium": k = max(k, settings.reference_baseline_top_k) elif interference_level == "maximum": k = max(k, settings.reference_baseline_top_k + 2) query = build_retrieval_query( section_label, observations, section_id=paragraph_section_id, ) hits = _search_tier( tenant_id, query, paragraph_section_id=paragraph_section_id, k=k, interference_level=interference_level, retrieval_level=retrieval_level, allowed_doc_keys=allowed_doc_keys, ) if property_context: from backend.core.property_blocklist import filter_hits_by_property hits = filter_hits_by_property(hits, property_context) hits = rerank_hits_by_observations(hits, observations) hits = rerank_hits_for_comprehensive_baseline(hits) return hits[:k] def find_paragraph_by_topic( tenant_id: str, observations: list[str], *, interference_level: InterferenceLevel = "minimum", top_k: int = 3, allowed_doc_keys: frozenset[str] | None = None, paragraph_section_id: str = "", ) -> list[SearchHit]: """Topic search scoped to ``paragraph_section_id`` when provided.""" parts = [o.strip() for o in observations if o and o.strip()] if not parts: return [] query = " ".join(sorted(parts, key=len, reverse=True)[:5]) store = get_rag_store() sid = (paragraph_section_id or "").strip() or None use_strict = bool(sid) if _uses_reference_tier(interference_level): hits = store.search_for_reference_mapping( tenant_id, query, section_id=sid, top_k=max(top_k * 5, top_k), section_strict=use_strict, allowed_doc_keys=allowed_doc_keys, ) else: hits = store.search_for_generation( tenant_id, query, section_id=sid, top_k=max(top_k * 5, top_k), section_strict=use_strict, ) hits = _filter_hits_to_section(hits, paragraph_section_id) ranked = rerank_hits_by_observations(hits, parts) return rerank_hits_for_comprehensive_baseline(ranked)[:top_k] def report_section_for_paragraph_id( schema: TemplateSchema, paragraph_section_id: str, ) -> str | None: if not paragraph_section_id: return None for report_id, para_id in schema.section_alias_map.items(): if para_id == paragraph_section_id: return report_id if schema.get_section(paragraph_section_id): return paragraph_section_id return None def guess_report_section_from_topic( schema: TemplateSchema, observations: list[str], reference_text: str = "", ) -> str | None: """Best-effort section when REFERENCE chunks lack section_id metadata.""" corpus = " ".join(observations) + " " + reference_text words = _content_tokens(corpus) if not words: return None best_id: str | None = None best_score = 0 for sec in schema.ordered_sections(): title_words = _content_tokens(f"{sec.id} {sec.title}") overlap = len(words & title_words) if overlap > best_score: best_score = overlap best_id = sec.id return best_id if best_score > 0 else None # Backward-compatible alias — always searches MASTER tier. def retrieve_master_paragraphs( tenant_id: str, *, section_label: str, paragraph_section_id: str, observations: list[str], top_k: int | None = None, ) -> list[SearchHit]: k = top_k or settings.retrieval_top_k query = build_retrieval_query( section_label, observations, section_id=paragraph_section_id, ) store = get_rag_store() hits = store.search_for_generation( tenant_id, query, section_id=paragraph_section_id or None, top_k=max(k * 3, k), section_strict=bool(paragraph_section_id), ) hits = _filter_hits_to_section(hits, paragraph_section_id) hits = rerank_hits_by_observations(hits, observations) hits = rerank_hits_for_comprehensive_baseline(hits) return hits[:k]