pdf-quiz-generator/backend/app/services/clinical_corpus.py

"""The clinical library, as source material for a draft.

An indexed shelf of reference works, searched over MCP by the same server the
clinical assistant uses. Same index, deliberately: a second one over the same
documents would be a copy that drifts. What differs is the budget — a chat
answer wants a few tight excerpts because somebody is waiting, and an article
synthesises a whole topic, so it wants more of them and longer.

Excerpts are working material. What comes back is passed to the model as
context to write *from*, in its own words, the way anybody writing from a
reference book does — the prompt says so, and the draft is reviewed by an
educator before it is published either way.

Never raises. A library that cannot be reached must not fail the article
somebody is writing; it just means they are writing without it, and the screen
says so.
"""
import json
import logging
import threading
import time
import urllib.request

from app.config import settings
from app.services import site_settings

logger = logging.getLogger(__name__)

TOOL = "clinical_semantic_search"
TIMEOUT = 60
#: Generous, but not unbounded. "No limit" only moves the ceiling to the
#: model's context window, where overflow truncates the middle of the prompt
#: silently — the worst possible place to lose source material.
DEFAULT_LIMIT = 24
DEFAULT_CONTEXT = 1200
#: An MCP session is a handshake plus an id; reusing it saves two round trips
#: per search. Ten minutes, after which the server may have forgotten us.
SESSION_TTL = 600

_session: dict = {}
_lock = threading.Lock()


def endpoint() -> str:
    return (site_settings.get_value("clinical_mcp_url")
            or getattr(settings, "CLINICAL_MCP_URL", "")
            or "").strip()


def is_enabled() -> bool:
    return bool(endpoint()) and site_settings.get_flag("clinical_library_enabled", False)


def _post(url: str, body: dict, session_id: str | None = None) -> tuple[str | None, str]:
    headers = {"Content-Type": "application/json",
               "Accept": "application/json, text/event-stream"}
    if session_id:
        headers["mcp-session-id"] = session_id
    request = urllib.request.Request(url, data=json.dumps(body).encode(), headers=headers)
    with urllib.request.urlopen(request, timeout=TIMEOUT) as response:
        return response.headers.get("mcp-session-id"), response.read().decode("utf-8", "replace")


def _payload(raw: str) -> dict:
    """The JSON out of an SSE body, or out of a plain one."""
    for line in raw.splitlines():
        if line.startswith("data: "):
            try:
                return json.loads(line[6:])
            except ValueError:
                continue
    try:
        return json.loads(raw)
    except ValueError:
        return {}


def _open_session(url: str) -> str:
    session_id, _ = _post(url, {
        "jsonrpc": "2.0", "id": 1, "method": "initialize",
        "params": {"protocolVersion": "2024-11-05", "capabilities": {},
                   "clientInfo": {"name": "pedshub", "version": "1.0"}},
    })
    if not session_id:
        raise RuntimeError("the library did not open a session")
    _post(url, {"jsonrpc": "2.0", "method": "notifications/initialized"}, session_id)
    return session_id


def _session_for(url: str) -> str:
    with _lock:
        fresh = (_session.get("url") == url
                 and _session.get("id")
                 and time.time() - _session.get("at", 0) < SESSION_TTL)
        if not fresh:
            _session.update({"url": url, "id": _open_session(url), "at": time.time()})
        return _session["id"]


def _drop_session() -> None:
    with _lock:
        _session.clear()


def search(query: str, *, limit: int = DEFAULT_LIMIT,
           context_chars: int = DEFAULT_CONTEXT) -> dict:
    """Excerpts from the library for a topic, or an empty list and a reason."""
    text = str(query or "").strip()
    if not text:
        return {"results": [], "reason": "empty query"}
    url = endpoint()
    if not url:
        return {"results": [], "reason": "No clinical library is configured"}
    if not site_settings.get_flag("clinical_library_enabled", False):
        return {"results": [], "reason": "The clinical library is turned off for this site"}

    body = {
        "jsonrpc": "2.0", "id": 2, "method": "tools/call",
        "params": {"name": TOOL, "arguments": {
            "query": text[:400], "limit": limit, "doc_types": ["file"],
            "score_threshold": 0, "fusion": "rrf",
            "include_context": True, "context_chars": context_chars,
        }},
    }
    for attempt in (1, 2):
        try:
            _, raw = _post(url, body, _session_for(url))
            payload = _payload(raw)
            if payload.get("error"):
                # A dead session looks like an error, not a connection failure,
                # so one retry starts a new one rather than giving up.
                if attempt == 1:
                    _drop_session()
                    continue
                return {"results": [], "reason": "The library refused the search"}
            content = (payload.get("result") or {}).get("content") or []
            found = json.loads(content[0]["text"]) if content else {}
            results = []
            for row in (found.get("results") or [])[:limit]:
                excerpt = " ".join(str(row.get("excerpt") or "").split())
                if not excerpt:
                    continue
                results.append({
                    "title": str(row.get("title") or "Untitled").replace("%20", " "),
                    "category": row.get("category") or "",
                    "excerpt": excerpt[:context_chars * 2],
                })
            return {"results": results, "reason": None if results else "nothing indexed on that"}
        except Exception as exc:
            _drop_session()
            if attempt == 1:
                continue
            logger.warning("Clinical library search failed: %s", exc)
            return {"results": [], "reason": "The clinical library could not be reached"}
    return {"results": [], "reason": "The clinical library could not be reached"}


def for_prompt(results: list[dict]) -> str:
    """The excerpts, each under the work it came from."""
    return "\n\n---\n\n".join(
        f"{row['title']}\n{row['excerpt']}" for row in results)


def sources(results: list[dict]) -> list[str]:
    """The distinct works the excerpts came from, in the order first seen."""
    seen: list[str] = []
    for row in results:
        title = row["title"]
        if title not in seen:
            seen.append(title)
    return seen