nace-haystack 0.1.0__py3-none-any.whl

haystack_integrations/common/nace/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: 2026-present Nace AI
+#
+# SPDX-License-Identifier: Apache-2.0

haystack_integrations/common/nace/client.py

@@ -0,0 +1,171 @@
+# SPDX-FileCopyrightText: 2026-present Nace AI
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+from importlib.metadata import PackageNotFoundError, version
+from typing import Any
+from urllib.parse import urlsplit
+
+import httpx
+import nace_sdk
+from haystack import logging
+from haystack.utils import Secret
+
+logger = logging.getLogger(__name__)
+
+try:
+    __version__ = version("nace-haystack")
+except PackageNotFoundError:  # a source checkout that isn't installed
+    __version__ = "0.0.0"
+
+CLIENT_NAME = f"nace-haystack/{__version__}"
+
+# The newest Drex model, used when a component is created without `model`. Bumped in a package release whenever
+# Drex ships a new model (see RELEASING.md); every other model id is sent as given and Drex decides.
+DEFAULT_MODEL = "drex-max-v1.1"
+
+# Released, pinned versions never gain image input, so this list can't go stale.
+TEXT_ONLY_MODELS = frozenset({"drex-v1.5", "drex-v2.0"})
+
+# Drex may take up to 55 seconds to answer, so shorter timeouts can cut off a valid request.
+MIN_TIMEOUT = 60.0
+
+# The longest wait between retries, in seconds, even when the server asks for a longer one
+MAX_RETRY_WAIT = 60.0
+
+
+class CappedRetryPolicy(nace_sdk.RetryPolicy):
+    """
+    nace-sdk's retry policy, with every wait capped at `MAX_RETRY_WAIT`.
+
+    The SDK honours `retry-after-ms` however long it is, so a misbehaving server or gateway could stall a pipeline
+    for hours. This overrides the SDK's private `_wait` until the SDK caps waits itself.
+    """
+
+    def _wait(self, attempt: int, error: BaseException) -> float:
+        return min(super()._wait(attempt, error), MAX_RETRY_WAIT)
+
+
+def is_text_only(model: str) -> bool:
+    """
+    Whether `model` is a released Drex version known to read text only.
+
+    :param model: A Drex model id, optionally with a provider prefix such as `nace-ai/` and an OpenRouter variant
+        suffix such as `:floor`.
+    :returns: `True` for `drex-v1.5` and `drex-v2.0`, `False` for every other id.
+    """
+    return model.rsplit("/", 1)[-1].lstrip("~").partition(":")[0] in TEXT_ONLY_MODELS
+
+
+def is_openrouter(base_url: str | None) -> bool:
+    """
+    Whether `base_url` points to OpenRouter, which takes a different request format.
+
+    :param base_url: The API base URL, or `None` for the default.
+    :returns: `True` if the host is openrouter.ai.
+    """
+    return bool(base_url) and (urlsplit(str(base_url)).hostname or "").endswith("openrouter.ai")
+
+
+def provider_name(base_url: str | None) -> str:
+    """
+    Names who answers requests sent to `base_url`, for error messages.
+
+    :param base_url: The API base URL, or `None` for the default.
+    :returns: `"OpenRouter"` or `"Drex"`.
+    """
+    return "OpenRouter" if is_openrouter(base_url) else "Drex"
+
+
+def resolve_base_url(api_base_url: str | None) -> str | None:
+    """
+    Resolves the API base URL from the argument or the `NACE_BASE_URL` environment variable.
+
+    :param api_base_url: The base URL passed to the component, or `None`.
+    :returns: The base URL, or `None` for the SDK default. A blank value (common in `.env` files and CI) means the
+        default. OpenRouter's documented `.../api/v1` is trimmed to `.../api`, because the SDK appends `/v1/...`.
+    """
+    base_url = api_base_url if api_base_url is not None else os.environ.get("NACE_BASE_URL")
+    base_url = (base_url or "").strip() or None
+    if base_url and is_openrouter(base_url):
+        base_url = base_url.rstrip("/").removesuffix("/v1")
+    return base_url
+
+
+def client_options(
+    api_key: Secret, api_base_url: str | None, timeout: float | None, max_retries: int | None
+) -> dict[str, Any]:
+    """
+    Builds the keyword arguments for `nace_sdk.NaceClient` and `nace_sdk.AsyncNaceClient`.
+
+    :param api_key: The Nace API key.
+    :param api_base_url: The base URL passed to the component, or `None`.
+    :param timeout: Timeout in seconds for each request attempt. `None` means 60 seconds.
+    :param max_retries: Maximum number of retries after a failed attempt. `None` means the SDK default of 2.
+    :returns: The client keyword arguments.
+    """
+    base_url = resolve_base_url(api_base_url)
+    if timeout is None:
+        timeout = MIN_TIMEOUT
+    elif timeout < MIN_TIMEOUT:
+        logger.warning(
+            "A timeout of {timeout} s can cut off Drex, which may take up to 55 s to answer. Use at least {minimum} s.",
+            timeout=timeout,
+            minimum=MIN_TIMEOUT,
+        )
+    headers = {}
+    if is_openrouter(base_url):
+        headers = {"HTTP-Referer": "https://nace.ai", "X-OpenRouter-Title": "Drex for Haystack"}
+    # The SDK's retry budget defaults to 30 seconds in total, which would skip retries after a long attempt.
+    # `timeout` already bounds each attempt, so the budget is disabled.
+    retry_kwargs: dict[str, Any] = {"timeout": None}
+    if max_retries is not None:
+        retry_kwargs["max_retries"] = max_retries
+    return {
+        "api_key": api_key.resolve_value(),
+        "base_url": base_url,
+        "timeout": timeout,
+        "retry": CappedRetryPolicy(**retry_kwargs),
+        "headers": headers,
+        "client_info": CLIENT_NAME,
+    }
+
+
+def _set_user_agent(request: httpx.Request) -> None:
+    # The SDK sends its own User-Agent; ours names this integration so traffic can be measured per channel.
+    request.headers["User-Agent"] = CLIENT_NAME
+
+
+async def _set_user_agent_async(request: httpx.Request) -> None:
+    _set_user_agent(request)
+
+
+def make_client(options: dict[str, Any]) -> nace_sdk.NaceClient:
+    """
+    Creates a synchronous Nace client that identifies this integration on every request.
+
+    :param options: Keyword arguments from `client_options`, plus an optional `transport` for tests.
+    :returns: The client.
+    """
+    options = dict(options)
+    http_client = httpx.Client(
+        timeout=options["timeout"], transport=options.pop("transport", None), event_hooks={"request": [_set_user_agent]}
+    )
+    return nace_sdk.NaceClient(**options, http_client=http_client)
+
+
+def make_async_client(options: dict[str, Any]) -> nace_sdk.AsyncNaceClient:
+    """
+    Creates an asynchronous Nace client that identifies this integration on every request.
+
+    :param options: Keyword arguments from `client_options`, plus an optional `transport` for tests.
+    :returns: The client.
+    """
+    options = dict(options)
+    http_client = httpx.AsyncClient(
+        timeout=options["timeout"],
+        transport=options.pop("transport", None),
+        event_hooks={"request": [_set_user_agent_async]},
+    )
+    return nace_sdk.AsyncNaceClient(**options, http_client=http_client)

