Update app.py

app.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 """
-Indian Legal AI Assistant with integrated web lookup.
+Indian Legal AI Assistant with integrated India Code web lookup.
 
 CPU-only Hugging Face Spaces app.
 
@@ -9,20 +9,26 @@ invincibleambuj/Ambuj-Tripathi-Indian-Legal-Llama-GGUF
 llama-3.2-1b-instruct.Q4_K_M.gguf
 
 Web step:
-The app automatically searches the web for each user query and injects
-the retrieved snippets into the model prompt. The user does not need to
-run a separate search.
+- 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.
 """
 
 import os
 import re
-import time
+import io
 import html
+import time
 import traceback
-from urllib.parse import quote_plus, urlparse
+from functools import lru_cache
+from urllib.parse import urljoin, urlparse, quote_plus, urldefrag
 
 import requests
 from bs4 import BeautifulSoup
+from pypdf import PdfReader
 
 import gradio as gr
 from huggingface_hub import hf_hub_download
@@ -42,30 +48,53 @@ 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. Keep this at 0.
+# CPU only.
 N_GPU_LAYERS = 0
 
 # Generation settings.
-MAX_TOKENS = int(os.getenv("MAX_TOKENS", "512"))
-TEMPERATURE = float(os.getenv("TEMPERATURE", "0.45"))
+MAX_TOKENS = int(os.getenv("MAX_TOKENS", "768"))
+TEMPERATURE = float(os.getenv("TEMPERATURE", "0.35"))
 TOP_P = float(os.getenv("TOP_P", "0.9"))
 
-# Web settings.
-WEB_RESULTS = int(os.getenv("WEB_RESULTS", "5"))
-WEB_TIMEOUT = int(os.getenv("WEB_TIMEOUT", "10"))
-
-# Use persistent storage if attached, otherwise local cache.
+# Model cache.
 MODEL_CACHE_DIR = os.getenv("MODEL_CACHE_DIR", "./models")
 
 llm = None
 
 
 # -------------------------------------------------
-# Utility helpers
+# India Code search configuration
+# -------------------------------------------------
+
+INDIACODE_HOME = "https://www.indiacode.nic.in/"
+ALLOWED_DOMAINS = {
+    "indiacode.nic.in",
+    "www.indiacode.nic.in",
+}
+
+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"))
+
+HEADERS = {
+    "User-Agent": (
+        "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 "
+        "(KHTML, like Gecko) Chrome/120.0 Safari/537.36"
+    )
+}
+
+
+# -------------------------------------------------
+# General helpers
 # -------------------------------------------------
 
 def clean_text(text):
-    """Normalize whitespace and remove noisy characters."""
+    """Normalize whitespace and decode HTML entities."""
     if not text:
         return ""
 
@@ -74,190 +103,521 @@ def clean_text(text):
     return text.strip()
 
 
-def domain_from_url(url):
-    """Return a readable domain name from a URL."""
+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)
-        return parsed.netloc.replace("www.", "")
+        domain = parsed.netloc.lower()
+        return domain in ALLOWED_DOMAINS
     except Exception:
-        return "source"
+        return False
+
 
+def looks_like_pdf_url(url):
+    return ".pdf" in url.lower()
 
-def is_probably_url(text):
-    """Basic URL detector."""
-    return text.startswith("http://") or text.startswith("https://")
+
+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",
+    }
+
+    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)
+
+    if not terms:
+        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
+
+    return score
+
+
+def make_snippet(text, query, max_chars=900):
+    """Create a short context snippet near query terms."""
+    text = clean_text(text)
+    if not text:
+        return ""
+
+    terms = query_terms(query)
+    lower = text.lower()
+
+    first_hit = None
+    for term in terms:
+        idx = lower.find(term)
+        if idx != -1:
+            first_hit = idx
+            break
+
+    if first_hit is None:
+        return text[:max_chars]
+
+    start = max(first_hit - 250, 0)
+    end = min(start + max_chars, len(text))
+    return text[start:end]
 
 
 # -------------------------------------------------
-# Integrated web search
+# HTTP helpers
 # -------------------------------------------------
 
-def duckduckgo_search(query, max_results=5):
-    """
-    Lightweight web search using DuckDuckGo HTML results.
+def safe_get(url, timeout=REQUEST_TIMEOUT):
+    """GET request with basic error handling."""
+    try:
+        response = requests.get(url, headers=HEADERS, timeout=timeout)
+        response.raise_for_status()
+        return response
+    except Exception as error:
+        print(f"GET failed: {url} :: {error}")
+        return None
 
-    This does not require an API key.
 
-    Note:
-    Search engines can sometimes block automated requests or return no results.
-    The app handles that gracefully.
-    """
-    results = []
+def content_type(response):
+    if response is None:
+        return ""
+    return response.headers.get("content-type", "").lower()
 
