agentinputlab 0.1.0a1__py3-none-any.whl

agentinputlab/data/checksums.py

@@ -0,0 +1,9 @@
+"""Small checksum helpers shared by download and preprocessing pipelines."""
+
+from __future__ import annotations
+
+import hashlib
+
+
+def sha256_bytes(content: bytes) -> str:
+    return hashlib.sha256(content).hexdigest()

agentinputlab/data/downloader.py

@@ -0,0 +1,235 @@
+"""Fail-closed artifact acquisition with an injected hardened transport.
+
+The core package deliberately does not ship a generic ``urlopen`` transport.
+An injected transport receives already-resolved targets, must connect to one of
+those addresses while validating TLS for the original hostname, and must not
+follow redirects itself. AgentInputLab validates every redirect hop and the
+reported peer before accepting bytes.
+"""
+
+from __future__ import annotations
+
+import ipaddress
+import socket
+import urllib.parse
+from collections.abc import Iterable, Mapping
+from contextlib import AbstractContextManager
+from dataclasses import dataclass
+from typing import Protocol, runtime_checkable
+
+from agentinputlab.artifacts import ArtifactManifest, ArtifactStore
+from agentinputlab.contracts.errors import ArtifactError
+
+
+@dataclass(frozen=True, slots=True)
+class DownloadPolicy:
+    allow_network: bool = False
+    allowed_hosts: tuple[str, ...] = ()
+    allow_private_addresses: bool = False
+    max_bytes: int = 2 * 1024 * 1024 * 1024
+    timeout_seconds: float = 120.0
+    max_redirects: int = 3
+
+    def __post_init__(self) -> None:
+        normalized = tuple(sorted({host.strip().lower() for host in self.allowed_hosts if host}))
+        object.__setattr__(self, "allowed_hosts", normalized)
+        if self.max_bytes < 1:
+            raise ArtifactError("download max_bytes must be positive")
+        if self.timeout_seconds <= 0:
+            raise ArtifactError("download timeout_seconds must be positive")
+        if self.max_redirects < 0:
+            raise ArtifactError("download max_redirects must be non-negative")
+
+
+@dataclass(frozen=True, slots=True)
+class ValidatedDownloadTarget:
+    url: str
+    hostname: str
+    port: int
+    resolved_addresses: tuple[str, ...]
+
+
+@dataclass(frozen=True, slots=True)
+class DownloadResponse:
+    status_code: int
+    headers: Mapping[str, str]
+    body: bytes
+    peer_address: str
+
+
+@dataclass(frozen=True, slots=True)
+class StreamingDownloadResponse:
+    status_code: int
+    headers: Mapping[str, str]
+    chunks: Iterable[bytes]
+    peer_address: str
+
+
+class StreamingDownloadTransport(Protocol):
+    def stream(
+        self, target: ValidatedDownloadTarget, *, timeout_seconds: float, max_bytes: int
+    ) -> AbstractContextManager[StreamingDownloadResponse]:
+        """Open one DNS-pinned hop; close it even when the consumer stops early."""
+
+
+@runtime_checkable
+class DownloadTransport(Protocol):
+    def request(
+        self,
+        target: ValidatedDownloadTarget,
+        *,
+        timeout_seconds: float,
+        max_bytes: int,
+    ) -> DownloadResponse:
+        """Fetch one hop without redirects or decompression expansion."""
+
+
+def validate_download_url(url: str, policy: DownloadPolicy) -> ValidatedDownloadTarget:
+    if not policy.allow_network:
+        raise ArtifactError("network acquisition is disabled")
+    if not policy.allowed_hosts:
+        raise ArtifactError("network acquisition requires a non-empty host allowlist")
+    parsed = urllib.parse.urlsplit(url)
+    if parsed.scheme != "https" or not parsed.hostname:
+        raise ArtifactError("artifact URLs must use absolute HTTPS")
+    if parsed.username is not None or parsed.password is not None:
+        raise ArtifactError("artifact URLs must not contain user information")
+    hostname = parsed.hostname.lower()
+    if hostname not in policy.allowed_hosts:
+        raise ArtifactError(f"host is not allowlisted: {hostname}")
+    try:
+        port = parsed.port or 443
+    except ValueError as exc:
+        raise ArtifactError("artifact URL contains an invalid port") from exc
+    try:
+        addresses = {
+            info[4][0].split("%", 1)[0]
+            for info in socket.getaddrinfo(hostname, port, type=socket.SOCK_STREAM)
+        }
+    except socket.gaierror as exc:
+        raise ArtifactError(f"cannot resolve artifact host: {hostname}") from exc
+    if not addresses:
+        raise ArtifactError(f"artifact host resolved to no addresses: {hostname}")
+    if not policy.allow_private_addresses:
+        for address in addresses:
+            ip = ipaddress.ip_address(address)
+            if not ip.is_global:
+                raise ArtifactError(f"artifact host resolves to disallowed address: {address}")
+    return ValidatedDownloadTarget(
+        url=urllib.parse.urlunsplit(parsed),
+        hostname=hostname,
+        port=port,
+        resolved_addresses=tuple(sorted(addresses)),
+    )
+
+
+def download_artifact(
+    manifest: ArtifactManifest,
+    store: ArtifactStore,
+    *,
+    policy: DownloadPolicy,
+    transport: DownloadTransport | None = None,
+) -> str:
+    if not manifest.sha256:
+        raise ArtifactError("downloaded artifacts require a pinned SHA-256")
+    if transport is None:
+        raise ArtifactError("network acquisition requires an explicit hardened transport")
+    target = validate_download_url(manifest.source_uri, policy)
+    response: DownloadResponse | None = None
+    for redirect_count in range(policy.max_redirects + 1):
+        response = transport.request(
+            target,
+            timeout_seconds=policy.timeout_seconds,
+            max_bytes=policy.max_bytes,
+        )
+        _validate_peer(response.peer_address, target)
+        if 300 <= response.status_code < 400:
+            if redirect_count >= policy.max_redirects:
+                raise ArtifactError("artifact download exceeded redirect limit")
+            location = _header(response.headers, "location")
+            if not location:
+                raise ArtifactError("artifact redirect is missing Location")
+            target = validate_download_url(
+                urllib.parse.urljoin(target.url, location),
+                policy,
+            )
+            continue
+        if not 200 <= response.status_code < 300:
+            raise ArtifactError(f"artifact download returned HTTP {response.status_code}")
+        break
+    if response is None:  # pragma: no cover - loop invariant
+        raise ArtifactError("artifact transport returned no response")
+    content_length = _header(response.headers, "content-length")
+    if content_length is not None:
+        try:
+            declared_size = int(content_length)
+        except ValueError as exc:
+            raise ArtifactError("artifact Content-Length is invalid") from exc
+        if declared_size < 0 or declared_size > policy.max_bytes:
+            raise ArtifactError("artifact exceeds configured size limit")
+    if len(response.body) > policy.max_bytes:
+        raise ArtifactError("artifact exceeds configured size limit")
+    return str(store.put_bytes(manifest, response.body))
+
+
+def _validate_peer(peer_address: str, target: ValidatedDownloadTarget) -> None:
+    try:
+        normalized = str(ipaddress.ip_address(peer_address.split("%", 1)[0]))
+    except ValueError as exc:
+        raise ArtifactError("download transport reported an invalid peer address") from exc
+    allowed = {str(ipaddress.ip_address(address)) for address in target.resolved_addresses}
+    if normalized not in allowed:
+        raise ArtifactError("download peer does not match the validated DNS result")
+
+
+def download_artifact_streamed(
+    manifest: ArtifactManifest,
+    store: ArtifactStore,
+    *,
+    policy: DownloadPolicy,
+    transport: StreamingDownloadTransport,
+) -> str:
+    """Acquire a pinned file with bounded RAM and validation at every hop."""
+    if not manifest.sha256:
+        raise ArtifactError("downloaded artifacts require a pinned SHA-256")
+    if manifest.size_bytes is not None and manifest.size_bytes > policy.max_bytes:
+        raise ArtifactError("artifact exceeds configured size limit before download")
+    target = validate_download_url(manifest.source_uri, policy)
+    for count in range(policy.max_redirects + 1):
+        with transport.stream(
+            target, timeout_seconds=policy.timeout_seconds, max_bytes=policy.max_bytes
+        ) as response:
+            _validate_peer(response.peer_address, target)
+            if 300 <= response.status_code < 400:
+                if count >= policy.max_redirects:
+                    raise ArtifactError("artifact download exceeded redirect limit")
+                location = _header(response.headers, "location")
+                if not location:
+                    raise ArtifactError("artifact redirect is missing Location")
+                target = validate_download_url(urllib.parse.urljoin(target.url, location), policy)
+                continue
+            if not 200 <= response.status_code < 300:
+                raise ArtifactError(f"artifact download returned HTTP {response.status_code}")
+            encoding = _header(response.headers, "content-encoding")
+            if encoding and encoding.lower() != "identity":
+                raise ArtifactError("artifact stream refuses encoded response bodies")
+            declared = _header(response.headers, "content-length")
+            if declared is not None:
+                try:
+                    size = int(declared)
+                except ValueError as exc:
+                    raise ArtifactError("artifact Content-Length is invalid") from exc
+                if size < 0 or size > policy.max_bytes:
+                    raise ArtifactError("artifact exceeds configured size limit")
+                if manifest.size_bytes is not None and size != manifest.size_bytes:
+                    raise ArtifactError("artifact declared size differs from pinned size")
+            return str(store.put_stream(manifest, response.chunks, max_bytes=policy.max_bytes))
+    raise ArtifactError("artifact download exceeded redirect limit")  # pragma: no cover
+
+
+def _header(headers: Mapping[str, str], name: str) -> str | None:
+    expected = name.lower()
+    for key, value in headers.items():
+        if key.lower() == expected:
+            return str(value)
+    return None

