Source code for mixle.task.llm

"""Provider-agnostic LLM adapters for task teachers and local model calls.

mixle has no LLM client of its own; this is the minimal adapter. An :class:`LLM` is anything with ``complete(prompt)
-> str``. :class:`CallableLLM` wraps a local function (a llama.cpp/transformers call, or a test double);
:class:`OpenAICompatLLM` posts to any OpenAI-compatible ``/v1/chat/completions`` endpoint (Ollama, vLLM, TGI,
llama.cpp server, a hosted API) using only the standard library -- no ``openai``/``requests`` dependency.

:func:`llm_labeler` turns an LLM into the teacher the rest of ``mixle.task``
distills from. ``teacher = llm_labeler(OpenAICompatLLM(...), ["spam", "ham"])`` plugs a hosted or local model
into :func:`mixle.task.distill.distill` / :func:`mixle.task.active.active_distill`: the LLM labels selected
examples and the calibrated local student serves requests that it can answer. The same adapter lets an LLM
propose a model specification (:mod:`mixle.task.design`).
"""

from __future__ import annotations

import inspect
import json
import math
import re
import urllib.parse
import urllib.request
from collections.abc import Callable, Sequence
from typing import Any, Protocol, runtime_checkable

_MAX_HTTP_RESPONSE_BYTES = 8 * 1024 * 1024
#: The only URL schemes an LLM endpoint may use. ``urlopen`` honors every scheme its opener knows --
#: ``file:``, ``ftp:`` and ``data:`` among them -- so an endpoint string that reached this module from
#: configuration, a CLI argument or an environment variable could name a local path rather than a
#: server, and ``urlopen`` would happily read it and hand the bytes back as the "model's" reply.
_ALLOWED_ENDPOINT_SCHEMES = frozenset({"http", "https"})




class LLMResponseError(ValueError):
    """Raised when an LLM response cannot satisfy the requested typed contract."""





class LLMParseError(LLMResponseError):
    """Raised when a required label or JSON value cannot be parsed unambiguously."""





@runtime_checkable
class LLM(Protocol):
    """Anything that can turn a prompt into text. The whole contract the rest of the package depends on."""



    def complete(self, prompt: str, *, system: str | None = None, **kwargs: Any) -> str:
        """Return model text for a prompt and optional system instruction."""
        ...






class CallableLLM:
    """Wrap a plain ``fn(prompt) -> str`` (or ``fn(prompt, system)``) as an :class:`LLM` -- local models and tests."""

    def __init__(self, fn: Callable[..., str]) -> None:
        self.fn = fn
        try:
            n_params = len(inspect.signature(fn).parameters)
        except (TypeError, ValueError):
            # a builtin/C callable or anything else signature() can't introspect: assume the
            # single-argument form rather than risk invoking `fn` twice on unrelated errors.
            n_params = 1
        self._accepts_system = n_params >= 2



    def complete(self, prompt: str, *, system: str | None = None, **kwargs: Any) -> str:
        """Call the wrapped Python function and return its text output."""
        # Dispatch on `fn`'s arity up front (checked once, at wrap time) rather than by catching
        # TypeError from the call itself -- a TypeError raised *inside* fn(prompt, system) for an
        # unrelated reason used to be misread as "wrong arity" and silently retried as fn(prompt),
        # invoking `fn` a second time and masking the real error.
        if self._accepts_system:
            return self.fn(prompt, system)
        return self.fn(prompt)






def validated_endpoint_url(url: str) -> str:
    """Return ``url`` unchanged if it names an ``http``/``https`` endpoint; raise ``ValueError`` otherwise.

    Every request this module makes goes through this check first (see :data:`_ALLOWED_ENDPOINT_SCHEMES`
    for why): the scheme is validated against an allowlist *before* a request object is built, so a
    ``file:``/``ftp:``/``data:`` "endpoint" is rejected rather than opened.
    """
    if not isinstance(url, str) or not url.strip():
        raise ValueError("LLM endpoint URL must be a non-empty string")
    scheme = urllib.parse.urlsplit(url).scheme.lower()
    if scheme not in _ALLOWED_ENDPOINT_SCHEMES:
        raise ValueError(
            f"LLM endpoint URL must use one of {sorted(_ALLOWED_ENDPOINT_SCHEMES)}, got "
            f"{scheme!r} in {url!r}. Only an HTTP(S) server can serve a chat-completions endpoint; "
            "any other scheme would make urlopen read a local or non-HTTP resource instead."
        )
    return url



def _http_post_json(url: str, headers: dict[str, str], payload: dict[str, Any], timeout: float) -> dict[str, Any]:
    """POST JSON and parse JSON back, with only the standard library (monkeypatched in tests)."""
    if not isinstance(timeout, (int, float)) or isinstance(timeout, bool) or not math.isfinite(timeout) or timeout <= 0:
        raise ValueError("timeout must be a finite positive number")
    url = validated_endpoint_url(url)
    data = json.dumps(payload).encode()
    req = urllib.request.Request(url, data=data, headers={"Content-Type": "application/json", **headers})
    with urllib.request.urlopen(req, timeout=timeout) as resp:  # nosec B310 # validated_endpoint_url allowlisted the scheme to http/https on the line above, so file:/ftp:/data: never reach urlopen
        status = getattr(resp, "status", 200)
        if not isinstance(status, int) or not 200 <= status < 300:
            raise LLMResponseError(f"LLM endpoint returned HTTP status {status!r}")
        raw = resp.read(_MAX_HTTP_RESPONSE_BYTES + 1)
        if len(raw) > _MAX_HTTP_RESPONSE_BYTES:
            raise LLMResponseError(f"LLM response exceeds {_MAX_HTTP_RESPONSE_BYTES} bytes")
        try:
            decoded = json.loads(raw.decode("utf-8"))
        except (UnicodeDecodeError, json.JSONDecodeError) as exc:
            raise LLMParseError("LLM endpoint returned invalid UTF-8 JSON") from exc
        if not isinstance(decoded, dict):
            raise LLMResponseError("LLM endpoint response must be a JSON object")
        return decoded




class OpenAICompatLLM:
    """An :class:`LLM` backed by any OpenAI-compatible ``/v1/chat/completions`` endpoint (stdlib HTTP only)."""

    def __init__(
        self,
        base_url: str,
        model: str,
        *,
        api_key: str | None = None,
        temperature: float = 0.0,
        max_tokens: int = 512,
        timeout: float = 60.0,
    ) -> None:
        if not isinstance(base_url, str) or not base_url.strip():
            raise ValueError("base_url must be a non-empty string")
        if not isinstance(model, str) or not model.strip():
            raise ValueError("model must be a non-empty string")
        if isinstance(max_tokens, bool) or not isinstance(max_tokens, int) or max_tokens <= 0:
            raise ValueError("max_tokens must be a positive integer")
        if isinstance(temperature, bool) or not isinstance(temperature, (int, float)) or not math.isfinite(temperature):
            raise ValueError("temperature must be finite")
        if (
            isinstance(timeout, bool)
            or not isinstance(timeout, (int, float))
            or not math.isfinite(timeout)
            or timeout <= 0
        ):
            raise ValueError("timeout must be a finite positive number")
        self.base_url = base_url.rstrip("/")
        self.model = model
        self.api_key = api_key
        self.temperature = temperature
        self.max_tokens = max_tokens
        self.timeout = timeout



    def complete(self, prompt: str, *, system: str | None = None, **kwargs: Any) -> str:
        """Call an OpenAI-compatible chat-completions endpoint and return message text."""
        if not isinstance(prompt, str) or not prompt:
            raise ValueError("prompt must be a non-empty string")
        if system is not None and not isinstance(system, str):
            raise ValueError("system must be a string or None")
        temperature = kwargs.get("temperature", self.temperature)
        max_tokens = kwargs.get("max_tokens", self.max_tokens)
        if isinstance(temperature, bool) or not isinstance(temperature, (int, float)) or not math.isfinite(temperature):
            raise ValueError("temperature must be finite")
        if isinstance(max_tokens, bool) or not isinstance(max_tokens, int) or max_tokens <= 0:
            raise ValueError("max_tokens must be a positive integer")
        messages = ([{"role": "system", "content": system}] if system else []) + [{"role": "user", "content": prompt}]
        payload = {
            "model": self.model,
            "messages": messages,
            "temperature": temperature,
            "max_tokens": max_tokens,
        }
        headers = {"Authorization": f"Bearer {self.api_key}"} if self.api_key else {}
        out = _http_post_json(f"{self.base_url}/chat/completions", headers, payload, self.timeout)
        choices = out.get("choices")
        if not isinstance(choices, list) or not choices or not isinstance(choices[0], dict):
            raise LLMResponseError("LLM response must contain a non-empty choices array")
        message = choices[0].get("message")
        if not isinstance(message, dict) or not isinstance(message.get("content"), str):
            raise LLMResponseError("LLM response choice must contain string message.content")
        return message["content"]