-    if not query or not query.strip():
-        return results
 
-    search_url = f"https://duckduckgo.com/html/?q={quote_plus(query)}"
+# -------------------------------------------------
+# HTML and PDF extraction
+# -------------------------------------------------
 
-    headers = {
-        "User-Agent": (
-            "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 "
-            "(KHTML, like Gecko) Chrome/120.0 Safari/537.36"
-        )
-    }
+def extract_links_from_html(html_text, base_url):
+    """Extract India Code links from HTML."""
+    links = []
 
     try:
-        response = requests.get(
-            search_url,
-            headers=headers,
-            timeout=WEB_TIMEOUT,
-        )
-        response.raise_for_status()
+        soup = BeautifulSoup(html_text, "html.parser")
+
+        for a in soup.find_all("a", href=True):
+            href = a.get("href", "")
+            text = clean_text(a.get_text(" "))
+            url = normalize_url(href, base_url)
+
+            if is_indiacode_url(url):
+                links.append(
+                    {
+                        "url": url,
+                        "anchor": text,
+                    }
+                )
+    except Exception as error:
+        print(f"Link extraction failed for {base_url}: {error}")
 
-        soup = BeautifulSoup(response.text, "html.parser")
+    return links
 
-        for result in soup.select(".result"):
-            title_el = result.select_one(".result__title a")
-            snippet_el = result.select_one(".result__snippet")
 
-            if not title_el:
-                continue
+def extract_text_from_html(html_text):
+    """Extract readable text from HTML."""
+    try:
+        soup = BeautifulSoup(html_text, "html.parser")
 
-            title = clean_text(title_el.get_text(" "))
-            url = title_el.get("href", "").strip()
-            snippet = clean_text(snippet_el.get_text(" ")) if snippet_el else ""
+        for tag in soup(["script", "style", "nav", "footer", "header", "aside", "form"]):
+            tag.decompose()
 
-            if not title or not url:
-                continue
+        title = clean_text(soup.title.get_text(" ")) if soup.title else ""
 
-            results.append(
-                {
-                    "title": title,
-                    "url": url,
-                    "domain": domain_from_url(url),
-                    "snippet": snippet,
-                }
-            )
+        parts = []
+        if title:
+            parts.append(title)
 
-            if len(results) >= max_results:
-                break
+        for tag in soup.find_all(["h1", "h2", "h3", "h4", "p", "li", "td", "th"], limit=300):
+            text = clean_text(tag.get_text(" "))
+            if len(text) >= 20:
+                parts.append(text)
 
-    except Exception as error:
-        print(f"Web search failed: {error}")
+        return clean_text(" ".join(parts))
 
-    return results
+    except Exception as error:
+        print(f"HTML extraction failed: {error}")
+        return ""
 
 
-def fetch_page_text(url):
+def extract_pdf_text(pdf_bytes, max_pages=MAX_PDF_PAGES):
     """
-    Fetch readable text from a page.
+    Extract text from a PDF.
 
-    This is optional enrichment after search results.
-    It keeps text short to avoid slowing inference.
+    This reads text-based PDFs. It will not OCR scanned image-only PDFs.
     """
-    if not url or not is_probably_url(url):
+    try:
+        reader = PdfReader(io.BytesIO(pdf_bytes))
+        parts = []
+
+        total_pages = len(reader.pages)
+        pages_to_read = min(total_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}")
+            except Exception as page_error:
+                print(f"PDF page extraction failed: {page_error}")
+
+        extracted = clean_text(" ".join(parts))
+
+        if not extracted:
+            return "[PDF detected, but no extractable text was found. The PDF may be scanned/image-based.]"
+
+        return extracted
+
+    except Exception as error:
+        print(f"PDF extraction failed: {error}")
         return ""
 
-    headers = {
-        "User-Agent": (
-            "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 "
-            "(KHTML, like Gecko) Chrome/120.0 Safari/537.36"
-        )
-    }
 