agentinputlab/data/http.py

@@ -0,0 +1,171 @@
+"""Pinned-address HTTPS transport for explicit benchmark acquisition.
+
+This transport is opt-in.  It connects to one of the addresses validated by
+``validate_download_url`` while retaining the original hostname for TLS SNI and
+certificate verification.  It does not use environment proxies, follow
+redirects, or decompress response bodies.
+"""
+
+from __future__ import annotations
+
+import http.client
+import socket
+import ssl
+import urllib.parse
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+
+from agentinputlab.contracts.errors import ArtifactError
+
+from .downloader import DownloadResponse, StreamingDownloadResponse, ValidatedDownloadTarget
+
+
+class _PinnedHTTPSConnection(http.client.HTTPSConnection):
+    def __init__(
+        self,
+        target: ValidatedDownloadTarget,
+        address: str,
+        *,
+        timeout: float,
+        context: ssl.SSLContext,
+    ) -> None:
+        super().__init__(
+            target.hostname,
+            target.port,
+            timeout=timeout,
+            context=context,
+        )
+        self._validated_address = address
+
+    def connect(self) -> None:
+        raw_socket = socket.create_connection(
+            (self._validated_address, self.port),
+            self.timeout,
+            self.source_address,
+        )
+        try:
+            self.sock = self._context.wrap_socket(
+                raw_socket,
+                server_hostname=self.host,
+            )
+        except BaseException:
+            raw_socket.close()
+            raise
+
+
+@dataclass(slots=True)
+class PinnedHTTPSDownloadTransport:
+    """One-hop stdlib HTTPS client honoring the downloader's DNS pin."""
+
+    user_agent: str = "AgentInputLab/0.1 benchmark-acquisition"
+    ssl_context: ssl.SSLContext = field(default_factory=ssl.create_default_context)
+
+    def request(
+        self,
+        target: ValidatedDownloadTarget,
+        *,
+        timeout_seconds: float,
+        max_bytes: int,
+    ) -> DownloadResponse:
+        with self.stream(target, timeout_seconds=timeout_seconds, max_bytes=max_bytes) as response:
+            return DownloadResponse(
+                status_code=response.status_code,
+                headers=response.headers,
+                body=b"".join(response.chunks),
+                peer_address=response.peer_address,
+            )
+
+    @contextmanager
+    def stream(
+        self,
+        target: ValidatedDownloadTarget,
+        *,
+        timeout_seconds: float,
+        max_bytes: int,
+    ):
+        """Stream one validated address, without redirect or body-read retries."""
+        if max_bytes < 1:
+            raise ArtifactError("download transport size limit must be positive")
+        parsed = urllib.parse.urlsplit(target.url)
+        request_target = urllib.parse.urlunsplit(("", "", parsed.path or "/", parsed.query, ""))
+        host_header = target.hostname if target.port == 443 else f"{target.hostname}:{target.port}"
+        failures: list[str] = []
+        for address in target.resolved_addresses:
+            connection = _PinnedHTTPSConnection(
+                target,
+                address,
+                timeout=timeout_seconds,
+                context=self.ssl_context,
+            )
+            try:
+                connection.request(
+                    "GET",
+                    request_target,
+                    headers={
+                        "Accept": "*/*",
+                        "Accept-Encoding": "identity",
+                        "Connection": "close",
+                        "Host": host_header,
+                        "User-Agent": self.user_agent,
+                    },
+                )
+                if connection.sock is None:  # pragma: no cover - http.client invariant
+                    raise ArtifactError("HTTPS transport did not establish a socket")
+                peer_address = str(connection.sock.getpeername()[0])
+                response = connection.getresponse()
+                headers = {key: value for key, value in response.getheaders()}
+                encoding = _header(headers, "content-encoding")
+                if encoding and encoding.lower() != "identity":
+                    raise ArtifactError("download transport refuses encoded response bodies")
+                declared = _header(headers, "content-length")
+                if declared is not None:
+                    try:
+                        declared_size = int(declared)
+                    except ValueError as exc:
+                        raise ArtifactError(
+                            "download response has an invalid Content-Length"
+                        ) from exc
+                    if declared_size < 0 or declared_size > max_bytes:
+                        raise ArtifactError("download response exceeds the transport size limit")
+            except ArtifactError:
+                connection.close()
+                raise
+            except (OSError, ssl.SSLError, http.client.HTTPException) as exc:
+                connection.close()
+                failures.append(f"{address}: {type(exc).__name__}")
+                continue
+            break
+        else:
+            detail = ", ".join(failures) if failures else "no validated addresses"
+            raise ArtifactError(f"HTTPS transport could not connect ({detail})")
+        try:
+            yield StreamingDownloadResponse(
+                status_code=response.status,
+                headers=headers,
+                chunks=_bounded_chunks(response, max_bytes=max_bytes),
+                peer_address=peer_address,
+            )
+        except (OSError, ssl.SSLError, http.client.HTTPException) as exc:
+            raise ArtifactError(f"artifact stream interrupted ({type(exc).__name__})") from exc
+        finally:
+            connection.close()
+
+
+def _bounded_chunks(response: http.client.HTTPResponse, *, max_bytes: int):
+    size = 0
+    while True:
+        chunk = response.read(min(1024 * 1024, max_bytes + 1 - size))
+        if not chunk:
+            return
+        size += len(chunk)
+        if size > max_bytes:
+            raise ArtifactError("download response exceeds the transport size limit")
+        yield chunk
+
+
+def _header(headers: dict[str, str], name: str) -> str | None:
+    expected = name.lower()
+    for key, value in headers.items():
+        if key.lower() == expected:
+            return value
+    return None

