diff --git a/.github/workflows/medcat-llm-components_ci.yml b/.github/workflows/medcat-llm-components_ci.yml new file mode 100644 index 000000000..a3a6b807d --- /dev/null +++ b/.github/workflows/medcat-llm-components_ci.yml @@ -0,0 +1,110 @@ +name: medcat-llm-components - CI (test | publish) + +on: + push: + branches: [ main ] + tags: + - 'medcat-llm-components/v*.*.*' + pull_request: + paths: + - 'medcat-plugins/llm-comps/**' + - '.github/workflows/medcat-llm-components**' + +permissions: + id-token: write + +defaults: + run: + working-directory: ./medcat-plugins/llm-comps + +jobs: + build: + runs-on: ubuntu-latest + strategy: + matrix: + python-version: [ '3.10', '3.11', '3.12' ] + max-parallel: 4 + steps: + - uses: actions/checkout@v7 + - name: Install uv for Python ${{ matrix.python-version }} + uses: astral-sh/setup-uv@v7 + with: + python-version: ${{ matrix.python-version }} + enable-cache: true + - name: Install the project + run: | + uv sync --all-extras --dev + uv run python -m ensurepip + uv run python -m pip install --upgrade pip + uv run python -m pip install "../../medcat-v2[spacy]" + - name: Check types + run: | + uv run python -m mypy --follow-imports=normal src/medcat_llm_components + - name: Ruff linting + run: | + uv run ruff check src/medcat_llm_components --preview + - name: Test + run: | + uv run python -m pytest tests + + publish-to-test-PyPI: + runs-on: ubuntu-latest + needs: build + steps: + - name: Checkout main + uses: actions/checkout@v7 + with: + fetch-depth: 0 # fetch all history + fetch-tags: true # fetch tags explicitly + + - name: Install uv for Python 3.10 + uses: astral-sh/setup-uv@v7 + with: + python-version: '3.10' + enable-cache: true + + - name: Install dependencies + run: | + uv run python -m ensurepip + + - name: Set timestamp-based dev version + run: | + TS=$(date -u +"%Y%m%d%H%M%S") + echo "SETUPTOOLS_SCM_PRETEND_VERSION_FOR_MEDCAT_LLM_COMPONENTS=0.2.2.dev${TS}" >> $GITHUB_ENV + + - name: Build package + run: | + uv build + + - name: Publish distribution to TestPyPI + uses: pypa/gh-action-pypi-publish@release/v1 + with: + repository_url: https://test.pypi.org/legacy/ + packages_dir: medcat-plugins/llm-comps/dist + + publish-to-PyPI: + runs-on: ubuntu-latest + if: startsWith(github.ref, 'refs/tags/') + needs: build + steps: + - name: Checkout main + uses: actions/checkout@v7 + + - name: Install uv for Python 3.10 + uses: astral-sh/setup-uv@v7 + with: + python-version: '3.10' + enable-cache: true + + - name: Install dependencies + run: | + uv run python -m ensurepip + + - name: Build client package + run: | + uv build + + - name: Publish production distribution to PyPI + uses: pypa/gh-action-pypi-publish@release/v1 + with: + packages_dir: medcat-plugins/llm-comps/dist diff --git a/medcat-plugins/llm-comps/pyproject.toml b/medcat-plugins/llm-comps/pyproject.toml new file mode 100644 index 000000000..dc1fe5240 --- /dev/null +++ b/medcat-plugins/llm-comps/pyproject.toml @@ -0,0 +1,63 @@ +[build-system] +requires = ["setuptools>=61.0", "wheel", "setuptools_scm>=8"] +build-backend = "setuptools.build_meta" + +[project] +name = "medcat_llm_components" +dynamic = ["version"] +description = "" +readme = "README.md" +license = { text = "Apache-2.0" } +authors = [ + { name="Mart Ratas", email="mart.ratas@kcl.ac.uk" } +] +requires-python = ">=3.10" + +keywords = ["MedCAT"] + +classifiers = [ + "Development Status :: 3 - Alpha", + "Intended Audience :: Science/Research", + "Topic :: Scientific/Engineering :: Artificial Intelligence", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "License :: OSI Approved :: Apache Software License" +] + +dependencies = [ + "medcat>=2.5", + "pydantic", + "openai~=3.3.1", +] + +[project.optional-dependencies] +dev = [ + "ruff", + "mypy", + "pytest", +] + +# entry-points to add onto medcat +[project.entry-points."medcat.plugins"] +llm_components = "medcat_llm_components" + +[project.urls] +Homepage = "https://github.com/CogStack/cogstack-nlp/tree/main/medcat-plugins/llm-comps" +Repository = "https://github.com/CogStack/cogstack-nlp/tree/main/medcat-plugins/llm-comps" +Issues = "https://github.com/CogStack/cogstack-nlp/issues" + +[tool.setuptools_scm] +root = "../.." +tag_regex = "^medcat-llm-components/v(?P\\d+(?:\\.\\d+)*)(?:[ab]\\d+|rc\\d+)?$" +version_scheme = "post-release" +local_scheme = "no-local-version" +git_describe_command = "git describe --dirty --tags --long --match 'medcat-llm-components/v*'" + +[tool.setuptools.packages.find] +where = ["src"] + +[tool.setuptools.package-data] +"medcat_llm_components" = ["py.typed"] diff --git a/medcat-plugins/llm-comps/src/medcat_llm_components/__init__.py b/medcat-plugins/llm-comps/src/medcat_llm_components/__init__.py new file mode 100644 index 000000000..1f1d2b174 --- /dev/null +++ b/medcat-plugins/llm-comps/src/medcat_llm_components/__init__.py @@ -0,0 +1,3 @@ +from .registration import do_registration as __register + +__register() diff --git a/medcat-plugins/llm-comps/src/medcat_llm_components/base.py b/medcat-plugins/llm-comps/src/medcat_llm_components/base.py new file mode 100644 index 000000000..4420d5b0a --- /dev/null +++ b/medcat-plugins/llm-comps/src/medcat_llm_components/base.py @@ -0,0 +1,151 @@ +"""LLM-based entity provider components for MedCAT (proof of concept). + +Targets the OpenAI-compatible chat-completions wire format +(`POST {base_url}/chat/completions`), since that's the lowest common +denominator for locally-hosted LLM servers - Ollama, vLLM, llama.cpp's +server, LM Studio, TGI, text-generation-webui - either natively or via +an OpenAI-compat mode. It also happens to cover hosted providers +(OpenAI, Groq, Together, OpenRouter, ...) for free, but that's a +secondary benefit, not the design target. + +Layout: + LLMConnectionConfig / AbstractLLMEntityComponent + - shared: client construction, retries, structured-output + negotiation with fallback, response-text cleanup + LLMNERConfig / MyLLMNER + - NER step (ents=None): freeform CSV prompt + span + reconciliation against the doc + LLMLinkConfig / MyLLMLinker + - linking step (ents given): structured-output-first, since + constraining the model to a candidate list is exactly what + it's good at +""" +# from __future__ import annotations + +import logging +import re +import time +from abc import ABC +from typing import Any + +from medcat.components.types import AbstractEntityProvidingComponent +from medcat.config.config import ComponentConfig +from openai import APIConnectionError, APIError, APITimeoutError, OpenAI + +logger = logging.getLogger(__name__) + + +class UnknownSpanException(ValueError): + """Raised when an LLM-reported span can't be reconciled with the source text.""" + + +class _StructuredOutputUnsupported(Exception): + """Internal signal: backend rejected response_format; retry freeform.""" + + +def _looks_like_unsupported_response_format(exc: Exception) -> bool: + # NOTE: heuristic. Backends don't agree on a dedicated error type for + # "I don't support response_format" - they just 400 with varying + # messages. This is best-effort, not a contract; if you hit a backend + # that phrases it differently, structured output will look like a + # hard failure instead of falling back. Worth tightening once you + # know which backends you actually need to support. + msg = str(exc).lower() + return any(s in msg for s in ( + "response_format", "json_schema", "unsupported", "not supported")) + + +# --------------------------------------------------------------------------- +# Shared plumbing +# --------------------------------------------------------------------------- + +class LLMConnectionConfig(ComponentConfig): + """Everything needed to talk to an OpenAI-compatible chat endpoint. + + Shared by every LLM-based component regardless of task. Task + configs (LLMNERConfig, LLMLinkConfig) inherit from this. + """ + base_url: str + api_key: str = "not-needed" # most local servers ignore it, but the SDK requires a non-empty string + model: str + timeout: float = 60.0 + retries: int = 1 + retry_backoff_seconds: float = 1.0 + temperature: float = 0.0 + use_structured_output: bool = True + + +class AbstractLLMEntityComponent(AbstractEntityProvidingComponent, ABC): + """Shared connection / chat / cleanup plumbing for LLM-based components. + + Subclasses own the prompt, the (optional) response schema, and + turning the model's response into MedCAT entities. + """ + + def __init__(self, cnf: LLMConnectionConfig) -> None: + super().__init__() + self.cnf = cnf + self._client = OpenAI(base_url=cnf.base_url, api_key=cnf.api_key) + # once a backend tells us it doesn't support structured output, + # don't keep paying a failed round-trip to rediscover that + self._structured_output_supported = cnf.use_structured_output + + def _chat(self, prompt: str, schema: dict[str, Any] | None = None) -> str: + """Send `prompt` as a single user message, return the raw text + response. `schema`, if given, requests structured output for + this call specifically (falls back to freeform if the backend + rejects it). Retries transient connection/timeout failures.""" + use_schema = schema if self._structured_output_supported else None + + last_exc: Exception | None = None + for attempt in range(self.cnf.retries + 1): + try: + return self._one_call(prompt, use_schema) + except _StructuredOutputUnsupported: + logger.warning( + "%s: backend rejected structured output, falling " + "back to freeform for the rest of this session", + self.cnf.comp_name) + self._structured_output_supported = False + use_schema = None + continue # retry immediately, don't burn a retry slot on this + except (APIConnectionError, APITimeoutError, APIError) as exc: + last_exc = exc + logger.warning( + "%s: LLM call failed (attempt %d/%d): %s", + self.cnf.comp_name, attempt + 1, self.cnf.retries + 1, exc) + if attempt < self.cnf.retries: + time.sleep(self.cnf.retry_backoff_seconds) + assert last_exc is not None + raise last_exc + + def _one_call(self, prompt: str, schema: dict[str, Any] | None) -> str: + kwargs: dict[str, Any] = { + "model": self.cnf.model, + "messages": [{"role": "user", "content": prompt}], + "temperature": self.cnf.temperature, + "timeout": self.cnf.timeout, + } + if schema is not None: + kwargs["response_format"] = { + "type": "json_schema", + "json_schema": {"name": "response", "schema": schema, "strict": True}, + } + try: + resp = self._client.chat.completions.create(**kwargs) + except APIError as exc: + if schema is not None and _looks_like_unsupported_response_format(exc): + raise _StructuredOutputUnsupported from exc + raise + return resp.choices[0].message.content or "" + + _FENCE_RE = re.compile(r"^```[a-zA-Z]*\n|\n```$") + + def _clean_response(self, raw: str) -> str: + text = raw.strip() + # models wrap output in ```csv/```json fences despite instructions not to + text = self._FENCE_RE.sub("", text).strip() + return text + +class MisconfiguredComponentException(ValueError): + pass diff --git a/medcat-plugins/llm-comps/src/medcat_llm_components/linker.py b/medcat-plugins/llm-comps/src/medcat_llm_components/linker.py new file mode 100644 index 000000000..76adfe9b5 --- /dev/null +++ b/medcat-plugins/llm-comps/src/medcat_llm_components/linker.py @@ -0,0 +1,120 @@ +from __future__ import annotations + +import json +import logging +from collections.abc import Callable +from typing import Any + +from medcat.cdb import CDB +from medcat.components.types import CoreComponentType +from medcat.config.config import ComponentConfig, Linking +from medcat.tokenizing.tokenizers import BaseTokenizer +from medcat.tokenizing.tokens import MutableDocument, MutableEntity +from medcat.vocab import Vocab + +from .base import ( + AbstractLLMEntityComponent, + LLMConnectionConfig, + MisconfiguredComponentException, +) + +logger = logging.getLogger(__name__) + + +CandidateFn = Callable[[str], list[tuple[str, str]]] + + +class LLMLinkConfig(LLMConnectionConfig): + comp_name: str = "llm_linker" + context_window: int = 200 + prompt: str = ( + "Given the surrounding text and a medical term found in it, pick " + "the single best matching concept from the candidates below, or " + "'NONE' if none fit.\n\nCONTEXT:\n%s\n\nTERM: %s\n\nCANDIDATES " + "(cui: name):\n%s" + ) + + +class LLMLinker(AbstractLLMEntityComponent): + def __init__(self, cnf: LLMLinkConfig, candidate_fn: CandidateFn) -> None: + super().__init__(cnf) + self.cnf: LLMLinkConfig = cnf + self.candidate_fn = candidate_fn + + def get_type(self) -> CoreComponentType: + return CoreComponentType.linking + + def _candidate_schema(self, candidates: list[tuple[str, str]]) -> dict[str, Any]: + return { + "type": "object", + "properties": { + "cui": {"type": "string", "enum": [c for c, _ in candidates] + ["NONE"]}, + }, + "required": ["cui"], + } + + def _extract_cui(self, raw: str) -> str: + text = self._clean_response(raw) + try: + return json.loads(text)["cui"] + except (json.JSONDecodeError, KeyError, TypeError): + return text # freeform fallback: model just replied with the CUI/NONE + + def _pick_cui( + self, context: str, name: str, candidates: list[tuple[str, str]] + ) -> str | None: + cand_str = "\n".join(f"{cui}: {pretty}" for cui, pretty in candidates) + prompt = self.cnf.prompt % (context, name, cand_str) + raw = self._chat(prompt, schema=self._candidate_schema(candidates)) + answer = self._extract_cui(raw) + valid = {cui for cui, _ in candidates} + return answer if answer in valid else None + + def predict_entities( + self, doc: MutableDocument, ents: list[MutableEntity] | None = None + ) -> list[MutableEntity]: + linked_ents: list[MutableEntity] = [] + if ents is None: + raise NotImplementedError( + "MyLLMLinker only implements the linking step (ents " + "required); use MyLLMNER for the NER step.") + text = doc.base.text + for ent in ents: + candidates = self.candidate_fn(ent.detected_name) + if not candidates: + continue + start = max(0, ent.base.start_char_index - self.cnf.context_window) + end = min(len(text), ent.base.end_char_index + self.cnf.context_window) + cui = self._pick_cui(text[start:end], ent.base.text, candidates) + if cui is not None: + ent.cui = cui # NOTE: attribute name is a guess - adjust to MutableEntity's real API + linked_ents.append(ent) + return linked_ents + + @classmethod + def create_new_component( + cls, + cnf: ComponentConfig, + tokenizer: BaseTokenizer, + cdb: CDB, + vocab: Vocab, + model_load_path: str | None, + ) -> LLMLinker: + + def get_candidates(name: str) -> list[tuple[str, str]]: + if name not in cdb.name2info: + return [] + return [(cui, cdb.get_name(cui)) for + cui in cdb.name2info[name]['per_cui_status']] + if not isinstance(cnf, Linking): + raise MisconfiguredComponentException( + "Wrong type of config on config.linking - " + f"Expected Linking, got {type(cnf).__name__}" + ) + llm_cnf = cnf.additional + if not isinstance(llm_cnf, LLMLinkConfig): + raise MisconfiguredComponentException( + "Wrong type of config on config.linking.additional - " + f"Expected LLMLinkConfig, got {type(llm_cnf).__name__}" + ) + return cls(llm_cnf, get_candidates) diff --git a/medcat-plugins/llm-comps/src/medcat_llm_components/ner.py b/medcat-plugins/llm-comps/src/medcat_llm_components/ner.py new file mode 100644 index 000000000..00f6dfe48 --- /dev/null +++ b/medcat-plugins/llm-comps/src/medcat_llm_components/ner.py @@ -0,0 +1,158 @@ +from __future__ import annotations + +import csv +import io +import logging +import re + +from medcat.cdb import CDB +from medcat.components.types import CoreComponentType +from medcat.config import Config +from medcat.config.config import ComponentConfig, Ner +from medcat.tokenizing.tokenizers import BaseTokenizer +from medcat.tokenizing.tokens import MutableDocument, MutableEntity +from medcat.vocab import Vocab + +from .base import ( + AbstractLLMEntityComponent, + LLMConnectionConfig, + MisconfiguredComponentException, + UnknownSpanException, +) + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# NER step +# --------------------------------------------------------------------------- + +class LLMNERConfig(LLMConnectionConfig): + comp_name: str = "llm_ner" + prompt: str = ( + "Given the following clinical text, list every medical term " + "(finding, symptom, disease, procedure, drug, etc). " + "Respond with ONLY a CSV with header 'entity,start,end' - no " + "prose, no markdown fences, no other commentary. 'start' and " + "'end' are the character offsets of the term in the text " + "below, copied exactly as they appear (same case, same " + "spelling).\n\nTEXT:\n%s" + ) + trust_llm_span: bool = False + span_tolerance_total: int = 10 + + +class LLMNER(AbstractLLMEntityComponent): + def __init__( + self, base_config: Config, + tokenizer: BaseTokenizer, cnf: LLMNERConfig + ) -> None: + super().__init__(cnf) + self.base_config = base_config + self.tokenizer = tokenizer + self.cnf: LLMNERConfig = cnf # narrow the type for the rest of this class + + def get_type(self) -> CoreComponentType: + return CoreComponentType.ner + + def _parse_csv(self, raw: str) -> list[tuple[str, int, int]]: + text = self._clean_response(raw) + if not text: + return [] + reader = csv.DictReader(io.StringIO(text)) + out: list[tuple[str, int, int]] = [] + for row in reader: + try: + name = row["entity"].strip() + start = int(row["start"]) + end = int(row["end"]) + except (KeyError, ValueError, AttributeError) as exc: + logger.warning("Skipping malformed LLM row %r: %s", row, exc) + continue + out.append((name, start, end)) + return out + + def _call_api_raw(self, text: str) -> list[tuple[str, int, int]]: + raw = self._chat(self.cnf.prompt % text) + return self._parse_csv(raw) + + def _get_real_start_end( + self, text: str, name: str, start: int, end: int + ) -> tuple[int, int]: + in_text = text[start:end] + if in_text == name: + return start, end + occurrences = [m.start() for m in re.finditer(re.escape(name), text)] + if not occurrences: + raise UnknownSpanException( + f"'{name}' does not appear verbatim in the document text " + f"(reported span [{start}:{end}] contained {in_text!r}).") + best_start = min(occurrences, key=lambda s: abs(s - start)) + best_end = best_start + len(name) + dist = abs(start - best_start) + abs(end - best_end) + if dist > self.cnf.span_tolerance_total: + raise UnknownSpanException( + f"Unable to find {name!r} near [{start}:{end}]. Nearest " + f"match at [{best_start}:{best_end}] is {dist} chars away " + "(over span_tolerance_total). If this is still correct, " + "raise the tolerance in the config.") + return best_start, best_end + + def _process_spans_into_ents( + self, doc: MutableDocument, raw_spans: list[tuple[str, int, int]] + ) -> list[MutableEntity]: + all_ents: list[MutableEntity] = [] + raw_text = doc.base.text + seen: set[tuple[int, int]] = set() + for name, start, end in raw_spans: + if not self.cnf.trust_llm_span: + try: + start, end = self._get_real_start_end(raw_text, name, start, end) + except UnknownSpanException as exc: + logger.warning("%s", exc) + continue + if (start, end) in seen: + continue # LLM sometimes repeats a term + tkns = doc.get_tokens(start, end) + if not tkns: + logger.warning( + "Unable to tokenize span [%d:%d] (%r)", start, end, name) + continue + entity = self.tokenizer.entity_from_tokens_in_doc(tkns, doc) + entity.detected_name = self.base_config.general.separator.join( + [tkn.base.text for tkn in tkns]) + all_ents.append(entity) + seen.add((start, end)) + return all_ents + + def predict_entities( + self, doc: MutableDocument, ents: list[MutableEntity] | None = None + ) -> list[MutableEntity]: + if ents is not None: + raise NotImplementedError( + "MyLLMNER only implements the NER step; use MyLLMLinker " + "for the ents-provided (linking) step.") + raw_spans = self._call_api_raw(doc.base.text) + return self._process_spans_into_ents(doc, raw_spans) + + @classmethod + def create_new_component( + cls, + cnf: ComponentConfig, + tokenizer: BaseTokenizer, + cdb: CDB, + vocab: Vocab, + model_load_path: str | None, + ) -> LLMNER: + if not isinstance(cnf, Ner): + raise MisconfiguredComponentException( + "Wrong type of config on config.ner - " + f"Expected Ner, got {type(cnf).__name__}" + ) + llm_cnf = cnf.custom_cnf + if not isinstance(llm_cnf, LLMNERConfig): + raise MisconfiguredComponentException( + "Wrong type of config on config.ner.custom_cnf - " + f"Expected LLMNERConfig, got {type(llm_cnf).__name__}" + ) + return cls(cdb.config, tokenizer, llm_cnf) diff --git a/medcat-plugins/llm-comps/src/medcat_llm_components/registration.py b/medcat-plugins/llm-comps/src/medcat_llm_components/registration.py new file mode 100644 index 000000000..a5ceda184 --- /dev/null +++ b/medcat-plugins/llm-comps/src/medcat_llm_components/registration.py @@ -0,0 +1,14 @@ +import logging + +from medcat.components.types import CoreComponentType, lazy_register_core_component + +logger = logging.getLogger(__name__) + + +def do_registration(): + lazy_register_core_component( + CoreComponentType.ner, "llm_ner", + "medcat_llm_components.ner", "LLMNER.create_new_component") + lazy_register_core_component( + CoreComponentType.linking, "llm_linker", + "medcat_llm_components.linker", "LLMLinker.create_new_component") diff --git a/medcat-plugins/llm-comps/tests/mock_llm_server.py b/medcat-plugins/llm-comps/tests/mock_llm_server.py new file mode 100644 index 000000000..860688830 --- /dev/null +++ b/medcat-plugins/llm-comps/tests/mock_llm_server.py @@ -0,0 +1,137 @@ +"""Zero-dependency mock OpenAI-compatible chat endpoint for exercising +MyLLMNER end-to-end without a real LLM, an API key, or internet access. + +Speaks the same shape as `POST {base_url}/chat/completions` that the +`openai` SDK sends and Ollama/vLLM/etc. accept - point the component +straight at it, same as you'd point it at a real local server: + + cnf = LLMNERConfig(base_url="http://localhost:8009", model="whatever") + # OpenAI() appends "/chat/completions" itself, given base_url above + +Run: + python mock_llm_server.py + +It string-matches a small fixed vocabulary against whatever text you +send it and returns CSV matching MyLLMNER's prompt contract, wrapped +in a chat-completion response envelope. It also deliberately nudges +one match's offsets off by 2 characters so you can watch +`_get_real_start_end`'s reconciliation (trust_llm_span=False) correct +it - useful to confirm that path actually works, not just the happy +path. + +This mock only implements the freeform path (what MyLLMNER uses). It +ignores `response_format` if sent, so it won't exercise MyLLMLinker's +structured-output path yet - extend `Handler.do_POST` if/when you need +that (return a JSON body matching the requested schema instead of the +CSV string). + +When you're ready to point this at a real local server instead: same +config, just change `base_url` - e.g. `http://localhost:11434/v1` for +Ollama, or wherever your institute's server lives. +""" +from __future__ import annotations + +import json +import re +import time +from contextlib import contextmanager +import threading +from http.server import BaseHTTPRequestHandler, HTTPServer + +from medcat.components.types import CoreComponentType + +VOCAB = [ + "diabetes", "hypertension", "metformin", + "chest pain", "aspirin", "kidney disease", + "kidney failure" +] + + +def _fake_ner_extract(text: str) -> str: + rows = ["entity,start,end"] + for i, term in enumerate(VOCAB): + for m in re.finditer(re.escape(term), text, re.IGNORECASE): + start, end = m.start(), m.end() + if i == 0: # deliberately mangle the first term's offsets + start, end = start + 2, end + 2 + rows.append(f"{text[m.start():m.end()]},{start},{end}") + return "\n".join(rows) + + +def _fake_linking_return(text: str) -> str: + # just always this one, whatever for now + return "C01" + + +class NERHandler(BaseHTTPRequestHandler): + + def _get_fake_response(self, text: str) -> str: + return _fake_ner_extract(text) + + def do_POST(self): + length = int(self.headers.get("Content-Length", 0)) + raw = self.rfile.read(length).decode("utf-8") + payload = json.loads(raw) + messages = payload.get("messages", []) + prompt = messages[-1]["content"] if messages else "" + + text = prompt.split("TEXT:\n", 1)[-1] + content = self._get_fake_response(text) + + response = { + "id": "chatcmpl-mock", + "object": "chat.completion", + "created": int(time.time()), + "model": payload.get("model", "mock"), + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": content}, + "finish_reason": "stop", + }], + } + body = json.dumps(response).encode("utf-8") + + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, fmt, *args): + pass # keep test output quiet + + +class LinkingHandler(NERHandler): + + def _get_fake_response(self, text: str) -> str: + return _fake_linking_return(text) + + +@contextmanager +def mock_llm_server( + mock_for: CoreComponentType = CoreComponentType.ner, + host: str = "localhost", port: int = 8009 +): + """Context manager to run the mock LLM server in a background thread.""" + handler = NERHandler if mock_for == CoreComponentType.ner else LinkingHandler + server = HTTPServer((host, port), handler) + server_thread = threading.Thread(target=server.serve_forever, daemon=True) + + # Start the server in the background + server_thread.start() + print(f"Mock OpenAI-compatible endpoint started on http://{host}:{port}") + + try: + yield server + finally: + # Shut down the server gracefully when exiting the `with` block + server.shutdown() + server.server_close() + server_thread.join(timeout=2) + print("Mock OpenAI-compatible endpoint stopped.") + + +if __name__ == "__main__": + server = HTTPServer(("localhost", 8009), NERHandler) + print("Mock OpenAI-compatible endpoint on http://localhost:8009 (Ctrl+C to stop)") + server.serve_forever() diff --git a/medcat-plugins/llm-comps/tests/test_registration.py b/medcat-plugins/llm-comps/tests/test_registration.py new file mode 100644 index 000000000..3d4692ac4 --- /dev/null +++ b/medcat-plugins/llm-comps/tests/test_registration.py @@ -0,0 +1,61 @@ +import pytest + +from medcat.components.types import CoreComponentType, create_core_component, CoreComponent +from medcat.config.config import Linking, Ner, Config +from medcat.cdb import CDB +from medcat_llm_components.ner import LLMNERConfig +from medcat_llm_components.linker import LLMLinkConfig + + +BASE_CNF_OPTS = { + "base_url": "https://www.example.com", + "model": "this-model-does-not-exist", +} + + +COMP_CREATORS = [ + # NOTE: for now, the args are just None all around + (CoreComponentType.ner, 'llm_ner', (Ner( + custom_cnf=LLMNERConfig(**BASE_CNF_OPTS), + ), None, None, None)), + (CoreComponentType.linking, 'llm_linker', (Linking( + additional=LLMLinkConfig(**BASE_CNF_OPTS), + ), None, None, None)), +] + + +@pytest.fixture +def base_cnf(): + return Config() + + +@pytest.fixture +def cdb(base_cnf): + return CDB(base_cnf) + + +@pytest.mark.parametrize("comp_type,comp_name,args", COMP_CREATORS) +def test_has_registered_components( + comp_type: CoreComponentType, comp_name: str, args: list, + cdb, +): + comp = create_core_component(comp_type, comp_name, *args[:2], cdb, *args[2:]) + assert comp + + +@pytest.mark.parametrize("comp_type,comp_name,args", COMP_CREATORS) +def test_components_are_core_components( + comp_type: CoreComponentType, comp_name: str, args: list, + cdb, +): + comp = create_core_component(comp_type, comp_name, *args[:2], cdb, *args[2:]) + assert isinstance(comp, CoreComponent) + + +@pytest.mark.parametrize("comp_type,comp_name,args", COMP_CREATORS) +def test_components_are_correct_type( + comp_type: CoreComponentType, comp_name: str, args: list, + cdb, +): + comp = create_core_component(comp_type, comp_name, *args[:2], cdb, *args[2:]) + assert comp.get_type() is comp_type diff --git a/medcat-plugins/llm-comps/tests/test_with_mock_server.py b/medcat-plugins/llm-comps/tests/test_with_mock_server.py new file mode 100644 index 000000000..8a75060de --- /dev/null +++ b/medcat-plugins/llm-comps/tests/test_with_mock_server.py @@ -0,0 +1,130 @@ +from mock_llm_server import mock_llm_server, VOCAB + +from medcat_llm_components.ner import LLMNER, LLMNERConfig +from medcat_llm_components.linker import LLMLinker, LLMLinkConfig + +from medcat.components.types import CoreComponentType +from medcat.config.config import Linking, Ner, Config +from medcat.tokenizing.tokenizers import create_tokenizer +from medcat.cdb import CDB +from medcat.preprocessors.cleaners import prepare_name + +import pytest +import re + + +BASE_URL = "http://localhost:8009" + + +@pytest.fixture +def has_ner_server(): + with mock_llm_server(CoreComponentType.ner): + yield + +@pytest.fixture +def has_linking_server(): + with mock_llm_server(CoreComponentType.linking): + yield + + +@pytest.fixture +def ner_cnf(): + return LLMNERConfig( + base_url=BASE_URL, + model="non-existant-model", + ) + + +@pytest.fixture +def linking_cnf(): + return LLMLinkConfig( + base_url=BASE_URL, + model="non-existant-model", + ) + + +@pytest.fixture +def base_cnf(): + cnf = Config() + cnf.general.nlp.provider = "regex" + return cnf + + +@pytest.fixture +def tokenizer(base_cnf): + return create_tokenizer(base_cnf.general.nlp.provider, base_cnf) + + +@pytest.fixture +def cdb(base_cnf, tokenizer): + cdb = CDB(base_cnf) + for num, name in enumerate(VOCAB): + prepped = prepare_name(name, tokenizer, {}, [ + base_cnf.general, base_cnf.preprocessing, base_cnf.cdb_maker]) + cui = f"C{num + 1:02d}" + cdb.add_names(cui, prepped) + return cdb + + +@pytest.fixture +def ner(ner_cnf, tokenizer, cdb): + ner_cnf = Ner( + custom_cnf=ner_cnf, + ) + return LLMNER.create_new_component(ner_cnf, tokenizer, cdb, None, None) + + +@pytest.fixture +def linker(linking_cnf, cdb): + linking_cnf = Linking( + additional=linking_cnf, + ) + return LLMLinker.create_new_component(linking_cnf, None, cdb, None, None) + + +@pytest.fixture +def example_text(): + return "Patient had diabetes and kidney failure" + + +@pytest.fixture +def example_doc(tokenizer, example_text): + return tokenizer(example_text) + + +def test_ner_can_predict(has_ner_server, ner, example_doc): + ents = ner.predict_entities(example_doc, None) + assert ents + assert all( + ent.detected_name for ent in ents + ) + + +def _generate_ents(doc) -> list: + out_ents = [] + raw_text = doc.text + for name in VOCAB: + if name not in raw_text: + continue + occurrences = [m.start() for m in re.finditer(re.escape(name), raw_text)] + for start in occurrences: + tkns = doc.get_tokens(start, start + len(name)) + if not tkns: + continue + tkn_start, tkn_end = tkns[0].index, tkns[-1].index + ent = doc[tkn_start: tkn_end + 1] + ent.detected_name = name + out_ents.append(ent) + return out_ents + + +def test_linker_can_predict( + has_linking_server, linker, example_doc +): + ner_ents = _generate_ents(example_doc) + assert ner_ents + linked_ents = linker.predict_entities(example_doc, ner_ents) + assert linked_ents + assert all( + ent.cui for ent in linked_ents + )