jev-align 0.1.0__py3-none-any.whl

jev_align/__init__.py

@@ -0,0 +1,13 @@
+"""CLI for building AI Functions with pluggable evaluation backends."""
+
+from importlib.metadata import PackageNotFoundError, version
+
+from .capture import Capture
+from .runtime import AIFunction
+
+__all__ = ["AIFunction", "Capture", "__version__"]
+
+try:
+    __version__ = version("jev-align")
+except PackageNotFoundError:
+    __version__ = "0+unknown"

jev_align/acquisition.py

@@ -0,0 +1,80 @@
+from __future__ import annotations
+
+import random
+from dataclasses import dataclass
+from statistics import mean, median
+
+from .models import Prediction
+
+
+def ambiguity(probability: float) -> float:
+    return 1.0 - 2.0 * abs(probability - 0.5)
+
+
+def prediction_ambiguity(prediction: Prediction) -> float:
+    if prediction.label_probabilities is not None:
+        return max(
+            ambiguity(value) for value in prediction.label_probabilities.values()
+        )
+    if prediction.confidence is not None:
+        return 1.0 - prediction.confidence
+    if prediction.probability is None:
+        raise ValueError(
+            "prediction has neither native confidence nor class probability"
+        )
+    return ambiguity(prediction.probability)
+
+
+@dataclass(frozen=True)
+class Acquisition:
+    prediction: Prediction
+    source: str
+
+
+def select_batch(
+    predictions: list[Prediction],
+    *,
+    batch_size: int = 5,
+    exploration_count: int = 1,
+    seed: int = 0,
+) -> list[Acquisition]:
+    if batch_size < 1:
+        raise ValueError("batch_size must be positive")
+    if not 0 <= exploration_count <= batch_size:
+        raise ValueError("exploration_count must be between zero and batch_size")
+    if len(predictions) < batch_size:
+        raise ValueError("not enough unlabeled rows for an acquisition batch")
+
+    ambiguous_count = batch_size - exploration_count
+    ranked = sorted(
+        predictions,
+        key=lambda item: (-prediction_ambiguity(item), item.story_id),
+    )
+    ambiguous = ranked[:ambiguous_count]
+    chosen = {item.story_id for item in ambiguous}
+    remainder = [item for item in predictions if item.story_id not in chosen]
+    rng = random.Random(seed)
+    exploration = rng.sample(remainder, exploration_count)
+    return [
+        *(Acquisition(item, "ambiguous") for item in ambiguous),
+        *(Acquisition(item, "exploration") for item in exploration),
+    ]
+
+
+def ambiguity_summary(predictions: list[Prediction]) -> dict[str, float | int]:
+    values = [prediction_ambiguity(item) for item in predictions]
+    if not values:
+        return {
+            "count": 0,
+            "mean": 0.0,
+            "median": 0.0,
+            "at_least_0_5": 0,
+            "at_least_0_8": 0,
+        }
+    return {
+        "count": len(values),
+        "mean": mean(values),
+        "median": median(values),
+        "at_least_0_5": sum(value >= 0.5 for value in values),
+        "at_least_0_8": sum(value >= 0.8 for value in values),
+    }

jev_align/backends.py

