atif-sql 0.2.0__py3-none-any.whl

atif_embed/infrastructure/settings.py

@@ -0,0 +1,154 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""Runtime configuration for atif-embed.
+
+Pydantic v2 ``BaseSettings`` populated from env vars prefixed with
+``ATIF_SQL_`` (workspace convention). Two providers share one vector-store
+contract, selected by ``ATIF_SQL_EMBED_PROVIDER``:
+
+* ``cohere`` (the default): Cohere Embed v4 on Bedrock via the global CRIS
+  profile, 1024-dim Matryoshka, int8 documents, batch 96, concurrency 8.
+* ``gemma``: EmbeddingGemma 2 on this machine, text-only, 768-dim (MRL
+  widths 512/256/128 too), through the optional ``local`` extra.
+
+Each provider keeps its own default store under the corpus root, so switching
+providers never meets the fail-loud ``(model, dim)`` guard on the other's
+vectors, and both stores can live side by side.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Literal, Self
+
+from pydantic import Field, model_validator
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+#: Matryoshka widths each provider can emit. Cohere Embed v4 takes these four
+#: as ``output_dimension``; EmbeddingGemma 2 is trained for 768 natively and
+#: for truncation to the other three (re-normalized after truncating).
+PROVIDER_DIMENSIONS: dict[str, tuple[int, ...]] = {
+    "cohere": (256, 512, 1024, 1536),
+    "gemma": (768, 512, 256, 128),
+}
+
+#: Width a provider emits when ``output_dimension`` is unset.
+DEFAULT_DIMENSION: dict[str, int] = {"cohere": 1024, "gemma": 768}
+
+#: Store directory under the corpus root, per provider. Cohere keeps the name
+#: every existing store already has.
+STORE_DIRNAME: dict[str, str] = {
+    "cohere": "embeddings_lance",
+    "gemma": "embeddings_lance_gemma",
+}
+
+
+class EmbedSettings(BaseSettings):
+    """Env-driven settings for the embedding pipeline."""
+
+    model_config = SettingsConfigDict(
+        env_prefix="ATIF_SQL_",
+        env_file=".env",
+        extra="ignore",
+    )
+
+    #: Which embedder writes and queries the store: Cohere on Bedrock, or
+    #: EmbeddingGemma 2 on this machine.
+    embed_provider: Literal["cohere", "gemma"] = "cohere"
+    #: Cohere Embed v4 global CRIS profile on bedrock-runtime.
+    embed_model_id: str = "global.cohere.embed-v4:0"
+    #: Matryoshka output width; ``None`` takes the provider's default (1024
+    #: for Cohere, 768 for Gemma). An ``int`` checked against the provider's
+    #: widths rather than a ``Literal``, because pydantic refuses the string
+    #: an environment variable carries for an integer ``Literal``.
+    output_dimension: int | None = None
+    #: Cohere document storage type. int8 documents + float queries is the
+    #: adapter's asymmetry: docs are float-widened on the way out for the
+    #: FLOAT[dim] Lance column, queries always embed as float for the distance
+    #: math.
+    embedding_type: Literal["int8", "float", "uint8", "binary", "ubinary"] = "int8"
+    #: Max texts per invoke_model call (Cohere Embed v4 batch ceiling).
+    batch_size: int = 96
+    #: Concurrent Bedrock batches. 8 × batch 96 ran without throttling in
+    #: testing — Cohere's TPM bucket is the binding constraint.
+    embed_concurrency: int = 8
+    #: AWS region for the bedrock-runtime client.
+    region: str = "us-east-1"
+    #: Hugging Face repository of the local model, also the identity stamped
+    #: on every row it writes.
+    gemma_model_id: str = "google/embeddinggemma-2"
+    #: Repository commit the local model loads, pinned so a later upload under
+    #: the same name cannot change the vectors a store already holds.
+    gemma_revision: str = "914f7f89142e33e77833254d9c9b90c3cef7303b"
+    #: Where the local model runs; ``auto`` takes CUDA, then Apple MPS, then CPU.
+    gemma_device: Literal["auto", "cpu", "cuda", "mps"] = "auto"
+    #: Texts per local forward pass. Peak memory grows with it times the
+    #: longest text in the batch.
+    gemma_batch_size: int = Field(default=32, ge=1)
+    #: Local LanceDB dataset URI. ``None`` (the default) resolves per corpus
+    #: and per provider via :meth:`resolve_lance_uri` — one vector store per
+    #: corpus, so re-pointing the source root never mixes two corpora's
+    #: vectors (the corpus-scoping rule).
+    lance_uri: Path | None = None
+    #: Distance metric for the IVF_HNSW_SQ index.
+    hnsw_metric: Literal["cosine", "l2", "dot"] = "cosine"
+
+    @model_validator(mode="after")
+    def _width_fits_provider(self) -> Self:
+        """Refuse a width the selected provider cannot emit."""
+        allowed = PROVIDER_DIMENSIONS[self.embed_provider]
+        if self.output_dimension is not None and self.output_dimension not in allowed:
+            msg = (
+                f"ATIF_SQL_OUTPUT_DIMENSION={self.output_dimension} is not a width the "
+                f"{self.embed_provider} provider emits; use one of {list(allowed)} "
+                f"or leave it unset for {DEFAULT_DIMENSION[self.embed_provider]}"
+            )
+            raise ValueError(msg)
+        return self
+
+    @property
+    def embedding_dim(self) -> int:
+        """The vector width the selected provider writes and queries."""
+        if self.output_dimension is None:
+            return DEFAULT_DIMENSION[self.embed_provider]
+        return int(self.output_dimension)
+
+    @property
+    def active_model_id(self) -> str:
+        """The model identity the selected provider stamps on every row."""
+        if self.embed_provider == "gemma":
+            return self.gemma_model_id
+        return self.embed_model_id
+
+    @property
+    def active_batch_size(self) -> int:
+        """Texts per provider call: one Bedrock request, or one local forward pass."""
+        return self.gemma_batch_size if self.embed_provider == "gemma" else self.batch_size
+
+    @property
+    def active_concurrency(self) -> int:
+        """Batches in flight at once; the local model runs one at a time."""
+        return 1 if self.embed_provider == "gemma" else self.embed_concurrency
+
+    def resolve_lance_uri(self, corpus_root: Path) -> Path:
+        """Effective Lance dataset directory for ``corpus_root``."""
+        if self.lance_uri is not None:
+            return self.lance_uri
+        return corpus_root / STORE_DIRNAME[self.embed_provider]
+
+    def expected_embedding_identity(self) -> tuple[str, int]:
+        """Return ``(model_id, dim)`` without building the embedder.
+
+        Dependency-free identity for the fail-loud store guard: callers that
+        only need to BIND the store (e.g. the search path) can enforce the
+        provider stamp without importing boto3 or torch.
+        """
+        return (self.active_model_id, self.embedding_dim)
+
+
+__all__ = [
+    "DEFAULT_DIMENSION",
+    "PROVIDER_DIMENSIONS",
+    "STORE_DIRNAME",
+    "EmbedSettings",
+]

