From 98fa4ec8cdefd764f98fca858d1307343ee6e96c Mon Sep 17 00:00:00 2001 From: Marius Ciepluch <11855163+norandom@users.noreply.github.com> Date: Thu, 1 Oct 2026 09:56:52 +0200 Subject: [PATCH] feat(learn): metered embeddings and the cache (decision-engine 6.2) plane/models/openrouter-embedding.yaml: openai/text-embedding-3-small via OpenRouter (embeddings only), input_per_m 0.02, cache/output 0, egress openrouter.ai; `openrouter` joins the provider extensions. plane/budget.py MeteredEmbeddingClient: prompt_tokens priced from the Model manifest, usd/calls ceilings stop the next call, a priced call that reports no tokens raises UnmeteredUsage (a silent zero is an instrument failure). OpenRouterEmbeddingClient keeps one usage row per call (parse_embedding_payload); its api key is out of repr. learn/embeddings.py + `ousast learn memory embed`: excerpts embedded in batches into embeddings//.json (float32 base64), cached by content hash so a rerun makes no call; the declared dimension is checked; --dry-run estimates tokens and cost without a call; a paid run needs --budget-usd and exits 3 (unfinished) at its ceiling. Not run against the provider: the store embedding waits for its budget go-ahead. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01EmM1j46qgKMh3zqD1QMwEo --- .kiro/specs/learned-decision-engine/tasks.md | 2 +- .../measurements/2026-09-08-module-audit.json | 21 ++- plane/models/openrouter-embedding.yaml | 24 +++ src/openultrasast/learn/cli.py | 48 +++++ src/openultrasast/learn/embeddings.py | 172 ++++++++++++++++++ src/openultrasast/plane/budget.py | 87 ++++++++- src/openultrasast/plane/manifests.py | 2 +- src/openultrasast/provider/openrouter.py | 16 +- tests/test_learn_embeddings.py | 133 ++++++++++++++ tests/test_plane_increment.py | 2 +- tests/test_plane_manifests.py | 2 +- 11 files changed, 498 insertions(+), 11 deletions(-) create mode 100644 plane/models/openrouter-embedding.yaml create mode 100644 src/openultrasast/learn/embeddings.py create mode 100644 tests/test_learn_embeddings.py diff --git a/.kiro/specs/learned-decision-engine/tasks.md b/.kiro/specs/learned-decision-engine/tasks.md index 47431fa..92db26c 100644 --- a/.kiro/specs/learned-decision-engine/tasks.md +++ b/.kiro/specs/learned-decision-engine/tasks.md @@ -46,7 +46,7 @@ a model or embedding call. 10% of a source has no excerpt); `example` added to `KINDS`; `put_blob`/`get_blob` on both backends. - Tests: `test_learn_excerpt.py`, example-builder refusals, blob round trip on FileStore (MinIO gated). - _Requirements: 3.2, 3.6, 7.2_ - - [ ] 6.2 Metered embeddings and the cache + - [x] 6.2 Metered embeddings and the cache - `plane/models/openrouter-embedding.yaml` (`input_per_m: 0.02`), `MeteredEmbeddingClient` in `plane/budget.py`, usage kept by `parse_embedding_response`; cache `embeddings//.json`; `ousast learn memory embed`; scripted-client tests. diff --git a/benchmarks/measurements/2026-09-08-module-audit.json b/benchmarks/measurements/2026-09-08-module-audit.json index 806da5d..c79105a 100644 --- a/benchmarks/measurements/2026-09-08-module-audit.json +++ b/benchmarks/measurements/2026-09-08-module-audit.json @@ -1,7 +1,7 @@ { - "modules": 146, + "modules": 147, "counts": { - "load_bearing": 139, + "load_bearing": 140, "standalone_capability": 7 }, "orphaned": [], @@ -353,6 +353,7 @@ "importers": [ "cli", "learn.cli", + "learn.embeddings", "learn.examples", "learn.excerpt", "learn.features", @@ -370,6 +371,14 @@ "cli" ] }, + { + "module": "learn.embeddings", + "classification": "load_bearing", + "reason": "reachable from gate, map_gate, pair_gate or the CLI", + "importers": [ + "learn.cli" + ] + }, { "module": "learn.examples", "classification": "load_bearing", @@ -383,6 +392,7 @@ "classification": "load_bearing", "reason": "reachable from gate, map_gate, pair_gate or the CLI", "importers": [ + "learn.embeddings", "learn.examples" ] }, @@ -752,6 +762,7 @@ "cli", "improve.memory", "learn.cli", + "learn.embeddings", "learn.examples", "plane.egress", "plane.engine_alerts", @@ -772,6 +783,8 @@ "classification": "load_bearing", "reason": "ai-service-plane Req 2.3/2.4: metered chat client with usd and calls ceilings; task integration in tasks 3, 6", "importers": [ + "learn.cli", + "learn.embeddings", "plane.generate", "plane.tasks.roles", "plane.tasks.verify" @@ -812,6 +825,7 @@ "classification": "load_bearing", "reason": "ai-service-plane Req 1.1-1.4: ax-shaped Task/Workspace/Model and the Run kind, validated before anything runs", "importers": [ + "learn.cli", "plane.egress", "plane.generate", "plane.reconciler", @@ -827,6 +841,7 @@ "importers": [ "improve.memory", "learn.cli", + "learn.embeddings", "learn.examples", "plane.generate", "plane.tasks.loop", @@ -1036,6 +1051,7 @@ "classification": "load_bearing", "reason": "reachable from gate, map_gate, pair_gate or the CLI", "importers": [ + "learn.cli", "model.endpoint", "tool_hunter" ] @@ -1045,6 +1061,7 @@ "classification": "load_bearing", "reason": "reachable from gate, map_gate, pair_gate or the CLI", "importers": [ + "learn.cli", "model.endpoint", "tool_hunter" ] diff --git a/plane/models/openrouter-embedding.yaml b/plane/models/openrouter-embedding.yaml new file mode 100644 index 0000000..2479cce --- /dev/null +++ b/plane/models/openrouter-embedding.yaml @@ -0,0 +1,24 @@ +# The embedding model of the decision engine's local memory (learned-decision-engine design 4.1-4.2, Req 3.2). +# OpenRouter serves embeddings only; the chat model stays DeepSeek. `openrouter` is not an ax provider, so the +# annotation declares the extension. `secretKey` names the Kubernetes Secret and key that hold OPENROUTER_API_KEY; +# the key itself is never in a manifest. Prices in USD per million tokens: embeddings bill input only, so the cache +# and output rates are 0 and `plane.budget.MeteredEmbeddingClient` prices the provider's `usage.prompt_tokens`. +apiVersion: ax.io/v1alpha1 +kind: Model +metadata: + name: openrouter-embedding + atespace: default # ax's own name for the namespace field (plane/manifests.py) + annotations: + openultrasast.io/provider-extension: "true" + openultrasast.io/egress-hosts: "openrouter.ai" # provider.openrouter base URL; the embed task's only model egress +spec: + provider: openrouter + model: openai/text-embedding-3-small + secretKey: + name: openrouter + key: OPENROUTER_API_KEY + parameters: + cache_hit_per_m: 0 + input_per_m: 0.02 + output_per_m: 0 + dimensions: 1536 diff --git a/src/openultrasast/learn/cli.py b/src/openultrasast/learn/cli.py index 5b2f72f..f792359 100644 --- a/src/openultrasast/learn/cli.py +++ b/src/openultrasast/learn/cli.py @@ -26,15 +26,63 @@ def add_commands(learn_sub: Any) -> None: build.add_argument("--memory", help="the store (default OUSAST_MEMORY, then results_root()/memory)") build.add_argument("--profile", choices=("static", "plane"), default="static") build.add_argument("--dry-run", action="store_true", help="count, write nothing") + embed = memory_sub.add_parser("embed", help="embed the stored excerpts (OpenRouter, metered; cached by excerpt sha)") + embed.add_argument("--memory", help="the store (default OUSAST_MEMORY, then results_root()/memory)") + embed.add_argument("--model", default="openai/text-embedding-3-small") + embed.add_argument("--manifest", type=Path, default=Path("plane/models/openrouter-embedding.yaml"), help="the Model with the prices") + embed.add_argument("--budget-usd", type=float, help="the ceiling; required unless --dry-run") + embed.add_argument("--budget-calls", type=int) + embed.add_argument("--batch", type=int, default=32) + embed.add_argument("--dry-run", action="store_true", help="count and estimate the tokens; no call") def run(args: argparse.Namespace) -> int: if args.learn_command == "memory" and args.memory_command == "build": return _memory_build(args) + if args.learn_command == "memory" and args.memory_command == "embed": + return _memory_embed(args) print(f"learn: unknown command {args.learn_command}", file=sys.stderr) return 2 +def _model_parameters(path: Path) -> dict[str, Any]: + from ..plane.manifests import load_manifests + + models = load_manifests([path]).models + if len(models) != 1: + raise SystemExit(f"{path}: expected one Model, found {sorted(models)}") + return dict(next(iter(models.values())).parameters) + + +def _memory_embed(args: argparse.Namespace) -> int: + """Exit 2 without a ceiling, 3 when the ceiling stops the run (``unfinished``; a rerun resumes from the cache).""" + from ..plane.budget import BudgetExhausted, MeteredEmbeddingClient + from ..plane.memory import open_store + from .embeddings import embed_excerpts + from .examples import EXAMPLE_KIND + + store = open_store(args.memory) + shas = sorted({str(r.row["excerpt_sha"]) for r in store.rows(kind=EXAMPLE_KIND)}) + client = None + if not args.dry_run: + if args.budget_usd is None: + print("learn memory embed: --budget-usd is required (a paid run starts only under a ceiling)", file=sys.stderr) + return 2 + from ..provider.openrouter import OpenRouterEmbeddingClient + + client = MeteredEmbeddingClient( + OpenRouterEmbeddingClient.from_env(), prices=_model_parameters(args.manifest), budget_usd=args.budget_usd, + budget_calls=args.budget_calls, + ) # fmt: skip + try: + report = embed_excerpts(store, shas, client, model=args.model, batch=args.batch) + except BudgetExhausted as exc: + print(f"learn memory embed: unfinished -- {exc}", file=sys.stderr) + return 3 + print(json.dumps(report.as_dict(), indent=2, sort_keys=True)) + return 0 + + def _memory_build(args: argparse.Namespace) -> int: """Exit 2 on a refused label (``origin: user``, a source outside the list) or a source with >10% unread excerpts.""" from ..plane.memory import open_store diff --git a/src/openultrasast/learn/embeddings.py b/src/openultrasast/learn/embeddings.py new file mode 100644 index 0000000..d408307 --- /dev/null +++ b/src/openultrasast/learn/embeddings.py @@ -0,0 +1,172 @@ +"""Code embeddings of excerpts, metered and cached by content (learned-decision-engine design 4.2, Req 3.2). + +OpenRouter serves embeddings only (``openai/text-embedding-3-small``, 1,536 dimensions, $0.02 per million input +tokens; ``plane/models/openrouter-embedding.yaml``). Every call goes through +:class:`..plane.budget.MeteredEmbeddingClient`, so a ceiling stops the next batch and a call reporting no tokens is +an error, not a free call. The cache key is the excerpt's sha (the content-addressed blob ``excerpts/.txt``) +per embedding model: ``embeddings//.json`` holds ``{model, dimensions, vector (float32, +base64, little-endian), tokens, created}``. An excerpt is embedded once; a rerun makes no call. +""" + +from __future__ import annotations + +import base64 +import json +import math +import re +import struct +from collections.abc import Iterable, Sequence +from dataclasses import asdict, dataclass +from datetime import date +from typing import Any + +from ..plane.budget import EmbeddingClient +from ..plane.memory import MemoryStore +from .excerpt import excerpt_sha, normalise + +DEFAULT_MODEL = "openai/text-embedding-3-small" +DEFAULT_DIMENSIONS = 1536 +INPUT_PER_M = 0.02 +CHARS_PER_TOKEN = 4.0 # an estimate for dry runs only; metered runs use the provider's usage + + +class EmbeddingError(RuntimeError): + """A provider answer of the wrong shape: a vector count or dimension that is not what was asked.""" + + +def model_slug(model: str) -> str: + return re.sub(r"[^A-Za-z0-9._-]+", "__", model) + + +def encode_vector(vector: Sequence[float]) -> str: + return base64.b64encode(struct.pack(f"<{len(vector)}f", *vector)).decode("ascii") + + +def decode_vector(text: str) -> list[float]: + data = base64.b64decode(text) + return list(struct.unpack(f"<{len(data) // 4}f", data)) + + +def cosine(a: Sequence[float], b: Sequence[float]) -> float: + if len(a) != len(b) or not a: + raise EmbeddingError(f"cosine of vectors of dimension {len(a)} and {len(b)}") + dot = sum(x * y for x, y in zip(a, b, strict=True)) + norm = math.sqrt(sum(x * x for x in a)) * math.sqrt(sum(y * y for y in b)) + return dot / norm if norm else 0.0 + + +def estimate_tokens(text: str) -> int: + return max(1, round(len(text) / CHARS_PER_TOKEN)) + + +class EmbeddingCache: + """The vectors of one embedding model in the store, by excerpt sha.""" + + def __init__(self, store: MemoryStore, model: str = DEFAULT_MODEL) -> None: + self.store, self.model = store, model + self.prefix = f"embeddings/{model_slug(model)}" + + def get(self, sha: str) -> list[float] | None: + data = self.store.get_blob(self.prefix, sha, verify=False) + if data is None: + return None + entry = json.loads(data) + if entry.get("model") != self.model: + raise EmbeddingError(f"{self.prefix}/{sha}: stored for model {entry.get('model')!r}, not {self.model!r}") + return decode_vector(str(entry["vector"])) + + def put(self, sha: str, vector: Sequence[float], *, tokens: int, created: str) -> None: + entry = {"created": created, "dimensions": len(vector), "model": self.model, "tokens": tokens, "vector": encode_vector(vector)} + self.store.put_blob(self.prefix, json.dumps(entry, sort_keys=True).encode("utf-8"), name=sha) + + def names(self) -> set[str]: + return set(self.store.blob_names(self.prefix)) + + +@dataclass +class EmbedReport: + requested: int = 0 + cached: int = 0 + embedded: int = 0 + missing_excerpt: int = 0 + calls: int = 0 + tokens: int = 0 + estimated_tokens: int = 0 + usd: float | None = None + + def as_dict(self) -> dict[str, Any]: + return asdict(self) + + +def embed_excerpts( + store: MemoryStore, + shas: Iterable[str], + client: EmbeddingClient | None, + *, + model: str = DEFAULT_MODEL, + dimensions: int | None = DEFAULT_DIMENSIONS, + batch: int = 32, + created: str | None = None, +) -> EmbedReport: + """Embed every excerpt in ``shas`` that has no cached vector, in batches. ``client=None`` is a dry run: it counts + what would be embedded and estimates the tokens (characters / 4), making no call.""" + cache = EmbeddingCache(store, model) + have = cache.names() + report = EmbedReport() + todo: list[tuple[str, str]] = [] + for sha in dict.fromkeys(shas): + report.requested += 1 + if sha in have: + report.cached += 1 + continue + data = store.get_blob("excerpts", sha) + if data is None: + report.missing_excerpt += 1 + continue + text = data.decode("utf-8") + report.estimated_tokens += estimate_tokens(text) + todo.append((sha, text)) + if client is None: + report.usd = report.estimated_tokens * INPUT_PER_M / 1_000_000 + return report + when = created or date.today().isoformat() + tokens_before = int(getattr(client, "tokens", 0) or 0) + for start in range(0, len(todo), batch): + chunk = todo[start : start + batch] + seen = int(getattr(client, "tokens", 0) or 0) + vectors = client.embed(model=model, inputs=[text for _, text in chunk]) + report.calls += 1 + if len(vectors) != len(chunk): + raise EmbeddingError(f"asked for {len(chunk)} embeddings, the provider returned {len(vectors)}") + spent = int(getattr(client, "tokens", 0) or 0) - seen + total_chars = sum(len(text) for _, text in chunk) or 1 + for (sha, text), vector in zip(chunk, vectors, strict=True): + if dimensions is not None and len(vector) != dimensions: + raise EmbeddingError(f"embedding of dimension {len(vector)}, declared {dimensions}: the model or route changed") + cache.put(sha, vector, tokens=round(spent * len(text) / total_chars), created=when) + report.embedded += 1 + report.tokens = int(getattr(client, "tokens", 0) or 0) - tokens_before + usd = getattr(client, "usd", None) + report.usd = float(usd) if isinstance(usd, int | float) else None + return report + + +def vector_for(text: str, cache: EmbeddingCache, client: EmbeddingClient | None, *, created: str | None = None) -> list[float] | None: + """The vector of a candidate's excerpt (cached by its sha), ``None`` without a client and no cached vector.""" + sha = excerpt_sha(text) + found = cache.get(sha) + if found is not None or client is None: + return found + seen = int(getattr(client, "tokens", 0) or 0) + vectors = client.embed(model=cache.model, inputs=[normalise(text)]) + if len(vectors) != 1: + raise EmbeddingError(f"asked for 1 embedding, the provider returned {len(vectors)}") + spent = int(getattr(client, "tokens", 0) or 0) - seen + cache.put(sha, vectors[0], tokens=spent, created=created or date.today().isoformat()) + return vectors[0] + + +__all__ = [ + "DEFAULT_DIMENSIONS", "DEFAULT_MODEL", "EmbedReport", "EmbeddingCache", "EmbeddingError", "cosine", "decode_vector", + "embed_excerpts", "encode_vector", "estimate_tokens", "model_slug", "vector_for", +] # fmt: skip diff --git a/src/openultrasast/plane/budget.py b/src/openultrasast/plane/budget.py index 7db2e68..5a95ceb 100644 --- a/src/openultrasast/plane/budget.py +++ b/src/openultrasast/plane/budget.py @@ -10,7 +10,7 @@ failure is `AccountError` (the task is `failed` and the run starts nothing furth from __future__ import annotations from collections.abc import Mapping -from typing import Any +from typing import Any, Protocol from ..model.endpoint import Prices from ..tool_hunter import ChatClient, ChatResponse @@ -160,4 +160,87 @@ def _tokens(value: object) -> int: return int(value) if isinstance(value, int | float) and not isinstance(value, bool) else 0 -__all__ = ["AccountError", "BudgetExhausted", "MeteredClient", "cost_of", "is_account_error", "prices_from"] +class EmbeddingClient(Protocol): + def embed(self, *, model: str, inputs: list[str], timeout_seconds: int = 60) -> list[list[float]]: ... + + +class UnmeteredUsage(RuntimeError): + """A priced embedding call whose provider reported no tokens: a working-looking zero, never a free call.""" + + +class MeteredEmbeddingClient: + """An embedding client metered like :class:`MeteredClient`: the wrapped client records one usage dict per call + in a ``usage`` list (``OpenRouterEmbeddingClient`` does), ``prompt_tokens`` are priced at the Model manifest's + ``input_per_m``, and the usd/calls ceilings stop the next call. A priced call that reports no tokens raises + :class:`UnmeteredUsage` (the 2026-09 lesson: a silent zero is an instrument failure).""" + + def __init__( + self, + inner: EmbeddingClient, + *, + prices: Mapping[str, object] | Prices | None = None, + budget_usd: float | None = None, + budget_calls: int | None = None, + require_usage: bool = True, + ) -> None: + self.inner = inner + self.prices = prices_from(prices) + if budget_usd is not None and self.prices is None: + raise ValueError("unpriced model cannot run under a usd budget") + self.budget_usd, self.budget_calls, self.require_usage = budget_usd, budget_calls, require_usage + self.calls = 0 + self.tokens = 0 + self._usd = 0.0 + self._seen_rows = len(self._inner_rows()) + + @property + def usd(self) -> float | None: + return self._usd if self.prices is not None else None + + def check(self) -> None: + if self.budget_usd is not None and self._usd >= self.budget_usd: + raise BudgetExhausted(f"usd budget exhausted: spent {self._usd:.6f} of {self.budget_usd:.6f}", usd=self._usd, calls=self.calls) + if self.budget_calls is not None and self.calls >= self.budget_calls: + raise BudgetExhausted(f"calls budget exhausted: {self.calls} of {self.budget_calls}", usd=self.usd, calls=self.calls) + + def embed(self, *, model: str, inputs: list[str], timeout_seconds: int = 60) -> list[list[float]]: + self.check() + try: + vectors = self.inner.embed(model=model, inputs=inputs, timeout_seconds=timeout_seconds) + except Exception as exc: + self._absorb_rows() + if is_account_error(exc): + raise AccountError(str(exc)) from exc + raise + self.calls += 1 + tokens = self._absorb_rows() + if self.require_usage and self.prices is not None and tokens == 0 and inputs: + raise UnmeteredUsage( + f"embedding call {self.calls} reported no prompt tokens for {len(inputs)} inputs: the meter cannot price it" + ) + return vectors + + def summary(self) -> dict[str, Any]: + return {"usd": self.usd, "calls": self.calls, "usage": {"prompt_tokens": self.tokens}, "priced": self.prices is not None} + + def _inner_rows(self) -> list[Any]: + rows = getattr(self.inner, "usage", None) + return rows if isinstance(rows, list) else [] + + def _absorb_rows(self) -> int: + rows = self._inner_rows() + tokens = 0 + for row in rows[self._seen_rows :]: + if isinstance(row, Mapping): + tokens += _tokens(row.get("prompt_tokens") or row.get("total_tokens")) + self._seen_rows = len(rows) + self.tokens += tokens + if self.prices is not None: + self._usd += tokens * self.prices.input_per_m / 1_000_000 + return tokens + + +__all__ = [ + "AccountError", "BudgetExhausted", "EmbeddingClient", "MeteredClient", "MeteredEmbeddingClient", "UnmeteredUsage", "cost_of", + "is_account_error", "prices_from", +] # fmt: skip diff --git a/src/openultrasast/plane/manifests.py b/src/openultrasast/plane/manifests.py index 3e448be..185c48b 100644 --- a/src/openultrasast/plane/manifests.py +++ b/src/openultrasast/plane/manifests.py @@ -59,7 +59,7 @@ GIT_COMMITS_ANNOTATION = "openultrasast.io/git-commits" EGRESS_HOSTS_ANNOTATION = "openultrasast.io/egress-hosts" # the Model API hosts a bound task may reach _DNS_NAME = re.compile(r"^(?=.{1,253}$)[a-z0-9]([a-z0-9-]{0,61}[a-z0-9])?(\.[a-z0-9]([a-z0-9-]{0,61}[a-z0-9])?)*$") AX_PROVIDERS = frozenset({"google", "anthropic"}) -EXTENSION_PROVIDERS = frozenset({"deepseek"}) +EXTENSION_PROVIDERS = frozenset({"deepseek", "openrouter"}) # openrouter: embeddings only (decision engine) _API_VERSIONS = {"Task": AX_API_VERSION, "Workspace": AX_API_VERSION, "Model": AX_API_VERSION, "Run": RUN_API_VERSION} diff --git a/src/openultrasast/provider/openrouter.py b/src/openultrasast/provider/openrouter.py index 1c8436b..3b538e3 100644 --- a/src/openultrasast/provider/openrouter.py +++ b/src/openultrasast/provider/openrouter.py @@ -8,7 +8,7 @@ import time import urllib.error import urllib.request from collections.abc import Callable, Mapping, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import TypeVar, cast @@ -135,10 +135,11 @@ class OpenRouterChatClient: @dataclass(frozen=True) class OpenRouterEmbeddingClient: - api_key: str + api_key: str = field(repr=False) base_url: str = "https://openrouter.ai/api/v1" max_attempts: int = 3 retry_base_delay: float = 0.5 + usage: list[dict[str, object]] = field(default_factory=list, compare=False, repr=False) # one per call, for the meter @classmethod def from_env(cls) -> OpenRouterEmbeddingClient: @@ -167,7 +168,16 @@ class OpenRouterEmbeddingClient: response_payload = call_with_retry(_do, attempts=self.max_attempts, base_delay=self.retry_base_delay) except (urllib.error.URLError, TimeoutError, http.client.HTTPException, json.JSONDecodeError) as exc: raise OpenRouterError(f"OpenRouter embedding request failed: {type(exc).__name__}: {exc}") from exc - return parse_embedding_response(response_payload) + vectors, usage = parse_embedding_payload(response_payload) + self.usage.append(usage) + return vectors + + +def parse_embedding_payload(payload: object) -> tuple[list[list[float]], dict[str, object]]: + """The vectors and the provider's ``usage`` (``{}`` when absent: the meter treats that as unmetered, not free).""" + vectors = parse_embedding_response(payload) + usage = payload.get("usage") if isinstance(payload, dict) else None + return vectors, dict(usage) if isinstance(usage, dict) else {} def parse_embedding_response(payload: object) -> list[list[float]]: diff --git a/tests/test_learn_embeddings.py b/tests/test_learn_embeddings.py new file mode 100644 index 0000000..6479c45 --- /dev/null +++ b/tests/test_learn_embeddings.py @@ -0,0 +1,133 @@ +"""Metered embeddings and the cache (learned-decision-engine task 6.2). Every client is scripted: no embedding call +leaves the process.""" + +from __future__ import annotations + +import json +from collections.abc import Iterator +from pathlib import Path + +import pytest +from test_plane_memory import FakeObjects + +from openultrasast.learn.embeddings import ( + EmbeddingCache, + EmbeddingError, + cosine, + decode_vector, + embed_excerpts, + encode_vector, + model_slug, + vector_for, +) +from openultrasast.learn.excerpt import excerpt_sha +from openultrasast.plane.budget import BudgetExhausted, MeteredEmbeddingClient, UnmeteredUsage +from openultrasast.plane.manifests import load_manifests +from openultrasast.plane.memory import FileStore, MemoryStore, MinioStore +from openultrasast.provider.openrouter import parse_embedding_payload + +PRICES = {"cache_hit_per_m": 0, "input_per_m": 0.02, "output_per_m": 0} + + +class ScriptedEmbedder: + """Answers each input with a 4-dimensional vector from its length; records usage like OpenRouterEmbeddingClient.""" + + def __init__(self, *, dims: int = 4, tokens_per_input: int = 100, report_usage: bool = True) -> None: + self.dims, self.tokens_per_input, self.report_usage = dims, tokens_per_input, report_usage + self.calls: list[list[str]] = [] + self.usage: list[dict[str, object]] = [] + + def embed(self, *, model: str, inputs: list[str], timeout_seconds: int = 60) -> list[list[float]]: + self.calls.append(list(inputs)) + self.usage.append({"prompt_tokens": self.tokens_per_input * len(inputs)} if self.report_usage else {}) + return [[float(len(text) % 7 + i) for i in range(self.dims)] for text in inputs] + + +@pytest.fixture(params=["file", "fake-minio"]) +def store(request: pytest.FixtureRequest, tmp_path: Path) -> Iterator[MemoryStore]: + yield FileStore(tmp_path / "memory") if request.param == "file" else MinioStore(FakeObjects("works"), "bucket", "p") + + +def _excerpts(store: MemoryStore, n: int) -> list[str]: + return [store.put_blob("excerpts", f" 1 def f{i}():\n 2 return {i}\n".encode()) for i in range(n)] + + +def test_embedding_cost_is_metered_from_the_provider_usage(store: MemoryStore) -> None: + shas = _excerpts(store, 5) + inner = ScriptedEmbedder(dims=4) + meter = MeteredEmbeddingClient(inner, prices=PRICES, budget_usd=1.0) + report = embed_excerpts(store, shas, meter, dimensions=4, batch=2, created="2026-10-01") + assert (report.embedded, report.calls, report.cached) == (5, 3, 0) and len(inner.calls) == 3 + assert report.tokens == 500 and meter.tokens == 500 and report.usd == pytest.approx(500 * 0.02 / 1e6) + assert meter.summary()["usage"] == {"prompt_tokens": 500} + + +def test_cache_by_content_hash_makes_a_rerun_free(store: MemoryStore) -> None: + shas = _excerpts(store, 3) + inner = ScriptedEmbedder() + embed_excerpts(store, shas, MeteredEmbeddingClient(inner, prices=PRICES), dimensions=4) + again = MeteredEmbeddingClient(inner, prices=PRICES) + report = embed_excerpts(store, shas + shas, again, dimensions=4) + assert report.cached == 3 and report.embedded == 0 and len(inner.calls) == 1 and again.calls == 0 and again.usd == 0 + cache = EmbeddingCache(store) + assert cache.prefix == "embeddings/openai__text-embedding-3-small" and cache.names() == set(shas) + stored = json.loads(store.get_blob(cache.prefix, shas[0], verify=False) or b"{}") + assert stored["dimensions"] == 4 and stored["model"] == "openai/text-embedding-3-small" and stored["tokens"] == 100 + + +def test_dry_run_estimates_without_a_call(store: MemoryStore) -> None: + shas = _excerpts(store, 2) + report = embed_excerpts(store, [*shas, "0" * 64], None) + assert report.embedded == 0 and report.calls == 0 and report.missing_excerpt == 1 and report.estimated_tokens > 0 + assert report.usd == pytest.approx(report.estimated_tokens * 0.02 / 1e6) + + +def test_ceiling_stops_the_next_call_and_a_rerun_resumes(store: MemoryStore) -> None: + shas = _excerpts(store, 4) + inner = ScriptedEmbedder() + with pytest.raises(BudgetExhausted): + embed_excerpts(store, shas, MeteredEmbeddingClient(inner, prices=PRICES, budget_calls=1), dimensions=4, batch=2) + assert len(EmbeddingCache(store).names()) == 2 + report = embed_excerpts(store, shas, MeteredEmbeddingClient(inner, prices=PRICES), dimensions=4, batch=2) + assert report.cached == 2 and report.embedded == 2 + + +def test_a_zero_token_answer_is_an_instrument_failure(store: MemoryStore) -> None: + shas = _excerpts(store, 1) + with pytest.raises(UnmeteredUsage): + embed_excerpts(store, shas, MeteredEmbeddingClient(ScriptedEmbedder(report_usage=False), prices=PRICES), dimensions=4) + + +def test_a_dimension_change_is_refused(store: MemoryStore) -> None: + shas = _excerpts(store, 1) + with pytest.raises(EmbeddingError, match="dimension"): + embed_excerpts(store, shas, MeteredEmbeddingClient(ScriptedEmbedder(dims=3), prices=PRICES), dimensions=4) + + +def test_candidate_vector_is_cached_by_its_excerpt_sha(store: MemoryStore) -> None: + inner = ScriptedEmbedder() + cache = EmbeddingCache(store) + first = vector_for("def g():\n pass\n", cache, MeteredEmbeddingClient(inner, prices=PRICES)) + second = vector_for("def g(): \n pass", cache, MeteredEmbeddingClient(inner, prices=PRICES)) + assert first == second and len(inner.calls) == 1 and cache.names() == {excerpt_sha("def g():\n pass\n")} + assert vector_for("other", cache, None) is None + + +def test_vectors_round_trip_and_cosine() -> None: + vector = [0.25, -1.5, 3.0] + assert decode_vector(encode_vector(vector)) == vector + assert cosine([1, 0], [1, 0]) == 1.0 and cosine([1, 0], [0, 1]) == 0.0 + assert model_slug("openai/text-embedding-3-small") == "openai__text-embedding-3-small" + + +def test_usage_is_kept_from_the_provider_payload() -> None: + vectors, usage = parse_embedding_payload({"data": [{"embedding": [1, 2]}], "usage": {"prompt_tokens": 7, "total_tokens": 7}}) + assert vectors == [[1.0, 2.0]] and usage["prompt_tokens"] == 7 + assert parse_embedding_payload({"data": [{"embedding": [1]}]})[1] == {} + + +def test_the_embedding_model_manifest_prices_input_only() -> None: + model = load_manifests([Path("plane/models/openrouter-embedding.yaml")]).models["openrouter-embedding"] + assert model.provider == "openrouter" and model.model == "openai/text-embedding-3-small" + assert model.parameters["input_per_m"] == 0.02 and model.parameters["output_per_m"] == 0 + assert model.metadata.annotations["openultrasast.io/egress-hosts"] == "openrouter.ai" diff --git a/tests/test_plane_increment.py b/tests/test_plane_increment.py index a3d0c11..1393f6e 100644 --- a/tests/test_plane_increment.py +++ b/tests/test_plane_increment.py @@ -39,7 +39,7 @@ def test_every_manifest_under_plane_loads() -> None: paths = sorted(PLANE.rglob("*.yaml")) assert len(paths) >= 36, paths manifests = load_manifests(paths) - assert set(manifests.models) == {"deepseek-flash"} + assert set(manifests.models) == {"deepseek-flash", "openrouter-embedding"} # the second: embeddings only (decision engine) assert set(TEMPLATES) <= set(manifests.tasks) assert len(manifests.workspaces) == 30 and "validation-46" in manifests.runs full = {t.name for t in manifests.runs["validation-46"].tasks} diff --git a/tests/test_plane_manifests.py b/tests/test_plane_manifests.py index c1e72ce..7225047 100644 --- a/tests/test_plane_manifests.py +++ b/tests/test_plane_manifests.py @@ -370,7 +370,7 @@ def test_deepseek_without_extension_annotation_is_rejected() -> None: def test_unknown_provider_is_rejected_even_with_annotation() -> None: document = doc(MODEL_YAML) document["spec"]["provider"] = "openai" - assert rejected(document, 'spec.provider "openai" is not one of anthropic, deepseek, google') + assert rejected(document, 'spec.provider "openai" is not one of anthropic, deepseek, google, openrouter') def test_model_requires_model_name() -> None: