Update app.py

app.py

@@ -1,20 +1,16 @@
 #!/usr/bin/env python3
 """
-Indian Legal AI Assistant with integrated India Code web lookup.
+Indian Legal AI Assistant with integrated India Code lookup.
 
 CPU-only Hugging Face Spaces app.
 
-Model:
-invincibleambuj/Ambuj-Tripathi-Indian-Legal-Llama-GGUF
-llama-3.2-1b-instruct.Q4_K_M.gguf
-
-Web step:
-- Uses India Code only: https://www.indiacode.nic.in/
-- Searches logically across India Code navigation, Central Acts browse pages,
-  State Acts links, Repealed Acts, Spent Acts, and discovered India Code pages.
-- Reads HTML pages.
-- Reads extractable text from PDFs using pypdf.
-- Injects retrieved India Code context into the model prompt.
+Features:
+- Uses only India Code as the legal web source.
+- Reads India Code HTML pages.
+- Reads text-based PDFs using pypdf.
+- Falls back to lightweight OCR for scanned/image PDFs using PyMuPDF + Tesseract.
+- Automatically trims prompt/context to avoid context-window overflow.
+- Supports follow-up questions using recent chat history.
 """
 
 import os
@@ -24,11 +20,14 @@ import html
 import time
 import traceback
 from functools import lru_cache
-from urllib.parse import urljoin, urlparse, quote_plus, urldefrag
+from urllib.parse import urljoin, urlparse, quote_plus, urldefrag, parse_qs
 
 import requests
 from bs4 import BeautifulSoup
+from PIL import Image
 from pypdf import PdfReader
+import fitz  # PyMuPDF
+import pytesseract
 
 import gradio as gr
 from huggingface_hub import hf_hub_download
@@ -42,28 +41,26 @@ from llama_cpp import Llama
 MODEL_REPO = "invincibleambuj/Ambuj-Tripathi-Indian-Legal-Llama-GGUF"
 MODEL_FILE = "llama-3.2-1b-instruct.Q4_K_M.gguf"
 
-# CPU-only defaults for Hugging Face CPU Spaces.
-N_CTX = int(os.getenv("N_CTX", "1024"))
+# CPU-only settings.
+# 2048 is safer than 1024 because India Code context + user question can exceed 1024.
+N_CTX = int(os.getenv("N_CTX", "2048"))
 N_THREADS = int(os.getenv("N_THREADS", "2"))
 N_THREADS_BATCH = int(os.getenv("N_THREADS_BATCH", "2"))
 N_BATCH = int(os.getenv("N_BATCH", "512"))
-
-# CPU only.
 N_GPU_LAYERS = 0
 
 # Generation settings.
-MAX_TOKENS = int(os.getenv("MAX_TOKENS", "768"))
+MAX_TOKENS = int(os.getenv("MAX_TOKENS", "512"))
 TEMPERATURE = float(os.getenv("TEMPERATURE", "0.35"))
 TOP_P = float(os.getenv("TOP_P", "0.9"))
 
-# Model cache.
 MODEL_CACHE_DIR = os.getenv("MODEL_CACHE_DIR", "./models")
 
 llm = None
 
 
 # -------------------------------------------------
-# India Code search configuration
+# India Code / extraction configuration
 # -------------------------------------------------
 
 INDIACODE_HOME = "https://www.indiacode.nic.in/"