haystack_integrations/common/nace/concurrency.py

@@ -0,0 +1,49 @@
+# SPDX-FileCopyrightText: 2026-present Nace AI
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from collections.abc import Callable, Iterator, Sequence
+from concurrent.futures import FIRST_COMPLETED, Future, ThreadPoolExecutor, wait
+from typing import TypeVar
+
+T = TypeVar("T")
+R = TypeVar("R")
+
+
+def map_in_threads(func: Callable[[T], R], items: Sequence[T], max_workers: int) -> list[R]:
+    """
+    Applies `func` to every item on up to `max_workers` threads and returns the results in input order.
+
+    Unlike `Executor.map`, which notices a failure only once every earlier item is done, the first exception from
+    any item stops the run: no item starts after it, so no further requests are sent.
+
+    :param func: The function to apply.
+    :param items: The inputs.
+    :param max_workers: The most items processed at once.
+    :returns: `func(item)` for each item, in the order of `items`.
+    :raises Exception: The first exception `func` raised, after the items already running have finished.
+    """
+    results: dict[int, R] = {}
+    pending: dict[Future[R], int] = {}
+    queue: Iterator[tuple[int, T]] = iter(enumerate(items))
+    # Items are handed out one at a time as workers free up, rather than queued all at once, so a worker can never
+    # pick up another item between a failure and the moment it is noticed
+    with ThreadPoolExecutor(max_workers=max_workers) as executor:
+
+        def start_next() -> None:
+            for index, item in queue:
+                pending[executor.submit(func, item)] = index
+                return
+
+        for _ in range(max_workers):
+            start_next()
+        while pending:
+            done, _ = wait(pending, return_when=FIRST_COMPLETED)
+            for future in done:
+                index = pending.pop(future)
+                error = future.exception()
+                if error is not None:
+                    raise error
+                results[index] = future.result()
+                start_next()
+    return [results[index] for index in range(len(items))]

