lorebase-inference 0.2.3__py3-none-any.whl

lorebase_inference/__init__.py

@@ -0,0 +1,70 @@
+"""Model-agnostic inference: interfaces, provider registry and a deterministic fake provider.
+
+Importing the package registers the `inference` settings section (with the `LLM_*` environment
+aliases) with `lorebase_platform`.
+"""
+
+from lorebase_inference.errors import (
+    DimensionMismatchError,
+    EndpointError,
+    InferenceError,
+    InvalidOutputError,
+    ProviderUnavailableError,
+)
+from lorebase_inference.interfaces import (
+    ChatModel,
+    ChatOptions,
+    ChatResult,
+    EmbeddingModel,
+    JsonSchema,
+    Label,
+    Message,
+    RelationExtractor,
+    RelationLabel,
+    Span,
+    SpanExtractor,
+    SpanRelation,
+    Usage,
+)
+from lorebase_inference.registry import (
+    ProviderRegistry,
+    UnknownProviderError,
+    create_chat_model,
+    create_embedding_model,
+    create_relation_extractor,
+    create_span_extractor,
+    validate_providers,
+)
+from lorebase_inference.settings import InferenceSettings, inference_settings
+
+__version__ = "0.2.3"
+
+__all__ = [
+    "ChatModel",
+    "ChatOptions",
+    "ChatResult",
+    "DimensionMismatchError",
+    "EmbeddingModel",
+    "EndpointError",
+    "InferenceError",
+    "InferenceSettings",
+    "InvalidOutputError",
+    "JsonSchema",
+    "Label",
+    "Message",
+    "ProviderRegistry",
+    "ProviderUnavailableError",
+    "RelationExtractor",
+    "RelationLabel",
+    "Span",
+    "SpanExtractor",
+    "SpanRelation",
+    "UnknownProviderError",
+    "Usage",
+    "create_chat_model",
+    "create_embedding_model",
+    "create_relation_extractor",
+    "create_span_extractor",
+    "inference_settings",
+    "validate_providers",
+]

lorebase_inference/__main__.py

@@ -0,0 +1,6 @@
+"""Run the command line: `python -m lorebase_inference check`."""
+
+from lorebase_inference.cli import main
+
+if __name__ == "__main__":
+    main()

lorebase_inference/chat.py