def pick_label(text: str, labels: Sequence[str]) -> str:
    """Map a reply to exactly one declared label, or raise :class:`LLMParseError`."""
    if not isinstance(text, str):
        raise LLMParseError("classification response must be text")
    labels = list(labels)
    if (
        not labels
        or any(not isinstance(label, str) or not label.strip() for label in labels)
        or len({label.casefold() for label in labels}) != len(labels)
    ):
        raise ValueError("labels must be a non-empty sequence of case-insensitively unique strings")
    normalized = text.strip().casefold()
    for label in labels:
        if normalized == label.casefold():
            return label
    matches = [
        label for label in labels if re.search(rf"(?<!\w){re.escape(label.casefold())}(?!\w)", normalized) is not None
    ]
    if len(matches) == 1:
        return matches[0]
    if matches:
        raise LLMParseError(f"classification response ambiguously names labels {matches!r}")
    raise LLMParseError("classification response does not name a declared label")





def llm_labeler(
    llm: LLM,
    labels: Sequence[str],
    *,
    instruction: str | None = None,
    system: str | None = None,
) -> Callable[[list[str]], list[str]]:
    """Turn an LLM into a label-constrained teacher ``texts -> [label]`` for distillation / active labeling.

    Each item is classified into ``labels`` by a constrained prompt; the reply is mapped back with
    :func:`pick_label`. The returned callable has the batched-teacher shape the rest of ``mixle.task`` expects.
    """
    labels = list(labels)
    # Validate once at construction, not after the provider has already been called.
    pick_label(str(labels[0]) if labels else "", labels)
    options = ", ".join(str(label) for label in labels)
    instr = instruction or "Classify the following text."
    sys = system or f"You are a precise classifier. Answer with exactly one of: {options}. Output only the label."

    def teacher(texts: list[str]) -> list[str]:
        out = []
        for t in texts:
            reply = llm.complete(f"{instr}\n\nText: {t}\n\nLabel ({options}):", system=sys)
            out.append(pick_label(reply, labels))
        return out

    return teacher





def extract_json_object(text: str) -> dict[str, Any]:
    """Decode the first valid JSON object in arbitrary surrounding text.

    ``JSONDecoder.raw_decode`` provides real JSON string/escape semantics, so braces
    inside strings do not corrupt the scan. A malformed earlier candidate does not
    hide a later valid object.
    """
    if not isinstance(text, str):
        raise LLMParseError("JSON response must be text")
    decoder = json.JSONDecoder()
    for match in re.finditer(r"\{", text):
        try:
            value, _end = decoder.raw_decode(text, match.start())
        except json.JSONDecodeError:
            continue
        if isinstance(value, dict):
            return value
    raise LLMParseError("LLM response contains no valid JSON object")



def _extract_json_object(text: str) -> dict[str, Any]:
    """Backward-compatible private alias for :func:`extract_json_object`."""
    return extract_json_object(text)




def llm_extractor(
    llm: LLM,
    fields: Sequence[str],
    *,
    instruction: str | None = None,
    system: str | None = None,
) -> Callable[[list[str]], list[dict[str, str]]]:
    """Turn an LLM into a field-extraction teacher ``texts -> [{field: value}]`` for :func:`mixle.task.extract.distill_extractor`.

    Each text is extracted into a JSON object over ``fields`` (values must be verbatim substrings so they align to
    token spans during distillation). The returned callable has the batched-teacher shape the extractor expects.
    """
    fields = list(fields)
    if (
        not fields
        or any(not isinstance(field, str) or not field for field in fields)
        or len(set(fields)) != len(fields)
    ):
        raise ValueError("fields must be a non-empty sequence of unique strings")
    field_list = ", ".join(fields)
    instr = instruction or "Extract the fields from the text."
    sys = system or (
        f"You extract fields from text. Return a single JSON object with keys {field_list}. Each value must be an "
        "exact substring of the text (copy it verbatim). Omit a key if the field is absent. Output only JSON."
    )

    def teacher(texts: list[str]) -> list[dict[str, str]]:
        out = []
        for t in texts:
            reply = llm.complete(f"{instr}\n\nFields: {field_list}\n\nText: {t}\n\nJSON:", system=sys)
            parsed = extract_json_object(reply)
            selected: dict[str, str] = {}
            for key, value in parsed.items():
                if key not in fields or value in (None, ""):
                    continue
                if not isinstance(value, str):
                    raise LLMResponseError(f"extracted field {key!r} must be a string")
                if value not in t:
                    raise LLMResponseError(f"extracted field {key!r} must be a verbatim input substring")
                selected[key] = value
            out.append(selected)
        return out

    return teacher