agentinputlab 0.1.0a1__py3-none-any.whl

agentinputlab/backends/llm/openai_compatible.py

@@ -0,0 +1,332 @@
+"""OpenAI-compatible backend for hosted APIs and self-hosted vLLM servers."""
+
+from __future__ import annotations
+
+import http.client
+import json
+import os
+import urllib.error
+import urllib.parse
+import urllib.request
+from collections.abc import Callable, Mapping
+from contextlib import suppress
+from dataclasses import dataclass, field
+from typing import Any
+
+from agentinputlab.contracts.errors import CapabilityError, ValidationError
+from agentinputlab.contracts.serialization import content_hash, freeze_mapping, to_jsonable
+
+from .base import (
+    GenerationRequest,
+    GenerationResponse,
+    LLMBackend,
+    ModelCapability,
+    Usage,
+)
+
+Transport = Callable[[str, Mapping[str, str], Mapping[str, Any], float], Mapping[str, Any]]
+_MAX_RESPONSE_BYTES = 16 * 1024 * 1024
+_SENSITIVE_EXTRA_HEADERS = frozenset(
+    {
+        "api-key",
+        "authorization",
+        "cookie",
+        "proxy-authorization",
+        "set-cookie",
+        "x-amz-security-token",
+        "x-api-key",
+        "x-auth-token",
+        "x-goog-api-key",
+    }
+)
+_TRANSPORT_ATTEMPT_POLICY = "single_transport_no_implicit_retry.v1"
+
+
+class _OpenAICompatibleHTTPStatusError(RuntimeError):
+    """Sanitized status-only failure without remote response content."""
+
+    def __init__(self, status: int) -> None:
+        self.status = status
+        super().__init__(f"OpenAI-compatible endpoint returned HTTP {status}")
+
+
+class _NoRedirectHandler(urllib.request.HTTPRedirectHandler):
+    def redirect_request(self, req, fp, code, msg, headers, newurl):
+        return None
+
+
+def _default_transport(
+    url: str,
+    headers: Mapping[str, str],
+    payload: Mapping[str, Any],
+    timeout_seconds: float,
+) -> Mapping[str, Any]:
+    body = json.dumps(payload).encode("utf-8")
+    request = urllib.request.Request(url, data=body, headers=dict(headers), method="POST")
+    opener = urllib.request.build_opener(_NoRedirectHandler())
+    http_status: int | None = None
+    endpoint_unavailable = False
+    protocol_failure = False
+    try:
+        with opener.open(request, timeout=timeout_seconds) as response:
+            raw = response.read(_MAX_RESPONSE_BYTES + 1)
+            if len(raw) > _MAX_RESPONSE_BYTES:
+                raise RuntimeError("OpenAI-compatible response exceeds the size limit")
+            return json.loads(raw.decode("utf-8"))
+    except urllib.error.HTTPError as exc:
+        # A remote endpoint can reflect Authorization values or other opaque
+        # request material in its error body.  Never copy that untrusted body
+        # into exceptions because experiment monitors may persist them.
+        http_status = exc.code
+        # Cleanup diagnostics are remote/transport-controlled too; the fixed
+        # status-only error below remains the sole public message.
+        with suppress(Exception):
+            exc.close()
+    except urllib.error.URLError:
+        # ``reason`` is likewise remote/transport-controlled and may contain a
+        # credential-bearing URL.  Keep the diagnostic intentionally generic.
+        endpoint_unavailable = True
+    except http.client.HTTPException:
+        # Invalid status lines and incomplete HTTP frames can embed raw,
+        # remote-controlled bytes in their exception text.
+        protocol_failure = True
+    # Raise after leaving the handler so neither ``__cause__`` nor
+    # ``__context__`` retains the remote-controlled exception for traceback
+    # loggers to render.
+    if http_status is not None:
+        raise _OpenAICompatibleHTTPStatusError(http_status)
+    if endpoint_unavailable:
+        raise RuntimeError("OpenAI-compatible endpoint is unavailable")
+    if protocol_failure:
+        raise RuntimeError("OpenAI-compatible endpoint returned an invalid HTTP response")
+    raise RuntimeError("OpenAI-compatible transport failed without a diagnostic")
+
+
+@dataclass(slots=True)
+class OpenAICompatibleBackend(LLMBackend):
+    model_id: str
+    base_url: str
+    deployment_identity: Mapping[str, Any] = field(default_factory=dict)
+    api_key: str | None = field(default=None, repr=False)
+    api_key_env: str = "OPENAI_API_KEY"
+    timeout_seconds: float = 120.0
+    extra_headers: Mapping[str, str] = field(default_factory=dict, repr=False)
+    extra_body: Mapping[str, Any] = field(default_factory=dict, repr=False)
+    capabilities: frozenset[ModelCapability] = frozenset(
+        {
+            ModelCapability.CHAT,
+            ModelCapability.COMPLETION,
+            ModelCapability.JSON_OBJECT,
+            ModelCapability.TOOL_CALLING,
+            ModelCapability.SEED,
+        }
+    )
+    transport: Transport = _default_transport
+
+    def __post_init__(self) -> None:
+        parsed = urllib.parse.urlparse(self.base_url)
+        if parsed.scheme not in {"http", "https"} or not parsed.netloc:
+            raise ValidationError("base_url must be an absolute http(s) URL")
+        if parsed.username or parsed.password or parsed.query or parsed.fragment:
+            raise ValidationError("base_url must not contain credentials, query, or fragment")
+        if not self.deployment_identity or not self.deployment_identity.get("revision"):
+            raise ValidationError("OpenAI-compatible backend requires deployment_identity.revision")
+        for name, value in self.extra_headers.items():
+            if not isinstance(name, str) or not isinstance(value, str):
+                raise ValidationError("OpenAI-compatible extra headers must contain strings")
+            if name.casefold() in _SENSITIVE_EXTRA_HEADERS:
+                raise ValidationError(
+                    f"sensitive extra header {name!r} is forbidden; use api_key or api_key_env"
+                )
+        self.base_url = self.base_url.rstrip("/")
+        self.deployment_identity = freeze_mapping(self.deployment_identity)
+        self.extra_headers = freeze_mapping(self.extra_headers)
+        self.extra_body = freeze_mapping(self.extra_body)
+        self.capabilities = frozenset(self.capabilities)
+
+    def generate(self, request: GenerationRequest) -> GenerationResponse:
+        if request.tools and ModelCapability.TOOL_CALLING not in self.capabilities:
+            raise CapabilityError(f"backend {self.model_id!r} does not support tools")
+        if request.seed is not None and ModelCapability.SEED not in self.capabilities:
+            raise CapabilityError(f"backend {self.model_id!r} does not support seed")
+        if request.logprobs and ModelCapability.LOGPROBS not in self.capabilities:
+            raise CapabilityError(f"backend {self.model_id!r} does not support logprobs")
+        if (
+            request.response_format
+            and ModelCapability.JSON_OBJECT not in self.capabilities
+            and ModelCapability.JSON_SCHEMA not in self.capabilities
+        ):
+            raise CapabilityError(f"backend {self.model_id!r} does not support response_format")
+        messages: list[dict[str, str]] = []
+        if request.system:
+            messages.append({"role": "system", "content": request.system})
+        if request.messages:
+            for message in request.messages:
+                item = {"role": message.role, "content": message.content}
+                if message.name:
+                    item["name"] = message.name
+                messages.append(item)
+        elif request.prompt:
+            messages.append({"role": "user", "content": request.prompt})
+
+        payload: dict[str, Any] = {
+            **dict(self.extra_body),
+            "model": self.model_id,
+            "messages": messages,
+            "temperature": request.temperature,
+            "max_tokens": request.max_tokens,
+        }
+        if request.seed is not None:
+            payload["seed"] = request.seed
+        if request.stop:
+            payload["stop"] = list(request.stop)
+        if request.logprobs:
+            payload["logprobs"] = True
+            if request.top_logprobs is not None:
+                payload["top_logprobs"] = request.top_logprobs
+        if request.response_format:
+            payload["response_format"] = dict(request.response_format)
+        if request.tools:
+            payload["tools"] = list(request.tools)
+
+        # ``freeze_mapping`` recursively protects backend/request contracts,
+        # so nested values may be MappingProxyType instances.  Materialize the
+        # complete payload—not just its outer mapping—before handing it to a
+        # custom transport or ``json.dumps``.
+        payload = dict(to_jsonable(payload))
+
+        key = self.api_key or os.environ.get(self.api_key_env, "")
+        headers = {"Content-Type": "application/json", **dict(self.extra_headers)}
+        if key:
+            headers["Authorization"] = f"Bearer {key}"
+        data: Mapping[str, Any] | None = None
+        transport_failed = False
+        try:
+            data = self.transport(
+                f"{self.base_url}/chat/completions",
+                headers,
+                payload,
+                self.timeout_seconds,
+            )
+        except Exception:
+            # Custom transports are an extension boundary and may raise
+            # exceptions containing request headers or provider bodies.  Do
+            # not let their text or traceback context enter run artifacts.
+            transport_failed = True
+        if transport_failed:
+            # A failed request may already have generated/billed tokens. Do
+            # not hide another model attempt inside one budget reservation,
+            # or replace its unknown usage with a successful retry's usage.
+            raise RuntimeError("OpenAI-compatible transport failed")
+        if key and _contains_exact_text(data, key):
+            raise RuntimeError("OpenAI-compatible response contains the configured credential")
+        response: GenerationResponse | None = None
+        schema_error = False
+        try:
+            response = _parse_generation_response(data, fallback_model=self.model_id)
+        except (AttributeError, IndexError, KeyError, OverflowError, TypeError, ValueError):
+            schema_error = True
+        # Leave the handler before raising so remote-controlled values cannot
+        # survive in an exception context rendered by a traceback logger.
+        if schema_error or response is None:
+            raise RuntimeError("OpenAI-compatible response schema is invalid")
+        return response
+
+    def fingerprint(self) -> Mapping[str, Any]:
+        parsed = urllib.parse.urlparse(self.base_url)
+        identity = {
+            "backend": "openai_compatible",
+            "model_id": self.model_id,
+            "endpoint_origin": f"{parsed.scheme}://{parsed.netloc}",
+            "base_url": self.base_url,
+            "deployment_identity_hash": content_hash(self.deployment_identity),
+            "capabilities": sorted(item.value for item in self.capabilities),
+            "extra_body_hash": content_hash(self.extra_body),
+            "extra_header_names": sorted(self.extra_headers),
+            "extra_headers_hash": content_hash(self.extra_headers),
+            "transport_attempt_policy": _TRANSPORT_ATTEMPT_POLICY,
+        }
+        revision = self.deployment_identity.get("revision")
+        if isinstance(revision, str) and revision.strip():
+            # Source-bound workers need the declared public revision as well as
+            # its deployment hash. Never expose the entire operator mapping.
+            identity["model_revision"] = revision
+        return identity
+
+
+def _parse_generation_response(
+    data: Any,
+    *,
+    fallback_model: str,
+) -> GenerationResponse:
+    if not isinstance(data, Mapping):
+        raise TypeError("response must be a mapping")
+    choices = data.get("choices")
+    if not isinstance(choices, list) or not choices:
+        raise ValueError("response choices must be a non-empty list")
+    choice = choices[0]
+    if not isinstance(choice, Mapping):
+        raise TypeError("response choice must be a mapping")
+    message = choice.get("message")
+    text: Any = None
+    if message is not None:
+        if not isinstance(message, Mapping):
+            raise TypeError("response message must be a mapping")
+        text = message.get("content")
+    if text is None:
+        text = choice.get("text", "")
+    if not isinstance(text, str):
+        raise TypeError("response text must be a string")
+    model = data.get("model", fallback_model)
+    if not isinstance(model, str):
+        raise TypeError("response model must be a string")
+    finish_reason = choice.get("finish_reason", "")
+    if finish_reason is None:
+        finish_reason = ""
+    if not isinstance(finish_reason, str):
+        raise TypeError("response finish_reason must be a string")
+    usage = data.get("usage", {})
+    if usage is None:
+        usage = {}
+    if not isinstance(usage, Mapping):
+        raise TypeError("response usage must be a mapping")
+    return GenerationResponse(
+        text=text,
+        model=model,
+        finish_reason=finish_reason,
+        usage=Usage(
+            input_tokens=_response_token_count(usage.get("prompt_tokens", 0)),
+            output_tokens=_response_token_count(usage.get("completion_tokens", 0)),
+        ),
+        raw=data,
+    )
+
+
+def _response_token_count(value: Any) -> int:
+    if isinstance(value, bool) or not isinstance(value, int) or value < 0:
+        raise ValueError("response token count must be a non-negative integer")
+    return value
+
+
+def _contains_exact_text(value: Any, needle: str, *, _seen: set[int] | None = None) -> bool:
+    if isinstance(value, str):
+        return needle in value
+    if isinstance(value, Mapping):
+        seen = set() if _seen is None else _seen
+        identity = id(value)
+        if identity in seen:
+            return False
+        seen.add(identity)
+        return any(
+            (isinstance(key, str) and needle in key)
+            or _contains_exact_text(item, needle, _seen=seen)
+            for key, item in value.items()
+        )
+    if isinstance(value, (list, tuple)):
+        seen = set() if _seen is None else _seen
+        identity = id(value)
+        if identity in seen:
+            return False
+        seen.add(identity)
+        return any(_contains_exact_text(item, needle, _seen=seen) for item in value)
+    return False