@@ -0,0 +1,101 @@
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from typing import Literal, Protocol
+
+from .models import (
+    BackendConfig,
+    MulticlassCandidateSpec,
+    MultilabelCandidateSpec,
+    Prediction,
+    ScoreCandidateSpec,
+    Story,
+    TaskSpec,
+)
+
+TaskKind = Literal["binary", "multiclass", "multilabel", "score"]
+UncertaintyKind = Literal["calibrated_probability", "native_confidence"]
+
+
+@dataclass(frozen=True)
+class BackendCapabilities:
+    """What an evaluation backend can reliably provide to active learning."""
+
+    task_types: frozenset[TaskKind]
+    uncertainty: Mapping[TaskKind, UncertaintyKind]
+
+
+class EvaluationBackend(Protocol):
+    backend_id: str
+    display_name: str
+    model: str
+    capabilities: BackendCapabilities
+
+    def evaluate_many(
+        self, candidate: TaskSpec, stories: Sequence[Story]
+    ) -> list[Prediction]: ...
+
+
+def create_backend(
+    config: BackendConfig, *, concurrency: int
+) -> EvaluationBackend:
+    """Construct the configured adapter without leaking providers into the core."""
+    if config.provider == "typesafe":
+        from .jev import TypeSafeJevEvaluator
+
+        return TypeSafeJevEvaluator(model=config.model, concurrency=concurrency)
+    raise ValueError(f"unsupported evaluation backend: {config.provider}")
+
+
+def backend_credential_error(
+    config: BackendConfig, environment: Mapping[str, str]
+) -> str | None:
+    if config.provider == "typesafe" and not environment.get("TYPESAFE_API_KEY"):
+        return "TYPESAFE_API_KEY must be set for the TypeSafe backend"
+    return None
+
+
+def backend_display_name(backend: EvaluationBackend) -> str:
+    """Give minimal third-party adapters a useful name without extra boilerplate."""
+    return str(getattr(backend, "display_name", backend.__class__.__name__))
+
+
+def task_kind(candidate: TaskSpec) -> TaskKind:
+    if isinstance(candidate, MultilabelCandidateSpec):
+        return "multilabel"
+    if isinstance(candidate, ScoreCandidateSpec):
+        return "score"
+    if isinstance(candidate, MulticlassCandidateSpec):
+        return "multiclass"
+    return "binary"
+
+
+def validate_backend_for_task(
+    backend: EvaluationBackend, candidate: TaskSpec
+) -> None:
+    """Reject adapters that cannot supply uncertainty for the requested task."""
+    kind = task_kind(candidate)
+    capabilities = backend.capabilities
+    if kind not in capabilities.task_types:
+        raise ValueError(
+            f"{backend_display_name(backend)} does not support {kind} AI Functions"
+        )
+    if kind not in capabilities.uncertainty:
+        raise ValueError(
+            f"{backend_display_name(backend)} does not provide usable uncertainty "
+            f"for {kind} AI Functions"
+        )
+
+
+def evaluate_with_progress(
+    backend: EvaluationBackend,
+    candidate: TaskSpec,
+    stories: Sequence[Story],
+    description: str,
+) -> list[Prediction]:
+    """Use progress when supported, while preserving lightweight test adapters."""
+    method = getattr(backend, "evaluate_many_with_progress", None)
+    if callable(method):
+        return method(candidate, stories, description)
+    return backend.evaluate_many(candidate, stories)

jev_align/capture.py