@@ -0,0 +1,447 @@
+"""Chat completions with schema-valid JSON over any compatible HTTP endpoint.
+
+One client covers every server and proxy, local or hosted, that speaks the chat completions
+protocol (see docs/inference-providers.md for the tested list): `POST {base_url}/chat/completions`
+with `Authorization: Bearer <key>` when a key is configured, plus the configured `headers`;
+`extra_body` fields are merged into every request. The model has no tools and nothing is
+streamed.
+
+Structured output is requested in decreasing order of strength and the answer is always
+validated against the schema locally:
+
+1. `json_schema`: `response_format` of type `json_schema` carrying the schema;
+2. `json_mode`: `response_format` of type `json_object`, with the schema in the system prompt;
+3. `prompt`: no `response_format`, schema and a strict instruction in the system prompt.
+
+The first mode is chosen by `structured_output`. When the server rejects a mode with HTTP
+400 or 422 and an error text that names the response format, the client steps down one mode
+and remembers it per endpoint, model and output settings, in the process and, when a
+`ModeMemory` is installed (`use_mode_memory`), across restarts, so later clients start at the
+mode that works. A parse or validation failure triggers up
+to `repair_attempts` repair calls that show the model its mistakes; usage of all calls is
+summed in the result.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import hashlib
+import hmac
+import json
+import logging
+import math
+import re
+import threading
+from collections import OrderedDict
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import Any, Protocol
+
+import httpx
+from jsonschema import exceptions as schema_exceptions
+from jsonschema import validators
+
+from lorebase_inference.errors import EndpointError, InvalidOutputError
+from lorebase_inference.interfaces import ChatOptions, ChatResult, JsonSchema, Message, Usage
+from lorebase_inference.settings import LlmSettings
+from lorebase_inference.transport import EndpointClient, Sleep
+
+logger = logging.getLogger("lorebase.inference")
+
+MODES = ("json_schema", "json_mode", "prompt")
+_UNSUPPORTED_STATUS = frozenset({400, 422})
+_FORMAT_HINT = re.compile(
+    r"response_format|json_schema|json_object|structured|grammar|schema", re.IGNORECASE
+)
+_MODE_CACHE_LIMIT = 256
+_FENCE = re.compile(r"^```[a-zA-Z0-9_-]*\s*\n?(.*?)\n?```\s*$", re.DOTALL)
+_THINK = re.compile(r"<think>.*?</think>", re.DOTALL)
+_ECHO_LIMIT = 4000
+_MAX_PROBLEMS = 5
+_CHARS_PER_TOKEN = 4
+
+_SCHEMA_INSTRUCTION = (
+    "Answer with a single JSON value that conforms to the JSON Schema below. "
+    "Output only the JSON: no prose, no comments, no code fences.\n\nJSON Schema:\n"
+)
+_REPAIR_INSTRUCTION = (
+    "Your previous answer is not valid: {problems}\n"
+    "Answer again with only the corrected JSON value that conforms to the schema."
+)
+
+
+_ModeKey = tuple[str, str, str, bool | None]
+_known_modes: OrderedDict[_ModeKey, int] = OrderedDict()
+_known_modes_lock = threading.Lock()
+
+
+def reset_known_modes() -> None:
+    """Forget the structured output modes learned from server rejections (for tests)."""
+    with _known_modes_lock:
+        _known_modes.clear()
+
+
+class ModeMemory(Protocol):
+    """Durable memory of the modes learned from rejections, keyed by `mode_digest`.
+
+    `max_age_s` is how long an entry counts: an older one reads as absent and is replaced by
+    the next `remember`. Failures of the memory never fail a call: the client falls back to
+    the in-process cache.
+    """
+
+    async def recall(self, digest: bytes, max_age_s: float) -> int | None: ...
+
+    async def remember(self, digest: bytes, index: int, max_age_s: float) -> None: ...
+
+
+@dataclass(frozen=True, slots=True)
+class _Installed:
+    memory: ModeMemory
+    secret: bytes | None
+    max_age_s: float
+
+
+_installed: _Installed | None = None
+_DIGEST_LABEL = b"lorebase/inference/structured-output-mode/v1"
+
+
+def use_mode_memory(
+    memory: ModeMemory | None, *, secret: str | None = None, max_age_s: float = 7 * 86400.0
+) -> None:
+    """Install (or, with None, remove) the durable memory shared by new and live clients.
+
+    `secret` keys the digests (`inference.mode_memory_key`); entries expire after `max_age_s`.
+    """
+    global _installed
+    _installed = (
+        None
+        if memory is None
+        else _Installed(memory, None if not secret else secret.encode(), max_age_s)
+    )
+
+
+def mode_digest(key: _ModeKey, secret: bytes | None = None) -> bytes:
+    """HMAC-SHA256 of the endpoint, model and output settings, keyed by a fixed label and the
+    optional secret: no address is stored in the clear and, with a secret, a stored digest
+    cannot be matched to a guessed address."""
+    hmac_key = _DIGEST_LABEL + b"\0" + secret if secret else _DIGEST_LABEL
+    return hmac.new(hmac_key, json.dumps(key).encode(), hashlib.sha256).digest()
+
+
+def _mode_key(llm: LlmSettings) -> _ModeKey:
+    return (llm.base_url, llm.model, llm.structured_output, llm.json_schema_strict)
+
+
+def _known_mode(key: _ModeKey) -> int | None:
+    with _known_modes_lock:
+        index = _known_modes.get(key)
+        if index is not None:
+            _known_modes.move_to_end(key)
+        return index
+
+
+def _remember_mode(key: _ModeKey, index: int) -> None:
+    with _known_modes_lock:
+        _known_modes[key] = max(index, _known_modes.get(key, 0))
+        _known_modes.move_to_end(key)
+        while len(_known_modes) > _MODE_CACHE_LIMIT:
+            _known_modes.popitem(last=False)
+
+
+@dataclass(frozen=True, slots=True)
+class _Problem:
+    path: str
+    keyword: str
+    message: str
+
+
+class EndpointChatModel:
+    """`ChatModel` for an HTTP chat completions endpoint."""
+
+    def __init__(
+        self,
+        llm: LlmSettings,
+        *,
+        transport: httpx.AsyncBaseTransport | None = None,
+        sleep: Sleep | None = None,
+    ) -> None:
+        self._llm = llm
+        self._transport = transport
+        self._sleep = sleep
+        self._mode_key = _mode_key(llm)
+        self._mode_index = max(MODES.index(llm.structured_output), _known_mode(self._mode_key) or 0)
+        self._recalled = False
+        self._recall_lock = asyncio.Lock()
+        self._client = EndpointClient(
+            base_url=llm.base_url,
+            api_key=llm.api_key,
+            timeout_s=llm.timeout_s,
+            retry=llm.retry,
+            trust_env=llm.trust_env,
+            max_concurrency=llm.max_concurrency,
+            headers=llm.headers,
+            transport=transport,
+            sleep=sleep,
+        )
+        self._routes: dict[str, EndpointChatModel] = {}
+
+    @property
+    def model(self) -> str:
+        """Configured model name."""
+        return self._llm.model
+
+    @property
+    def output_mode(self) -> str:
+        """Structured output mode in use; lower than configured after a server rejection."""
+        return MODES[self._mode_index]
+
+    async def complete_json(
+        self,
+        messages: Sequence[Message],
+        schema: JsonSchema,
+        options: ChatOptions | None = None,
+    ) -> ChatResult:
+        """Answer `messages` with JSON that conforms to `schema`.
+
+        Raises `InvalidOutputError` when no attempt satisfies the schema, `EndpointError` when
+        the endpoint keeps failing, and `ValueError` when `schema` is not a valid JSON Schema.
+        """
+        options = options or ChatOptions()
+        if options.task and options.task in self._llm.tasks:
+            return await self._route(options.task).complete_json(messages, schema, options)
+        validator = _validator(schema)
+        conversation = list(messages)
+        usage = Usage()
+        problems: list[_Problem] = []
+        attempts = self._llm.repair_attempts + 1
+        for attempt in range(1, attempts + 1):
+            text, call_usage, model, truncated = await self._call(conversation, schema, options)
+            usage = Usage(
+                usage.input_tokens + call_usage.input_tokens,
+                usage.output_tokens + call_usage.output_tokens,
+            )
+            data, problems = _check(text, validator)
+            if not problems:
+                return ChatResult(data=data, usage=usage, model=model)
+            if truncated:
+                break
+            if attempt < attempts:
+                conversation = [
+                    *messages,
+                    Message("assistant", text[:_ECHO_LIMIT]),
+                    Message("user", _REPAIR_INSTRUCTION.format(problems=_describe(problems))),
+                ]
+        reason = ""
+        if truncated:
+            reason = (
+                " (answer cut off by the token limit)"
+                if text.strip()
+                else " (token limit reached before any answer: a reasoning model may have spent "
+                "it on reasoning; raise max_tokens or turn reasoning off through extra_body)"
+            )
+        raise InvalidOutputError(
+            f"model {self._llm.model!r} did not produce JSON satisfying the schema after "
+            f"{attempt} attempt(s){reason}: "
+            + "; ".join(f"{p.path or '/'} [{p.keyword}]" for p in problems[:_MAX_PROBLEMS]),
+            problems=[f"{p.path or '/'}: {p.keyword}" for p in problems],
+            attempts=attempt,
+            usage=usage,
+        )
+
+    async def aclose(self) -> None:
+        """Release connections, including those of per-task routes."""
+        await self._client.aclose()
+        for route in self._routes.values():
+            await route.aclose()
+
+    def _route(self, task: str) -> EndpointChatModel:
+        if task not in self._routes:
+            self._routes[task] = EndpointChatModel(
+                self._llm.for_task(task), transport=self._transport, sleep=self._sleep
+            )
+        return self._routes[task]
+
+    async def _call(
+        self, conversation: list[Message], schema: JsonSchema, options: ChatOptions
+    ) -> tuple[str, Usage, str, bool]:
+        """One completion; steps down the structured output mode when the server rejects it."""
+        await self._recall()
+        while True:
+            payload = self._payload(conversation, schema, options, MODES[self._mode_index])
+            try:
+                body = await self._client.post("/chat/completions", payload)
+            except EndpointError as exc:
+                if (
+                    exc.status in _UNSUPPORTED_STATUS
+                    and self._mode_index < len(MODES) - 1
+                    and _FORMAT_HINT.search(exc.detail)
+                ):
+                    previous = MODES[self._mode_index]
+                    self._mode_index += 1
+                    _remember_mode(self._mode_key, self._mode_index)
+                    await self._persist()
+                    logger.info(
+                        "endpoint rejected structured output mode %s; using %s",
+                        previous,
+                        MODES[self._mode_index],
+                    )
+                    continue
+                raise
+            return self._read(body, conversation)
+
+    async def _recall(self) -> None:
+        """Start at the mode a previous process learned, once per client."""
+        installed = _installed
+        if self._recalled or installed is None:
+            return
+        async with self._recall_lock:
+            if self._recalled:
+                return
+            try:
+                index = await installed.memory.recall(
+                    mode_digest(self._mode_key, installed.secret), installed.max_age_s
+                )
+            except Exception:
+                logger.warning("structured output mode memory unavailable", exc_info=True)
+                self._recalled = True
+                return
+            self._recalled = True
+            if index is not None and 0 <= index < len(MODES):
+                self._mode_index = max(self._mode_index, index)
+                _remember_mode(self._mode_key, self._mode_index)
+
+    async def _persist(self) -> None:
+        installed = _installed
+        if installed is None:
+            return
+        try:
+            await installed.memory.remember(
+                mode_digest(self._mode_key, installed.secret),
+                self._mode_index,
+                installed.max_age_s,
+            )
+        except Exception:
+            logger.warning("structured output mode memory unavailable", exc_info=True)
+
+    def _payload(
+        self, conversation: list[Message], schema: JsonSchema, options: ChatOptions, mode: str
+    ) -> dict[str, Any]:
+        messages = [{"role": m.role, "content": m.content} for m in conversation]
+        payload: dict[str, Any] = {
+            **self._llm.extra_body,
+            "model": self._llm.model,
+            "messages": messages,
+            "stream": False,
+        }
+        temperature = self._llm.temperature if options.temperature is None else options.temperature
+        if temperature is not None:
+            payload["temperature"] = temperature
+        max_tokens = options.max_tokens or self._llm.max_tokens
+        if max_tokens:
+            payload["max_tokens"] = max_tokens
+        if options.seed is not None:
+            payload["seed"] = options.seed
+        if mode == "json_schema":
+            json_schema: dict[str, Any] = {"name": "answer", "schema": dict(schema)}
+            if self._llm.json_schema_strict is not None:
+                json_schema["strict"] = self._llm.json_schema_strict
+            payload["response_format"] = {"type": "json_schema", "json_schema": json_schema}
+        else:
+            _add_instruction(messages, schema)
+            if mode == "json_mode":
+                payload["response_format"] = {"type": "json_object"}
+        return payload
+
+    def _read(
+        self, body: dict[str, Any], conversation: list[Message]
+    ) -> tuple[str, Usage, str, bool]:
+        choices = body.get("choices")
+        choice = choices[0] if isinstance(choices, list) and choices else None
+        message = choice.get("message") if isinstance(choice, dict) else None
+        content = message.get("content") if isinstance(message, dict) else None
+        truncated = isinstance(choice, dict) and choice.get("finish_reason") == "length"
+        if content is None and truncated:
+            content = ""
+        if not isinstance(content, str):
+            raise EndpointError("/chat/completions: response has no text content")
+        reported = body.get("usage")
+        raw_usage: dict[str, Any] = reported if isinstance(reported, dict) else {}
+        usage = Usage(
+            input_tokens=_count(
+                raw_usage.get("prompt_tokens"), sum(len(m.content) for m in conversation)
+            ),
+            output_tokens=_count(raw_usage.get("completion_tokens"), len(content)),
+        )
+        model = body.get("model")
+        return (
+            content,
+            usage,
+            model if isinstance(model, str) and model else self._llm.model,
+            truncated,
+        )
+
+
+def _count(reported: object, chars: int) -> int:
+    """Reported token count; estimated from the text length when the server omits it."""
+    if isinstance(reported, int) and not isinstance(reported, bool) and reported >= 0:
+        return reported
+    return math.ceil(chars / _CHARS_PER_TOKEN)
+
+
+def _add_instruction(messages: list[dict[str, Any]], schema: JsonSchema) -> None:
+    """Put the schema instruction into the first system message, creating it if needed."""
+    instruction = _SCHEMA_INSTRUCTION + json.dumps(schema, ensure_ascii=False)
+    if messages and messages[0]["role"] == "system":
+        messages[0] = {"role": "system", "content": f"{messages[0]['content']}\n\n{instruction}"}
+    else:
+        messages.insert(0, {"role": "system", "content": instruction})
+
+
+def _validator(schema: JsonSchema) -> Any:
+    document = dict(schema)
+    cls = validators.validator_for(document)
+    try:
+        cls.check_schema(document)
+    except schema_exceptions.SchemaError as exc:
+        raise ValueError(f"invalid JSON Schema: {exc.message}") from None
+    return cls(document)
+
+
+def parse_json_text(text: str) -> Any:
+    """Parse a model answer as JSON, tolerating reasoning blocks, code fences and prose.
+
+    Raises `ValueError` when no JSON value can be found.
+    """
+    text = _THINK.sub("", text).strip()
+    fenced = _FENCE.match(text)
+    if fenced:
+        text = fenced.group(1).strip()
+    try:
+        return json.loads(text)
+    except ValueError:
+        pass
+    decoder = json.JSONDecoder()
+    for match in re.finditer(r"[{\[]", text):
+        try:
+            value, _end = decoder.raw_decode(text, match.start())
+        except ValueError:
+            continue
+        return value
+    raise ValueError("no JSON value found in the answer")
+
+
+def _check(text: str, validator: Any) -> tuple[Any, list[_Problem]]:
+    try:
+        data = parse_json_text(text)
+    except ValueError as exc:
+        return None, [_Problem("", "json", str(exc))]
+    errors = sorted(validator.iter_errors(data), key=lambda e: [str(p) for p in e.absolute_path])
+    problems = [
+        _Problem("/" + "/".join(str(p) for p in e.absolute_path), str(e.validator), e.message)
+        for e in errors
+    ]
+    return data, problems
+
+
+def _describe(problems: list[_Problem]) -> str:
+    return "; ".join(f"{p.path or '/'}: {p.message[:200]}" for p in problems[:_MAX_PROBLEMS])