agentinputlab/backends/llm/transformers.py

@@ -0,0 +1,171 @@
+"""Optional local Transformers backend loaded without a hard runtime dependency."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from dataclasses import dataclass, field
+from typing import Any
+
+from agentinputlab.contracts.errors import CapabilityError, ExternalDependencyError, ValidationError
+from agentinputlab.contracts.serialization import freeze_mapping
+
+from .base import GenerationRequest, GenerationResponse, LLMBackend, ModelCapability, Usage
+
+
+@dataclass(slots=True)
+class TransformersBackend(LLMBackend):
+    model_id: str
+    model_revision: str
+    tokenizer_revision: str
+    tokenizer: Any = None
+    model: Any = None
+    generation_kwargs: Mapping[str, Any] = field(default_factory=dict)
+    cache_dir: str | None = None
+    offline: bool = False
+    capabilities: frozenset[ModelCapability] = frozenset(
+        {ModelCapability.CHAT, ModelCapability.COMPLETION}
+    )
+
+    def __post_init__(self) -> None:
+        if not isinstance(self.offline, bool):
+            raise ValidationError("TransformersBackend offline must be boolean")
+        if self.cache_dir is not None and (
+            not isinstance(self.cache_dir, str) or not self.cache_dir.strip()
+        ):
+            raise ValidationError("TransformersBackend cache_dir must be a non-empty path string")
+        if not self.model_revision or not self.tokenizer_revision:
+            raise ValidationError(
+                "TransformersBackend requires explicit model and tokenizer revisions"
+            )
+        self.generation_kwargs = freeze_mapping(self.generation_kwargs)
+        num_return_sequences = self.generation_kwargs.get("num_return_sequences", 1)
+        if (
+            isinstance(num_return_sequences, bool)
+            or not isinstance(num_return_sequences, int)
+            or num_return_sequences != 1
+        ):
+            raise ValidationError(
+                "TransformersBackend requires num_return_sequences to be the integer 1"
+            )
+        self.capabilities = frozenset(self.capabilities)
+
+    def _load(self) -> None:
+        if self.model is not None and self.tokenizer is not None:
+            return
+        try:
+            from transformers import AutoModelForCausalLM, AutoTokenizer
+        except ImportError as exc:
+            raise ExternalDependencyError(
+                "TransformersBackend requires `pip install agentinputlab[local]`"
+            ) from exc
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            self.model_id,
+            revision=self.tokenizer_revision,
+            trust_remote_code=False,
+            cache_dir=self.cache_dir,
+            local_files_only=self.offline,
+        )
+        self.model = AutoModelForCausalLM.from_pretrained(
+            self.model_id,
+            revision=self.model_revision,
+            trust_remote_code=False,
+            cache_dir=self.cache_dir,
+            local_files_only=self.offline,
+        )
+
+    def generate(self, request: GenerationRequest) -> GenerationResponse:
+        if request.seed is not None and ModelCapability.SEED not in self.capabilities:
+            raise CapabilityError(
+                "TransformersBackend does not implement deterministic seed control"
+            )
+        if request.tools and ModelCapability.TOOL_CALLING not in self.capabilities:
+            raise CapabilityError("TransformersBackend does not support structured tool calling")
+        if request.response_format and not {
+            ModelCapability.JSON_OBJECT,
+            ModelCapability.JSON_SCHEMA,
+        }.intersection(self.capabilities):
+            raise CapabilityError("TransformersBackend does not support response_format")
+        self._load()
+        if request.messages and hasattr(self.tokenizer, "apply_chat_template"):
+            messages = []
+            if request.system:
+                messages.append({"role": "system", "content": request.system})
+            messages.extend(
+                {"role": item.role, "content": item.content} for item in request.messages
+            )
+            prompt = self.tokenizer.apply_chat_template(
+                messages,
+                tokenize=False,
+                add_generation_prompt=True,
+            )
+        elif request.messages:
+            raise CapabilityError(
+                "Transformers tokenizer has no chat template for message-based generation"
+            )
+        else:
+            prompt = f"{request.system}\n\n{request.prompt}" if request.system else request.prompt
+        inputs = self.tokenizer(prompt, return_tensors="pt")
+        outputs = self.model.generate(
+            **inputs,
+            max_new_tokens=request.max_tokens,
+            do_sample=request.temperature > 0,
+            temperature=max(request.temperature, 1e-5),
+            **dict(self.generation_kwargs),
+        )
+        sequences = outputs.sequences if hasattr(outputs, "sequences") else outputs
+        try:
+            input_shape = inputs["input_ids"].shape
+            if len(input_shape) != 2 or int(input_shape[0]) != 1:
+                raise ValueError("input_ids must contain exactly one two-dimensional sequence")
+            input_width = int(input_shape[1])
+            if input_width < 1:
+                raise ValueError("input_ids sequence must not be empty")
+
+            sequence_shape = sequences.shape
+            if len(sequence_shape) != 2 or int(sequence_shape[0]) != 1:
+                raise ValueError("generate must return exactly one two-dimensional sequence")
+            sequence_width = int(sequence_shape[1])
+            if sequence_width < input_width:
+                raise ValueError("generated sequence is shorter than its input")
+
+            generated_sequence = sequences[0]
+            new_tokens = generated_sequence[input_width:]
+            output_tokens = sequence_width - input_width
+        except (AttributeError, IndexError, KeyError, OverflowError, TypeError, ValueError) as exc:
+            raise ValidationError(
+                "Transformers model.generate returned an unsupported sequence shape"
+            ) from exc
+        attention_mask = inputs.get("attention_mask")
+        if attention_mask is None:
+            input_tokens = input_width
+        else:
+            try:
+                mask_shape = attention_mask.shape
+                if (
+                    len(mask_shape) != 2
+                    or int(mask_shape[0]) != 1
+                    or int(mask_shape[1]) != input_width
+                ):
+                    raise ValueError("attention_mask shape must match the single input sequence")
+                input_tokens = int(attention_mask.sum())
+                if input_tokens < 0 or input_tokens > input_width:
+                    raise ValueError("attention_mask token count is outside the input width")
+            except (AttributeError, OverflowError, TypeError, ValueError) as exc:
+                raise ValidationError(
+                    "Transformers tokenizer returned an unsupported attention_mask"
+                ) from exc
+        text = self.tokenizer.decode(new_tokens, skip_special_tokens=True)
+        return GenerationResponse(
+            text=text,
+            model=self.model_id,
+            usage=Usage(input_tokens=input_tokens, output_tokens=output_tokens),
+        )
+
+    def fingerprint(self) -> Mapping[str, Any]:
+        return {
+            "backend": "transformers",
+            "model_id": self.model_id,
+            "model_revision": self.model_revision,
+            "tokenizer_revision": self.tokenizer_revision,
+            "generation_kwargs": dict(self.generation_kwargs),
+        }

agentinputlab/backends/retrieval/__init__.py

@@ -0,0 +1,35 @@
+"""Retrieval backend contracts and small deterministic implementations."""
+
+from .base import Document, RetrievalResult, Retriever
+from .dense import (
+    DENSE_RETRIEVER_PROFILE_IDS,
+    RETRIEVER_PROFILES,
+    BatchTextEncoder,
+    DenseIndexSummary,
+    EmbeddingBackendEncoder,
+    RetrieverProfile,
+    SentenceTransformerDenseRetriever,
+    build_retriever,
+    retriever_profile,
+)
+from .memory import FunctionRetriever, InMemoryBM25Retriever, default_tokenize
+from .sqlite_fts import SQLiteFTS5Retriever
+
+__all__ = [
+    "DENSE_RETRIEVER_PROFILE_IDS",
+    "RETRIEVER_PROFILES",
+    "BatchTextEncoder",
+    "DenseIndexSummary",
+    "Document",
+    "EmbeddingBackendEncoder",
+    "FunctionRetriever",
+    "InMemoryBM25Retriever",
+    "RetrievalResult",
+    "Retriever",
+    "RetrieverProfile",
+    "SQLiteFTS5Retriever",
+    "SentenceTransformerDenseRetriever",
+    "build_retriever",
+    "default_tokenize",
+    "retriever_profile",
+]

agentinputlab/backends/retrieval/base.py

@@ -0,0 +1,50 @@
+"""Standard knowledge retrieval interfaces."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from dataclasses import dataclass, field
+from typing import Any, Protocol, runtime_checkable
+
+from agentinputlab.contracts.errors import ValidationError
+from agentinputlab.contracts.serialization import freeze_mapping
+
+
+@dataclass(frozen=True, slots=True)
+class Document:
+    doc_id: str
+    text: str
+    title: str = ""
+    corpus_revision: str = ""
+    metadata: Mapping[str, Any] = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        if not self.doc_id:
+            raise ValidationError("document requires a non-empty doc_id")
+        object.__setattr__(self, "metadata", freeze_mapping(self.metadata))
+
+
+@dataclass(frozen=True, slots=True)
+class RetrievalResult:
+    document: Document
+    score: float
+    rank: int
+    features: Mapping[str, float] = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        if not isinstance(self.document, Document):
+            raise ValidationError("retrieval result requires a Document")
+        if self.rank < 1:
+            raise ValidationError("retrieval result rank must be positive")
+        object.__setattr__(self, "features", freeze_mapping(self.features))
+
+
+@runtime_checkable
+class Retriever(Protocol):
+    retriever_id: str
+
+    def retrieve(self, query: str, *, top_k: int = 5) -> tuple[RetrievalResult, ...]:
+        """Return deterministic ranked candidates."""
+
+    def fingerprint(self) -> Mapping[str, Any]:
+        """Return corpus/index/retriever identity for balance auditing."""