import os import re import time import asyncio import logging from typing import List, Tuple, Optional import faiss from sentence_transformers import SentenceTransformer from PyPDF2 import PdfReader from docx import Document import pytesseract from PIL import Image import io import openpyxl import pandas as pd from duckduckgo_search import DDGS from duckduckgo_search.exceptions import RatelimitException, DuckDuckGoSearchException from fastapi import UploadFile logger = logging.getLogger(__name__) _EMBED_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2" _embedder: Optional[SentenceTransformer] = None def _get_embedder() -> SentenceTransformer: global _embedder if _embedder is None: logger.info(f"Loading embedding model: {_EMBED_MODEL_NAME}") _embedder = SentenceTransformer(_EMBED_MODEL_NAME) return _embedder def extract_text(file: UploadFile) -> str: ext = os.path.splitext(file.filename)[1].lower() content = file.file.read() file_bytes = io.BytesIO(content) if ext == ".pdf": try: reader = PdfReader(file_bytes) return "\n".join(page.extract_text() or "" for page in reader.pages) except Exception as e: logger.error(f"PDF extract failed: {e}") return "" elif ext == ".docx": try: doc = Document(file_bytes) return "\n".join(p.text for p in doc.paragraphs if p.text) except Exception as e: logger.error(f"DOCX extract failed: {e}") return "" elif ext in [".xlsx", ".xls"]: try: wb = openpyxl.load_workbook(file_bytes, read_only=True, data_only=True) text = [] for sheet in wb: for row in sheet.iter_rows(values_only=True): text.append(" ".join(str(cell) for cell in row if cell is not None)) return "\n".join(text) except Exception as e: logger.error(f"Excel extract failed: {e}") return "" elif ext == ".csv": try: df = pd.read_csv(file_bytes) return df.to_string() except Exception as e: logger.error(f"CSV extract failed: {e}") return "" # NOTE: no image branch here anymore — image files (.jpg/.jpeg/.png/.gif) # are intercepted earlier in build_context_from_files() and routed to # encode_image_to_data_url() for the vision model instead. If this # function receives an image, something upstream isn't filtering it # out correctly — that'd be a bug worth catching, not silently OCR'ing. else: try: return content.decode("utf-8", errors="ignore") except Exception as e: logger.error(f"Text extract failed: {e}") return "" def clean_text(text: str) -> str: t = re.sub(r"[ \t]+", " ", text) t = re.sub(r"\n{3,}", "\n\n", t) return t.strip() def chunk_text(text: str, max_tokens: int = 400, overlap: int = 50) -> List[str]: text = clean_text(text) if not text: return [] words = text.split() chunks, start = [], 0 while start < len(words): end = min(len(words), start + max_tokens) chunk = " ".join(words[start:end]).strip() if chunk: chunks.append(chunk) if end == len(words): break start = max(0, end - overlap) return chunks class RagIndex: def __init__(self, index: faiss.IndexFlatIP, dim: int, chunks: List[str]): self.index = index self.dim = dim self.chunks = chunks def build_faiss_index(chunks: List[str]) -> RagIndex: emb = _get_embedder() vectors = emb.encode(chunks, convert_to_numpy=True, normalize_embeddings=True) dim = vectors.shape[1] index = faiss.IndexFlatIP(dim) index.add(vectors) return RagIndex(index=index, dim=dim, chunks=chunks) def search(index: RagIndex, query: str, top_k: int = 6) -> List[Tuple[str, float]]: emb = _get_embedder() q = emb.encode([query], convert_to_numpy=True, normalize_embeddings=True) D, I = index.index.search(q, top_k) hits = [] for score, idx in zip(D[0], I[0]): if idx == -1: continue hits.append((index.chunks[idx], float(score))) return hits def build_context_from_files(files: List[UploadFile], prompt: str, top_k: int = 6) -> str: all_text = [] for file in files: txt = extract_text(file) if txt: all_text.append(txt) file.file.seek(0) big_text = "\n\n".join(all_text) chunks = chunk_text(big_text, max_tokens=450, overlap=80) if not chunks: return "" idx = build_faiss_index(chunks) hits = search(idx, prompt, top_k=top_k) context_sections = [f"[DOC#{i} score={score:.3f}]\n{chunk}" for i, (chunk, score) in enumerate(hits, 1)] return "\n\n".join(context_sections) def _web_search_sync(query: str, max_results: int = 5) -> str: """ DDGS is an unofficial DuckDuckGo scraper with no SLA — it rate-limits datacenter/cloud IPs (like Hugging Face Spaces) intermittently and without warning. Retrying with backoff smooths over most transient hits; it can't fix a sustained block, which is a hosting-environment limitation, not a code bug. """ last_error = None for attempt in range(3): try: with DDGS() as ddgs: results = list(ddgs.text(query, max_results=max_results, backend="lite")) if not results: return "No results found." sections = [ f"[WEB#{i}] Title: {r['title']}\nSnippet: {r['body']}\nURL: {r['href']}" for i, r in enumerate(results, 1) ] return "\n\n".join(sections) except RatelimitException as e: last_error = e logger.warning(f"Web search rate-limited, attempt {attempt + 1}/3: {e}") time.sleep(15 * (attempt + 1)) except DuckDuckGoSearchException as e: last_error = e logger.error(f"Web search failed: {e}") break except Exception as e: last_error = e logger.error(f"Web search unexpected error: {e}") break logger.error(f"Web search exhausted retries: {last_error}") return f"[Web search unavailable: {type(last_error).__name__ if last_error else 'unknown error'}]" async def web_search(query: str) -> str: # DDGS is synchronous/blocking — offload it so it can't freeze # the event loop (and therefore every other user's request) while waiting. return await asyncio.to_thread(_web_search_sync, query)