@@ -0,0 +1,183 @@
+"""Best-effort, bounded local capture. Never makes evaluation requests."""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+import random
+import threading
+from datetime import datetime, timezone
+from pathlib import Path
+from queue import Empty, Full, Queue
+from time import monotonic
+from uuid import uuid4
+
+from .acquisition import prediction_ambiguity
+from .models import Prediction, Story
+
+logger = logging.getLogger(__name__)
+
+
+class Capture:
+    """One bounded queue and background JSONL writer per instance.
+
+    Use as a context manager or call close() at application shutdown. Create
+    instances inside each worker process, after forking. Captures are disposable
+    observations, never human labels or changes to an alignment run.
+    """
+
+    def __init__(
+        self,
+        directory: str | Path = ".jev-align/captures",
+        *,
+        ambiguity_threshold: float = 0.8,
+        audit_rate: float = 0.05,
+        queue_size: int = 256,
+        max_record_bytes: int = 64 * 1024,
+        max_file_bytes: int = 64 * 1024 * 1024,
+    ) -> None:
+        if not 0 <= ambiguity_threshold <= 1 or not 0 <= audit_rate <= 1:
+            raise ValueError("ambiguity_threshold and audit_rate must be in [0, 1]")
+        if min(queue_size, max_record_bytes, max_file_bytes) < 1:
+            raise ValueError("capture size limits must be positive")
+        self.path = Path(directory).resolve() / f"{uuid4().hex}.jsonl"
+        self.ambiguity_threshold = ambiguity_threshold
+        self.audit_rate = audit_rate
+        self.max_record_bytes = max_record_bytes
+        self.max_file_bytes = max_file_bytes
+        self._queue: Queue[bytes] = Queue(maxsize=queue_size)
+        self._lock = threading.Lock()
+        self._stop = threading.Event()
+        self._pid = os.getpid()
+        self._reserved_bytes = 0
+        self._written = 0
+        self._dropped = 0
+        self._error: str | None = None
+        self._thread = threading.Thread(
+            target=self._write_loop, name="jev-align-capture", daemon=True
+        )
+        self._thread.start()
+
+    @property
+    def stats(self) -> dict[str, int | str | None]:
+        """Counters for selected records; filtered-out calls are not drops."""
+        if os.getpid() != self._pid:
+            return {"written": 0, "dropped": 0, "error": "create Capture after fork"}
+        with self._lock:
+            return {
+                "written": self._written,
+                "dropped": self._dropped,
+                "error": self._error,
+            }
+
+    def record(
+        self,
+        story: Story,
+        prediction: Prediction,
+        *,
+        run_id: str,
+        definition: dict[str, object],
+        fingerprint: str,
+    ) -> None:
+        # An inherited writer has no running thread. Avoid its inherited locks.
+        if os.getpid() != self._pid:
+            return
+        uncertainty = prediction_ambiguity(prediction)
+        reason = "ambiguous"
+        if uncertainty < self.ambiguity_threshold:
+            if random.random() >= self.audit_rate:
+                return
+            reason = "audit"
+        with self._lock:
+            if (
+                self._stop.is_set()
+                or self._queue.full()
+                or self._reserved_bytes >= self.max_file_bytes
+                or sum(len(value) for value in story.fields.values())
+                > self.max_record_bytes
+            ):
+                self._dropped += 1
+                return
+        payload = {
+            "version": 1,
+            "created_at": datetime.now(timezone.utc).isoformat(),
+            "run_id": run_id,
+            "definition": definition,
+            "fingerprint": fingerprint,
+            "input": story.fields,
+            "prediction": prediction.model_dump(mode="json"),
+            "ambiguity": uncertainty,
+            "reason": reason,
+        }
+        line = (json.dumps(payload) + "\n").encode("utf-8")
+        with self._lock:
+            if (
+                self._stop.is_set()
+                or len(line) > self.max_record_bytes
+                or self._reserved_bytes + len(line) > self.max_file_bytes
+            ):
+                self._dropped += 1
+                return
+            try:
+                self._queue.put_nowait(line)
+            except Full:
+                self._dropped += 1
+            else:
+                self._reserved_bytes += len(line)
+
+    def _write_loop(self) -> None:
+        batch: list[bytes] = []
+        try:
+            self.path.parent.mkdir(parents=True, exist_ok=True)
+            with self.path.open("xb") as handle:
+                deadline = monotonic() + 0.25
+                while not self._stop.is_set() or not self._queue.empty():
+                    try:
+                        batch.append(
+                            self._queue.get(timeout=max(0, deadline - monotonic()))
+                        )
+                    except Empty:
+                        pass
+                    if len(batch) >= 64 or monotonic() >= deadline:
+                        handle.write(b"".join(batch))
+                        handle.flush()
+                        with self._lock:
+                            self._written += len(batch)
+                        batch.clear()
+                        deadline = monotonic() + 0.25
+                handle.write(b"".join(batch))
+                handle.flush()
+                with self._lock:
+                    self._written += len(batch)
+                batch.clear()
+        except OSError as error:
+            with self._lock:
+                self._error = str(error)
+                self._stop.set()
+                self._dropped += len(batch)
+                while True:
+                    try:
+                        self._queue.get_nowait()
+                        self._dropped += 1
+                    except Empty:
+                        break
+            logger.warning("JEV capture disabled after a local write error: %s", error)
+
+    def close(self, timeout: float = 5.0) -> None:
+        """Stop accepting records and wait up to timeout seconds for the writer.
+
+        Buffered records may be lost on abrupt exit or if shutdown times out.
+        No fsync or durability guarantee is added to the request path.
+        """
+        if os.getpid() != self._pid:
+            return
+        with self._lock:
+            self._stop.set()
+        self._thread.join(timeout=timeout)
+
+    def __enter__(self) -> Capture:
+        return self
+
+    def __exit__(self, *_args: object) -> None:
+        self.close()

