From 33eee603dcc29586c111a51659e9a50db22b9d21 Mon Sep 17 00:00:00 2001 From: fabbrik <22822543+fabbrik@users.noreply.github.com> Date: Sat, 17 Oct 2026 11:49:48 -0500 Subject: [PATCH] feat(go): semantic retrieval, contract 1.6 (story 23.3) Co-Authored-By: Claude Sonnet 5.5 Claude-Session: https://claude.ai/code/session_01PUvmgPgFMAAc5aCaVqTdMF --- .github/workflows/ci.yml | 3 +- CHANGELOG.md | 23 +- README.md | 2 +- go/README.md | 71 +- go/contract_test.go | 143 +++- go/doc.go | 6 +- go/embedding.go | 516 +++++++++++++ go/embedding_test.go | 212 ++++++ go/inmemory/store.go | 149 +++- go/inmemory/vectors_test.go | 247 ++++++ go/internal/contracttest/contracttest.go | 95 ++- go/maf/capture.go | 2 + go/maf/inject.go | 28 +- go/maf/maf.go | 56 +- go/maf/record.go | 21 +- go/maf/vectors.go | 294 ++++++++ go/maf/vectors_postgres_test.go | 173 +++++ go/maf/vectors_test.go | 705 ++++++++++++++++++ go/postgres/pgtest/pgtest.go | 43 +- go/postgres/pgtest/vectors_postgres_test.go | 180 +++++ go/postgres/sql/search-experience-vector.sql | 40 + .../sql/write-experience-embedding.sql | 31 + go/postgres/store.go | 240 +++++- go/postgres/store_test.go | 73 +- go/store.go | 14 + 25 files changed, 3296 insertions(+), 71 deletions(-) create mode 100644 go/embedding.go create mode 100644 go/embedding_test.go create mode 100644 go/inmemory/vectors_test.go create mode 100644 go/maf/vectors.go create mode 100644 go/maf/vectors_postgres_test.go create mode 100644 go/maf/vectors_test.go create mode 100644 go/postgres/pgtest/vectors_postgres_test.go create mode 100644 go/postgres/sql/search-experience-vector.sql create mode 100644 go/postgres/sql/write-experience-embedding.sql diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index ea6b7308..2bb2d64a 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -505,7 +505,8 @@ jobs: # stable. Then the PostgreSQL leg (the postgres build tag, race detector # on), which writes every finalize fixture through the Go store and reads it back row for row, runs every search # fixture, and uses the store from many goroutines at once, on the oldest and newest majors the contract tests cover. - # The leg sets AGENTEXPERIENCE_REQUIRE_POSTGRES, so it fails rather than skips without Docker, and + # The leg uses the pgvector/pgvector image (contract 1.6's vector fixtures need the extension). The leg sets + # AGENTEXPERIENCE_REQUIRE_POSTGRES, so it fails rather than skips without Docker, and # AGENTEXPERIENCE_POSTGRES_MAJOR, which the tests assert the server reports. testcontainers-go prepends the # workflow's TESTCONTAINERS_HUB_IMAGE_NAME_PREFIX as a path, so it needs no slash. go: diff --git a/CHANGELOG.md b/CHANGELOG.md index 4b4626a7..6abc4da3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -22,7 +22,7 @@ AgentExperience.NET is a **preview**. It is not production ready, and public API contract test keeps the .NET privileges call equal to the 1.0, 1.4 and 1.6 files together. No library behaviour or public API changed. - **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. + families by name until stories 23.2 and 23.3 implement them (23.2, Python, and 23.3, Go, below). ### Python semantic retrieval (story 23.2) @@ -39,7 +39,26 @@ AgentExperience.NET is a **preview**. It is not production ready, and public API 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. + PostgreSQL leg. + +### Go semantic retrieval (story 23.3) + +- **The Go port passes contract 1.6.** `maf.Options.Embedder` takes your embedding object (an `ae.Embedder`: + `ModelID`, `Dimension` and `Embed(ctx, 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 `Rank`, 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 `OnVectorFallback` and `InjectionResult.VectorFallback`. The vector channel never fails + the call and is bounded by `VectorTimeout` (500 ms by default); the text channel's failure still fails it. + Embedder and vector-store calls are bounded too: at most 16 in flight, abandoned ones included, so a hung provider + cannot starve the text search. `MinVectorRelevance` drops weak vector candidates, never text ones. +- **Stores.** `inmemory.Store` and `postgres.Store` (PostgreSQL with pgvector, running the contract's two new SQL files + as written) both serve the channel through `WriteEmbedding` and `SearchVectors` (the new optional `ae.EmbeddingIndex`). + 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 `OnCaptureFailure` as the new stage `Index` and never fails capture + (`maf.WaitForIndexWrites` awaits the writes still running). +- Without an embedder nothing changes. The Go PostgreSQL tests now need a pgvector image (`pgvector/pgvector`). All + three ports (.NET, Python, Go) now pass contract 1.6. ## 0.1.0-preview.14 diff --git a/README.md b/README.md index 13f2cdb4..f93445f4 100644 --- a/README.md +++ b/README.md @@ -245,7 +245,7 @@ lesson was injected) is published as [a language-neutral contract](contract/READ 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.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 +agent learned is injected into a Python agent, and the other way round. A Go core and PostgreSQL store that pass contract 1.6 (including retrieval by meaning, with an embedder you register), 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, `0.1.0aN`, for `0.1.0-preview.N`), and are previews in the same sense: diff --git a/go/README.md b/go/README.md index be444246..a39c50f9 100644 --- a/go/README.md +++ b/go/README.md @@ -4,8 +4,8 @@ A Go core of AgentExperience that implements the [language-neutral contract](../ learned by a Go program is stored in the format the .NET and Python libraries read, and the other way round, with a Microsoft Agent Framework for Go integration. -**Status: preview, contract 1.5, with a Microsoft Agent Framework integration (the `maf` module).** It passes every -fixture of contract 1.5 (and so every 1.0, 1.1, 1.2, 1.3 and 1.4 fixture). Every release of this repository tags each module at the +**Status: preview, contract 1.6, with a Microsoft Agent Framework integration (the `maf` module).** It passes every +fixture of contract 1.6 (and so every 1.0, 1.1, 1.2, 1.3, 1.4 and 1.5 fixture). Every release of this repository tags each module at the same version as the .NET packages (`v0.1.0-preview.N`), and the Go module proxy serves it from there. Like the .NET packages, it claims no production readiness, and its API may change between previews. It is Apache-2.0, as the [repository](../LICENSE) is. @@ -81,6 +81,47 @@ the identity names its application, the recorded environment is fixed so the out and it prints each run and the block. Its output is checked in as [`quickstart.expected.txt`](maf/examples/quickstart/quickstart.expected.txt), and a test compares them. +### Semantic retrieval: registering an embedder + +Without an embedder retrieval is text only, and a reworded query finds nothing it shares no words with. Implement +`ae.Embedder` over any model (none is a dependency of this repository) and set `maf.Options.Embedder`, and retrieval +becomes hybrid: + +```go +type myEmbedder struct{ client *myprovider.Client } + +func (myEmbedder) ModelID() string { return "my-embedding-model" } // stored with every vector; another model's vectors are never compared +func (myEmbedder) Dimension() int { return 1536 } // the vector's length, known without calling the provider + +func (e myEmbedder) Embed(ctx context.Context, text string) ([]float32, error) { // exactly Dimension() finite numbers + return e.client.Embed(ctx, text) +} + +err := maf.UseAgentExperience(&config, maf.Options{Store: store, Identity: identity, Verify: verify, + Embedder: myEmbedder{client}}) +``` + +- **In-memory and PostgreSQL stores.** `inmemory.Store` compares 32-bit floats by cosine distance. `postgres.Store` + 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); it runs the + contract's two new SQL files as written (`WriteEmbedding`, `SearchVectors`). Both are an `ae.EmbeddingIndex`. +- **Fail-open.** The vector channel never fails or delays past `VectorTimeout` (default 500 ms, covering the embedding + and the search) the text answer, and at most `maf.MaxVectorCalls` (16) embedder and vector-store calls are in flight at + once, abandoned ones included, so a hung provider cannot starve the text search; past that a call fails at once. When + the channel does not contribute the injection answers from text and reports the reason through `OnVectorFallback` + and `InjectionResult.VectorFallback`: `NotConfigured` (no embedder, or a store that is not an `ae.EmbeddingIndex`), + `ProviderUnavailable` (the embedder failed, panicked, 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. `MinVectorRelevance` (0 to 1) drops vector candidates below it; text candidates are never + dropped. +- **Writing embeddings.** After a record is stored, a goroutine 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 `OnCaptureFailure` with the stage `maf.StageIndex` (`"Index"`). + `maf.WaitForIndexWrites(ctx, &config)` (and `ContextProviderOf(&config).Close(ctx)`) 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. + ### What the one call wires - **Capture**: agent middleware and function invocation middleware, each put first so it wraps everything else on @@ -113,6 +154,9 @@ and it prints each run and the block. Its output is checked in as overruns it is abandoned. `OnInjectionRecorded` reports how it ended (`Recorded`, `NotTracked`, `NotRecorded` or `OutcomeUnknown`, as .NET's `record_injection` operation), and `maf.ContextProviderOf(&config).Close(ctx)` awaits the records still running, for a host shutting down. +- **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. - **The task text, and where history sits.** Capture records the task text of the invocation's own input. In the Go SDK a session's history provider always runs before the context providers, so injection sees the replayed history, then the new input, and derives the task text from both, as .NET's `ChatClientAgent`-level provider does: a request @@ -177,9 +221,13 @@ Python agents on the same database. - **The error class.** `ErrorClass` reduces a captured error to its short class. - **Search and ranking.** `SearchParameters` and `AuthorizationBounds` are the contract statements' parameters; `Rank` applies the contract's safety checks, truncation and ranking to the rows a search returns. +- **Semantic retrieval** (contract 1.6). `Embedder` is your embedding object; `RetrievalSummary`, `RecordSummary` and + `ContentHash` are what is embedded and how it is hashed; `EmbeddingWrite` and `VectorQuery` validate before any + statement; `MergeChannels` and `RankHybrid` turn the text and vector candidates into one ranked list (`Rank` itself is + unchanged, and ranks every merged candidate, up to 100). - **PostgreSQL** (the `postgres` package). `Store` writes a finalized record and its initial event as the application - role, reads records back, runs the contract's text search, and records injections (`RecordInjection`, contract - 1.4), every transaction first declaring its bounds. The contract's three statement files are embedded and run as + role, reads records back, runs the contract's text search, records injections (`RecordInjection`, contract + 1.4) and writes and searches embeddings (`WriteEmbedding` and `SearchVectors`, contract 1.6), every transaction first declaring its bounds. The contract's three statement files are embedded and run as written (a test keeps the copies in `postgres/sql/` identical to `contract/sql/`); only their `@name` placeholders are numbered for the driver. The store needs a schema the .NET migrator applied, with row-level security on; it never creates or alters a schema object. The application role's privileges are the .NET privileges call's, or the @@ -202,7 +250,7 @@ Python agents on the same database. an invocation runs under and `DevelopmentIdentity` (see [the quick start](#quick-start-with-microsoft-agent-framework)), in the core, so they need no Agent Framework type. -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. It writes no logs at all, so no record content is ever logged. [`example_test.go`](example_test.go) captures a run, finalizes it and renders the block; `go test` runs it and checks @@ -271,7 +319,7 @@ In the core, `go test ./...` runs the unit tests and the conformance runner ([`c every fixture family from [`../contract/fixtures`](../contract/fixtures) and compares this package's output exactly: blocks and task text byte for byte, JSON by value, with a number written as an integer never equal to one written with a fraction or an exponent. It checks every contract file against `contract/lock.json`, and fails with a clear -message if the lock names a contract version this package does not support (1.0 to 1.5 are). The verifier and +message if the lock names a contract version this package does not support (1.0 to 1.6 are). The verifier and classifier tests port the cases of .NET's and Python's. In `postgres/`, `go test ./...` validates every record of the fixtures, and every record the core writes for a finalize fixture, against the contract's JSON Schema (there, because it needs a validator and the core module requires none). The @@ -304,15 +352,18 @@ cancellation, parallel calls), runs the quick start and compares its output **PostgreSQL** (the `postgres` build tag, in `postgres/pgtest/` and `maf/`, Go 1.26): in `postgres/pgtest/`, each finalize fixture is written through the Go store and read back row for row, every search fixture is run and ranked, -every record-injection fixture is run, and sealed rows, tombstones, the exact scope, row-level security, `recorded_at` and concurrent use on a pool are +every record-injection fixture is run, every embedding-write and vector-search fixture is run, and sealed rows, tombstones, the exact scope, row-level security, `recorded_at` and concurrent use on a pool are checked; in `maf/`, the cross-language tests run. Both use the server helper the `pgtest` package is (its roles are made under an advisory lock, so test binaries may share one server). Each test database (dropped when its test ends) is migrated with the .NET migrator's own scripts from [`src/AgentExperience.Storage.Postgres/Migrations`](../src/AgentExperience.Storage.Postgres/Migrations), in order, and given the application role's privileges from the contract's [`sql/application-role-privileges.sql`](../contract/sql/application-role-privileges.sql) then -[`sql/application-role-privileges-1.4.sql`](../contract/sql/application-role-privileges-1.4.sql). The server is a -testcontainers-go `postgres` container (`AGENTEXPERIENCE_POSTGRES_MAJOR`, 16 by default), or the server `PGHOST`, -`PGPORT`, `PGUSER` and `PGPASSWORD` name (a superuser). With neither the tests skip, unless +[`sql/application-role-privileges-1.4.sql`](../contract/sql/application-role-privileges-1.4.sql) and +[`sql/application-role-privileges-1.6.sql`](../contract/sql/application-role-privileges-1.6.sql); from contract 1.6 the +vectors package's migrations (`src/AgentExperience.Storage.Postgres.Vectors/Migrations`) follow, after a superuser ran +`CREATE EXTENSION vector`. The server is a testcontainers-go `pgvector/pgvector` container +(`AGENTEXPERIENCE_POSTGRES_MAJOR`, 16 by default), or the server `PGHOST`, `PGPORT`, `PGUSER` and `PGPASSWORD` name (a +superuser, with the pgvector extension available). With neither the tests skip, unless `AGENTEXPERIENCE_REQUIRE_POSTGRES=1`, when they fail. On a Mac with Rancher Desktop or Colima, testcontainers-go may need: diff --git a/go/contract_test.go b/go/contract_test.go index c09c5c1e..1c770dd9 100644 --- a/go/contract_test.go +++ b/go/contract_test.go @@ -13,10 +13,11 @@ import ( "go/parser" "go/token" "io/fs" + "math" "os" "path/filepath" - "slices" "sort" + "strconv" "strings" "testing" "unicode/utf16" @@ -43,10 +44,10 @@ func TestTheContractVersionIsOneThisPackageSupports(t *testing.T) { if err != nil { t.Fatal(err) } - if version != ae.ContractVersion && !slices.Contains(ct.PendingContractVersions, version) { + if version != ae.ContractVersion { t.Fatalf("the lock is contract %s, this package implements %s", version, ae.ContractVersion) } - if strings.Join(ae.SupportedContractVersions, ", ") != "1.0, 1.1, 1.2, 1.3, 1.4, 1.5" { + if strings.Join(ae.SupportedContractVersions, ", ") != "1.0, 1.1, 1.2, 1.3, 1.4, 1.5, 1.6" { t.Fatalf("supported versions: %v", ae.SupportedContractVersions) } } @@ -54,7 +55,7 @@ func TestTheContractVersionIsOneThisPackageSupports(t *testing.T) { func TestAnUnsupportedContractVersionFailsWithAClearMessage(t *testing.T) { for _, version := range []any{"2.0", "1.7", nil} { _, err := ct.CheckContractVersion(ae.NewObject(ae.Member{Key: "contractVersion", Value: version})) - if !errors.Is(err, ct.ErrUnsupportedContractVersion) || !strings.Contains(err.Error(), "supports only 1.0, 1.1, 1.2, 1.3, 1.4, 1.5") { + if !errors.Is(err, ct.ErrUnsupportedContractVersion) || !strings.Contains(err.Error(), "supports only 1.0, 1.1, 1.2, 1.3, 1.4, 1.5, 1.6") { t.Fatalf("version %v: %v", version, err) } } @@ -124,19 +125,137 @@ func TestEveryFixtureFamilyHasCases(t *testing.T) { } } -// TestPendingVectorFamilies skips contract 1.6's vector families by name: this package does not implement them yet -// (stories 23.2 and 23.3 do, and remove ct.PendingVectorFamilies). Each must exist, so a renamed family fails here. -func TestPendingVectorFamilies(t *testing.T) { - for _, family := range ct.PendingVectorFamilies { - t.Run(family, func(t *testing.T) { - if _, err := os.Stat(fixture(family)); err != nil { - t.Fatal(err) +// ---- Contract 1.6: the vector families that need no database ---------------------------------------------------- + +func TestEmbeddingSummary(t *testing.T) { + cases := ct.MustReadJSON(fixture("embedding-summary", "cases.json")).([]any) + if len(cases) == 0 { + t.Fatal("no cases") + } + text := func(o *ae.Object, key string) *string { + value, _ := o.Get(key) + s, ok := value.(string) + if !ok { + return nil + } + return &s + } + for i, item := range cases { + c := item.(*ae.Object) + description := *text(c, "description") + t.Run(fmt.Sprintf("%d", i), func(t *testing.T) { + summary := ae.RetrievalSummary(text(c, "taskId"), text(c, "taskSummary"), text(c, "lesson")) + if summary != *text(c, "summary") { + t.Fatalf("%s: summary %q, expected %q", description, summary, *text(c, "summary")) + } + if hash := ae.ContentHash(*text(c, "modelId"), summary); hash != *text(c, "contentHash") { + t.Fatalf("%s: hash %s, expected %s", description, hash, *text(c, "contentHash")) } - t.Skipf("%s: a contract 1.6 vector family, pending in this package", family) }) } } +func relevanceOf(t *testing.T, value any) float64 { + t.Helper() + if s, ok := value.(string); ok { + if s != "NaN" { + t.Fatalf("relevance %q", s) + } + return math.NaN() + } + f, err := strconv.ParseFloat(string(value.(json.Number)), 64) + if err != nil { + t.Fatal(err) + } + return f +} + +func TestHybridMerge(t *testing.T) { + cases := ct.MustReadJSON(fixture("hybrid-merge", "cases.json")).([]any) + if len(cases) == 0 { + t.Fatal("no cases") + } + type candidate struct { + id ae.UUID + relevance float64 + } + entries := func(c *ae.Object, key string) []candidate { + value, _ := c.Get(key) + var result []candidate + for _, item := range value.([]any) { + o := item.(*ae.Object) + id, _ := o.Get("experienceId") + relevance, _ := o.Get("relevance") + result = append(result, candidate{ae.MustParseUUID(id.(string)), relevanceOf(t, relevance)}) + } + return result + } + for i, item := range cases { + c := item.(*ae.Object) + t.Run(fmt.Sprintf("%d", i), func(t *testing.T) { + var floor *float64 + if value, _ := c.Get("minimumVectorRelevance"); value != nil { + f := relevanceOf(t, value) + floor = &f + } + merged := ae.MergeChannels(entries(c, "text"), entries(c, "vector"), + func(x candidate) ae.UUID { return x.id }, func(x candidate) float64 { return x.relevance }, floor) + expectedValue, _ := c.Get("merged") + expected := expectedValue.([]any) + description, _ := c.Get("description") + if len(merged) != len(expected) { + t.Fatalf("%v: %d merged, expected %d", description, len(merged), len(expected)) + } + for j, entry := range merged { + e := expected[j].(*ae.Object) + id, _ := e.Get("experienceId") + from, _ := e.Get("from") + relevance, _ := e.Get("relevance") + want := relevanceOf(t, relevance) + same := entry.Relevance == want || (math.IsNaN(want) && math.IsNaN(entry.Relevance)) + if entry.Item.id.String() != id.(string) || string(entry.Source) != from.(string) || !same { + t.Fatalf("%v: entry %d is %v %v %v, expected %v %v %v", description, j, entry.Item.id, + entry.Source, entry.Relevance, id, from, want) + } + } + }) + } +} + +func TestALoneSurrogateIsHashedAsTheReplacementCharacter(t *testing.T) { + // A lone surrogate in three-byte form (what text carrying one looks like in a Go string), and any invalid byte. + for _, bad := range []string{"a\xed\xa0\x80b", "a\xffb"} { + if ae.ContentHash("m", bad) != ae.ContentHash("m", "a\uFFFDb") { + t.Fatalf("%q is not hashed as U+FFFD", bad) + } + } + if ae.ContentHash("m", "a\xed\xa0\x80\xed\xb0\x80b") != ae.ContentHash("m", "a\uFFFD\uFFFDb") { + t.Fatal("two lone surrogates are two U+FFFD") + } +} + +func TestTheSummaryCapBacksOffASplitPairAndATrailingHighSurrogate(t *testing.T) { + lesson := func(s string) string { return ae.RetrievalSummary(nil, nil, &s) } + // 8,191 units and then an astral character (two units): the pair does not fit, so it is dropped whole. + astral := strings.Repeat("a", 8191) + "\U0001F600" + if got := lesson(astral); got != strings.Repeat("a", 8191) { + t.Fatalf("%d units kept", ae.UTF16Len(got)) + } + // 8,191 units and then a lone high surrogate and one more unit: the cut ends on the surrogate, which goes too. + high := strings.Repeat("a", 8191) + "\xed\xa0\x80" + "b" + if got := lesson(high); got != strings.Repeat("a", 8191) { + t.Fatalf("%d units kept", ae.UTF16Len(got)) + } + // A cut that leaves trailing spaces drops them. + spaced := strings.Repeat("a", 8190) + " b c" + if got := lesson(spaced); got != strings.Repeat("a", 8190)+" b" { + t.Fatalf("%q", got[len(got)-4:]) + } + if got := lesson(strings.Repeat("a", 8192)); ae.UTF16Len(got) != 8192 { + t.Fatal("a summary of exactly 8,192 units is kept whole") + } +} + // ---- finalize -------------------------------------------------------------------------------------------------- func TestFinalize(t *testing.T) { diff --git a/go/doc.go b/go/doc.go index 9b78ed77..4749ece8 100644 --- a/go/doc.go +++ b/go/doc.go @@ -1,5 +1,5 @@ // Package agentexperience is a Go core of AgentExperience that implements the AgentExperience.NET language-neutral -// contract 1.5 (the repository's contract directory). +// contract 1.6 (the repository's contract directory). // // 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 computes the text-search parameters and the @@ -18,7 +18,7 @@ package agentexperience // ContractVersion is the contract version this package implements. -const ContractVersion = "1.5" +const ContractVersion = "1.6" // SupportedContractVersions are the contract versions this package passes: every fixture of each. -var SupportedContractVersions = []string{"1.0", "1.1", "1.2", "1.3", "1.4", "1.5"} +var SupportedContractVersions = []string{"1.0", "1.1", "1.2", "1.3", "1.4", "1.5", "1.6"} diff --git a/go/embedding.go b/go/embedding.go new file mode 100644 index 00000000..1568482d --- /dev/null +++ b/go/embedding.go @@ -0,0 +1,516 @@ +package agentexperience + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "errors" + "fmt" + "math" + "sort" + "strconv" + "strings" + "unicode/utf8" +) + +// 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 module: an Embedder is the +// caller's own object. + +// Embedder turns text into a vector: the caller's object, registered with the agent integration's Options.Embedder. +// ModelID and Dimension are known without calling the provider; the vector Embed returns has exactly Dimension +// finite components. Embeddings of different models or dimensions are never compared. Embed should honour ctx: the +// integration cancels it when the vector channel's timeout passes, but never waits for a call that does not return. +type Embedder interface { + ModelID() string + Dimension() int + Embed(ctx context.Context, text string) ([]float32, error) +} + +// Limits of contract 1.6. +const ( + // MaxSummaryLength is the longest summary that is embedded, in UTF-16 code units. + MaxSummaryLength = 8192 + // MaxModelIDLength is the longest model ID, in UTF-16 code units. + MaxModelIDLength = 256 + // MaxDimension is the most components a vector has (pgvector's vector type accepts up to 16,000). + MaxDimension = 16_000 +) + +// VectorFallbackReason is why an injection answered from text alone (contract 1.6). +type VectorFallbackReason string + +// The fallback reasons, named exactly as the contract and .NET name them. +const ( + FallbackNotConfigured VectorFallbackReason = "NotConfigured" + FallbackProviderUnavailable VectorFallbackReason = "ProviderUnavailable" + FallbackModelMismatch VectorFallbackReason = "ModelMismatch" + FallbackDimensionMismatch VectorFallbackReason = "DimensionMismatch" + FallbackVectorSearchFailed VectorFallbackReason = "VectorSearchFailed" +) + +// VectorFallback says the vector channel did not contribute: the contract's reason, a sentence that never quotes +// record or vector content, and the error, when there was one. +type VectorFallback struct { + Reason VectorFallbackReason + Detail string + Err error +} + +// ---- What is embedded --------------------------------------------------------------------------------------------- + +// decodeUnits is s as UTF-16 semantics see it: a lone surrogate (written in its three-byte form, as .NET text with +// one is) stays one unit in the surrogate range, and any other invalid byte is U+FFFD. +func decodeUnits(s string) []rune { + out := make([]rune, 0, len(s)) + for i := 0; i < len(s); { + r, size := utf8.DecodeRuneInString(s[i:]) + if r == utf8.RuneError && size <= 1 { + if i+2 < len(s) && s[i] == 0xED && s[i+1] >= 0xA0 && s[i+1] <= 0xBF && s[i+2]&0xC0 == 0x80 { + out = append(out, rune(0xD000|int(s[i+1]&0x3F)<<6|int(s[i+2]&0x3F))) + i += 3 + continue + } + out = append(out, utf8.RuneError) + i++ + continue + } + out = append(out, r) + i += size + } + return out +} + +// validUTF8 is s as .NET's UTF-8 encoder writes it: a lone surrogate, and any other invalid byte, is U+FFFD. +func validUTF8(s string) string { return string(decodeUnits(s)) } + +func unitSize(r rune) int { + if r > 0xFFFF { + return 2 + } + return 1 +} + +func isSeparator(r rune) bool { + switch category(r) { + case "Cc", "Zs", "Zl", "Zp": + return true + } + return false +} + +// RetrievalSummary is the only text of a record that is embedded: its task ID, task summary and lesson (a nil field +// contributes nothing), each with every run of separators (Unicode categories Cc, Zs, Zl and Zp, such as controls, +// spaces and line separators; format characters are kept) as one space, joined by one space, and at most +// MaxSummaryLength UTF-16 code units: cut there, a high surrogate left last dropped too (so a pair is never split), +// then trailing spaces dropped. The result is valid UTF-8 (a lone surrogate is U+FFFD). +func RetrievalSummary(taskID, taskSummary, lesson *string) string { + var parts [][]rune + for _, field := range []*string{taskID, taskSummary, lesson} { + if field == nil { + continue + } + var words []rune + pending := false + for _, r := range decodeUnits(*field) { + if isSeparator(r) { + pending = len(words) > 0 + continue + } + if pending { + words = append(words, ' ') + pending = false + } + words = append(words, r) + } + if len(words) > 0 { + parts = append(parts, words) + } + } + var text []rune + for i, part := range parts { + if i > 0 { + text = append(text, ' ') + } + text = append(text, part...) + } + units := 0 + for _, r := range text { + units += unitSize(r) + } + if units <= MaxSummaryLength { + return string(text) + } + used, keep := 0, 0 + for keep < len(text) && used+unitSize(text[keep]) <= MaxSummaryLength { + used += unitSize(text[keep]) + keep++ + } + kept := text[:keep] + if n := len(kept); n > 0 && kept[n-1] >= 0xD800 && kept[n-1] <= 0xDBFF { + kept = kept[:n-1] + } + for n := len(kept); n > 0 && kept[n-1] == ' '; n = len(kept) { + kept = kept[:n-1] + } + return string(kept) +} + +// RecordSummary is RetrievalSummary of a record. +func RecordSummary(record *ExperienceRecord) string { + var lesson *string + if record.Reflection != nil { + lesson = &record.Reflection.Lesson + } + return RetrievalSummary(&record.TaskID, record.TaskSummary, lesson) +} + +// ContentHash is the lowercase hex SHA-256 of the UTF-8 bytes of the model ID, U+001F and the summary; text that is +// not valid (a lone surrogate) is encoded as U+FFFD, as .NET encodes it. +func ContentHash(modelID, summary string) string { + digest := sha256.Sum256([]byte(validUTF8(modelID + "\u001f" + summary))) + return hex.EncodeToString(digest[:]) +} + +// ---- Validation (before any statement) ----------------------------------------------------------------------------- + +// ValidateModelID refuses a blank model ID, and one longer than MaxModelIDLength UTF-16 code units. +func ValidateModelID(modelID string) error { + if isNullOrWhitespace(modelID) { + return errors.New("the model ID must not be blank") + } + if UTF16Len(modelID) > MaxModelIDLength { + return fmt.Errorf("the model ID is longer than %d UTF-16 code units", MaxModelIDLength) + } + return nil +} + +// ValidateVector refuses a vector that does not have 1 to MaxDimension finite components, and, when dimension is +// positive, exactly that many. +func ValidateVector(vector []float32, dimension int) error { + if len(vector) < 1 || len(vector) > MaxDimension { + return fmt.Errorf("a vector has 1 to %d components", MaxDimension) + } + if dimension > 0 && len(vector) != dimension { + return errors.New("the vector's length is not its dimension") + } + for _, component := range vector { + if math.IsNaN(float64(component)) || math.IsInf(float64(component), 0) { + return errors.New("a vector component must be finite") + } + } + return nil +} + +// VectorLiteral is a pgvector literal, [x,y,...], each component its shortest round-trip text. +func VectorLiteral(vector []float32) string { + var out strings.Builder + out.WriteByte('[') + for i, component := range vector { + if i > 0 { + out.WriteByte(',') + } + out.WriteString(strconv.FormatFloat(float64(component), 'g', -1, 32)) + } + out.WriteByte(']') + return out.String() +} + +// EmbeddingWrite is 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. +type EmbeddingWrite struct { + ExperienceID UUID + Scope Scope + ModelID string + Dimension int + ContentHash string + SourceRevision int64 + Vector []float32 + // Now is the writer's clock, which becomes created_at and updated_at. + Now Instant +} + +// Validate refuses a write contract 1.6 refuses: a blank or overlong model ID, a dimension outside 1 to +// MaxDimension, a vector that is not Dimension finite components, a blank content hash, a negative revision. +func (w EmbeddingWrite) Validate() error { + if err := ValidateModelID(w.ModelID); err != nil { + return err + } + if w.Dimension < 1 || w.Dimension > MaxDimension { + return fmt.Errorf("the dimension must be 1 to %d", MaxDimension) + } + if err := ValidateVector(w.Vector, w.Dimension); err != nil { + return err + } + if isNullOrWhitespace(w.ContentHash) { + return errors.New("the content hash must not be blank") + } + if w.SourceRevision < 0 { + return errors.New("the source revision must not be negative") + } + return nil +} + +// VectorQuery is a nearest-neighbour search in exactly Scope, over embeddings of ModelID and the vector's length. +type VectorQuery struct { + Scope Scope + ModelID string + Vector []float32 + // Statuses are EligibleStatuses() when nil. + Statuses []ExperienceStatus + // MinConfidence is DefaultMinConfidence when nil. + MinConfidence *float64 + // Limit is StatementLimit when 0. + Limit int + ExcludeModelAuthored bool +} + +// Validate refuses a search contract 1.6 refuses: a blank or overlong model ID, a vector that is not 1 to +// MaxDimension finite components, an empty or unknown status, a minimum confidence outside 0 to 1, a negative limit. +func (q VectorQuery) Validate() error { + if err := ValidateModelID(q.ModelID); err != nil { + return err + } + if err := ValidateVector(q.Vector, 0); err != nil { + return err + } + if q.Statuses != nil && len(q.Statuses) == 0 { + return errors.New("the statuses must not be empty: leave them nil for the eligible statuses") + } + for _, status := range q.Statuses { + if !contains(experienceStatuses, string(status)) { + return errors.New("the statuses must be experience status member names") + } + } + if q.MinConfidence != nil && !(*q.MinConfidence >= 0 && *q.MinConfidence <= 1) { + return errors.New("the minimum confidence must be a number from 0 to 1") + } + if q.Limit < 0 { + return errors.New("the limit must be strictly positive") + } + return nil +} + +// EffectiveStatuses are the statuses searched. +func (q VectorQuery) EffectiveStatuses() []ExperienceStatus { + if q.Statuses == nil { + return EligibleStatuses() + } + return q.Statuses +} + +// EffectiveMinConfidence is the minimum confidence searched. +func (q VectorQuery) EffectiveMinConfidence() float64 { + if q.MinConfidence == nil { + return DefaultMinConfidence + } + return *q.MinConfidence +} + +// EffectiveLimit is the number of rows the search may return. +func (q VectorQuery) EffectiveLimit() int { + if q.Limit == 0 { + return StatementLimit + } + return q.Limit +} + +// EmbeddingWriteOutcome is how a write ended. +type EmbeddingWriteOutcome string + +// The write outcomes. +const ( + // EmbeddingWritten: the row now holds this vector and descriptor. + EmbeddingWritten EmbeddingWriteOutcome = "Written" + // EmbeddingStale: the record moved on since the summary was read (or a newer embedding is stored); nothing changed. + EmbeddingStale EmbeddingWriteOutcome = "Stale" + // EmbeddingMissing: an unknown, other-scope or erased record; nothing was created. + EmbeddingMissing EmbeddingWriteOutcome = "Missing" +) + +// EmbeddingWriteResult is the outcome, and, for Stale, the record's current revision. +type EmbeddingWriteResult struct { + Outcome EmbeddingWriteOutcome + CurrentRevision int64 +} + +// VectorSearchOutcome is how a vector search ended. +type VectorSearchOutcome string + +// The search outcomes. +const ( + VectorFound VectorSearchOutcome = "Found" + VectorModelMismatch VectorSearchOutcome = "ModelMismatch" + VectorDimensionMismatch VectorSearchOutcome = "DimensionMismatch" +) + +// RelevanceOf is clamp(1 - distance / 2, 0, 1), and 0 for NaN. +func RelevanceOf(distance float64) float64 { + if math.IsNaN(distance) { + return 0 + } + return math.Min(math.Max(1.0-distance/2.0, 0), 1) +} + +// VectorHit is one record the vector search returned, and its cosine distance (NaN for a zero-magnitude vector). +type VectorHit struct { + Record ExperienceRecord + Distance float64 +} + +// Relevance is RelevanceOf the hit's distance. +func (h VectorHit) Relevance() float64 { return RelevanceOf(h.Distance) } + +// VectorSearchResult is the outcome, and the hits nearest first (ties by experience ID, NaN last); no hits unless +// the outcome is VectorFound. +type VectorSearchResult struct { + Outcome VectorSearchOutcome + Hits []VectorHit +} + +// OrderHits orders hits nearest first; equal distances by experience ID as text, ordinally; NaN after every number, +// equal to each other. It sorts hits in place and returns them. +func OrderHits(hits []VectorHit) []VectorHit { + sort.SliceStable(hits, func(i, j int) bool { + a, b := hits[i].Distance, hits[j].Distance + aNaN, bNaN := math.IsNaN(a), math.IsNaN(b) + switch { + case aNaN != bNaN: + return bNaN + case !aNaN && a != b: + return a < b + } + return hits[i].Record.ExperienceID.String() < hits[j].Record.ExperienceID.String() + }) + return hits +} + +// ---- Merging the two channels -------------------------------------------------------------------------------------- + +// Channel names where a merged candidate came from. +type Channel string + +// The channels. +const ( + ChannelText Channel = "text" + ChannelVector Channel = "vector" +) + +// Merged is one entry of the merged list: the candidate kept (the first channel's), the relevance kept, and the +// channel the candidate came from. +type Merged[T any] struct { + Item T + Relevance float64 + Source Channel +} + +func clampedRelevance(relevance float64) float64 { + if math.IsNaN(relevance) { + return 0 + } + return math.Min(math.Max(relevance, 0), 1) +} + +// MergeChannels is contract 1.6's steps 1 and 3. The floor drops vector candidates below minimumVectorRelevance +// (nil: no floor; 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. +func MergeChannels[T any](text, vector []T, key func(T) UUID, relevance func(T) float64, + minimumVectorRelevance *float64) []Merged[T] { + merged := make([]Merged[T], 0, len(text)+len(vector)) + positions := make(map[UUID]int, len(text)+len(vector)) + for _, item := range text { + if _, held := positions[key(item)]; !held { + positions[key(item)] = len(merged) + } + merged = append(merged, Merged[T]{item, relevance(item), ChannelText}) + } + for _, item := range vector { + value := relevance(item) + if minimumVectorRelevance != nil && !(value >= *minimumVectorRelevance) { + continue + } + position, held := positions[key(item)] + switch { + case !held: + positions[key(item)] = len(merged) + merged = append(merged, Merged[T]{item, value, ChannelVector}) + case clampedRelevance(value) > clampedRelevance(merged[position].Relevance): + merged[position].Relevance = value + } + } + return merged +} + +// CheckChannel is contract 1.6's step 2 for one channel: over its first CandidateLimit candidates, one outside scope +// or a repeated experience ID is an error wrapping ErrSearchSafety; it returns those first candidates. +func CheckChannel(rows []SearchRow, scope Scope) ([]SearchRow, error) { + considered := rows + if len(considered) > CandidateLimit { + considered = considered[:CandidateLimit] + } + seen := map[UUID]bool{} + for i := range considered { + row := &considered[i] + if row.SharedByGrant || !row.Record.Scope.Equal(scope) { + return nil, errors.Join(ErrSearchSafety, + errors.New("a row lies outside the requested scope; nothing is returned")) + } + if seen[row.Record.ExperienceID] { + return nil, errors.Join(ErrSearchSafety, + errors.New("the same record came back more than once; nothing is returned")) + } + seen[row.Record.ExperienceID] = true + } + return considered, nil +} + +// RankHybrid is the whole of contract 1.6's steps 1 to 4: the floor, the per-channel checks and truncation, the +// merge, then Rank's scoring 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). +func RankHybrid(textRows, vectorRows []SearchRow, scope Scope, now Instant, preferredEnvironment map[string]string, + minimumVectorRelevance *float64, excludeModelAuthored bool) (RankedResult, error) { + var floored []SearchRow + for i := range vectorRows { + if i > CandidateLimit { + break + } + if minimumVectorRelevance == nil || vectorRows[i].Relevance >= *minimumVectorRelevance { + floored = append(floored, vectorRows[i]) + } + } + // The checks are over each channel's first 50 candidates as returned (the floor only drops vector candidates). + text, err := CheckChannel(textRows, scope) + if err != nil { + return RankedResult{}, err + } + vector, err := CheckChannel(floored, scope) + if err != nil { + return RankedResult{}, err + } + merged := MergeChannels(text, vector, func(row SearchRow) UUID { return row.Record.ExperienceID }, + func(row SearchRow) float64 { return row.Relevance }, nil) + rows := make([]SearchRow, len(merged)) + for i, entry := range merged { + rows[i] = SearchRow{Record: entry.Item.Record, Relevance: entry.Relevance} + } + ranked := []RankedRecord{} + for start := 0; start < len(rows); start += CandidateLimit { + result, err := Rank(rows[start:min(len(rows), start+CandidateLimit)], scope, now, preferredEnvironment, + excludeModelAuthored) + if err != nil { + return RankedResult{}, err + } + ranked = append(ranked, result.Ranked...) + } + sort.SliceStable(ranked, func(i, j int) bool { + if ranked[i].Score != ranked[j].Score { + return ranked[i].Score > ranked[j].Score + } + return ranked[i].Record.ExperienceID.String() < ranked[j].Record.ExperienceID.String() + }) + return RankedResult{Ranked: ranked, Truncated: len(textRows) > CandidateLimit || len(floored) > CandidateLimit}, nil +} diff --git a/go/embedding_test.go b/go/embedding_test.go new file mode 100644 index 00000000..0ba30e1c --- /dev/null +++ b/go/embedding_test.go @@ -0,0 +1,212 @@ +package agentexperience + +import ( + "errors" + "math" + "strings" + "testing" +) + +// Contract 1.6's pure rules beyond the fixtures: validation, the literal, relevance, ordering, and the hybrid ranking. + +func TestVectorValidation(t *testing.T) { + nan, inf := float32(math.NaN()), float32(math.Inf(-1)) + for name, vector := range map[string][]float32{"empty": nil, "nan": {1, nan}, "infinite": {inf}, "too long": make([]float32, MaxDimension+1)} { + if ValidateVector(vector, 0) == nil { + t.Errorf("%s was accepted", name) + } + } + if ValidateVector([]float32{1, 2}, 3) == nil || ValidateVector([]float32{1, 2}, 2) != nil || + ValidateVector(make([]float32, MaxDimension), 0) != nil || ValidateVector([]float32{0, 0}, 0) != nil { + t.Fatal("dimension and zero-magnitude rules") + } + if ValidateModelID("") == nil || ValidateModelID(" \u0085") == nil || ValidateModelID(strings.Repeat("m", 257)) == nil || + ValidateModelID(strings.Repeat("m", 256)) != nil || ValidateModelID(strings.Repeat("\U0001F600", 129)) == nil { + t.Fatal("model ID rules (256 UTF-16 code units)") + } +} + +func TestTheVectorLiteralIsPgvectorsFormWithShortestRoundTripComponents(t *testing.T) { + if got := VectorLiteral([]float32{1, 0.1, -2.5, 1e-7, 3e10}); got != "[1,0.1,-2.5,1e-07,3e+10]" { + t.Fatal(got) + } +} + +func TestRelevanceAndHitOrder(t *testing.T) { + if RelevanceOf(math.NaN()) != 0 || RelevanceOf(-1) != 1 || RelevanceOf(3) != 0 || RelevanceOf(1) != 0.5 { + t.Fatal("relevance is clamp(1 - d/2, 0, 1), NaN as 0") + } + records := seedRows(t) + hit := func(i int, distance float64) VectorHit { return VectorHit{Record: records[i], Distance: distance} } + ordered := OrderHits([]VectorHit{hit(0, math.NaN()), hit(2, 0.5), hit(1, 0.5), hit(3, math.NaN()), hit(4, 0.1)}) + want := []ExperienceRecord{records[4]} + pair := []ExperienceRecord{records[1], records[2]} + if pair[0].ExperienceID.String() > pair[1].ExperienceID.String() { + pair[0], pair[1] = pair[1], pair[0] + } + want = append(want, pair...) + nans := []ExperienceRecord{records[0], records[3]} + if nans[0].ExperienceID.String() > nans[1].ExperienceID.String() { + nans[0], nans[1] = nans[1], nans[0] + } + want = append(want, nans...) + for i, h := range ordered { + if h.Record.ExperienceID != want[i].ExperienceID { + t.Fatalf("hit %d is %s, not %s", i, h.Record.ExperienceID, want[i].ExperienceID) + } + } +} + +func rowsOf(records []ExperienceRecord, relevance float64) []SearchRow { + rows := make([]SearchRow, len(records)) + for i, record := range records { + rows[i] = SearchRow{Record: record, Relevance: relevance} + } + return rows +} + +// manyRecords are n copies of a seed record, each with an ID of its own. +func manyRecords(t *testing.T, n int) []ExperienceRecord { + t.Helper() + base := seedRows(t)[0] + records := make([]ExperienceRecord, n) + for i := range records { + records[i] = base + records[i].ExperienceID = UUID{0: 0xab, 6: 0x40, 8: 0x80, 14: byte(i >> 8), 15: byte(i)} + } + return records +} + +func TestRankHybridScoresEveryMergedCandidateUpToOneHundred(t *testing.T) { + records := manyRecords(t, 100) + scope := records[0].Scope + result, err := RankHybrid(rowsOf(records[:50], 0.2), rowsOf(records[50:], 0.9), scope, t0, nil, nil, false) + if err != nil || len(result.Ranked) != 100 || result.Truncated { + t.Fatalf("%v %d %v", err, len(result.Ranked), result.Truncated) + } + // The vector candidates' higher relevance puts them first, whichever slice of 50 they were scored in. + for i, entry := range result.Ranked { + if (i < 50) != (entry.Record.ExperienceID[15] >= 50) { + t.Fatalf("entry %d is record %d", i, entry.Record.ExperienceID[15]) + } + } + for i := 1; i < len(result.Ranked); i++ { + a, b := result.Ranked[i-1], result.Ranked[i] + if a.Score < b.Score || (a.Score == b.Score && a.Record.ExperienceID.String() > b.Record.ExperienceID.String()) { + t.Fatalf("entries %d and %d are out of order", i-1, i) + } + } + // A record in both channels is one entry, with the higher relevance. + both, err := RankHybrid(rowsOf(records[:3], 0.2), rowsOf(records[2:4], 0.9), scope, t0, nil, nil, false) + if err != nil || len(both.Ranked) != 4 { + t.Fatalf("%v %d", err, len(both.Ranked)) + } + // Truncated when either channel returned 51. + more, err := RankHybrid(rowsOf(manyRecords(t, 51), 0.2), nil, scope, t0, nil, nil, false) + if err != nil || !more.Truncated || len(more.Ranked) != 50 { + t.Fatalf("%v %v %d", err, more.Truncated, len(more.Ranked)) + } +} + +func TestRankHybridWithNoVectorCandidatesIsRank(t *testing.T) { + records := manyRecords(t, 20) + scope := records[0].Scope + for i := range records { + records[i].ReuseConfidence = float64(i) / 20 + } + rows := rowsOf(records, 0.4) + hybrid, err := RankHybrid(rows, nil, scope, t0, nil, nil, false) + plain, plainErr := Rank(rows, scope, t0, nil, false) + if err != nil || plainErr != nil || len(hybrid.Ranked) != len(plain.Ranked) { + t.Fatalf("%v %v", err, plainErr) + } + for i := range plain.Ranked { + if hybrid.Ranked[i].Record.ExperienceID != plain.Ranked[i].Record.ExperienceID || hybrid.Ranked[i].Score != plain.Ranked[i].Score { + t.Fatalf("entry %d differs", i) + } + } +} + +func TestRankHybridFloorsOnlyVectorCandidates(t *testing.T) { + records := manyRecords(t, 4) + scope := records[0].Scope + floor := 0.5 + text := rowsOf(records[:2], 0.1) // below the floor, never floored + vector := []SearchRow{{Record: records[2], Relevance: 0.5}, {Record: records[3], Relevance: 0.4999}} + result, err := RankHybrid(text, vector, scope, t0, nil, &floor, false) + if err != nil || len(result.Ranked) != 3 { + t.Fatalf("%v %d", err, len(result.Ranked)) + } + for _, entry := range result.Ranked { + if entry.Record.ExperienceID == records[3].ExperienceID { + t.Fatal("a vector candidate below the floor was kept") + } + } +} + +func TestRankHybridFailsTheWholeResultForARowOutsideTheScopeOrADuplicateInAChannel(t *testing.T) { + records := manyRecords(t, 4) + scope := records[0].Scope + foreign := records[3] + foreign.Scope.AgentID = Ptr("other") + for name, run := range map[string]func() error{ + "foreign vector row": func() error { + _, err := RankHybrid(rowsOf(records[:1], 0.1), []SearchRow{{Record: foreign, Relevance: 0.9}}, scope, t0, nil, nil, false) + return err + }, + "foreign text row": func() error { + _, err := RankHybrid([]SearchRow{{Record: foreign, Relevance: 0.9}}, rowsOf(records[:1], 0.1), scope, t0, nil, nil, false) + return err + }, + "shared by grant": func() error { + _, err := RankHybrid(nil, []SearchRow{{Record: records[0], Relevance: 0.9, SharedByGrant: true}}, scope, t0, nil, nil, false) + return err + }, + "repeated in the vector channel": func() error { + _, err := RankHybrid(nil, rowsOf([]ExperienceRecord{records[0], records[0]}, 0.9), scope, t0, nil, nil, false) + return err + }, + } { + if err := run(); !errors.Is(err, ErrSearchSafety) { + t.Errorf("%s: %v", name, err) + } + } +} + +func TestEmbeddingWriteAndQueryValidate(t *testing.T) { + good := EmbeddingWrite{ExperienceID: NewRandomUUID(), Scope: scope, ModelID: "m", Dimension: 2, ContentHash: "h", + SourceRevision: 0, Vector: []float32{1, 0}, Now: t0} + if err := good.Validate(); err != nil { + t.Fatal(err) + } + for name, mutate := range map[string]func(*EmbeddingWrite){ + "blank model": func(w *EmbeddingWrite) { w.ModelID = " " }, + "zero dimension": func(w *EmbeddingWrite) { w.Dimension = 0; w.Vector = nil }, + "wrong length": func(w *EmbeddingWrite) { w.Dimension = 3 }, + "huge dimension": func(w *EmbeddingWrite) { w.Dimension = MaxDimension + 1 }, + "blank hash": func(w *EmbeddingWrite) { w.ContentHash = "" }, + "negative source": func(w *EmbeddingWrite) { w.SourceRevision = -1 }, + } { + w := good + mutate(&w) + if w.Validate() == nil { + t.Errorf("%s was accepted", name) + } + } + confidence := 1.5 + for name, q := range map[string]VectorQuery{ + "empty statuses": {Scope: scope, ModelID: "m", Vector: []float32{1}, Statuses: []ExperienceStatus{}}, + "bad status": {Scope: scope, ModelID: "m", Vector: []float32{1}, Statuses: []ExperienceStatus{"Bogus"}}, + "confidence": {Scope: scope, ModelID: "m", Vector: []float32{1}, MinConfidence: &confidence}, + "limit": {Scope: scope, ModelID: "m", Vector: []float32{1}, Limit: -1}, + } { + if q.Validate() == nil { + t.Errorf("%s was accepted", name) + } + } + q := VectorQuery{Scope: scope, ModelID: "m", Vector: []float32{1}} + if q.Validate() != nil || q.EffectiveLimit() != StatementLimit || q.EffectiveMinConfidence() != DefaultMinConfidence || + len(q.EffectiveStatuses()) != 2 { + t.Fatal("defaults") + } +} diff --git a/go/inmemory/store.go b/go/inmemory/store.go index bc4885c8..2626fb51 100644 --- a/go/inmemory/store.go +++ b/go/inmemory/store.go @@ -1,5 +1,6 @@ // Package inmemory is an in-memory store for development and tests: no PostgreSQL. It keeps only what the contract -// needs (insert a finalized record, read one back, the text search, and when each record was last injected), +// needs (insert a finalized record, read one back, the text search, when each record was last injected, and the +// embeddings of contract 1.6 with their nearest-neighbour search), // honouring scope and the authorization // bounds the way the PostgreSQL store's row-level security does. Its text search is the .NET in-memory store's: // words normalized to Unicode form KC and lower-cased, the English stopwords PostgreSQL drops left out of a query, @@ -12,6 +13,7 @@ package inmemory import ( "context" "errors" + "math" "sort" "sync" "time" @@ -176,6 +178,17 @@ type Store struct { order []ae.UUID events map[ae.UUID]Event injections map[ae.UUID]Injection + embeddings map[ae.UUID]Embedding +} + +// Embedding is a record's stored embedding (contract 1.6): the model, the dimension, the content hash of the summary +// that was embedded, the record revision it was read at, and the vector, its components 32-bit floats. +type Embedding struct { + ModelID string + Dimension int + ContentHash string + SourceRevision int64 + Vector []float32 } // Injection is when a record was last injected in its own scope, and how many injected blocks delivered it. @@ -186,7 +199,8 @@ type Injection struct { // New is a store holding records (each experience ID once). func New(records ...ae.ExperienceRecord) (*Store, error) { - s := &Store{records: map[ae.UUID]*indexed{}, events: map[ae.UUID]Event{}, injections: map[ae.UUID]Injection{}} + s := &Store{records: map[ae.UUID]*indexed{}, events: map[ae.UUID]Event{}, injections: map[ae.UUID]Injection{}, + embeddings: map[ae.UUID]Embedding{}} for _, record := range records { if _, held := s.records[record.ExperienceID]; held { return nil, errors.New("two records share one experience ID") @@ -465,7 +479,138 @@ func (s *Store) SetInjection(id ae.UUID, injection Injection) { s.injections[id] = injection } +// WriteEmbedding 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 its 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. A write that fails Validate is an error, and stores nothing. +func (s *Store) WriteEmbedding(ctx context.Context, bounds ae.AuthorizationBounds, + write ae.EmbeddingWrite) (ae.EmbeddingWriteResult, error) { + if err := ctx.Err(); err != nil { + return ae.EmbeddingWriteResult{}, err + } + if err := write.Validate(); err != nil { + return ae.EmbeddingWriteResult{}, err + } + if !Permits(bounds, write.Scope) { + return ae.EmbeddingWriteResult{Outcome: ae.EmbeddingMissing}, nil + } + s.mu.Lock() + defer s.mu.Unlock() + held, ok := s.records[write.ExperienceID] + if !ok || !held.record.Scope.Equal(write.Scope) { + return ae.EmbeddingWriteResult{Outcome: ae.EmbeddingMissing}, nil + } + revision := held.record.Revision + if existing, stored := s.embeddings[write.ExperienceID]; revision != write.SourceRevision || + (stored && existing.SourceRevision > write.SourceRevision) { + return ae.EmbeddingWriteResult{Outcome: ae.EmbeddingStale, CurrentRevision: revision}, nil + } + s.embeddings[write.ExperienceID] = Embedding{ModelID: write.ModelID, Dimension: write.Dimension, + ContentHash: write.ContentHash, SourceRevision: write.SourceRevision, + Vector: append([]float32(nil), write.Vector...)} + return ae.EmbeddingWriteResult{Outcome: ae.EmbeddingWritten}, nil +} + +// Embedding is a record's stored embedding, for tests; ok is false when it has none. +func (s *Store) Embedding(id ae.UUID) (Embedding, bool) { + s.mu.RLock() + defer s.mu.RUnlock() + held, ok := s.embeddings[id] + held.Vector = append([]float32(nil), held.Vector...) + return held, ok +} + +// CosineDistance is 1 - cos of two vectors of equal length, clamped to [0, 2]: the components are 32-bit floats, +// the arithmetic double. NaN when either has no magnitude (a zero vector). +func CosineDistance(left, right []float32) float64 { + var dot, leftSquares, rightSquares float64 + for i := range left { + a, b := float64(left[i]), float64(right[i]) + dot += a * b + leftSquares += a * a + rightSquares += b * b + } + norm := math.Sqrt(leftSquares) * math.Sqrt(rightSquares) + if norm == 0 || math.IsNaN(norm) { + return math.NaN() + } + return math.Min(math.Max(1.0-dot/norm, 0), 2) +} + +// SearchVectors is 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, ties by experience ID, NaN last, at most the 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. A query that fails Validate is an error. +func (s *Store) SearchVectors(ctx context.Context, bounds ae.AuthorizationBounds, + query ae.VectorQuery) (ae.VectorSearchResult, error) { + if err := ctx.Err(); err != nil { + return ae.VectorSearchResult{}, err + } + if err := query.Validate(); err != nil { + return ae.VectorSearchResult{}, err + } + scope := query.Scope + if !Permits(bounds, scope) { + return ae.VectorSearchResult{Outcome: ae.VectorFound}, nil + } + eligible := map[ae.ExperienceStatus]bool{} + for _, status := range query.EffectiveStatuses() { + eligible[status] = true + } + minConfidence := query.EffectiveMinConfidence() + + type candidate struct { + record *indexed + embedding Embedding + } + s.mu.RLock() + var held []candidate + for _, id := range s.order { + embedding, ok := s.embeddings[id] + if !ok { + continue + } + x := s.records[id] + record := &x.record + if !record.Scope.Equal(scope) || !eligible[record.Status] || record.ReuseConfidence < minConfidence || + (query.ExcludeModelAuthored && ae.IsModelAuthored(record.Reflection)) { + continue + } + held = append(held, candidate{x, embedding}) + } + s.mu.RUnlock() + + var hits []ae.VectorHit + modelHeld := false + for _, c := range held { + if c.embedding.ModelID != query.ModelID { + continue + } + modelHeld = true + if c.embedding.Dimension != len(query.Vector) || len(c.embedding.Vector) != len(query.Vector) { + continue + } + hits = append(hits, ae.VectorHit{Record: mustClone(c.record.record), + Distance: CosineDistance(c.embedding.Vector, query.Vector)}) + } + switch { + case len(hits) > 0: + hits = ae.OrderHits(hits) + if limit := query.EffectiveLimit(); len(hits) > limit { + hits = hits[:limit] + } + return ae.VectorSearchResult{Outcome: ae.VectorFound, Hits: hits}, nil + case len(held) == 0: + return ae.VectorSearchResult{Outcome: ae.VectorFound}, nil + case !modelHeld: + return ae.VectorSearchResult{Outcome: ae.VectorModelMismatch}, nil + } + return ae.VectorSearchResult{Outcome: ae.VectorDimensionMismatch}, nil +} + var ( _ ae.Store = (*Store)(nil) _ ae.InjectionRecorder = (*Store)(nil) + _ ae.EmbeddingIndex = (*Store)(nil) ) diff --git a/go/inmemory/vectors_test.go b/go/inmemory/vectors_test.go new file mode 100644 index 00000000..5a2eb279 --- /dev/null +++ b/go/inmemory/vectors_test.go @@ -0,0 +1,247 @@ +package inmemory + +// Contract 1.6 on the in-memory store: the fixtures that need no database, and the rules the fixtures do not state. + +import ( + "context" + "errors" + "math" + "path/filepath" + "sync" + "testing" + + ae "github.com/fabbrik/AgentExperience.NET/go" + ct "github.com/fabbrik/AgentExperience.NET/go/internal/contracttest" +) + +const vectorTolerance = 1e-6 + +func fixtureSteps(t *testing.T, family, name string) ([]ae.ExperienceRecord, []any, []any) { + t.Helper() + get := func(o *ae.Object, key string) any { value, _ := o.Get(key); return value } + var records []ae.ExperienceRecord + for _, stored := range ct.MustReadJSON(filepath.Join(ct.Fixtures, family, name, "seed.json")).([]any) { + record, err := ae.FromStored(stored) + if err != nil { + t.Fatal(err) + } + records = append(records, record) + } + steps := get(ct.MustReadJSON(filepath.Join(ct.Fixtures, family, name, "steps.json")).(*ae.Object), "steps") + expected := get(ct.MustReadJSON(filepath.Join(ct.Fixtures, family, name, "expected.json")).(*ae.Object), "steps") + return records, steps.([]any), expected.([]any) +} + +func TestTheVectorFixturesThatNeedNoDatabaseGiveTheirOutcomes(t *testing.T) { + get := func(o *ae.Object, key string) any { value, _ := o.Get(key); return value } + ran := 0 + for _, family := range []string{"embedding-write", "vector-search"} { + for _, name := range ct.Cases(family) { + records, steps, expected := fixtureSteps(t, family, name) + needsDatabase := false + for _, item := range steps { + step := item.(*ae.Object) + needsDatabase = needsDatabase || get(step, "erase") != nil || get(step, "setRevision") != nil + } + if needsDatabase { + continue + } + ran++ + t.Run(family+"/"+name, func(t *testing.T) { + store, err := New(records...) + if err != nil { + t.Fatal(err) + } + ctx := context.Background() + for i, item := range steps { + step := item.(*ae.Object) + var actual any + if write := get(step, "write"); write != nil { + w := ct.EmbeddingWriteOf(write) + result, err := store.WriteEmbedding(ctx, ae.BoundsOf(w.Scope), w) + if err != nil { + t.Fatal(err) + } + actual = ct.WriteOutcome(result) + } else { + q := ct.VectorQueryOf(get(step, "search")) + result, err := store.SearchVectors(ctx, ae.BoundsOf(q.Scope), q) + if err != nil { + t.Fatal(err) + } + actual = ct.SearchOutcome(result) + } + if difference := ct.Difference(expected[i], ct.AsWritten(actual), "$", + ct.Options{Tolerance: vectorTolerance, IgnoreKind: true}); difference != "" { + t.Fatalf("step %d: %s", i, difference) + } + } + }) + } + } + if ran == 0 { + t.Fatal("no vector fixture ran") + } +} + +func embedded(t *testing.T, records ...ae.ExperienceRecord) *Store { + t.Helper() + store, err := New(records...) + if err != nil { + t.Fatal(err) + } + return store +} + +func writeOf(r ae.ExperienceRecord, model string, revision int64, vector ...float32) ae.EmbeddingWrite { + return ae.EmbeddingWrite{ExperienceID: r.ExperienceID, Scope: r.Scope, ModelID: model, Dimension: len(vector), + ContentHash: "h", SourceRevision: revision, Vector: vector, Now: at} +} + +func TestANewerEmbeddingIsNeverReplacedByAnOlderOneAndAMovedRecordIsStale(t *testing.T) { + r := record(1, "rotate-keys", "Rotate the keys", "Rotate first.", ae.StatusValidated, 0.9) + r.Revision = 2 + store := embedded(t, r) + ctx, bounds := context.Background(), ae.BoundsOf(scope) + + result, err := store.WriteEmbedding(ctx, bounds, writeOf(r, "m", 2, 0, 1)) + if err != nil || result.Outcome != ae.EmbeddingWritten { + t.Fatalf("%+v, %v", result, err) + } + result, err = store.WriteEmbedding(ctx, bounds, writeOf(r, "m", 1, 1, 0)) + if err != nil || result.Outcome != ae.EmbeddingStale || result.CurrentRevision != 2 { + t.Fatalf("%+v, %v", result, err) + } + if held, _ := store.Embedding(r.ExperienceID); held.SourceRevision != 2 || held.Vector[1] != 1 { + t.Fatalf("%+v", held) + } + other := scope + other.AgentID = ae.Ptr("other") + missing := writeOf(r, "m", 2, 0, 1) + missing.Scope = other + if result, err = store.WriteEmbedding(ctx, ae.BoundsOf(other), missing); err != nil || result.Outcome != ae.EmbeddingMissing { + t.Fatalf("%+v, %v", result, err) + } + // Outside the bounds the host declared, the record is not there. + if result, err = store.WriteEmbedding(ctx, ae.BoundsOf(other), writeOf(r, "m", 2, 0, 1)); err != nil || + result.Outcome != ae.EmbeddingMissing { + t.Fatalf("%+v, %v", result, err) + } +} + +func TestVectorsAreCompared32BitAndTiesAreOrderedByIDAndNaNIsLast(t *testing.T) { + a := record(1, "a", "A summary", "Lesson a", ae.StatusValidated, 0.9) + b := record(2, "b", "B summary", "Lesson b", ae.StatusValidated, 0.9) + c := record(3, "c", "C summary", "Lesson c", ae.StatusValidated, 0.9) + store := embedded(t, c, b, a) + ctx, bounds := context.Background(), ae.BoundsOf(scope) + for _, r := range []ae.ExperienceRecord{c, b} { + if _, err := store.WriteEmbedding(ctx, bounds, writeOf(r, "m", 1, 1, 0)); err != nil { + t.Fatal(err) + } + } + if _, err := store.WriteEmbedding(ctx, bounds, writeOf(a, "m", 1, 0, 0)); err != nil { // no magnitude + t.Fatal(err) + } + result, err := store.SearchVectors(ctx, bounds, ae.VectorQuery{Scope: scope, ModelID: "m", Vector: []float32{2, 0}}) + if err != nil || result.Outcome != ae.VectorFound || len(result.Hits) != 3 { + t.Fatalf("%+v, %v", result, err) + } + if result.Hits[0].Record.ExperienceID != b.ExperienceID || result.Hits[1].Record.ExperienceID != c.ExperienceID || + result.Hits[2].Record.ExperienceID != a.ExperienceID || !math.IsNaN(result.Hits[2].Distance) || + result.Hits[2].Relevance() != 0 || result.Hits[0].Distance != 0 { + t.Fatalf("%+v", result.Hits) + } + // A limit keeps the nearest. + limited, _ := store.SearchVectors(ctx, bounds, ae.VectorQuery{Scope: scope, ModelID: "m", Vector: []float32{2, 0}, Limit: 1}) + if len(limited.Hits) != 1 || limited.Hits[0].Record.ExperienceID != b.ExperienceID { + t.Fatalf("%+v", limited.Hits) + } + // The hits are copies: changing one changes nothing stored. + result.Hits[0].Record.TaskID = "changed" + if again, _ := store.SearchVectors(ctx, bounds, ae.VectorQuery{Scope: scope, ModelID: "m", Vector: []float32{2, 0}}); again.Hits[0].Record.TaskID != "b" { + t.Fatal("a hit shares memory with the store") + } + // Nothing outside the bounds. + narrow := ae.BoundsOf(ae.Scope{TenantID: "acme", ApplicationID: "support", ProjectID: "refunds", AgentID: ae.Ptr("someone-else")}) + if found, err := store.SearchVectors(ctx, narrow, ae.VectorQuery{Scope: scope, ModelID: "m", Vector: []float32{2, 0}}); err != nil || + len(found.Hits) != 0 { + t.Fatalf("%+v, %v", found, err) + } +} + +func TestAnEmptyVectorSearchSaysWhy(t *testing.T) { + r := record(1, "a", "A summary", "Lesson", ae.StatusValidated, 0.9) + store := embedded(t, r) + ctx, bounds := context.Background(), ae.BoundsOf(scope) + search := func(model string, vector ...float32) ae.VectorSearchOutcome { + result, err := store.SearchVectors(ctx, bounds, ae.VectorQuery{Scope: scope, ModelID: model, Vector: vector}) + if err != nil || len(result.Hits) != 0 { + t.Fatalf("%+v, %v", result, err) + } + return result.Outcome + } + if outcome := search("m", 1, 0); outcome != ae.VectorFound { + t.Fatalf("nothing stored: %s", outcome) + } + if _, err := store.WriteEmbedding(ctx, bounds, writeOf(r, "m", 1, 1, 0)); err != nil { + t.Fatal(err) + } + if outcome := search("another", 1, 0); outcome != ae.VectorModelMismatch { + t.Fatalf("another model: %s", outcome) + } + if outcome := search("m", 1, 0, 0); outcome != ae.VectorDimensionMismatch { + t.Fatalf("another dimension: %s", outcome) + } +} + +func TestVectorCallsRefuseInvalidInputAndADoneContext(t *testing.T) { + r := record(1, "a", "A summary", "Lesson", ae.StatusValidated, 0.9) + store := embedded(t, r) + ctx, bounds := context.Background(), ae.BoundsOf(scope) + bad := writeOf(r, "m", 1, 1, 0) + bad.Vector = []float32{1, float32(math.NaN())} + if _, err := store.WriteEmbedding(ctx, bounds, bad); err == nil { + t.Fatal("a non-finite vector was stored") + } + if _, ok := store.Embedding(r.ExperienceID); ok { + t.Fatal("a refused write stored something") + } + if _, err := store.SearchVectors(ctx, bounds, ae.VectorQuery{Scope: scope, ModelID: "m"}); err == nil { + t.Fatal("an empty query vector was searched") + } + done, cancel := context.WithCancel(ctx) + cancel() + if _, err := store.WriteEmbedding(done, bounds, writeOf(r, "m", 1, 1, 0)); !errors.Is(err, context.Canceled) { + t.Fatalf("write: %v", err) + } + if _, err := store.SearchVectors(done, bounds, ae.VectorQuery{Scope: scope, ModelID: "m", Vector: []float32{1}}); !errors.Is(err, context.Canceled) { + t.Fatalf("search: %v", err) + } +} + +func TestVectorCallsAreSafeToRunTogether(t *testing.T) { + var records []ae.ExperienceRecord + for n := byte(1); n <= 20; n++ { + records = append(records, record(n, "task", "Summary", "Lesson", ae.StatusValidated, 0.9)) + } + store := embedded(t, records...) + ctx, bounds := context.Background(), ae.BoundsOf(scope) + var wg sync.WaitGroup + for _, r := range records { + wg.Add(2) + go func() { + defer wg.Done() + _, _ = store.WriteEmbedding(ctx, bounds, writeOf(r, "m", 1, 1, 0)) + }() + go func() { + defer wg.Done() + _, _ = store.SearchVectors(ctx, bounds, ae.VectorQuery{Scope: scope, ModelID: "m", Vector: []float32{1, 0}}) + }() + } + wg.Wait() + result, err := store.SearchVectors(ctx, bounds, ae.VectorQuery{Scope: scope, ModelID: "m", Vector: []float32{1, 0}}) + if err != nil || len(result.Hits) != 20 { + t.Fatalf("%d hits, %v", len(result.Hits), err) + } +} diff --git a/go/internal/contracttest/contracttest.go b/go/internal/contracttest/contracttest.go index 31fe5f2c..c9ed8197 100644 --- a/go/internal/contracttest/contracttest.go +++ b/go/internal/contracttest/contracttest.go @@ -33,6 +33,8 @@ var ( Fixtures = filepath.Join(Contract, "fixtures") Schema = filepath.Join(Contract, "schema", "experience-record.v1.json") Migrations = filepath.Join(Repository, "src", "AgentExperience.Storage.Postgres", "Migrations") + // VectorMigrations are the vectors package's migrations (contract 1.6), which need the pgvector extension. + VectorMigrations = filepath.Join(Repository, "src", "AgentExperience.Storage.Postgres.Vectors", "Migrations") ) // ReadJSON reads a JSON file, keeping every number exactly as written and every object's member order. @@ -72,21 +74,12 @@ func Cases(family string) []string { // ErrUnsupportedContractVersion is a lock naming a contract version this package does not support. var ErrUnsupportedContractVersion = errors.New("unsupported contract version") -// PendingContractVersions are the contract versions the lock may name although this package does not pass them yet: -// contract 1.6 only adds PendingVectorFamilies, so every 1.0 to 1.5 fixture still runs under its lock. -var PendingContractVersions = []string{"1.6"} - -// PendingVectorFamilies are contract 1.6's vector fixture families, which this package does not implement yet. They -// are skipped by name (TestPendingVectorFamilies), never ignored. Stories 23.2 and 23.3 implement them and remove this. -var PendingVectorFamilies = []string{"embedding-summary", "embedding-write", "hybrid-merge", "vector-search"} - -// CheckContractVersion is the lock's contract version, when this package supports it or it is pending; a clear -// failure otherwise. +// CheckContractVersion is the lock's contract version, when this package supports it; a clear failure otherwise. func CheckContractVersion(lock any) (string, error) { object, _ := lock.(*ae.Object) value, _ := object.Get("contractVersion") version, _ := value.(string) - for _, supported := range append(append([]string{}, ae.SupportedContractVersions...), PendingContractVersions...) { + for _, supported := range ae.SupportedContractVersions { if version == supported { return version, nil } @@ -438,3 +431,83 @@ func RunFinalizeCase(input any) (ae.FinalizedExperience, *ae.Object, error) { }) return finalized, answers, err } + +// ---- Contract 1.6's vector fixtures ------------------------------------------------------------------------------ + +func vectorOf(value any) []float32 { + var vector []float32 + for _, item := range list(value) { + n, _ := item.(json.Number) + f, err := strconv.ParseFloat(string(n), 32) + if err != nil { + panic(fmt.Sprintf("contracttest: %v is not a vector component", item)) + } + vector = append(vector, float32(f)) + } + return vector +} + +// EmbeddingWriteOf reads a write step of the embedding-write and vector-search fixtures. +func EmbeddingWriteOf(value any) ae.EmbeddingWrite { + o := object(value) + return ae.EmbeddingWrite{ + ExperienceID: must(ae.ParseUUID(text(get(o, "experienceId")))), + Scope: ScopeOf(get(o, "scope")), + ModelID: text(get(o, "modelId")), + Dimension: int(integer(get(o, "dimension"))), + ContentHash: text(get(o, "contentHash")), + SourceRevision: integer(get(o, "sourceRevision")), + Vector: vectorOf(get(o, "vector")), + Now: must(ae.ParseInstant(text(get(o, "now")))), + } +} + +// VectorQueryOf reads a search step of the vector-search fixtures. +func VectorQueryOf(value any) ae.VectorQuery { + o := object(value) + var statuses []ae.ExperienceStatus + for _, status := range list(get(o, "statuses")) { + statuses = append(statuses, ae.ExperienceStatus(text(status))) + } + number, _ := get(o, "minConfidence").(json.Number) + minConfidence, err := strconv.ParseFloat(string(number), 64) + if err != nil { + panic(fmt.Sprintf("contracttest: %v is not a confidence", number)) + } + excluding, _ := get(o, "excludeModelAuthored").(bool) + return ae.VectorQuery{ + Scope: ScopeOf(get(o, "scope")), + ModelID: text(get(o, "modelId")), + Vector: vectorOf(get(o, "vector")), + Statuses: statuses, + MinConfidence: &minConfidence, + Limit: int(integer(get(o, "limit"))), + ExcludeModelAuthored: excluding, + } +} + +// WriteOutcome is a write's result in the fixtures' form: {"outcome": ...} and, for Stale, "currentRevision". +func WriteOutcome(result ae.EmbeddingWriteResult) *ae.Object { + o := ae.NewObject(ae.Member{Key: "outcome", Value: string(result.Outcome)}) + if result.Outcome == ae.EmbeddingStale { + o.Set("currentRevision", json.Number(strconv.FormatInt(result.CurrentRevision, 10))) + } + return o +} + +// SearchOutcome is a search's result in the fixtures' form: the outcome, and the rows with each distance (null for +// NaN) and relevance. +func SearchOutcome(result ae.VectorSearchResult) *ae.Object { + rows := []any{} + for _, hit := range result.Hits { + row := ae.NewObject(ae.Member{Key: "experienceId", Value: hit.Record.ExperienceID.String()}) + if math.IsNaN(hit.Distance) { + row.Set("distance", nil) + } else { + row.Set("distance", json.Number(strconv.FormatFloat(hit.Distance, 'g', -1, 64))) + } + row.Set("relevance", json.Number(strconv.FormatFloat(hit.Relevance(), 'g', -1, 64))) + rows = append(rows, row) + } + return ae.NewObject(ae.Member{Key: "outcome", Value: string(result.Outcome)}, ae.Member{Key: "rows", Value: rows}) +} diff --git a/go/maf/capture.go b/go/maf/capture.go index 5bb007ac..138b47a9 100644 --- a/go/maf/capture.go +++ b/go/maf/capture.go @@ -290,6 +290,8 @@ func (c *captureScope) finishWork(ctx context.Context, status ae.RunExecutionSta if err != nil { return report(StageStore, "Storing the record failed; the run is not stored.", err) } + // The embedding is derived data: written after the record's own commit, in the background, never failing it. + c.startIndex(ctx, finalized.Record) // Stored. Told only while the invocation still waits: an abandoned work was reported as such. if ctx.Err() != nil { return false diff --git a/go/maf/inject.go b/go/maf/inject.go index df8d8377..4c87ed9c 100644 --- a/go/maf/inject.go +++ b/go/maf/inject.go @@ -86,14 +86,37 @@ func (p *ContextProvider) inject(ctx context.Context, messages []*message.Messag } parameters := ae.NewSearchParameters(identity.Scope, taskText) parameters.ExcludeModelAuthored = s.ExcludeModelAuthored + // Both channels run together; the vector channel never fails (it answers with a fallback instead), and is bounded + // by VectorTimeout, so waiting for it after the text search is bounded too. + type vectorAnswer struct { + rows []ae.SearchRow + fallback *ae.VectorFallback + } + vectorDone := make(chan vectorAnswer, 1) + if _, fallback := s.indexOf(); fallback != nil { + vectorDone <- vectorAnswer{fallback: fallback} + } else { + go func() { + vectorRows, fallback := s.vectorChannel(ctx, identity, taskText) + vectorDone <- vectorAnswer{vectorRows, fallback} + }() + } rows, err := s.Store.Search(ctx, identity.Authorization(), parameters) + vector := <-vectorDone + s.reportVectorFallback(vector.fallback) if err != nil { s.reportInjection(InjectionFailure{Reason: "The search failed; no block was added.", Err: err}) return messages } // The instant the block is built at: ranking's "now", and the injection's instant. now := s.Clock() - ranked, err := ae.Rank(rows, identity.Scope, now, s.PreferredEnvironment, s.ExcludeModelAuthored) + var ranked ae.RankedResult + if len(vector.rows) > 0 { + ranked, err = ae.RankHybrid(rows, vector.rows, identity.Scope, now, s.PreferredEnvironment, + s.MinVectorRelevance, s.ExcludeModelAuthored) + } else { + ranked, err = ae.Rank(rows, identity.Scope, now, s.PreferredEnvironment, s.ExcludeModelAuthored) + } if err != nil { s.reportInjection(InjectionFailure{Reason: "Ranking failed; no block was added.", Err: err}) return messages @@ -132,7 +155,8 @@ func (p *ContextProvider) inject(ctx context.Context, messages []*message.Messag if s.OnInjected != nil { func() { defer func() { _ = recover() }() - s.OnInjected(InjectionResult{ExperienceIDs: ids, ByteCount: block.ByteCount, Truncated: ranked.Truncated}) + s.OnInjected(InjectionResult{ExperienceIDs: ids, ByteCount: block.ByteCount, Truncated: ranked.Truncated, + VectorFallback: vector.fallback}) }() } diff --git a/go/maf/maf.go b/go/maf/maf.go index 7ec5d48d..4a0453bc 100644 --- a/go/maf/maf.go +++ b/go/maf/maf.go @@ -27,16 +27,25 @@ // panic, recorded by its value's type and re-raised; or a returned call the classifier named an error class // for). When the invocation ends, Verify decides from the host's own facts; a verdict is finalized through the // core (the deterministic lesson, the derived IDs) and stored. Without Verify nothing is stored. +// // - Injection, as a context provider. It derives the task text from the invocation's messages (the contract's // rule), searches the store in the identity's scope, ranks with the contract's weights, renders the Historical // Reference block with BlockOptions, and adds it as one user-role message (MessageRole system sends it as a // system message) just before the new input, stamped AdditionalProperties[HistoricalReferenceKey] = true. +// // - Injection tracking (contract 1.4). Once the block is handed over, a store that records injections (an // ae.InjectionRecorder: the PostgreSQL store, the in-memory store) is told which records it delivered, in a // goroutine bounded by InjectionRecordTimeout, so the .NET lessons view shows when each lesson was last injected. // It never delays or fails the invocation; the outcome is reported through OnInjectionRecorded, and // ContextProviderOf(config).Close awaits the records still running, for a host shutting down. // +// - Semantic retrieval (contract 1.6), with Options.Embedder. Injection then runs the text and the vector channel +// together, merges them as the contract says and ranks as before; every stored record's summary is embedded in +// the background (a failure is reported as StageIndex and never fails the capture). Without an embedder, or when +// the vector channel fails, the injection answers from text alone and says why (OnVectorFallback). To register +// one, implement ae.Embedder (ModelID, Dimension, Embed) over your provider's client and set Options.Embedder; no +// embedding model is a dependency of this module. +// // Nothing here returns an error or panics into the agent: a store, identity, classifier or verifier failure means no // block, or no stored record, and is reported through OnInjectionFailure or OnCaptureFailure. Only the invocation's // own errors and cancellation, and the tool's own errors and panics, reach the caller, unchanged. A panic through the @@ -143,6 +152,9 @@ const ( StageVerify CaptureStage = "Verify" StageFinalize CaptureStage = "Finalize" StageStore CaptureStage = "Store" + // StageIndex is the background embedding write after a record was stored (contract 1.6): the record stays + // stored, only its vector is missing. + StageIndex CaptureStage = "Index" ) // CaptureFailure is a capture problem, reported once per stage per run. Reason never quotes record content; Err is @@ -163,11 +175,13 @@ type InjectionFailure struct { } // InjectionResult is what one injection delivered: the records in the block, best first, the block's size in UTF-8 -// bytes, and whether the search returned more candidates than the ranking considers. +// bytes, whether the search returned more candidates than the ranking considers, and, when the vector channel did not +// contribute (contract 1.6), why: nil when it did. type InjectionResult struct { - ExperienceIDs []ae.UUID - ByteCount int - Truncated bool + ExperienceIDs []ae.UUID + ByteCount int + Truncated bool + VectorFallback *ae.VectorFallback } // IdentityContext is what an identity resolver is given: the agent and the caller's session. Never model output. @@ -220,6 +234,25 @@ type Options struct { Clock func() ae.Instant // NewID is ae.NewRandomUUID when nil. NewID func() ae.UUID + // Embedder (contract 1.6) turns text into a vector, for hybrid retrieval and for embedding each stored record. + // With one, injection runs the text and the vector channel together (a store that is an ae.EmbeddingIndex serves + // it: the in-memory store, and the PostgreSQL store with pgvector), merges them and ranks as before, and every + // stored record's summary is embedded in the background. Without one, or when the vector channel fails, the + // injection answers from text alone and says why (OnVectorFallback, InjectionResult.VectorFallback). The + // embedder's calls honour ctx, are bounded by VectorTimeout, and at most MaxVectorCalls run at once, abandoned + // ones included. + Embedder ae.Embedder + // MinVectorRelevance, when set (0 to 1), drops vector candidates whose relevance is below it; text candidates + // never are. + MinVectorRelevance *float64 + // VectorTimeout bounds the vector channel of one injection in all (embedding the task text, then searching): + // past it the injection answers from text alone and reports ProviderUnavailable or VectorSearchFailed. + // DefaultVectorTimeout when 0. + VectorTimeout time.Duration + // OnVectorFallback is told, once per injection, why the vector channel did not contribute (NotConfigured, + // ProviderUnavailable, ModelMismatch, DimensionMismatch, VectorSearchFailed); optional, and a panic in it is + // recovered and ignored. Its detail never quotes record or vector content. + OnVectorFallback func(ae.VectorFallback) // FinalizationTimeout bounds how long an invocation's end waits for finalization (append, verify, finalize, // store), on its own clock, never the caller's context: past it the work is abandoned and reported, and the // invocation returns. DefaultFinalizationTimeout when 0. @@ -249,6 +282,10 @@ func DefaultEnvironment() ae.EnvironmentFingerprint { type settings struct { Options environment ae.EnvironmentFingerprint + // vectorSlots bound the embedder and vector-store calls in flight, abandoned ones included. + vectorSlots chan struct{} + // index counts the background embedding writes started and not yet ended. + index tracker } func newSettings(options Options) (*settings, error) { @@ -302,13 +339,22 @@ func newSettings(options Options) (*settings, error) { if options.InjectionRecordTimeout == 0 { options.InjectionRecordTimeout = DefaultInjectionRecordTimeout } + if options.VectorTimeout < 0 { + return nil, errors.New("maf: the vector timeout must not be negative") + } + if options.VectorTimeout == 0 { + options.VectorTimeout = DefaultVectorTimeout + } + if options.MinVectorRelevance != nil && !(*options.MinVectorRelevance >= 0 && *options.MinVectorRelevance <= 1) { + return nil, errors.New("maf: the minimum vector relevance must be a number from 0 to 1") + } if options.Clock == nil { options.Clock = ae.Now } if options.NewID == nil { options.NewID = ae.NewRandomUUID } - s := &settings{Options: options} + s := &settings{Options: options, vectorSlots: make(chan struct{}, MaxVectorCalls)} if options.Environment != nil { s.environment = *options.Environment } else { diff --git a/go/maf/record.go b/go/maf/record.go index 9feab036..ffde01d4 100644 --- a/go/maf/record.go +++ b/go/maf/record.go @@ -58,7 +58,8 @@ func ContextProviderOf(config *agent.Config) *ContextProvider { return nil } -// Close waits for every injection record the provider started to end, or for ctx to be done, whichever is first, +// Close waits for every injection record the provider started, and every background embedding write its +// configuration started, to end, or for ctx to be done, whichever is first, // and returns ctx's error in the second case. Each record is bounded by InjectionRecordTimeout, so a host shutting // down can call it with a context of a few seconds. A store write that overran that bound was abandoned when its // record ended and is not awaited: it may still be running (AbandonedInjectionRecords counts them). Close may run @@ -66,18 +67,20 @@ func ContextProviderOf(config *agent.Config) *ContextProvider { // be awaited. func (p *ContextProvider) Close(ctx context.Context) error { p.mu.Lock() + idle := p.idle if p.running == 0 { - p.mu.Unlock() - return nil + idle = nil } - idle := p.idle p.mu.Unlock() - select { - case <-idle: - return nil - case <-ctx.Done(): - return ctx.Err() + if idle != nil { + select { + case <-idle: + case <-ctx.Done(): + return ctx.Err() + } } + // The background embedding writes of the records stored (contract 1.6) too. + return p.settings.index.wait(ctx) } // beginRecording counts a record started. diff --git a/go/maf/vectors.go b/go/maf/vectors.go new file mode 100644 index 00000000..9737babd --- /dev/null +++ b/go/maf/vectors.go @@ -0,0 +1,294 @@ +package maf + +import ( + "context" + "errors" + "fmt" + "sync" + "time" + + "github.com/microsoft/agent-framework-go/agent" + + ae "github.com/fabbrik/AgentExperience.NET/go" +) + +// Semantic retrieval (contract 1.6): the vector channel of an injection, and the background embedding write after a +// record is stored. Nothing here fails the invocation or the capture. + +// DefaultVectorTimeout is how long the vector channel of one injection (embedding the task text, then searching the +// stored vectors) may take by default; past it the injection answers from text alone. +const DefaultVectorTimeout = 500 * time.Millisecond + +// MaxVectorCalls is how many embedder or vector-store calls may be in flight at once, per UseAgentExperience call, +// abandoned (timed-out) ones included: a goroutine cannot be stopped, so a provider that never returns holds its slot. +// Past it a call fails at once (ProviderUnavailable, VectorSearchFailed), so a hung provider can never starve the text +// search or pile up goroutines. +const MaxVectorCalls = 16 + +// errTooManyVectorCalls is the error of a call not started because MaxVectorCalls are already in flight. +var errTooManyVectorCalls = errors.New("too many earlier embedding or vector-search calls are still running") + +// callBounded runs fn in a goroutine of its own and waits for it, or for ctx to be done, whichever is first. It holds +// one of slots until fn returns (even when it was abandoned), and fails at once, running nothing, when none is free. +// A panic in fn is an error. +func callBounded[T any](ctx context.Context, slots chan struct{}, fn func(context.Context) (T, error)) (T, error) { + var zero T + select { + case slots <- struct{}{}: + default: + return zero, errTooManyVectorCalls + } + type outcome struct { + value T + err error + } + done := make(chan outcome, 1) + go func() { + defer func() { <-slots }() + defer func() { + if recovered := recover(); recovered != nil { + done <- outcome{err: panicError{recovered}} + } + }() + value, err := fn(ctx) + done <- outcome{value, err} + }() + select { + case result := <-done: + return result.value, result.err + case <-ctx.Done(): + select { + case result := <-done: + return result.value, result.err + default: + } + return zero, ctx.Err() + } +} + +// tracker counts background work started and not yet ended, so a host shutting down can wait for it. +type tracker struct { + mu sync.Mutex + running int + idle chan struct{} +} + +func (t *tracker) begin() { + t.mu.Lock() + defer t.mu.Unlock() + if t.running == 0 { + t.idle = make(chan struct{}) + } + t.running++ +} + +func (t *tracker) end() { + t.mu.Lock() + defer t.mu.Unlock() + t.running-- + if t.running == 0 { + close(t.idle) + } +} + +// wait returns when nothing is running, or with ctx's error when it is done first. +func (t *tracker) wait(ctx context.Context) error { + t.mu.Lock() + if t.running == 0 { + t.mu.Unlock() + return nil + } + idle := t.idle + t.mu.Unlock() + select { + case <-idle: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +// WaitForIndexWrites waits for every background embedding write the configuration's agent started (after a record was +// stored) to end, or for ctx to be done, whichever is first, and returns ctx's error in the second case. Each write is +// bounded by FinalizationTimeout. It does nothing for a configuration UseAgentExperience was not called on. +func WaitForIndexWrites(ctx context.Context, config *agent.Config) error { + if config == nil { + return nil + } + for _, middleware := range config.Middlewares { + if capture, ok := middleware.(*captureMiddleware); ok { + return capture.settings.index.wait(ctx) + } + } + return nil +} + +func (s *settings) reportVectorFallback(fallback *ae.VectorFallback) { + if fallback == nil || s.OnVectorFallback == nil { + return + } + defer func() { _ = recover() }() + s.OnVectorFallback(*fallback) +} + +// embedded is a text's vector, and the model that made it. +type embedded struct { + modelID string + vector []float32 +} + +// embed asks the embedder for text's vector and checks what it answered: a non-blank model ID, and a vector of +// exactly its dimension and finite. It runs inside a bounded call. +func (s *settings) embed(ctx context.Context, text string) (embedded, error) { + modelID, dimension := s.Embedder.ModelID(), s.Embedder.Dimension() + if err := ae.ValidateModelID(modelID); err != nil { + return embedded{}, err + } + if dimension < 1 || dimension > ae.MaxDimension { + return embedded{}, fmt.Errorf("the embedder declares a dimension outside 1 to %d", ae.MaxDimension) + } + vector, err := s.Embedder.Embed(ctx, text) + if err != nil { + return embedded{}, err + } + if err := ae.ValidateVector(vector, dimension); err != nil { + return embedded{}, err + } + return embedded{modelID, vector}, nil +} + +// indexOf is the store's embedding index, when it is one and an embedder is registered. +func (s *settings) indexOf() (ae.EmbeddingIndex, *ae.VectorFallback) { + if s.Embedder == nil { + return nil, &ae.VectorFallback{Reason: ae.FallbackNotConfigured, + Detail: "No embedder is registered, so only text retrieval ran."} + } + index, ok := s.Store.(ae.EmbeddingIndex) + if !ok { + return nil, &ae.VectorFallback{Reason: ae.FallbackNotConfigured, + Detail: "The store does not serve vector search, so only text retrieval ran."} + } + return index, nil +} + +// vectorChannel is the vector channel (contract 1.6): embed the task text, search the stored vectors, both within +// VectorTimeout. It never panics and never fails the invocation: whatever goes wrong is the fallback it answers with +// (and no rows), and the fallback's text never quotes record or vector content. +func (s *settings) vectorChannel(ctx context.Context, identity *ae.Identity, taskText string) (rows []ae.SearchRow, + fallback *ae.VectorFallback) { + defer func() { + if recovered := recover(); recovered != nil { + rows, fallback = nil, &ae.VectorFallback{Reason: ae.FallbackVectorSearchFailed, + Detail: "The vector channel failed unexpectedly.", Err: panicError{recovered}} + } + }() + index, fallback := s.indexOf() + if fallback != nil { + return nil, fallback + } + ctx, cancel := context.WithTimeout(ctx, s.VectorTimeout) + defer cancel() + + query, err := callBounded(ctx, s.vectorSlots, func(ctx context.Context) (ae.VectorQuery, error) { + e, err := s.embed(ctx, taskText) + if err != nil { + return ae.VectorQuery{}, err + } + query := ae.VectorQuery{Scope: identity.Scope, ModelID: e.modelID, Vector: e.vector, + ExcludeModelAuthored: s.ExcludeModelAuthored} + return query, query.Validate() + }) + if err != nil { + if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) { + return nil, &ae.VectorFallback{Reason: ae.FallbackProviderUnavailable, + Detail: "The embedder did not answer in time.", Err: err} + } + return nil, &ae.VectorFallback{Reason: ae.FallbackProviderUnavailable, + Detail: fmt.Sprintf("The embedder failed or returned an unusable vector (%s).", errorName(err)), Err: err} + } + + result, err := callBounded(ctx, s.vectorSlots, func(ctx context.Context) (ae.VectorSearchResult, error) { + return index.SearchVectors(ctx, identity.Authorization(), query) + }) + if err != nil { + if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) { + return nil, &ae.VectorFallback{Reason: ae.FallbackVectorSearchFailed, + Detail: "The vector search did not answer in time.", Err: err} + } + return nil, &ae.VectorFallback{Reason: ae.FallbackVectorSearchFailed, + Detail: fmt.Sprintf("The vector search failed (%s).", errorName(err)), Err: err} + } + switch result.Outcome { + case ae.VectorFound: + case ae.VectorModelMismatch: + return nil, &ae.VectorFallback{Reason: ae.FallbackModelMismatch, + Detail: "Every stored embedding in this scope is of another model."} + case ae.VectorDimensionMismatch: + return nil, &ae.VectorFallback{Reason: ae.FallbackDimensionMismatch, + Detail: "Every stored embedding in this scope is of another width."} + default: + return nil, &ae.VectorFallback{Reason: ae.FallbackVectorSearchFailed, + Detail: "The vector search answered with an unknown outcome."} + } + rows = make([]ae.SearchRow, len(result.Hits)) + for i, hit := range result.Hits { + rows[i] = ae.SearchRow{Record: hit.Record, Relevance: hit.Relevance()} + } + return rows, nil +} + +// startIndex starts embedding a stored record's summary in the background: derived data, written after the record's +// own commit, never inside it, and never failing it. It never blocks and never panics. +func (c *captureScope) startIndex(ctx context.Context, record ae.ExperienceRecord) { + s := c.settings + index, fallback := s.indexOf() + if fallback != nil { + return + } + // The invocation's values, never its cancellation: the invocation ends long before the write may. + base := context.WithoutCancel(ctx) + s.index.begin() + go func() { + defer s.index.end() + c.writeEmbedding(base, index, &record) + }() +} + +// writeEmbedding embeds record's summary and stores the vector, all within FinalizationTimeout. A failure is reported +// as StageIndex (once per run); a Stale or Missing write is not one: the record moved on, or is gone. +func (c *captureScope) writeEmbedding(ctx context.Context, index ae.EmbeddingIndex, record *ae.ExperienceRecord) { + s := c.settings + defer func() { + if recovered := recover(); recovered != nil { + c.report(StageIndex, "Embedding the record panicked; its text retrieval is unaffected.", + panicError{recovered}) + } + }() + ctx, cancel := context.WithTimeout(ctx, s.FinalizationTimeout) + defer cancel() + fail := func(err error) { + if errors.Is(err, context.DeadlineExceeded) { + c.report(StageIndex, fmt.Sprintf("Embedding the record did not finish within %v.", s.FinalizationTimeout), err) + return + } + c.report(StageIndex, fmt.Sprintf("Embedding the record failed (%s); its text retrieval is unaffected.", + errorName(err)), err) + } + summary := ae.RecordSummary(record) + e, err := callBounded(ctx, s.vectorSlots, func(ctx context.Context) (embedded, error) { + return s.embed(ctx, summary) + }) + if err != nil { + fail(err) + return + } + write := ae.EmbeddingWrite{ExperienceID: record.ExperienceID, Scope: record.Scope, ModelID: e.modelID, + Dimension: len(e.vector), ContentHash: ae.ContentHash(e.modelID, summary), SourceRevision: record.Revision, + Vector: e.vector, Now: s.Clock()} + _, err = callBounded(ctx, s.vectorSlots, func(ctx context.Context) (ae.EmbeddingWriteResult, error) { + return index.WriteEmbedding(ctx, c.identity.Authorization(), write) + }) + if err != nil { + fail(err) + } +} diff --git a/go/maf/vectors_postgres_test.go b/go/maf/vectors_postgres_test.go new file mode 100644 index 00000000..76959f27 --- /dev/null +++ b/go/maf/vectors_postgres_test.go @@ -0,0 +1,173 @@ +//go:build postgres + +package maf + +// Story 23.3's cross-language proof on PostgreSQL with pgvector (contract 1.6): +// +// - .NET writes, Go reads: the records of the vector-search/nearest-first fixture are embedded by the contract's own +// write statement (write-experience-embedding.sql, run as written, not by the Go store), and a Go agent whose +// embedder returns the fixture's query vector finds them; without the embedder it finds nothing and says why. +// - Go writes, Go reads: a Go agent stores a lesson, its embedding is written in the background through the Go store, +// and a later reworded query finds it. + +import ( + "context" + "path/filepath" + "testing" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" + + ae "github.com/fabbrik/AgentExperience.NET/go" + ct "github.com/fabbrik/AgentExperience.NET/go/internal/contracttest" + "github.com/fabbrik/AgentExperience.NET/go/postgres" + "github.com/fabbrik/AgentExperience.NET/go/postgres/pgtest" +) + +func TestARecordEmbeddedByTheContractsSQLIsFoundByAGoAgent(t *testing.T) { + admin := pgtest.Require(t) + ctx := context.Background() + database := createDatabase(t, admin, "vector_nearest_first") + store := newStore(t, admin, database) + for _, record := range seedRecords(t, "vector-search", "nearest-first") { + record := record + if err := store.InsertRecord(ctx, ae.BoundsOf(record.Scope), &record); err != nil { + t.Fatal(err) + } + } + + // The embeddings are written by the fixture's own statement, not by the Go store. + pool, err := pgxpool.New(ctx, admin.Application(database)) + if err != nil { + t.Fatal(err) + } + defer pool.Close() + statements, err := postgres.ContractStatements(postgres.WriteEmbeddingSQL) + if err != nil { + t.Fatal(err) + } + write := postgres.ToDriverPlaceholders(statements["write"]) + declareSQL, err := postgres.ContractSQL(postgres.DeclareBoundsSQL) + if err != nil { + t.Fatal(err) + } + declare := postgres.ToDriverPlaceholders(declareSQL) + get := func(o *ae.Object, key string) any { value, _ := o.Get(key); return value } + steps := get(ct.MustReadJSON(filepath.Join(ct.Fixtures, "vector-search", "nearest-first", "steps.json")).(*ae.Object), + "steps").([]any) + var query ae.VectorQuery + written := 0 + for _, item := range steps { + step := item.(*ae.Object) + if search := get(step, "search"); search != nil { + query = ct.VectorQueryOf(search) + continue + } + w := ct.EmbeddingWriteOf(get(step, "write")) + declared := map[string]any{} + for name, value := range ae.BoundsOf(w.Scope).Parameters() { + declared[name] = value + } + declareArgs, err := declare.Args(declared) + if err != nil { + t.Fatal(err) + } + args, err := write.Args(map[string]any{ + "experience_id": [16]byte(w.ExperienceID), "tenant_id": w.Scope.TenantID, + "application_id": w.Scope.ApplicationID, "project_id": w.Scope.ProjectID, "team_id": w.Scope.TeamID, + "agent_id": w.Scope.AgentID, "user_id": w.Scope.UserID, "model_id": w.ModelID, + "dimension": int32(w.Dimension), "content_hash": w.ContentHash, "source_revision": w.SourceRevision, + "embedding": ae.VectorLiteral(w.Vector), "now": w.Now.Time(), + }) + if err != nil { + t.Fatal(err) + } + err = pgx.BeginFunc(ctx, pool, func(tx pgx.Tx) error { + if _, err := tx.Exec(ctx, declare.SQL, declareArgs...); err != nil { + return err + } + tag, err := tx.Exec(ctx, write.SQL, args...) + if err == nil && tag.RowsAffected() != 1 { + t.Fatalf("the contract's write wrote %d rows", tag.RowsAffected()) + } + return err + }) + if err != nil { + t.Fatal(err) + } + written++ + } + if written != 3 || query.ModelID == "" { + t.Fatalf("%d writes, query %+v", written, query) + } + + embedder := newFixed() + embedder.vector, embedder.model = query.Vector, query.ModelID + v := &vhost{host: newHost(t)} + a, _, _ := wire(t, v, store, Options{Embedder: embedder}) + runText(t, a, paraphrase) + if len(v.reasons()) != 0 || len(v.injectionFailures) != 0 { + t.Fatalf("fallbacks %v, failures %+v", v.reasons(), v.injectionFailures) + } + if got := numbers(onlyInjection(t, v)); !equalInts(got, []int{3, 2, 1}) { + t.Fatalf("%v", got) + } + + // Without the embedder the same query finds nothing, and says why. + plain := &vhost{host: newHost(t)} + b, _, _ := wire(t, plain, store, Options{}) + runText(t, b, paraphrase) + if reasons := plain.reasons(); len(plain.injections) != 0 || len(reasons) != 1 || reasons[0] != ae.FallbackNotConfigured { + t.Fatalf("injections %+v, fallbacks %v", plain.injections, reasons) + } + + // A model the rows were not embedded with is a ModelMismatch, not a failure. + other := newFixed() + other.vector, other.model = query.Vector, "another-model" + mismatch := &vhost{host: newHost(t)} + c, _, _ := wire(t, mismatch, store, Options{Embedder: other}) + runText(t, c, paraphrase) + if reasons := mismatch.reasons(); len(mismatch.injections) != 0 || len(reasons) != 1 || reasons[0] != ae.FallbackModelMismatch { + t.Fatalf("injections %+v, fallbacks %v", mismatch.injections, reasons) + } +} + +func TestAGoAgentStoresARecordAndItsEmbeddingInPostgreSQL(t *testing.T) { + admin := pgtest.Require(t) + ctx := context.Background() + database := createDatabase(t, admin, "vector_roundtrip") + store := newStore(t, admin, database) + v := &vhost{host: newHost(t)} + a, config := storing(t, v, store, Options{Embedder: newFixed(), + Identity: ae.Identity{Principal: "go-principal", Scope: testScope}}) + runText(t, a, ticket) + waitForIndex(t, config) + if len(v.captureFailures) != 0 || len(v.records()) != 1 { + t.Fatalf("failures %+v, %d records", v.captureFailures, len(v.records())) + } + + conn, err := pgx.Connect(ctx, admin.Admin(database)) + if err != nil { + t.Fatal(err) + } + defer conn.Close(ctx) + var model, embedding, hash string + var dimension, revision int + if err := conn.QueryRow(ctx, "SELECT model_id, dimension, source_revision, embedding::text, content_hash "+ + "FROM agent_experience.experience_embeddings").Scan(&model, &dimension, &revision, &embedding, &hash); err != nil { + t.Fatal(err) + } + record := v.records()[0] + if model != vectorModel || dimension != 4 || int64(revision) != record.Revision || embedding != "[1,0,0,0]" || + hash != ae.ContentHash(vectorModel, ae.RecordSummary(&record)) { + t.Fatalf("%s %d %d %s %s", model, dimension, revision, embedding, hash) + } + + later := &vhost{host: newHost(t)} + b, _, _ := wire(t, later, store, Options{Embedder: newFixed(), + Identity: ae.Identity{Principal: "go-principal", Scope: testScope}}) + runText(t, b, "Zebra giraffe antelope") + if len(later.injections) != 1 || len(later.reasons()) != 0 { + t.Fatalf("injections %+v, fallbacks %v", later.injections, later.reasons()) + } +} diff --git a/go/maf/vectors_test.go b/go/maf/vectors_test.go new file mode 100644 index 00000000..0c04ae4d --- /dev/null +++ b/go/maf/vectors_test.go @@ -0,0 +1,705 @@ +package maf + +// Story 23.3: semantic retrieval in Go (contract 1.6), through a real agent.Agent and the in-memory store. The +// Python port's cases (python/tests/test_vectors.py): every row of the fallback matrix, the floor, the bounded calls, +// and the background embedding write after a record is stored. The PostgreSQL cases are in vectors_postgres_test.go. + +import ( + "context" + "errors" + "math" + "path/filepath" + "strings" + "sync" + "testing" + "time" + + "github.com/microsoft/agent-framework-go/agent" + "github.com/microsoft/agent-framework-go/tool" + + ae "github.com/fabbrik/AgentExperience.NET/go" + "github.com/fabbrik/AgentExperience.NET/go/inmemory" + ct "github.com/fabbrik/AgentExperience.NET/go/internal/contracttest" +) + +const ( + vectorModel = "contract-embedding" + paraphrase = "Renew the cryptographic credentials immediately" + textQuery = "Rotate the signing keys for the service" +) + +var vectorScope = ae.Scope{TenantID: "contract-tenant", ApplicationID: "contract-app", ProjectID: "contract-project", + AgentID: ae.Ptr("refund-agent")} + +// ---- The caller's embedder, and stores that misbehave ---------------------------------------------------------- + +// fixedEmbedder answers one vector for every text, unless it is told to misbehave. +type fixedEmbedder struct { + vector []float32 + model string + dimension int // the declared dimension; len(vector) when 0 + fail bool // returns an error whose message holds record content + panics bool // panics + delay time.Duration // sleeps, ignoring its context + block chan struct{} // blocks until closed, ignoring its context + honour bool // blocks until its context is done, then returns its error + + mu sync.Mutex + texts []string + calls int + cancelled chan struct{} +} + +func newFixed() *fixedEmbedder { + return &fixedEmbedder{vector: []float32{1, 0, 0, 0}, model: vectorModel, cancelled: make(chan struct{}, 64)} +} + +func (e *fixedEmbedder) ModelID() string { return e.model } + +func (e *fixedEmbedder) Dimension() int { + if e.dimension != 0 { + return e.dimension + } + return len(e.vector) +} + +func (e *fixedEmbedder) Embed(ctx context.Context, text string) ([]float32, error) { + e.mu.Lock() + e.texts = append(e.texts, text) + e.calls++ + e.mu.Unlock() + switch { + case e.panics: + panic("secret-record-content") + case e.fail: + return nil, errors.New("the provider is down: secret-record-content") + case e.block != nil: + <-e.block + case e.honour: + <-ctx.Done() + e.cancelled <- struct{}{} + return nil, ctx.Err() + case e.delay > 0: + time.Sleep(e.delay) + } + return append([]float32(nil), e.vector...), nil +} + +func (e *fixedEmbedder) seen() []string { + e.mu.Lock() + defer e.mu.Unlock() + return append([]string(nil), e.texts...) +} + +func (e *fixedEmbedder) called() int { + e.mu.Lock() + defer e.mu.Unlock() + return e.calls +} + +// vectorSpy is an in-memory store that counts the vector searches it is asked to run, and can be made to fail. +type vectorSpy struct { + *inmemory.Store + mu sync.Mutex + searches int + searchErr error + searchDelay time.Duration + writeErr error +} + +func (s *vectorSpy) SearchVectors(ctx context.Context, bounds ae.AuthorizationBounds, + query ae.VectorQuery) (ae.VectorSearchResult, error) { + s.mu.Lock() + s.searches++ + s.mu.Unlock() + if s.searchDelay > 0 { + time.Sleep(s.searchDelay) + } + if s.searchErr != nil { + return ae.VectorSearchResult{}, s.searchErr + } + return s.Store.SearchVectors(ctx, bounds, query) +} + +func (s *vectorSpy) WriteEmbedding(ctx context.Context, bounds ae.AuthorizationBounds, + write ae.EmbeddingWrite) (ae.EmbeddingWriteResult, error) { + if s.writeErr != nil { + return ae.EmbeddingWriteResult{}, s.writeErr + } + return s.Store.WriteEmbedding(ctx, bounds, write) +} + +func (s *vectorSpy) searchCount() int { + s.mu.Lock() + defer s.mu.Unlock() + return s.searches +} + +// textOnlyStore has the text methods only: no vector methods. +type textOnlyStore struct{ inner *inmemory.Store } + +func (s textOnlyStore) InsertFinalized(ctx context.Context, bounds ae.AuthorizationBounds, f *ae.FinalizedExperience, + actor string, recordedAt *time.Time) error { + return s.inner.InsertFinalized(ctx, bounds, f, actor, recordedAt) +} + +func (s textOnlyStore) Search(ctx context.Context, bounds ae.AuthorizationBounds, + p ae.SearchParameters) ([]ae.SearchRow, error) { + return s.inner.Search(ctx, bounds, p) +} + +// failingText has a text search that fails, and a vector channel that works. +type failingText struct{ *inmemory.Store } + +func (failingText) Search(context.Context, ae.AuthorizationBounds, ae.SearchParameters) ([]ae.SearchRow, error) { + return nil, errors.New("the text index is down") +} + +// ---- The host --------------------------------------------------------------------------------------------------- + +type vhost struct { + *host + fmu sync.Mutex + fallbacks []ae.VectorFallback +} + +func (v *vhost) options(o Options) Options { + if o.Identity.Principal == "" && o.ResolveIdentity == nil { + o.Identity = ae.Identity{Principal: "go-principal", Scope: vectorScope} + } + o = v.host.options(o) + o.OnVectorFallback = func(f ae.VectorFallback) { v.fmu.Lock(); v.fallbacks = append(v.fallbacks, f); v.fmu.Unlock() } + return o +} + +func (v *vhost) reasons() []ae.VectorFallbackReason { + v.fmu.Lock() + defer v.fmu.Unlock() + var reasons []ae.VectorFallbackReason + for _, f := range v.fallbacks { + reasons = append(reasons, f.Reason) + } + return reasons +} + +func seedRecords(t *testing.T, family, name string) []ae.ExperienceRecord { + t.Helper() + var records []ae.ExperienceRecord + for _, stored := range ct.MustReadJSON(filepath.Join(ct.Fixtures, family, name, "seed.json")).([]any) { + record, err := ae.FromStored(stored) + if err != nil { + t.Fatal(err) + } + records = append(records, record) + } + return records +} + +// embeddedStore holds the three 'rotate-keys' records of the nearest-first fixture, embedded as it does: 3 is the +// query's own vector, 2 is near, 1 is further. +func embeddedStore(t *testing.T) *inmemory.Store { + t.Helper() + records := seedRecords(t, "vector-search", "nearest-first") + store, err := inmemory.New(records...) + if err != nil { + t.Fatal(err) + } + vectors := map[byte][]float32{1: {1, 1, 0, 0}, 2: {1, 0.1, 0, 0}, 3: {1, 0, 0, 0}} + for _, record := range records { + result, err := store.WriteEmbedding(context.Background(), ae.BoundsOf(vectorScope), ae.EmbeddingWrite{ + ExperienceID: record.ExperienceID, Scope: vectorScope, ModelID: vectorModel, Dimension: 4, + ContentHash: "h", SourceRevision: record.Revision, Vector: vectors[record.ExperienceID[15]], + Now: ae.MustParseInstant("2026-03-02T10:00:00+00:00")}) + if err != nil || result.Outcome != ae.EmbeddingWritten { + t.Fatalf("%+v, %v", result, err) + } + } + return store +} + +func wire(t *testing.T, v *vhost, store ae.Store, o Options) (*agent.Agent, *agent.Config, *scripted) { + t.Helper() + o.Store = store + model := newScripted(0) + config := agent.Config{Name: "refund-agent"} + if err := UseAgentExperience(&config, v.options(o)); err != nil { + t.Fatal(err) + } + return newAgent(model, config), &config, model +} + +func numbers(result InjectionResult) []int { + numbers := make([]int, len(result.ExperienceIDs)) + for i, id := range result.ExperienceIDs { + numbers[i] = int(id[15]) + } + return numbers +} + +func sorted(numbers []int) []int { + out := append([]int(nil), numbers...) + for i := range out { + for j := i + 1; j < len(out); j++ { + if out[j] < out[i] { + out[i], out[j] = out[j], out[i] + } + } + } + return out +} + +func equalInts(a, b []int) bool { + if len(a) != len(b) { + return false + } + for i := range a { + if a[i] != b[i] { + return false + } + } + return true +} + +func onlyInjection(t *testing.T, v *vhost) InjectionResult { + t.Helper() + if len(v.injections) != 1 { + t.Fatalf("%d injections, failures %+v", len(v.injections), v.injectionFailures) + } + return v.injections[0] +} + +// ---- Hybrid retrieval ------------------------------------------------------------------------------------------- + +func TestOnlyTheVectorFindsARewordedQueryAndWithoutAnEmbedderItDoesNot(t *testing.T) { + v := &vhost{host: newHost(t)} + a, _, model := wire(t, v, embeddedStore(t), Options{Embedder: newFixed()}) + runText(t, a, paraphrase) + if len(v.reasons()) != 0 || len(v.injectionFailures) != 0 { + t.Fatalf("fallbacks %v, failures %+v", v.fallbacks, v.injectionFailures) + } + result := onlyInjection(t, v) + if !equalInts(numbers(result), []int{3, 2, 1}) || result.VectorFallback != nil { + t.Fatalf("%v, fallback %+v", numbers(result), result.VectorFallback) + } + blocks := model.blocks() + if len(blocks) != 1 || !strings.Contains(blocks[0].String(), "--- RECORD 1: rotate-keys ---") { + t.Fatalf("%d blocks", len(blocks)) + } + + without := &vhost{host: newHost(t)} + a, _, model = wire(t, without, embeddedStore(t), Options{}) + runText(t, a, paraphrase) + if len(without.injections) != 0 || len(model.blocks()) != 0 { + t.Fatalf("injections %+v", without.injections) + } + if reasons := without.reasons(); len(reasons) != 1 || reasons[0] != ae.FallbackNotConfigured { + t.Fatalf("fallbacks %v", reasons) + } +} + +func TestATextMatchAndAVectorMatchAreOneEntry(t *testing.T) { + v := &vhost{host: newHost(t)} + a, _, _ := wire(t, v, embeddedStore(t), Options{Embedder: newFixed()}) + runText(t, a, textQuery) + if got := sorted(numbers(onlyInjection(t, v))); !equalInts(got, []int{1, 2, 3}) { + t.Fatalf("%v", got) + } +} + +func TestAStoreWithoutVectorMethodsAnswersFromTextAndReportsNotConfigured(t *testing.T) { + v := &vhost{host: newHost(t)} + embedder := newFixed() + a, _, _ := wire(t, v, textOnlyStore{embeddedStore(t)}, Options{Embedder: embedder}) + runText(t, a, textQuery) + if reasons := v.reasons(); len(reasons) != 1 || reasons[0] != ae.FallbackNotConfigured { + t.Fatalf("fallbacks %v", reasons) + } + result := onlyInjection(t, v) + if got := sorted(numbers(result)); !equalInts(got, []int{1, 2, 3}) { + t.Fatalf("%v", got) + } + if result.VectorFallback == nil || result.VectorFallback.Reason != ae.FallbackNotConfigured || embedder.called() != 0 { + t.Fatalf("%+v, %d embeds", result.VectorFallback, embedder.called()) + } +} + +func TestABadEmbedderAnswersFromTextAndNoStatementRuns(t *testing.T) { + for name, configure := range map[string]func(*fixedEmbedder){ + "raises": func(e *fixedEmbedder) { e.fail = true }, + "panics": func(e *fixedEmbedder) { e.panics = true }, + "times-out": func(e *fixedEmbedder) { e.delay = 400 * time.Millisecond }, + "honours-ctx": func(e *fixedEmbedder) { e.honour = true }, + "blank-model": func(e *fixedEmbedder) { e.model = " " }, + "empty": func(e *fixedEmbedder) { e.vector = nil }, + "wrong-length": func(e *fixedEmbedder) { e.dimension = 3 }, + "bad-dimension": func(e *fixedEmbedder) { e.dimension = -1 }, + "nan": func(e *fixedEmbedder) { e.vector = []float32{1, float32(math.NaN()), 0, 0} }, + "infinite": func(e *fixedEmbedder) { e.vector = []float32{1, float32(math.Inf(1)), 0, 0} }, + "long-model-id": func(e *fixedEmbedder) { e.model = strings.Repeat("m", 300) }, + } { + t.Run(name, func(t *testing.T) { + embedder := newFixed() + configure(embedder) + store := &vectorSpy{Store: embeddedStore(t)} + v := &vhost{host: newHost(t)} + a, _, _ := wire(t, v, store, Options{Embedder: embedder, VectorTimeout: 100 * time.Millisecond}) + start := time.Now() + runText(t, a, textQuery) + if elapsed := time.Since(start); elapsed > 350*time.Millisecond { + t.Fatalf("the vector channel held the injection for %v", elapsed) + } + if reasons := v.reasons(); len(reasons) != 1 || reasons[0] != ae.FallbackProviderUnavailable { + t.Fatalf("fallbacks %v", reasons) + } + if strings.Contains(v.fallbacks[0].Detail, "secret-record-content") { + t.Fatalf("the detail quotes content: %q", v.fallbacks[0].Detail) + } + if store.searchCount() != 0 || len(v.injectionFailures) != 0 { + t.Fatalf("%d searches, failures %+v", store.searchCount(), v.injectionFailures) + } + if got := sorted(numbers(onlyInjection(t, v))); !equalInts(got, []int{1, 2, 3}) { + t.Fatalf("%v", got) + } + }) + } +} + +func TestAnEmbedderThatHonoursItsContextIsCancelledWhenTheTimeoutPasses(t *testing.T) { + embedder := newFixed() + embedder.honour = true + v := &vhost{host: newHost(t)} + a, _, _ := wire(t, v, embeddedStore(t), Options{Embedder: embedder, VectorTimeout: 50 * time.Millisecond}) + runText(t, a, textQuery) + select { + case <-embedder.cancelled: + case <-time.After(2 * time.Second): + t.Fatal("the embedder's context was never cancelled") + } +} + +func TestAMismatchAnswersFromTextAndSaysWhich(t *testing.T) { + for name, tc := range map[string]struct { + configure func(*fixedEmbedder) + reason ae.VectorFallbackReason + }{ + "model": {func(e *fixedEmbedder) { e.model = "another-model" }, ae.FallbackModelMismatch}, + "dimension": {func(e *fixedEmbedder) { e.vector = []float32{1, 0, 0} }, ae.FallbackDimensionMismatch}, + } { + t.Run(name, func(t *testing.T) { + embedder := newFixed() + tc.configure(embedder) + v := &vhost{host: newHost(t)} + a, _, _ := wire(t, v, embeddedStore(t), Options{Embedder: embedder}) + runText(t, a, paraphrase) + if reasons := v.reasons(); len(reasons) != 1 || reasons[0] != tc.reason { + t.Fatalf("fallbacks %v", reasons) + } + if len(v.injections) != 0 { + t.Fatalf("injections %+v", v.injections) + } + }) + } +} + +func TestAFailingSearchAnswersFromTextAndReportsIt(t *testing.T) { + for name, store := range map[string]*vectorSpy{ + "raises": {searchErr: errors.New("the index is down")}, + "times-out": {searchDelay: 500 * time.Millisecond}, + } { + t.Run(name, func(t *testing.T) { + store.Store = embeddedStore(t) + v := &vhost{host: newHost(t)} + a, _, _ := wire(t, v, store, Options{Embedder: newFixed(), VectorTimeout: 100 * time.Millisecond}) + start := time.Now() + runText(t, a, textQuery) + if elapsed := time.Since(start); elapsed > 400*time.Millisecond { + t.Fatalf("the vector channel held the injection for %v", elapsed) + } + if reasons := v.reasons(); len(reasons) != 1 || reasons[0] != ae.FallbackVectorSearchFailed { + t.Fatalf("fallbacks %v", reasons) + } + if len(v.injectionFailures) != 0 { + t.Fatalf("failures %+v", v.injectionFailures) + } + if got := sorted(numbers(onlyInjection(t, v))); !equalInts(got, []int{1, 2, 3}) { + t.Fatalf("%v", got) + } + }) + } +} + +func TestAFailingTextSearchFailsTheInjectionWhateverTheVectorFound(t *testing.T) { + v := &vhost{host: newHost(t)} + a, _, model := wire(t, v, failingText{embeddedStore(t)}, Options{Embedder: newFixed()}) + runText(t, a, paraphrase) + if len(v.injectionFailures) != 1 || v.injectionFailures[0].Reason != "The search failed; no block was added." { + t.Fatalf("failures %+v", v.injectionFailures) + } + if len(v.injections) != 0 || len(model.blocks()) != 0 { + t.Fatalf("injections %+v", v.injections) + } +} + +func TestTheFloorDropsWeakVectorCandidatesAndNeverATextOne(t *testing.T) { + floor := 0.9 + v := &vhost{host: newHost(t)} + a, _, _ := wire(t, v, embeddedStore(t), Options{Embedder: newFixed(), MinVectorRelevance: &floor}) + runText(t, a, paraphrase) + if got := numbers(onlyInjection(t, v)); !equalInts(got, []int{3, 2}) { // record 1's relevance is about 0.85 + t.Fatalf("%v", got) + } + + // A text match is kept although its vector is below the floor. + one := 1.0 + far := newFixed() + far.vector = []float32{0, 0, 0, 1} + v = &vhost{host: newHost(t)} + a, _, _ = wire(t, v, embeddedStore(t), Options{Embedder: far, MinVectorRelevance: &one}) + runText(t, a, textQuery) + if got := sorted(numbers(onlyInjection(t, v))); !equalInts(got, []int{1, 2, 3}) { + t.Fatalf("%v", got) + } +} + +func TestOptionsRefuseABadFloorOrTimeout(t *testing.T) { + store, _ := inmemory.New() + base := Options{Store: store, Identity: ae.Identity{Principal: "p", Scope: vectorScope}} + for name, mutate := range map[string]func(*Options){ + "floor above one": func(o *Options) { o.MinVectorRelevance = floatPtr(1.5) }, + "floor below zero": func(o *Options) { o.MinVectorRelevance = floatPtr(-0.1) }, + "floor NaN": func(o *Options) { o.MinVectorRelevance = floatPtr(math.NaN()) }, + "negative timeout": func(o *Options) { o.VectorTimeout = -time.Second }, + } { + options := base + mutate(&options) + if _, err := newSettings(options); err == nil { + t.Errorf("%s was accepted", name) + } + } + if s, err := newSettings(base); err != nil || s.VectorTimeout != DefaultVectorTimeout { + t.Fatalf("%+v, %v", s, err) + } +} + +// ---- Bounded calls ------------------------------------------------------------------------------------------------ + +func TestHungVectorCallsFailFastAndNeverStarveTheTextSearch(t *testing.T) { + release := make(chan struct{}) + embedder := newFixed() + embedder.block = release + v := &vhost{host: newHost(t)} + a, config, _ := wire(t, v, embeddedStore(t), Options{Embedder: embedder, VectorTimeout: 10 * time.Millisecond}) + slots := settingsOf(config).vectorSlots + + start := time.Now() + for range MaxVectorCalls + 4 { + runText(t, a, textQuery) + } + if elapsed := time.Since(start); elapsed > 5*time.Second { + t.Fatalf("the injections took %v", elapsed) + } + // Every injection was answered from text, and only MaxVectorCalls embedder calls ever started. + if len(v.injections) != MaxVectorCalls+4 || embedder.called() != MaxVectorCalls || len(slots) != MaxVectorCalls { + t.Fatalf("%d injections, %d embedder calls, %d slots held", len(v.injections), embedder.called(), len(slots)) + } + for _, f := range v.fallbacks { + if f.Reason != ae.FallbackProviderUnavailable { + t.Fatalf("fallback %+v", f) + } + } + if last := v.fallbacks[len(v.fallbacks)-1]; !errors.Is(last.Err, errTooManyVectorCalls) { + t.Fatalf("the call past the limit failed with %v", last.Err) + } + + close(release) + deadline := time.Now().Add(5 * time.Second) + for len(slots) != 0 && time.Now().Before(deadline) { + time.Sleep(5 * time.Millisecond) + } + if len(slots) != 0 { + t.Fatalf("%d slots are still held after the calls returned", len(slots)) + } +} + +func TestCallBoundedHoldsASlotUntilTheCallReturnsAndRecoversAPanic(t *testing.T) { + slots := make(chan struct{}, 2) + release := make(chan struct{}) + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) + defer cancel() + for range 2 { + _, err := callBounded(ctx, slots, func(context.Context) (int, error) { <-release; return 1, nil }) + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("a hung call answered %v", err) + } + } + ran := false + _, err := callBounded(context.Background(), slots, func(context.Context) (int, error) { ran = true; return 1, nil }) + if !errors.Is(err, errTooManyVectorCalls) || ran { + t.Fatalf("a call past the limit ran (%v, %v)", ran, err) + } + close(release) + for len(slots) != 0 { + time.Sleep(time.Millisecond) + } + if _, err := callBounded(context.Background(), slots, func(context.Context) (int, error) { panic("boom") }); err == nil || + errorName(err) != "github.com/fabbrik/AgentExperience.NET/go/maf.panicError" { + t.Fatalf("a panic answered %v", err) + } + if value, err := callBounded(context.Background(), slots, func(context.Context) (int, error) { return 7, nil }); err != nil || + value != 7 { + t.Fatalf("%v, %v", value, err) + } +} + +// ---- The background write after a record is stored -------------------------------------------------------------- + +func settingsOf(config *agent.Config) *settings { + for _, middleware := range config.Middlewares { + if capture, ok := middleware.(*captureMiddleware); ok { + return capture.settings + } + } + return nil +} + +// storing is an agent that stores a lesson from its one passing tool call. +func storing(t *testing.T, v *vhost, store ae.Store, o Options) (*agent.Agent, *agent.Config) { + t.Helper() + o.Store = store + o.Verify = ae.ToolSucceeded("run_tests") + o.ClassifyToolResult = ae.ExitCode("run_tests") + config := agent.Config{Name: "triage-ticket", Tools: []tool.Tool{exits("exit=0")}} + if err := UseAgentExperience(&config, v.options(o)); err != nil { + t.Fatal(err) + } + return newAgent(newScripted(1), config), &config +} + +func waitForIndex(t *testing.T, config *agent.Config) { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + if err := WaitForIndexWrites(ctx, config); err != nil { + t.Fatal(err) + } +} + +func TestAStoredRecordIsEmbeddedInTheBackgroundFromItsSummary(t *testing.T) { + store, err := inmemory.New() + if err != nil { + t.Fatal(err) + } + v, embedder := &vhost{host: newHost(t)}, newFixed() + a, config := storing(t, v, store, Options{Embedder: embedder}) + runText(t, a, ticket) + waitForIndex(t, config) + + records := store.Records() + if len(records) != 1 || len(v.captureFailures) != 0 { + t.Fatalf("%d records, failures %+v", len(records), v.captureFailures) + } + record := records[0] + summary := ae.RecordSummary(&record) + if summary != ae.RetrievalSummary(&record.TaskID, record.TaskSummary, &record.Reflection.Lesson) { + t.Fatal("the summary is not the task ID, task summary and lesson") + } + if texts := embedder.seen(); texts[len(texts)-1] != summary { // after the query's own embedding + t.Fatalf("the embedder was given %q, not the summary %q", texts[len(texts)-1], summary) + } + stored, ok := store.Embedding(record.ExperienceID) + if !ok || stored.ModelID != vectorModel || stored.Dimension != 4 || stored.SourceRevision != record.Revision || + stored.ContentHash != ae.ContentHash(vectorModel, summary) || len(stored.Vector) != 4 || stored.Vector[0] != 1 { + t.Fatalf("%+v (%v)", stored, ok) + } +} + +func TestAFailedEmbeddingWriteIsReportedAsIndexAndTheRecordStaysStored(t *testing.T) { + for name, tc := range map[string]struct { + embedder func() *fixedEmbedder + store func(*inmemory.Store) ae.Store + }{ + "raises": {func() *fixedEmbedder { e := newFixed(); e.fail = true; return e }, func(s *inmemory.Store) ae.Store { return s }}, + "panics": {func() *fixedEmbedder { e := newFixed(); e.panics = true; return e }, func(s *inmemory.Store) ae.Store { return s }}, + "wrong-length": {func() *fixedEmbedder { e := newFixed(); e.dimension = 9; return e }, func(s *inmemory.Store) ae.Store { return s }}, + "store-fails": {newFixed, func(s *inmemory.Store) ae.Store { + return &vectorSpy{Store: s, writeErr: errors.New("the index is down: secret-record-content")} + }}, + } { + t.Run(name, func(t *testing.T) { + inner, err := inmemory.New() + if err != nil { + t.Fatal(err) + } + v := &vhost{host: newHost(t)} + a, config := storing(t, v, tc.store(inner), Options{Embedder: tc.embedder()}) + runText(t, a, ticket) + waitForIndex(t, config) + records := inner.Records() + if len(records) != 1 { + t.Fatalf("%d records", len(records)) + } + if stages := v.stages(); len(stages) != 1 || stages[0] != StageIndex { + t.Fatalf("stages %v", stages) + } + if strings.Contains(v.captureFailures[0].Reason, "secret-record-content") { + t.Fatalf("the reason quotes content: %q", v.captureFailures[0].Reason) + } + if _, ok := inner.Embedding(records[0].ExperienceID); ok { + t.Fatal("an embedding was stored") + } + }) + } +} + +func TestNoEmbeddingIsWrittenWithoutAnEmbedderOrAStoreThatTakesIt(t *testing.T) { + store, _ := inmemory.New() + v := &vhost{host: newHost(t)} + a, config := storing(t, v, store, Options{}) + runText(t, a, ticket) + waitForIndex(t, config) + if _, ok := store.Embedding(store.Records()[0].ExperienceID); ok { + t.Fatal("an embedding was written without an embedder") + } + + inner, _ := inmemory.New() + embedder := newFixed() + v = &vhost{host: newHost(t)} + a, config = storing(t, v, textOnlyStore{inner}, Options{Embedder: embedder}) + runText(t, a, ticket) + waitForIndex(t, config) + if len(inner.Records()) != 1 || embedder.called() != 0 || len(v.captureFailures) != 0 { + t.Fatalf("%d records, %d embeds, failures %+v", len(inner.Records()), embedder.called(), v.captureFailures) + } +} + +func TestCloseAwaitsTheBackgroundEmbeddingWrites(t *testing.T) { + store, _ := inmemory.New() + embedder := newFixed() + embedder.delay = 150 * time.Millisecond + v := &vhost{host: newHost(t)} + a, config := storing(t, v, store, Options{Embedder: embedder, VectorTimeout: time.Second}) + runText(t, a, ticket) + closeRecordings(t, config) // ContextProviderOf(config).Close + if _, ok := store.Embedding(store.Records()[0].ExperienceID); !ok { + t.Fatal("Close returned before the embedding was written") + } +} + +func TestALessonStoredAndEmbeddedIsFoundByARewordedQueryLater(t *testing.T) { + store, _ := inmemory.New() + floor := 0.5 + v := &vhost{host: newHost(t)} + a, config := storing(t, v, store, Options{Embedder: newFixed(), MinVectorRelevance: &floor}) + runText(t, a, ticket) + waitForIndex(t, config) + + later := &vhost{host: newHost(t)} + b, _, model := wire(t, later, store, Options{Embedder: newFixed(), Identity: ae.Identity{Principal: "go-principal", + Scope: vectorScope}}) + runText(t, b, "Zebra giraffe antelope") + if len(later.injections) != 1 || len(later.reasons()) != 0 || len(model.blocks()) != 1 { + t.Fatalf("injections %+v, fallbacks %v", later.injections, later.reasons()) + } +} + +func floatPtr(value float64) *float64 { return &value } diff --git a/go/postgres/pgtest/pgtest.go b/go/postgres/pgtest/pgtest.go index b5a41b75..9d239aa1 100644 --- a/go/postgres/pgtest/pgtest.go +++ b/go/postgres/pgtest/pgtest.go @@ -7,10 +7,14 @@ // (src/AgentExperience.Storage.Postgres/Migrations/*.sql) in name order, one transaction per script, as the owner // role, which is what ExperienceSchemaMigrator.MigrateAsync does. The application role's privileges are the // contract's sql/application-role-privileges.sql then, from contract 1.4, sql/application-role-privileges-1.4.sql (the -// injection table's), both modes (row-level security on, what a contract 1.x port runs against). +// injection table's) and, from contract 1.6, sql/application-role-privileges-1.6.sql (the embeddings table's), both +// modes (row-level security on, what a contract 1.x port runs against). From contract 1.6 the schema is also given +// the vectors package's migrations (src/AgentExperience.Storage.Postgres.Vectors/Migrations), after a superuser ran +// CREATE EXTENSION vector (pgvector is not a trusted extension), so the server must have pgvector: the container is +// the pgvector/pgvector image; a server named by PGHOST must have the extension installed. // // The server comes from PGHOST (with PGPORT, PGUSER, PGPASSWORD: a superuser connection) when set, else from a -// Testcontainers postgres container (AGENTEXPERIENCE_POSTGRES_MAJOR, default 16). With neither, the tests skip, unless +// Testcontainers pgvector container (AGENTEXPERIENCE_POSTGRES_MAJOR, default 16). With neither, the tests skip, unless // AGENTEXPERIENCE_REQUIRE_POSTGRES=1, when they fail. testcontainers-go applies TESTCONTAINERS_HUB_IMAGE_NAME_PREFIX // to every image it starts, the Reaper's (Ryuk) included, and recognises the Reaper by its name's suffix, so a // prefix (CI's Docker Hub mirror) needs no workaround here, unlike testcontainers-python (python/tests/_postgres.py). @@ -135,9 +139,9 @@ func start(ctx context.Context) (Server, func(), error) { if major == "" { major = "16" } - container, err := runContainer(ctx, "postgres:"+major) + container, err := runContainer(ctx, "pgvector/pgvector:pg"+major) if err != nil { - return Server{}, nil, fmt.Errorf("no PGHOST, and no Docker for a postgres:%s container: %w", major, err) + return Server{}, nil, fmt.Errorf("no PGHOST, and no Docker for a pgvector/pgvector:pg%s container: %w", major, err) } stop := func() { _ = testcontainers.TerminateContainer(container) } host, err := container.Host(ctx) @@ -225,17 +229,21 @@ func prepareRoles(ctx context.Context, admin Server) error { var roleName = regexp.MustCompile(`^[a-z_][a-z0-9_]*$`) -// The contract's privileges files: 1.0's, and 1.4's for the injection table. +// The contract's privileges files: 1.0's, 1.4's for the injection table and 1.6's for the embeddings table. const ( PrivilegesFile = "application-role-privileges.sql" Privileges14File = "application-role-privileges-1.4.sql" + Privileges16File = "application-role-privileges-1.6.sql" ) -// ApplicationRolePrivileges are the statements of the contract's two privileges files for role (a plain -// identifier), the 1.0 file then the 1.4 file: what the .NET privileges call grants. +// ApplicationRolePrivileges are the statements of the contract's three privileges files for role (a plain +// identifier), the 1.0 file, the 1.4 file then the 1.6 file: what the .NET privileges call grants on a schema with +// the vectors migrations applied. func ApplicationRolePrivileges(t *testing.T, role string) []string { t.Helper() - return append(PrivilegesOf(t, role, PrivilegesFile), PrivilegesOf(t, role, Privileges14File)...) + privileges := PrivilegesOf(t, role, PrivilegesFile) + privileges = append(privileges, PrivilegesOf(t, role, Privileges14File)...) + return append(privileges, PrivilegesOf(t, role, Privileges16File)...) } // PrivilegesOf are the statements of one contract privileges file (under contract/sql/) for role (a plain @@ -284,7 +292,8 @@ func DatabaseName(label string) string { return "go_" + prefix + "_" + hex.EncodeToString(digest[:])[:12] } -// CreateDatabase is a fresh database, migrated by the .NET scripts and with the application role's privileges, +// CreateDatabase is a fresh database, migrated by the .NET scripts (the vectors package's too) and with the +// application role's privileges, // dropped when the test ends (after the stores that use it close). Give each module's labels their own prefix: a // server may be shared. func CreateDatabase(t *testing.T, admin Server, label string) string { @@ -330,6 +339,22 @@ func CreateDatabaseWith(t *testing.T, admin Server, label string, privileges []s t.Fatalf("no migration scripts under %s", ct.Migrations) } sort.Strings(scripts) + vectorScripts, err := filepath.Glob(filepath.Join(ct.VectorMigrations, "*.sql")) + if err != nil || len(vectorScripts) == 0 { + t.Fatalf("no migration scripts under %s", ct.VectorMigrations) + } + sort.Strings(vectorScripts) + scripts = append(scripts, vectorScripts...) + // pgvector is not a trusted extension: a superuser installs it, as an operator does (contract 1.6). + superuser, err := pgx.Connect(ctx, admin.Admin(name)) + if err != nil { + t.Fatal(err) + } + _, err = superuser.Exec(ctx, "CREATE EXTENSION IF NOT EXISTS vector") + superuser.Close(ctx) + if err != nil { + t.Fatal(err) + } owner, err := pgx.Connect(ctx, admin.Owner(name)) if err != nil { t.Fatal(err) diff --git a/go/postgres/pgtest/vectors_postgres_test.go b/go/postgres/pgtest/vectors_postgres_test.go new file mode 100644 index 00000000..21e306db --- /dev/null +++ b/go/postgres/pgtest/vectors_postgres_test.go @@ -0,0 +1,180 @@ +//go:build postgres + +package pgtest_test + +// Contract 1.6's embedding-write and vector-search fixtures, each step through the store's WriteEmbedding and +// SearchVectors (the embedded write-experience-embedding.sql and search-experience-vector.sql), then every row of the +// embeddings table as the owner reads it. + +import ( + "context" + "encoding/json" + "strings" + "testing" + + "github.com/jackc/pgx/v5" + + ae "github.com/fabbrik/AgentExperience.NET/go" + ct "github.com/fabbrik/AgentExperience.NET/go/internal/contracttest" +) + +const vectorTolerance = 1e-9 + +// runVectorSteps seeds a fixture case, runs its steps and compares them and the embeddings table with expected.json. +func runVectorSteps(t *testing.T, family, name string) { + t.Helper() + admin := requireServer(t) + ctx := context.Background() + get := func(o *ae.Object, key string) any { value, _ := o.Get(key); return value } + seed := ct.MustReadJSON(fixture(family, name, "seed.json")).([]any) + steps := get(ct.MustReadJSON(fixture(family, name, "steps.json")).(*ae.Object), "steps").([]any) + expected := ct.MustReadJSON(fixture(family, name, "expected.json")).(*ae.Object) + + database := createDatabase(t, admin, family+"_"+name) + store := newStore(t, admin, database) + owner := ownerConn(t, admin, database) + for _, stored := range seed { + record, err := ae.FromStored(stored) + if err != nil { + t.Fatal(err) + } + if err := store.InsertRecord(ctx, ae.BoundsOf(record.Scope), &record); err != nil { + t.Fatal(err) + } + } + actual := []any{} + for _, item := range steps { + step := item.(*ae.Object) + switch { + case get(step, "write") != nil: + write := ct.EmbeddingWriteOf(get(step, "write")) + result, err := store.WriteEmbedding(ctx, ae.BoundsOf(write.Scope), write) + if err != nil { + t.Fatal(err) + } + actual = append(actual, ct.WriteOutcome(result)) + case get(step, "search") != nil: + query := ct.VectorQueryOf(get(step, "search")) + result, err := store.SearchVectors(ctx, ae.BoundsOf(query.Scope), query) + if err != nil { + t.Fatal(err) + } + actual = append(actual, ct.SearchOutcome(result)) + case get(step, "setRevision") != nil: + revision := get(step, "setRevision").(*ae.Object) + if _, err := owner.Exec(ctx, "UPDATE agent_experience.experience_records SET revision = $1::text::bigint "+ + "WHERE experience_id = $2::text::uuid", string(get(revision, "revision").(json.Number)), + get(revision, "experienceId")); err != nil { + t.Fatal(err) + } + actual = append(actual, nil) + default: + eraseStep(t, owner, get(step, "erase").(*ae.Object)) + actual = append(actual, nil) + } + } + if difference := ct.Difference(get(expected, "steps"), ct.AsWritten(actual), "$.steps", + ct.Options{Tolerance: vectorTolerance, IgnoreKind: true}); difference != "" { + t.Fatal(difference) + } + + expectedRows := get(expected, "experience_embeddings").([]any) + columns := []string{"experience_id"} + if len(expectedRows) > 0 { + columns = expectedRows[0].(*ae.Object).Keys() + } + if _, err := owner.Exec(ctx, "SET TIME ZONE 'UTC'"); err != nil { + t.Fatal(err) + } + selected := make([]string, len(columns)) + for i, column := range columns { + selected[i] = pgx.Identifier{column}.Sanitize() + "::text" + } + rows, err := owner.Query(ctx, "SELECT "+strings.Join(selected, ", ")+ + " FROM agent_experience.experience_embeddings ORDER BY experience_id") + if err != nil { + t.Fatal(err) + } + defer rows.Close() + stored := []any{} + for rows.Next() { + values := make([]*string, len(columns)) + targets := make([]any, len(columns)) + for i := range values { + targets[i] = &values[i] + } + if err := rows.Scan(targets...); err != nil { + t.Fatal(err) + } + row := ae.NewObject() + for i, column := range columns { + if values[i] == nil { + row.Set(column, nil) + } else { + row.Set(column, *values[i]) + } + } + stored = append(stored, row) + } + if err := rows.Err(); err != nil { + t.Fatal(err) + } + if difference := ct.Difference(expectedRows, stored, "$.experience_embeddings", ct.Options{}); difference != "" { + t.Fatal(difference) + } +} + +func TestEmbeddingWrite(t *testing.T) { + for _, name := range ct.Cases("embedding-write") { + t.Run(name, func(t *testing.T) { runVectorSteps(t, "embedding-write", name) }) + } +} + +func TestVectorSearch(t *testing.T) { + for _, name := range ct.Cases("vector-search") { + t.Run(name, func(t *testing.T) { runVectorSteps(t, "vector-search", name) }) + } +} + +func TestAnInvalidEmbeddingWriteOrSearchRunsNoStatementAgainstPostgres(t *testing.T) { + admin := requireServer(t) + ctx := context.Background() + store := newStore(t, admin, createDatabase(t, admin, "vector_invalid")) + scope := ae.Scope{TenantID: "contract-tenant", ApplicationID: "contract-app", ProjectID: "contract-project"} + if _, err := store.WriteEmbedding(ctx, ae.BoundsOf(scope), ae.EmbeddingWrite{ExperienceID: ae.NewRandomUUID(), + Scope: scope, ModelID: "m", Dimension: 3, ContentHash: "h", Vector: []float32{1, 2}, Now: ae.Now()}); err == nil { + t.Fatal("a vector of the wrong length was written") + } + if _, err := store.SearchVectors(ctx, ae.BoundsOf(scope), + ae.VectorQuery{Scope: scope, ModelID: " ", Vector: []float32{1}}); err == nil { + t.Fatal("a blank model ID was searched") + } +} + +func TestAWriteOutsideTheBoundsIsMissing(t *testing.T) { + admin := requireServer(t) + ctx := context.Background() + store := newStore(t, admin, createDatabase(t, admin, "vector_bounds")) + finalized, _, err := ct.RunFinalizeCase(ct.MustReadJSON(fixture("finalize", "verified-after-one-failed-attempt", + "input.json"))) + if err != nil { + t.Fatal(err) + } + record := finalized.Record + if err := store.InsertFinalized(ctx, ae.BoundsOf(record.Scope), &finalized, actor, nil); err != nil { + t.Fatal(err) + } + other := record.Scope + other.TenantID = "another-tenant" + write := ae.EmbeddingWrite{ExperienceID: record.ExperienceID, Scope: record.Scope, ModelID: "m", Dimension: 2, + ContentHash: "h", SourceRevision: record.Revision, Vector: []float32{1, 0}, Now: ae.Now()} + // The record's scope is outside the bounds the host declared: row-level security hides it. + result, err := store.WriteEmbedding(ctx, ae.BoundsOf(other), write) + if err != nil || result.Outcome != ae.EmbeddingMissing { + t.Fatalf("%+v, %v", result, err) + } + result, err = store.WriteEmbedding(ctx, ae.BoundsOf(record.Scope), write) + if err != nil || result.Outcome != ae.EmbeddingWritten { + t.Fatalf("%+v, %v", result, err) + } +} diff --git a/go/postgres/sql/search-experience-vector.sql b/go/postgres/sql/search-experience-vector.sql new file mode 100644 index 00000000..702302d8 --- /dev/null +++ b/go/postgres/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/go/postgres/sql/write-experience-embedding.sql b/go/postgres/sql/write-experience-embedding.sql new file mode 100644 index 00000000..2ebb5e05 --- /dev/null +++ b/go/postgres/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/go/postgres/store.go b/go/postgres/store.go index 809b008a..e6d90e3c 100644 --- a/go/postgres/store.go +++ b/go/postgres/store.go @@ -1,5 +1,6 @@ // Package postgres is the PostgreSQL store: it writes finalized records and their initial event, reads records -// back, runs the contract's text search, and records injections, all as the application role inside a transaction +// back, runs the contract's text search, records injections, and (contract 1.6, with pgvector) writes and searches +// embeddings, all as the application role inside a transaction // that first declared its bounds. // // It runs against a schema the .NET migrator applied (ExperienceSchemaMigrator.MigrateAsync, then @@ -30,6 +31,10 @@ const ( DeclareBoundsSQL = "declare-bounds.sql" SearchSQL = "search-experience-text.sql" RecordInjectionSQL = "record-injection.sql" + // WriteEmbeddingSQL and SearchVectorSQL are contract 1.6's: each holds several statements, each starting at its + // "-- statement: " line (ContractStatements). + WriteEmbeddingSQL = "write-experience-embedding.sql" + SearchVectorSQL = "search-experience-vector.sql" ) const ( @@ -38,6 +43,7 @@ const ( ) //go:embed sql/declare-bounds.sql sql/search-experience-text.sql sql/record-injection.sql +//go:embed sql/write-experience-embedding.sql sql/search-experience-vector.sql var contractFiles embed.FS // ContractSQL is one of the contract's SQL files, as this package embeds it (DeclareBoundsSQL, SearchSQL, @@ -50,6 +56,33 @@ func ContractSQL(name string) (string, error) { return string(data), nil } +// ContractStatements are the statements of a contract SQL file that holds several, by name: each starts at its +// "-- statement: " line and runs to the next one (or the end); the header before the first is not a statement. +func ContractStatements(name string) (map[string]string, error) { + text, err := ContractSQL(name) + if err != nil { + return nil, err + } + lines := map[string][]string{} + current := "" + started := false + for _, line := range strings.Split(strings.ReplaceAll(text, "\r\n", "\n"), "\n") { + if label, ok := strings.CutPrefix(line, "-- statement: "); ok { + current, started = strings.TrimSpace(label), true + lines[current] = nil + continue + } + if started { + lines[current] = append(lines[current], line) + } + } + statements := make(map[string]string, len(lines)) + for key, body := range lines { + statements[key] = strings.TrimSpace(strings.Join(body, "\n")) + } + return statements, nil +} + // Statement is a SQL text with its @name placeholders numbered for the driver, and the parameter each number // stands for. type Statement struct { @@ -224,6 +257,12 @@ type Store struct { declare Statement search Statement record Statement + // Contract 1.6's: the embedding write (its two statements), and the vector search statements, which keep their + // {dim} until the query vector's length is known. + writeEmbedding Statement + probeRevision Statement + vectorSQL map[string]string + probes map[string]Statement } // NewStore is a store that runs its transactions on db, which the host controls (and closes). @@ -243,8 +282,33 @@ func NewStore(db Beginner) (*Store, error) { if err != nil { return nil, err } - return &Store{db: db, declare: ToDriverPlaceholders(declare), search: ToDriverPlaceholders(search), - record: ToDriverPlaceholders(record)}, nil + writes, err := ContractStatements(WriteEmbeddingSQL) + if err != nil { + return nil, err + } + vectors, err := ContractStatements(SearchVectorSQL) + if err != nil { + return nil, err + } + store := &Store{db: db, declare: ToDriverPlaceholders(declare), search: ToDriverPlaceholders(search), + record: ToDriverPlaceholders(record), writeEmbedding: ToDriverPlaceholders(writes["write"]), + probeRevision: ToDriverPlaceholders(writes["probe-revision"]), vectorSQL: vectors, + probes: map[string]Statement{}} + for _, name := range []string{"write", "probe-revision"} { + if writes[name] == "" { + return nil, fmt.Errorf("%s has no %q statement", WriteEmbeddingSQL, name) + } + } + for _, name := range []string{"search", "search-excluding-model-authored", "compatibility", + "compatibility-excluding-model-authored"} { + if vectors[name] == "" { + return nil, fmt.Errorf("%s has no %q statement", SearchVectorSQL, name) + } + } + for _, name := range []string{"compatibility", "compatibility-excluding-model-authored"} { + store.probes[name] = ToDriverPlaceholders(vectors[name]) + } + return store, nil } // Open is a store on a new connection pool for a connection string (as the application role). Close it with Close. @@ -521,6 +585,176 @@ func (s *Store) RecordInjection(ctx context.Context, bounds ae.AuthorizationBoun var _ ae.InjectionRecorder = (*Store)(nil) +// WriteEmbedding stores a record's embedding (contract 1.6): in one transaction that declared bounds, "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. The write is validated +// first (an error, and nothing runs). Needs the grants of sql/application-role-privileges-1.6.sql and the vectors +// migrations. +func (s *Store) WriteEmbedding(ctx context.Context, bounds ae.AuthorizationBounds, + write ae.EmbeddingWrite) (ae.EmbeddingWriteResult, error) { + if err := write.Validate(); err != nil { + return ae.EmbeddingWriteResult{}, err + } + scope := write.Scope + values := map[string]any{ + "experience_id": [16]byte(write.ExperienceID), + "tenant_id": scope.TenantID, + "application_id": scope.ApplicationID, + "project_id": scope.ProjectID, + "team_id": scope.TeamID, + "agent_id": scope.AgentID, + "user_id": scope.UserID, + "model_id": write.ModelID, + "dimension": int32(write.Dimension), + "content_hash": write.ContentHash, + "source_revision": write.SourceRevision, + "embedding": ae.VectorLiteral(write.Vector), + "now": write.Now.TruncateToMicroseconds().Time(), + } + writeArgs, err := s.writeEmbedding.Args(values) + if err != nil { + return ae.EmbeddingWriteResult{}, err + } + probeArgs, err := s.probeRevision.Args(values) + if err != nil { + return ae.EmbeddingWriteResult{}, err + } + result := ae.EmbeddingWriteResult{Outcome: ae.EmbeddingWritten} + err = s.transaction(ctx, bounds, func(tx pgx.Tx) error { + tag, err := tx.Exec(ctx, s.writeEmbedding.SQL, writeArgs...) + if err != nil { + return err + } + if tag.RowsAffected() == 1 { + return nil + } + var revision int64 + switch err := tx.QueryRow(ctx, s.probeRevision.SQL, probeArgs...).Scan(&revision); { + case errors.Is(err, pgx.ErrNoRows): + result = ae.EmbeddingWriteResult{Outcome: ae.EmbeddingMissing} + case err != nil: + return err + default: + result = ae.EmbeddingWriteResult{Outcome: ae.EmbeddingStale, CurrentRevision: revision} + } + return nil + }) + if err != nil { + return ae.EmbeddingWriteResult{}, err + } + return result, nil +} + +// SearchVectors is the nearest-neighbour search (contract 1.6): in one transaction that declared 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. Hits come back nearest first, ties by experience ID, +// NaN last; a sealed row (payload_version 2) is skipped. The query is validated first (an error, and nothing runs). +func (s *Store) SearchVectors(ctx context.Context, bounds ae.AuthorizationBounds, + query ae.VectorQuery) (ae.VectorSearchResult, error) { + if err := query.Validate(); err != nil { + return ae.VectorSearchResult{}, err + } + searchName, probeName := "search", "compatibility" + if query.ExcludeModelAuthored { + searchName, probeName = "search-excluding-model-authored", "compatibility-excluding-model-authored" + } + search := ToDriverPlaceholders(strings.ReplaceAll(s.vectorSQL[searchName], "{dim}", + strconv.Itoa(len(query.Vector)))) + probe := s.probes[probeName] + scope := query.Scope + statuses := query.EffectiveStatuses() + names := make([]string, len(statuses)) + for i, status := range statuses { + names[i] = string(status) + } + values := map[string]any{ + "tenant_id": scope.TenantID, + "application_id": scope.ApplicationID, + "project_id": scope.ProjectID, + "team_id": scope.TeamID, + "agent_id": scope.AgentID, + "user_id": scope.UserID, + "statuses": names, + "min_confidence": query.EffectiveMinConfidence(), + "model_id": query.ModelID, + "query_vector": ae.VectorLiteral(query.Vector), + "limit": int32(query.EffectiveLimit()), + } + searchArgs, err := search.Args(values) + if err != nil { + return ae.VectorSearchResult{}, err + } + probeArgs, err := probe.Args(values) + if err != nil { + return ae.VectorSearchResult{}, err + } + result := ae.VectorSearchResult{Outcome: ae.VectorFound} + err = s.transaction(ctx, bounds, func(tx pgx.Tx) error { + rows, err := tx.Query(ctx, search.SQL, searchArgs...) + if err != nil { + return err + } + defer rows.Close() + fields := rows.FieldDescriptions() + count := len(ae.ColumnNames) + if len(fields) != count+4 || fields[count+3].Name != "distance" { + return errors.New("the vector search returned columns in an unexpected order") + } + for i, name := range ae.ColumnNames { + if fields[i].Name != name { + return errors.New("the vector search returned columns in an unexpected order") + } + } + returned := 0 + for rows.Next() { + returned++ + var row recordRow + var distance *float64 + targets := append(row.targets(), new(any), new(any), new(any), &distance) + if err := rows.Scan(targets...); err != nil { + return err + } + if row.payloadVersion != ae.PayloadVersion { + continue + } + if distance == nil { + return errors.New("the vector search returned a row with no distance") + } + record, err := row.record() + if err != nil { + return err + } + result.Hits = append(result.Hits, ae.VectorHit{Record: record, Distance: *distance}) + } + if err := rows.Err(); err != nil { + return err + } + rows.Close() + if returned > 0 { + return nil + } + var comparable, sameModel bool + if err := tx.QueryRow(ctx, probe.SQL, probeArgs...).Scan(&comparable, &sameModel); err != nil { + return err + } + switch { + case !comparable: + case sameModel: + result.Outcome = ae.VectorDimensionMismatch + default: + result.Outcome = ae.VectorModelMismatch + } + return nil + }) + if err != nil { + return ae.VectorSearchResult{}, err + } + result.Hits = ae.OrderHits(result.Hits) + return result, nil +} + +var _ ae.EmbeddingIndex = (*Store)(nil) + // recordRow is the eighteen record columns as the driver reads them. type recordRow struct { experienceID, sourceRunID [16]byte diff --git a/go/postgres/store_test.go b/go/postgres/store_test.go index bb0ec171..0a22748f 100644 --- a/go/postgres/store_test.go +++ b/go/postgres/store_test.go @@ -2,6 +2,7 @@ package postgres import ( "context" + "math" "os" "path/filepath" "reflect" @@ -15,7 +16,7 @@ import ( ) func TestTheEmbeddedSQLIsTheContractSQL(t *testing.T) { - for _, name := range []string{DeclareBoundsSQL, SearchSQL, RecordInjectionSQL} { + for _, name := range []string{DeclareBoundsSQL, SearchSQL, RecordInjectionSQL, WriteEmbeddingSQL, SearchVectorSQL} { embedded, err := ContractSQL(name) if err != nil { t.Fatal(err) @@ -133,3 +134,73 @@ func TestInsertFinalizedRefusesWhatIsNotARecordAndItsOwnEvent(t *testing.T) { t.Error("a store with no connection source") } } + +func TestTheContractStatementsAreSplitAtTheirMarkers(t *testing.T) { + writes, err := ContractStatements(WriteEmbeddingSQL) + if err != nil { + t.Fatal(err) + } + if len(writes) != 2 || !strings.HasPrefix(writes["write"], "INSERT INTO agent_experience.experience_embeddings") || + !strings.HasPrefix(writes["probe-revision"], "SELECT r.revision") || strings.Contains(writes["write"], "-- ") { + t.Fatalf("write statements: %v", writes) + } + vectors, err := ContractStatements(SearchVectorSQL) + if err != nil { + t.Fatal(err) + } + for _, name := range []string{"search", "search-excluding-model-authored", "compatibility", + "compatibility-excluding-model-authored"} { + if vectors[name] == "" { + t.Fatalf("no %s statement", name) + } + } + if !strings.Contains(vectors["search"], "{dim}") || strings.Contains(vectors["compatibility"], "{dim}") { + t.Fatal("only the searches have {dim}") + } + names := ToDriverPlaceholders(strings.ReplaceAll(vectors["search"], "{dim}", "4")).Names + want := []string{"query_vector", "tenant_id", "application_id", "project_id", "team_id", "agent_id", "user_id", + "statuses", "min_confidence", "model_id", "limit"} + if !reflect.DeepEqual(names, want) { + t.Fatalf("search parameters: %v", names) + } +} + +func TestAnInvalidWriteOrSearchIsRefusedBeforeAnyStatement(t *testing.T) { + // No connection source can be used: a refused call never reaches it. + store, err := NewStore(nil) + if err == nil || store != nil { + t.Fatal("a store needs a connection source") + } + store, err = NewStore(neverBegins{}) + if err != nil { + t.Fatal(err) + } + scope := ae.Scope{TenantID: "t", ApplicationID: "a", ProjectID: "p"} + good := ae.EmbeddingWrite{ExperienceID: ae.NewRandomUUID(), Scope: scope, ModelID: "m", Dimension: 2, + ContentHash: "h", SourceRevision: 1, Vector: []float32{1, 0}, Now: ae.Now()} + bad := []ae.EmbeddingWrite{good, good, good, good, good} + bad[0].ModelID = " " + bad[1].Vector = []float32{1} + bad[2].Vector = []float32{1, float32(math.NaN())} + bad[3].ContentHash = "" + bad[4].SourceRevision = -1 + for i, write := range bad { + if _, err := store.WriteEmbedding(context.Background(), ae.BoundsOf(scope), write); err == nil { + t.Fatalf("write %d was accepted", i) + } + } + for i, query := range []ae.VectorQuery{{Scope: scope, ModelID: "", Vector: []float32{1}}, + {Scope: scope, ModelID: "m"}, {Scope: scope, ModelID: "m", Vector: []float32{float32(math.Inf(1))}}, + {Scope: scope, ModelID: "m", Vector: []float32{1}, Limit: -1}} { + if _, err := store.SearchVectors(context.Background(), ae.BoundsOf(scope), query); err == nil { + t.Fatalf("query %d was accepted", i) + } + } +} + +// neverBegins fails the test's premise if a statement is ever run: it panics when asked for a transaction. +type neverBegins struct{} + +func (neverBegins) Begin(context.Context) (pgx.Tx, error) { + panic("a refused call began a transaction") +} diff --git a/go/store.go b/go/store.go index 055bc995..6017c60b 100644 --- a/go/store.go +++ b/go/store.go @@ -30,6 +30,20 @@ type InjectionRecorder interface { injectedAt Instant) error } +// EmbeddingIndex is a Store that keeps embeddings and searches them by cosine distance (contract 1.6), for semantic +// retrieval. Optional: the agent integration runs the vector channel only through a store that is one, and reports +// NotConfigured for one that is not. The PostgreSQL store (with pgvector) and the in-memory store both are. +type EmbeddingIndex interface { + // WriteEmbedding stores a record's embedding, only for a record live in exactly the write's scope and still at + // exactly its source revision: Written; Stale (with the record's revision) when it moved on, or a newer + // embedding is stored; Missing for an unknown, other-scope or erased record. A write that fails Validate is an + // error and runs nothing. + WriteEmbedding(ctx context.Context, bounds AuthorizationBounds, write EmbeddingWrite) (EmbeddingWriteResult, error) + // SearchVectors runs the nearest-neighbour search and returns its hits nearest first (ties by experience ID, + // NaN last). A query that fails Validate is an error and runs nothing. + SearchVectors(ctx context.Context, bounds AuthorizationBounds, query VectorQuery) (VectorSearchResult, error) +} + // MaxRecordInjectionIDs is the most records one RecordInjection call names: the .NET store's limit // (ExperienceRecordGetManyResult.MaxCount). The agent integration splits a longer list into chunks. const MaxRecordInjectionIDs = 200