atif_models/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""atif-models: model alias registry + structured-output LLM client.
+
+Owns CONTRACT-V2.md §Model registry and §Client contract — size aliases
+(small/medium/large) resolve to concrete Bedrock GLOBAL inference-profile
+ids per family (openai default, anthropic escape hatch), and the
+``LlmStructuredProvider`` port turns (system, prompt, pydantic schema)
+into a validated schema instance. NO other atif-sql package hardcodes a
+model id. Hexagonal: ``domain`` (registry, port, schema transform, usage
+accounting — pure) < ``infrastructure`` (bedrock-runtime adapters, env
+settings). Independence contract: only atif-cli and atif-analytics may
+import this package.
+"""

atif_models/domain/__init__.py

@@ -0,0 +1,9 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""Pure domain layer for atif-models.
+
+The model registry (frozen value objects, no I/O), the
+``LlmStructuredProvider`` port + error taxonomy + usage accounting, and
+the pydantic → OpenAI strict-mode schema transform. No boto3, no env,
+no clock — infrastructure adapts to these shapes.
+"""

atif_models/domain/ports.py

@@ -0,0 +1,143 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""The structured-output provider port + error taxonomy + usage accounting.
+
+The seam is narrow on purpose: a system prompt, a user prompt, and a
+pydantic schema in; a validated instance of that schema out. Everything
+provider-specific (model id, region, reasoning effort, the boto3 client,
+the concurrency limiter) is owned by the adapter and fixed at construction.
+
+Error taxonomy follows the atif-converter ``DomainError`` pattern:
+``RefusalError`` is terminal (content filter / model refusal — do not
+retry), ``ProviderUnavailable`` is retryable (the caller's retry queue
+owns the second chance after the adapter's in-process tenacity budget is
+exhausted).
+
+Usage accounting: each call yields one frozen :class:`CallUsage`;
+:class:`UsageAccumulator` sums them thread-safely per pipeline. A
+``threading.Lock`` (not an ``anyio.CapacityLimiter``) guards the counters
+because adapters dispatch blocking ``invoke_model`` calls via
+``anyio.to_thread`` — two worker threads can land in
+:meth:`~UsageAccumulator.add` concurrently, and a lock is the right
+primitive for a critical section.
+"""
+
+from __future__ import annotations
+
+import threading
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Protocol, TypeVar, runtime_checkable
+
+if TYPE_CHECKING:
+    from pydantic import BaseModel
+
+#: A structured-output schema, bound to ``BaseModel`` so the adapter can
+#: both derive the wire schema and ``model_validate`` the response.
+SchemaT = TypeVar("SchemaT", bound="BaseModel")
+
+
+class DomainError(Exception):
+    """Base for every atif-models domain error."""
+
+
+class RefusalError(DomainError):
+    """Terminal: the model refused (content filter / refusal finish reason).
+
+    Never retried — the input itself is the problem. Callers record the
+    unit as refused and move on.
+    """
+
+
+class ProviderUnavailable(DomainError):  # noqa: N818 — names a provider state, not an "*Error"
+    """Retryable: transport/service failure or an unusable response.
+
+    Raised after the adapter's in-process tenacity budget is exhausted,
+    or immediately for non-retryable-but-recoverable shapes (e.g. a
+    ``length`` finish reason — bump ``max_completion_tokens`` or shrink
+    the prompt). Callers enqueue the unit on the retry queue.
+    """
+
+
+@dataclass(frozen=True, slots=True)
+class CallUsage:
+    """Token accounting for one structured-output call.
+
+    ``reasoning_tokens`` is the GPT-5.6 ``completion_tokens_details.reasoning_tokens``
+    count (a subset of ``output_tokens``, already billed as output);
+    ``cached_tokens`` is ``prompt_tokens_details.cached_tokens`` (a subset
+    of ``input_tokens`` billed at the cached rate). Both are 0 when the
+    provider does not report them.
+    """
+
+    input_tokens: int = 0
+    output_tokens: int = 0
+    reasoning_tokens: int = 0
+    cached_tokens: int = 0
+
+
+class UsageAccumulator:
+    """Thread-safe per-pipeline accumulation of :class:`CallUsage`.
+
+    One instance per pipeline run; adapters call :meth:`add` after every
+    successful response, and the pipeline reads :meth:`summary` at the
+    end for the cost-guard log line.
+    """
+
+    def __init__(self) -> None:
+        self._lock = threading.Lock()
+        self._calls = 0
+        self._input_tokens = 0
+        self._output_tokens = 0
+        self._reasoning_tokens = 0
+        self._cached_tokens = 0
+
+    def add(self, usage: CallUsage) -> None:
+        """Fold one call's usage into the totals (safe from any thread)."""
+        with self._lock:
+            self._calls += 1
+            self._input_tokens += usage.input_tokens
+            self._output_tokens += usage.output_tokens
+            self._reasoning_tokens += usage.reasoning_tokens
+            self._cached_tokens += usage.cached_tokens
+
+    def summary(self) -> dict[str, int]:
+        """Snapshot of the totals: calls + the four token counters."""
+        with self._lock:
+            return {
+                "calls": self._calls,
+                "input_tokens": self._input_tokens,
+                "output_tokens": self._output_tokens,
+                "reasoning_tokens": self._reasoning_tokens,
+                "cached_tokens": self._cached_tokens,
+            }
+
+
+@runtime_checkable
+class LlmStructuredProvider(Protocol):
+    """Port: one structured-output call. One adapter per backend.
+
+    The seam is deliberately narrow — everything provider-specific is
+    fixed at adapter construction.
+    """
+
+    async def classify_structured(
+        self, *, system: str, prompt: str, schema: type[SchemaT]
+    ) -> SchemaT:
+        """Run one structured-output call and return a validated ``schema`` instance.
+
+        ``system`` goes in as a ``system``-role message; ``prompt`` is the
+        per-call user payload. Raises :class:`RefusalError` (terminal) or
+        :class:`ProviderUnavailable` (retryable).
+        """
+        ...
+
+
+__all__ = [
+    "CallUsage",
+    "DomainError",
+    "LlmStructuredProvider",
+    "ProviderUnavailable",
+    "RefusalError",
+    "SchemaT",
+    "UsageAccumulator",
+]