@@ -74,12 +71,21 @@ ALLOWED_DOMAINS = {
 
 REQUEST_TIMEOUT = int(os.getenv("REQUEST_TIMEOUT", "12"))
 
-# Keep these conservative for CPU Spaces.
 MAX_DISCOVERY_RESULTS = int(os.getenv("MAX_DISCOVERY_RESULTS", "12"))
-MAX_CRAWL_PAGES = int(os.getenv("MAX_CRAWL_PAGES", "28"))
-MAX_CONTEXT_DOCS = int(os.getenv("MAX_CONTEXT_DOCS", "6"))
-MAX_PDF_PAGES = int(os.getenv("MAX_PDF_PAGES", "8"))
-MAX_TEXT_PER_DOC = int(os.getenv("MAX_TEXT_PER_DOC", "3500"))
+MAX_CRAWL_PAGES = int(os.getenv("MAX_CRAWL_PAGES", "24"))
+MAX_CONTEXT_DOCS = int(os.getenv("MAX_CONTEXT_DOCS", "5"))
+
+# PDF text extraction.
+MAX_PDF_TEXT_PAGES = int(os.getenv("MAX_PDF_TEXT_PAGES", "8"))
+
+# OCR fallback. Keep this small for CPU Spaces.
+MAX_OCR_PAGES = int(os.getenv("MAX_OCR_PAGES", "3"))
+OCR_DPI_SCALE = float(os.getenv("OCR_DPI_SCALE", "1.5"))
+MIN_PDF_TEXT_CHARS_BEFORE_OCR = int(os.getenv("MIN_PDF_TEXT_CHARS_BEFORE_OCR", "250"))
+
+# Context trimming.
+MAX_TEXT_PER_DOC = int(os.getenv("MAX_TEXT_PER_DOC", "2400"))
+PROMPT_SAFETY_MARGIN = int(os.getenv("PROMPT_SAFETY_MARGIN", "96"))
 
 HEADERS = {
     "User-Agent": (
@@ -90,41 +96,28 @@ HEADERS = {
 
 
 # -------------------------------------------------
-# General helpers
+# Basic helpers
 # -------------------------------------------------
 
 def clean_text(text):
-    """Normalize whitespace and decode HTML entities."""
     if not text:
         return ""
-
     text = html.unescape(text)
     text = re.sub(r"\s+", " ", text)
     return text.strip()
 
 
 def normalize_url(url, base=INDIACODE_HOME):
-    """Resolve, defragment, and normalize a URL."""
     if not url:
         return ""
-
     url = urljoin(base, url)
     url, _fragment = urldefrag(url)
     return url.strip()
 
 
-def get_domain(url):
-    try:
-        return urlparse(url).netloc.lower().replace("www.", "")
-    except Exception:
-        return ""
-
-
 def is_indiacode_url(url):
-    """Allow only India Code URLs."""
     try:
-        parsed = urlparse(url)
-        domain = parsed.netloc.lower()
+        domain = urlparse(url).netloc.lower()
         return domain in ALLOWED_DOMAINS
     except Exception:
         return False
@@ -135,19 +128,17 @@ def looks_like_pdf_url(url):
 
 
 def query_terms(query):
-    """Extract useful query terms for scoring."""
     stopwords = {
         "the", "a", "an", "and", "or", "of", "in", "on", "to", "for", "with",
-        "under", "section", "sections", "act", "law", "laws", "what", "is",
-        "are", "explain", "about", "current", "latest", "india", "indian",
+        "under", "section", "sections", "act", "acts", "law", "laws", "what",
+        "is", "are", "explain", "about", "current", "latest", "india", "indian",
+        "tell", "me", "please", "does", "do",
     }
-
     terms = re.findall(r"[a-zA-Z0-9]+", query.lower())
     return [t for t in terms if len(t) >= 3 and t not in stopwords]
 
 
 def score_text_against_query(text, query):
-    """Simple lexical scoring for relevance."""
     text_l = (text or "").lower()
     terms = query_terms(query)
 
@@ -155,13 +146,11 @@ def score_text_against_query(text, query):
         return 0
 
     score = 0
-
     for term in terms:
         count = text_l.count(term)
         if count:
             score += min(count, 5)
 
-    # Boost exact phrase match.
     q = clean_text(query).lower()
     if q and q in text_l:
         score += 10
@@ -169,14 +158,13 @@ def score_text_against_query(text, query):
     return score
 
 
-def make_snippet(text, query, max_chars=900):
-    """Create a short context snippet near query terms."""
+def make_snippet(text, query, max_chars=MAX_TEXT_PER_DOC):
     text = clean_text(text)
     if not text:
         return ""
 
-    terms = query_terms(query)
     lower = text.lower()
+    terms = query_terms(query)
 
     first_hit = None
     for term in terms:
@@ -188,19 +176,14 @@ def make_snippet(text, query, max_chars=900):
     if first_hit is None:
         return text[:max_chars]
 
-    start = max(first_hit - 250, 0)
+    start = max(first_hit - 350, 0)
     end = min(start + max_chars, len(text))
     return text[start:end]
 
 
-# -------------------------------------------------
-# HTTP helpers
-# -------------------------------------------------
-
-def safe_get(url, timeout=REQUEST_TIMEOUT):
-    """GET request with basic error handling."""
+def safe_get(url):
     try:
-        response = requests.get(url, headers=HEADERS, timeout=timeout)
+        response = requests.get(url, headers=HEADERS, timeout=REQUEST_TIMEOUT)
         response.raise_for_status()
         return response
     except Exception as error:
@@ -208,18 +191,11 @@ def safe_get(url, timeout=REQUEST_TIMEOUT):
         return None
 
 
-def content_type(response):
-    if response is None:
-        return ""
-    return response.headers.get("content-type", "").lower()
-
-
 # -------------------------------------------------
-# HTML and PDF extraction
+# HTML extraction
 # -------------------------------------------------
 
 def extract_links_from_html(html_text, base_url):
-    """Extract India Code links from HTML."""
     links = []
 
     try:
@@ -231,12 +207,8 @@ def extract_links_from_html(html_text, base_url):
             url = normalize_url(href, base_url)
 
             if is_indiacode_url(url):
-                links.append(
-                    {
-                        "url": url,
-                        "anchor": text,
-                    }
-                )
+                links.append({"url": url, "anchor": text})
+
     except Exception as error:
         print(f"Link extraction failed for {base_url}: {error}")
 
@@ -244,20 +216,18 @@ def extract_links_from_html(html_text, base_url):
 
 
 def extract_text_from_html(html_text):
-    """Extract readable text from HTML."""
     try:
         soup = BeautifulSoup(html_text, "html.parser")
 
         for tag in soup(["script", "style", "nav", "footer", "header", "aside", "form"]):
             tag.decompose()
 
-        title = clean_text(soup.title.get_text(" ")) if soup.title else ""
-
         parts = []
-        if title:
-            parts.append(title)
 
-        for tag in soup.find_all(["h1", "h2", "h3", "h4", "p", "li", "td", "th"], limit=300):
+        if soup.title:
+            parts.append(clean_text(soup.title.get_text(" ")))
+
+        for tag in soup.find_all(["h1", "h2", "h3", "h4", "p", "li", "td", "th"], limit=350):
             text = clean_text(tag.get_text(" "))
             if len(text) >= 20:
                 parts.append(text)
@@ -269,46 +239,118 @@ def extract_text_from_html(html_text):
         return ""
 
 
-def extract_pdf_text(pdf_bytes, max_pages=MAX_PDF_PAGES):
-    """
-    Extract text from a PDF.
+# -------------------------------------------------
+# PDF extraction with OCR fallback
+# -------------------------------------------------
 
-    This reads text-based PDFs. It will not OCR scanned image-only PDFs.
-    """
+def extract_pdf_text_with_pypdf(pdf_bytes, max_pages=MAX_PDF_TEXT_PAGES):
     try:
         reader = PdfReader(io.BytesIO(pdf_bytes))
         parts = []
 
-        total_pages = len(reader.pages)
-        pages_to_read = min(total_pages, max_pages)
+        pages_to_read = min(len(reader.pages), max_pages)
 
         for page_index in range(pages_to_read):
             try:
                 text = reader.pages[page_index].extract_text() or ""
                 text = clean_text(text)
                 if text:
-                    parts.append(f"[PDF page {page_index + 1}] {text}")
+                    parts.append(f"[PDF text page {page_index + 1}] {text}")
             except Exception as page_error:
-                print(f"PDF page extraction failed: {page_error}")
+                print(f"pypdf page extraction failed: {page_error}")
+
+        return clean_text(" ".join(parts))
+
+    except Exception as error:
+        print(f"pypdf extraction failed: {error}")
+        return ""
+
+
+def ocr_pdf_with_tesseract(pdf_bytes, max_pages=MAX_OCR_PAGES):
+    """
+    Lightweight OCR fallback for scanned PDFs.
+
+    Uses:
+    - PyMuPDF to render PDF pages to images.
+    - Tesseract via pytesseract to OCR those images.
+
+    Kept intentionally small because Hugging Face CPU Spaces have limited CPU.
+    """
+    try:
+        doc = fitz.open(stream=pdf_bytes, filetype="pdf")
+        parts = []
+
+        pages_to_read = min(len(doc), max_pages)
+
+        for page_index in range(pages_to_read):
+            try:
+                page = doc.load_page(page_index)
+
+                matrix = fitz.Matrix(OCR_DPI_SCALE, OCR_DPI_SCALE)
+                pix = page.get_pixmap(matrix=matrix, alpha=False)
+
+                img = Image.frombytes(
+                    "RGB",
+                    [pix.width, pix.height],
+                    pix.samples,
+                )
 
-        extracted = clean_text(" ".join(parts))
+                # Convert to grayscale to reduce OCR work.
+                img = img.convert("L")
 
-        if not extracted:
-            return "[PDF detected, but no extractable text was found. The PDF may be scanned/image-based.]"
+                text = pytesseract.image_to_string(img, lang="eng")
+                text = clean_text(text)
+
+                if text:
+                    parts.append(f"[OCR PDF page {page_index + 1}] {text}")
+
+            except Exception as page_error:
+                print(f"OCR failed on page {page_index + 1}: {page_error}")
 
-        return extracted
+        doc.close()
+
+        return clean_text(" ".join(parts))
 
     except Exception as error:
-        print(f"PDF extraction failed: {error}")
+        print(f"OCR PDF extraction failed: {error}")
         return ""
 
 
-@lru_cache(maxsize=256)
-def fetch_document_text(url):
+def extract_pdf_text(pdf_bytes):
     """
-    Fetch and extract text from an India Code HTML or PDF document.
-    Cached to speed up follow-up questions.
+    First try normal PDF text extraction.
+    If too little text is found, fall back to OCR.
     """
+    text = extract_pdf_text_with_pypdf(pdf_bytes)
+
+    if len(text) >= MIN_PDF_TEXT_CHARS_BEFORE_OCR:
+        return text, "pdf-text"
+
+    print("PDF appears scanned or has too little extractable text. Running OCR fallback...")
+
+    ocr_text = ocr_pdf_with_tesseract(pdf_bytes)
+
+    if ocr_text:
+        if text:
+            return text + "\n\n" + ocr_text, "pdf-text-plus-ocr"
+        return ocr_text, "pdf-ocr"
+
+    if text:
+        return text, "pdf-text-low"
+
+    return (
+        "[PDF detected, but no readable text could be extracted. "
+        "The PDF may be scanned, low quality, encrypted, or OCR failed.]",
+        "pdf-unreadable",
+    )
+
+
+# -------------------------------------------------
+# Document fetching
+# -------------------------------------------------
+
+@lru_cache(maxsize=256)
+def fetch_document_text(url):
     if not is_indiacode_url(url):
         return {
             "url": url,
@@ -319,6 +361,7 @@ def fetch_document_text(url):
         }
 
     response = safe_get(url)
+
     if response is None:
         return {
             "url": url,
@@ -328,16 +371,16 @@ def fetch_document_text(url):
             "type": "failed",
         }
 
-    ctype = content_type(response)
+    ctype = response.headers.get("content-type", "").lower()
 
     if "application/pdf" in ctype or looks_like_pdf_url(url):
-        text = extract_pdf_text(response.content)
+        text, pdf_type = extract_pdf_text(response.content)
         return {
             "url": url,
             "title": url.split("/")[-1] or "India Code PDF",
             "text": text,
             "links": [],
-            "type": "pdf",
+            "type": pdf_type,
         }
 
     html_text = response.text
@@ -362,19 +405,10 @@ def fetch_document_text(url):
 
 
 # -------------------------------------------------
-# India Code discovery
+# India Code discovery / crawling
 # -------------------------------------------------
 
 def india_code_seed_urls():
-    """
-    Core India Code entry points.
-
-    These are logical browse points exposed by India Code:
-    - Home
-    - Central Acts browse pages
-    - Repealed Acts
-    - Spent Acts
-    """
     return [
         INDIACODE_HOME,
 
@@ -386,19 +420,17 @@ def india_code_seed_urls():
         "https://www.indiacode.nic.in/handle/123456789/1362/browse?type=ministry",
         "https://www.indiacode.nic.in/handle/123456789/1362/browse?type=department",
 
-        # Repealed and spent Acts.
+        # Repealed / spent Acts.
         "https://www.indiacode.nic.in/repealed-act/repealed-act.jsp",
         "https://www.indiacode.nic.in/spent-act/spent-act.jsp",
     ]
 
 
 def discover_from_indiacode_home():
-    """
-    Discover State Act and other India Code navigation links from homepage.
-    """
     discovered = []
 
     home_doc = fetch_document_text(INDIACODE_HOME)
+
     for link in home_doc.get("links", []):
         url = link.get("url", "")
         if is_indiacode_url(url):
@@ -409,18 +441,16 @@ def discover_from_indiacode_home():
 
 def duckduckgo_site_discovery(query):
     """
-    URL discovery using a site-restricted query.
+    Site-restricted discovery only.
 
-    Important:
-    This is not used as a content source. It only discovers India Code URLs.
-    The app fetches and reads the resulting India Code pages directly.
+    Content source remains India Code only. This only helps discover India Code URLs.
     """
     discovered = []
-
     search_query = f"site:indiacode.nic.in {query}"
     search_url = f"https://duckduckgo.com/html/?q={quote_plus(search_query)}"
 
     response = safe_get(search_url)
+
     if response is None:
         return discovered
 
@@ -430,10 +460,8 @@ def duckduckgo_site_discovery(query):
         for a in soup.select(".result__title a"):
             url = a.get("href", "").strip()
 
-            # DuckDuckGo may wrap links. Keep only direct India Code URLs.
             if "uddg=" in url:
                 try:
-                    from urllib.parse import parse_qs
                     parsed = urlparse(url)
                     qs = parse_qs(parsed.query)
                     if "uddg" in qs:
@@ -450,21 +478,12 @@ def duckduckgo_site_discovery(query):
                 break
 
     except Exception as error:
-        print(f"Site discovery parsing failed: {error}")
+        print(f"Site discovery failed: {error}")
 
     return discovered
 
 
 def relevant_link_filter(link, query):
-    """
-    Decide whether to crawl a link.
-
-    We keep:
-    - links with query terms in anchor or URL
-    - PDF links
-    - handle/item/bitstream links, because India Code DSpace pages often use them
-    - browse pages
-    """
     url = link.get("url", "")
     anchor = link.get("anchor", "")
 
@@ -483,12 +502,15 @@ def relevant_link_filter(link, query):
         "/browse",
         "repealed-act",
         "spent-act",
+        "download",
+        "pdf",
     ]
 
     if any(pattern in url_l for pattern in important_patterns):
         return True
 
     terms = query_terms(query)
+
     if any(term in url_l or term in anchor_l for term in terms):
         return True
 
@@ -496,22 +518,11 @@ def relevant_link_filter(link, query):
 
 
 def crawl_indiacode_for_query(query):
-    """
-    Crawl India Code pages to find relevant documents.
-
-    Strategy:
-    1. Start with logical India Code seed URLs.
-    2. Add homepage-discovered India Code links, including State Acts links.
-    3. Add site-restricted discovered India Code URLs.
-    4. Fetch pages, score text, follow relevant India Code links.
-    5. Include readable PDFs.
-    """
     seeds = []
     seeds.extend(india_code_seed_urls())
     seeds.extend(discover_from_indiacode_home())
     seeds.extend(duckduckgo_site_discovery(query))
 
-    # Preserve order while removing duplicates.
     queue = []
     seen = set()
 
@@ -537,10 +548,10 @@ def crawl_indiacode_for_query(query):
         title = doc.get("title", "India Code document")
         doc_type = doc.get("type", "html")
 
-        combined_for_score = f"{title} {url} {text}"
-        score = score_text_against_query(combined_for_score, query)
+        combined = f"{title} {url} {text}"
+        score = score_text_against_query(combined, query)
 
-        if score > 0 or doc_type == "pdf":
+        if score > 0 or doc_type.startswith("pdf"):
             scored_docs.append(
                 {
                     "url": url,
@@ -551,7 +562,6 @@ def crawl_indiacode_for_query(query):
                 }
             )
 
-        # Follow relevant India Code links from HTML pages.
         for link in doc.get("links", []):
             link_url = normalize_url(link.get("url", ""), url)
 
@@ -564,34 +574,36 @@ def crawl_indiacode_for_query(query):
 
     scored_docs.sort(key=lambda item: item["score"], reverse=True)
 
-    # Keep top docs with real text.
     useful_docs = []
+
     for doc in scored_docs:
         if doc.get("text"):
             useful_docs.append(doc)
+
         if len(useful_docs) >= MAX_CONTEXT_DOCS:
             break
 
     return useful_docs
 
 
-def build_indiacode_context(query):
-    """
-    Build compact India Code context for the model.
-    """
+def build_indiacode_context(query, max_text_per_doc=MAX_TEXT_PER_DOC, max_docs=MAX_CONTEXT_DOCS):
     docs = crawl_indiacode_for_query(query)
+    docs = docs[:max_docs]
 
     if not docs:
         return (
-            "No directly relevant readable content was retrieved from India Code for this query. "
-            "The answer should clearly state that India Code verification was not available.",
+            "No directly relevant readable content was retrieved from India Code for this query.",
             [],
         )
 
-    context_blocks = []
+    blocks = []
 
     for index, doc in enumerate(docs, start=1):
-        snippet = make_snippet(doc.get("text", ""), query, max_chars=MAX_TEXT_PER_DOC)
+        snippet = make_snippet(
+            doc.get("text", ""),
+            query,
+            max_chars=max_text_per_doc,
+        )
 
         block = (
             f"[India Code Source {index}]\n"
@@ -601,13 +613,12 @@ def build_indiacode_context(query):
             f"Relevant excerpt:\n{snippet}\n"
         )
 
-        context_blocks.append(block)
+        blocks.append(block)
 
-    return "\n\n".join(context_blocks), docs
+    return "\n\n".join(blocks), docs
 
 
 def format_sources(docs):
-    """Append India Code source links to the answer."""
     if not docs:
         return "\n\nIndia Code sources checked: No readable India Code source was retrieved."
 
@@ -627,11 +638,6 @@ def format_sources(docs):
 # -------------------------------------------------
 
 def load_model():
-    """
-    Download and load the GGUF model once.
-
-    The model is loaded lazily on the first real user message.
-    """
     global llm
 
     if llm is not None:
@@ -664,17 +670,10 @@ def load_model():
 
 
 # -------------------------------------------------
-# Chat history and prompt construction
+# Prompt management / context-window safety
 # -------------------------------------------------
 
-def extract_recent_history(history, max_turns=4):
-    """
-    Keep recent conversation history for follow-up questions.
-
-    Supports:
-    - list of dicts: {"role": "...", "content": "..."}
-    - list of tuples/lists: (user, assistant)
-    """
+def extract_recent_history(history, max_turns=3):
     if not history:
         return ""
 
@@ -698,14 +697,10 @@ def extract_recent_history(history, max_turns=4):
     return conversation
 
 
-def build_prompt(message, history, indiacode_context):
-    """
-    Build a detailed prompt with integrated India Code context.
-    """
-    conversation = extract_recent_history(history, max_turns=4)
+def base_prompt_template(message, conversation, indiacode_context):
     current_date = time.strftime("%Y-%m-%d")
 
-    prompt = f"""You are an Indian legal AI assistant.
+    return f"""You are an Indian legal AI assistant.
 
 Current date: {current_date}
 
@@ -716,13 +711,12 @@ You must answer using:
 
 Critical rules:
 - The India Code lookup has already been performed automatically.
-- Do not describe the lookup as a separate action the user must do.
 - Use India Code context as the primary legal source.
-- If India Code context is missing, weak, or unreadable, clearly say what could not be verified from India Code.
+- If India Code context is missing, weak, OCR-based, or unreadable, clearly say what could not be verified.
+- If source text came from OCR, mention that OCR can contain recognition errors.
 - Do not invent legal provisions, case names, dates, citations, or section text.
-- If the user asks for the latest/current position, rely only on the India Code context provided.
-- If the context contains a PDF extraction warning, mention that the PDF may be scanned or unreadable.
-- Provide detailed, practical explanations.
+- If the user asks for latest/current law, rely only on the India Code context provided.
+- Provide a detailed, practical explanation.
 - Support follow-up questions using the recent conversation.
 - Do not claim to be a lawyer.
 - Do not present the answer as formal legal advice.
@@ -739,7 +733,88 @@ User question:
 
 Answer:"""
 
-    return prompt
+
+def count_tokens(model, prompt):
+    try:
+        return len(model.tokenize(prompt.encode("utf-8"), add_bos=True))
+    except Exception:
+        # Fallback estimate: roughly 4 chars per token.
+        return max(1, len(prompt) // 4)
+
+
+def build_safe_prompt(model, message, history, source_docs):
+    """
+    Build a prompt that fits within the model context window.
+
+    This prevents:
+    Requested tokens (...) exceed context window (...)
+    """
+    max_prompt_tokens = max(128, N_CTX - MAX_TOKENS - PROMPT_SAFETY_MARGIN)
+
+    history_options = [3, 2, 1, 0]
+    doc_options = [5, 4, 3, 2, 1]
+    chars_options = [2400, 1800, 1200, 800, 500]
+
+    for history_turns in history_options:
+        conversation = extract_recent_history(history, max_turns=history_turns)
+
+        for doc_count in doc_options:
+            docs = source_docs[:doc_count]
+
+            for chars_per_doc in chars_options:
+                blocks = []
+
+                for index, doc in enumerate(docs, start=1):
+                    snippet = make_snippet(
+                        doc.get("text", ""),
+                        message,
+                        max_chars=chars_per_doc,
+                    )
+
+                    blocks.append(
+                        f"[India Code Source {index}]\n"
+                        f"Title: {doc.get('title', 'India Code document')}\n"
+                        f"Type: {doc.get('type', 'html')}\n"
+                        f"URL: {doc.get('url')}\n"
+                        f"Relevant excerpt:\n{snippet}\n"
+                    )
+
+                context = "\n\n".join(blocks)
+
+                if not context:
+                    context = "No directly relevant readable content was retrieved from India Code for this query."
+
+                prompt = base_prompt_template(
+                    message=message,
+                    conversation=conversation,
+                    indiacode_context=context,
+                )
+
+                prompt_tokens = count_tokens(model, prompt)
+
+                if prompt_tokens <= max_prompt_tokens:
+                    return prompt, docs, prompt_tokens
+
+    # Final emergency fallback.
+    prompt = base_prompt_template(
+        message=message,
+        conversation="",
+        indiacode_context=(
+            "India Code context was retrieved but had to be heavily reduced "
+            "because it exceeded the local model context window."
+        ),
+    )
+
+    return prompt, source_docs[:1], count_tokens(model, prompt)
+
+
+def safe_generation_max_tokens(prompt_tokens):
+    available = N_CTX - prompt_tokens - PROMPT_SAFETY_MARGIN
+
+    if available < 96:
+        return 96
+
+    return min(MAX_TOKENS, available)
 
 
 # -------------------------------------------------
@@ -747,33 +822,37 @@ Answer:"""
 # -------------------------------------------------
 
 def chat(message, history):
-    """
-    Main Gradio chat function.
-
-    Integrated flow:
-    user message -> India Code lookup -> PDF/HTML extraction -> prompt -> local LLM answer.
-    """
     if not message or not message.strip():
         return "Please enter a question."
 
     try:
         user_query = message.strip()
 
-        # Integrated India Code web step.
-        indiacode_context, source_docs = build_indiacode_context(user_query)
+        # India Code lookup first.
+        _raw_context, source_docs = build_indiacode_context(user_query)
 
-        # Load local model lazily.
+        # Load model.
         model = load_model()
 
-        prompt = build_prompt(
+        # Build prompt safely within context window.
+        prompt, used_docs, prompt_tokens = build_safe_prompt(
+            model=model,
             message=user_query,
             history=history,
-            indiacode_context=indiacode_context,
+            source_docs=source_docs,
+        )
+
+        generation_tokens = safe_generation_max_tokens(prompt_tokens)
+
+        print(
+            f"Prompt tokens: {prompt_tokens}, "
+            f"generation tokens: {generation_tokens}, "
+            f"context window: {N_CTX}"
         )
 
         response = model(
             prompt,
-            max_tokens=MAX_TOKENS,
+            max_tokens=generation_tokens,
             temperature=TEMPERATURE,
             top_p=TOP_P,
             echo=False,
@@ -785,16 +864,17 @@ def chat(message, history):
         if not answer:
             answer = "I could not generate a response. Please try rephrasing your question."
 
-        answer += format_sources(source_docs)
+        answer += format_sources(used_docs)
 
         return answer
 
     except Exception as error:
-        print("Error during India Code lookup or generation:")
+        print("Error during India Code lookup, OCR, or generation:")
         traceback.print_exc()
 
         return (
-            "The app encountered an error while searching India Code or running the model.\n\n"
+            "The app encountered an error while searching India Code, reading a PDF, "
+            "running OCR, or generating the response.\n\n"
             f"Error details: {str(error)}"
         )
 
@@ -806,21 +886,21 @@ def chat(message, history):
 description = """
 # 🏛️ Indian Legal AI Assistant
 
-Ask questions about Indian laws, Acts, legal procedures, sections, rules, and follow-up questions.
+Ask questions about Indian laws, Acts, legal sections, rules, and follow-up questions.
 
 This app uses:
 - a local GGUF model through `llama-cpp-python`
-- an integrated India Code lookup step
-- readable HTML extraction
-- extractable-text PDF reading through `pypdf`
+- integrated India Code lookup
+- HTML extraction
+- text-based PDF extraction
+- lightweight OCR fallback for scanned PDFs
 - recent chat history for follow-up questions
 
-The India Code lookup runs inside each answer.  
-You do not need to run a separate search.
+The India Code lookup runs inside each answer.
 
 ---
 
-**Primary source used by the web step:** https://www.indiacode.nic.in/
+**Primary source:** https://www.indiacode.nic.in/
 
 **Disclaimer:** This assistant provides general legal information only.  
 It is not a substitute for advice from a qualified legal professional.  
@@ -855,11 +935,13 @@ if __name__ == "__main__":
     print(f"Using model: {MODEL_REPO}/{MODEL_FILE}")
     print("CPU-only mode enabled.")
     print("Integrated India Code lookup enabled.")
+    print("PDF OCR fallback enabled.")
     print(
         f"N_CTX={N_CTX}, "
         f"N_THREADS={N_THREADS}, "
         f"N_THREADS_BATCH={N_THREADS_BATCH}, "
-        f"N_BATCH={N_BATCH}"
+        f"N_BATCH={N_BATCH}, "
+        f"MAX_TOKENS={MAX_TOKENS}"
     )
 
     demo.launch(