haystack_integrations/common/nace/errors.py

@@ -0,0 +1,138 @@
+# SPDX-FileCopyrightText: 2026-present Nace AI
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import json
+import re
+from collections.abc import Sequence
+from typing import Any
+
+import nace_sdk
+
+BILLING_URL = "https://console.nace.ai/billing"
+OPENROUTER_CREDITS_URL = "https://openrouter.ai/settings/credits"
+
+# Longer messages are cut, so an unexpected response body can't flood logs and tracebacks
+MAX_MESSAGE_CHARS = 500
+
+_IMAGE_PATH = re.compile(r"^images\.(\d+)")
+
+
+def readable_error(error: nace_sdk.NaceError, provider: str, *, image_types: Sequence[str] = ()) -> nace_sdk.NaceError:
+    """
+    Rewrites an API error so its message says who refused the request and what to change.
+
+    The result is an instance of the same nace-sdk class, so `except nace_sdk.NaceError` (or a narrower class such
+    as `nace_sdk.NaceRateLimitError`) still catches it, and its `status`, `body`, `issues`, `request_id` and
+    `idempotency_key` are kept. A 404 or 405 names the URL, which is how a wrong base URL shows.
+
+    :param error: The error nace-sdk raised.
+    :param provider: Who answered: `"Drex"`, or `"OpenRouter"` on that backend.
+    :param image_types: The MIME type of each image sent, in order, so a refused PDF can be named.
+    :returns: The readable error, or `error` itself if it isn't an HTTP error response.
+    """
+    if not isinstance(error, nace_sdk.NaceAPIError) or isinstance(error, nace_sdk.NaceAPIResponseValidationError):
+        return error
+    message, found = _unwrap(error.message) if error.message else (None, [])
+    if message:
+        message = _shorten(message, provider, error.status)
+    issues = error.issues or found
+    # nace-sdk prints errors as "<status> <message> (request_id=...)"
+    readable = type(error)(
+        error.status,
+        error.body,
+        error.headers,
+        message=f"from {provider}: {_describe(error, provider, message, issues, image_types)}",
+    )
+    readable.idempotency_key = error.idempotency_key
+    return readable
+
+
+def _unwrap(message: str) -> tuple[str, list[dict[str, str]]]:
+    # Gateways such as OpenRouter pass upstream errors through as text like `HTTP 422: {"detail": "..."}` or a JSON
+    # list of schema issues. Return a short message and any `{"path", "message"}` issues found in it.
+    text = message.strip()
+    wrapped = re.match(r"^HTTP \d{3}: (.*)$", text, re.DOTALL)
+    if wrapped:
+        text = wrapped.group(1).strip()
+    try:
+        data: Any = json.loads(text)
+    except ValueError:
+        return text, []
+    if isinstance(data, dict):
+        inner = data.get("detail") or data.get("message") or data.get("error")
+        if isinstance(inner, str):
+            return inner, []
+        if isinstance(inner, dict) and isinstance(inner.get("message"), str):
+            return inner["message"], []
+        data = inner
+    if isinstance(data, list):
+        issues = []
+        for item in data:
+            if isinstance(item, dict) and item.get("message"):
+                path = item.get("path", "")
+                if isinstance(path, list):
+                    path = ".".join(str(part) for part in path)
+                issues.append({"path": str(path), "message": str(item["message"])})
+        if issues:
+            return "the request is invalid", issues
+    return text, []
+
+
+def _shorten(message: str, provider: str, status: int) -> str:
+    # Gateways and proxies answer some failures with an HTML page, which says nothing useful in a Python error
+    if message.lstrip()[:15].lower().startswith(("<!doctype", "<html")):
+        return f"{provider} returned an HTML error page (HTTP {status})"
+    if len(message) > MAX_MESSAGE_CHARS:
+        return message[:MAX_MESSAGE_CHARS] + "…"
+    return message
+
+
+def _describe(
+    error: nace_sdk.NaceAPIError,
+    provider: str,
+    message: str | None,
+    issues: list[dict[str, str]],
+    image_types: Sequence[str],
+) -> str:
+    if error.status == 413:  # noqa: PLR2004
+        # Often an HTML page from the gateway rather than a Drex error
+        if not image_types:
+            return "the request is too large. Shorten the text."
+        text = (
+            "the request is too large. Send fewer or smaller images, or pass big images as files uploaded to Drex "
+            "with nace-sdk."
+        )
+        if "application/pdf" in image_types:
+            text += (
+                "\nA PDF can't be shrunk to fit: convert its pages to images first, for example with Haystack's "
+                "PDFToImageContent."
+            )
+        return text
+    text = message or "the request failed"
+    lines = [f"{issue.get('path') or '(request)'}: {issue['message']}" for issue in issues if issue.get("message")]
+    # Drex's top-level message often repeats a single issue; don't print it twice
+    repeated = len(issues) == 1 and (issues[0].get("message") == message or lines == [message])
+    if lines and not repeated:
+        text += "\n" + "\n".join(f"  {line}" for line in lines)
+    for issue in issues:
+        refused = _IMAGE_PATH.match(issue.get("path") or "")
+        if (
+            refused
+            and int(refused.group(1)) < len(image_types)
+            and image_types[int(refused.group(1))] == "application/pdf"
+        ):
+            number = int(refused.group(1)) + 1
+            text += (
+                f"\nImage {number} is a PDF, and Drex reads images. Convert PDF pages to images first, for example "
+                "with Haystack's PDFToImageContent."
+            )
+    if error.status in (404, 405) and error.endpoint:
+        # Usually a wrong base URL (`api_base_url` or `NACE_BASE_URL`)
+        text += f"\n{error.endpoint} doesn't exist. Check the base URL."
+    if error.status == 402:  # noqa: PLR2004
+        if provider == "OpenRouter":
+            text += f"\nAdd credits at {OPENROUTER_CREDITS_URL}"
+        else:
+            text += f"\nAdd funds at {BILLING_URL}"
+    return text