atif_models/domain/registry.py

@@ -0,0 +1,151 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""The model alias registry (CONTRACT-V2 §Model registry).
+
+Size aliases (small/medium/large) resolve to concrete Bedrock GLOBAL
+inference-profile ids per family. This module is the ONLY place in the
+atif-sql workspace where a model id is written down — every pipeline
+names a (family, size) pair and lets :func:`resolve` pick the id.
+
+Pricing (USD per 1M tokens, on-demand; Bedrock charges OpenAI/Anthropic
+list parity, so these track the vendors' published rates):
+
+* openai gpt-5.6, post the 2026-07-30 Bedrock price cut: luna $0.20/$1.20,
+  terra $2/$12, sol $5/$30.
+* anthropic: haiku-4-5 $1/$5, sonnet-5 $2/$10, opus-5 $5/$25.
+
+Pure domain: frozen value objects, no I/O, no env. Env-driven family /
+size selection lives in :mod:`atif_models.infrastructure.settings`.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Literal
+
+Family = Literal["openai", "anthropic"]
+Size = Literal["small", "medium", "large"]
+
+#: GPT-5.6 reasoning levels per the Bedrock GA post; the anthropic entries
+#: reuse the same vocabulary.
+ReasoningEffort = Literal["none", "low", "medium", "high", "xhigh", "max"]
+
+DEFAULT_FAMILY: Family = "openai"
+DEFAULT_REASONING_EFFORT: ReasoningEffort = "high"
+DEFAULT_MAX_COMPLETION_TOKENS = 32_000
+
+
+@dataclass(frozen=True, slots=True)
+class ModelSpec:
+    """One concrete model behind a (family, size) alias.
+
+    ``pricing_in`` / ``pricing_out`` are USD per 1M tokens (``None`` when
+    unknown — :func:`estimate_cost` tolerates that by returning ``None``).
+    ``supports_strict_json`` marks native OpenAI strict structured outputs
+    (``response_format.json_schema.strict``); the anthropic entries are
+    ``False`` because Bedrock-Anthropic enforces schemas through
+    ``output_config``, a different contract.
+    """
+
+    family: Family
+    size: Size
+    model_id: str
+    pricing_in: float | None
+    pricing_out: float | None
+    max_completion_tokens: int = DEFAULT_MAX_COMPLETION_TOKENS
+    reasoning_effort: ReasoningEffort = DEFAULT_REASONING_EFFORT
+    supports_strict_json: bool = True
+
+
+REGISTRY: dict[tuple[Family, Size], ModelSpec] = {
+    ("openai", "small"): ModelSpec(
+        family="openai",
+        size="small",
+        model_id="global.openai.gpt-5.6-luna",
+        pricing_in=0.20,
+        pricing_out=1.20,
+    ),
+    ("openai", "medium"): ModelSpec(
+        family="openai",
+        size="medium",
+        model_id="global.openai.gpt-5.6-terra",
+        pricing_in=2.00,
+        pricing_out=12.00,
+    ),
+    ("openai", "large"): ModelSpec(
+        family="openai",
+        size="large",
+        model_id="global.openai.gpt-5.6-sol",
+        pricing_in=5.00,
+        pricing_out=30.00,
+    ),
+    ("anthropic", "small"): ModelSpec(
+        family="anthropic",
+        size="small",
+        model_id="global.anthropic.claude-haiku-4-5",
+        pricing_in=1.00,
+        pricing_out=5.00,
+        supports_strict_json=False,
+    ),
+    ("anthropic", "medium"): ModelSpec(
+        family="anthropic",
+        size="medium",
+        model_id="global.anthropic.claude-sonnet-5",
+        pricing_in=2.00,
+        pricing_out=10.00,
+        supports_strict_json=False,
+    ),
+    ("anthropic", "large"): ModelSpec(
+        family="anthropic",
+        size="large",
+        model_id="global.anthropic.claude-opus-5",
+        pricing_in=5.00,
+        pricing_out=25.00,
+        supports_strict_json=False,
+    ),
+}
+
+
+def resolve(size: Size, family: Family = DEFAULT_FAMILY) -> ModelSpec:
+    """Return the :class:`ModelSpec` behind a (family, size) alias.
+
+    Default family is openai: GPT-5.6 supports native strict structured
+    outputs, and it is the only family with a provider adapter. The
+    anthropic column records model ids and pricing for a future adapter;
+    :data:`atif_models.infrastructure.settings.RUNNABLE_FAMILIES` is the
+    gate that keeps it from being selected before one exists.
+    """
+    try:
+        return REGISTRY[(family, size)]
+    except KeyError:  # pragma: no cover — Literal types make this unreachable from typed code
+        msg = f"no model registered for family={family!r} size={size!r}"
+        raise KeyError(msg) from None
+
+
+def estimate_cost(spec: ModelSpec, *, input_tokens: int, output_tokens: int) -> float | None:
+    """USD estimate for one call (or an accumulated pipeline) on ``spec``.
+
+    Returns ``None`` when either price is unknown — callers must render
+    that as "pricing unavailable", never as $0. Cache-read / reasoning
+    discounts are deliberately NOT modeled here: reasoning tokens bill as
+    output (already inside ``output_tokens`` for the OpenAI usage shape)
+    and cached input still bills at a nonzero rate, so this is an upper
+    bound suitable for the dry-run cost guard.
+    """
+    if spec.pricing_in is None or spec.pricing_out is None:
+        return None
+    return (input_tokens * spec.pricing_in + output_tokens * spec.pricing_out) / 1_000_000
+
+
+__all__ = [
+    "DEFAULT_FAMILY",
+    "DEFAULT_MAX_COMPLETION_TOKENS",
+    "DEFAULT_REASONING_EFFORT",
+    "REGISTRY",
+    "Family",
+    "ModelSpec",
+    "ReasoningEffort",
+    "Size",
+    "estimate_cost",
+    "resolve",
+]

atif_models/domain/schema.py

@@ -0,0 +1,109 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""Pydantic → OpenAI strict-mode JSON Schema transform (CONTRACT-V2 §Client contract).
+
+OpenAI's strict structured outputs (``response_format.json_schema.strict``)
+accept a JSON Schema subset with two hard rules this module enforces:
+
+* ``additionalProperties: false`` on EVERY object level.
+* EVERY property listed in ``required`` — optionality is expressed by
+  making the field's type nullable (``["string", "null"]`` / an ``anyOf``
+  null branch), never by omitting it from ``required``.
+
+``$defs``/``$ref`` are KEPT: OpenAI strict accepts them, including a
+``$defs``-referenced enum, on all three gpt-5.6 sizes. Ref inlining is a
+per-provider wire-format concern, so it belongs in whichever adapter needs
+it rather than here.
+
+``default`` is dropped (it contradicts required-everywhere) and ``title``
+is dropped as noise. Pydantic ``model_validate`` on the response is the
+second gate, so ge/le style constraints still apply even if the model
+ignores them.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from pydantic import BaseModel
+
+#: Keys whose value is a mapping of NAME -> subschema (names are user data,
+#: never schema keywords — a field literally named "default" must survive).
+_MAP_OF_SCHEMAS = ("$defs", "definitions", "properties", "patternProperties")
+#: Keys whose value is a list of subschemas.
+_LIST_OF_SCHEMAS = ("anyOf", "oneOf", "allOf", "prefixItems")
+#: Keys whose value is a single subschema.
+_SINGLE_SCHEMA = ("items", "contains", "propertyNames", "not")
+
+
+def to_openai_strict(schema: type[BaseModel]) -> dict[str, Any]:
+    """Return ``schema.model_json_schema()`` transformed to OpenAI strict mode.
+
+    Raises ``ValueError`` for shapes strict mode cannot represent (an
+    object with a schema-valued ``additionalProperties``, i.e. a
+    ``dict[str, X]`` field) — better a loud failure at body-build time
+    than a Bedrock 400 mid-pipeline.
+    """
+    return _strictify(schema.model_json_schema())
+
+
+def _strictify(node: dict[str, Any]) -> dict[str, Any]:
+    """Recursively apply the strict-mode rules to one schema node."""
+    out: dict[str, Any] = {}
+    for key, val in node.items():
+        if key in ("title", "default"):
+            continue
+        if key in _MAP_OF_SCHEMAS and isinstance(val, dict):
+            out[key] = {name: _strictify(sub) for name, sub in val.items()}
+        elif key in _LIST_OF_SCHEMAS and isinstance(val, list):
+            out[key] = [_strictify(sub) if isinstance(sub, dict) else sub for sub in val]
+        elif key in _SINGLE_SCHEMA and isinstance(val, dict):
+            out[key] = _strictify(val)
+        elif key == "additionalProperties" and isinstance(val, dict):
+            msg = (
+                "OpenAI strict mode cannot represent open mappings "
+                "(schema-valued additionalProperties, e.g. a dict[str, X] field); "
+                "model the keys explicitly or use a list of key/value pairs"
+            )
+            raise ValueError(msg)
+        else:
+            out[key] = val
+    if out.get("type") == "object" or "properties" in out:
+        props: dict[str, Any] = out.get("properties") or {}
+        already_required = set(out.get("required") or [])
+        for name, sub in props.items():
+            if name not in already_required:
+                props[name] = _nullable(sub)
+        if props:
+            out["required"] = list(props)
+        out["additionalProperties"] = False
+    return out
+
+
+def _nullable(sub: dict[str, Any]) -> dict[str, Any]:
+    """Rewrite one property subschema so ``null`` is an accepted value."""
+    if "$ref" in sub:
+        rest = {k: v for k, v in sub.items() if k != "$ref"}
+        return {**rest, "anyOf": [{"$ref": sub["$ref"]}, {"type": "null"}]}
+    if "anyOf" in sub:
+        branches: list[Any] = sub["anyOf"]
+        if any(isinstance(b, dict) and b.get("type") == "null" for b in branches):
+            return sub
+        return {**sub, "anyOf": [*branches, {"type": "null"}]}
+    type_val = sub.get("type")
+    if isinstance(type_val, str):
+        new = {**sub, "type": [type_val, "null"]}
+    elif isinstance(type_val, list):
+        new = sub if "null" in type_val else {**sub, "type": [*type_val, "null"]}
+    else:
+        # No type and no ref/anyOf (e.g. a bare enum or empty schema):
+        # wrap rather than guess.
+        return {"anyOf": [sub, {"type": "null"}]}
+    enum_val = new.get("enum")
+    if isinstance(enum_val, list) and None not in enum_val:
+        new = {**new, "enum": [*enum_val, None]}
+    return new
+
+
+__all__ = ["to_openai_strict"]

atif_models/infrastructure/__init__.py

@@ -0,0 +1,8 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""Infrastructure layer for atif-models.
+
+bedrock-runtime adapters (OpenAI chat-completions default, Anthropic
+``output_config`` experimental stub) and the ``ATIF_SQL_``-prefixed env
+settings. Everything here imports DOWN into ``domain`` only.
+"""