diff --git a/CHANGELOG.md b/CHANGELOG.md index d6a3e556..4b4626a7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -24,6 +24,23 @@ AgentExperience.NET is a **preview**. It is not production ready, and public API - **The Python and Go ports still pass contract 1.5**, not 1.6. Their suites accept a 1.6 lock and skip the four vector families by name until stories 23.2 and 23.3 implement them. +### Python semantic retrieval (story 23.2) + +- **The Python port passes contract 1.6.** `use_agent_experience(..., embedder=...)` takes your embedding callable (an + object with `model_id`, `dimension` and a sync or async `embed(text)`; no model is a dependency). With one, injection + runs the text and the vector channel together and merges them as contract 1.6 says, then ranks with the unchanged + ranking, so a reworded query finds a lesson it shares no words with. Without one, or when the vector channel fails, + it answers from text alone and reports why: `NotConfigured`, `ProviderUnavailable`, `ModelMismatch`, + `DimensionMismatch` or `VectorSearchFailed`, through `on_vector_fallback` and `InjectionResult.vector_fallback`. The + vector channel never fails the call and is bounded by `vector_timeout` (0.5 seconds by default); the text channel's + failure still fails it. `min_vector_relevance` drops weak vector candidates, never text ones. +- **Stores.** `InMemoryExperienceStore` and `PostgresExperienceStore` (PostgreSQL with pgvector, running the contract's + two new SQL files as written) both serve the channel through `write_embedding` and `search_vectors`. After a record is + stored, its summary's embedding is written in the background, outside the record's own commit; a failure is reported + to `on_capture_failure` as the new stage `Index` and never fails capture. +- Without an embedder nothing changes. The Python test suite now needs a pgvector image (`pgvector/pgvector`) for its + PostgreSQL leg. The Go port still passes contract 1.5. + ## 0.1.0-preview.14 This is the first preview that publishes the Python and Go ports. The same tag now releases the six NuGet packages, the diff --git a/README.md b/README.md index 036172e5..13f2cdb4 100644 --- a/README.md +++ b/README.md @@ -244,7 +244,7 @@ lesson was injected) is published as [a language-neutral contract](contract/READ [what 1.5 adds](contract/README-1.5.md), [what 1.6 adds](contract/README-1.6.md): retrieval by meaning, with embeddings): a JSON Schema, the SQL to run, and fixtures generated from and checked against this implementation on every build. A pure-Python core that passes -contract 1.5, with a one-call Python Agent Framework integration, is in [`python/`](python/README.md). A lesson a .NET +contract 1.6 (including retrieval by meaning, with an embedder you register), with a one-call Python Agent Framework integration, is in [`python/`](python/README.md). A lesson a .NET agent learned is injected into a Python agent, and the other way round. A Go core and PostgreSQL store that pass contract 1.5, with a one-call Go Agent Framework integration, are in [`go/`](go/README.md). A lesson a .NET agent learned is injected into a Go agent, and the other way round. Both are published with each release, at the same version as the NuGet packages (a Python alpha, diff --git a/python/README.md b/python/README.md index 06aa7ada..9d4a9646 100644 --- a/python/README.md +++ b/python/README.md @@ -3,8 +3,8 @@ A pure-Python core of AgentExperience that implements the [language-neutral contract](https://github.com/fabbrik/AgentExperience.NET/blob/main/contract/README.md), so a lesson learned by a Python agent is stored in the format the .NET library reads, and the other way round. -**Status: alpha, contract 1.5, with a Microsoft Agent Framework integration (the `maf` extra).** It passes every -fixture of contract 1.5. It is published to PyPI as `agenticexperience` (the import name stays `agentexperience`), with every +**Status: alpha, contract 1.6, with a Microsoft Agent Framework integration (the `maf` extra).** It passes every +fixture of contract 1.6. It is published to PyPI as `agenticexperience` (the import name stays `agentexperience`), with every release of this repository and at the same version as the .NET packages: .NET `0.1.0-preview.N` is `agenticexperience` `0.1.0aN` (a PEP 440 alpha; while every release is an alpha, pip and uv install the newest one). Like the .NET packages, it claims no production readiness, and its API may change between alphas. @@ -64,6 +64,9 @@ prints each run and the block. Its output is checked in as `additional_properties["AgentExperience.HistoricalReference"] = True` (`message_role="system"` sends it as a system message). An agent that already holds an `ExperienceContextProvider` gets no second one; `inject=False` captures without injecting. +- **Semantic retrieval** (contract 1.6), with an embedder you register (below). Injection then runs the text and the + vector channel together, merges them and ranks as before; each stored record's summary is embedded in the + background. - **Injection tracking** (contract 1.4). Once a block is handed over, a store with `record_injection` (the PostgreSQL and in-memory stores) is told which records it delivered, so the .NET lessons view shows when each lesson was last injected, by a Python agent too. It runs in a background task on the event loop, bounded by @@ -96,6 +99,40 @@ prints each run and the block. Its output is checked in as (`agentexperience.classifiers`) have the semantics of `ExperienceToolResultClassifiers`: `exit_code(*tools)` and `http_status(*tools)` read a JSON object's top-level fields, else the last match in the text. +### Semantic retrieval: registering an embedder + +Without an embedder retrieval is text only, and a reworded query finds nothing it shares no words with. Register an +embedding callable (any model: none is a dependency) and retrieval becomes hybrid: + +```python +class Embedder: + model_id = "my-embedding-model" # stored with every vector; vectors of another model are never compared + dimension = 1536 # the vector's length, known without calling the provider + + def embed(self, text: str) -> list[float]: # or `async def`; exactly `dimension` finite numbers + return my_provider.embed(text) + +agent = use_agent_experience(agent, store=store, identity=identity, verify=verify, embedder=Embedder()) +``` + +- **In-memory and PostgreSQL stores.** `InMemoryExperienceStore` compares 32-bit floats by cosine distance. + `PostgresExperienceStore` needs PostgreSQL with pgvector, the .NET vectors package's migrations (`0004`, `0020`) and + the grants of `sql/application-role-privileges-1.6.sql` after the 1.0 and 1.4 files (or the .NET privileges call). +- **Fail-open.** The vector channel never fails or delays past `vector_timeout` (seconds, default 0.5, covering the + embedding and the search) the text answer. When it does not contribute the injection answers from text and reports + the reason through `on_vector_fallback(VectorFallback)` and `InjectionResult.vector_fallback`: `NotConfigured` (no + embedder, or a store without `write_embedding` and `search_vectors`), `ProviderUnavailable` (the embedder raised, ran + out of time, or returned a blank model ID or an empty, wrong-length or non-finite vector; no statement runs), + `ModelMismatch` or `DimensionMismatch` (every stored vector is of another model or width) and `VectorSearchFailed`. + A failure of the text search still fails the injection, as without an embedder. `min_vector_relevance` (0 to 1) drops + vector candidates below it; text candidates are never dropped. +- **Writing embeddings.** After a record is stored, a background task embeds its summary (task ID, task summary and + lesson) and writes it, only while the record is still at that revision. It never runs inside the record's own commit + and never fails capture: a failure is reported to `on_capture_failure` with the stage `Index`. + `await settings.wait_for_index_writes()` (also on the provider) awaits the ones still running. Records stored before + an embedder was registered have no vector until one is written; the .NET library's re-index does it for a shared + database. + ### In production The identity must come from your own authentication, never from model output: pass a function of an @@ -128,7 +165,7 @@ which shares lessons with .NET agents on the same database. `sql/application-role-privileges.sql`, then `sql/application-role-privileges-1.4.sql` (without the second, every injection is still made and its record is reported `NotRecorded`). -It does nothing outside contract 1.x: no grants, crypto-shredding, provenance signing or vectors, and no lifecycle +It does nothing outside contract 1.x: no grants, crypto-shredding or provenance signing, and no lifecycle after finalization (reuse evidence, reinforcement, contest, revocation). It writes no logs at all, so no record content is ever logged. @@ -258,7 +295,7 @@ uv run mypy --strict src tests uv run pytest ``` -`uv run pytest` runs the unit tests and the conformance runner (`tests/test_contract.py`), which loads every fixture +`uv run pytest` (its PostgreSQL leg uses a `pgvector/pgvector` image) runs the unit tests and the conformance runner (`tests/test_contract.py`), which loads every fixture family from [`../contract/fixtures`](https://github.com/fabbrik/AgentExperience.NET/tree/main/contract/fixtures) and compares this package's output exactly: blocks and task text byte for byte, JSON by value. It validates every record against the contract's JSON Schema, checks every contract file against `contract/lock.json`, and fails with a clear message if the lock names a contract version this diff --git a/python/src/agentexperience/__init__.py b/python/src/agentexperience/__init__.py index ea785d85..0f67e55c 100644 --- a/python/src/agentexperience/__init__.py +++ b/python/src/agentexperience/__init__.py @@ -1,5 +1,5 @@ """ -AgentExperience for Python: a pure-Python core that implements the AgentExperience.NET language-neutral contract 1.5. +AgentExperience for Python: a pure-Python core that implements the AgentExperience.NET language-neutral contract 1.6. Pre-release. It finalizes a captured run into a stored Experience Record, renders the injected Historical Reference block, derives the task text, classifies errors, and (with the ``postgres`` extra) writes records to and searches them @@ -13,6 +13,19 @@ from .block import Block, BlockOptions, RankedRecord, RankingComponent, render_block from .capture import CapturedRun, RawAttempt, RawToolCall, RunCapture +from .embedding import ( + Embedder, + EmbeddingWrite, + EmbeddingWriteResult, + FallbackReason, + VectorFallback, + VectorHit, + VectorQuery, + VectorSearchResult, + content_hash, + merge_channels, + retrieval_summary, +) from .encoding import Duration, Instant from .error_class import error_class from .finalize import FinalizedExperience, InitialLifecycleEvent, finalize @@ -49,11 +62,11 @@ except _metadata.PackageNotFoundError: # running from a source tree that was never installed __version__ = "0.0.0+unknown" -CONTRACT_VERSION = "1.5" -"""The contract version this package passes: every fixture of contract 1.5, and so every 1.0, 1.1, 1.2, 1.3 and 1.4 -fixture.""" +CONTRACT_VERSION = "1.6" +"""The contract version this package passes: every fixture of contract 1.6, and so every 1.0, 1.1, 1.2, 1.3, 1.4 and +1.5 fixture.""" -SUPPORTED_CONTRACT_VERSIONS: tuple[str, ...] = ("1.0", "1.1", "1.2", "1.3", "1.4", "1.5") +SUPPORTED_CONTRACT_VERSIONS: tuple[str, ...] = ("1.0", "1.1", "1.2", "1.3", "1.4", "1.5", "1.6") """Every contract version whose fixtures this package passes.""" __all__ = [ @@ -67,9 +80,13 @@ "CapturedRun", "ClosedVerificationRound", "Duration", + "Embedder", + "EmbeddingWrite", + "EmbeddingWriteResult", "EnvironmentFingerprint", "Evidence", "ExperienceRecord", + "FallbackReason", "FinalizedExperience", "InitialLifecycleEvent", "Instant", @@ -95,14 +112,21 @@ "SearchRow", "SearchSafetyError", "ToolCall", + "VectorFallback", + "VectorHit", + "VectorQuery", + "VectorSearchResult", + "content_hash", "derive_task_text", "error_class", "experience_id_for", "finalize", "from_stored", "initial_event_id_for", + "merge_channels", "rank", "reflection_id_for", "render_block", + "retrieval_summary", "to_stored", ] diff --git a/python/src/agentexperience/embedding.py b/python/src/agentexperience/embedding.py new file mode 100644 index 00000000..1702bf48 --- /dev/null +++ b/python/src/agentexperience/embedding.py @@ -0,0 +1,447 @@ +""" +Semantic retrieval (contract 1.6): the embedder the caller brings, what is embedded and how it is hashed, the +validation that runs before any statement, the write and search types a store serves, and the pure functions that turn +two channels' candidates into one list. + +No embedding model is a dependency of this package: an :class:`Embedder` is the caller's own object. +""" + +from __future__ import annotations + +import asyncio +import hashlib +import inspect +import math +import struct +import threading +import unicodedata +import uuid +from collections.abc import Awaitable, Callable, Mapping, Sequence +from dataclasses import dataclass +from typing import Generic, Literal, Protocol, TypeVar, runtime_checkable + +from ._text import format_round_trip, is_null_or_whitespace, utf16_len +from .block import RankedRecord +from .encoding import Instant +from .records import ExperienceRecord, ExperienceStatus, Scope +from .search import ( + CANDIDATE_LIMIT, + DEFAULT_MIN_CONFIDENCE, + ELIGIBLE_STATUSES, + STATEMENT_LIMIT, + RankedResult, + SearchRow, + SearchSafetyError, + rank, +) + +MAX_SUMMARY_LENGTH = 8192 +"""The longest summary that is embedded, in UTF-16 code units.""" + +MAX_MODEL_ID_LENGTH = 256 +MAX_DIMENSION = 16_000 +"""The most components a vector has (pgvector's ``vector`` type accepts up to 16,000).""" + +_SEPARATORS = frozenset({"Cc", "Zs", "Zl", "Zp"}) + + +@runtime_checkable +class Embedder(Protocol): + """ + Turns text into a vector: the caller's callable, registered with ``use_agent_experience(embedder=...)``. + + ``model_id`` and ``dimension`` are known without calling the provider; the vector ``embed`` returns has exactly + ``dimension`` finite components. ``embed`` may be a plain function (it then runs on a worker thread) or a coroutine + function. Embeddings of different ``model_id`` or ``dimension`` are never compared. + """ + + @property + def model_id(self) -> str: ... + + @property + def dimension(self) -> int: ... + + def embed(self, text: str) -> Sequence[float] | Awaitable[Sequence[float]]: ... + + +FallbackReason = Literal[ + "NotConfigured", "ProviderUnavailable", "ModelMismatch", "DimensionMismatch", "VectorSearchFailed" +] +"""Why an injection answered from text alone (contract 1.6).""" + + +@dataclass(frozen=True) +class VectorFallback: + """The vector channel did not contribute: the contract's reason, a sentence that never quotes record or vector + content, and the exception, when there was one.""" + + reason: FallbackReason + detail: str + exception: BaseException | None = None + + +# ---- What is embedded --------------------------------------------------------------------------------------------- + + +def _append(parts: list[str], value: str | None) -> None: + if value is None: + return + words: list[str] = [] + word: list[str] = [] + for ch in value: + if unicodedata.category(ch) in _SEPARATORS: + if word: + words.append("".join(word)) + word = [] + else: + word.append(ch) + if word: + words.append("".join(word)) + if words: + parts.append(" ".join(words)) + + +def retrieval_summary(task_id: str | None, task_summary: str | None, lesson: str | None) -> str: + """ + The only text of a record that is embedded: task ID, task summary and lesson, each with every run of separators + (Unicode categories Cc, Zs, Zl, Zp) as one space, joined by one space, at most 8,192 UTF-16 code units. + """ + parts: list[str] = [] + for value in (task_id, task_summary, lesson): + _append(parts, value) + text = " ".join(parts) + if utf16_len(text) <= MAX_SUMMARY_LENGTH: + return text + # Cut at 8,192 code units; a cut that would end on a high surrogate drops it too; then trailing spaces go. + kept: list[str] = [] + used = 0 + for ch in text: + size = 2 if ord(ch) > 0xFFFF else 1 + if used + size > MAX_SUMMARY_LENGTH: + break + kept.append(ch) + used += size + return "".join(kept).rstrip(" ") + + +def record_summary(record: ExperienceRecord) -> str: + """:func:`retrieval_summary` of a record.""" + return retrieval_summary( + record.task_id, record.task_summary, record.reflection.lesson if record.reflection else None + ) + + +def content_hash(model_id: str, summary: str) -> str: + """Lowercase hex SHA-256 of the UTF-8 bytes of the model ID, U+001F and the summary.""" + text = model_id + "\u001f" + summary + # A lone surrogate is U+FFFD in UTF-8, as .NET encodes it. + data = text.encode("utf-16-le", "surrogatepass").decode("utf-16-le", "replace").encode("utf-8") + return hashlib.sha256(data).hexdigest() + + +MAX_VECTOR_THREADS = 16 +"""How many embedder or vector-store calls may run at once, abandoned (timed-out) ones included. They run on daemon +threads of their own, never the event loop's default executor that the text search uses, so a hung provider can never +starve the text channel; past this a call fails fast.""" + +_vector_threads = threading.BoundedSemaphore(MAX_VECTOR_THREADS) + + +async def run_on_vector_thread(function: Callable[..., object], *args: object) -> object: + """Runs ``function`` on a daemon thread of its own and awaits it. Raises ``RuntimeError`` at once when + :data:`MAX_VECTOR_THREADS` calls are already running. A call the caller stops waiting for keeps its thread + (a thread cannot be stopped) until it ends.""" + if not _vector_threads.acquire(blocking=False): + raise RuntimeError("Too many earlier embedding or vector-search calls are still running.") + loop = asyncio.get_running_loop() + future: asyncio.Future[object] = loop.create_future() + + def settle(value: object, error: BaseException | None) -> None: + if not future.done(): + if error is None: + future.set_result(value) + else: + future.set_exception(error) + + def run() -> None: + try: + try: + value, error = function(*args), None + except BaseException as ex: + value, error = None, ex + try: + loop.call_soon_threadsafe(settle, value, error) + except RuntimeError: # the loop closed first + pass + finally: + _vector_threads.release() + + threading.Thread(target=run, name="agentexperience-vector", daemon=True).start() + return await future + + +async def embed_text(embedder: Embedder, text: str) -> Sequence[float]: + """The embedder's vector for ``text``: a coroutine function is awaited, a plain function runs on a worker thread + (so a blocking provider never stalls the event loop).""" + if inspect.iscoroutinefunction(embedder.embed): + result = await embedder.embed(text) + else: + result = await run_on_vector_thread(embedder.embed, text) + if inspect.isawaitable(result): + result = await result + return result # type: ignore[no-any-return] + + +# ---- Validation (before any statement) ----------------------------------------------------------------------------- + + +def validate_model_id(model_id: str) -> None: + if not isinstance(model_id, str) or is_null_or_whitespace(model_id): + raise ValueError("The model ID must not be blank.") + if utf16_len(model_id) > MAX_MODEL_ID_LENGTH: + raise ValueError(f"The model ID is longer than {MAX_MODEL_ID_LENGTH} UTF-16 code units.") + + +def float32(value: float) -> float: + """``value`` rounded to the nearest 32-bit float; ``ValueError`` when that is not finite.""" + try: + rounded = float(struct.unpack(" tuple[float, ...]: + """The vector as floats; ``ValueError`` unless it has 1 to 16,000 finite components (and ``dimension`` of them, + when given).""" + try: + components = tuple(float(component) for component in vector) + except (TypeError, ValueError, OverflowError) as bad: + raise ValueError("A vector must be a sequence of numbers.") from bad + if not 1 <= len(components) <= MAX_DIMENSION: + raise ValueError(f"A vector has 1 to {MAX_DIMENSION} components.") + if dimension is not None and len(components) != dimension: + raise ValueError("The vector's length is not its dimension.") + for component in components: + float32(component) + if not math.isfinite(component): + raise ValueError("A vector component must be finite.") + return components + + +def vector_literal(vector: Sequence[float]) -> str: + """A pgvector literal, ``[x,y,...]``, each component its shortest round-trip text.""" + return "[" + ",".join(format_round_trip(float(component)) for component in vector) + "]" + + +@dataclass(frozen=True) +class EmbeddingWrite: + """One record's embedding to store: the record (its ID, exact scope and the revision the summary was read at) and + what the embedding is.""" + + experience_id: uuid.UUID + scope: Scope + model_id: str + dimension: int + content_hash: str + source_revision: int + vector: tuple[float, ...] + now: Instant + + def __post_init__(self) -> None: + validate_model_id(self.model_id) + if ( + not isinstance(self.dimension, int) + or isinstance(self.dimension, bool) + or not 1 <= self.dimension <= MAX_DIMENSION + ): + raise ValueError(f"The dimension must be 1 to {MAX_DIMENSION}.") + validate_vector(self.vector, self.dimension) + if is_null_or_whitespace(self.content_hash): + raise ValueError("The content hash must not be blank.") + if not isinstance(self.source_revision, int) or self.source_revision < 0: + raise ValueError("The source revision must not be negative.") + + +@dataclass(frozen=True) +class VectorQuery: + """A nearest-neighbour search in exactly ``scope``, over embeddings of ``model_id`` and the vector's length.""" + + scope: Scope + model_id: str + vector: tuple[float, ...] + statuses: tuple[ExperienceStatus, ...] = ELIGIBLE_STATUSES + min_confidence: float = DEFAULT_MIN_CONFIDENCE + limit: int = STATEMENT_LIMIT + exclude_model_authored: bool = False + + def __post_init__(self) -> None: + validate_model_id(self.model_id) + validate_vector(self.vector) + if self.limit < 1: + raise ValueError("The limit must be strictly positive.") + + +EmbeddingWriteOutcome = Literal["Written", "Stale", "Missing"] + + +@dataclass(frozen=True) +class EmbeddingWriteResult: + """``Written``; ``Stale`` with the record's current revision (it moved on, nothing changed); or ``Missing`` (an + unknown, other-scope or erased record, nothing created).""" + + outcome: EmbeddingWriteOutcome + current_revision: int | None = None + + +VectorSearchOutcome = Literal["Found", "ModelMismatch", "DimensionMismatch"] + + +def relevance_of(distance: float) -> float: + """``clamp(1 - distance / 2, 0, 1)``, and 0 for NaN.""" + if math.isnan(distance): + return 0.0 + return min(max(1.0 - distance / 2.0, 0.0), 1.0) + + +@dataclass(frozen=True) +class VectorHit: + """One record the vector search returned, and its cosine distance (NaN for a zero-magnitude vector).""" + + record: ExperienceRecord + distance: float + + @property + def relevance(self) -> float: + return relevance_of(self.distance) + + +@dataclass(frozen=True) +class VectorSearchResult: + """The outcome, and the hits nearest first (ties by experience ID ordinal, NaN last); no hits unless ``Found``.""" + + outcome: VectorSearchOutcome + hits: tuple[VectorHit, ...] = () + + +def order_hits(hits: Sequence[VectorHit]) -> list[VectorHit]: + """Nearest first; equal distances by experience ID as text, ordinally; NaN after every number, equal to each + other.""" + return sorted( + hits, + key=lambda hit: ( + math.isnan(hit.distance), + 0.0 if math.isnan(hit.distance) else hit.distance, + str(hit.record.experience_id), + ), + ) + + +# ---- Merging the two channels -------------------------------------------------------------------------------------- + +T = TypeVar("T") + +Channel = Literal["text", "vector"] + + +@dataclass(frozen=True) +class Merged(Generic[T]): + """One entry of the merged list: the candidate kept (the first channel's), the relevance kept, and its channel.""" + + item: T + relevance: float + source: Channel + + +def _clamped(relevance: float) -> float: + return 0.0 if math.isnan(relevance) else min(max(relevance, 0.0), 1.0) + + +def merge_channels( + text: Sequence[T], + vector: Sequence[T], + key: Callable[[T], uuid.UUID], + relevance: Callable[[T], float], + minimum_vector_relevance: float | None = None, +) -> list[Merged[T]]: + """ + Contract 1.6 steps 1 and 3. The floor drops vector candidates below ``minimum_vector_relevance`` (one exactly at it + stays; a NaN relevance is below any floor; text is never floored). The merge lists the text candidates in order, + then the vector candidates in theirs; a record already listed keeps its position and the text channel's candidate, + and takes the vector relevance, as given, only when its clamped value (NaN as 0) is higher. + """ + merged: list[Merged[T]] = [] + positions: dict[uuid.UUID, int] = {} + for item in text: + positions.setdefault(key(item), len(merged)) + merged.append(Merged(item, relevance(item), "text")) + for item in vector: + value = relevance(item) + if minimum_vector_relevance is not None and not value >= minimum_vector_relevance: + continue + position = positions.get(key(item)) + if position is None: + positions[key(item)] = len(merged) + merged.append(Merged(item, value, "vector")) + elif _clamped(value) > _clamped(merged[position].relevance): + merged[position] = Merged(merged[position].item, value, merged[position].source) + return merged + + +def check_channel(rows: Sequence[SearchRow], scope: Scope) -> list[SearchRow]: + """Contract 1.6 step 2 for one channel: over its first 50 candidates, one outside ``scope`` or a repeated + ``experience_id`` raises :class:`SearchSafetyError`; returns the first 50.""" + considered = list(rows[:CANDIDATE_LIMIT]) + seen: set[uuid.UUID] = set() + for row in considered: + if row.shared_by_grant or row.record.scope != scope: + raise SearchSafetyError("A row lies outside the requested scope; nothing is returned.") + if row.record.experience_id in seen: + raise SearchSafetyError("The same record came back more than once; nothing is returned.") + seen.add(row.record.experience_id) + return considered + + +def rank_hybrid( + text_rows: Sequence[SearchRow], + vector_rows: Sequence[SearchRow], + scope: Scope, + now: Instant, + preferred_environment: Mapping[str, str] | None = None, + *, + minimum_vector_relevance: float | None = None, + exclude_model_authored: bool = False, +) -> RankedResult: + """ + The whole of steps 1 to 4: floor, per-channel checks and truncation, merge, then the existing + :func:`~agentexperience.search.rank` over the merged candidates (each with the relevance kept). ``rank`` looks at + 50 rows, so the merged list (up to 100) is ranked in slices of 50: a row's score does not depend on the others, and + the slices are put back in the order ``rank`` itself uses (score descending, then experience ID). + """ + floored = [ + row + for row in vector_rows[: CANDIDATE_LIMIT + 1] + if minimum_vector_relevance is None or row.relevance >= minimum_vector_relevance + ] + # Checks are over each channel's first 50 candidates as returned (the floor only drops vector candidates). + text = check_channel(text_rows, scope) + vector = check_channel([row for row in floored], scope) + merged = merge_channels(text, vector, lambda row: row.record.experience_id, lambda row: row.relevance) + rows = [SearchRow(entry.item.record, entry.relevance) for entry in merged] + ranked: list[RankedRecord] = [] + for start in range(0, len(rows), CANDIDATE_LIMIT): + ranked.extend( + rank( + rows[start : start + CANDIDATE_LIMIT], + scope, + now, + preferred_environment, + exclude_model_authored=exclude_model_authored, + ).ranked + ) + ranked.sort(key=lambda entry: str(entry.record.experience_id)) + ranked.sort(key=lambda entry: entry.score, reverse=True) + truncated = len(text_rows) > CANDIDATE_LIMIT or len(floored) > CANDIDATE_LIMIT + return RankedResult(tuple(ranked), truncated) diff --git a/python/src/agentexperience/inmemory.py b/python/src/agentexperience/inmemory.py index 95188d55..5629656b 100644 --- a/python/src/agentexperience/inmemory.py +++ b/python/src/agentexperience/inmemory.py @@ -8,6 +8,7 @@ from __future__ import annotations +import math import threading import unicodedata import uuid @@ -15,6 +16,15 @@ from dataclasses import dataclass from typing import Protocol, runtime_checkable +from .embedding import ( + EmbeddingWrite, + EmbeddingWriteResult, + VectorHit, + VectorQuery, + VectorSearchResult, + float32, + order_hits, +) from .encoding import Instant from .finalize import FinalizedExperience, InitialLifecycleEvent from .records import ExperienceRecord, Scope @@ -65,6 +75,19 @@ def record_injection( ) -> None: ... +@runtime_checkable +class EmbeddingIndex(Protocol): + """ + A store that keeps embeddings and searches them by cosine distance (contract 1.6). Optional: the agent integration + runs the vector channel only through a store that has both methods, and reports ``NotConfigured`` for one that + does not. + """ + + def write_embedding(self, bounds: AuthorizationBounds, write: EmbeddingWrite) -> EmbeddingWriteResult: ... + + def search_vectors(self, bounds: AuthorizationBounds, query: VectorQuery) -> VectorSearchResult: ... + + MAX_RECORD_INJECTION_IDS = 200 """The most records one :meth:`InjectionRecorder.record_injection` call names: the .NET store's limit (``ExperienceRecordGetManyResult.MaxCount``). The agent integration splits a longer list into chunks.""" @@ -133,6 +156,24 @@ def query_terms(text: str) -> list[str]: return list(dict.fromkeys(word for word in _words(text) if len(word) >= 2 and word not in _STOPWORDS)) +@dataclass(frozen=True) +class _Embedding: + model_id: str + dimension: int + content_hash: str + source_revision: int + vector: tuple[float, ...] # 32-bit float values + + +def cosine_distance(left: Sequence[float], right: Sequence[float]) -> float: + """Cosine distance (``1 - cos``, in [0, 2]) of two vectors of equal length; NaN when either has no magnitude.""" + dot = sum(a * b for a, b in zip(left, right, strict=True)) + norm = math.sqrt(sum(a * a for a in left)) * math.sqrt(sum(b * b for b in right)) + if norm == 0.0 or math.isnan(norm): + return math.nan + return min(max(1.0 - dot / norm, 0.0), 2.0) + + @dataclass(frozen=True) class _Indexed: record: ExperienceRecord @@ -177,6 +218,7 @@ def __init__(self, records: Iterable[ExperienceRecord] = (), *, clock: Callable[ self._records: dict[uuid.UUID, _Indexed] = {} self._events: dict[uuid.UUID, tuple[InitialLifecycleEvent, str]] = {} self._injections: dict[uuid.UUID, InjectionSummary] = {} + self._embeddings: dict[uuid.UUID, _Embedding] = {} for record in records: if record.experience_id in self._records: raise ValueError("Two records share one experience ID.") @@ -273,3 +315,78 @@ def search(self, bounds: AuthorizationBounds, parameters: SearchParameters) -> l matches.sort(key=lambda row: str(row.record.experience_id)) matches.sort(key=lambda row: row.relevance, reverse=True) return matches[: parameters.limit] + + def write_embedding(self, bounds: AuthorizationBounds, write: EmbeddingWrite) -> EmbeddingWriteResult: + """ + Stores a record's embedding as ``sql/write-experience-embedding.sql`` does: only for a record held in exactly + the write's scope and still at exactly ``source_revision`` (``Written``); otherwise ``Stale`` with the record's + revision, or ``Missing`` for an unknown record or one in another scope or outside the bounds. A row stored + from a newer revision is never replaced by an older one. Components are kept as 32-bit floats. + """ + vector = tuple(float32(component) for component in write.vector) + if not permits(bounds, write.scope): + return EmbeddingWriteResult("Missing") + with self._lock: + indexed = self._records.get(write.experience_id) + if indexed is None or indexed.record.scope != write.scope: + return EmbeddingWriteResult("Missing") + revision = indexed.record.revision + existing = self._embeddings.get(write.experience_id) + if revision != write.source_revision or ( + existing is not None and existing.source_revision > write.source_revision + ): + return EmbeddingWriteResult("Stale", revision) + self._embeddings[write.experience_id] = _Embedding( + write.model_id, write.dimension, write.content_hash, write.source_revision, vector + ) + return EmbeddingWriteResult("Written") + + def embedding(self, experience_id: uuid.UUID) -> tuple[str, int, str, int, tuple[float, ...]] | None: + """A record's stored embedding (model ID, dimension, content hash, source revision, vector), for tests.""" + with self._lock: + stored = self._embeddings.get(experience_id) + return ( + None + if stored is None + else (stored.model_id, stored.dimension, stored.content_hash, stored.source_revision, stored.vector) + ) + + def search_vectors(self, bounds: AuthorizationBounds, query: VectorQuery) -> VectorSearchResult: + """ + The nearest-neighbour search: embeddings of the query's model and dimension, of records in exactly the query's + scope with an eligible status and confidence (and not model-authored, when excluded), by cosine distance + (components compared as 32-bit floats), ties by experience ID, NaN last, at most ``limit``. With none, the + outcome says why: ``ModelMismatch`` or ``DimensionMismatch`` when eligible embeddings exist but none is + comparable, else ``Found`` and empty. Nothing outside the bounds. + """ + scope = query.scope + if not permits(bounds, scope): + return VectorSearchResult("Found") + vector = tuple(float32(component) for component in query.vector) + statuses = set(query.statuses) + with self._lock: + held = [ + (self._records[experience_id].record, embedding) + for experience_id, embedding in self._embeddings.items() + if experience_id in self._records + ] + eligible = [ + (record, embedding) + for record, embedding in held + if record.scope == scope + and record.status in statuses + and record.reuse_confidence >= query.min_confidence + and not (query.exclude_model_authored and is_model_authored(record.reflection)) + ] + hits = [ + VectorHit(record, cosine_distance(embedding.vector, vector)) + for record, embedding in eligible + if embedding.model_id == query.model_id and embedding.dimension == len(vector) + ] + if hits: + return VectorSearchResult("Found", tuple(order_hits(hits)[: query.limit])) + if not eligible: + return VectorSearchResult("Found") + if not any(embedding.model_id == query.model_id for _, embedding in eligible): + return VectorSearchResult("ModelMismatch") + return VectorSearchResult("DimensionMismatch") diff --git a/python/src/agentexperience/maf.py b/python/src/agentexperience/maf.py index f24951fc..4fc33ffb 100644 --- a/python/src/agentexperience/maf.py +++ b/python/src/agentexperience/maf.py @@ -24,6 +24,9 @@ searches the store in the identity's scope, ranks with the contract's weights, renders the Historical Reference block with the block options, and adds it as one user-role message (``message_role="system"`` sends it as a system message), stamped ``additional_properties["AgentExperience.HistoricalReference"] = True``. +- **Semantic retrieval** (contract 1.6), with ``embedder=``. Injection then runs the text and the vector channel + together, merges them and ranks as before; every stored record's summary is embedded in the background. Without an + embedder, or when the vector channel fails, it answers from text alone and reports why (``on_vector_fallback``). - **Injection tracking** (contract 1.4). Once the block is handed over, a store that records injections (the PostgreSQL store, the in-memory store) is told which records it delivered, in a background task bounded by ``injection_record_timeout``, so the .NET lessons view shows when each lesson was last injected. It never delays or @@ -69,6 +72,20 @@ exit_code, http_status, ) +from .embedding import ( + Embedder, + EmbeddingWrite, + VectorFallback, + VectorQuery, + VectorSearchResult, + content_hash, + embed_text, + rank_hybrid, + record_summary, + run_on_vector_thread, + validate_model_id, + validate_vector, +) from .encoding import Duration, Instant from .error_class import is_single_token from .finalize import FinalizedExperience, finalize @@ -80,10 +97,16 @@ is_development_environment, reduce_name, ) -from .inmemory import MAX_RECORD_INJECTION_IDS, ExperienceStore, InjectionRecorder, InMemoryExperienceStore +from .inmemory import ( + MAX_RECORD_INJECTION_IDS, + EmbeddingIndex, + ExperienceStore, + InjectionRecorder, + InMemoryExperienceStore, +) from .records import EnvironmentFingerprint, Provenance, RunExposure from .sanitization import Sanitization -from .search import SearchParameters, rank +from .search import SearchParameters, SearchRow, rank from .task_text import CHAT_HISTORY, EXTERNAL, derive_task_text from .task_text import Message as TaskMessage from .task_text import MessagePart as TaskMessagePart @@ -115,6 +138,10 @@ DEFAULT_FINALIZATION_TIMEOUT = 5.0 """How long, in seconds, an invocation's end waits for finalization by default: .NET's ``FinalizationTimeout``.""" +DEFAULT_VECTOR_TIMEOUT = 0.5 +"""How long, in seconds, the vector channel (embedding the task text, then searching the stored vectors) may take by +default; past it the injection answers from text alone.""" + DEFAULT_INJECTION_RECORD_TIMEOUT = 5.0 """How long, in seconds, recording one injection may take by default, all chunks together: .NET's.""" @@ -150,6 +177,7 @@ class _Write: "Verify", "Finalize", "Store", + "Index", ] AgentT = TypeVar("AgentT") @@ -181,6 +209,8 @@ class InjectionResult: experience_ids: tuple[uuid.UUID, ...] byte_count: int truncated: bool + vector_fallback: VectorFallback | None = None + """Why the injection answered from text alone (contract 1.6), or ``None`` when the vector channel contributed.""" InjectionRecordOutcome = Literal["Recorded", "NotTracked", "NotRecorded", "OutcomeUnknown"] @@ -246,6 +276,11 @@ class ExperienceSettings: on_injected: Callable[[InjectionResult], None] | None = None on_finalized: Callable[[FinalizedExperience], None] | None = None on_injection_recorded: Callable[[InjectionRecordResult], None] | None = None + embedder: Embedder | None = None + """Turns text into a vector, for hybrid retrieval and for embedding each stored record (contract 1.6).""" + min_vector_relevance: float | None = None + """When set (0 to 1), vector candidates whose relevance is below it are dropped; text candidates never are.""" + on_vector_fallback: Callable[[VectorFallback], None] | None = None clock: Callable[[], Instant] = Instant.now new_id: Callable[[], uuid.UUID] = uuid.uuid4 finalization_timeout: float = DEFAULT_FINALIZATION_TIMEOUT @@ -254,6 +289,10 @@ class ExperienceSettings: injection_record_timeout: float = DEFAULT_INJECTION_RECORD_TIMEOUT """Seconds recording one injected block may take in the background, all chunks together; past it the write is abandoned and reported ``OutcomeUnknown``. The invocation never waits for it.""" + vector_timeout: float = DEFAULT_VECTOR_TIMEOUT + """Seconds the vector channel may take in all (embedding the task text and searching); past it the injection + answers from text alone and reports ``ProviderUnavailable`` or ``VectorSearchFailed``.""" + _index_tasks: set[asyncio.Task[None]] = field(default_factory=set, init=False, repr=False, compare=False) def __post_init__(self) -> None: if self.store is None: @@ -276,6 +315,7 @@ def __post_init__(self) -> None: "on_injected", "on_finalized", "on_injection_recorded", + "on_vector_fallback", "clock", "new_id", ): @@ -286,6 +326,16 @@ def __post_init__(self) -> None: raise ValueError("finalization_timeout must be a positive number of seconds.") if not self.injection_record_timeout > 0: raise ValueError("injection_record_timeout must be a positive number of seconds.") + if not self.vector_timeout > 0: + raise ValueError("vector_timeout must be a positive number of seconds.") + if self.min_vector_relevance is not None and not 0.0 <= self.min_vector_relevance <= 1.0: + raise ValueError("min_vector_relevance must be between 0 and 1.") + if self.embedder is not None and not ( + isinstance(getattr(self.embedder, "model_id", None), str) + and hasattr(self.embedder, "dimension") + and callable(getattr(self.embedder, "embed", None)) + ): + raise TypeError("embedder must have model_id, dimension and an embed(text) method.") if self.message_role not in ("user", "system"): raise ValueError("message_role must be 'user' or 'system'.") if self.task_id is not None and not self.task_id.strip(): @@ -314,6 +364,21 @@ def report_injection(self, failure: InjectionFailure) -> None: except Exception: pass + def report_vector_fallback(self, fallback: VectorFallback) -> None: + if self.on_vector_fallback is None: + return + try: + self.on_vector_fallback(fallback) + except Exception: + pass + + async def wait_for_index_writes(self) -> None: + """Completes once every background embedding write this configuration started on the running event loop has + ended. Each is bounded by ``finalization_timeout``; never raises.""" + loop = asyncio.get_running_loop() + while pending := [task for task in self._index_tasks if task.get_loop() is loop and not task.done()]: + await asyncio.gather(*pending, return_exceptions=True) + def report_injection_recorded(self, result: InjectionRecordResult) -> None: if self.on_injection_recorded is None: return @@ -509,12 +574,63 @@ async def _finish(self, status: RunExecutionStatus, result: str | None, error: s self.report("Store", f"Storing the record raised {_type_name(ex)}; the run is not stored.", ex) return + # The embedding is derived data: written after the record's own commit, in the background, never failing it. + self._start_index(finalized, identity) + if settings.on_finalized is not None: try: settings.on_finalized(finalized) except Exception: pass + def _start_index(self, finalized: FinalizedExperience, identity: ExperienceIdentity) -> None: + """Starts embedding the stored record's summary in the background; never raises, never waits.""" + settings = self.settings + if settings.embedder is None or not isinstance(settings.store, EmbeddingIndex): + return + try: + task = asyncio.get_running_loop().create_task(self._index(finalized, identity)) + except Exception as ex: + self.report("Index", f"Starting the embedding write raised {_type_name(ex)}.", ex) + return + settings._index_tasks.add(task) + task.add_done_callback(settings._index_tasks.discard) + + async def _index(self, finalized: FinalizedExperience, identity: ExperienceIdentity) -> None: + settings = self.settings + try: + await asyncio.wait_for(self._write_embedding(finalized, identity), settings.finalization_timeout) + except (asyncio.TimeoutError, TimeoutError) as ex: + self.report( + "Index", f"Embedding the record did not finish within {settings.finalization_timeout} seconds.", ex + ) + except asyncio.CancelledError as ex: + self.report("Index", "Embedding the record was cancelled.", ex) + raise + except Exception as ex: + self.report("Index", f"Embedding the record raised {_type_name(ex)}; its text retrieval is unaffected.", ex) + + async def _write_embedding(self, finalized: FinalizedExperience, identity: ExperienceIdentity) -> None: + settings = self.settings + embedder, store = settings.embedder, settings.store + assert embedder is not None and isinstance(store, EmbeddingIndex) + record = finalized.record + summary = record_summary(record) + model_id, dimension = embedder.model_id, embedder.dimension + validate_model_id(model_id) + vector = validate_vector(await embed_text(embedder, summary), dimension) + write = EmbeddingWrite( + record.experience_id, + record.scope, + model_id, + dimension, + content_hash(model_id, summary), + record.revision, + vector, + settings.clock(), + ) + await run_on_vector_thread(store.write_embedding, identity.authorization, write) + _current_scope: ContextVar[_CaptureScope | None] = ContextVar("agentexperience_capture_scope", default=None) @@ -748,6 +864,64 @@ def _abandon(self, state: _Write) -> None: state.abandoned = True self._abandoned_records += 1 + async def wait_for_index_writes(self) -> None: + """Completes once the embedding writes started after records were stored have ended; see + :meth:`ExperienceSettings.wait_for_index_writes`.""" + await self.settings.wait_for_index_writes() + + async def _vector_channel( + self, task_text: str, identity: ExperienceIdentity + ) -> tuple[list[SearchRow], VectorFallback | None]: + """ + The vector channel (contract 1.6): embed the task text, search the stored vectors, both within + ``vector_timeout``. Never raises (apart from the caller's cancellation): whatever goes wrong is the fallback it + answers with, and its text never quotes record or vector content. + """ + settings = self.settings + embedder, store = settings.embedder, settings.store + if embedder is None: + return [], VectorFallback("NotConfigured", "No embedder is registered, so only text retrieval ran.") + if not isinstance(store, EmbeddingIndex): + return [], VectorFallback( + "NotConfigured", "The store does not serve vector search, so only text retrieval ran." + ) + loop = asyncio.get_running_loop() + deadline = loop.time() + settings.vector_timeout + try: + model_id, dimension = embedder.model_id, embedder.dimension + raw = await asyncio.wait_for(embed_text(embedder, task_text), max(deadline - loop.time(), 0.0)) + if not isinstance(model_id, str) or not model_id.strip(): + return [], VectorFallback("ProviderUnavailable", "The embedder reported no model ID.") + validate_model_id(model_id) + vector = validate_vector(raw, dimension) + query = VectorQuery( + identity.scope, model_id, vector, exclude_model_authored=settings.exclude_model_authored + ) + except (asyncio.TimeoutError, TimeoutError) as ex: + return [], VectorFallback("ProviderUnavailable", "The embedder did not answer in time.", ex) + except Exception as ex: + return [], VectorFallback( + "ProviderUnavailable", f"The embedder failed or returned an unusable vector ({_type_name(ex)}).", ex + ) + try: + result = await asyncio.wait_for( + run_on_vector_thread(store.search_vectors, identity.authorization, query), + max(deadline - loop.time(), 0.0), + ) + except (asyncio.TimeoutError, TimeoutError) as ex: + return [], VectorFallback("VectorSearchFailed", "The vector search did not answer in time.", ex) + except Exception as ex: + return [], VectorFallback("VectorSearchFailed", f"The vector search raised {_type_name(ex)}.", ex) + if not isinstance(result, VectorSearchResult): + return [], VectorFallback("VectorSearchFailed", "The store answered the vector search with something else.") + if result.outcome == "ModelMismatch": + return [], VectorFallback("ModelMismatch", "Every stored embedding in this scope is of another model.") + if result.outcome == "DimensionMismatch": + return [], VectorFallback("DimensionMismatch", "Every stored embedding in this scope is of another width.") + if result.outcome != "Found": + return [], VectorFallback("VectorSearchFailed", "The vector search answered with an unknown outcome.") + return [SearchRow(hit.record, hit.relevance) for hit in result.hits], None + async def before_run( self, *, @@ -787,16 +961,41 @@ async def before_run( parameters = SearchParameters( identity.scope, task_text, exclude_model_authored=settings.exclude_model_authored ) - rows = await asyncio.to_thread(settings.store.search, identity.authorization, parameters) + text_search = asyncio.ensure_future( + asyncio.to_thread(settings.store.search, identity.authorization, parameters) + ) + vector_search = asyncio.ensure_future(self._vector_channel(task_text, identity)) + # Both channels run together; the vector channel never raises (it answers with a fallback instead). + text_found, vector_found = await asyncio.gather(text_search, vector_search, return_exceptions=True) + if isinstance(vector_found, BaseException): + vector_found = ([], VectorFallback("VectorSearchFailed", "The vector channel failed unexpectedly.")) + vector_rows, fallback = vector_found + if fallback is not None: + settings.report_vector_fallback(fallback) + if isinstance(text_found, BaseException): + # The text channel's failure fails the call, whatever the vector channel found. + raise text_found + rows = text_found # The instant the block is built at: ranking's "now", and the injection's instant. now = settings.clock() - ranked = rank( - rows, - identity.scope, - now, - settings.preferred_environment, - exclude_model_authored=settings.exclude_model_authored, - ) + if vector_rows: + ranked = rank_hybrid( + rows, + vector_rows, + identity.scope, + now, + settings.preferred_environment, + minimum_vector_relevance=settings.min_vector_relevance, + exclude_model_authored=settings.exclude_model_authored, + ) + else: + ranked = rank( + rows, + identity.scope, + now, + settings.preferred_environment, + exclude_model_authored=settings.exclude_model_authored, + ) block = render_block(ranked.ranked, settings.block_options) if block.is_empty: return @@ -817,7 +1016,7 @@ async def before_run( delivered = tuple(record.experience_id for record in block.included) if settings.on_injected is not None: try: - settings.on_injected(InjectionResult(delivered, block.byte_count, ranked.truncated)) + settings.on_injected(InjectionResult(delivered, block.byte_count, ranked.truncated, fallback)) except Exception: pass @@ -1016,6 +1215,10 @@ def use_agent_experience( new_id: Callable[[], uuid.UUID] | None = None, finalization_timeout: float = DEFAULT_FINALIZATION_TIMEOUT, injection_record_timeout: float = DEFAULT_INJECTION_RECORD_TIMEOUT, + embedder: Embedder | None = None, + min_vector_relevance: float | None = None, + vector_timeout: float = DEFAULT_VECTOR_TIMEOUT, + on_vector_fallback: Callable[[VectorFallback], None] | None = None, ) -> AgentT: """ Wires capture and injection into a MAF agent (an ``Agent``, or any agent with ``context_providers`` and @@ -1035,6 +1238,12 @@ def use_agent_experience( ``on_injection_recorded`` reports how that ended. On PostgreSQL the application role needs the grants of the contract's ``sql/application-role-privileges-1.4.sql`` (or the .NET privileges call); without them every injection is still made and reported ``NotRecorded``. + + ``embedder`` (contract 1.6) turns text into a vector: with one, injection runs the text and the vector channel + together (in-memory and PostgreSQL+pgvector stores serve it) and merges them, and every stored record's summary is + embedded in the background. Without one, or when the vector channel fails, the injection answers from text and + reports why (``on_vector_fallback``, ``InjectionResult.vector_fallback``). ``vector_timeout`` bounds the embedding + and the search of one injection; ``min_vector_relevance`` drops weaker vector candidates. """ providers = getattr(agent, "context_providers", None) if not isinstance(providers, list): @@ -1066,6 +1275,10 @@ def use_agent_experience( new_id=new_id or uuid.uuid4, finalization_timeout=finalization_timeout, injection_record_timeout=injection_record_timeout, + embedder=embedder, + min_vector_relevance=min_vector_relevance, + vector_timeout=vector_timeout, + on_vector_fallback=on_vector_fallback, ) # Capture first, so it wraps everything else on the agent: injection runs inside it and reuses its identity. @@ -1078,6 +1291,7 @@ def use_agent_experience( __all__ = [ "DEFAULT_INJECTION_RECORD_TIMEOUT", + "DEFAULT_VECTOR_TIMEOUT", "HISTORICAL_REFERENCE_KEY", "MAX_ABANDONED_INJECTION_RECORDS", "PROVENANCE_SOURCE", @@ -1087,6 +1301,7 @@ def use_agent_experience( "Approval", "CaptureFailure", "DevelopmentOnlyError", + "Embedder", "ExperienceCaptureMiddleware", "ExperienceContextProvider", "ExperienceIdentity", @@ -1100,6 +1315,7 @@ def use_agent_experience( "InjectionRecordResult", "InjectionResult", "ToolResultContext", + "VectorFallback", "Verification", "VerificationContext", "all_of", diff --git a/python/src/agentexperience/postgres.py b/python/src/agentexperience/postgres.py index c495baab..2cf46a41 100644 --- a/python/src/agentexperience/postgres.py +++ b/python/src/agentexperience/postgres.py @@ -21,6 +21,15 @@ from typing import TYPE_CHECKING, Any from ._text import is_null_or_whitespace +from .embedding import ( + EmbeddingWrite, + EmbeddingWriteResult, + VectorHit, + VectorQuery, + VectorSearchResult, + order_hits, + vector_literal, +) from .encoding import Instant from .finalize import FinalizedExperience from .inmemory import AuthorizationDenied, injection_ids, permits @@ -118,9 +127,26 @@ def until(end: str, start: int) -> int: _DOLLAR_TAG = re.compile(r"\$(?:[A-Za-z_][A-Za-z0-9_]*)?\$") +def contract_statements(name: str) -> dict[str, str]: + """ + The statements of a contract SQL file that holds several, each starting at its ``-- statement: `` line and + running to the next one (or the end), by name; the header before the first is not a statement. + """ + statements: dict[str, list[str]] = {} + current: list[str] | None = None + for line in contract_sql(name).splitlines(): + if line.startswith("-- statement: "): + current = statements.setdefault(line.removeprefix("-- statement: ").strip(), []) + elif current is not None: + current.append(line) + return {key: "\n".join(lines).strip() for key, lines in statements.items()} + + DECLARE_BOUNDS_SQL = "declare-bounds.sql" SEARCH_SQL = "search-experience-text.sql" RECORD_INJECTION_SQL = "record-injection.sql" +WRITE_EMBEDDING_SQL = "write-experience-embedding.sql" +SEARCH_VECTOR_SQL = "search-experience-vector.sql" _INSERT_RECORD = ( f"INSERT INTO {TABLE} ({', '.join(COLUMN_NAMES)}) VALUES ({', '.join(f'%({name})s' for name in COLUMN_NAMES)})" @@ -177,6 +203,11 @@ def opener() -> psycopg.Connection[Any]: self._declare = to_driver_placeholders(contract_sql(DECLARE_BOUNDS_SQL)) self._search = to_driver_placeholders(contract_sql(SEARCH_SQL)) self._record_injection = to_driver_placeholders(contract_sql(RECORD_INJECTION_SQL)) + self._write_statements = { + key: to_driver_placeholders(text) for key, text in contract_statements(WRITE_EMBEDDING_SQL).items() + } + # The vector search statements keep their ``{dim}`` until the query vector's length is known. + self._vector_statements = contract_statements(SEARCH_VECTOR_SQL) @contextmanager def _transaction(self, bounds: AuthorizationBounds) -> Iterator[psycopg.Cursor[Any]]: @@ -288,6 +319,92 @@ def search(self, bounds: AuthorizationBounds, parameters: SearchParameters) -> l ) return result + def write_embedding(self, bounds: AuthorizationBounds, write: EmbeddingWrite) -> EmbeddingWriteResult: + """ + Stores a record's embedding (contract 1.6): in one transaction that declared the bounds for the record's scope, + ``write`` from ``sql/write-experience-embedding.sql`` and, when it wrote no row, ``probe-revision``. + ``Written``; + ``Stale`` with the record's current revision; ``Missing`` for an unknown ID, a record in another scope or an + erased one. ``write`` is validated first (``ValueError``); nothing runs when it is refused. Needs the grants of + ``sql/application-role-privileges-1.6.sql`` and the vectors migrations. + """ + scope = write.scope + parameters: dict[str, Any] = { + "experience_id": write.experience_id, + "tenant_id": scope.tenant_id, + "application_id": scope.application_id, + "project_id": scope.project_id, + "team_id": scope.team_id, + "agent_id": scope.agent_id, + "user_id": scope.user_id, + "model_id": write.model_id, + "dimension": write.dimension, + "content_hash": write.content_hash, + "source_revision": write.source_revision, + "embedding": vector_literal(write.vector), + "now": write.now.truncate_to_microseconds().to_datetime(), + } + with self._transaction(bounds) as cursor: + cursor.execute(self._write_statements["write"], parameters) + if cursor.rowcount == 1: + return EmbeddingWriteResult("Written") + cursor.execute(self._write_statements["probe-revision"], parameters) + row = cursor.fetchone() + return EmbeddingWriteResult("Missing") if row is None else EmbeddingWriteResult("Stale", int(row[0])) + + def search_vectors(self, bounds: AuthorizationBounds, query: VectorQuery) -> VectorSearchResult: + """ + The nearest-neighbour search (contract 1.6): in one transaction that declared the bounds, ``search`` (or + ``search-excluding-model-authored``) from ``sql/search-experience-vector.sql`` with ``{dim}`` the query + vector's length and, when it returns no row, the matching compatibility probe. Rows come back nearest first, + ties by experience ID, NaN last. A sealed row (``payload_version`` 2) is skipped. ``query`` is validated first. + """ + from psycopg.types.numeric import Int4 + + scope = query.scope + excluding = query.exclude_model_authored + search = self._vector_statements["search-excluding-model-authored" if excluding else "search"] + probe = self._vector_statements["compatibility-excluding-model-authored" if excluding else "compatibility"] + search = to_driver_placeholders(search.replace("{dim}", str(len(query.vector)))) + probe = to_driver_placeholders(probe) + parameters: dict[str, Any] = { + "tenant_id": scope.tenant_id, + "application_id": scope.application_id, + "project_id": scope.project_id, + "team_id": scope.team_id, + "agent_id": scope.agent_id, + "user_id": scope.user_id, + "statuses": list(query.statuses), + "min_confidence": float(query.min_confidence), + "model_id": query.model_id, + "query_vector": vector_literal(query.vector), + "limit": Int4(query.limit), + } + with self._transaction(bounds) as cursor: + cursor.execute(search, parameters) + names = [column.name for column in cursor.description or ()] + rows = cursor.fetchall() + compatible: tuple[Any, ...] | None = None + if not rows: + cursor.execute(probe, {key: value for key, value in parameters.items() if key != "limit"}) + compatible = cursor.fetchone() + if not rows: + if compatible is None or not compatible[0]: + return VectorSearchResult("Found") + return VectorSearchResult("DimensionMismatch" if compatible[1] else "ModelMismatch") + count = len(COLUMN_NAMES) + if names[:count] != list(COLUMN_NAMES) or names[-1:] != ["distance"]: + raise RuntimeError("The vector search returned columns in an unexpected order.") + hits: list[VectorHit] = [] + for row in rows: + columns = dict(zip(names, row, strict=True)) + if columns["payload_version"] != PAYLOAD_VERSION: + continue + if columns["distance"] is None: + raise RuntimeError("The vector search returned a row with no distance.") + hits.append(VectorHit(from_columns(columns, columns["payload"]), float(columns["distance"]))) + return VectorSearchResult("Found", tuple(order_hits(hits))) + def record_injection( self, bounds: AuthorizationBounds, diff --git a/python/src/agentexperience/sql/search-experience-vector.sql b/python/src/agentexperience/sql/search-experience-vector.sql new file mode 100644 index 00000000..702302d8 --- /dev/null +++ b/python/src/agentexperience/sql/search-experience-vector.sql @@ -0,0 +1,40 @@ +-- contract 1.6: the nearest-neighbour search over stored embeddings, in a transaction that ran declare-bounds.sql for +-- the reader's scope. Four statements, each starting at its "-- statement:" line and running to the next one (or the +-- end): run "search", or "search-excluding-model-authored" when model-written lessons are excluded; only when it +-- returns no row, run the matching compatibility probe ("compatibility" or "compatibility-excluding-model-authored") +-- in the same transaction. Parameters: +-- @tenant_id, @application_id, @project_id text, the reader's scope +-- @team_id, @agent_id, @user_id text, NULL when the scope leaves the field unset +-- @statuses text[], {Validated,Reinforced} +-- @min_confidence double precision, 0.5 +-- @model_id text, the query vector's model ID: non-blank, at most 256 UTF-16 +-- code units +-- @query_vector text, the query vector as a pgvector literal, [x,y,...] +-- @limit integer, 51 (the candidate limit, 50, plus one); search only +-- {dim} is not a parameter: replace every {dim} with the query vector's length (1 to 16000) in decimal digits before +-- running the statement. A pgvector type modifier is part of the type, so it cannot be bound. The probes have no {dim}. +-- Rows come back nearest first; the statement has no tie-break, so order rows of equal distance by experience_id +-- ascending (the lower-case hyphenated text, ordinally). distance is double precision, NaN for a zero-magnitude vector; +-- relevance = clamp(1 - distance / 2, 0, 1), and 0 for NaN. Shared rows are out of contract 1.x scope: the statement +-- matches the reader's own scope only, as the text search does with @with_grants false. +-- An empty search's outcome is the probe's: both false, or the first false, is Found (nothing comparable is stored); +-- the first true and the second false is ModelMismatch; both true is DimensionMismatch. +-- The application role needs the grants in application-role-privileges-1.6.sql. +-- statement: search +SELECT r.experience_id, r.source_run_id, r.tenant_id, r.application_id, r.project_id, r.team_id, r.agent_id, r.user_id, r.task_id, r.status, r.reuse_confidence, r.supporting_validations, r.contradictions, r.revision, r.created_at, r.updated_at, r.payload_version, r.payload, false AS shared_by_grant, NULL::uuid AS permitting_grant_id, NULL::text AS permitting_grant_disclosure, (e.embedding::vector({dim}) <=> CAST(@query_vector AS vector({dim}))) AS distance +FROM agent_experience.experience_embeddings e JOIN agent_experience.experience_records r ON r.experience_id = e.experience_id +WHERE (e.tenant_id = @tenant_id AND e.application_id = @application_id AND e.project_id = @project_id AND e.team_id IS NOT DISTINCT FROM @team_id AND e.agent_id IS NOT DISTINCT FROM @agent_id AND e.user_id IS NOT DISTINCT FROM @user_id AND r.tenant_id = @tenant_id AND r.application_id = @application_id AND r.project_id = @project_id AND r.team_id IS NOT DISTINCT FROM @team_id AND r.agent_id IS NOT DISTINCT FROM @agent_id AND r.user_id IS NOT DISTINCT FROM @user_id) AND r.deleted_at IS NULL AND r.status = ANY(@statuses) AND r.reuse_confidence >= @min_confidence +AND e.model_id = @model_id AND e.dimension = {dim} +ORDER BY (e.embedding::vector({dim}) <=> CAST(@query_vector AS vector({dim}))) LIMIT @limit +-- statement: search-excluding-model-authored +SELECT r.experience_id, r.source_run_id, r.tenant_id, r.application_id, r.project_id, r.team_id, r.agent_id, r.user_id, r.task_id, r.status, r.reuse_confidence, r.supporting_validations, r.contradictions, r.revision, r.created_at, r.updated_at, r.payload_version, r.payload, false AS shared_by_grant, NULL::uuid AS permitting_grant_id, NULL::text AS permitting_grant_disclosure, (e.embedding::vector({dim}) <=> CAST(@query_vector AS vector({dim}))) AS distance +FROM agent_experience.experience_embeddings e JOIN agent_experience.experience_records r ON r.experience_id = e.experience_id +WHERE (e.tenant_id = @tenant_id AND e.application_id = @application_id AND e.project_id = @project_id AND e.team_id IS NOT DISTINCT FROM @team_id AND e.agent_id IS NOT DISTINCT FROM @agent_id AND e.user_id IS NOT DISTINCT FROM @user_id AND r.tenant_id = @tenant_id AND r.application_id = @application_id AND r.project_id = @project_id AND r.team_id IS NOT DISTINCT FROM @team_id AND r.agent_id IS NOT DISTINCT FROM @agent_id AND r.user_id IS NOT DISTINCT FROM @user_id) AND r.deleted_at IS NULL AND r.status = ANY(@statuses) AND r.reuse_confidence >= @min_confidence +AND r.reflection_model_authored IS FALSE AND e.model_id = @model_id AND e.dimension = {dim} +ORDER BY (e.embedding::vector({dim}) <=> CAST(@query_vector AS vector({dim}))) LIMIT @limit +-- statement: compatibility +SELECT EXISTS (SELECT 1 FROM agent_experience.experience_embeddings e JOIN agent_experience.experience_records r ON r.experience_id = e.experience_id WHERE (e.tenant_id = @tenant_id AND e.application_id = @application_id AND e.project_id = @project_id AND e.team_id IS NOT DISTINCT FROM @team_id AND e.agent_id IS NOT DISTINCT FROM @agent_id AND e.user_id IS NOT DISTINCT FROM @user_id AND r.tenant_id = @tenant_id AND r.application_id = @application_id AND r.project_id = @project_id AND r.team_id IS NOT DISTINCT FROM @team_id AND r.agent_id IS NOT DISTINCT FROM @agent_id AND r.user_id IS NOT DISTINCT FROM @user_id) AND r.deleted_at IS NULL AND r.status = ANY(@statuses) AND r.reuse_confidence >= @min_confidence), +EXISTS (SELECT 1 FROM agent_experience.experience_embeddings e JOIN agent_experience.experience_records r ON r.experience_id = e.experience_id WHERE (e.tenant_id = @tenant_id AND e.application_id = @application_id AND e.project_id = @project_id AND e.team_id IS NOT DISTINCT FROM @team_id AND e.agent_id IS NOT DISTINCT FROM @agent_id AND e.user_id IS NOT DISTINCT FROM @user_id AND r.tenant_id = @tenant_id AND r.application_id = @application_id AND r.project_id = @project_id AND r.team_id IS NOT DISTINCT FROM @team_id AND r.agent_id IS NOT DISTINCT FROM @agent_id AND r.user_id IS NOT DISTINCT FROM @user_id) AND r.deleted_at IS NULL AND r.status = ANY(@statuses) AND r.reuse_confidence >= @min_confidence AND e.model_id = @model_id) +-- statement: compatibility-excluding-model-authored +SELECT EXISTS (SELECT 1 FROM agent_experience.experience_embeddings e JOIN agent_experience.experience_records r ON r.experience_id = e.experience_id WHERE (e.tenant_id = @tenant_id AND e.application_id = @application_id AND e.project_id = @project_id AND e.team_id IS NOT DISTINCT FROM @team_id AND e.agent_id IS NOT DISTINCT FROM @agent_id AND e.user_id IS NOT DISTINCT FROM @user_id AND r.tenant_id = @tenant_id AND r.application_id = @application_id AND r.project_id = @project_id AND r.team_id IS NOT DISTINCT FROM @team_id AND r.agent_id IS NOT DISTINCT FROM @agent_id AND r.user_id IS NOT DISTINCT FROM @user_id) AND r.deleted_at IS NULL AND r.status = ANY(@statuses) AND r.reuse_confidence >= @min_confidence AND r.reflection_model_authored IS FALSE), +EXISTS (SELECT 1 FROM agent_experience.experience_embeddings e JOIN agent_experience.experience_records r ON r.experience_id = e.experience_id WHERE (e.tenant_id = @tenant_id AND e.application_id = @application_id AND e.project_id = @project_id AND e.team_id IS NOT DISTINCT FROM @team_id AND e.agent_id IS NOT DISTINCT FROM @agent_id AND e.user_id IS NOT DISTINCT FROM @user_id AND r.tenant_id = @tenant_id AND r.application_id = @application_id AND r.project_id = @project_id AND r.team_id IS NOT DISTINCT FROM @team_id AND r.agent_id IS NOT DISTINCT FROM @agent_id AND r.user_id IS NOT DISTINCT FROM @user_id) AND r.deleted_at IS NULL AND r.status = ANY(@statuses) AND r.reuse_confidence >= @min_confidence AND r.reflection_model_authored IS FALSE AND e.model_id = @model_id) diff --git a/python/src/agentexperience/sql/write-experience-embedding.sql b/python/src/agentexperience/sql/write-experience-embedding.sql new file mode 100644 index 00000000..2ebb5e05 --- /dev/null +++ b/python/src/agentexperience/sql/write-experience-embedding.sql @@ -0,0 +1,31 @@ +-- contract 1.6: stores one record's embedding, conditionally, in a transaction that ran declare-bounds.sql for the same +-- scope. Two statements, each starting at its "-- statement:" line and running to the next one (or the end); run +-- "write", and only when it reports 0 rows written, "probe-revision" in the same transaction, then commit. Parameters: +-- @experience_id uuid, the record whose summary was embedded +-- @tenant_id, @application_id, @project_id text, the record's scope (exactly; a write never crosses a scope) +-- @team_id, @agent_id, @user_id text, NULL when the scope leaves the field unset +-- @model_id text, the embedding model's ID: non-blank, at most 256 UTF-16 code units +-- @dimension integer, the vector's length, 1 to 16000 +-- @content_hash text, the content hash of the summary that was embedded (README-1.6.md) +-- @source_revision bigint, the record revision the summary was read at +-- @embedding text, the vector as a pgvector literal, [x,y,...]; any text pgvector +-- parses as a vector of @dimension finite components +-- @now timestamptz, the writer's clock, truncated to whole microseconds in UTC +-- "write" inserts or replaces the row only for a record that is live (not erased) in exactly that scope and still at +-- exactly @source_revision, its scope columns copied from the record, never from the parameters; the record row is +-- locked FOR KEY SHARE, so an erasure racing it either waits and then removes the vector, or runs first and leaves a +-- tombstone this statement does not select. A row already stored from a newer revision is never replaced by an older +-- one. Outcomes: 1 row written is Written. 0 rows: "probe-revision" returns the record's current revision (Stale, with +-- that revision) or no row at all (Missing: an unknown ID, a record in another scope, or an erased record). +-- The application role needs the grants in application-role-privileges-1.6.sql. +-- statement: write +INSERT INTO agent_experience.experience_embeddings AS t (experience_id, tenant_id, application_id, project_id, team_id, agent_id, user_id, model_id, dimension, content_hash, source_revision, embedding, created_at, updated_at) +SELECT r.experience_id, r.tenant_id, r.application_id, r.project_id, r.team_id, r.agent_id, r.user_id, @model_id, @dimension, @content_hash, @source_revision, CAST(@embedding AS vector), @now, @now +FROM agent_experience.experience_records r +WHERE r.experience_id = @experience_id AND r.tenant_id = @tenant_id AND r.application_id = @application_id AND r.project_id = @project_id AND r.team_id IS NOT DISTINCT FROM @team_id AND r.agent_id IS NOT DISTINCT FROM @agent_id AND r.user_id IS NOT DISTINCT FROM @user_id AND r.deleted_at IS NULL AND r.revision = @source_revision +FOR KEY SHARE OF r +ON CONFLICT (experience_id) DO UPDATE SET model_id = EXCLUDED.model_id, dimension = EXCLUDED.dimension, content_hash = EXCLUDED.content_hash, source_revision = EXCLUDED.source_revision, embedding = EXCLUDED.embedding, updated_at = EXCLUDED.updated_at +WHERE t.source_revision <= EXCLUDED.source_revision +-- statement: probe-revision +SELECT r.revision FROM agent_experience.experience_records r +WHERE r.experience_id = @experience_id AND r.tenant_id = @tenant_id AND r.application_id = @application_id AND r.project_id = @project_id AND r.team_id IS NOT DISTINCT FROM @team_id AND r.agent_id IS NOT DISTINCT FROM @agent_id AND r.user_id IS NOT DISTINCT FROM @user_id AND r.deleted_at IS NULL diff --git a/python/tests/_contract.py b/python/tests/_contract.py index cb9952c8..2a086340 100644 --- a/python/tests/_contract.py +++ b/python/tests/_contract.py @@ -43,23 +43,14 @@ def cases(family: str) -> list[str]: return sorted(entry.name for entry in (FIXTURES / family).iterdir() if entry.is_dir()) -PENDING_CONTRACT_VERSIONS: tuple[str, ...] = ("1.6",) -"""Contract versions the lock may name although this package does not pass them yet: contract 1.6 only adds -``PENDING_VECTOR_FAMILIES``, so every 1.0 to 1.5 fixture still runs under its lock.""" - -PENDING_VECTOR_FAMILIES: tuple[str, ...] = ("embedding-summary", "embedding-write", "hybrid-merge", "vector-search") -"""Contract 1.6's vector fixture families, which this package does not implement yet. They are skipped by name -(``test_pending_vector_family``), never ignored. Stories 23.2 and 23.3 implement them and remove this.""" - - class UnsupportedContractVersion(Exception): pass def check_contract_version(lock: dict[str, Any]) -> str: - """The lock's contract version, when this package supports it or it is pending; a clear failure otherwise.""" + """The lock's contract version, when this package supports it; a clear failure otherwise.""" version = lock.get("contractVersion") - if version not in SUPPORTED_CONTRACT_VERSIONS + PENDING_CONTRACT_VERSIONS: + if version not in SUPPORTED_CONTRACT_VERSIONS: raise UnsupportedContractVersion( f"contract/lock.json is contract version {version!r}, but this package supports only " f"{', '.join(SUPPORTED_CONTRACT_VERSIONS)}. Implement the new version before running its fixtures." diff --git a/python/tests/_postgres.py b/python/tests/_postgres.py index 1d807bb2..572e48ef 100644 --- a/python/tests/_postgres.py +++ b/python/tests/_postgres.py @@ -35,6 +35,8 @@ PRIVILEGES_SQL = CONTRACT / "sql" / "application-role-privileges.sql" PRIVILEGES_1_4_SQL = CONTRACT / "sql" / "application-role-privileges-1.4.sql" +PRIVILEGES_1_6_SQL = CONTRACT / "sql" / "application-role-privileges-1.6.sql" +VECTOR_MIGRATIONS = MIGRATIONS.parent.parent / "AgentExperience.Storage.Postgres.Vectors" / "Migrations" _ROLE_PLACEHOLDER = ':"application_role"' @@ -73,7 +75,13 @@ def application_role_privileges( APPLICATION_ROLE_PRIVILEGES = APPLICATION_ROLE_PRIVILEGES_1_0 + application_role_privileges( APPLICATION, path=PRIVILEGES_1_4_SQL ) -"""The 1.0 file, then the 1.4 file: what the .NET privileges call grants.""" +"""The 1.0 file, then the 1.4 file: what the .NET privileges call grants on the base schema.""" + +APPLICATION_ROLE_PRIVILEGES_1_6 = APPLICATION_ROLE_PRIVILEGES + application_role_privileges( + APPLICATION, path=PRIVILEGES_1_6_SQL +) +"""The 1.0 file, then the 1.4 file, then the 1.6 file (the embeddings table's): what the .NET privileges call grants +on a schema with the vectors migrations applied.""" @dataclass(frozen=True) @@ -119,7 +127,8 @@ def server() -> Generator[Server, None, None]: from testcontainers.postgres import PostgresContainer major = os.environ.get("AGENTEXPERIENCE_POSTGRES_MAJOR") or "16" - image = _apply_hub_image_prefix(f"postgres:{major}") + # pgvector's image: the vector channel (contract 1.6) needs the extension. + image = _apply_hub_image_prefix(f"pgvector/pgvector:pg{major}") with PostgresContainer(image, username="postgres", password="postgres", dbname="postgres") as pg: yield Server(pg.get_container_host_ip(), int(pg.get_exposed_port(5432)), "postgres", "postgres") @@ -168,8 +177,9 @@ def database_name(label: str) -> str: def create_database(admin: Server, label: str, *, privileges: list[str] | None = None) -> str: - """A fresh database, migrated by the .NET scripts and with the application role's privileges (both contract - files unless ``privileges`` says otherwise); its name.""" + """A fresh database, migrated by the .NET scripts (the vectors package's too, after ``CREATE EXTENSION vector`` as a + superuser) and with the application role's privileges (the three contract files unless ``privileges`` says + otherwise); its name.""" import psycopg from psycopg import sql @@ -180,12 +190,17 @@ def create_database(admin: Server, label: str, *, privileges: list[str] | None = scripts = sorted(MIGRATIONS.glob("*.sql"), key=lambda path: path.name) assert scripts, f"No migration scripts under {MIGRATIONS}." + vector_scripts = sorted(VECTOR_MIGRATIONS.glob("*.sql"), key=lambda path: path.name) + assert vector_scripts, f"No migration scripts under {VECTOR_MIGRATIONS}." + with psycopg.connect(admin.conninfo(name), autocommit=True) as superuser: + # pgvector is not a trusted extension: a superuser installs it, as an operator does (contract 1.6). + superuser.execute("CREATE EXTENSION IF NOT EXISTS vector") with psycopg.connect(admin.conninfo(name, OWNER, PASSWORD)) as owner: - for script in scripts: + for script in [*scripts, *vector_scripts]: with owner.transaction(): owner.execute(script.read_text(encoding="utf-8").encode("utf-8")) with owner.transaction(): - for statement in APPLICATION_ROLE_PRIVILEGES if privileges is None else privileges: + for statement in APPLICATION_ROLE_PRIVILEGES_1_6 if privileges is None else privileges: owner.execute(statement.encode("utf-8")) return name diff --git a/python/tests/test_contract.py b/python/tests/test_contract.py index 2598671b..e8d0911a 100644 --- a/python/tests/test_contract.py +++ b/python/tests/test_contract.py @@ -8,6 +8,7 @@ import copy import hashlib import json +import math import re import uuid from collections.abc import Iterator @@ -20,8 +21,6 @@ from _contract import ( CONTRACT, FIXTURES, - PENDING_CONTRACT_VERSIONS, - PENDING_VECTOR_FAMILIES, SCHEMA, UnsupportedContractVersion, cases, @@ -46,10 +45,13 @@ RequiredCheck, RunCapture, Scope, + content_hash, derive_task_text, error_class, from_stored, + merge_channels, render_block, + retrieval_summary, to_stored, ) from agentexperience.block import ranked_from_json @@ -108,12 +110,12 @@ def assert_json_equal(expected: Any, actual: Any, what: str) -> None: def test_the_contract_version_is_one_this_package_supports() -> None: version = check_contract_version(read_json(CONTRACT / "lock.json")) - assert version == CONTRACT_VERSION or version in PENDING_CONTRACT_VERSIONS + assert version == CONTRACT_VERSION @pytest.mark.parametrize("version", ["2.0", "1.7", None]) def test_an_unsupported_contract_version_fails_with_a_clear_message(version: str | None) -> None: - with pytest.raises(UnsupportedContractVersion, match=r"supports only 1\.0, 1\.1, 1\.2, 1\.3, 1\.4, 1\.5"): + with pytest.raises(UnsupportedContractVersion, match=r"supports only 1\.0, 1\.1, 1\.2, 1\.3, 1\.4, 1\.5, 1\.6"): check_contract_version({"contractVersion": version}) @@ -132,7 +134,16 @@ def test_every_contract_file_matches_the_lock() -> None: assert hashlib.sha256((CONTRACT / name).read_bytes()).hexdigest() == digest, name -@pytest.mark.parametrize("name", ["declare-bounds.sql", "search-experience-text.sql", "record-injection.sql"]) +@pytest.mark.parametrize( + "name", + [ + "declare-bounds.sql", + "search-experience-text.sql", + "record-injection.sql", + "write-experience-embedding.sql", + "search-experience-vector.sql", + ], +) def test_the_packaged_sql_is_the_contract_sql(name: str) -> None: assert contract_sql(name) == (CONTRACT / "sql" / name).read_text(encoding="utf-8") @@ -156,12 +167,45 @@ def test_every_fixture_family_has_cases() -> None: assert cases(family), family -@pytest.mark.parametrize("family", PENDING_VECTOR_FAMILIES) -def test_pending_vector_family(family: str) -> None: - # Contract 1.6's vector families, skipped by name: this package does not implement them yet (stories 23.2 and 23.3 - # do, and remove PENDING_VECTOR_FAMILIES). Each must exist, so a renamed family fails here rather than going unseen. - assert (FIXTURES / family).is_dir(), family - pytest.skip(f"{family}: a contract 1.6 vector family, pending in this package") +# ---- Contract 1.6: the vector families that need no database ----------------------------------------------------- + + +@pytest.mark.parametrize("index", range(len(read_json(FIXTURES / "embedding-summary" / "cases.json")))) +def test_embedding_summary(index: int) -> None: + case = read_json(FIXTURES / "embedding-summary" / "cases.json")[index] + summary = retrieval_summary(case["taskId"], case["taskSummary"], case["lesson"]) + assert summary == case["summary"], case["description"] + assert content_hash(case["modelId"], summary) == case["contentHash"], case["description"] + + +def _nan(value: Any) -> float: + return float("nan") if value == "NaN" else float(value) + + +def _same_relevance(expected: Any, actual: float) -> bool: + return math.isnan(actual) if expected == "NaN" else actual == expected + + +@pytest.mark.parametrize("index", range(len(read_json(FIXTURES / "hybrid-merge" / "cases.json")))) +def test_hybrid_merge(index: int) -> None: + case = read_json(FIXTURES / "hybrid-merge" / "cases.json")[index] + + def entries(channel: list[dict[str, Any]]) -> list[tuple[uuid.UUID, float]]: + return [(uuid.UUID(item["experienceId"]), _nan(item["relevance"])) for item in channel] + + merged = merge_channels( + entries(case["text"]), + entries(case["vector"]), + lambda item: item[0], + lambda item: item[1], + case["minimumVectorRelevance"], + ) + assert [str(entry.item[0]) for entry in merged] == [item["experienceId"] for item in case["merged"]], case[ + "description" + ] + assert [entry.source for entry in merged] == [item["from"] for item in case["merged"]] + for entry, expected in zip(merged, case["merged"], strict=True): + assert _same_relevance(expected["relevance"], entry.relevance), (case["description"], expected, entry) # ---- finalize --------------------------------------------------------------------------------------------------- diff --git a/python/tests/test_vectors.py b/python/tests/test_vectors.py new file mode 100644 index 00000000..392fa61b --- /dev/null +++ b/python/tests/test_vectors.py @@ -0,0 +1,709 @@ +""" +Story 23.2: semantic retrieval in Python (contract 1.6). The vector fixture families through the PostgreSQL store (and +the in-memory store where a case needs nothing only a database does), and the agent: every row of the fallback matrix, +the floor, the background embedding write, and a record embedded by the contract's own SQL found by a Python agent. +""" + +from __future__ import annotations + +import asyncio +import math +import time +import uuid +from collections.abc import Sequence +from dataclasses import replace +from typing import Any + +import psycopg +import pytest +from _contract import FIXTURES, cases, json_difference, read_json, scope_of + +pytest.importorskip("agent_framework") + +from _maf import ScriptedClient +from _postgres import OWNER, PASSWORD, Server, application_conninfo, create_database +from _sync import run +from agent_framework import Agent, tool +from test_postgres import _erase + +from agentexperience import ( + AuthorizationBounds, + EmbeddingWrite, + ExperienceRecord, + Instant, + Scope, + SearchParameters, + VectorQuery, + content_hash, + from_stored, + retrieval_summary, +) +from agentexperience.embedding import record_summary, relevance_of +from agentexperience.inmemory import EmbeddingIndex, InMemoryExperienceStore +from agentexperience.maf import ( + CaptureFailure, + ExperienceContextProvider, + ExperienceIdentity, + ExperienceSettings, + InjectionFailure, + InjectionResult, + VectorFallback, + tool_succeeded, + use_agent_experience, +) +from agentexperience.postgres import ( + DECLARE_BOUNDS_SQL, + PostgresExperienceStore, + contract_sql, + contract_statements, + to_driver_placeholders, +) + +TOLERANCE = 1e-9 +SCOPE = Scope("contract-tenant", "contract-app", "contract-project", agent_id="refund-agent") +BOUNDS = AuthorizationBounds.of(SCOPE) +MODEL = "contract-embedding" +PARAPHRASE = "Renew the cryptographic credentials immediately" +NOW = Instant.parse("2026-03-02T10:00:00+00:00") + + +def seed_records(family: str, name: str) -> list[ExperienceRecord]: + return [from_stored(item) for item in read_json(FIXTURES / family / name / "seed.json")] + + +def write_of(step: dict[str, Any]) -> EmbeddingWrite: + return EmbeddingWrite( + uuid.UUID(step["experienceId"]), + scope_of(step["scope"]), + step["modelId"], + step["dimension"], + step["contentHash"], + step["sourceRevision"], + tuple(float(c) for c in step["vector"]), + Instant.parse(step["now"]), + ) + + +def query_of(step: dict[str, Any]) -> VectorQuery: + return VectorQuery( + scope_of(step["scope"]), + step["modelId"], + tuple(float(c) for c in step["vector"]), + tuple(step["statuses"]), + step["minConfidence"], + step["limit"], + step["excludeModelAuthored"], + ) + + +def search_outcome(result: Any) -> dict[str, Any]: + return { + "outcome": result.outcome, + "rows": [ + { + "experienceId": str(hit.record.experience_id), + "distance": None if math.isnan(hit.distance) else hit.distance, + "relevance": hit.relevance, + } + for hit in result.hits + ], + } + + +def write_outcome(result: Any) -> dict[str, Any]: + return {"outcome": result.outcome} | ( + {"currentRevision": result.current_revision} if result.outcome == "Stale" else {} + ) + + +# ---- The fixtures through the PostgreSQL store --------------------------------------------- + + +def _set_revision(owner: str, step: dict[str, Any]) -> None: + with psycopg.connect(owner) as connection: + connection.execute( + "UPDATE agent_experience.experience_records SET revision = %s WHERE experience_id = %s::uuid", + (step["revision"], step["experienceId"]), + ) + + +def run_steps_on_postgres(pg_server: Server, family: str, name: str) -> tuple[list[Any], str]: + directory = FIXTURES / family / name + steps = read_json(directory / "steps.json")["steps"] + database = create_database(pg_server, f"{family}_{name}") + owner = pg_server.conninfo(database, OWNER, PASSWORD) + store = PostgresExperienceStore(application_conninfo(pg_server, database)) + for record in seed_records(family, name): + store.insert_record(AuthorizationBounds.of(record.scope), record) + actual: list[Any] = [] + for step in steps: + if "write" in step: + write = write_of(step["write"]) + actual.append(write_outcome(store.write_embedding(AuthorizationBounds.of(write.scope), write))) + elif "search" in step: + query = query_of(step["search"]) + actual.append(search_outcome(store.search_vectors(AuthorizationBounds.of(query.scope), query))) + elif "setRevision" in step: + _set_revision(owner, step["setRevision"]) + actual.append(None) + else: + _erase(owner, step["erase"]) + actual.append(None) + return actual, owner + + +def stored_embeddings(owner: str, columns: Sequence[str]) -> list[dict[str, Any]]: + with psycopg.connect(owner) as connection: + connection.execute("SET TIME ZONE 'UTC'") + rows = connection.execute( + "SELECT " + ", ".join(f"{c}::text" for c in columns) + " FROM agent_experience.experience_embeddings " + "ORDER BY experience_id" + ).fetchall() + return [dict(zip(columns, row, strict=True)) for row in rows] + + +@pytest.mark.postgres +@pytest.mark.parametrize("name", cases("embedding-write")) +def test_embedding_write(pg_server: Server, name: str) -> None: + expected = read_json(FIXTURES / "embedding-write" / name / "expected.json") + actual, owner = run_steps_on_postgres(pg_server, "embedding-write", name) + assert json_difference(expected["steps"], actual, "$.steps", TOLERANCE) is None, actual + columns = list(expected["experience_embeddings"][0]) if expected["experience_embeddings"] else ["experience_id"] + rows = stored_embeddings(owner, columns) + assert json_difference(expected["experience_embeddings"], rows, "$.experience_embeddings") is None, rows + + +@pytest.mark.postgres +@pytest.mark.parametrize("name", cases("vector-search")) +def test_vector_search(pg_server: Server, name: str) -> None: + expected = read_json(FIXTURES / "vector-search" / name / "expected.json") + actual, owner = run_steps_on_postgres(pg_server, "vector-search", name) + assert json_difference(expected["steps"], actual, "$.steps", TOLERANCE) is None, actual + columns = list(expected["experience_embeddings"][0]) if expected["experience_embeddings"] else ["experience_id"] + rows = stored_embeddings(owner, columns) + assert json_difference(expected["experience_embeddings"], rows, "$.experience_embeddings") is None, rows + + +@pytest.mark.parametrize(("family", "name"), [(f, n) for f in ("embedding-write", "vector-search") for n in cases(f)]) +def test_the_in_memory_store_gives_the_fixtures_outcomes(family: str, name: str) -> None: + steps = read_json(FIXTURES / family / name / "steps.json")["steps"] + if any("erase" in step or "setRevision" in step for step in steps): + pytest.skip("needs a database: an erasure or a lifecycle commit") + expected = read_json(FIXTURES / family / name / "expected.json")["steps"] + store = InMemoryExperienceStore(seed_records(family, name)) + for step, outcome in zip(steps, expected, strict=True): + if "write" in step: + write = write_of(step["write"]) + assert write_outcome(store.write_embedding(AuthorizationBounds.of(write.scope), write)) == outcome + else: + query = query_of(step["search"]) + actual = search_outcome(store.search_vectors(AuthorizationBounds.of(query.scope), query)) + assert json_difference(outcome, actual, "$", 1e-6) is None, actual + + +def test_the_in_memory_store_never_replaces_a_newer_revision_with_an_older_one() -> None: + records = [replace(record, revision=2) for record in seed_records("vector-search", "nearest-first")][:1] + store = InMemoryExperienceStore(records) + experience_id = records[0].experience_id + base = EmbeddingWrite(experience_id, SCOPE, MODEL, 2, "h2", 2, (0.0, 1.0), NOW) + assert store.write_embedding(BOUNDS, base).outcome == "Written" + older = replace(base, source_revision=1, content_hash="h1", vector=(1.0, 0.0)) + stale = store.write_embedding(BOUNDS, older) + assert (stale.outcome, stale.current_revision) == ("Stale", 2) + assert store.embedding(experience_id) == (MODEL, 2, "h2", 2, (0.0, 1.0)) + other = replace(base, scope=replace(SCOPE, agent_id="other")) + assert store.write_embedding(AuthorizationBounds.of(other.scope), other).outcome == "Missing" + + +def test_the_in_memory_store_compares_32_bit_floats_and_orders_ties_by_id() -> None: + records = seed_records("vector-search", "an-exact-tie-is-ordered-by-id") + store = InMemoryExperienceStore(records) + for record, vector in zip( + sorted(records, key=lambda r: str(r.experience_id)), ((1.0, 0, 0, 0), (2.0, 0, 0, 0)), strict=False + ): + store.write_embedding( + BOUNDS, EmbeddingWrite(record.experience_id, SCOPE, MODEL, 4, "h", record.revision, vector, NOW) + ) + result = store.search_vectors(BOUNDS, VectorQuery(SCOPE, MODEL, (1.0, 0.0, 0.0, 0.0))) + assert [str(hit.record.experience_id) for hit in result.hits] == sorted( + str(h.record.experience_id) for h in result.hits + ) + assert relevance_of(float("nan")) == 0 and relevance_of(-1) == 1 and relevance_of(3) == 0 + # Outside the bounds nothing is found and nothing is written. + narrow = AuthorizationBounds.of(replace(SCOPE, agent_id="someone-else")) + assert store.search_vectors(narrow, VectorQuery(SCOPE, MODEL, (1.0, 0.0, 0.0, 0.0))).hits == () + + +# ---- Validation: refused before any statement --------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "build", + [ + lambda: EmbeddingWrite(uuid.uuid4(), SCOPE, " ", 2, "h", 1, (1.0, 0.0), NOW), + lambda: EmbeddingWrite(uuid.uuid4(), SCOPE, "m" * 257, 2, "h", 1, (1.0, 0.0), NOW), + lambda: EmbeddingWrite(uuid.uuid4(), SCOPE, "m", 3, "h", 1, (1.0, 0.0), NOW), + lambda: EmbeddingWrite(uuid.uuid4(), SCOPE, "m", 0, "h", 1, (), NOW), + lambda: EmbeddingWrite(uuid.uuid4(), SCOPE, "m", 2, "h", 1, (float("nan"), 0.0), NOW), + lambda: EmbeddingWrite(uuid.uuid4(), SCOPE, "m", 2, "h", 1, (float("inf"), 0.0), NOW), + lambda: EmbeddingWrite(uuid.uuid4(), SCOPE, "m", 2, " ", 1, (1.0, 0.0), NOW), + lambda: EmbeddingWrite(uuid.uuid4(), SCOPE, "m", 2, "h", -1, (1.0, 0.0), NOW), + lambda: EmbeddingWrite(uuid.uuid4(), SCOPE, "m", 16_001, "h", 1, (0.0,) * 16_001, NOW), + lambda: VectorQuery(SCOPE, "", (1.0,)), + lambda: VectorQuery(SCOPE, "m", ()), + lambda: VectorQuery(SCOPE, "m", (float("nan"),)), + lambda: VectorQuery(SCOPE, "m", (1e300,)), + ], +) +def test_an_invalid_write_or_search_is_refused(build: Any) -> None: + with pytest.raises(ValueError): + build() + + +def test_a_zero_magnitude_vector_is_valid() -> None: + assert VectorQuery(SCOPE, "m", (0.0, 0.0)).vector == (0.0, 0.0) + + +# ---- The agent ------------------------------------------------------------------------------------------------------ + + +class FixedEmbedder: + """The caller's embedder: one vector for every text, unless it is told to misbehave.""" + + def __init__( + self, + vector: Sequence[float] = (1.0, 0.0, 0.0, 0.0), + *, + model_id: str = MODEL, + dimension: int | None = None, + fail: bool = False, + delay: float = 0.0, + ) -> None: + self.vector = list(vector) + self.model_id = model_id + self.dimension = len(self.vector) if dimension is None else dimension + self.fail = fail + self.delay = delay + self.texts: list[str] = [] + + def embed(self, text: str) -> list[float]: + self.texts.append(text) + if self.delay: + time.sleep(self.delay) + if self.fail: + raise RuntimeError("the provider is down: secret-record-content") + return list(self.vector) + + +class AsyncEmbedder(FixedEmbedder): + async def embed(self, text: str) -> list[float]: # type: ignore[override] + self.texts.append(text) + await asyncio.sleep(0) + return list(self.vector) + + +class Spy(InMemoryExperienceStore): + """An in-memory store that counts the vector statements it is asked to run, and can be made to fail them.""" + + def __init__(self, *args: Any, search_error: bool = False, search_delay: float = 0.0, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self.vector_searches = 0 + self.search_error = search_error + self.search_delay = search_delay + + def search_vectors(self, bounds: AuthorizationBounds, query: VectorQuery) -> Any: + self.vector_searches += 1 + if self.search_delay: + time.sleep(self.search_delay) + if self.search_error: + raise RuntimeError("the index is down") + return super().search_vectors(bounds, query) + + +class TextOnlyStore: + """A store with the text methods only: no vector methods.""" + + def __init__(self, inner: InMemoryExperienceStore) -> None: + self.inner = inner + + def insert_finalized(self, bounds: AuthorizationBounds, finalized: Any, *, actor: str) -> None: + self.inner.insert_finalized(bounds, finalized, actor=actor) + + def search(self, bounds: AuthorizationBounds, parameters: SearchParameters) -> Any: + return self.inner.search(bounds, parameters) + + +class Seen: + def __init__(self) -> None: + self.injections: list[InjectionResult] = [] + self.fallbacks: list[VectorFallback] = [] + self.injection_failures: list[InjectionFailure] = [] + self.capture_failures: list[CaptureFailure] = [] + self.client = ScriptedClient(calls=0) + + +def make_agent(store: Any, seen: Seen, **options: Any) -> Agent: + return use_agent_experience( + Agent(client=seen.client, name="refund-agent"), + store=store, + identity=ExperienceIdentity("python-principal", SCOPE), + on_injected=seen.injections.append, + on_vector_fallback=seen.fallbacks.append, + on_injection_failure=seen.injection_failures.append, + on_capture_failure=seen.capture_failures.append, + **options, + ) + + +def embedded_store(cls: type[InMemoryExperienceStore] = InMemoryExperienceStore, **kwargs: Any) -> Any: + """Three 'rotate-keys' records, embedded as the nearest-first fixture does: 3 is the query's own vector.""" + records = seed_records("vector-search", "nearest-first") + store = cls(records, **kwargs) + vectors = {1: (1.0, 1.0, 0.0, 0.0), 2: (1.0, 0.1, 0.0, 0.0), 3: (1.0, 0.0, 0.0, 0.0)} + for record in records: + number = int(str(record.experience_id)[-1]) + store.write_embedding( + BOUNDS, EmbeddingWrite(record.experience_id, SCOPE, MODEL, 4, "h", record.revision, vectors[number], NOW) + ) + return store + + +def ids(result: InjectionResult) -> list[int]: + return [int(str(experience_id)[-1]) for experience_id in result.experience_ids] + + +def test_only_the_vector_finds_a_reworded_query_and_without_an_embedder_it_does_not() -> None: + with_embedder = Seen() + run(make_agent(embedded_store(), with_embedder, embedder=FixedEmbedder()).run(PARAPHRASE)) + assert with_embedder.fallbacks == [] and with_embedder.injection_failures == [] + assert ids(with_embedder.injections[0]) == [3, 2, 1] + assert with_embedder.injections[0].vector_fallback is None + assert "--- RECORD 1: rotate-keys ---" in with_embedder.client.blocks()[0] + + without = Seen() + run(make_agent(embedded_store(), without).run(PARAPHRASE)) + assert without.injections == [] and without.client.blocks() == [] + assert [fallback.reason for fallback in without.fallbacks] == ["NotConfigured"] + + +def test_a_text_match_and_a_vector_match_are_one_entry() -> None: + seen = Seen() + run(make_agent(embedded_store(), seen, embedder=FixedEmbedder()).run("Rotate the signing keys for the service")) + assert sorted(ids(seen.injections[0])) == [1, 2, 3] + + +def test_an_async_embedder_works() -> None: + seen = Seen() + run(make_agent(embedded_store(), seen, embedder=AsyncEmbedder()).run(PARAPHRASE)) + assert ids(seen.injections[0]) == [3, 2, 1] + + +def test_a_store_without_vector_methods_answers_from_text_and_reports_not_configured() -> None: + inner = embedded_store() + store = TextOnlyStore(inner) + assert not isinstance(store, EmbeddingIndex) + embedder = FixedEmbedder() + seen = Seen() + run(make_agent(store, seen, embedder=embedder).run("Rotate the signing keys for the service")) + assert [fallback.reason for fallback in seen.fallbacks] == ["NotConfigured"] + assert sorted(ids(seen.injections[0])) == [1, 2, 3] + assert seen.injections[0].vector_fallback == seen.fallbacks[0] + assert embedder.texts == [] + + +@pytest.mark.parametrize( + "embedder", + [ + FixedEmbedder(fail=True), + FixedEmbedder(delay=0.4), + FixedEmbedder(model_id=" "), + FixedEmbedder(vector=[]), + FixedEmbedder(dimension=3), + FixedEmbedder(vector=[1.0, float("nan"), 0.0, 0.0]), + FixedEmbedder(vector=[1.0, float("inf"), 0.0, 0.0]), + FixedEmbedder(model_id="m" * 300), + ], + ids=["raises", "times-out", "blank-model", "empty", "wrong-length", "nan", "infinite", "long-model"], +) +def test_a_bad_embedder_answers_from_text_and_no_statement_runs(embedder: FixedEmbedder) -> None: + store = embedded_store(Spy) + seen = Seen() + run(make_agent(store, seen, embedder=embedder, vector_timeout=0.1).run("Rotate the signing keys for the service")) + assert [fallback.reason for fallback in seen.fallbacks] == ["ProviderUnavailable"] + assert "secret-record-content" not in seen.fallbacks[0].detail + assert store.vector_searches == 0 + assert seen.injection_failures == [] + assert sorted(ids(seen.injections[0])) == [1, 2, 3] + + +def test_a_mismatch_answers_from_text_and_says_which() -> None: + seen = Seen() + run(make_agent(embedded_store(), seen, embedder=FixedEmbedder(model_id="another-model")).run(PARAPHRASE)) + assert [fallback.reason for fallback in seen.fallbacks] == ["ModelMismatch"] + assert seen.injections == [] + + seen = Seen() + run(make_agent(embedded_store(), seen, embedder=FixedEmbedder(vector=[1.0, 0.0, 0.0])).run(PARAPHRASE)) + assert [fallback.reason for fallback in seen.fallbacks] == ["DimensionMismatch"] + + +@pytest.mark.parametrize("kind", ["raises", "times-out"]) +def test_a_failing_search_answers_from_text_and_reports_it(kind: str) -> None: + store = embedded_store(Spy, search_error=kind == "raises", search_delay=0.5 if kind == "times-out" else 0.0) + seen = Seen() + run( + make_agent(store, seen, embedder=FixedEmbedder(), vector_timeout=0.1).run( + "Rotate the signing keys for the service" + ) + ) + assert [fallback.reason for fallback in seen.fallbacks] == ["VectorSearchFailed"] + assert seen.injection_failures == [] + assert sorted(ids(seen.injections[0])) == [1, 2, 3] + + +def test_a_failing_text_search_fails_the_injection_whatever_the_vector_found() -> None: + class Broken(InMemoryExperienceStore): + def search(self, bounds: AuthorizationBounds, parameters: SearchParameters) -> Any: + raise RuntimeError("the text index is down") + + store = embedded_store(Broken) + seen = Seen() + run(make_agent(store, seen, embedder=FixedEmbedder()).run(PARAPHRASE)) + assert [failure.reason for failure in seen.injection_failures] == [ + "Injection raised RuntimeError; no block was added." + ] + assert seen.injections == [] and seen.client.blocks() == [] + + +def test_the_floor_drops_weak_vector_candidates_and_never_a_text_one() -> None: + seen = Seen() + run(make_agent(embedded_store(), seen, embedder=FixedEmbedder(), min_vector_relevance=0.9).run(PARAPHRASE)) + assert ids(seen.injections[0]) == [3, 2] # record 1's relevance is about 0.85 + + # A text match is kept although its vector is below the floor. + seen = Seen() + run( + make_agent( + embedded_store(), seen, embedder=FixedEmbedder(vector=[0.0, 0.0, 0.0, 1.0]), min_vector_relevance=1.0 + ).run("Rotate the signing keys for the service") + ) + assert sorted(ids(seen.injections[0])) == [1, 2, 3] + + +def test_settings_refuse_a_bad_embedder_floor_or_timeout() -> None: + base: dict[str, Any] = {"store": InMemoryExperienceStore(), "identity": ExperienceIdentity("p", SCOPE)} + with pytest.raises(TypeError): + ExperienceSettings(**base, embedder=object()) # type: ignore[arg-type] + with pytest.raises(ValueError): + ExperienceSettings(**base, min_vector_relevance=1.5) + with pytest.raises(ValueError): + ExperienceSettings(**base, vector_timeout=0) + + +# ---- The background write after a record is stored --------------------------------------------------------------- + + +def stored_agent(store: Any, seen: Seen, embedder: Any, **options: Any) -> Agent: + @tool + def run_tests(suite: str) -> str: + """Runs a test suite.""" + return "ok" + + seen.client = ScriptedClient(calls=1, arguments={"suite": "unit"}) + return use_agent_experience( + Agent(client=seen.client, name="refund-agent", tools=[run_tests]), + store=store, + identity=ExperienceIdentity("python-principal", SCOPE), + verify=tool_succeeded("run_tests", artifact_revision="build-1"), + on_capture_failure=seen.capture_failures.append, + embedder=embedder, + **options, + ) + + +def settings_of(agent: Agent) -> ExperienceSettings: + provider = next(p for p in agent.context_providers if isinstance(p, ExperienceContextProvider)) + return provider.settings + + +def test_a_stored_record_is_embedded_in_the_background_from_its_summary() -> None: + store = InMemoryExperienceStore() + seen = Seen() + embedder = FixedEmbedder() + agent = stored_agent(store, seen, embedder) + + async def go() -> None: + await agent.run("Refund ticket 43 is stuck in pending") + await settings_of(agent).wait_for_index_writes() + + run(go()) + (record,) = store.records() + stored = store.embedding(record.experience_id) + assert stored is not None + summary = record_summary(record) + assert summary == retrieval_summary( + record.task_id, record.task_summary, record.reflection.lesson if record.reflection else None + ) + assert embedder.texts[-1] == summary # after the query's own embedding + assert stored == (MODEL, 4, content_hash(MODEL, summary), record.revision, (1.0, 0.0, 0.0, 0.0)) + assert seen.capture_failures == [] + + +@pytest.mark.parametrize( + "embedder", [FixedEmbedder(fail=True), FixedEmbedder(dimension=9)], ids=["raises", "wrong-length"] +) +def test_a_failed_embedding_write_is_reported_as_index_and_the_record_stays_stored(embedder: FixedEmbedder) -> None: + store = InMemoryExperienceStore() + seen = Seen() + agent = stored_agent(store, seen, embedder) + + async def go() -> None: + await agent.run("Refund ticket 43 is stuck in pending") + await settings_of(agent).wait_for_index_writes() + + run(go()) + assert len(store.records()) == 1 + assert [failure.stage for failure in seen.capture_failures] == ["Index"] + assert "secret-record-content" not in seen.capture_failures[0].reason + assert store.embedding(store.records()[0].experience_id) is None + + +def test_no_embedding_is_written_without_an_embedder_or_a_store_that_takes_it() -> None: + store = InMemoryExperienceStore() + seen = Seen() + run(stored_agent(store, seen, None).run("Refund ticket 43 is stuck in pending")) + assert store.embedding(store.records()[0].experience_id) is None + + inner = InMemoryExperienceStore() + seen = Seen() + embedder = FixedEmbedder() + run(stored_agent(TextOnlyStore(inner), seen, embedder).run("Refund ticket 43 is stuck in pending")) + assert len(inner.records()) == 1 and embedder.texts == [] and seen.capture_failures == [] + + +def test_a_lesson_stored_and_embedded_is_found_by_a_reworded_query_later() -> None: + store = InMemoryExperienceStore() + seen = Seen() + agent = stored_agent(store, seen, FixedEmbedder(), min_vector_relevance=0.5) + + async def go() -> None: + await agent.run("Refund ticket 43 is stuck in pending") + await settings_of(agent).wait_for_index_writes() + + run(go()) + later = Seen() + run(make_agent(store, later, embedder=FixedEmbedder()).run("Zebra giraffe antelope")) + assert len(later.injections) == 1 and later.fallbacks == [] + + +# ---- Against PostgreSQL, with a record embedded by the contract's own SQL ----------------- + + +@pytest.mark.postgres +def test_a_record_embedded_by_the_contracts_sql_is_found_by_a_python_agent(pg_server: Server) -> None: + case = FIXTURES / "vector-search" / "nearest-first" + steps = read_json(case / "steps.json")["steps"] + database = create_database(pg_server, "agent_vector_nearest_first") + conninfo = application_conninfo(pg_server, database) + store = PostgresExperienceStore(conninfo) + for record in seed_records("vector-search", "nearest-first"): + store.insert_record(AuthorizationBounds.of(record.scope), record) + + # The embeddings are written by the fixture's own statement, not by the Python store. + write_sql = to_driver_placeholders(contract_statements("write-experience-embedding.sql")["write"]) + declare = to_driver_placeholders(contract_sql(DECLARE_BOUNDS_SQL)) + for step in (s["write"] for s in steps if "write" in s): + scope = scope_of(step["scope"]) + with psycopg.connect(conninfo) as connection, connection.transaction(): + connection.execute(declare, AuthorizationBounds.of(scope).parameters()) + written = connection.execute( + write_sql, + { + "experience_id": uuid.UUID(step["experienceId"]), + **dict( + zip( + ("tenant_id", "application_id", "project_id", "team_id", "agent_id", "user_id"), + scope.fields(), + strict=True, + ) + ), + "model_id": step["modelId"], + "dimension": step["dimension"], + "content_hash": step["contentHash"], + "source_revision": step["sourceRevision"], + "embedding": "[" + ",".join(str(c) for c in step["vector"]) + "]", + "now": Instant.parse(step["now"]).to_datetime(), + }, + ) + assert written.rowcount == 1 + + query = next(s["search"] for s in steps if "search" in s) + seen = Seen() + run(make_agent(store, seen, embedder=FixedEmbedder(query["vector"], model_id=query["modelId"])).run(PARAPHRASE)) + assert seen.fallbacks == [] and seen.injection_failures == [] + assert ids(seen.injections[0]) == [3, 2, 1] + + # Without the embedder the same query finds nothing, and says why. + plain = Seen() + run(make_agent(store, plain).run(PARAPHRASE)) + assert plain.injections == [] and [f.reason for f in plain.fallbacks] == ["NotConfigured"] + + +@pytest.mark.postgres +def test_a_python_agent_stores_a_record_and_its_embedding_in_postgres(pg_server: Server) -> None: + database = create_database(pg_server, "agent_vector_roundtrip") + store = PostgresExperienceStore(application_conninfo(pg_server, database)) + seen = Seen() + agent = stored_agent(store, seen, FixedEmbedder()) + + async def go() -> None: + await agent.run("Refund ticket 43 is stuck in pending") + await settings_of(agent).wait_for_index_writes() + + run(go()) + assert seen.capture_failures == [] + owner = pg_server.conninfo(database, OWNER, PASSWORD) + rows = stored_embeddings(owner, ["model_id", "dimension", "source_revision", "embedding", "content_hash"]) + assert len(rows) == 1 + assert rows[0]["model_id"] == MODEL and rows[0]["embedding"] == "[1,0,0,0]" + + later = Seen() + run(make_agent(store, later, embedder=FixedEmbedder()).run("Zebra giraffe antelope")) + assert len(later.injections) == 1 and later.fallbacks == [] + + +def test_a_lone_surrogate_hashes_as_the_replacement_character() -> None: + from agentexperience.embedding import content_hash + + assert content_hash("m", "a\ud800b") == content_hash("m", "a�b") + + +def test_hung_vector_calls_fail_fast_and_never_take_the_default_executor() -> None: + import asyncio + import threading + + from agentexperience import embedding as embedding_module + from agentexperience.embedding import MAX_VECTOR_THREADS, run_on_vector_thread + + release = threading.Event() + + async def scenario() -> None: + # Threads that earlier tests abandoned (a timed-out embedder) may still hold slots: wait for them to end. + for _ in range(200): + free = [embedding_module._vector_threads.acquire(blocking=False) for _ in range(MAX_VECTOR_THREADS)] + for taken in free: + if taken: + embedding_module._vector_threads.release() + if all(free): + break + await asyncio.sleep(0.05) + hung = [asyncio.ensure_future(run_on_vector_thread(release.wait)) for _ in range(MAX_VECTOR_THREADS)] + await asyncio.sleep(0.05) + with pytest.raises(RuntimeError): + await run_on_vector_thread(lambda: 1) + assert await asyncio.to_thread(lambda: 2) == 2 # the default executor is untouched + release.set() + await asyncio.gather(*hung) + assert await run_on_vector_thread(lambda: 3) == 3 + + asyncio.run(scenario())