# utils.py import os import pickle import time from typing import List, Tuple, Dict, Optional, Set from collections import deque from urllib.parse import urlparse, urljoin import requests from bs4 import BeautifulSoup import trafilatura import fitz # PyMuPDF from pdf2image import convert_from_path import pytesseract import numpy as np import faiss from sentence_transformers import SentenceTransformer # ---------- Configuration ---------- EMBED_MODEL_NAME = "all-MiniLM-L6-v2" PDF_OCR_DPI = 200 CRAWL_TIMEOUT = 8 DEFAULT_MAX_PAGES = 150 DEFAULT_MAX_DEPTH = 2 CHUNK_SIZE_WORDS = 220 CHUNK_OVERLAP = 40 # Keep a module-level model holder to avoid reloading repeatedly _embed_model = None def load_embed_model(): global _embed_model if _embed_model is None: _embed_model = SentenceTransformer(EMBED_MODEL_NAME) return _embed_model # ---------- PDF extraction (selectable + OCR fallback) ---------- def extract_text_from_pdf(pdf_path: str) -> str: """ Extract text from PDF using PyMuPDF; fallback to OCR via pdf2image + pytesseract for pages with little text. """ if not os.path.exists(pdf_path): return "" texts = [] try: doc = fitz.open(pdf_path) except Exception: return "" for p in range(len(doc)): try: page = doc[p] page_text = page.get_text("text") except Exception: page_text = "" if page_text and len(page_text.strip()) > 60: texts.append(page_text) else: # OCR fallback for scanned page or images try: images = convert_from_path(pdf_path, first_page=p+1, last_page=p+1, dpi=PDF_OCR_DPI) ocr_text = "" for img in images: ocr_text += pytesseract.image_to_string(img) texts.append(ocr_text) except Exception: # if OCR fails, append whatever page_text had if page_text: texts.append(page_text) return "\n".join(texts) # ---------- Robust HTML extraction ---------- def extract_text_from_html(raw_html: str) -> str: """ Use trafilatura for robust extraction; fallback to BeautifulSoup text. """ try: res = trafilatura.extract(raw_html, include_comments=False, include_tables=False) if res: return res except Exception: pass soup = BeautifulSoup(raw_html, "html.parser") for s in soup(["script", "style", "noscript"]): s.decompose() return soup.get_text(separator="\n", strip=True) def is_internal(base_domain: str, url: str) -> bool: try: parsed = urlparse(url) if not parsed.netloc: return True return base_domain in parsed.netloc except Exception: return False def normalize_url(base: str, href: str) -> Optional[str]: try: return urljoin(base, href) except Exception: return None # ---------- Dynamic crawler (breadth-first) ---------- def crawl_site_dynamic(start_url: str, max_pages: int = DEFAULT_MAX_PAGES, max_depth: int = DEFAULT_MAX_DEPTH, timeout: int = CRAWL_TIMEOUT) -> Tuple[str, List[str]]: """ Crawl internal pages breadth-first starting from start_url. Returns concatenated extracted text and list of visited URLs (for reference). """ parsed = urlparse(start_url) base_domain = parsed.netloc visited: Set[str] = set() q = deque() q.append((start_url, 0)) headers = {"User-Agent": "Mozilla/5.0 (compatible; DUET-Indexer/1.0)"} extracted_texts = [] visited_urls = [] while q and len(visited) < max_pages: url, depth = q.popleft() if url in visited: continue if depth > max_depth: continue try: r = requests.get(url, timeout=timeout, headers=headers) if r.status_code != 200: visited.add(url) continue raw_html = r.text text = extract_text_from_html(raw_html) if text and len(text.strip()) > 30: extracted_texts.append(text) visited_urls.append(url) visited.add(url) if depth < max_depth: soup = BeautifulSoup(raw_html, "html.parser") for a in soup.find_all("a", href=True): href = a["href"] full = normalize_url(url, href) if not full: continue if "#" in full: full = full.split("#")[0] if full in visited: continue if is_internal(base_domain, full): q.append((full, depth + 1)) except Exception: visited.add(url) continue return "\n\n".join(extracted_texts), visited_urls # ---------- Chunking ---------- def chunk_text(text: str, chunk_size_words: int = CHUNK_SIZE_WORDS, overlap: int = CHUNK_OVERLAP) -> List[str]: words = text.split() if not words: return [] chunks = [] i = 0 while i < len(words): chunk_words = words[i:i + chunk_size_words] chunks.append(" ".join(chunk_words)) i += (chunk_size_words - overlap) return chunks # ---------- Embeddings & Index helpers ---------- def compute_embeddings(chunks: List[str]) -> np.ndarray: """ Compute normalized embeddings for chunks (shape N x dim). """ model = load_embed_model() vectors = model.encode(chunks, convert_to_numpy=True, show_progress_bar=False) norms = np.linalg.norm(vectors, axis=1, keepdims=True) norms[norms == 0] = 1.0 vectors = vectors / norms return vectors.astype("float32") def build_faiss_index(vectors: np.ndarray) -> faiss.IndexFlatIP: dim = vectors.shape[1] idx = faiss.IndexFlatIP(dim) idx.add(vectors) return idx def retrieve_top_k(query: str, chunks: List[str], vectors: np.ndarray, top_k: int = 5) -> Tuple[List[str], List[float], List[int]]: if vectors is None or len(vectors) == 0: return [], [], [] model = load_embed_model() qv = model.encode([query], convert_to_numpy=True) qv = qv / (np.linalg.norm(qv, axis=1, keepdims=True) + 1e-10) idx = build_faiss_index(vectors) D, I = idx.search(qv.astype("float32"), top_k) indices = I[0].tolist() scores = D[0].tolist() retrieved = [chunks[i] for i in indices if i < len(chunks)] return retrieved, scores, indices # ---------- Save / Load embeddings and metadata ---------- def save_embeddings(path: str, chunks: List[str], vectors: np.ndarray, sources: Optional[List[str]] = None): payload = { "chunks": chunks, "vectors": vectors, "sources": sources or [] } with open(path, "wb") as f: pickle.dump(payload, f) def load_embeddings(path: str): with open(path, "rb") as f: payload = pickle.load(f) return payload.get("chunks", []), payload.get("vectors", None), payload.get("sources", [])