feat: unified search, Claude BYOK, expanded legislation coverage

Search pipeline:
- Three-layer search in chat: keyword cache → semantic vectors → live Lex API
- Semantic search now also used in chat tab (was Scenario Matcher only)
- Section Lookup falls back to live Lex API on cache miss
- Scenario Matcher falls back to live Lex API when semantic returns nothing

LLM:
- Replace Gemini with Claude (claude-sonnet-4-6) via BYOK
- Users enter their own Anthropic API key (session only, never stored)
- Graceful fallback to raw statutory text when no key provided
- Consistent with other CQAI tools (nursing-research-writer pattern)

Dependencies:
- Remove torch (replaced numpy cosine similarity — saves ~2GB memory on CPU)
- Add anthropic>=0.40
- local_search.py now initialises embedding model at startup, not first request

Legislation coverage:
- Add Mental Health Act 2007, Mental Capacity (Amendment) Act 2019,
Health and Social Care Act 2008, Misuse of Drugs Act 1971,
Data Protection Act 2018 to NURSING_ACTS and lex_client.py
- Expand keyword map from 15 to 50+ entries:
CTOs, AMHP, nearest relative, DoLS, LPS, duty of candour,
data protection, controlled drugs, reasonable adjustments,
human rights, care needs assessment, court of protection, LPA

Bug fixes:
- Resolve git merge conflict in lex_client.py (duplicate file content removed)
- lex_client.py now imported and used in app.py (was dead code)

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

app.py

@@ -2,8 +2,10 @@
 NurseLex — Legal Literacy Agent for All Nurses and Nursing Students
 Architecture:
   1. Local legislation.parquet — 219K health/social care Acts & SIs for browsing
-  2. cached_legislation.py — 1,128 sections loaded from nursing_sections.json
-  3. Gemini Flash REST API — Plain English explanations (with retry logic)
+  2. cached_legislation.py — 1,128 sections (keyword + BM25-style search)
+  3. local_search.py — semantic vector search via i-dot-ai fine-tuned MiniLM
+  4. lex_client.py — live i.AI Lex API fallback for cache misses
+  5. Claude (BYOK) via Anthropic SDK — plain English explanations
 """
 import os
 import asyncio
@@ -11,9 +13,11 @@ import httpx
 import logging
 import pandas as pd
 import gradio as gr
+from anthropic import AsyncAnthropic
 
 from cached_legislation import search_cached
 from local_search import search_scenarios_locally
+from lex_client import search_legislation_sections
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -30,10 +34,13 @@ except Exception as e:
 # --- Key nursing legislation IDs ---
 NURSING_ACTS = {
     "Mental Health Act 1983": "ukpga/1983/20",
+    "Mental Health Act 2007": "ukpga/2007/12",
     "Mental Capacity Act 2005": "ukpga/2005/9",
+    "Mental Capacity (Amendment) Act 2019": "ukpga/2019/17",
     "Care Act 2014": "ukpga/2014/23",
     "Human Rights Act 1998": "ukpga/1998/42",
     "Equality Act 2010": "ukpga/2010/15",
+    "Health and Social Care Act 2008": "ukpga/2008/14",
     "Health and Social Care Act 2012": "ukpga/2012/7",
     "Mental Health Units (Use of Force) Act 2018": "ukpga/2018/27",
     "Autism Act 2009": "ukpga/2009/15",
@@ -41,29 +48,26 @@ NURSING_ACTS = {
     "Children Act 2004": "ukpga/2004/31",
     "Safeguarding Vulnerable Groups Act 2006": "ukpga/2006/47",
     "Health and Care Act 2022": "ukpga/2022/31",
+    "Misuse of Drugs Act 1971": "ukpga/1971/38",
+    "Data Protection Act 2018": "ukpga/2018/12",
 }
 REVERSE_ACTS = {v: k for k, v in NURSING_ACTS.items()}
 
-# --- Gemini REST API ---
-GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "")
-GEMINI_BASE = "https://generativelanguage.googleapis.com/v1beta/models"
-GEMINI_MODELS = ["gemini-2.0-flash-lite", "gemini-2.0-flash"]
-
 SYSTEM_PROMPT = """You are NurseLex, a legal literacy assistant for all UK nurses and nursing students.
 
 Your role:
 1. Answer legal questions using ONLY the legislation text provided in the context.
 2. Explain the law in clear, plain English suitable for all nurses and nursing students.
 3. Always cite the specific Act, section number, and year.
-4. If the context doesn't contain enough information, say so clearly.
+4. If the context doesn't contain enough information, say so clearly — do not invent statutory text.
 5. Add practical nursing implications (e.g., "In practice, this means...").
-6. Include professional reminders (e.g., NMC Code, duty of care).
+6. Include professional reminders (e.g., NMC Code, duty of care, trust policy).
 
-Disclaimers to include:
+Always include:
 - "This is for educational purposes only — always consult your trust's legal team for specific cases."
 - "This reflects the legislation as written — local trust policies may add additional requirements."
 