agentinputlab/data/index.py

@@ -0,0 +1,21 @@
+"""Versioned index identity used by retrieval balance checks."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from agentinputlab.contracts.serialization import content_hash
+
+
+@dataclass(frozen=True, slots=True)
+class IndexManifest:
+    corpus_revision: str
+    document_filter_hash: str
+    chunker_hash: str
+    embedder_revision: str
+    index_revision: str
+    metric: str
+
+    @property
+    def lineage_hash(self) -> str:
+        return content_hash(self)

agentinputlab/data/overlap.py

@@ -0,0 +1,28 @@
+"""Dataset split-overlap records used by contamination-aware evaluation."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True, slots=True)
+class OverlapRecord:
+    left_split: str
+    right_split: str
+    shared_ids: tuple[str, ...]
+
+    @property
+    def count(self) -> int:
+        return len(self.shared_ids)
+
+
+def detect_id_overlap(
+    left_ids: Iterable[str],
+    right_ids: Iterable[str],
+    *,
+    left_split: str,
+    right_split: str,
+) -> OverlapRecord:
+    shared = tuple(sorted(set(left_ids) & set(right_ids)))
+    return OverlapRecord(left_split, right_split, shared)

agentinputlab/data/preprocess.py

@@ -0,0 +1,32 @@
+"""Deterministic JSONL preprocessing primitives."""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Callable, Iterable
+from typing import Any
+
+from agentinputlab.contracts.serialization import canonical_json
+
+
+def normalize_records(
+    records: Iterable[dict[str, Any]],
+    transform: Callable[[dict[str, Any]], dict[str, Any]],
+    *,
+    id_field: str = "id",
+) -> tuple[dict[str, Any], ...]:
+    normalized = [transform(dict(record)) for record in records]
+    if any(id_field not in record for record in normalized):
+        raise ValueError(f"all normalized records require {id_field!r}")
+    normalized.sort(key=lambda record: str(record[id_field]))
+    if len({str(record[id_field]) for record in normalized}) != len(normalized):
+        raise ValueError("normalized record IDs must be unique")
+    return tuple(normalized)
+
+
+def to_jsonl(records: Iterable[dict[str, Any]]) -> str:
+    return "\n".join(canonical_json(record) for record in records) + "\n"
+
+
+def from_jsonl(content: str) -> tuple[dict[str, Any], ...]:
+    return tuple(json.loads(line) for line in content.splitlines() if line.strip())