haystack_integrations/common/nace/images.py

@@ -0,0 +1,204 @@
+# SPDX-FileCopyrightText: 2026-present Nace AI
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import base64
+import io
+import math
+from collections.abc import Mapping, Sequence
+from typing import Any, TypeAlias
+
+import filetype
+import nace_sdk
+from haystack import logging
+from haystack.dataclasses import ByteStream, ImageContent
+from haystack.lazy_imports import LazyImport
+
+with LazyImport("Run 'pip install pillow'") as pillow_import:
+    from PIL import Image, ImageOps
+
+logger = logging.getLogger(__name__)
+
+ImageInput: TypeAlias = (
+    ImageContent | ByteStream | nace_sdk.UploadedFile | nace_sdk.WorkspaceFileSource | Mapping[str, str]
+)
+
+# Drex's documented request size limit. It is used only to decide when inline images must be shrunk; Drex enforces
+# the real limit and explains it when a request is too large.
+MAX_REQUEST_BYTES = 4 * 1024 * 1024
+
+_JPEG_QUALITY = 90
+_MAX_SHRINK_ROUNDS = 6
+# Below about 1 megapixel, small print on a scanned page stops being readable
+_MIN_SHRUNK_PIXELS = 1_000_000
+
+
+def _sniff(data: bytes) -> str:
+    guess = filetype.guess(data)
+    return guess.mime if guess else "application/octet-stream"
+
+
+def _declared(mime_type: str | None) -> str | None:
+    # `application/octet-stream` says only "some bytes", which is what loaders such as a plain file read produce
+    return None if mime_type == "application/octet-stream" else mime_type
+
+
+def _reference(image: Any) -> dict[str, str] | None:
+    if isinstance(image, (nace_sdk.UploadedFile, nace_sdk.WorkspaceFileSource)):
+        return {"type": "workspace_file", "workspace_id": str(image.workspace_id), "file_id": str(image.file_id)}
+    if isinstance(image, Mapping) and image.get("type") == "workspace_file":
+        return {"type": "workspace_file", "workspace_id": str(image["workspace_id"]), "file_id": str(image["file_id"])}
+    return None
+
+
+def to_wire_image(image: ImageInput, *, openrouter: bool) -> str | dict[str, str]:
+    """
+    Converts an image to the form Drex reads.
+
+    Images are sent with their own MIME type; Drex decides which types it reads and explains a refusal.
+
+    :param image: An `ImageContent`, a `ByteStream`, or a reference to a file uploaded to Drex with nace-sdk.
+    :param openrouter: Whether the request goes to OpenRouter, which takes inline images only.
+    :returns: A `data:` URL, or a `workspace_file` reference.
+    :raises ValueError: If the image is a file reference on OpenRouter, or isn't a supported input.
+    """
+    if isinstance(image, ImageContent):
+        mime_type = _declared(image.mime_type) or _sniff(base64.b64decode(image.base64_image))
+        return f"data:{mime_type};base64,{image.base64_image}"
+    if isinstance(image, ByteStream):
+        mime_type = _declared(image.mime_type) or _sniff(image.data)
+        return f"data:{mime_type};base64,{base64.b64encode(image.data).decode()}"
+    reference = _reference(image)
+    if reference is None:
+        msg = (
+            f"Unsupported image {type(image).__name__}. Pass an ImageContent, a ByteStream, or a file uploaded to "
+            "Drex with nace-sdk (an UploadedFile, a WorkspaceFileSource, or a "
+            '{"type": "workspace_file", "workspace_id": ..., "file_id": ...} dict).'
+        )
+        raise ValueError(msg)
+    if openrouter:
+        msg = "OpenRouter accepts inline images only, not files uploaded to Drex. Pass the image itself instead."
+        raise ValueError(msg)
+    return reference
+
+
+def image_types(images: Sequence[ImageInput]) -> list[str]:
+    """
+    Lists the MIME type of each image, so error messages can name the one Drex refused.
+
+    :param images: The images as the caller passed them.
+    :returns: One MIME type per image, or `"workspace_file"` for a reference to an uploaded file.
+    """
+    types = []
+    for image in images:
+        if isinstance(image, ImageContent):
+            types.append(_declared(image.mime_type) or _sniff(base64.b64decode(image.base64_image)))
+        elif isinstance(image, ByteStream):
+            types.append(_declared(image.mime_type) or _sniff(image.data))
+        else:
+            types.append("workspace_file")
+    return types
+
+
+def fit_images(images: list[str | dict[str, str]], budget: int) -> list[str | dict[str, str]]:
+    """
+    Shrinks inline images that don't fit in one request, keeping everything else as it is.
+
+    Drex Max downsizes every image to about 1.6 megapixels itself, so a shrunk image that still fits keeps the
+    detail the model reads. Images that fit are sent untouched, so a model that reads more detail isn't held back.
+    An image is never shrunk below about 1 megapixel, where small print stops being readable: when the text leaves
+    too little room, a clear error says so instead.
+
+    :param images: `data:` URLs and `workspace_file` references, in request order.
+    :param budget: How many characters the inline images may take in the request in total.
+    :returns: The images in the same order. Data URLs over their share of `budget` are re-encoded as JPEG and scaled
+        down until they fit. References, and data Pillow can't open (such as a PDF), are returned unchanged.
+    :raises ValueError: If an image must be shrunk and Pillow isn't installed, or can't fit without going below
+        the quality floor.
+    """
+    inline = [index for index, image in enumerate(images) if isinstance(image, str)]
+    if not inline or sum(len(images[index]) for index in inline) <= budget:
+        return list(images)
+    share = budget // len(inline)
+    oversized = [index for index in inline if len(images[index]) > share]
+    left = budget - sum(len(images[index]) for index in inline if index not in oversized)
+    share = left // len(oversized)
+    fitted = list(images)
+    for index in oversized:
+        fitted[index] = _shrink(str(images[index]), share, number=index + 1)
+    return fitted
+
+
+def _no_room(number: int) -> ValueError:
+    msg = (
+        f"Image {number} is too big to send with this much text: it would have to shrink below about 1 megapixel to "
+        "fit Drex's request limit. Shorten the text, send fewer images, or pass big images as files uploaded to Drex "
+        "with nace-sdk, which don't count toward the request size."
+    )
+    return ValueError(msg)
+
+
+def _shrink(data_url: str, share: int, *, number: int) -> str:
+    try:
+        pillow_import.check()
+    except ImportError as error:
+        msg = (
+            "An image is too big to send to Drex in one request, and shrinking it needs Pillow. "
+            "Run 'pip install pillow', or downscale images before this component, for example with "
+            "ImageFileToImageContent(size=(2048, 2048)) or PDFToImageContent(size=(2048, 2048))."
+        )
+        raise ValueError(msg) from error
+
+    raw = base64.b64decode(data_url.partition(",")[2])
+    try:
+        with Image.open(io.BytesIO(raw)) as opened:
+            # Phone photos store their rotation in EXIF, which JPEG re-encoding drops: apply it to the pixels first
+            original = _flatten(ImageOps.exif_transpose(opened))
+    except OSError:
+        # Not an image Pillow can read, such as a PDF. Drex refuses it with a readable error.
+        return data_url
+    if share <= 0:
+        raise _no_room(number)
+
+    scale = 1.0
+    image = original
+    shrunk = _jpeg(image)
+    for _ in range(_MAX_SHRINK_ROUNDS):
+        if len(shrunk) <= share:
+            break
+        scale *= math.sqrt(share / len(shrunk)) * 0.95
+        size = (max(1, round(original.width * scale)), max(1, round(original.height * scale)))
+        if size[0] * size[1] < _MIN_SHRUNK_PIXELS:
+            raise _no_room(number)
+        image = original.resize(size, Image.Resampling.LANCZOS)
+        shrunk = _jpeg(image)
+    if len(shrunk) > share:
+        raise _no_room(number)
+    logger.warning(
+        "Image {original} ({before:.1f} MB) -> {size} ({after:.1f} MB), shrunk to fit Drex's request limit.",
+        original=f"{original.width}x{original.height}",
+        before=len(raw) / 1e6,
+        size=f"{image.width}x{image.height}",
+        after=len(shrunk) * 3 / 4 / 1e6,
+    )
+    return shrunk
+
+
+def _flatten(image: "Image.Image") -> "Image.Image":
+    # JPEG has no transparency: paint transparent areas white rather than letting them turn black.
+    if image.mode in ("RGBA", "LA") or (image.mode == "P" and "transparency" in image.info):
+        rgba = image.convert("RGBA")
+        page = Image.new("RGB", rgba.size, (255, 255, 255))
+        page.paste(rgba, mask=rgba.getchannel("A"))
+        return page
+    # 16-bit and 32-bit grayscale (common in scanner output) would be clipped at 255, turning the page white.
+    # Scale it down to 8 bits instead.
+    if image.mode.startswith("I;16") or image.mode == "I":
+        image = image.convert("I").point(lambda value: value / 256).convert("L")
+    return image.convert("RGB")
+
+
+def _jpeg(image: "Image.Image") -> str:
+    buffer = io.BytesIO()
+    image.save(buffer, "JPEG", quality=_JPEG_QUALITY)
+    return "data:image/jpeg;base64," + base64.b64encode(buffer.getvalue()).decode()