Spaces:
Sleeping
Sleeping
| from __future__ import annotations | |
| import os | |
| import re | |
| MAX_CHAT_PAPERS = int(os.getenv("MAX_CHAT_PAPERS", "4")) | |
| CHAT_SUMMARY_MAX_CHARS = int(os.getenv("CHAT_SUMMARY_MAX_CHARS", "14000")) | |
| _SECTION_HEADING = re.compile(r"^##\s+\d+\.", re.MULTILINE) | |
| _QUESTION_STOPWORDS = frozenset( | |
| { | |
| "about", | |
| "also", | |
| "based", | |
| "compare", | |
| "could", | |
| "does", | |
| "first", | |
| "from", | |
| "have", | |
| "paper", | |
| "papers", | |
| "please", | |
| "provide", | |
| "range", | |
| "second", | |
| "source", | |
| "summarize", | |
| "summary", | |
| "that", | |
| "the", | |
| "their", | |
| "these", | |
| "third", | |
| "this", | |
| "those", | |
| "was", | |
| "were", | |
| "what", | |
| "when", | |
| "which", | |
| "would", | |
| "your", | |
| } | |
| ) | |
| _ORDINAL_PAPER = re.compile( | |
| r"\b(?P<ordinal>first|1st|second|2nd|third|3rd|fourth|4th)\s+paper\b", | |
| re.I, | |
| ) | |
| _ORDINAL_INDEX = { | |
| "first": 0, | |
| "1st": 0, | |
| "second": 1, | |
| "2nd": 1, | |
| "third": 2, | |
| "3rd": 2, | |
| "fourth": 3, | |
| "4th": 3, | |
| } | |
| def focus_document_ids_for_question(message: str, document_ids: list[str]) -> list[str]: | |
| """When the user names a specific paper, answer from that paper only.""" | |
| if len(document_ids) <= 1: | |
| return document_ids | |
| match = _ORDINAL_PAPER.search(message or "") | |
| if not match: | |
| return document_ids | |
| index = _ORDINAL_INDEX.get(match.group("ordinal").lower()) | |
| if index is None or index >= len(document_ids): | |
| return document_ids | |
| return [document_ids[index]] | |
| def enhance_ordinal_question(message: str, sources: list[dict]) -> str: | |
| """When the user says 'second paper', inject that paper's title into the prompt.""" | |
| if not _ORDINAL_PAPER.search(message or "") or len(sources) != 1: | |
| return message | |
| title = (sources[0].get("title") or "").strip() | |
| if not title: | |
| return message | |
| return f"{message}\n\nAnswer specifically about the paper titled: {title}" | |
| def _question_terms(question: str) -> list[str]: | |
| terms: list[str] = [] | |
| lowered = (question or "").lower() | |
| if "ph" in lowered: | |
| terms.append("ph") | |
| for token in re.findall(r"[a-z0-9]+", lowered): | |
| if token in _QUESTION_STOPWORDS or len(token) < 3: | |
| continue | |
| terms.append(token) | |
| return list(dict.fromkeys(terms)) | |
| def _extract_section(text: str, section_number: int, section_label: str) -> str: | |
| pattern = re.compile( | |
| rf"^##\s*{section_number}\.\s*{re.escape(section_label)}.*?(?=^##\s+\d+\.|\Z)", | |
| re.MULTILINE | re.DOTALL | re.IGNORECASE, | |
| ) | |
| match = pattern.search(text) | |
| return match.group(0).strip() if match else "" | |
| def _matching_lines(text: str, terms: list[str]) -> str: | |
| if not terms: | |
| return "" | |
| lines: list[str] = [] | |
| for line in text.splitlines(): | |
| lowered = line.lower() | |
| if any(term in lowered for term in terms): | |
| lines.append(line.rstrip()) | |
| return "\n".join(lines).strip() | |
| def _fit_blocks(blocks: list[str], budget: int) -> str: | |
| selected: list[str] = [] | |
| used = 0 | |
| for block in blocks: | |
| block = block.strip() | |
| if not block: | |
| continue | |
| extra = len(block) + (2 if selected else 0) | |
| if used + extra > budget: | |
| remaining = budget - used - (2 if selected else 0) | |
| if remaining > 120: | |
| selected.append(block[:remaining].rstrip() + "…") | |
| break | |
| selected.append(block) | |
| used += extra | |
| return "\n\n".join(selected) | |
| def excerpt_summary(text: str, *, question: str = "", n_papers: int = 1) -> str: | |
| """Prefer identity + key results (+ question matches) instead of blind head truncation.""" | |
| per_paper = max(2500, CHAT_SUMMARY_MAX_CHARS // max(n_papers, 1)) | |
| if len(text) <= per_paper: | |
| return text | |
| blocks: list[str] = [] | |
| identity = _extract_section(text, 1, "Paper Identity") | |
| key_results = _extract_section(text, 5, "Key Results") | |
| takeaway = _extract_section(text, 10, "Concise Takeaway") | |
| terms = _question_terms(question) | |
| matched = _matching_lines(text, terms) | |
| if identity: | |
| blocks.append(identity) | |
| if key_results: | |
| blocks.append(key_results) | |
| elif matched: | |
| blocks.append(matched) | |
| if takeaway and takeaway not in blocks: | |
| blocks.append(takeaway) | |
| if matched and matched not in blocks: | |
| blocks.append(matched) | |
| excerpt = _fit_blocks(blocks, per_paper) | |
| if excerpt: | |
| return excerpt | |
| return text[:per_paper] + "\n\n[Additional sections omitted for length.]" | |
| def truncate_summary(text: str, *, n_papers: int = 1, question: str = "") -> str: | |
| return excerpt_summary(text, question=question, n_papers=n_papers) | |