| 1 | """Best-window extraction for rerankable evidence snippets.""" |
| 2 | |
| 3 | from __future__ import annotations |
| 4 | |
| 5 | from . import relevance, schema |
| 6 | |
| 7 | |
| 8 | def _truncate_words(text: str, max_words: int) -> str: |
| 9 | words = text.split() |
| 10 | if len(words) <= max_words: |
| 11 | return text.strip() |
| 12 | return " ".join(words[:max_words]).strip() + "..." |
| 13 | |
| 14 | |
| 15 | def _windows(words: list[str], size: int, overlap: int) -> list[str]: |
| 16 | if not words: |
| 17 | return [] |
| 18 | if len(words) <= size: |
| 19 | return [" ".join(words)] |
| 20 | step = max(1, size - overlap) |
| 21 | return [ |
| 22 | " ".join(words[start:start + size]) |
| 23 | for start in range(0, len(words), step) |
| 24 | ] |
| 25 | |
| 26 | |
| 27 | def extract_best_snippet( |
| 28 | item: schema.SourceItem, |
| 29 | ranking_query: "str | relevance.PreparedQuery", |
| 30 | max_words: int = 120, |
| 31 | ) -> str: |
| 32 | """Prefer existing snippets, else extract the best matching evidence window.""" |
| 33 | preferred = item.snippet.strip() |
| 34 | if preferred: |
| 35 | return _truncate_words(preferred, max_words) |
| 36 | |
| 37 | body = item.body.strip() |
| 38 | if not body: |
| 39 | return _truncate_words(item.title, max_words) |
| 40 | |
| 41 | words = body.split() |
| 42 | candidates = _windows(words, size=min(max_words, 110), overlap=30) |
| 43 | if not candidates: |
| 44 | return _truncate_words(body, max_words) |
| 45 | |
| 46 | prepared_query = ranking_query if isinstance(ranking_query, relevance.PreparedQuery) else relevance.PreparedQuery(ranking_query) |
| 47 | best = max( |
| 48 | candidates, |
| 49 | key=lambda candidate: relevance.token_overlap_relevance(prepared_query, candidate), |
| 50 | ) |
| 51 | return _truncate_words(best, max_words) |
| 52 |