+@lru_cache(maxsize=256)
+def fetch_document_text(url):
+    """
+    Fetch and extract text from an India Code HTML or PDF document.
+    Cached to speed up follow-up questions.
+    """
+    if not is_indiacode_url(url):
+        return {
+            "url": url,
+            "title": "Blocked non-India-Code URL",
+            "text": "",
+            "links": [],
+            "type": "blocked",
+        }
+
+    response = safe_get(url)
+    if response is None:
+        return {
+            "url": url,
+            "title": "Fetch failed",
+            "text": "",
+            "links": [],
+            "type": "failed",
+        }
+
+    ctype = content_type(response)
+
+    if "application/pdf" in ctype or looks_like_pdf_url(url):
+        text = extract_pdf_text(response.content)
+        return {
+            "url": url,
+            "title": url.split("/")[-1] or "India Code PDF",
+            "text": text,
+            "links": [],
+            "type": "pdf",
+        }
+
+    html_text = response.text
+    links = extract_links_from_html(html_text, url)
+    text = extract_text_from_html(html_text)
+
+    title = "India Code page"
     try:
-        response = requests.get(url, headers=headers, timeout=WEB_TIMEOUT)
-        response.raise_for_status()
+        soup = BeautifulSoup(html_text, "html.parser")
+        if soup.title:
+            title = clean_text(soup.title.get_text(" "))
+    except Exception:
+        pass
+
+    return {
+        "url": url,
+        "title": title,
+        "text": text,
+        "links": links,
+        "type": "html",
+    }
+
+
+# -------------------------------------------------
+# India Code discovery
+# -------------------------------------------------
+
+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,
+
+        # Central Acts browse pages.
+        "https://www.indiacode.nic.in/handle/123456789/1362/browse?type=shorttitle",
+        "https://www.indiacode.nic.in/handle/123456789/1362/browse?type=actno",
+        "https://www.indiacode.nic.in/handle/123456789/1362/browse?type=actyear",
+        "https://www.indiacode.nic.in/handle/123456789/1362/browse?type=enactmentdate",
+        "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.
+        "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):
+            discovered.append(url)
+
+    return discovered
 
-        content_type = response.headers.get("content-type", "").lower()
-        if "text/html" not in content_type:
-            return ""
 
+def duckduckgo_site_discovery(query):
+    """
+    URL discovery using a site-restricted query.
+
+    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.
+    """
+    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
+
+    try:
         soup = BeautifulSoup(response.text, "html.parser")
 
-        for tag in soup(["script", "style", "nav", "footer", "header", "aside"]):
-            tag.decompose()
+        for a in soup.select(".result__title a"):
+            url = a.get("href", "").strip()
 
-        paragraphs = []
-        for p in soup.find_all(["p", "li"], limit=20):
-            text = clean_text(p.get_text(" "))
-            if len(text) > 40:
-                paragraphs.append(text)
+            # 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:
+                        url = qs["uddg"][0]
+                except Exception:
+                    pass
 
-        page_text = " ".join(paragraphs)
-        return page_text[:1800]
+            url = normalize_url(url)
+
+            if is_indiacode_url(url):
+                discovered.append(url)
+
+            if len(discovered) >= MAX_DISCOVERY_RESULTS:
+                break
 
     except Exception as error:
-        print(f"Page fetch failed for {url}: {error}")
-        return ""
+        print(f"Site discovery parsing failed: {error}")
+
+    return discovered
 
 
-def build_web_context(query):
+def relevant_link_filter(link, query):
     """
-    Search the web and build compact context for the model.
+    Decide whether to crawl a link.
 
-    The web step is integrated into the prompt, not shown as a separate workflow.
+    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
     """
-    results = duckduckgo_search(query, max_results=WEB_RESULTS)
+    url = link.get("url", "")
+    anchor = link.get("anchor", "")
 
-    if not results:
-        return "No web results were retrieved for this query.", []
+    if not is_indiacode_url(url):
+        return False
 
-    context_blocks = []
-    enriched_results = []
+    url_l = url.lower()
+    anchor_l = anchor.lower()
 
-    for index, result in enumerate(results, start=1):
-        page_text = ""
+    if looks_like_pdf_url(url):
+        return True
 
-        # Fetch only top 2 pages to keep things fast.
-        if index <= 2:
-            page_text = fetch_page_text(result["url"])
+    important_patterns = [
+        "/handle/",
+        "/bitstream/",
+        "/browse",
+        "repealed-act",
+        "spent-act",
+    ]
 
-        enriched = dict(result)
-        enriched["page_text"] = page_text
-        enriched_results.append(enriched)
+    if any(pattern in url_l for pattern in important_patterns):
+        return True
 