-Format with clear headings, bullet points, and bold key terms."""
+Format responses with clear headings, bullet points, and bold key terms."""
 
 QUICK_QUESTIONS = [
     "What is Section 5(4) of the Mental Health Act and when can a nurse use it?",
@@ -73,72 +77,99 @@ QUICK_QUESTIONS = [
     "What does Section 117 aftercare mean and who is entitled?",
     "What are a nurse's legal duties under the Care Act 2014 for safeguarding?",
     "What is Deprivation of Liberty and when do DoLS apply?",
-    "What rights does a patient have under Section 136?",
+    "What is a Community Treatment Order and when can it be used?",
 ]
 
-async def call_gemini(prompt: str) -> str:
-    """Call Gemini via REST API with retry logic and model fallback."""
-    if not GEMINI_API_KEY:
+
+# --- Claude BYOK ---
+
+async def call_claude(prompt: str, api_key: str) -> str:
+    """Call Claude via the Anthropic SDK with the user's own API key."""
+    if not api_key or not api_key.strip().startswith("sk-"):
         return ""
 
-    payload = {
-        "system_instruction": {"parts": [{"text": SYSTEM_PROMPT}]},
-        "contents": [{"parts": [{"text": prompt}]}],
-        "generationConfig": {"temperature": 0.3, "maxOutputTokens": 2048},
-    }
-
-    async with httpx.AsyncClient(timeout=60.0) as client:
-        for model in GEMINI_MODELS:
-            url = f"{GEMINI_BASE}/{model}:generateContent?key={GEMINI_API_KEY}"
-            for attempt in range(3):
-                try:
-                    resp = await client.post(url, json=payload)
-                    if resp.status_code == 429:
-                        wait = 2 ** (attempt + 1)
-                        logger.warning(f"Rate limited ({model}), retrying in {wait}s")
-                        await asyncio.sleep(wait)
-                        continue
-                    resp.raise_for_status()
-                    data = resp.json()
-                    return data["candidates"][0]["content"]["parts"][0]["text"]
-                except httpx.HTTPStatusError as e:
-                    if e.response.status_code == 429:
-                        wait = 2 ** (attempt + 1)
-                        logger.warning(f"Rate limited ({model}), retrying in {wait}s")
-                        await asyncio.sleep(wait)
-                        continue
-                    logger.error(f"Gemini API error ({model}): {e.response.status_code}")
-                    break
-                except Exception as e:
-                    logger.error(f"Gemini error ({model}): {type(e).__name__}")
-                    break
-            logger.info(f"Model {model} exhausted, trying next...")
+    try:
+        client = AsyncAnthropic(api_key=api_key.strip())
+        message = await client.messages.create(
+            model="claude-sonnet-4-6",
+            max_tokens=2048,
+            system=SYSTEM_PROMPT,
+            messages=[{"role": "user", "content": prompt}],
+        )
+        return message.content[0].text
+    except Exception as e:
+        err = str(e).lower()
+        logger.error(f"Claude API error: {type(e).__name__}: {e}")
+        if "authentication" in err or "401" in err or "invalid x-api-key" in err:
+            return "__AUTH_ERROR__"
+        if "rate" in err or "429" in err:
+            return "__RATE_LIMIT__"
+        return ""
+
+
+# --- Unified search pipeline ---
+
+async def _gather_sections(user_question: str) -> list[dict]:
+    """
+    Three-layer search: keyword/BM25 cache → semantic vector search → live Lex API.
+    Results are deduplicated by (legislation_id, section_number).
+    """
+    seen: set[tuple] = set()
+    sections: list[dict] = []
+
+    def add(s: dict) -> bool:
+        key = (s.get("legislation_id"), s.get("number"))
+        if key not in seen:
+            seen.add(key)
+            sections.append(s)
+            return True
+        return False
+
+    # Layer 1: keyword + BM25 cache (instant)
+    for s in search_cached(user_question, max_results=5):
+        add(s)
+    logger.info(f"Cache returned {len(sections)} sections")
+
+    # Layer 2: local semantic vector search
+    semantic = search_scenarios_locally(user_question, top_k=5)
+    for s in semantic:
+        add(s)
+    logger.info(f"After semantic search: {len(sections)} sections")
+
+    # Layer 3: live Lex API fallback when local sources are sparse
+    if len(sections) < 3:
+        logger.info("Falling back to live Lex API...")
+        try:
+            live = await search_legislation_sections(user_question, size=5)
+            for s in live:
+                add(s)
+            logger.info(f"After live Lex API: {len(sections)} sections")
+        except Exception as e:
+            logger.warning(f"Live Lex API fallback failed: {e}")
+
+    return sections
 
-    logger.error("All Gemini models failed")
-    return ""
 
 def search_legislation_index(query: str, max_results: int = 10) -> pd.DataFrame:
     """Search the full legislation index parquet by title."""
     if LEG_DF.empty:
         return pd.DataFrame()
-
     mask = LEG_DF["title"].str.contains(query, case=False, na=False)
-    results = LEG_DF[mask].sort_values("year", ascending=False).head(max_results)
-    return results
+    return LEG_DF[mask].sort_values("year", ascending=False).head(max_results)
 
-async def query_and_respond(user_question: str, history: list) -> str:
-    """Main RAG pipeline: local cached sections (1,128) + Gemini explanation."""
+
+async def query_and_respond(user_question: str, history: list, api_key: str = "") -> str:
+    """Main RAG pipeline: three-layer search + Claude explanation."""
     if not user_question.strip():
         return "Please enter a question about UK healthcare legislation."
 
-    # Step 1: Search local legislation sections
-    sections = search_cached(user_question, max_results=5)
-    logger.info(f"Local search returned {len(sections)} sections for: {user_question[:60]}")
+    # Gather sections from all three layers
+    sections = await _gather_sections(user_question)
 
-    # Step 2: Search parquet index for related Acts
+    # Search parquet index for related Acts
     related_acts = search_legislation_index(user_question, max_results=5)
 
-    # Step 3: Build context
+    # Build context string for the LLM
     context_parts = []
     for section in sections:
         title = section.get("title", "Untitled")
@@ -147,35 +178,45 @@ async def query_and_respond(user_question: str, history: list) -> str:
         num = section.get("number", "")
         context_parts.append(f"### {title}\n**Source:** {leg_id}, Section {num}\n\n{text}\n")
 
-    context = "\n---\n".join(context_parts) if context_parts else "No matching legislation sections found in cache."
+    context = "\n---\n".join(context_parts) if context_parts else "No matching legislation sections found."
+
+    prompt = (
+        f"## Nurse's Question\n{user_question}\n\n"
+        f"## Relevant UK Legislation\n{context}\n\n"
+        f"Please answer the nurse's question using the legislation above."
+    )
 
-    # Step 4: Generate Gemini response
-    prompt = f"## Nurse's Question\n{user_question}\n\n## Relevant UK Legislation\n{context}\n\nPlease answer the nurse's question using the legislation above."
-    
-    answer = await call_gemini(prompt)
+    answer = await call_claude(prompt, api_key)
 
-    if not answer:
-        # Fallback: show raw legislation if Gemini fails or is missing key
+    # Handle API errors gracefully
+    if answer == "__AUTH_ERROR__":
+        answer = (
+            "❌ **Invalid API key.** Please check your Anthropic API key in the sidebar.\n\n"
+            + _build_fallback(user_question, sections)
+        )
+    elif answer == "__RATE_LIMIT__":
+        answer = (
+            "⏳ **Rate limited.** Please wait a moment and try again.\n\n"
+            + _build_fallback(user_question, sections)
+        )
+    elif not answer:
         answer = _build_fallback(user_question, sections)
-        if not GEMINI_API_KEY:
-            answer += "\n\n⚠️ *Set `GEMINI_API_KEY` in Space secrets for AI-powered plain English explanations.*"
-        elif "rate limit" in answer.lower():
-             answer += "\n\n⚠️ *Gemini is currently rate limited, falling back to raw legislation.*"
-
-    # Add source citations
-    source_acts = set()
-    for s in sections:
-        leg_id = s.get("legislation_id", "")
-        if leg_id:
-            source_acts.add(leg_id)
+        if not api_key or not api_key.strip().startswith("sk-"):
+            answer += (
+                "\n\n💡 *Enter your Anthropic API key in the sidebar for AI-powered plain English explanations. "
+                "Get a free key at [console.anthropic.com](https://console.anthropic.com).*"
+            )
 
+    # Append source citations
+    source_acts = {s.get("legislation_id") for s in sections if s.get("legislation_id")}
     if source_acts:
         answer += "\n\n---\n📚 **Sources:** "
-        answer += " | ".join(f"[{sid}](https://www.legislation.gov.uk/id/{sid})" for sid in sorted(source_acts))
+        answer += " | ".join(
+            f"[{sid}](https://www.legislation.gov.uk/id/{sid})" for sid in sorted(source_acts)
+        )
 
-    # Add related Acts from parquet
+    # Append related Acts from parquet
     if not related_acts.empty:
-        answer += "\n\n📖 **Related legislation:** "
         act_links = []
         for _, row in related_acts.head(3).iterrows():
             uri = row.get("uri", "")
@@ -183,23 +224,24 @@ async def query_and_respond(user_question: str, history: list) -> str:
             if uri and title:
                 act_links.append(f"[{title}]({uri})")
         if act_links:
-            answer += " | ".join(act_links)
+            answer += "\n\n📖 **Related legislation:** " + " | ".join(act_links)
 
     answer += "\n\n🏛️ *Data from [legislation.gov.uk](https://www.legislation.gov.uk/) — Crown Copyright, OGL v3.0*"
     answer += "\n\n> 📌 **Always verify:** Check the linked source text above before relying on any explanation for practice."
     return answer
 
+
 def _build_fallback(question: str, sections: list) -> str:
-    """Show raw legislation without LLM."""
+    """Display raw legislation text when no LLM is available."""
     response = f"## Legislation relevant to: *{question}*\n\n"
 
     if not sections:
         response += (
-            "No matching sections found in cache. Try searching the full **Browse Legislation** tab for the Act title, or try specific terms like:\n"
+            "No matching sections found. Try the **Browse Legislation** tab, or use specific terms like:\n"
             "- **\"Section 5(4)\"** or **\"nurse holding power\"**\n"
             "- **\"best interests\"** or **\"capacity\"**\n"
             "- **\"safeguarding\"** or **\"Section 42\"**\n"
-            "- **\"Section 136\"** or **\"place of safety\"**\n"
+            "- **\"community treatment order\"** or **\"DoLS\"**\n"
         )
         return response
 
@@ -219,11 +261,12 @@ def _build_fallback(question: str, sections: list) -> str:
 
     return response
 
+
 async def section_lookup(act_name: str, section_input: str) -> str:
-    """Look up sections from cached legislation."""
+    """Look up sections from cached legislation, with live Lex API fallback."""
     legislation_id = NURSING_ACTS.get(act_name)
     if not legislation_id:
-        return f"❌ Act not found in NurseLex."
+        return "❌ Act not found in NurseLex."
 
     cache_query = f"{act_name} section {section_input}" if section_input.strip() else act_name
     sections = search_cached(cache_query, max_results=10)
@@ -231,13 +274,28 @@ async def section_lookup(act_name: str, section_input: str) -> str:
 
     if section_input.strip() and sections:
         try:
-            target_num = int(section_input.strip().replace("Section ", "").replace("s.", "").replace("S.", ""))
+            target_num = int(
+                section_input.strip()
+                .replace("Section ", "")
+                .replace("s.", "")
+                .replace("S.", "")
+            )
             matching = [s for s in sections if s.get("number") == target_num]
             if matching:
                 sections = matching
         except ValueError:
             pass
 
+    # Live Lex API fallback if cache misses
+    if not sections:
+        logger.info(f"Section lookup cache miss — trying live Lex API for {act_name} s.{section_input}")
+        try:
+            query = f"{act_name} section {section_input}" if section_input.strip() else act_name
+            live = await search_legislation_sections(query, legislation_id=legislation_id, size=5)
+            sections = live
+        except Exception as e:
+            logger.warning(f"Live Lex API lookup failed: {e}")
+
     if not sections:
         return (
             f"⏳ Section not found in cache for **{act_name}**.\n\n"
@@ -260,23 +318,20 @@ async def section_lookup(act_name: str, section_input: str) -> str:
     result += "\n🏛️ *Crown Copyright, OGL v3.0*"
     return result
 
+
 async def fetch_explanatory_note(act_name: str, section_input: str) -> str:
-    """Dynamically fetch Explanatory Notes from the i.AI Lex API."""
+    """Fetch Explanatory Notes from the i.AI Lex API."""
     if not section_input.strip():
         return "Please specify a section number to view its Explanatory Note."
-        
+
     try:
-        # Extract the digits
         section_number = "".join([c for c in section_input if c.isdigit()])
         if not section_number:
             return "Please enter a valid section number."
-            
-        url = 'https://lex.lab.i.ai.gov.uk/explanatory_note/section/search'
-        payload = {
-            'query': f'"{act_name}" Section {section_number}',
-            'limit': 5
-        }
-        
+
+        url = "https://lex.lab.i.ai.gov.uk/explanatory_note/section/search"
+        payload = {"query": f'"{act_name}" Section {section_number}', "limit": 5}
+
         async with httpx.AsyncClient() as client:
             r = await client.post(url, json=payload, timeout=10.0)
             if r.status_code == 200:
@@ -284,57 +339,70 @@ async def fetch_explanatory_note(act_name: str, section_input: str) -> str:
                 if isinstance(data, list):
                     parent_id = NURSING_ACTS.get(act_name, "")
                     for note in data:
-                        if parent_id and parent_id in note.get('legislation_id', ''):
-                            text = note.get('text', '')
+                        if parent_id and parent_id in note.get("legislation_id", ""):
+                            text = note.get("text", "")
                             if text:
                                 return f"### Official Explanatory Note\n\n{text}\n\n*Source: i.AI Lex API*"
-                                
-        return f"No official Explanatory Note found for {act_name} Section {section_number}.\n\n*(Note: Acts passed prior to 1999 generally do not have Explanatory Notes).*."
+
+        return (
+            f"No official Explanatory Note found for {act_name} Section {section_number}.\n\n"
+            "*(Acts passed prior to 1999 generally do not have Explanatory Notes.)*"
+        )
     except httpx.TimeoutException:
-        return "⏳ API Timeout while fetching Explanatory Note."
+        return "⏳ API timeout while fetching Explanatory Note."
     except Exception as e:
         return f"Error fetching note: {str(e)}"
 
+
 async def scenario_search(scenario_text: str) -> str:
-    """Use local i-dot-ai vector search to map a clinical scenario to legal sections."""
+    """Semantic search: map a clinical scenario to relevant UK law sections."""
     if not scenario_text.strip():
         return "Please describe a clinical scenario."
-        
+
     try:
         results = search_scenarios_locally(scenario_text, top_k=5)
-        
+
+        # Fallback to live Lex API if semantic search returns nothing
+        if not results:
+            logger.info("Scenario semantic search returned nothing — trying live Lex API...")
+            try:
+                live = await search_legislation_sections(scenario_text, size=5)
+                results = live
+            except Exception as e:
+                logger.warning(f"Live Lex API scenario fallback failed: {e}")
+
         if not results:
-            return "No matching legislation found for this scenario in the local cache."
-            
+            return "No matching legislation found for this scenario. Try the **Chat** tab for a broader search."
+
         result = f"## ⚖️ Probable Legislation Matches for:\n*{scenario_text}*\n\n"
-        
+
         for i, n in enumerate(results, 1):
             leg_id = n.get("legislation_id", "")
-            
-            # 1. Use the act_name from known mapping
+
             act_name = ""
             for known_id, known_name in REVERSE_ACTS.items():
                 if known_id in leg_id:
                     act_name = known_name
                     break
-            
-            # 2. Final fallback: extract from the legislation_id URL
             if not act_name:
                 act_name = leg_id.split("/id/")[-1] if "/id/" in leg_id else leg_id or "Legislation"
-                    
+
             sec_num = n.get("number", "??")
             title = n.get("title", "Untitled Section")
             text = n.get("text", "")
             uri = n.get("uri", f"https://www.legislation.gov.uk/id/{leg_id}/section/{sec_num}")
             score = n.get("score", 0.0)
-            
-            result += f"### {i}. {act_name} — Section {sec_num}: {title} (Match Score: {score:.2f})\n"
-            result += f"{text[:800]}...\n\n"
+
+            result += f"### {i}. {act_name} — Section {sec_num}: {title}"
+            if score:
+                result += f" *(Match: {score:.2f})*"
+            result += f"\n{text[:800]}...\n\n"
             result += f"🔗 [Read full text on legislation.gov.uk]({uri})\n\n---\n\n"
-            
+
         return result
     except Exception as e:
-        return f"Error during local scenario search: {str(e)}"
+        return f"Error during scenario search: {str(e)}"
+
 
 def browse_legislation(search_term: str, act_type: str) -> str:
     """Browse the legislation index from the parquet file."""
@@ -344,7 +412,13 @@ def browse_legislation(search_term: str, act_type: str) -> str:
     filtered = LEG_DF.copy()
 
     if act_type != "All":
-        type_map = {"Primary Acts": "ukpga", "Statutory Instruments": "uksi", "Scottish SIs": "ssi", "NI SRs": "nisr", "Welsh SIs": "wsi"}
+        type_map = {
+            "Primary Acts": "ukpga",
+            "Statutory Instruments": "uksi",
+            "Scottish SIs": "ssi",
+            "NI SRs": "nisr",
+            "Welsh SIs": "wsi",
+        }
         if act_type in type_map:
             filtered = filtered[filtered["type"] == type_map[act_type]]
 
@@ -357,7 +431,6 @@ def browse_legislation(search_term: str, act_type: str) -> str:
         return f"No legislation found matching '{search_term}'."
 
     result = f"## 📖 Legislation Index ({len(filtered)} results)\n\n| Year | Title | Type |\n|---|---|---|\n"
-
     for _, row in filtered.iterrows():
         year = row.get("year", "—")
         title = row.get("title", "Untitled")
@@ -366,11 +439,13 @@ def browse_legislation(search_term: str, act_type: str) -> str:
         title_link = f"[{title}]({uri})" if uri else title
         result += f"| {year} | {title_link} | {leg_type} |\n"
 
-    result += f"\n\n*Showing top 50 of {len(LEG_DF)} health & social care entries — {len(LEG_DF[LEG_DF['type']=='ukpga'])} Primary Acts*"
+    result += f"\n\n*Showing top 50 of {len(LEG_DF)} health & social care entries — {len(LEG_DF[LEG_DF['type'] == 'ukpga'])} Primary Acts*"
     result += "\n\n🏛️ *Data from i.AI Lex bulk downloads — Crown Copyright, OGL v3.0*"
     return result
 
+
 # --- Gradio UI ---
+
 THEME = gr.themes.Soft(
     primary_hue="indigo",
     secondary_hue="violet",
@@ -417,17 +492,39 @@ with gr.Blocks(theme=THEME, css=CSS, title="NurseLex — UK Law for All Nurses")
     </div>
     """)
 
+    # --- API Key sidebar ---
+    with gr.Accordion("🔑 Anthropic API Key (for AI explanations)", open=False):
+        gr.Markdown(
+            "Enter your own [Anthropic API key](https://console.anthropic.com) to enable AI-powered plain English explanations "
+            "via **Claude**. Your key is used only for your session and is never stored."
+        )
+        api_key_input = gr.Textbox(
+            label="API Key",
+            placeholder="sk-ant-...",
+            type="password",
+            show_label=False,
+        )
+        gr.Markdown(
+            "*Without a key, NurseLex still shows the raw statutory text from legislation.gov.uk.*"
+        )
+
     with gr.Tabs():
         # --- Tab 1: Chat ---
         with gr.TabItem("💬 Ask a Legal Question", id="chat"):
-            gr.Markdown("Ask about UK healthcare legislation — answers are grounded in **real statutory text**. (Cache: 1,128 Sections + 219K Acts)")
+            gr.Markdown(
+                "Ask about UK healthcare legislation — answers are grounded in **real statutory text** "
+                "from three sources: cached sections, semantic vector search, and the live i.AI Lex API."
+            )
 
             chatbot = gr.Chatbot(
                 label="NurseLex",
                 height=480,
                 type="messages",
                 show_copy_button=True,
-                avatar_images=(None, "https://em-content.zobj.net/source/twitter/376/classical-building_1f3db-fe0f.png"),
+                avatar_images=(
+                    None,
+                    "https://em-content.zobj.net/source/twitter/376/classical-building_1f3db-fe0f.png",
+                ),
             )
             msg = gr.Textbox(
                 label="Your question",
@@ -454,19 +551,23 @@ with gr.Blocks(theme=THEME, css=CSS, title="NurseLex — UK Law for All Nurses")
                         variant="secondary",
                     ).click(fn=lambda q=QUICK_QUESTIONS[i]: q, outputs=msg)
 
-            async def respond(message, history):
+            async def respond(message, history, api_key):
                 history = history or []
                 history.append({"role": "user", "content": message})
-                answer = await query_and_respond(message, history)
+                answer = await query_and_respond(message, history, api_key)
                 history.append({"role": "assistant", "content": answer})
                 return "", history
 
-            submit_btn.click(respond, [msg, chatbot], [msg, chatbot])
-            msg.submit(respond, [msg, chatbot], [msg, chatbot])
+            submit_btn.click(respond, [msg, chatbot, api_key_input], [msg, chatbot])
+            msg.submit(respond, [msg, chatbot, api_key_input], [msg, chatbot])
 
         # --- Tab 2: Section Lookup ---
         with gr.TabItem("📖 Section Lookup", id="lookup"):
-            gr.Markdown("Look up a **specific section** of key nursing Acts. Includes **Official Explanatory Notes** where available.")
+            gr.Markdown(
+                "Look up a **specific section** of key nursing Acts. "
+                "Includes **Official Explanatory Notes** where available. "
+                "Falls back to the live Lex API when not in local cache."
+            )
 
             with gr.Row():
                 act_dropdown = gr.Dropdown(
@@ -476,11 +577,11 @@ with gr.Blocks(theme=THEME, css=CSS, title="NurseLex — UK Law for All Nurses")
                 )
                 section_input_box = gr.Textbox(
                     label="Section number",
-                    placeholder="e.g., 5 or 117 or 136",
+                    placeholder="e.g., 5 or 117 or 136 or 17A",
                 )
 
             lookup_btn = gr.Button("🔍 Look Up Law & Notes", variant="primary")
-            
+
             with gr.Row():
                 lookup_output = gr.Markdown(label="Statutory Text")
                 note_output = gr.Markdown(label="Official Explanatory Note")
@@ -490,23 +591,29 @@ with gr.Blocks(theme=THEME, css=CSS, title="NurseLex — UK Law for All Nurses")
 
         # --- Tab 3: Scenario Matcher ---
         with gr.TabItem("🧠 Scenario Matcher", id="scenario"):
-            gr.Markdown("Describe a clinical scenario in plain English, and the **Lex Vector Search Engine** will map it to the most relevant UK laws.")
-            
+            gr.Markdown(
+                "Describe a clinical scenario in plain English. "
+                "The **semantic vector search engine** maps it to the most relevant UK laws, "
+                "with a live Lex API fallback for scenarios not in local cache."
+            )
+
             with gr.Row():
                 scenario_input = gr.Textbox(
                     label="Clinical Scenario",
                     placeholder="e.g. 'Patient wants to leave the ward but lacks capacity' or 'Doctor orders restraint without DoLS'",
-                    lines=3
+                    lines=3,
                 )
-            
+
             scenario_btn = gr.Button("🤖 Find Relevant Law", variant="primary")
             scenario_output = gr.Markdown(label="Semantic Search Results")
-            
+
             scenario_btn.click(scenario_search, [scenario_input], scenario_output)
 
         # --- Tab 4: Browse Legislation ---
         with gr.TabItem("📚 Browse Legislation", id="browse"):
-            gr.Markdown(f"Browse **219,678** health & social care Acts and Statutory Instruments from the i.AI Lex dataset.")
+            gr.Markdown(
+                "Browse **219,678** health & social care Acts and Statutory Instruments from the i.AI Lex dataset."
+            )
 
             with gr.Row():
                 browse_search = gr.Textbox(
@@ -524,9 +631,9 @@ with gr.Blocks(theme=THEME, css=CSS, title="NurseLex — UK Law for All Nurses")
 
             browse_btn.click(browse_legislation, [browse_search, browse_type], browse_output)
 
-        # --- Tab 4: About ---
+        # --- Tab 5: About ---
         with gr.TabItem("ℹ️ About", id="about"):
-            gr.Markdown(f"""
+            gr.Markdown("""
 ## About NurseLex
 
 **NurseLex** is a universal legal literacy tool for **all nurses and nursing students**.
@@ -534,23 +641,33 @@ with gr.Blocks(theme=THEME, css=CSS, title="NurseLex — UK Law for All Nurses")
 ### How It Works
 
 1. **You ask a question** about UK healthcare law
-2. **Cached legislation** provides the actual statutory text instantly
-3. **Gemini Flash** explains it in plain English with practical nursing implications
-4. **Every answer cites** the specific Act, section, and year
+2. **Three search layers** retrieve actual statutory text:
+   - Local cached sections (1,128 sections — instant)
+   - Semantic vector search using an i.AI fine-tuned MiniLM model
+   - Live i.AI Lex API fallback for cache misses
+3. **Claude** (your own API key) explains it in plain English with practical nursing implications
+4. **Every answer cites** the specific Act, section, and year with a direct link to legislation.gov.uk
 
 ### Data
 
 - **219,678 legislation entries** from the [i.AI Lex](https://lex.lab.i.ai.gov.uk/) bulk dataset
-- **1,128 key sections** pre-cached with full text (MHA 1983, MCA 2005, Care Act 2014)
+- **1,128 key sections** pre-cached with full text (MHA 1983, MCA 2005, Care Act 2014, and more)
 - **Crown Copyright** — Open Government Licence v3.0
 
 ### Key Acts Covered
 
 | Act | Key Sections | Nursing Relevance |
 |---|---|---|
-| Mental Health Act 1983 | S.2, S.3, S.4, S.5(2), S.5(4), S.17, S.117, S.135, S.136 | Detention, holding powers, leave, aftercare |
-| Mental Capacity Act 2005 | S.1 (Principles), S.2-3 (Capacity), S.4 (Best Interests), S.5 | Capacity assessments, best interests, DoLS |
-| Care Act 2014 | S.42 (Safeguarding), S.67 (Advocacy) | Safeguarding adults, independent advocates |
+| Mental Health Act 1983 | S.2, S.3, S.4, S.5(2), S.5(4), S.17, S.17A, S.26, S.114, S.117, S.135, S.136 | Detention, holding powers, CTOs, AMHP, aftercare |
+| Mental Health Act 2007 | Amends MHA 1983 — CTOs, AMHP role, Supervised Community Treatment | Community treatment, amended criteria |
+| Mental Capacity Act 2005 | S.1–5 (Principles/Capacity), S.4A (DoLS), S.9 (LPA), S.45 (CoP) | Capacity, best interests, DoLS, LPA |
+| Mental Capacity (Amendment) Act 2019 | Liberty Protection Safeguards | Replaces DoLS (implementation pending) |
+| Care Act 2014 | S.9 (Assessment), S.13 (Eligibility), S.42 (Safeguarding), S.67 (Advocacy) | Safeguarding adults, advocacy, care assessment |
+| Health and Social Care Act 2008 | S.20 (Duty of Candour) | Notifiable safety incidents, openness |
+| Misuse of Drugs Act 1971 | S.2 (Schedules), S.4 (Supply) | Controlled drug management |
+| Data Protection Act 2018 | S.1, S.45 (Subject access) | Patient data, records, consent |
+| Equality Act 2010 | S.4 (Characteristics), S.6 (Disability), S.20 (Adjustments) | Reasonable adjustments, discrimination |
+| Human Rights Act 1998 | Article 5 (Liberty), Article 8 (Privacy) | Detention rights, confidentiality |
 
 ### Built By
 
@@ -565,4 +682,3 @@ app.queue()
 
 if __name__ == "__main__":
     app.launch(server_name="0.0.0.0", server_port=7860)
-

cached_legislation.py

@@ -25,13 +25,24 @@ except Exception as e:
     logger.error(f"Error loading {JSON_PATH}: {e}")
 
 # --- Keyword Map for Natural Language Shortcuts ---
-# Direct section links for common nursing search terms
+# Maps common nursing search terms to (legislation_id, section_number).
+# CTOs and AMHP were inserted into MHA 1983 by MHA 2007, so their
+# legislation_id remains ukpga/1983/20.
 KEYWORD_MAP = {
-    "nurse holding power": ("ukpga/1983/20", "5(4)"),
-    "doctor holding power": ("ukpga/1983/20", "5(2)"),
+    # MHA 1983 — holding powers
+    "nurse holding power": ("ukpga/1983/20", "5"),
+    "doctor holding power": ("ukpga/1983/20", "5"),
+    "section 5": ("ukpga/1983/20", "5"),
+
+    # MHA 1983 — admission / detention
     "section 2": ("ukpga/1983/20", "2"),
+    "admission for assessment": ("ukpga/1983/20", "2"),
     "section 3": ("ukpga/1983/20", "3"),
+    "admission for treatment": ("ukpga/1983/20", "3"),
     "section 4": ("ukpga/1983/20", "4"),
+    "emergency admission": ("ukpga/1983/20", "4"),
+
+    # MHA 1983 — leave, aftercare, place of safety
     "aftercare": ("ukpga/1983/20", "117"),
     "section 117": ("ukpga/1983/20", "117"),
     "leave of absence": ("ukpga/1983/20", "17"),
@@ -39,77 +50,139 @@ KEYWORD_MAP = {
     "place of safety": ("ukpga/1983/20", "136"),
     "section 136": ("ukpga/1983/20", "136"),
     "section 135": ("ukpga/1983/20", "135"),
-    "best interests": ("ukpga/2005/9", "4"),
+    "warrant to search": ("ukpga/1983/20", "135"),
+
+    # MHA 1983 — community treatment orders (inserted by MHA 2007, s.32)
+    "community treatment order": ("ukpga/1983/20", "17a"),
+    "cto": ("ukpga/1983/20", "17a"),
+    "supervised community treatment": ("ukpga/1983/20", "17a"),
+    "conditional discharge": ("ukpga/1983/20", "17a"),
+    "recall": ("ukpga/1983/20", "17e"),
+
+    # MHA 1983 — roles (AMHP inserted by MHA 2007, s.21; nearest relative s.26)
+    "nearest relative": ("ukpga/1983/20", "26"),
+    "approved mental health professional": ("ukpga/1983/20", "114"),
+    "amhp": ("ukpga/1983/20", "114"),
+    "responsible clinician": ("ukpga/1983/20", "34"),
+    "rc": ("ukpga/1983/20", "34"),
+
+    # MCA 2005 — capacity principles
+    "mca principles": ("ukpga/2005/9", "1"),
     "capacity test": ("ukpga/2005/9", "3"),
     "functional test": ("ukpga/2005/9", "3"),
-    "mca principles": ("ukpga/2005/9", "1"),
+    "best interests": ("ukpga/2005/9", "4"),
+    "lasting power of attorney": ("ukpga/2005/9", "9"),
+    "lpa": ("ukpga/2005/9", "9"),
+    "court of protection": ("ukpga/2005/9", "45"),
+
+    # MCA 2005 — DoLS (Schedule A1, accessed via s.4a)
+    "deprivation of liberty": ("ukpga/2005/9", "4a"),
+    "dols": ("ukpga/2005/9", "4a"),
+    "deprivation of liberty safeguards": ("ukpga/2005/9", "4a"),
+    "standard authorisation": ("ukpga/2005/9", "4a"),
+    "urgent authorisation": ("ukpga/2005/9", "4a"),
+
+    # Mental Capacity (Amendment) Act 2019 — LPS
+    "liberty protection safeguards": ("ukpga/2019/17", "1"),
+    "lps": ("ukpga/2019/17", "1"),
+
+    # Care Act 2014
     "safeguarding": ("ukpga/2014/23", "42"),
     "section 42": ("ukpga/2014/23", "42"),
+    "safeguarding adult review": ("ukpga/2014/23", "44"),
     "advocacy": ("ukpga/2014/23", "67"),
+    "care needs assessment": ("ukpga/2014/23", "9"),
+    "eligible needs": ("ukpga/2014/23", "13"),
+
+    # Health and Social Care Act 2008 — Duty of Candour
+    "duty of candour": ("ukpga/2008/14", "20"),
+    "candour": ("ukpga/2008/14", "20"),
+    "notifiable safety incident": ("ukpga/2008/14", "20"),
+
+    # Data Protection Act 2018 / UK GDPR
+    "data protection": ("ukpga/2018/12", "1"),
+    "gdpr": ("ukpga/2018/12", "1"),
+    "patient data": ("ukpga/2018/12", "1"),
+    "subject access request": ("ukpga/2018/12", "45"),
+
+    # Misuse of Drugs Act 1971 — controlled drugs
+    "controlled drugs": ("ukpga/1971/38", "2"),
+    "schedule 2": ("ukpga/1971/38", "2"),
+    "cd cupboard": ("ukpga/1971/38", "2"),
+    "misuse of drugs": ("ukpga/1971/38", "4"),
+
+    # Equality Act 2010
+    "reasonable adjustments": ("ukpga/2010/15", "20"),
+    "disability discrimination": ("ukpga/2010/15", "6"),
+    "protected characteristic": ("ukpga/2010/15", "4"),
+
+    # Human Rights Act 1998
+    "article 5": ("ukpga/1998/42", "1"),
+    "right to liberty": ("ukpga/1998/42", "1"),
+    "article 8": ("ukpga/1998/42", "1"),
+    "right to privacy": ("ukpga/1998/42", "1"),
+    "human rights": ("ukpga/1998/42", "1"),
 }
 
+
 def search_cached(query: str, max_results: int = 5) -> list:
     """
-    Search local sections by keyword, title, or legislation ID.
+    Search local sections by keyword map, title match, or text content.
     Returns a list of section dictionaries.
     """
     if not query:
         return []
-    
+
     query = query.lower().strip()
     results = []
-    
-    # 1. Check Keyword Map First (High Precision)
+
+    # 1. Keyword Map — high precision shortcuts
     for kw, (leg_id, sec_num) in KEYWORD_MAP.items():
         if kw in query:
-            # Find the specific section in our list
             for s in CACHED_SECTIONS:
                 if s.get("legislation_id") == leg_id and str(s.get("number")) == sec_num:
                     if s not in results:
                         results.append(s)
-            # Find closest matches if exact number not found (e.g. 5(4) vs 5)
+            # Closest match if exact section number not cached (e.g. 5(4) vs 5)
             if not results:
-                 for s in CACHED_SECTIONS:
-                    if s.get("legislation_id") == leg_id and str(s.get("number")).startswith(sec_num.split('(')[0]):
+                for s in CACHED_SECTIONS:
+                    if s.get("legislation_id") == leg_id and str(s.get("number")).startswith(sec_num.split("(")[0]):
                         if s not in results:
-                             results.append(s)
+                            results.append(s)
 
-    # 2. Text-based Search
-    # Sort sections by relevance (title match > text match)
+    # 2. Scored text search — title and content
     scored_results = []
     for s in CACHED_SECTIONS:
         score = 0
         title = s.get("title", "").lower()
         text = s.get("text", "").lower()
-        leg_id = s.get("legislation_id", "").lower()
         num = str(s.get("number", "")).lower()
-        
-        # Exact section reference (e.g. "Section 5")
+
+        # Exact section reference
         if f"section {num}" in query or f"s.{num}" in query or f"s {num}" in query:
-             score += 100
-             
-        # Title matches
+            score += 100
+
+        # Full query in title
         if query in title:
             score += 50
-        
-        # Word matches in title
+
+        # Individual words in title
         for word in query.split():
             if len(word) > 3 and word in title:
                 score += 10
-        
-        # Content matches
+
+        # Full query in text
         if query in text:
             score += 5
-            
+
         if score > 0:
             scored_results.append((score, s))
-            
-    # Sort and add to results
+
     scored_results.sort(key=lambda x: x[0], reverse=True)
     for _, s in scored_results:
         if s not in results:
             results.append(s)
             if len(results) >= max_results:
                 break
-                
+
     return results[:max_results]

lex_client.py

@@ -14,17 +14,22 @@ LEX_TIMEOUT = 60.0  # Lex API can be slow for semantic search
 # Key legislation IDs for mental health & learning disability nursing
 NURSING_LEGISLATION = {
     "Mental Health Act 1983": "ukpga/1983/20",
+    "Mental Health Act 2007": "ukpga/2007/12",
     "Mental Capacity Act 2005": "ukpga/2005/9",
+    "Mental Capacity (Amendment) Act 2019": "ukpga/2019/17",
     "Care Act 2014": "ukpga/2014/23",
     "Human Rights Act 1998": "ukpga/1998/42",
     "Equality Act 2010": "ukpga/2010/15",
     "Health and Social Care Act 2012": "ukpga/2012/7",
+    "Health and Social Care Act 2008": "ukpga/2008/14",
     "Children Act 1989": "ukpga/1989/41",
     "Children Act 2004": "ukpga/2004/31",
     "Safeguarding Vulnerable Groups Act 2006": "ukpga/2006/47",
     "Mental Health Units (Use of Force) Act 2018": "ukpga/2018/27",
     "Health and Care Act 2022": "ukpga/2022/31",
     "Autism Act 2009": "ukpga/2009/15",
+    "Misuse of Drugs Act 1971": "ukpga/1971/38",
+    "Data Protection Act 2018": "ukpga/2018/12",
 }
 
 
@@ -188,195 +193,3 @@ def format_sections_for_context(sections: list[dict], max_chars: int = 6000) ->
         total_chars += len(entry)
 
     return "".join(context_parts) if context_parts else "No relevant legislation sections found."
-=======
-"""
-NurseLex — Lex API Client
-Wraps the i.AI Lex API for nursing-focused UK legislation search.
-"""
-import httpx
-import logging
-from typing import Optional
-
-logger = logging.getLogger(__name__)
-
-LEX_API_BASE = "https://lex.lab.i.ai.gov.uk"
-LEX_TIMEOUT = 60.0  # Lex API can be slow for semantic search
-
-# Key legislation IDs for mental health & learning disability nursing
-NURSING_LEGISLATION = {
-    "Mental Health Act 1983": "ukpga/1983/20",
-    "Mental Capacity Act 2005": "ukpga/2005/9",
-    "Care Act 2014": "ukpga/2014/23",
-    "Human Rights Act 1998": "ukpga/1998/42",
-    "Equality Act 2010": "ukpga/2010/15",
-    "Health and Social Care Act 2012": "ukpga/2012/7",
-    "Children Act 1989": "ukpga/1989/41",
-    "Children Act 2004": "ukpga/2004/31",
-    "Safeguarding Vulnerable Groups Act 2006": "ukpga/2006/47",
-    "Mental Health Units (Use of Force) Act 2018": "ukpga/2018/27",
-    "Health and Care Act 2022": "ukpga/2022/31",
-    "Autism Act 2009": "ukpga/2009/15",
-}
-
-
-async def _post(endpoint: str, payload: dict) -> dict | list:
-    """Make a POST request to the Lex API with retry logic."""
-    url = f"{LEX_API_BASE}{endpoint}"
-    for attempt in range(3):
-        try:
-            async with httpx.AsyncClient(timeout=LEX_TIMEOUT) as client:
-                resp = await client.post(url, json=payload)
-                resp.raise_for_status()
-                return resp.json()
-        except httpx.TimeoutException:
-            logger.warning(f"Lex API timeout (attempt {attempt + 1}/3): {endpoint}")
-            if attempt == 2:
-                raise
-        except httpx.HTTPStatusError as e:
-            logger.error(f"Lex API error {e.response.status_code}: {endpoint}")
-            raise
-    return []
-
-
-async def search_legislation_sections(
-    query: str,
-    legislation_id: Optional[str] = None,
-    size: int = 5,
-) -> list[dict]:
-    """Semantic search across legislation sections."""
-    payload = {
-        "query": query,
-        "size": size,
-        "include_text": True,
-    }
-    if legislation_id:
-        payload["legislation_id"] = legislation_id
-
-    try:
-        return await _post("/legislation/section/search", payload)
-    except Exception as e:
-        logger.error(f"Section search failed: {e}")
-        return []
-
-
-async def search_legislation_acts(
-    query: str,
-    limit: int = 5,
-) -> dict:
-    """Search for Acts and Statutory Instruments."""
-    payload = {
-        "query": query,
-        "limit": limit,
-        "include_text": True,
-    }
-
-    try:
-        return await _post("/legislation/search", payload)
-    except Exception as e:
-        logger.error(f"Act search failed: {e}")
-        return {"results": [], "total": 0, "offset": 0, "limit": limit}
-
-
-async def lookup_legislation(legislation_id: str) -> dict:
-    """Look up a specific Act by its ID (e.g., 'ukpga/1983/20')."""
-    parts = legislation_id.split("/")
-    payload = {
-        "legislation_type": parts[0],
-        "year": int(parts[1]),
-        "number": int(parts[2]),
-    }
-
-    return await _post("/legislation/lookup", payload)
-
-
-async def get_legislation_full_text(
-    legislation_id: str,
-    include_schedules: bool = False,
-) -> dict:
-    """Get the full text of a piece of legislation."""
-    payload = {
-        "legislation_id": legislation_id,
-        "include_schedules": include_schedules,
-    }
-
-    return await _post("/legislation/text", payload)
-
-
-async def get_sections_for_legislation(
-    legislation_id: str,
-    limit: int = 200,
-) -> list[dict]:
-    """Get all sections for a specific piece of legislation."""
-    payload = {
-        "legislation_id": legislation_id,
-        "limit": limit,
-    }
-
-    try:
-        return await _post("/legislation/section/lookup", payload)
-    except Exception as e:
-        logger.error(f"Section lookup failed: {e}")
-        return []
-
-
-async def search_explanatory_notes(
-    query: str,
-    legislation_id: Optional[str] = None,
-    size: int = 5,
-) -> list[dict]:
-    """Search explanatory notes for legislation."""
-    payload = {
-        "query": query,
-        "size": size,
-    }
-    if legislation_id:
-        payload["legislation_id"] = legislation_id
-
-    try:
-        return await _post("/explanatory_note/section/search", payload)
-    except Exception as e:
-        logger.error(f"Explanatory note search failed: {e}")
-        return []
-
-
-async def search_amendments(
-    legislation_id: str,
-    search_amended: bool = True,
-    size: int = 20,
-) -> list[dict]:
-    """Search for amendments to or by a piece of legislation."""
-    payload = {
-        "legislation_id": legislation_id,
-        "search_amended": search_amended,
-        "size": size,
-    }
-
-    try:
-        return await _post("/amendment/search", payload)
-    except Exception as e:
-        logger.error(f"Amendment search failed: {e}")
-        return []
-
-
-def format_sections_for_context(sections: list[dict], max_chars: int = 6000) -> str:
-    """Format legislation sections into a readable context string for the LLM."""
-    context_parts = []
-    total_chars = 0
-
-    for section in sections:
-        title = section.get("title", "Untitled")
-        text = section.get("text", "")
-        leg_id = section.get("legislation_id", "")
-        section_num = section.get("number", "")
-
-        entry = f"### {title}\n"
-        entry += f"**Source:** {leg_id}, Section {section_num}\n\n"
-        entry += f"{text}\n\n---\n\n"
-
-        if total_chars + len(entry) > max_chars:
-            break
-        context_parts.append(entry)
-        total_chars += len(entry)
-
-    return "".join(context_parts) if context_parts else "No relevant legislation sections found."
->>>>>>> a4e257b16d56f80612b7c9ac6d2e7c198fef5bb6

local_search.py

@@ -3,110 +3,119 @@ local_search.py
 
 Locally loads the `i-dot-ai/all-miniLM-L6-v2-UKPGA-6k-finetune` model to encode
 the cached `nursing_sections.json` into semantic embeddings for fast, reliable local searches.
+
+Uses numpy for cosine similarity — no torch dependency required.
+Initialises at module load time so the first user request is instant.
 """
 import json
 import os
 import logging
-import torch
 import numpy as np
-from sentence_transformers import SentenceTransformer, util
+from sentence_transformers import SentenceTransformer
 
 logger = logging.getLogger(__name__)
 
-# Constants
 MODEL_NAME = "i-dot-ai/all-miniLM-L6-v2-UKPGA-6k-finetune"
 CACHE_FILE = os.path.join(os.path.dirname(__file__), "nursing_sections.json")
 EMBEDDINGS_FILE = os.path.join(os.path.dirname(__file__), "nursing_sections_embeddings.npy")
 
-# Global variables to hold the model and embeddings in memory
 _model = None
-_corpus_embeddings = None
+_corpus_embeddings = None  # shape: (N, D) numpy array
 _sections = []
 
+
 def init_local_search():
-    """Initializes the model and computes embeddings for all cached sections."""
+    """Initialises the model and loads/computes embeddings for all cached sections."""
     global _model, _corpus_embeddings, _sections
-    
+
     if _model is not None:
-        return  # Already initialized
-        
+        return  # Already initialised
+
     try:
         logger.info(f"Loading local embedding model: {MODEL_NAME}...")
         _model = SentenceTransformer(MODEL_NAME)
-        
+
         if not os.path.exists(CACHE_FILE):
             logger.error(f"Cache file not found at {CACHE_FILE}")
             return
-            
+
         with open(CACHE_FILE, "r", encoding="utf-8") as f:
             _sections = json.load(f)
-            
+
         if not _sections:
             logger.warning("No sections found in cache.")
             return
 
         if os.path.exists(EMBEDDINGS_FILE):
-            logger.info("Loading precomputed numpy embeddings from disk (Instant)...")
-            np_embeddings = np.load(EMBEDDINGS_FILE)
-            # Convert back to tensor for cosine similarity
-            _corpus_embeddings = torch.from_numpy(np_embeddings)
-            logger.info("Local semantic search engine ready.")
+            logger.info("Loading precomputed numpy embeddings from disk (instant)...")
+            _corpus_embeddings = np.load(EMBEDDINGS_FILE)
+            logger.info(f"Local semantic search engine ready ({len(_sections)} sections).")
             return
 
-        logger.info(f"Computing embeddings for {len(_sections)} cached sections. This may take a minute on first run...")
-        # Prepare text for embedding: combine legislation title, section title, and text
+        logger.info(f"Computing embeddings for {len(_sections)} cached sections...")
         corpus_texts = []
         for s in _sections:
-            # Reconstruct the act name roughly from the URL to give the model context
             leg_id = s.get("legislation_id", "")
             act_name = leg_id.split("/")[-2] if "/" in leg_id else leg_id
-            
-            # Create a rich text representation for the vector search
             content = f"Act: {act_name}. Section {s.get('number', '')}: {s.get('title', '')}. {s.get('text', '')}"
             corpus_texts.append(content)
-            
-        # Encode all sections
-        _corpus_embeddings = _model.encode(corpus_texts, convert_to_tensor=True, show_progress_bar=False)
-        logger.info("Saving computed numpy embeddings for future use...")
+
+        embeddings = _model.encode(corpus_texts, convert_to_tensor=False, show_progress_bar=False)
+        _corpus_embeddings = np.array(embeddings, dtype=np.float32)
+
+        logger.info("Saving computed embeddings to disk for future use...")
         try:
-            np.save(EMBEDDINGS_FILE, _corpus_embeddings.cpu().numpy())
+            np.save(EMBEDDINGS_FILE, _corpus_embeddings)
         except Exception as save_err:
             logger.warning(f"Failed to save embeddings cache: {save_err}")
-            
-        logger.info("Local semantic search engine ready.")
-        
+
+        logger.info(f"Local semantic search engine ready ({len(_sections)} sections).")
+
     except Exception as e:
-        logger.error(f"Failed to initialize local search engine: {e}")
-        _model = None  # Reset on failure
+        logger.error(f"Failed to initialise local search engine: {e}")
+        _model = None
+
+
+def _cosine_similarity(query_emb: np.ndarray, corpus_emb: np.ndarray) -> np.ndarray:
+    """Cosine similarity between a single query vector and a corpus matrix."""
+    query_norm = query_emb / (np.linalg.norm(query_emb) + 1e-10)
+    corpus_norms = np.linalg.norm(corpus_emb, axis=1, keepdims=True)
+    corpus_normalized = corpus_emb / np.where(corpus_norms == 0, 1.0, corpus_norms)
+    return corpus_normalized @ query_norm
+
 
 def search_scenarios_locally(query: str, top_k: int = 5) -> list[dict]:
     """Semantic search over the local cached sections using cosine similarity."""
-    global _model, _corpus_embeddings, _sections
-    
     if _model is None or _corpus_embeddings is None:
+        logger.warning("Local search engine not ready — attempting re-init.")
         init_local_search()
-        
+
     if _model is None or _corpus_embeddings is None:
-        logger.error("Local search engine is unavailable.")
+        logger.error("Local search engine unavailable.")
         return []
-        
+
     try:
-        query_embedding = _model.encode(query, convert_to_tensor=True)
-        # Compute cosine similarities
-        cos_scores = util.cos_sim(query_embedding, _corpus_embeddings)[0]
-        
-        # Find the top_k scores
-        top_results = torch.topk(cos_scores, k=min(top_k, len(_sections)))
-        
+        query_embedding = np.array(
+            _model.encode(query, convert_to_tensor=False), dtype=np.float32
+        )
+        scores = _cosine_similarity(query_embedding, _corpus_embeddings)
+
+        top_indices = np.argsort(scores)[::-1][:top_k]
+
         results = []
-        for score, idx in zip(top_results[0], top_results[1]):
-            # Only return highly relevant matches (tune this threshold if needed)
-            if score.item() > 0.4:
+        for idx in top_indices:
+            score = float(scores[idx])
+            if score > 0.4:
                 match = _sections[idx].copy()
-                match["score"] = score.item()
+                match["score"] = score
                 results.append(match)
-                
+
         return results
+
     except Exception as e:
         logger.error(f"Error during local scenario search: {e}")
         return []
+
+
+# Initialise at module load time so first request is fast
+init_local_search()

requirements.txt

@@ -3,4 +3,5 @@ httpx>=0.27
 pandas>=2.0
 pyarrow>=14.0
 sentence-transformers>=2.7.0
-torch>=2.2.0
+numpy>=1.24
+anthropic>=0.40