jev_align/captured_inputs.py

@@ -0,0 +1,127 @@
+"""Read deduplicated captured inputs without model calls."""
+
+from __future__ import annotations
+
+import hashlib
+import heapq
+import json
+import logging
+import os
+from collections.abc import Iterable
+from pathlib import Path
+
+from .models import RunState, Story
+
+logger = logging.getLogger(__name__)
+
+
+def captured_story_id(fields: dict[str, str]) -> str:
+    payload = json.dumps(fields, sort_keys=True).encode("utf-8")
+    return "capture-" + hashlib.sha256(payload).hexdigest()
+
+
+def sample_captured_inputs(
+    directories: Iterable[Path],
+    state: RunState,
+    *,
+    excluded: Iterable[Story],
+    saved: Iterable[Story] = (),
+    required: Iterable[Story] = (),
+    limit: int | None = None,
+) -> list[Story]:
+    """Keep required inputs plus all other unique inputs by default.
+
+    An optional limit is retained for callers that explicitly request sampling;
+    required inputs do not count toward it. Stable hash priorities avoid
+    favoring frequent duplicates or the first file scanned when a limit is used.
+    Read only complete lines present when each file is opened, so active writers
+    cannot extend the scan.
+    """
+    if limit is not None and limit < 1:
+        raise ValueError("capture sample limit must be positive")
+    excluded_ids = {captured_story_id(story.fields) for story in excluded}
+    required_by_id = {
+        story.id: story for story in required if story.id not in excluded_ids
+    }
+    selected: dict[str, Story] = {}
+    heap: list[tuple[int, str]] = []
+
+    def consider(story: Story) -> None:
+        if (
+            story.id in excluded_ids
+            or story.id in required_by_id
+            or story.id in selected
+        ):
+            return
+        if limit is None:
+            selected[story.id] = story
+            return
+        priority = int(
+            hashlib.sha256(
+                f"{state.seed}:{state.round_number}:{story.id}".encode()
+            ).hexdigest(),
+            16,
+        )
+        if len(heap) == limit:
+            if priority >= -heap[0][0]:
+                return
+            _, removed = heapq.heappop(heap)
+            del selected[removed]
+        heapq.heappush(heap, (-priority, story.id))
+        selected[story.id] = story
+
+    for story in saved:
+        consider(story)
+    paths = {path for directory in directories for path in directory.glob("*.jsonl")}
+    for path in sorted(paths):
+        try:
+            with path.open("rb") as handle:
+                boundary = os.fstat(handle.fileno()).st_size
+                oversized = False
+                while handle.tell() < boundary:
+                    line = handle.readline(min(1024 * 1024, boundary - handle.tell()))
+                    if not line:
+                        break
+                    if not line.endswith(b"\n"):
+                        oversized = True
+                        continue
+                    if oversized:
+                        oversized = False
+                        continue
+                    try:
+                        row = json.loads(line)
+                    except (ValueError, UnicodeError):
+                        continue
+                    if not isinstance(row, dict) or row.get("version") != 1:
+                        continue
+                    if row.get("run_id") != state.run_id:
+                        continue
+                    fields = row.get("input")
+                    if (
+                        not isinstance(fields, dict)
+                        or not fields
+                        or not all(
+                            isinstance(key, str) and isinstance(value, str)
+                            for key, value in fields.items()
+                        )
+                    ):
+                        continue
+                    expected = (
+                        ["content"]
+                        if state.column_mode == "all_concatenated"
+                        else state.selected_columns
+                    )
+                    if expected is not None and set(fields) != set(expected):
+                        continue
+                    consider(
+                        Story(id=captured_story_id(fields), row_number=0, fields=fields)
+                    )
+        except OSError as error:
+            logger.warning("Could not read capture file %s: %s", path, error)
+    selected_ids = (
+        sorted(selected) if limit is None else [story_id for _, story_id in sorted(heap)]
+    )
+    return [
+        *(required_by_id[story_id] for story_id in sorted(required_by_id)),
+        *(selected[story_id] for story_id in selected_ids),
+    ]