-        block = (
-            f"[Source {index}]\n"
-            f"Title: {result['title']}\n"
-            f"Domain: {result['domain']}\n"
-            f"URL: {result['url']}\n"
-            f"Search snippet: {result['snippet']}\n"
+    terms = query_terms(query)
+    if any(term in url_l or term in anchor_l for term in terms):
+        return True
+
+    return False
+
+
+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()
+
+    for url in seeds:
+        url = normalize_url(url)
+        if is_indiacode_url(url) and url not in seen:
+            queue.append(url)
+            seen.add(url)
+
+    visited = set()
+    scored_docs = []
+
+    while queue and len(visited) < MAX_CRAWL_PAGES:
+        url = queue.pop(0)
+
+        if url in visited:
+            continue
+
+        visited.add(url)
+
+        doc = fetch_document_text(url)
+        text = doc.get("text", "")
+        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)
+
+        if score > 0 or doc_type == "pdf":
+            scored_docs.append(
+                {
+                    "url": url,
+                    "title": title,
+                    "type": doc_type,
+                    "score": score,
+                    "text": text,
+                }
+            )
+
+        # Follow relevant India Code links from HTML pages.
+        for link in doc.get("links", []):
+            link_url = normalize_url(link.get("url", ""), url)
+
+            if link_url in seen:
+                continue
+
+            if relevant_link_filter(link, query):
+                queue.append(link_url)
+                seen.add(link_url)
+
+    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.
+    """
+    docs = crawl_indiacode_for_query(query)
+
+    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.",
+            [],
         )
 
-        if page_text:
-            block += f"Page text excerpt: {page_text}\n"
+    context_blocks = []
 
-        context_blocks.append(block)
+    for index, doc in enumerate(docs, start=1):
+        snippet = make_snippet(doc.get("text", ""), query, max_chars=MAX_TEXT_PER_DOC)
 
-    web_context = "\n".join(context_blocks)
-    return web_context, enriched_results
+        block = (
+            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_blocks.append(block)
 
-def format_sources(results):
-    """Format sources for appending to the assistant answer."""
-    if not results:
-        return ""
+    return "\n\n".join(context_blocks), docs
 
-    lines = ["\n\nSources checked:"]
 
-    for index, result in enumerate(results, start=1):
-        title = result.get("title", "Source")
-        url = result.get("url", "")
-        domain = result.get("domain", "source")
+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."
 
-        if url:
-            lines.append(f"{index}. {title} — {domain}\n   {url}")
+    lines = ["\n\nIndia Code sources checked:"]
+
+    for index, doc in enumerate(docs, start=1):
+        title = doc.get("title", "India Code document")
+        url = doc.get("url", "")
+        dtype = doc.get("type", "html")
+        lines.append(f"{index}. {title} [{dtype}]\n   {url}")
 
     return "\n".join(lines)
 
@@ -304,14 +664,14 @@ def load_model():
 
 
 # -------------------------------------------------
-# Prompt construction
+# Chat history and prompt construction
 # -------------------------------------------------
 
 def extract_recent_history(history, max_turns=4):
     """
-    Keep only recent conversation history for speed and follow-up support.
+    Keep recent conversation history for follow-up questions.
 
-    Supports both:
+    Supports:
     - list of dicts: {"role": "...", "content": "..."}
     - list of tuples/lists: (user, assistant)
     """
@@ -338,12 +698,11 @@ def extract_recent_history(history, max_turns=4):
     return conversation
 
 
-def build_prompt(message, history, web_context):
+def build_prompt(message, history, indiacode_context):
     """
-    Build a detailed prompt with integrated web context.
+    Build a detailed prompt with integrated India Code context.
     """
     conversation = extract_recent_history(history, max_turns=4)
-
     current_date = time.strftime("%Y-%m-%d")
 
     prompt = f"""You are an Indian legal AI assistant.
@@ -352,28 +711,28 @@ Current date: {current_date}
 
 You must answer using:
 1. The user's question.
-2. The recent chat history.
-3. The integrated web context below.
-
-Important behaviour:
-- The web context has already been retrieved automatically.
-- Do not say "I searched separately" or ask the user to search.
-- Give a detailed, practical answer.
-- Use web context for current or time-sensitive facts.
-- If the web context is weak, say what you could not verify.
-- For Indian law questions, explain the legal concept clearly.
+2. The recent conversation.
+3. The integrated India Code context below.
+
+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.
+- 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.
+- Support follow-up questions using the recent conversation.
 - Do not claim to be a lawyer.
 - Do not present the answer as formal legal advice.
 - For specific legal matters, advise consulting a qualified lawyer.
