lex_client.py

"""
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."
=======
"""
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