Spaces:
Running
Running
| from __future__ import annotations | |
| import re | |
| _CONTINUED_PREFIX_RE = re.compile(r"^\[continued[^\]]*\]\n") | |
| def chunk_snippet_from_doc_text(doc_text: str, *, max_chars: int = 160) -> str: | |
| """Recover the original PDF-page text from a stored RAG document's text. | |
| build_pdf_documents() (src/pdf_manual.py) prepends a synthetic header | |
| ("OEM Owner's Manual: ...", "Pages: X-Y", "---") that isn't actually in the | |
| PDF, and prefixes continuation chunks with "[continued p.X-Y]" -- strip | |
| both so what's left is text that actually appears on the rendered page, | |
| suitable as a page.search_for() query. | |
| """ | |
| text = doc_text or "" | |
| marker = "---\n" | |
| idx = text.find(marker) | |
| body = text[idx + len(marker):] if idx >= 0 else text | |
| body = _CONTINUED_PREFIX_RE.sub("", body) | |
| return body.strip()[:max_chars] | |
| def _highlight_rects(pg, snippet: str, *, window: int = 5, max_words: int = 90) -> list: | |
| """Rectangles covering ``snippet`` on page ``pg``. | |
| The passage text is stored with the PDF's hard line wraps removed, while | |
| page.search_for() matches within the page's own line structure -- so a long | |
| verbatim probe often misses. Instead search short word windows (5 words, | |
| sliding by 4) and union the hits, which survives line wraps and small | |
| extraction differences and covers the *whole* passage, not just its start. | |
| """ | |
| words = snippet.split()[:max_words] | |
| if len(words) < window: | |
| words_probe = [" ".join(words)] if words else [] | |
| else: | |
| words_probe = [" ".join(words[i : i + window]) for i in range(0, len(words) - window + 1, window - 1)] | |
| rects: list = [] | |
| seen: set[tuple] = set() | |
| for probe in words_probe: | |
| if len(probe) < 8: | |
| continue | |
| for r in pg.search_for(probe): | |
| key = (round(r.x0), round(r.y0), round(r.x1), round(r.y1)) | |
| if key not in seen: | |
| seen.add(key) | |
| rects.append(r) | |
| return rects | |
| def render_pdf_page_png( | |
| pdf_path: str, | |
| *, | |
| page: int, | |
| highlight_snippet: str | None = None, | |
| zoom: float = 2.0, | |
| ) -> bytes: | |
| """Render one PDF page to PNG, best-effort highlighting the given snippet. | |
| page.search_for() is an exact-substring match against extracted text, so a | |
| minor whitespace/ligature difference between the stored chunk text and | |
| fitz's own extraction can make the full snippet miss; shrink it until | |
| something resolves rather than requiring a verbatim match. No match found | |
| at all just means a plain, unhighlighted page render -- not an error. | |
| """ | |
| import fitz # type: ignore | |
| doc = fitz.open(pdf_path) | |
| try: | |
| page_idx = max(0, min(page - 1, doc.page_count - 1)) | |
| pg = doc.load_page(page_idx) | |
| if highlight_snippet: | |
| rects = _highlight_rects(pg, highlight_snippet) | |
| for r in rects: | |
| pg.draw_rect(r, color=(0.85, 0.55, 0.1), fill=(1, 0.94, 0.75), fill_opacity=0.4, width=1.4) | |
| pix = pg.get_pixmap(matrix=fitz.Matrix(zoom, zoom)) | |
| return pix.tobytes("png") | |
| finally: | |
| doc.close() | |