-- Support follow-up questions by preserving continuity from recent chat history.
-- When sources conflict, mention the conflict.
-- Do not invent case law, sections, dates, or citations.
 
 Recent conversation:
 {conversation}
 
-Integrated web context:
-{web_context}
+Integrated India Code context:
+{indiacode_context}
 
 User question:
 {message}
@@ -384,15 +743,15 @@ Answer:"""
 
 
 # -------------------------------------------------
-# Chat function
+# Main chat function
 # -------------------------------------------------
 
 def chat(message, history):
     """
     Main Gradio chat function.
 
-    The web step is integrated here:
-    user message -> web search -> prompt with web context -> local LLM answer.
+    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."
@@ -400,8 +759,8 @@ def chat(message, history):
     try:
         user_query = message.strip()
 
-        # Integrated web step.
-        web_context, web_results = build_web_context(user_query)
+        # Integrated India Code web step.
+        indiacode_context, source_docs = build_indiacode_context(user_query)
 
         # Load local model lazily.
         model = load_model()
@@ -409,7 +768,7 @@ def chat(message, history):
         prompt = build_prompt(
             message=user_query,
             history=history,
-            web_context=web_context,
+            indiacode_context=indiacode_context,
         )
 
         response = model(
@@ -426,17 +785,16 @@ def chat(message, history):
         if not answer:
             answer = "I could not generate a response. Please try rephrasing your question."
 
-        # Append sources so the user can verify.
-        answer += format_sources(web_results)
+        answer += format_sources(source_docs)
 
         return answer
 
     except Exception as error:
-        print("Error during generation:")
+        print("Error during India Code lookup or generation:")
         traceback.print_exc()
 
         return (
-            "The app encountered an error while loading, searching, or running the model.\n\n"
+            "The app encountered an error while searching India Code or running the model.\n\n"
             f"Error details: {str(error)}"
         )
 
@@ -448,18 +806,22 @@ def chat(message, history):
 description = """
 # 🏛️ Indian Legal AI Assistant
 
-Ask questions about Indian laws, acts, legal procedures, and current legal information.
+Ask questions about Indian laws, Acts, legal procedures, sections, rules, and follow-up questions.
 
 This app uses:
 - a local GGUF model through `llama-cpp-python`
-- an integrated web step for up-to-date information
+- an integrated India Code lookup step
+- readable HTML extraction
+- extractable-text PDF reading through `pypdf`
 - recent chat history for follow-up questions
 
-The web step runs automatically inside each answer.
+The India Code lookup runs inside each answer.  
 You do not need to run a separate search.
 
 ---
 
+**Primary source used by the web step:** 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.  
 For specific legal matters, please consult a lawyer.
@@ -470,15 +832,15 @@ demo = gr.ChatInterface(
     title="Indian Legal AI Assistant",
     description=description,
     textbox=gr.Textbox(
-        placeholder="Ask about Indian laws, recent legal updates, procedures, rights, cases, or follow-up questions...",
+        placeholder="Ask about Indian laws, Acts, sections, rules, or follow-up questions...",
         lines=3,
         label="Your Question",
     ),
     examples=[
-        "What is the Indian Penal Code?",
-        "What changed after the new criminal laws came into force in India?",
-        "Explain Section 377 of IPC and its current legal status.",
+        "What is the current status of Section 377 under Indian law?",
+        "Explain the Bharatiya Nyaya Sanhita in detail.",
         "What are the grounds for divorce under the Hindu Marriage Act?",
+        "Find the latest India Code position on the Right to Information Act.",
     ],
     cache_examples=False,
 )
@@ -492,7 +854,13 @@ if __name__ == "__main__":
     print("Starting Indian Legal AI Assistant...")
     print(f"Using model: {MODEL_REPO}/{MODEL_FILE}")
     print("CPU-only mode enabled.")
-    print(f"N_CTX={N_CTX}, N_THREADS={N_THREADS}, N_THREADS_BATCH={N_THREADS_BATCH}, N_BATCH={N_BATCH}")
+    print("Integrated India Code lookup enabled.")
+    print(
+        f"N_CTX={N_CTX}, "
+        f"N_THREADS={N_THREADS}, "
+        f"N_THREADS_BATCH={N_THREADS_BATCH}, "
+        f"N_BATCH={N_BATCH}"
+    )
 
     demo.launch(
         server_name="0.0.0.0",