agentinputlab/data/storage.py

@@ -0,0 +1,57 @@
+"""Read-only destination-space checks for bounded data preparation.
+
+These checks are snapshots, not filesystem reservations. Independent writers can
+consume space later; callers must still bound each write and clean partial data.
+"""
+
+from __future__ import annotations
+
+import shutil
+from collections.abc import Iterable
+from pathlib import Path
+
+from agentinputlab.contracts import ArtifactError, ValidationError
+
+DEFAULT_PREPARED_LIMIT = 512 * 1024 * 1024
+DEFAULT_DISK_RESERVE = 1024**3
+
+
+def positive_byte_limit(value: int, *, name: str) -> int:
+    if isinstance(value, bool) or not isinstance(value, int) or value < 1:
+        raise ValidationError(f"{name} must be a positive integer")
+    return value
+
+
+def require_storage_capacity(
+    allocations: Iterable[tuple[str | Path, int]],
+    *,
+    reserve_bytes: int = DEFAULT_DISK_RESERVE,
+) -> None:
+    """Check all planned bytes, summing paths on the same destination device.
+
+    No directory or temporary file is created. A missing destination is checked
+    using its nearest existing ancestor, not the current working directory.
+    ``required`` is an upper bound on *additional* source/prepared bytes; cached
+    sources do not need a second allowance. Unknown or negative sizes are errors.
+    """
+    positive_byte_limit(reserve_bytes, name="reserve_bytes")
+    devices: dict[int, tuple[Path, int]] = {}
+    for destination, required in allocations:
+        if isinstance(required, bool) or not isinstance(required, int) or required < 0:
+            raise ValidationError("planned storage bytes must be a non-negative integer")
+        ancestor = Path(destination).expanduser().resolve()
+        while not ancestor.exists():
+            ancestor = ancestor.parent
+        if not ancestor.is_dir():
+            raise ArtifactError(f"storage destination ancestor is not a directory: {ancestor}")
+        device = ancestor.stat().st_dev
+        previous = devices.get(device, (ancestor, 0))
+        devices[device] = (previous[0], previous[1] + required)
+    for ancestor, required in devices.values():
+        free = shutil.disk_usage(ancestor).free
+        if free < required + reserve_bytes:
+            raise ArtifactError(
+                f"insufficient storage at {ancestor}: need {required:,} additional bytes "
+                f"plus {reserve_bytes:,} reserve, found {free:,} free; "
+                "defer this source or use a source-side bounded sample"
+            )