mcp_server.py

import asyncio
import httpx
import pandas as pd
import json
import os
from mcp.server.fastmcp import FastMCP

# Initialize FastMCP server
mcp = FastMCP("NurseLex-LexAPI")

# Core Constants
BASE_URL = 'https://lex.lab.i.ai.gov.uk'
NURSING_ACTS = {
    "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",
    "Mental Health Units (Use of Force) Act 2018": "ukpga/2018/27",
    "Autism Act 2009": "ukpga/2009/15",
    "Children Act 1989": "ukpga/1989/41",
    "Children Act 2004": "ukpga/2004/31",
    "Safeguarding Vulnerable Groups Act 2006": "ukpga/2006/47",
    "Health and Care Act 2022": "ukpga/2022/31",
}
REVERSE_ACTS = {v: k for k, v in NURSING_ACTS.items()}

# Load Cache (we need absolute paths since MCP might run from a different CWD)
DB_DIR = os.path.dirname(os.path.abspath(__file__))
CACHE_FILE = os.path.join(DB_DIR, "nursing_sections.json")

def _load_sections():
    try:
        with open(CACHE_FILE, 'r', encoding='utf-8') as f:
            return json.load(f)
    except FileNotFoundError:
        return []

SECTIONS_CACHE = _load_sections()

@mcp.tool()
def search_local_nursing_cache(query: str, limit: int = 5) -> str:
    """
    Search the local, curated cache of 1,128 critical nursing legislation sections 
    (from the Mental Health Act, Care Act, etc.) for a specific keyword or section number.
    Returns the exact statutory text.
    """
    if not SECTIONS_CACHE:
        return "Error: Local cache not found."
        
    query_lower = query.lower()
    results = []
    
    for section in SECTIONS_CACHE:
        act_name = section.get('act_name', '').lower()
        title = section.get('title', '').lower()
        text = section.get('text', '').lower()
        num_str = str(section.get('number', ''))
        
        score = 0
        if query_lower in act_name: score += 1
        if query_lower in title: score += 3
        if query_lower in text: score += 1
        if query_lower == f"section {num_str}" or query_lower == num_str: score += 5
        
        if score > 0:
            results.append((score, section))
            
    # Sort and take top matches
    results.sort(key=lambda x: x[0], reverse=True)
    
    if not results:
        return "No sections found matching the query in the local cache."
        
    out = f"## 📚 Local Cache Results for '{query}'\n\n"
    for r in results[:limit]:
        sec = r[1]
        out += f"**{sec.get('act_name')} — Section {sec.get('number')}**\n"
        out += f"*{sec.get('title')}*\n"
        out += f"{sec.get('text')}\n\n---\n\n"
        
    return out

@mcp.tool()
async def vector_search_lex_api(clinical_scenario: str) -> str:
    """
    Translates a plain-English clinical scenario (e.g. 'Patient wants to leave but lacks capacity') 
    into relevant UK legislation by querying the i.AI Lex API semantic vector search engine.
    This searches across the entire legislation database, not just the local cache.
    """
    url = f'{BASE_URL}/legislation/section/search'
    payload = {
        'query': clinical_scenario,
        'limit': 5
    }
    
    try:
        async with httpx.AsyncClient() as client:
            r = await client.post(url, json=payload, timeout=15.0)
            if r.status_code != 200:
                return f"Lex API Vector Search Failed: Status Code {r.status_code}"
                
            data = r.json()
            if not isinstance(data, list) or not data:
                return "No semantic matches found for this scenario."
                
            out = f"## ⚖️ Vector Matches for Scenario:\n*{clinical_scenario}*\n\n"
            for i, n in enumerate(data, 1):
                leg_id = n.get("legislation_id", "")
                
                # 1. Use the act_name from the API response if available
                act_name = n.get("act_name", "")
                
                # 2. If not, try our known mapping
                if not act_name:
                    for known_id, known_name in REVERSE_ACTS.items():
                        if known_id in leg_id:
                            act_name = known_name
                            break
                
                # 3. 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", "")
                
                out += f"### {i}. {act_name} — Section {sec_num}: {title}\n"
                out += f"{text[:600]}...\n\n"
                out += f"Source URI: {n.get('uri', f'https://www.legislation.gov.uk/id/{leg_id}/section/{sec_num}')}\n\n"
                
            return out
    except Exception as e:
        return f"Error querying Lex Vector API: {str(e)}"

@mcp.tool()
async def get_official_explanatory_note(act_name: str, section_number: str) -> str:
    """
    Dynamically fetches the Official Government Explanatory Note for a specific Act and section.
    Explanatory Notes are plain English explainers written by the government.
    Note: Acts passed prior to 1999 (e.g., Mental Health Act 1983) generally do not have these.
    
    Args:
        act_name: The full name of the Act (e.g., 'Mental Capacity Act 2005').
        section_number: The specific section number as a string (e.g., '3').
    """
    url = f'{BASE_URL}/explanatory_note/section/search'
    payload = {
        'query': f'"{act_name}" Section {section_number}',
        'limit': 3
    }
    
    try:
        async with httpx.AsyncClient() as client:
            r = await client.post(url, json=payload, timeout=10.0)
            if r.status_code == 200:
                data = r.json()
                if isinstance(data, list):
                    parent_id = NURSING_ACTS.get(act_name, "")
                    for note in data:
                        # Match the parent ID to ensure this note belongs to the right Act
                        if parent_id and parent_id in note.get('legislation_id', ''):
                            text = note.get('text', '')
                            if text:
                                return f"### Official Explanatory Note ({act_name} S.{section_number})\n\n{text}"
                                
        return f"No Official Explanatory Note could be found for '{act_name}' Section {section_number}. The Act may pre-date the 1999 introduction of Explanatory Notes."
    except Exception as e:
        return f"Error fetching Explanatory Note: {str(e)}"

if __name__ == "__main__":
    # Ensure this runs correctly when started via cursor/claude
    mcp.run()