diff --git a/README.md b/README.md index 74ab11ddc..d8570ac08 100644 --- a/README.md +++ b/README.md @@ -60,10 +60,12 @@ For a detailed .NET-to-Go feature comparison, see the [.NET and Go SDK feature c - [Skills examples](./examples/02-agents/skills/) | [Skills package](./agent/skills/) - **Text Search (RAG)**: Retrieve and cite external knowledge automatically before agent runs or through an on-demand search tool, with optional session-backed recent-message memory. - [Text search package](./agent/textsearch/) | [RAG examples](./examples/02-agents/rag/) +- **Chat History Memory**: Store conversations in a vector store and retrieve scoped memories automatically or through an on-demand tool. + - [Memory package](./agent/memory/) | [Memory examples](./examples/02-agents/memory/) - **Embeddings**: Provider-independent batch generation, numeric and binary vectors, request options, usage metadata, and single-input/vector/batch-pair helpers, with an OpenAI/Foundry adapter and opt-in OpenTelemetry instrumentation. - [Embedding package](./embedding/) | [OpenAI embedding generator](./provider/openaiprovider/embedding.go) | [Foundry embedding generator](./provider/foundryprovider/embedding.go) -- **Vector Data**: Provider-independent store and collection contracts, typed and dynamic record schemas, CRUD helpers, and vector/hybrid search options with portable filters. Includes an in-memory connector with exact vector search, batch embedding generation, and JSON snapshots. - - [Vector data package](./vectordata/) | [In-memory connector](./vectordata/inmemory/) | [Experimental provider services](./vectordata/providerservices/) +- **Vector Data**: Vector stores with typed and dynamic collections, CRUD helpers, and vector/hybrid search options with portable filters. Includes Azure AI Search and an in-memory connector with exact vector search, batch embedding generation, and JSON snapshots. + - [Vector data package](./vectordata/) | [Azure AI Search](./provider/azureprovider/) | [In-memory connector](./vectordata/inmemory/) | [Experimental provider services](./vectordata/providerservices/) - **AF Labs**: Not implemented yet in the Go SDK. - **DevUI**: Not implemented yet in the Go SDK. diff --git a/agent/memory/provider.go b/agent/memory/provider.go new file mode 100644 index 000000000..66f6ab378 --- /dev/null +++ b/agent/memory/provider.go @@ -0,0 +1,524 @@ +// Copyright (c) Microsoft. All rights reserved. + +// Package memory provides vector-backed chat history memory for agents. +package memory + +import ( + "context" + "errors" + "fmt" + "io" + "log/slog" + "reflect" + "runtime" + "strings" + "sync" + "sync/atomic" + "time" + + "github.com/google/uuid" + "github.com/microsoft/agent-framework-go/agent" + "github.com/microsoft/agent-framework-go/internal/telemetry" + "github.com/microsoft/agent-framework-go/message" + "github.com/microsoft/agent-framework-go/message/messagefilter" + "github.com/microsoft/agent-framework-go/tool/functool" + "github.com/microsoft/agent-framework-go/vectordata" + "golang.org/x/sync/semaphore" +) + +const ( + defaultSourceID = "ChatHistoryMemoryProvider" + defaultContextPrompt = "## Memories\nConsider the following memories when answering user questions:" + defaultMaxResults = 3 + defaultFunctionToolName = "Search" + defaultFunctionToolDescription = "Allows searching for related previous chat history to help answer the user question." +) + +// SearchBehavior controls when chat history is searched. +type SearchBehavior int + +const ( + // BeforeInvoke searches before each agent run and adds results as a user message. + BeforeInvoke SearchBehavior = iota + + // OnDemandFunctionCalling exposes a search tool for the model to call when needed. + OnDemandFunctionCalling +) + +// Scope selects which application, agent, user, and session a memory belongs to. +// A nil ID leaves that dimension unrestricted when searching. An explicit empty +// string matches only an empty ID. +type Scope struct { + ApplicationID *string `json:"applicationId,omitempty"` + AgentID *string `json:"agentId,omitempty"` + UserID *string `json:"userId,omitempty"` + SessionID *string `json:"sessionId,omitempty"` +} + +// State holds a provider's per-session scopes. The provider stores a *State under +// the key returned by [HistoryProvider.StateKeys]. Synchronize edits with runs and tool +// calls, and use [agent.Session.Set] to persist changes to restored state. +type State struct { + StorageScope Scope `json:"storageScope"` + + // SearchScope defaults to StorageScope when nil. + SearchScope *Scope `json:"searchScope,omitempty"` +} + +// Config configures the provider created by [NewHistoryProvider]. +type Config struct { + // SearchTime defaults to [BeforeInvoke]. + SearchTime SearchBehavior + + // FunctionToolName names the on-demand search tool. Nil defaults to "Search". + FunctionToolName *string + + // FunctionToolDescription describes the on-demand search tool. Nil uses a + // default description about searching related previous chat history. + FunctionToolDescription *string + + // ContextPrompt prefixes retrieved messages. Nil uses a default memory prompt. + // An explicit empty string suppresses the prompt text. + ContextPrompt *string + + // MaxResults limits the number of search results. Nil defaults to 3. + // An explicit value must be positive. + MaxResults *int + + // StateKey identifies provider state in the session. Nil defaults to + // "ChatHistoryMemoryProvider". Explicit empty or whitespace-only keys are + // invalid. Use distinct keys for providers that need independent scopes. + StateKey *string + + // SearchInputMessageFilter filters automatic search input. + // Nil defaults to [messagefilter.ExternalOnly]. + SearchInputMessageFilter messagefilter.Filter + + // StorageInputRequestMessageFilter filters request messages before storage. + // Nil defaults to [messagefilter.ExternalOnly]. + StorageInputRequestMessageFilter messagefilter.Filter + + // StorageInputResponseMessageFilter filters response messages before storage. + // Nil passes all response messages through. + StorageInputResponseMessageFilter messagefilter.Filter + + // Logger receives provider diagnostics. When nil, logging is disabled. + Logger *slog.Logger + + // LogSensitiveData includes user IDs, queries, and retrieved text in logs. + // The default is false. When true, Redactor is ignored. + LogSensitiveData bool + + // Redactor transforms sensitive log data. When nil and LogSensitiveData is + // false, nonempty sensitive values are replaced with "". + Redactor func(string) string +} + +// HistoryProvider stores chat messages in a vector collection and retrieves related +// messages as context. Create it with [NewHistoryProvider]. Do not copy it after use. +type HistoryProvider struct { + collection vectordata.Collection[any, map[string]any] + config Config + stateInitializer func(*agent.Session) *State + stateKey string + maxResults int + contextPrompt string + toolConfig functool.Config + contextProvider agent.ContextProvider + stateMu sync.Mutex + initialized atomic.Bool + initialization *semaphore.Weighted + cleanup runtime.Cleanup +} + +// NewHistoryProvider creates a chat history memory provider. +// +// The store must support dynamic map records, UUID keys, and embedding generation +// from strings. The provider obtains a collection handle at construction and +// creates the collection lazily on the first search or successful run. Failed +// creation is retried. The store is borrowed. +// +// stateInitializer supplies storage and search scopes when a session has no +// readable state. It must return a non-nil state. Without a session, it is called with +// nil for each lifecycle hook and no state is retained. On-demand tools capture +// the current invocation's search scope, not a later run's scope. +// +// Attach the provider using [agent.Config.ContextProviders]. Automatic searches +// join nonblank request texts with newlines and append nonblank results as one +// source-attributed user message. On-demand searches use only the tool's +// userQuestion argument. Use automatic tool invocation middleware when the +// underlying provider does not execute tools. +// +// Successful runs store filtered request messages followed by response messages, +// including messages without text. Storage and automatic-search failures are +// logged when Logger is configured and otherwise ignored, except cancellation +// errors while the caller's context is canceled. On-demand search, message-filter, +// and session-state errors are returned. +// +// Stored messages may contain sensitive data. Configure storage access controls +// and appropriate scopes; an empty scope searches across all stored conversations. +// Retrieved content and model-generated queries are untrusted. +// +// It panics if store or stateInitializer is nil, collectionName or StateKey is +// blank, or vectorDimensions or an explicit MaxResults is not positive. +// Collection configuration errors are returned. +func NewHistoryProvider(store *vectordata.Store, collectionName string, vectorDimensions int, stateInitializer func(*agent.Session) *State, config Config) (*HistoryProvider, error) { + if store == nil { + panic("store is required") + } + if stateInitializer == nil { + panic("stateInitializer is required") + } + if strings.TrimSpace(collectionName) == "" { + panic("collectionName must not be blank") + } + if vectorDimensions < 1 { + panic("vectorDimensions must be positive") + } + p := &HistoryProvider{ + config: config, + stateInitializer: stateInitializer, + stateKey: defaultSourceID, + maxResults: defaultMaxResults, + contextPrompt: defaultContextPrompt, + toolConfig: functool.Config{ + Name: defaultFunctionToolName, + Description: defaultFunctionToolDescription, + }, + initialization: semaphore.NewWeighted(1), + } + if config.StateKey != nil { + p.stateKey = *config.StateKey + } + if strings.TrimSpace(p.stateKey) == "" { + panic("StateKey must not be blank") + } + if config.MaxResults != nil { + p.maxResults = *config.MaxResults + if p.maxResults < 1 { + panic("MaxResults must be positive") + } + } + if config.ContextPrompt != nil { + p.contextPrompt = *config.ContextPrompt + } + if config.FunctionToolName != nil { + p.toolConfig.Name = *config.FunctionToolName + } + if config.FunctionToolDescription != nil { + p.toolConfig.Description = *config.FunctionToolDescription + } + collection, err := store.DynamicCollection(collectionName, + &vectordata.Definition{Properties: []vectordata.Property{ + &vectordata.KeyProperty{Name: "Key", Type: reflect.TypeFor[uuid.UUID]()}, + &vectordata.DataProperty{Name: "Role", Type: reflect.TypeFor[string](), IsIndexed: true}, + &vectordata.DataProperty{Name: "MessageId", Type: reflect.TypeFor[string](), IsIndexed: true}, + &vectordata.DataProperty{Name: "AuthorName", Type: reflect.TypeFor[string]()}, + &vectordata.DataProperty{Name: "ApplicationId", Type: reflect.TypeFor[string](), IsIndexed: true}, + &vectordata.DataProperty{Name: "AgentId", Type: reflect.TypeFor[string](), IsIndexed: true}, + &vectordata.DataProperty{Name: "UserId", Type: reflect.TypeFor[string](), IsIndexed: true}, + &vectordata.DataProperty{Name: "SessionId", Type: reflect.TypeFor[string](), IsIndexed: true}, + &vectordata.DataProperty{Name: "Content", Type: reflect.TypeFor[string](), IsFullTextIndexed: true}, + &vectordata.DataProperty{Name: "CreatedAt", Type: reflect.TypeFor[string](), IsIndexed: true}, + &vectordata.VectorProperty{Name: "ContentEmbedding", Type: reflect.TypeFor[string](), Dimensions: vectorDimensions}, + }}, + ) + if err != nil { + return nil, err + } + p.collection = collection + p.contextProvider = agent.NewContextProvider(agent.ContextProviderConfig{ + SourceID: defaultSourceID, + ProvideInputMessageFilter: config.SearchInputMessageFilter, + StoreInputRequestMessageFilter: config.StorageInputRequestMessageFilter, + StoreInputResponseMessageFilter: config.StorageInputResponseMessageFilter, + Provide: p.provide, + Store: p.store, + }) + if closer, ok := collection.(io.Closer); ok { + logger := config.Logger + p.cleanup = runtime.AddCleanup(p, func(closer io.Closer) { + if err := closer.Close(); err != nil && logger != nil { + logger.LogAttrs(context.Background(), slog.LevelError, "chat history collection cleanup failed", slog.Any("error", err)) + } + }, closer) + } + return p, nil +} + +// StateKeys returns the session key used by this provider. The returned slice +// belongs to the caller. +func (p *HistoryProvider) StateKeys() []string { + return []string{p.stateKey} +} + +// Invoking adds retrieved memories or a search tool scoped to this invocation. +func (p *HistoryProvider) Invoking(ctx context.Context, invoking agent.InvokingContext) ([]*message.Message, []agent.Option, error) { + defer runtime.KeepAlive(p) + return p.contextProvider.Invoking(ctx, invoking) +} + +// Invoked stores eligible request and response messages after successful runs. +func (p *HistoryProvider) Invoked(ctx context.Context, invoked agent.InvokedContext) error { + defer runtime.KeepAlive(p) + return p.contextProvider.Invoked(ctx, invoked) +} + +// Close is an optimization that releases the collection early if it implements +// [io.Closer]. A successful call cancels the [runtime.AddCleanup] fallback, which +// attempts to close the collection after the provider becomes unreachable. +// Runtime cleanup is asynchronous and is not guaranteed to run before program exit. +// +// Close never closes the borrowed store. Repeated calls follow the collection's +// Close behavior. Do not overlap Close with runs or tool calls. +func (p *HistoryProvider) Close() error { + defer runtime.KeepAlive(p) + if closer, ok := p.collection.(io.Closer); ok { + if err := closer.Close(); err != nil { + return err + } + } + p.cleanup.Stop() + return nil +} + +func (p *HistoryProvider) provide(ctx context.Context, invoking agent.InvokingContext) ([]*message.Message, []agent.Option, error) { + telemetry.MarkUsed(telemetry.FeatureChatHistoryMemoryProvider) + session, _ := agent.GetOption(invoking.Options, agent.WithSession) + state, err := p.loadState(session) + if err != nil { + return nil, nil, err + } + scope := state.SearchScope + if scope == nil { + scope = &state.StorageScope + } + if p.config.SearchTime == OnDemandFunctionCalling { + search, err := functool.New(p.toolConfig, func(ctx context.Context, args struct { + UserQuestion string `json:"userQuestion"` + }, + ) (string, error) { + return p.search(ctx, args.UserQuestion, scope) + }) + if err != nil { + return nil, nil, err + } + return nil, []agent.Option{agent.WithTool(search)}, nil + } + var parts []string + for _, msg := range invoking.Messages { + if msg != nil { + if text := msg.String(); strings.TrimSpace(text) != "" { + parts = append(parts, text) + } + } + } + formatted, err := p.search(ctx, strings.Join(parts, "\n"), scope) + if err != nil { + return nil, nil, p.handleFailure(ctx, "chat history search failed", scope, err) + } + if strings.TrimSpace(formatted) == "" { + return nil, nil, nil + } + return []*message.Message{message.NewText(formatted)}, nil, nil +} + +func (p *HistoryProvider) store(ctx context.Context, invoked agent.InvokedContext) error { + session, _ := agent.GetOption(invoked.Options, agent.WithSession) + state, err := p.loadState(session) + if err != nil { + return err + } + if err := p.storeMessages(ctx, invoked, &state.StorageScope); err != nil { + return p.handleFailure(ctx, "chat history storage failed", &state.StorageScope, err) + } + return nil +} + +func (p *HistoryProvider) storeMessages(ctx context.Context, invoked agent.InvokedContext, scope *Scope) error { + if err := p.ensureCollection(ctx); err != nil { + return err + } + records := make([]map[string]any, 0, len(invoked.RequestMessages)+len(invoked.ResponseMessages)) + for _, messages := range [][]*message.Message{invoked.RequestMessages, invoked.ResponseMessages} { + for _, msg := range messages { + if msg == nil { + return errors.New("memory: cannot store a nil message") + } + key, err := uuid.NewRandom() + if err != nil { + return err + } + createdAt := msg.CreatedAt + if createdAt.IsZero() { + createdAt = time.Now().UTC() + } + text := msg.String() + records = append(records, map[string]any{ + "Key": key, + "Role": string(msg.Role), + "MessageId": msg.ID, + "AuthorName": msg.AuthorName, + "ApplicationId": nullableString(scope.ApplicationID), + "AgentId": nullableString(scope.AgentID), + "UserId": nullableString(scope.UserID), + "SessionId": nullableString(scope.SessionID), + "Content": text, + "CreatedAt": createdAt.Format(time.RFC3339Nano), + "ContentEmbedding": text, + }) + } + } + if len(records) == 0 { + return nil + } + return p.collection.Upsert(ctx, records...) +} + +func (p *HistoryProvider) search(ctx context.Context, query string, scope *Scope) (string, error) { + defer runtime.KeepAlive(p) + if strings.TrimSpace(query) == "" { + return "", nil + } + if err := p.ensureCollection(ctx); err != nil { + return "", err + } + var records []map[string]any + for result, err := range p.collection.Search(ctx, query, p.maxResults, &vectordata.SearchOptions{Filter: scope.filter()}) { + if err != nil { + return "", err + } + records = append(records, result.Record) + } + if logger := p.config.Logger; logger != nil && logger.Enabled(ctx, slog.LevelInfo) { + logger.LogAttrs(ctx, slog.LevelInfo, "chat history search completed", + append(p.scopeLogAttrs(scope), slog.Int("count", len(records)))...) + } + var parts []string + for i, record := range records { + content, ok := record["Content"] + if !ok { + return "", fmt.Errorf("memory: search result %d has no Content field", i) + } + if content == nil { + continue + } + text, ok := content.(string) + if !ok { + return "", fmt.Errorf("memory: search result %d has Content of type %T, want string", i, content) + } + if strings.TrimSpace(text) != "" { + parts = append(parts, text) + } + } + if len(parts) == 0 { + return "", nil + } + formatted := p.contextPrompt + "\n" + strings.Join(parts, "\n") + if logger := p.config.Logger; logger != nil && logger.Enabled(ctx, slog.LevelDebug) { + logger.LogAttrs(ctx, slog.LevelDebug, "chat history search results", + append(p.scopeLogAttrs(scope), + slog.String("query", p.sanitizeLogData(query)), + slog.String("context", p.sanitizeLogData(formatted)))...) + } + return formatted, nil +} + +func (p *HistoryProvider) loadState(session *agent.Session) (*State, error) { + p.stateMu.Lock() + defer p.stateMu.Unlock() + var state *State + if ok, err := session.Get(p.stateKey, &state); err != nil { + return nil, err + } else if !ok || state == nil { + state = p.stateInitializer(session) + if state == nil { + return nil, errors.New("memory: stateInitializer returned nil") + } + session.Set(p.stateKey, state) + } + return state, nil +} + +func (p *HistoryProvider) ensureCollection(ctx context.Context) error { + if p.initialized.Load() { + return nil + } + if err := p.initialization.Acquire(ctx, 1); err != nil { + return err + } + defer p.initialization.Release(1) + if p.initialized.Load() { + return nil + } + if err := p.collection.EnsureExists(ctx); err != nil { + return err + } + p.initialized.Store(true) + return nil +} + +func (s *Scope) filter() *vectordata.Filter { + var filters []vectordata.Filter + for _, field := range []struct { + name string + value *string + }{ + {"ApplicationId", s.ApplicationID}, + {"AgentId", s.AgentID}, + {"UserId", s.UserID}, + {"SessionId", s.SessionID}, + } { + if field.value != nil { + filters = append(filters, vectordata.Filter{Operator: vectordata.Equal, Property: field.name, Value: *field.value}) + } + } + switch len(filters) { + case 0: + return nil + case 1: + return &filters[0] + default: + return &vectordata.Filter{Operator: vectordata.And, Operands: filters} + } +} + +func (p *HistoryProvider) handleFailure(ctx context.Context, text string, scope *Scope, err error) error { + if ctx.Err() != nil && (errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded)) { + return err + } + if logger := p.config.Logger; logger != nil && logger.Enabled(ctx, slog.LevelError) { + logger.LogAttrs(ctx, slog.LevelError, text, append(p.scopeLogAttrs(scope), slog.Any("error", err))...) + } + return nil +} + +func (p *HistoryProvider) scopeLogAttrs(scope *Scope) []slog.Attr { + userID := "" + if scope.UserID != nil { + userID = p.sanitizeLogData(*scope.UserID) + } + return []slog.Attr{ + slog.Any("application_id", nullableString(scope.ApplicationID)), + slog.Any("agent_id", nullableString(scope.AgentID)), + slog.Any("session_id", nullableString(scope.SessionID)), + slog.String("user_id", userID), + } +} + +func (p *HistoryProvider) sanitizeLogData(data string) string { + if p.config.LogSensitiveData || data == "" { + return data + } + if p.config.Redactor != nil { + return p.config.Redactor(data) + } + return "" +} + +func nullableString(value *string) any { + if value != nil { + return *value + } + return nil +} diff --git a/agent/memory/provider_test.go b/agent/memory/provider_test.go new file mode 100644 index 000000000..1c8d25a54 --- /dev/null +++ b/agent/memory/provider_test.go @@ -0,0 +1,1737 @@ +// Copyright (c) Microsoft. All rights reserved. + +package memory_test + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "iter" + "log/slog" + "reflect" + "runtime" + "slices" + "strings" + "sync" + "sync/atomic" + "testing" + "testing/synctest" + "time" + + "github.com/google/uuid" + "github.com/microsoft/agent-framework-go/agent" + "github.com/microsoft/agent-framework-go/agent/harness/toolautocall" + "github.com/microsoft/agent-framework-go/agent/memory" + "github.com/microsoft/agent-framework-go/embedding" + "github.com/microsoft/agent-framework-go/internal/agenttest" + "github.com/microsoft/agent-framework-go/message" + "github.com/microsoft/agent-framework-go/message/messagefilter" + "github.com/microsoft/agent-framework-go/tool" + "github.com/microsoft/agent-framework-go/tool/functool" + "github.com/microsoft/agent-framework-go/vectordata" + "github.com/microsoft/agent-framework-go/vectordata/inmemory" + "github.com/microsoft/agent-framework-go/vectordata/providerservices" +) + +const ( + stateKey = "ChatHistoryMemoryProvider" + contextPrompt = "## Memories\nConsider the following memories when answering user questions:" +) + +type record = map[string]any + +type testStore struct { + vectordata.StoreProvider + collection vectordata.Collection[any, any] + name string + definition *vectordata.Definition + err error + closeCalls int +} + +func (s *testStore) OpenCollection(name string, _, _ reflect.Type, definition *vectordata.Definition) (vectordata.Collection[any, any], error) { + s.name, s.definition = name, definition + return s.collection, s.err +} + +func (s *testStore) Close() error { + s.closeCalls++ + return nil +} + +type testCollection struct { + vectordata.Collection[any, any] + ensure func(context.Context) error + search func(context.Context, any, int, *vectordata.SearchOptions) iter.Seq2[vectordata.SearchResult[record], error] + upsert func(context.Context, ...record) error + close func() error +} + +func (c *testCollection) EnsureExists(ctx context.Context) error { + if c.ensure != nil { + return c.ensure(ctx) + } + return nil +} + +func (c *testCollection) Search(ctx context.Context, query any, top int, options *vectordata.SearchOptions) iter.Seq2[vectordata.SearchResult[any], error] { + return func(yield func(vectordata.SearchResult[any], error) bool) { + if c.search != nil { + for result, err := range c.search(ctx, query, top, options) { + if !yield(vectordata.SearchResult[any]{Record: result.Record, Score: result.Score}, err) { + return + } + } + } + } +} + +func (c *testCollection) Upsert(ctx context.Context, records ...any) error { + if c.upsert != nil { + typed := make([]record, len(records)) + for i, value := range records { + typed[i] = value.(record) + } + return c.upsert(ctx, typed...) + } + return nil +} + +func (c *testCollection) Close() error { + if c.close != nil { + return c.close() + } + return nil +} + +func searchResults(records ...record) iter.Seq2[vectordata.SearchResult[record], error] { + return func(yield func(vectordata.SearchResult[record], error) bool) { + for _, r := range records { + if !yield(vectordata.SearchResult[record]{Record: r}, nil) { + return + } + } + } +} + +func newProvider(t *testing.T, collection vectordata.Collection[any, any], initialize func(*agent.Session) *memory.State, config memory.Config) *memory.HistoryProvider { + t.Helper() + if initialize == nil { + initialize = func(*agent.Session) *memory.State { return new(memory.State) } + } + provider, err := memory.NewHistoryProvider(vectordata.NewStore(&testStore{collection: collection}), "history", 2, initialize, config) + if err != nil { + t.Fatal(err) + } + return provider +} + +func invokeProvider(provider agent.ContextProvider, ctx context.Context, messages []*message.Message, options ...agent.Option) ([]*message.Message, []agent.Option, error) { + return provider.Invoking(ctx, agent.InvokingContext{Messages: messages, Options: options}) +} + +func storeMessages(t *testing.T, provider agent.ContextProvider, session *agent.Session, request, response []*message.Message) { + t.Helper() + if err := provider.Invoked(t.Context(), agent.InvokedContext{ + RequestMessages: request, ResponseMessages: response, Options: []agent.Option{agent.WithSession(session)}, + }); err != nil { + t.Fatal(err) + } +} + +func searchTool(t *testing.T, provider agent.ContextProvider, options ...agent.Option) tool.FuncTool { + t.Helper() + request := message.NewText("not an automatic query") + messages, returned, err := invokeProvider(provider, t.Context(), []*message.Message{request}, options...) + if err != nil { + t.Fatal(err) + } + if len(messages) != 1 || messages[0] != request { + t.Fatalf("on-demand mode changed input messages: %v", messages) + } + tools := slices.Collect(agent.AllOptions(returned, agent.WithTool)) + if len(tools) != 1 { + t.Fatalf("tools = %d, want 1", len(tools)) + } + search, ok := tools[0].(tool.FuncTool) + if !ok { + t.Fatalf("search tool has type %T", tools[0]) + } + return search +} + +func TestChatHistoryMemoryProvider_ConfigValidation(t *testing.T) { + for _, name := range []string{ + "nil store", "nil initializer", "empty collection", "blank collection", + "zero dimensions", "negative dimensions", "zero maximum", "negative maximum", + "empty state key", "blank state key", + } { + t.Run(name, func(t *testing.T) { + store := vectordata.NewStore(&testStore{collection: &testCollection{}}) + initialize := func(*agent.Session) *memory.State { return new(memory.State) } + collectionName, dimensions := "history", 2 + config := memory.Config{} + switch name { + case "nil store": + store = nil + case "nil initializer": + initialize = nil + case "empty collection": + collectionName = "" + case "blank collection": + collectionName = " \t" + case "zero dimensions": + dimensions = 0 + case "negative dimensions": + dimensions = -5 + case "zero maximum": + config.MaxResults = new(0) + case "negative maximum": + config.MaxResults = new(-1) + case "empty state key": + config.StateKey = new("") + case "blank state key": + config.StateKey = new(" \t") + } + defer func() { + if recover() == nil { + t.Fatal("expected a configuration panic") + } + }() + _, _ = memory.NewHistoryProvider(store, collectionName, dimensions, initialize, config) + }) + } +} + +func TestChatHistoryMemoryProvider_CollectionDefinition(t *testing.T) { + store := &testStore{collection: &testCollection{ensure: func(context.Context) error { + t.Fatal("collection was created eagerly") + return nil + }}} + provider, err := memory.NewHistoryProvider(vectordata.NewStore(store), "history", 1536, + func(*agent.Session) *memory.State { return new(memory.State) }, memory.Config{}) + if err != nil || provider == nil || store.name != "history" { + t.Fatalf("NewHistoryProvider = (%v, %v), name = %q", provider, err, store.name) + } + want := &vectordata.Definition{Properties: []vectordata.Property{ + &vectordata.KeyProperty{Name: "Key", Type: reflect.TypeFor[uuid.UUID]()}, + &vectordata.DataProperty{Name: "Role", Type: reflect.TypeFor[string](), IsIndexed: true}, + &vectordata.DataProperty{Name: "MessageId", Type: reflect.TypeFor[string](), IsIndexed: true}, + &vectordata.DataProperty{Name: "AuthorName", Type: reflect.TypeFor[string]()}, + &vectordata.DataProperty{Name: "ApplicationId", Type: reflect.TypeFor[string](), IsIndexed: true}, + &vectordata.DataProperty{Name: "AgentId", Type: reflect.TypeFor[string](), IsIndexed: true}, + &vectordata.DataProperty{Name: "UserId", Type: reflect.TypeFor[string](), IsIndexed: true}, + &vectordata.DataProperty{Name: "SessionId", Type: reflect.TypeFor[string](), IsIndexed: true}, + &vectordata.DataProperty{Name: "Content", Type: reflect.TypeFor[string](), IsFullTextIndexed: true}, + &vectordata.DataProperty{Name: "CreatedAt", Type: reflect.TypeFor[string](), IsIndexed: true}, + &vectordata.VectorProperty{Name: "ContentEmbedding", Type: reflect.TypeFor[string](), Dimensions: 1536}, + }} + if !reflect.DeepEqual(store.definition, want) { + t.Fatalf("collection definition = %#v, want %#v", store.definition, want) + } + wantErr := errors.New("unsupported collection") + store.err = wantErr + if provider, err := memory.NewHistoryProvider(vectordata.NewStore(store), "history", 2, + func(*agent.Session) *memory.State { return new(memory.State) }, memory.Config{}); provider != nil || err != wantErr { + t.Fatalf("collection failure = (%v, %v)", provider, err) + } +} + +func TestChatHistoryMemoryProvider_StateKeysAndInitialization(t *testing.T) { + for _, key := range []*string{nil, new("custom-key")} { + wantKey := stateKey + if key != nil { + wantKey = *key + } + t.Run(wantKey, func(t *testing.T) { + session := agenttest.CreateSession() + want := &memory.State{StorageScope: memory.Scope{UserID: new("user")}} + calls := 0 + provider := newProvider(t, &testCollection{}, func(got *agent.Session) *memory.State { + if got != session { + t.Fatal("initializer received a different session") + } + calls++ + return want + }, memory.Config{StateKey: key}) + keys := provider.StateKeys() + if !slices.Equal(keys, []string{wantKey}) { + t.Fatalf("StateKeys = %v", keys) + } + keys[0] = "changed" + if key != nil { + *key = "changed" + } + if provider.StateKeys()[0] != wantKey { + t.Fatal("state key changed with caller-owned configuration") + } + for range 2 { + if _, _, err := invokeProvider(provider, t.Context(), nil, agent.WithSession(session)); err != nil { + t.Fatal(err) + } + storeMessages(t, provider, session, nil, nil) + } + var got *memory.State + if ok, err := session.Get(wantKey, &got); err != nil || !ok || got != want || calls != 1 { + t.Fatalf("state = (%v, %v, %v), initializer calls = %d", got, ok, err, calls) + } + }) + } +} + +func TestChatHistoryMemoryProvider_StateRestoration(t *testing.T) { + for _, explicitSearch := range []bool{false, true} { + t.Run(fmt.Sprintf("explicit-search=%t", explicitSearch), func(t *testing.T) { + initial := &memory.State{StorageScope: memory.Scope{ + ApplicationID: new("app"), AgentID: new("agent"), UserID: new("user"), SessionID: new(""), + }} + wantSearch := initial.StorageScope + if explicitSearch { + initial.SearchScope = &memory.Scope{UserID: new("search-user")} + wantSearch = *initial.SearchScope + } + calls := 0 + collection := &testCollection{} + provider := newProvider(t, collection, func(*agent.Session) *memory.State { + calls++ + return initial + }, memory.Config{}) + session := agenttest.CreateSession() + if _, _, err := invokeProvider(provider, t.Context(), nil, agent.WithSession(session)); err != nil { + t.Fatal(err) + } + data, err := json.Marshal(session) + if err != nil { + t.Fatal(err) + } + restored := new(agent.Session) + if err := json.Unmarshal(data, restored); err != nil { + t.Fatal(err) + } + var got *memory.State + if ok, err := restored.Get(stateKey, &got); err != nil || !ok || !reflect.DeepEqual(got, initial) { + t.Fatalf("restored state = (%#v, %v, %v)", got, ok, err) + } + collection.search = func(_ context.Context, _ any, _ int, options *vectordata.SearchOptions) iter.Seq2[vectordata.SearchResult[record], error] { + assertScopeFilter(t, options.Filter, wantSearch) + return searchResults() + } + collection.upsert = func(_ context.Context, records ...record) error { + if len(records) != 1 || records[0]["UserId"] != "user" || records[0]["SessionId"] != "" { + t.Fatalf("restored storage scope = %#v", records) + } + return nil + } + if _, _, err := invokeProvider(provider, t.Context(), []*message.Message{message.NewText("query")}, agent.WithSession(restored)); err != nil { + t.Fatal(err) + } + storeMessages(t, provider, restored, []*message.Message{message.NewText("stored")}, nil) + if calls != 1 { + t.Fatalf("initializer calls = %d, want no reinitialization after restore", calls) + } + }) + } +} + +func TestChatHistoryMemoryProvider_StateErrors(t *testing.T) { + session := agenttest.CreateSession() + var logs bytes.Buffer + collection := &testCollection{ensure: func(context.Context) error { + t.Fatal("state failure reached the collection") + return nil + }} + provider := newProvider(t, collection, func(*agent.Session) *memory.State { + return nil + }, memory.Config{Logger: slog.New(slog.NewJSONHandler(&logs, nil))}) + if _, _, err := invokeProvider(provider, t.Context(), nil, agent.WithSession(session)); err == nil { + t.Fatal("Invoking did not return a state error") + } + if err := provider.Invoked(t.Context(), agent.InvokedContext{Options: []agent.Option{agent.WithSession(session)}}); err == nil { + t.Fatal("Invoked did not return a state error") + } + if logs.Len() != 0 { + t.Fatal("state errors were treated as best-effort storage/search failures") + } +} + +func TestChatHistoryMemoryProvider_MissingAndNilSessionState(t *testing.T) { + for _, name := range []string{"missing", "typed nil", "null", "other type", "unreadable restored state", "no session"} { + t.Run(name, func(t *testing.T) { + session := agenttest.CreateSession() + switch name { + case "typed nil": + session.Set(stateKey, (*memory.State)(nil)) + case "null": + if err := json.Unmarshal([]byte(`{"State":{"ChatHistoryMemoryProvider":null}}`), session); err != nil { + t.Fatal(err) + } + case "other type": + session.Set(stateKey, 42) + case "unreadable restored state": + if err := json.Unmarshal([]byte(`{"State":{"ChatHistoryMemoryProvider":{"storageScope":42}}}`), session); err != nil { + t.Fatal(err) + } + case "no session": + session = nil + } + calls := 0 + provider := newProvider(t, &testCollection{}, func(got *agent.Session) *memory.State { + if got != session { + t.Fatal("initializer session mismatch") + } + calls++ + return new(memory.State) + }, memory.Config{}) + for range 2 { + if _, _, err := invokeProvider(provider, t.Context(), nil, agent.WithSession(session)); err != nil { + t.Fatal(err) + } + storeMessages(t, provider, session, nil, nil) + } + want := 1 + if session == nil { + want = 4 + } + if calls != want { + t.Fatalf("initializations = %d, want %d", calls, want) + } + }) + } +} + +func TestChatHistoryMemoryProvider_StoresMessages(t *testing.T) { + var stored []record + ensureCalls, upsertCalls := 0, 0 + collection := &testCollection{ + ensure: func(ctx context.Context) error { + if ctx != t.Context() { + t.Fatal("EnsureExists context changed") + } + ensureCalls++ + return nil + }, + upsert: func(ctx context.Context, records ...record) error { + if ctx != t.Context() { + t.Fatal("Upsert context changed") + } + upsertCalls++ + stored = append(stored, records...) + return nil + }, + } + provider := newProvider(t, collection, func(*agent.Session) *memory.State { + return &memory.State{StorageScope: memory.Scope{ + ApplicationID: new("app"), AgentID: new("agent"), UserID: new("user"), SessionID: new("session"), + }} + }, memory.Config{}) + request := message.NewText("request text") + request.ID, request.AuthorName = "req-1", "user1" + request.CreatedAt = time.Date(2000, 1, 1, 0, 0, 0, 123456789, time.FixedZone("offset", 3600)) + response := message.NewText("response text") + response.Role, response.ID, response.AuthorName = message.RoleAssistant, "resp-1", "assistant" + empty := message.New(&message.FunctionCallContent{CallID: "call", Name: "tool"}) + empty.Role = message.RoleAssistant + before := time.Now().UTC() + storeMessages(t, provider, agenttest.CreateSession(), + []*message.Message{request, message.NewText("request without metadata")}, + []*message.Message{response, empty}) + after := time.Now().UTC() + if len(stored) != 4 || ensureCalls != 1 || upsertCalls != 1 { + t.Fatalf("stored = %d, ensure calls = %d, upsert calls = %d", len(stored), ensureCalls, upsertCalls) + } + keys := make(map[uuid.UUID]bool) + for i, want := range []*message.Message{request, message.NewText("request without metadata"), response, empty} { + got := stored[i] + key, ok := got["Key"].(uuid.UUID) + if !ok || key == uuid.Nil || keys[key] || key.Version() != 4 { + t.Fatalf("invalid or repeated key: %v", got["Key"]) + } + keys[key] = true + if len(got) != 11 || got["Content"] != want.String() || got["ContentEmbedding"] != want.String() || + got["MessageId"] != want.ID || got["AuthorName"] != want.AuthorName || got["Role"] != string(want.Role) { + t.Fatalf("stored message %d = %#v", i, got) + } + for field, value := range map[string]string{"ApplicationId": "app", "AgentId": "agent", "UserId": "user", "SessionId": "session"} { + if got[field] != value { + t.Fatalf("stored %s = %v, want %q", field, got[field], value) + } + } + created, err := time.Parse(time.RFC3339Nano, got["CreatedAt"].(string)) + if err != nil { + t.Fatal(err) + } + if !want.CreatedAt.IsZero() { + if got["CreatedAt"] != want.CreatedAt.Format(time.RFC3339Nano) { + t.Fatalf("creation time = %q, want original offset and precision", got["CreatedAt"]) + } + } else if created.Before(before) || created.After(after) { + t.Fatalf("generated creation time %v not in [%v, %v]", created, before, after) + } + } + if !response.CreatedAt.IsZero() || !empty.CreatedAt.IsZero() || request.ID != "req-1" { + t.Fatal("storage changed input messages") + } + storeMessages(t, provider, nil, []*message.Message{request}, nil) + if len(stored) != 5 || ensureCalls != 1 || upsertCalls != 2 || stored[4]["Key"] == stored[0]["Key"] { + t.Fatal("repeated storage did not create a fresh record without reinitializing the collection") + } +} + +func TestChatHistoryMemoryProvider_EmptyStorageAndScopes(t *testing.T) { + ensureCalls, upsertCalls := 0, 0 + provider := newProvider(t, &testCollection{ + ensure: func(context.Context) error { ensureCalls++; return nil }, + upsert: func(_ context.Context, records ...record) error { + upsertCalls++ + if len(records) != 1 || records[0]["ApplicationId"] != nil || records[0]["AgentId"] != nil || + records[0]["SessionId"] != nil || records[0]["UserId"] != "" { + t.Fatalf("nil/empty storage scopes = %#v", records) + } + return nil + }, + }, func(*agent.Session) *memory.State { + return &memory.State{StorageScope: memory.Scope{UserID: new("")}} + }, memory.Config{}) + storeMessages(t, provider, nil, nil, nil) + if ensureCalls != 1 || upsertCalls != 0 { + t.Fatalf("empty store: ensure = %d, upsert = %d", ensureCalls, upsertCalls) + } + storeMessages(t, provider, nil, []*message.Message{message.NewText("")}, nil) + if upsertCalls != 1 { + t.Fatal("message without text was dropped") + } +} + +func TestChatHistoryMemoryProvider_FailedRunSkipsStorage(t *testing.T) { + provider := newProvider(t, &testCollection{}, func(*agent.Session) *memory.State { + t.Fatal("failed run initialized state") + return nil + }, memory.Config{StorageInputRequestMessageFilter: func(context.Context, []*message.Message) ([]*message.Message, error) { + t.Fatal("failed run applied a storage filter") + return nil, nil + }}) + if err := provider.Invoked(t.Context(), agent.InvokedContext{ + RequestMessages: []*message.Message{message.NewText("request")}, Err: errors.New("run failed"), + }); err != nil { + t.Fatal(err) + } +} + +func TestChatHistoryMemoryProvider_AutomaticSearch(t *testing.T) { + for _, tc := range []struct { + name string + config memory.Config + wantTop int + prompt string + }{ + {name: "defaults", wantTop: 3, prompt: contextPrompt}, + {name: "custom", config: memory.Config{MaxResults: new(2), ContextPrompt: new("Context:\n")}, wantTop: 2, prompt: "Context:\n"}, + {name: "empty prompt", config: memory.Config{ContextPrompt: new("")}, wantTop: 3}, + {name: "unknown behavior", config: memory.Config{SearchTime: memory.SearchBehavior(42)}, wantTop: 3, prompt: contextPrompt}, + } { + t.Run(tc.name, func(t *testing.T) { + calls := 0 + provider := newProvider(t, &testCollection{ + search: func(ctx context.Context, query any, top int, options *vectordata.SearchOptions) iter.Seq2[vectordata.SearchResult[record], error] { + calls++ + if ctx != t.Context() || query != " first \nsecond" || top != tc.wantTop || options.Filter != nil || + options.Skip != 0 || options.IncludeVectors || options.VectorProperty != "" { + t.Fatalf("search = (%v, %v, %d, %#v)", ctx, query, top, options) + } + return searchResults( + record{"Content": "one"}, record{"Content": nil}, record{"Content": " \t"}, record{"Content": " two "}, + ) + }, + }, nil, tc.config) + request := []*message.Message{message.NewText(" first "), nil, message.NewText("\t "), message.NewText("second")} + original := slices.Clone(request) + options := []agent.Option{agent.WithInstructions("keep these instructions")} + messages, returned, err := invokeProvider(provider, t.Context(), request, options...) + if err != nil || calls != 1 || len(messages) != len(request)+1 { + t.Fatalf("Invoking = (%v, %v), calls = %d", messages, err, calls) + } + if !slices.Equal(messages[:len(request)], original) || !slices.Equal(request, original) { + t.Fatal("search changed the input messages") + } + if got, _ := agent.GetOption(returned, agent.WithInstructions); got != "keep these instructions" || len(returned) != len(options) { + t.Fatal("search changed input options") + } + result := messages[len(request)] + if result.String() != tc.prompt+"\none\n two " || result.Role != message.RoleUser || + result.Source != (message.Source{Type: agent.SourceTypeContextProvider, ID: stateKey}) { + t.Fatalf("retrieved context = %#v (%q)", result, result.String()) + } + }) + } +} + +func TestChatHistoryMemoryProvider_EmptySearch(t *testing.T) { + for _, mode := range []memory.SearchBehavior{memory.BeforeInvoke, memory.OnDemandFunctionCalling} { + t.Run(fmt.Sprint(mode), func(t *testing.T) { + provider := newProvider(t, &testCollection{ensure: func(context.Context) error { + t.Fatal("blank query initialized the collection") + return nil + }}, nil, memory.Config{SearchTime: mode}) + session := agenttest.CreateSession() + if mode == memory.BeforeInvoke { + request := []*message.Message{nil, message.NewText(" \t\n")} + messages, _, err := invokeProvider(provider, t.Context(), request, agent.WithSession(session)) + if err != nil || !slices.Equal(messages, request) { + t.Fatalf("blank search = (%v, %v)", messages, err) + } + } else { + result, err := searchTool(t, provider, agent.WithSession(session)).Call(t.Context(), `{"userQuestion":" \t\n"}`) + if err != nil || result != "" { + t.Fatalf("blank tool search = (%v, %v)", result, err) + } + } + var state *memory.State + if ok, err := session.Get(stateKey, &state); err != nil || !ok || state == nil { + t.Fatalf("blank query did not initialize session state: (%v, %v, %v)", state, ok, err) + } + }) + } + for _, records := range [][]record{nil, {{"Content": nil}, {"Content": ""}, {"Content": " \n"}}} { + provider := newProvider(t, &testCollection{ + search: func(context.Context, any, int, *vectordata.SearchOptions) iter.Seq2[vectordata.SearchResult[record], error] { + return searchResults(records...) + }, + }, nil, memory.Config{}) + request := []*message.Message{message.NewText("query")} + messages, _, err := invokeProvider(provider, t.Context(), request) + if err != nil || !slices.Equal(messages, request) { + t.Fatalf("empty results produced context: (%v, %v)", messages, err) + } + } +} + +func assertScopeFilter(t *testing.T, filter *vectordata.Filter, scope memory.Scope) { + t.Helper() + values := map[string]*string{ + "ApplicationId": scope.ApplicationID, "AgentId": scope.AgentID, "UserId": scope.UserID, "SessionId": scope.SessionID, + } + matching := make(record) + restricted := false + for name, value := range values { + if value != nil { + matching[name] = *value + restricted = true + } else { + matching[name] = "unrestricted" + } + } + if !restricted && filter != nil || restricted && filter == nil { + t.Fatalf("scope filter = %#v for %#v", filter, scope) + } + predicate, err := providerservices.CompileFilter(filter, func(name string) (providerservices.FilterProperty, error) { + if _, ok := values[name]; !ok { + return providerservices.FilterProperty{}, fmt.Errorf("unexpected scope property %q", name) + } + return providerservices.FilterProperty{ + Type: reflect.TypeFor[string](), + Value: func(input any) (any, bool, error) { + value, ok := input.(record)[name] + return value, ok, nil + }, + }, nil + }) + if err != nil { + t.Fatal(err) + } + if ok, err := predicate(matching); err != nil || !ok { + t.Fatalf("matching record rejected: (%v, %v)", ok, err) + } + for name, value := range values { + original := matching[name] + for _, mismatch := range []any{"different", nil} { + matching[name] = mismatch + if ok, err := predicate(matching); err != nil || ok != (value == nil) { + t.Fatalf("%s=%v matched: (%v, %v), restricted=%t", name, mismatch, ok, err, value != nil) + } + } + matching[name] = original + } +} + +func TestChatHistoryMemoryProvider_ScopeFilters(t *testing.T) { + for mask := range 16 { + for _, value := range []string{"scope-value", ""} { + t.Run(fmt.Sprintf("mask=%d/value=%q", mask, value), func(t *testing.T) { + scope := memory.Scope{} + for i, field := range []**string{&scope.ApplicationID, &scope.AgentID, &scope.UserID, &scope.SessionID} { + if mask&(1< 0 && redactions == 0 { + t.Fatalf("redactor calls = %d", redactions) + } + }) + } + } + } +} + +func TestChatHistoryMemoryProvider_DisabledLogging(t *testing.T) { + for _, withLogger := range []bool{false, true} { + t.Run(fmt.Sprintf("logger=%t", withLogger), func(t *testing.T) { + var logs bytes.Buffer + config := memory.Config{Redactor: func(string) string { + t.Fatal("disabled log invoked the redactor") + return "" + }} + if withLogger { + config.Logger = slog.New(slog.NewJSONHandler(&logs, &slog.HandlerOptions{Level: slog.Level(100)})) + } + provider := newProvider(t, &testCollection{ + search: func(context.Context, any, int, *vectordata.SearchOptions) iter.Seq2[vectordata.SearchResult[record], error] { + return searchResults(record{"Content": "memory"}) + }, + upsert: func(context.Context, ...record) error { return errors.New("storage failed") }, + }, func(*agent.Session) *memory.State { + return &memory.State{StorageScope: memory.Scope{UserID: new("private user")}} + }, config) + request := []*message.Message{message.NewText("query")} + if _, _, err := invokeProvider(provider, t.Context(), request); err != nil { + t.Fatal(err) + } + storeMessages(t, provider, nil, request, nil) + if logs.Len() != 0 { + t.Fatalf("disabled logging produced %s", logs.String()) + } + }) + } +} + +func TestChatHistoryMemoryProvider_NilStoredMessage(t *testing.T) { + var logs bytes.Buffer + provider := newProvider(t, &testCollection{ + upsert: func(context.Context, ...record) error { + t.Fatal("invalid message batch was partially stored") + return nil + }, + }, nil, memory.Config{Logger: slog.New(slog.NewJSONHandler(&logs, nil))}) + storeMessages(t, provider, nil, []*message.Message{message.NewText("valid"), nil}, nil) + if !strings.Contains(logs.String(), "cannot store a nil message") { + t.Fatalf("invalid message was not logged: %s", logs.String()) + } +} + +type textGenerator struct { + calls [][]string +} + +func (g *textGenerator) Generate(_ context.Context, inputs []string, _ *embedding.Options) (*embedding.Generated[*embedding.Of[float32]], error) { + g.calls = append(g.calls, slices.Clone(inputs)) + result := &embedding.Generated[*embedding.Of[float32]]{} + for _, input := range inputs { + vector := embedding.VectorOf[float32]{1, 0} + if input == "beta" { + vector = embedding.VectorOf[float32]{0, 1} + } + result.Embeddings = append(result.Embeddings, &embedding.Of[float32]{Vector: vector}) + } + return result, nil +} + +func TestChatHistoryMemoryProvider_InMemoryStore(t *testing.T) { + generator := new(textGenerator) + store := inmemory.New(inmemory.Config{EmbeddingGenerator: embedding.Bind(generator)}) + provider, err := memory.NewHistoryProvider(store.Store, "history", 2, + func(*agent.Session) *memory.State { return new(memory.State) }, memory.Config{MaxResults: new(1)}) + if err != nil { + t.Fatal(err) + } + for _, seed := range []struct { + user string + session string + text string + }{{"user-a", "one", "alpha"}, {"user-a", "two", "beta"}, {"user-b", "one", "alpha"}} { + session := agenttest.CreateSession() + session.Set(stateKey, &memory.State{StorageScope: memory.Scope{UserID: new(seed.user), SessionID: new(seed.session)}}) + storeMessages(t, provider, session, []*message.Message{message.NewText(seed.text)}, nil) + } + reader := agenttest.CreateSession() + state := &memory.State{ + StorageScope: memory.Scope{UserID: new("user-a"), SessionID: new("one")}, + SearchScope: &memory.Scope{UserID: new("user-a")}, + } + reader.Set(stateKey, state) + for _, query := range []string{"alpha", "beta"} { + messages, _, err := invokeProvider(provider, t.Context(), []*message.Message{message.NewText(query)}, agent.WithSession(reader)) + if err != nil || len(messages) != 2 || messages[1].String() != contextPrompt+"\n"+query { + t.Fatalf("cross-session scoped search %q = (%v, %v)", query, messages, err) + } + } + state.SearchScope = nil + messages, _, err := invokeProvider(provider, t.Context(), []*message.Message{message.NewText("beta")}, agent.WithSession(reader)) + if err != nil || len(messages) != 2 || messages[1].String() != contextPrompt+"\nalpha" { + t.Fatalf("default storage-scoped search = (%v, %v)", messages, err) + } + if want := [][]string{{"alpha"}, {"beta"}, {"alpha"}, {"alpha"}, {"beta"}, {"beta"}}; !reflect.DeepEqual(generator.calls, want) { + t.Fatalf("embedding inputs = %v, want %v", generator.calls, want) + } +} + +func TestChatHistoryMemoryProvider_AgentLifecycle(t *testing.T) { + for _, stream := range []bool{false, true} { + t.Run(fmt.Sprintf("stream=%t", stream), func(t *testing.T) { + var queries []string + var batches [][]string + provider := newProvider(t, &testCollection{ + search: func(_ context.Context, query any, _ int, _ *vectordata.SearchOptions) iter.Seq2[vectordata.SearchResult[record], error] { + queries = append(queries, query.(string)) + return searchResults(record{"Content": "retrieved memory"}) + }, + upsert: func(_ context.Context, records ...record) error { + var texts []string + for _, r := range records { + texts = append(texts, r["Content"].(string)) + } + batches = append(batches, texts) + return nil + }, + }, nil, memory.Config{}) + assertContext := func(_ context.Context, messages []*message.Message, _ ...agent.Option) { + last := messages[len(messages)-1] + if last.Source.ID != stateKey || last.String() != contextPrompt+"\nretrieved memory" { + t.Fatalf("agent did not receive memory context: %#v", last) + } + } + runner := &agenttest.Runner{Responses: agenttest.NewResponseBuilder(assertContext). + AddText("reply ").AddText("A").NewTurn(assertContext).AddText("reply ").AddText("B").Build()} + a := agent.New(agent.ProviderConfig{Run: runner.Run}, agent.Config{ContextProviders: []agent.ContextProvider{provider}}) + session, err := a.CreateSession(t.Context()) + if err != nil { + t.Fatal(err) + } + for _, question := range []string{"A", "B"} { + response, err := a.RunText(t.Context(), question, agent.WithSession(session), agent.Stream(stream)).Collect() + if err != nil || response.String() != "reply "+question { + t.Fatalf("agent response = (%v, %v)", response, err) + } + } + if !slices.Equal(queries, []string{"A", "B"}) || + !reflect.DeepEqual(batches, [][]string{{"A", "reply A"}, {"B", "reply B"}}) { + t.Fatalf("queries=%q, stored=%q; history or retrieved context was duplicated", queries, batches) + } + }) + } +} + +func TestChatHistoryMemoryProvider_AgentOnDemandToolExecution(t *testing.T) { + for _, stream := range []bool{false, true} { + t.Run(fmt.Sprintf("stream=%t", stream), func(t *testing.T) { + searchCalls := 0 + var stored []record + provider := newProvider(t, &testCollection{ + search: func(_ context.Context, query any, _ int, options *vectordata.SearchOptions) iter.Seq2[vectordata.SearchResult[record], error] { + searchCalls++ + if query != "model query" { + t.Fatalf("query = %v", query) + } + assertScopeFilter(t, options.Filter, memory.Scope{UserID: new("user")}) + return searchResults(record{"Content": "retrieved"}) + }, + upsert: func(_ context.Context, records ...record) error { stored = records; return nil }, + }, func(*agent.Session) *memory.State { + return &memory.State{StorageScope: memory.Scope{UserID: new("user")}} + }, memory.Config{SearchTime: memory.OnDemandFunctionCalling}) + runner := &agenttest.Runner{Responses: agenttest.NewResponseBuilder(). + AddFunctionCall("call", "Search", `{"userQuestion":"model query"}`). + NewTurn(func(_ context.Context, messages []*message.Message, _ ...agent.Option) { + result, ok := messages[len(messages)-1].Contents[0].(*message.FunctionResultContent) + if !ok || result.CallID != "call" || result.Result != contextPrompt+"\nretrieved" { + t.Fatalf("tool result = %#v", result) + } + }).AddText("answer").Build()} + a := agent.New(agent.ProviderConfig{ + Run: runner.Run, Middlewares: []agent.Middleware{toolautocall.New(toolautocall.Config{})}, + }, agent.Config{ContextProviders: []agent.ContextProvider{provider}}) + response, err := a.RunText(t.Context(), "question", agent.Stream(stream)).Collect() + if err != nil || response.String() != "answer" || searchCalls != 1 { + t.Fatalf("run = (%v, %v), search calls = %d", response, err, searchCalls) + } + var contents, roles []string + for _, r := range stored { + contents, roles = append(contents, r["Content"].(string)), append(roles, r["Role"].(string)) + } + if !slices.Equal(contents, []string{"question", "", "", "answer"}) || + !slices.Equal(roles, []string{"user", "assistant", "tool", "assistant"}) { + t.Fatalf("stored tool conversation: contents=%q, roles=%q", contents, roles) + } + }) + } +} + +func TestChatHistoryMemoryProvider_UnfinishedRunsDoNotStore(t *testing.T) { + for _, earlyStop := range []bool{false, true} { + t.Run(fmt.Sprintf("early-stop=%t", earlyStop), func(t *testing.T) { + provider := newProvider(t, &testCollection{upsert: func(context.Context, ...record) error { + t.Fatal("unfinished run stored messages") + return nil + }}, nil, memory.Config{}) + released := false + wantErr := errors.New("model failed") + a := agent.New(agent.ProviderConfig{ + Run: func(context.Context, []*message.Message, ...agent.Option) iter.Seq2[*agent.ResponseUpdate, error] { + return func(yield func(*agent.ResponseUpdate, error) bool) { + defer func() { released = true }() + if !yield(&agent.ResponseUpdate{Role: message.RoleAssistant, Contents: []message.Content{&message.TextContent{Text: "partial"}}}, nil) { + return + } + yield(nil, wantErr) + } + }, + }, agent.Config{ContextProviders: []agent.ContextProvider{provider}}) + if earlyStop { + for _, err := range a.RunText(t.Context(), "query", agent.Stream(true)) { + if err != nil { + t.Fatal(err) + } + break + } + } else if _, err := a.RunText(t.Context(), "query", agent.Stream(true)).Collect(); !errors.Is(err, wantErr) { + t.Fatalf("run error = %v", err) + } + if !released { + t.Fatal("underlying run was not released") + } + }) + } +} diff --git a/docs/dotnet-go-sdk-feature-comparison.md b/docs/dotnet-go-sdk-feature-comparison.md index eb2870bd7..8463601df 100644 --- a/docs/dotnet-go-sdk-feature-comparison.md +++ b/docs/dotnet-go-sdk-feature-comparison.md @@ -193,7 +193,9 @@ The [vector data package](../vectordata/) is based on [Microsoft.Extensions.VectorData.Abstractions 10.10.0](https://github.com/dotnet/extensions/tree/294665a5ec4eceaa245388cfb70883ea01881081/src/Libraries/Microsoft.Extensions.VectorData.Abstractions). Go differences: -- Stores fix key/record types through generics. `Get` returns +- Generic collection methods require Go 1.27; `GetCollection` and + `DynamicCollection` also work on Go 1.26. Provider dispatch uses runtime types. +- `Get` returns `(record, found, error)`; `iter.Seq2` replaces async enumerables, and variadic `Upsert` replaces separate single/batch methods. - Struct tags replace attributes; dynamic records use `map[string]any`. @@ -213,16 +215,36 @@ The [in-memory connector](../vectordata/inmemory/) ports (tag `mevd/InMemory/v1.0.1`). Additional differences: - Plain structs follow Go value semantics; use pointers or maps to observe - generated UUIDs on input. Keys follow Go equality and reject nil and NaN. + generated UUIDs on input. Keys use their declared schema type for Go equality + and reject nil and NaN. - Strings sort lexicographically. Custom comparisons use `Compare(T) int` and `Equal(T) bool`; equality otherwise uses `reflect.DeepEqual`. - Batch dictionary updates are atomic; mutations to input records are not rolled back on failure. -- Snapshots use Go JSON, not the .NET format. NaN scores are omitted; other scores +- Snapshots use Go JSON and the generic `inmemory.Import` function. NaN scores are omitted; other scores may differ through floating-point rounding. +The [Azure AI Search connector](../provider/azureprovider/) ports the +[Toolkit connector at the same revision](https://github.com/CommunityToolkit/AI/tree/7a48d308df63d85427a772f031616889905a0fdf/MEVD/src/AzureAISearch). +Additional Go differences: + +- Uses the `azcore` REST pipeline instead of a native Search client. +- VectorData storage names replace serializer naming configuration; only schema + properties are mapped. Go `int` maps to `Edm.Int64`. +- Hybrid `IncludeVectors` is honored, and all partial batch failures are returned. +- Nil collections serialize as empty arrays; membership filters retain empty strings. + These packages remain outside the catalog's extraction scope. +The [chat-history memory provider](../agent/memory/) ports +[.NET 1.25.0](https://github.com/microsoft/agent-framework/blob/95511ff025d81cc546c44e653d4be65de5bd7f35/dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProvider.cs). +Go differences: + +- Uses the context-provider interface without a separate message-only overload. +- Scope IDs are optional pointers; a nil `State.SearchScope` uses `StorageScope`. +- Stored timestamps use RFC3339Nano; absent message IDs/authors use empty strings. +- `Close` disposes the collection only when it implements `io.Closer`. + The separate [dotnetsymbols command](../_catalog/cmd/dotnetsymbols/README.md) extracts an inventory from compiled .NET assemblies using `go-winmd`. By default, `go -C _catalog run ./cmd/dotnetsymbols` resolves the latest stable MAF release and downloads its prebuilt core NuGet packages, without Git or a .NET SDK. Pin a version with `-release 1.22.0` for reproducible reruns. Select additional packages with `-package` and an exact target framework with `-framework`. Package versions, hashes, source metadata, and assembly identities are recorded; the release is not assumed to match this catalog's commit baseline. For unpublished source, build the relevant upstream library projects separately and pass their reference or implementation assemblies with `-assembly`. Both modes include types, declared members, overloads, events, generic constraints, and experimental/compiler-generated markers without executing assembly code. Extraction omits metadata not used for reconciliation, such as parameter names/defaults, constant values, raw attribute lists and unused API flags. No per-symbol paths or line numbers are stored. diff --git a/docs/feature-usage-telemetry.md b/docs/feature-usage-telemetry.md index 753b8dde5..d3165d427 100644 --- a/docs/feature-usage-telemetry.md +++ b/docs/feature-usage-telemetry.md @@ -87,6 +87,7 @@ These switches do not configure OpenTelemetry or change service telemetry. | 0 | `core.agent` | Agent invocation | | 2 | `core.workflow` | Successful workflow build | | 3 | `core.tool_approval` | Tool-approval middleware participation | +| 4 | `core.chat_history_memory_provider` | Automatic retrieval or on-demand search-tool provisioning | | 5 | `core.file_memory_provider` | File-memory context provider participation | | 6 | `core.text_search_provider` | Automatic search or on-demand search-tool provisioning | | 8 | `core.skills_provider` | Skills context provider participation | diff --git a/embedding/embedding_test.go b/embedding/embedding_test.go index 5189a7152..630e9df74 100644 --- a/embedding/embedding_test.go +++ b/embedding/embedding_test.go @@ -317,7 +317,7 @@ func TestEmbedding_JsonDeserialization_KnownPayload(t *testing.T) { data := []byte(`{ "Type":"float32", "vector":[1,2,3], "createdAt":"2024-01-01T00:00:00Z", - "modelId":"text-embedding-ada-002", + "modelId":"text-embedding-3-large", "additionalProperties":{"key":"val"} }`) var e embedding.Of[float32] @@ -325,7 +325,7 @@ func TestEmbedding_JsonDeserialization_KnownPayload(t *testing.T) { t.Fatal(err) } if !slices.Equal(e.Vector, []float32{1, 2, 3}) || - e.ModelID != "text-embedding-ada-002" || + e.ModelID != "text-embedding-3-large" || !e.CreatedAt.Equal(time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC)) || e.AdditionalProperties["key"] != "val" { t.Fatalf("embedding = %+v", e) @@ -334,7 +334,7 @@ func TestEmbedding_JsonDeserialization_KnownPayload(t *testing.T) { if err != nil { t.Fatal(err) } - want := `{"Type":"float32","CreatedAt":"2024-01-01T00:00:00Z","ModelID":"text-embedding-ada-002","AdditionalProperties":{"key":"val"},"Vector":[1,2,3]}` + want := `{"Type":"float32","CreatedAt":"2024-01-01T00:00:00Z","ModelID":"text-embedding-3-large","AdditionalProperties":{"key":"val"},"Vector":[1,2,3]}` if string(data) != want { t.Fatalf("Marshal = %s, want %s", data, want) } diff --git a/embedding/options_test.go b/embedding/options_test.go index fccfee10e..5bbf7eed6 100644 --- a/embedding/options_test.go +++ b/embedding/options_test.go @@ -88,9 +88,9 @@ func TestEmbeddingGenerationOptions_JsonSerialization_Roundtrips(t *testing.T) { func TestEmbeddingGenerationOptions_JsonDeserialization_KnownPayload(t *testing.T) { var options embedding.Options - err := json.Unmarshal([]byte(`{"dimensions":1536,"modelId":"text-embedding-ada-002","additionalProperties":{"key":"val"}}`), &options) - if err != nil || options.Dimensions == nil || *options.Dimensions != 1536 || - options.ModelID == nil || *options.ModelID != "text-embedding-ada-002" || + err := json.Unmarshal([]byte(`{"dimensions":3072,"modelId":"text-embedding-3-large","additionalProperties":{"key":"val"}}`), &options) + if err != nil || options.Dimensions == nil || *options.Dimensions != 3072 || + options.ModelID == nil || *options.ModelID != "text-embedding-3-large" || options.AdditionalProperties["key"] != "val" { t.Fatalf("options = %+v, %v", options, err) } diff --git a/examples/02-agents/agents/step22_foundry_memory/main.go b/examples/02-agents/agents/step22_foundry_memory/main.go index 6f2d91c20..aa342dd43 100644 --- a/examples/02-agents/agents/step22_foundry_memory/main.go +++ b/examples/02-agents/agents/step22_foundry_memory/main.go @@ -19,7 +19,7 @@ import ( var ( chatDeployment = demo.FoundryModel memoryStoreName = cmp.Or(os.Getenv("AZURE_AI_MEMORY_STORE_ID"), "memory-store-sample") - embeddingDeployment = cmp.Or(os.Getenv("AZURE_AI_EMBEDDING_DEPLOYMENT_NAME"), "text-embedding-ada-002") + embeddingDeployment = cmp.Or(os.Getenv("AZURE_AI_EMBEDDING_DEPLOYMENT_NAME"), "text-embedding-3-large") ) const memoryScope = "sample-user-123" diff --git a/examples/02-agents/memory/README.md b/examples/02-agents/memory/README.md new file mode 100644 index 000000000..f26804e0f --- /dev/null +++ b/examples/02-agents/memory/README.md @@ -0,0 +1,44 @@ +# Memory + +Ports of the [.NET memory samples](https://github.com/microsoft/agent-framework/tree/95511ff025d81cc546c44e653d4be65de5bd7f35/dotnet/samples/02-agents/AgentWithMemory). +Step numbers follow upstream. + +| Example | Demonstrates | +| --- | --- | +| [Step 01: Chat history memory](./step01_chat_history_memory/) | Remembering a user's preferences across two independent sessions | +| [Step 05: Bounded chat history](./step05_bounded_chat_history/) | Keeping four recent messages and recalling older messages from vector-store overflow | + +Both use `memory.HistoryProvider` and an in-memory vector store. Memories survive +new sessions, but not process restarts. The fixed `UID1` scope is for demonstration; +applications should use the authenticated user's identifier. + +## Run + +Sign in with a credential supported by `DefaultAzureCredential`, such as +`az login`, and configure deployed Foundry models: + +```powershell +$env:FOUNDRY_PROJECT_ENDPOINT = "https://.services.ai.azure.com/api/projects/" +$env:FOUNDRY_MODEL = "gpt-5.4-mini" +$env:FOUNDRY_EMBEDDING_MODEL = "text-embedding-3-large" + +go run .\examples\02-agents\memory\step01_chat_history_memory +go run .\examples\02-agents\memory\step05_bounded_chat_history +``` + +These samples use **3072-dimensional embeddings**, matching the shared demo +default, `text-embedding-3-large`. If your embedding deployment produces a +different vector size, update the dimension argument to `memory.NewHistoryProvider` +in the sample. Runs invoke live services and may incur charges. + +Step 05 preserves the first system message outside the four-message window. +Function-call and function-result messages are excluded from both retained history +and archived memory. Retrieved memories are prepended to recent history and are +not stored again. + +## Local validation + +```powershell +go build .\examples\02-agents\memory\... +go vet .\examples\02-agents\memory\... +``` diff --git a/examples/02-agents/memory/step01_chat_history_memory/main.go b/examples/02-agents/memory/step01_chat_history_memory/main.go new file mode 100644 index 000000000..f96e97de3 --- /dev/null +++ b/examples/02-agents/memory/step01_chat_history_memory/main.go @@ -0,0 +1,71 @@ +// Copyright (c) Microsoft. All rights reserved. + +// This sample recalls chat messages from earlier sessions for the same user. +package main + +import ( + "context" + "fmt" + "log/slog" + + "github.com/google/uuid" + "github.com/microsoft/agent-framework-go/agent" + "github.com/microsoft/agent-framework-go/agent/memory" + "github.com/microsoft/agent-framework-go/embedding" + "github.com/microsoft/agent-framework-go/examples/internal/demo" + "github.com/microsoft/agent-framework-go/provider/foundryprovider" + "github.com/microsoft/agent-framework-go/vectordata/inmemory" +) + +var logger = demo.NewLogger( + "Chat History Memory", + "Demonstrates scoped recall across separate conversations.", + "Model", demo.FoundryModel, + "Embedding Model", demo.FoundryEmbeddingModel, +) + +func main() { + ctx := context.Background() + credential := demo.FoundryTokenCredential() + generator := foundryprovider.NewEmbeddingGenerator(demo.FoundryProjectEndpoint, credential, + foundryprovider.EmbeddingGeneratorConfig{Model: demo.FoundryEmbeddingModel}) + store := inmemory.New(inmemory.Config{EmbeddingGenerator: embedding.Bind(generator)}) + + memories, err := memory.NewHistoryProvider(store.Store, "chathistory", 3072, + func(*agent.Session) *memory.State { + return &memory.State{ + StorageScope: memory.Scope{UserID: new("UID1"), SessionID: new(uuid.NewString())}, + SearchScope: &memory.Scope{UserID: new("UID1")}, + } + }, memory.Config{Logger: slog.New(logger)}) + if err != nil { + demo.Panic(err) + } + + a := foundryprovider.NewAgent(demo.FoundryProjectEndpoint, credential, + foundryprovider.ModelDeployment(demo.FoundryModel), foundryprovider.AgentConfig{ + Instructions: "You are good at telling jokes.", + DisableStoreOutput: true, + Config: agent.Config{ + Name: "Joker", + ContextProviders: []agent.ContextProvider{memories}, + Middlewares: []agent.Middleware{logger}, + }, + }) + + for i, question := range []string{ + "I like jokes about Pirates. Tell me a joke about a pirate.", + "Tell me a joke that I might like.", + } { + session, err := a.CreateSession(ctx) + if err != nil { + demo.Panic(err) + } + fmt.Printf(">> Session %d\n\n", i+1) + response, err := a.RunText(ctx, question, agent.WithSession(session)).Collect() + if err != nil { + demo.Panic(err) + } + demo.Response(response, nil) + } +} diff --git a/examples/02-agents/memory/step05_bounded_chat_history/history.go b/examples/02-agents/memory/step05_bounded_chat_history/history.go new file mode 100644 index 000000000..14d1502cf --- /dev/null +++ b/examples/02-agents/memory/step05_bounded_chat_history/history.go @@ -0,0 +1,99 @@ +// Copyright (c) Microsoft. All rights reserved. + +package main + +import ( + "context" + "strings" + + "github.com/microsoft/agent-framework-go/agent" + "github.com/microsoft/agent-framework-go/agent/memory" + "github.com/microsoft/agent-framework-go/message" + "github.com/microsoft/agent-framework-go/message/messagefilter" +) + +const boundedHistoryStateKey = "bounded-chat-history" + +func newBoundedHistoryProvider(maxMessages int, memories *memory.HistoryProvider) agent.HistoryProvider { + if maxMessages <= 0 { + panic("maxMessages must be positive") + } + history := agent.NewInMemoryHistoryProvider(agent.InMemoryHistoryProviderConfig{ + StateKey: boundedHistoryStateKey, + StoreInputRequestMessageFilter: messagefilter.PassThrough, + }) + return agent.NewHistoryProvider(agent.HistoryProviderConfig{ + SourceID: boundedHistoryStateKey, + Provide: func(ctx context.Context, invoking agent.InvokingContext) ([]*message.Message, error) { + recent, err := history.Invoking(ctx, agent.InvokingContext{Options: invoking.Options}) + if err != nil { + return nil, err + } + result, _, err := memories.Invoking(ctx, invoking) + if err != nil { + return nil, err + } + var texts []string + for _, msg := range result { + if msg != nil && msg.Source.Type == agent.SourceTypeContextProvider && strings.TrimSpace(msg.String()) != "" { + texts = append(texts, msg.String()) + } + } + if len(texts) == 0 { + return recent, nil + } + contextMessage := message.NewText("The following are memories from earlier in this conversation. Use them to inform your responses:\n" + + strings.Join(texts, "\n")) + return append([]*message.Message{contextMessage}, recent...), nil + }, + Store: func(ctx context.Context, invoked agent.InvokedContext) error { + if err := history.Invoked(ctx, invoked); err != nil { + return err + } + session, _ := agent.GetOption(invoked.Options, agent.WithSession) + messages, err := session.InMemoryHistory(boundedHistoryStateKey) + if err != nil { + return err + } + retained, removed := truncateHistory(messages, maxMessages) + if err := session.SetInMemoryHistory(boundedHistoryStateKey, retained); err != nil { + return err + } + if len(removed) == 0 { + return nil + } + return memories.Invoked(ctx, agent.InvokedContext{ + RequestMessages: removed, + Options: invoked.Options, + }) + }, + }) +} + +func truncateHistory(messages []*message.Message, maxMessages int) (retained, removed []*message.Message) { + var system *message.Message + var conversation []*message.Message +nextMessage: + for _, msg := range messages { + if msg.Role == message.RoleSystem { + if system == nil { + system = msg + } + continue + } + for _, content := range msg.Contents { + switch content.(type) { + case *message.FunctionCallContent, *message.FunctionResultContent: + continue nextMessage + } + } + conversation = append(conversation, msg) + } + overflow := max(0, len(conversation)-maxMessages) + removed = conversation[:overflow] + if system != nil { + retained = append(retained, system) + } + retained = append(retained, conversation[overflow:]...) + return retained, removed +} diff --git a/examples/02-agents/memory/step05_bounded_chat_history/main.go b/examples/02-agents/memory/step05_bounded_chat_history/main.go new file mode 100644 index 000000000..67fbb0600 --- /dev/null +++ b/examples/02-agents/memory/step05_bounded_chat_history/main.go @@ -0,0 +1,82 @@ +// Copyright (c) Microsoft. All rights reserved. + +// This sample archives older chat messages and recalls them from a vector store. +package main + +import ( + "context" + "fmt" + "log/slog" + + "github.com/google/uuid" + "github.com/microsoft/agent-framework-go/agent" + "github.com/microsoft/agent-framework-go/agent/memory" + "github.com/microsoft/agent-framework-go/embedding" + "github.com/microsoft/agent-framework-go/examples/internal/demo" + "github.com/microsoft/agent-framework-go/message/messagefilter" + "github.com/microsoft/agent-framework-go/provider/foundryprovider" + "github.com/microsoft/agent-framework-go/vectordata/inmemory" +) + +var logger = demo.NewLogger( + "Bounded Chat History", + "Demonstrates a four-message window with vector-backed overflow memory.", + "Model", demo.FoundryModel, + "Embedding Model", demo.FoundryEmbeddingModel, +) + +func main() { + ctx := context.Background() + credential := demo.FoundryTokenCredential() + generator := foundryprovider.NewEmbeddingGenerator(demo.FoundryProjectEndpoint, credential, + foundryprovider.EmbeddingGeneratorConfig{Model: demo.FoundryEmbeddingModel}) + store := inmemory.New(inmemory.Config{EmbeddingGenerator: embedding.Bind(generator)}) + + sessionID := uuid.NewString() + memories, err := memory.NewHistoryProvider(store.Store, "chathistory-overflow", 3072, + func(*agent.Session) *memory.State { + return &memory.State{ + StorageScope: memory.Scope{UserID: new("UID1"), SessionID: &sessionID}, + SearchScope: &memory.Scope{UserID: new("UID1")}, + } + }, memory.Config{ + SearchInputMessageFilter: messagefilter.PassThrough, + // Overflow includes assistant messages, not just external user input. + StorageInputRequestMessageFilter: messagefilter.PassThrough, + Logger: slog.New(logger), + }) + if err != nil { + demo.Panic(err) + } + + a := foundryprovider.NewAgent(demo.FoundryProjectEndpoint, credential, + foundryprovider.ModelDeployment(demo.FoundryModel), foundryprovider.AgentConfig{ + Instructions: "You are a helpful assistant. Answer questions concisely.", + DisableStoreOutput: true, + Config: agent.Config{ + Name: "Assistant", + HistoryProvider: newBoundedHistoryProvider(4, memories), + Middlewares: []agent.Middleware{logger}, + }, + }) + + session, err := a.CreateSession(ctx) + if err != nil { + demo.Panic(err) + } + for _, turn := range []struct{ heading, question string }{ + {"Filling the session window (4 messages max)", "My favorite color is blue."}, + {"", "I have a dog named Max."}, + {"Next exchange will trigger overflow to vector store", "What is the capital of France?"}, + {"Asking about overflowed information", "What is my favorite color?"}, + } { + if turn.heading != "" { + fmt.Printf(">> %s\n\n", turn.heading) + } + response, err := a.RunText(ctx, turn.question, agent.WithSession(session)).Collect() + if err != nil { + demo.Panic(err) + } + demo.Response(response, nil) + } +} diff --git a/examples/02-agents/providers/azure/ai_search/README.md b/examples/02-agents/providers/azure/ai_search/README.md new file mode 100644 index 000000000..b2f60ec10 --- /dev/null +++ b/examples/02-agents/providers/azure/ai_search/README.md @@ -0,0 +1,28 @@ +# Azure AI Search + +Stores text embeddings in a temporary Azure AI Search index, then runs vector and +hybrid searches with a metadata filter. The index is deleted on exit; the Search +service and Foundry deployments are left untouched. + +## Run + +Sign in with `az login` and configure existing resources: + +```powershell +$env:AZURE_AI_SEARCH_ENDPOINT = "https://.search.windows.net" +$env:FOUNDRY_PROJECT_ENDPOINT = "https://.services.ai.azure.com/api/projects/" +$env:FOUNDRY_EMBEDDING_MODEL = "text-embedding-3-large" + +go run .\examples\02-agents\providers\azure\ai_search +``` + +Authentication uses the shared `demo.FoundryTokenCredential` helper +(`DefaultAzureCredential`), like the neighboring Foundry examples. + +Search access requires permission to manage indexes and read/write their +documents, such as **Search Service Contributor** and **Search Index Data +Contributor**, along with access to the Foundry embedding deployment. + +The schema expects **3072-dimensional embeddings**. Runs incur embedding usage; +the Search service may also incur charges depending on its tier. A Free Search +service is sufficient for this sample. diff --git a/examples/02-agents/providers/azure/ai_search/main.go b/examples/02-agents/providers/azure/ai_search/main.go new file mode 100644 index 000000000..3ee5fa4f7 --- /dev/null +++ b/examples/02-agents/providers/azure/ai_search/main.go @@ -0,0 +1,118 @@ +// Copyright (c) Microsoft. All rights reserved. + +// This sample stores generated embeddings in Azure AI Search and runs scoped +// vector and hybrid queries. Its temporary index is deleted before exit. +package main + +import ( + "context" + "errors" + "fmt" + "os" + "strings" + "time" + + "github.com/google/uuid" + "github.com/microsoft/agent-framework-go/embedding" + "github.com/microsoft/agent-framework-go/examples/internal/demo" + "github.com/microsoft/agent-framework-go/provider/azureprovider" + "github.com/microsoft/agent-framework-go/provider/foundryprovider" + "github.com/microsoft/agent-framework-go/vectordata" +) + +type document struct { + ID uuid.UUID `vectordata:"key,storage=id"` + Topic string `vectordata:"data,indexed,storage=topic"` + Text string `vectordata:"data,fulltext,storage=text"` + Content string `vectordata:"vector,dimensions=3072,storage=vector"` +} + +func main() { + ctx := context.Background() + demo.NewLogger( + "Azure AI Search", + "Demonstrates generated embeddings, metadata filters, and hybrid search.", + "Embedding Model", demo.FoundryEmbeddingModel, + ) + endpoint := strings.TrimSpace(os.Getenv("AZURE_AI_SEARCH_ENDPOINT")) + if endpoint == "" { + demo.Panic("AZURE_AI_SEARCH_ENDPOINT must be set") + } + credential := demo.FoundryTokenCredential() + generator := foundryprovider.NewEmbeddingGenerator(demo.FoundryProjectEndpoint, credential, + foundryprovider.EmbeddingGeneratorConfig{Model: demo.FoundryEmbeddingModel}) + search, err := azureprovider.NewAISearch(endpoint, azureprovider.AISearchConfig{ + Credential: credential, EmbeddingGenerator: embedding.Bind(generator), + }) + if err != nil { + demo.Panic(err) + } + name := "afgo-example-" + uuid.NewString() + collection, err := vectordata.GetCollection[uuid.UUID, *document](search.Store, name, nil) + if err != nil { + demo.Panic(err) + } + defer func() { + cleanupCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 30*time.Second) + err = errors.Join(err, collection.EnsureDeleted(cleanupCtx)) + cancel() + if err != nil { + demo.Panic(err) + } + }() + fmt.Printf("Temporary index: %s\n\n", name) + if err = collection.EnsureExists(ctx); err != nil { + return + } + var records []*document + for _, seed := range []struct{ topic, text string }{ + {"humor", "I like pirate jokes, especially jokes about parrots and ships."}, + {"travel", "I enjoy hiking trips in the mountains and scenic viewpoints."}, + {"cooking", "I like cooking pasta with fresh tomatoes and basil."}, + } { + records = append(records, &document{Topic: seed.topic, Text: seed.text, Content: seed.text}) + } + if err = collection.Upsert(ctx, records...); err != nil { + return + } + filter := vectordata.Filter{Operator: vectordata.Equal, Property: "Topic", Value: "humor"} + if err = waitForIndex(ctx, collection, filter); err != nil { + return + } + + fmt.Println("Vector search, scoped to humor:") + for result, searchErr := range collection.Search(ctx, "Tell me a pirate joke.", 2, &vectordata.SearchOptions{Filter: &filter}) { + if searchErr != nil { + err = searchErr + return + } + fmt.Println(result.Record.Text) + } + fmt.Println("\nHybrid search, scoped to humor:") + hybrid := collection.(vectordata.KeywordHybridSearcher[*document]) + for result, searchErr := range hybrid.HybridSearch(ctx, "pirate jokes", []string{"pirate", "parrot"}, 2, + &vectordata.HybridSearchOptions{Filter: &filter}) { + if searchErr != nil { + err = searchErr + return + } + fmt.Println(result.Record.Text) + } +} + +func waitForIndex(ctx context.Context, collection vectordata.Collection[uuid.UUID, *document], filter vectordata.Filter) error { + ctx, cancel := context.WithTimeout(ctx, time.Minute) + defer cancel() + ticker := time.NewTicker(500 * time.Millisecond) + defer ticker.Stop() + for { + for _, err := range collection.GetFiltered(ctx, filter, 1, nil) { + return err + } + select { + case <-ctx.Done(): + return fmt.Errorf("waiting for document indexing: %w", ctx.Err()) + case <-ticker.C: + } + } +} diff --git a/examples/02-agents/rag/README.md b/examples/02-agents/rag/README.md index 069074505..8d6b58001 100644 --- a/examples/02-agents/rag/README.md +++ b/examples/02-agents/rag/README.md @@ -26,10 +26,8 @@ $env:FOUNDRY_EMBEDDING_MODEL = "text-embedding-3-large" The examples use the shared `demo.FoundryModel`, which defaults to `gpt-4o-mini` when `FOUNDRY_MODEL` is unset or blank. Set it as above to match the upstream sample's `gpt-5.4-mini` choice. Steps 01 and 02 also use -`FOUNDRY_EMBEDDING_MODEL`. Set it to `text-embedding-3-large` as above for their -3072-dimensional vector schemas. The shared demo default, `text-embedding-ada-002`, -produces 1536-dimensional vectors; using it requires changing the corresponding -store/schema dimensions to 1536. Deployment names must match those in your resource. +`FOUNDRY_EMBEDDING_MODEL`, which defaults to `text-embedding-3-large` for their +3072-dimensional vector schemas. Deployment names must match those in your resource. For production, choose a specific credential appropriate to the deployment, rather than the development-oriented `DefaultAzureCredential` chain. diff --git a/examples/02-agents/rag/step01_basic_text_rag/main.go b/examples/02-agents/rag/step01_basic_text_rag/main.go index 5427b07ba..386c4ec14 100644 --- a/examples/02-agents/rag/step01_basic_text_rag/main.go +++ b/examples/02-agents/rag/step01_basic_text_rag/main.go @@ -32,9 +32,7 @@ func main() { credential := demo.FoundryTokenCredential() generator := foundryprovider.NewEmbeddingGenerator(demo.FoundryProjectEndpoint, credential, foundryprovider.EmbeddingGeneratorConfig{Model: demo.FoundryEmbeddingModel}) - collection, err := inmemory.NewCollection[uuid.UUID, TextSearchDocument]("documents", vectordata.CollectionOptions{ - EmbeddingGenerator: embedding.Bind(generator), - }) + collection, err := inmemory.NewCollection[uuid.UUID, TextSearchDocument]("documents", &vectordata.Definition{EmbeddingGenerator: embedding.Bind(generator)}) if err != nil { demo.Panic(err) } diff --git a/examples/internal/demo/foundry.go b/examples/internal/demo/foundry.go index d4253c78d..84d15fd9b 100644 --- a/examples/internal/demo/foundry.go +++ b/examples/internal/demo/foundry.go @@ -13,7 +13,7 @@ import ( var ( FoundryProjectEndpoint = strings.TrimSpace(os.Getenv("FOUNDRY_PROJECT_ENDPOINT")) FoundryModel = cmp.Or(strings.TrimSpace(os.Getenv("FOUNDRY_MODEL")), DefaultDeployment) - FoundryEmbeddingModel = cmp.Or(os.Getenv("FOUNDRY_EMBEDDING_MODEL"), "text-embedding-ada-002") + FoundryEmbeddingModel = cmp.Or(os.Getenv("FOUNDRY_EMBEDDING_MODEL"), "text-embedding-3-large") ) func FoundryTokenCredential() *azidentity.DefaultAzureCredential { diff --git a/internal/telemetry/featureusage.go b/internal/telemetry/featureusage.go index fd1e1772b..f7bd2bd83 100644 --- a/internal/telemetry/featureusage.go +++ b/internal/telemetry/featureusage.go @@ -18,32 +18,33 @@ import ( // Feature indexes match the .NET version-1 registry. Only supported capabilities // registered there are tracked; indexes must not be repurposed for Go-only features. const ( - FeatureAgent = 0 - FeatureWorkflow = 2 - FeatureToolApproval = 3 - FeatureFileMemoryProvider = 5 - FeatureTextSearchProvider = 6 - FeatureInMemoryHistoryProvider = 13 - FeatureCompactionProvider = 9 - FeatureSkillsProvider = 8 - FeatureFileSkillsSource = 15 - FeatureInMemorySkillsSource = 16 - FeatureMCP = 14 - FeatureTodoProvider = 10 - FeatureAgentModeProvider = 11 - FeatureSequentialOrchestration = 32 - FeatureConcurrentOrchestration = 33 - FeatureGroupChatOrchestration = 34 - FeatureFoundryChatClient = 48 - FeatureFoundryAgent = 49 - FeatureFoundryMemory = 50 - FeatureOpenAI = 54 - FeatureAnthropic = 55 - FeatureGitHubCopilot = 57 - FeatureA2A = 62 - FeatureHostingAGUI = 63 - FeatureShell = 69 - FeatureHostingA2A = 73 + FeatureAgent = 0 + FeatureWorkflow = 2 + FeatureToolApproval = 3 + FeatureChatHistoryMemoryProvider = 4 + FeatureFileMemoryProvider = 5 + FeatureTextSearchProvider = 6 + FeatureInMemoryHistoryProvider = 13 + FeatureCompactionProvider = 9 + FeatureSkillsProvider = 8 + FeatureFileSkillsSource = 15 + FeatureInMemorySkillsSource = 16 + FeatureMCP = 14 + FeatureTodoProvider = 10 + FeatureAgentModeProvider = 11 + FeatureSequentialOrchestration = 32 + FeatureConcurrentOrchestration = 33 + FeatureGroupChatOrchestration = 34 + FeatureFoundryChatClient = 48 + FeatureFoundryAgent = 49 + FeatureFoundryMemory = 50 + FeatureOpenAI = 54 + FeatureAnthropic = 55 + FeatureGitHubCopilot = 57 + FeatureA2A = 62 + FeatureHostingAGUI = 63 + FeatureShell = 69 + FeatureHostingA2A = 73 ) const featureMaskDisabledEnvVar = "AGENT_FRAMEWORK_FEATURE_MASK_DISABLED" diff --git a/internal/telemetry/featureusage_activation_test.go b/internal/telemetry/featureusage_activation_test.go index f0fe21c4b..8599db70b 100644 --- a/internal/telemetry/featureusage_activation_test.go +++ b/internal/telemetry/featureusage_activation_test.go @@ -40,9 +40,11 @@ import ( "github.com/microsoft/agent-framework-go/agent/harness/todo" "github.com/microsoft/agent-framework-go/agent/harness/toolapproval" "github.com/microsoft/agent-framework-go/agent/harness/toolautocall" + "github.com/microsoft/agent-framework-go/agent/memory" "github.com/microsoft/agent-framework-go/agent/skills" "github.com/microsoft/agent-framework-go/agent/skills/fsskills" "github.com/microsoft/agent-framework-go/agent/textsearch" + "github.com/microsoft/agent-framework-go/embedding" "github.com/microsoft/agent-framework-go/internal/agenttest" "github.com/microsoft/agent-framework-go/internal/telemetry" "github.com/microsoft/agent-framework-go/message" @@ -57,6 +59,7 @@ import ( "github.com/microsoft/agent-framework-go/tool/functool" "github.com/microsoft/agent-framework-go/tool/mcptool" "github.com/microsoft/agent-framework-go/tool/shelltool" + "github.com/microsoft/agent-framework-go/vectordata/inmemory" "github.com/microsoft/agent-framework-go/workflow" "github.com/microsoft/agent-framework-go/workflow/agentworkflow" "github.com/modelcontextprotocol/go-sdk/mcp" @@ -74,6 +77,8 @@ var featureActivationCases = []struct { {"agent-session", []int{0, 13}, nil}, {"history-provide", []int{13}, nil}, {"history-store", []int{13}, nil}, + {"chat-history-memory", []int{4}, nil}, + {"chat-history-memory-on-demand", []int{4}, nil}, {"text-search", []int{6}, nil}, {"text-search-on-demand", []int{6}, nil}, {"compaction-context", []int{9}, nil}, @@ -320,6 +325,16 @@ func (activationCompactionStrategy) Compact(context.Context, *compaction.Message return false, nil } +type activationMemoryGenerator struct{} + +func (activationMemoryGenerator) Generate(_ context.Context, inputs []string, _ *embedding.Options) (*embedding.Generated[embedding.VectorOf[float32]], error) { + result := &embedding.Generated[embedding.VectorOf[float32]]{} + for range inputs { + result.Embeddings = append(result.Embeddings, embedding.VectorOf[float32]{1}) + } + return result, nil +} + func runFeatureActivation(t *testing.T, name string) { t.Helper() if strings.HasPrefix(name, "port-provider-") { @@ -420,6 +435,21 @@ func runFeatureActivation(t *testing.T, name string) { _, err := p.Invoking(ctx, invoking) return err } + case "chat-history-memory", "chat-history-memory-on-demand": + config := memory.Config{} + if strings.HasSuffix(name, "on-demand") { + config.SearchTime = memory.OnDemandFunctionCalling + } + store := inmemory.New(inmemory.Config{EmbeddingGenerator: embedding.Bind(activationMemoryGenerator{})}) + p, err := memory.NewHistoryProvider(store.Store, "history", 1, + func(*agent.Session) *memory.State { return new(memory.State) }, config) + if err != nil { + t.Fatal(err) + } + activate = func() error { + _, _, err := p.Invoking(ctx, invoking) + return err + } case "text-search", "text-search-on-demand": config := textsearch.Config{} if strings.HasSuffix(name, "on-demand") { diff --git a/provider/azureprovider/README.md b/provider/azureprovider/README.md new file mode 100644 index 000000000..506d51398 --- /dev/null +++ b/provider/azureprovider/README.md @@ -0,0 +1,86 @@ +# Azure AI Search + +Azure AI Search indexes back typed and dynamic VectorData collections, including +the [chat-history memory provider](../../agent/memory/). + +```go +search, err := azureprovider.NewAISearch(endpoint, azureprovider.AISearchConfig{ + Credential: credential, + EmbeddingGenerator: embedding.Bind(generator), +}) +if err != nil { + return err +} + +collection, err := vectordata.GetCollection[uuid.UUID, *Document]( + search.Store, "documents", nil) +if err != nil { + return err +} +if err := collection.EnsureExists(ctx); err != nil { + return err +} +``` + +Here `endpoint` is `https://.search.windows.net`, `credential` is an +`azcore.TokenCredential`, and `generator` produces float32 embeddings. Set +`AISearchConfig.APIKey` to `azcore.NewKeyCredential(key)` instead of `Credential` +for API-key authentication. Azure Identity access to documents requires an +appropriate data role; index creation also requires permission to manage indexes. + +```go +type Document struct { + ID uuid.UUID `vectordata:"key,storage=id"` + Text string `vectordata:"data,indexed,fulltext,storage=text"` + Content string `vectordata:"vector,dimensions=3072,storage=vector"` +} +``` + +The example schema matches the default dimensions of `text-embedding-3-large`. +An upsert stores `Text` as text and generates `Content`'s vector. Generators run +once per vector property for each batch; property and definition settings override +the store default. Native float32 vectors need no generation. + +- Use `search.DynamicCollection(name, definition)` for `map[string]any` records. + On Go 1.27, `search.Collection[K, R](name, definition)` is also available. +- UUID keys are generated by default. Use pointer or map records to observe them + on the input; disable generation through the key definition when needed. +- Created indexes support filtering and sorting by the key field. +- Storage names come from VectorData definitions and tags, not record `json` tags. + Only schema properties are uploaded. Go `int` maps to `Edm.Int64`. +- Nil collection fields, including native vectors, are uploaded as empty arrays + without changing input records. Nullable scalar fields remain null. +- HNSW and cosine similarity are defaults; flat indexes, dot product, and + Euclidean distance are also supported. +- String queries without a generator use Azure AI Search vectorization. This + requires an existing index with a configured vectorizer; the connector does not + provision vectorizers. +- Vector fields are excluded from retrieval by default. Generated fields cannot + be retrieved with `IncludeVectors`. Hybrid search supports full-text data fields. + Score thresholds filter service-ranked candidates on the client. +- Batch writes can partially succeed. All per-document failures are returned; + generated input keys are not rolled back. Stores and collections do not close + borrowed credentials, generators, or transports. + +The connector uses the Azure `azcore` pipeline with REST API `2025-09-01`. +Behavior follows the +[pinned .NET connector](https://github.com/CommunityToolkit/AI/tree/7a48d308df63d85427a772f031616889905a0fdf/MEVD/src/AzureAISearch). + +See the [runnable example](../../examples/02-agents/providers/azure/ai_search/). + +## Live validation + +The integration tests create uniquely named indexes and delete only those indexes. +They require permission to manage indexes and documents. + +```powershell +$env:RUN_AZURE_AI_SEARCH_INTEGRATION_TESTS = "true" +$env:AZURE_SUBSCRIPTION_ID = "" +$env:AZURE_AI_SEARCH_ENDPOINT = "https://.search.windows.net" +go test -count=1 .\provider\azureprovider -run "^TestE2E_AISearch" +``` + +Set `AZURE_AI_SEARCH_API_KEY` to test key authentication instead of Entra ID. +Generated-vector and chat-history memory tests additionally use +`FOUNDRY_PROJECT_ENDPOINT` and a `FOUNDRY_EMBEDDING_MODEL` deployment producing +3072-dimensional vectors (default `text-embedding-3-large`). diff --git a/provider/azureprovider/aisearch.go b/provider/azureprovider/aisearch.go new file mode 100644 index 000000000..005dbfa34 --- /dev/null +++ b/provider/azureprovider/aisearch.go @@ -0,0 +1,174 @@ +// Copyright (c) Microsoft. All rights reserved. + +// Package azureprovider provides Azure service integrations. +package azureprovider + +import ( + "context" + "encoding/json" + "errors" + "iter" + "net/http" + "net/url" + "reflect" + "strings" + + "github.com/Azure/azure-sdk-for-go/sdk/azcore" + "github.com/microsoft/agent-framework-go/embedding" + "github.com/microsoft/agent-framework-go/vectordata" +) + +// AISearchConfig configures Azure AI Search. Set exactly one of Credential and APIKey. +type AISearchConfig struct { + Credential azcore.TokenCredential + APIKey *azcore.KeyCredential + + // Audience is the token audience. Empty uses https://search.azure.com. + Audience string + + // ClientOptions configures the Azure HTTP pipeline. Default transports and + // supplied *http.Client values reject cross-origin redirects. Other custom + // transports are responsible for their own redirect handling. + ClientOptions azcore.ClientOptions + + // EmbeddingGenerator supplies the default generator for collections. + // Definition and property generators take precedence. + EmbeddingGenerator *embedding.BoundGenerator +} + +// AISearch manages Azure AI Search indexes through its embedded [vectordata.Store]. +// Create it with [NewAISearch]. Pass the embedded field to provider-independent +// APIs. Methods support concurrent use. +// +// Keys are strings or uuid.UUID values. UUID keys are generated by default. +// Created indexes support filtering and sorting by the key field. +// Data fields support strings, booleans, signed integers, float32, float64, +// time.Time, nullable pointers to these types, and slices or arrays of scalars. +// Native vectors are []float32, embedding.VectorOf[float32], embedding.Of[float32], +// and *embedding.Of[float32]. Non-native vector fields require an embedding +// generator for upserts. String search inputs without a configured generator +// are sent to the index's server-side vectorizer. +// +// Property storage names, including vectordata storage tags, define index field +// names. Record-level json tags do not override them. Only schema properties are +// uploaded. Nil collection fields are uploaded as empty arrays; nullable scalar +// fields remain null. Retrieval creates independent records and excludes vectors +// by default. IncludeVectors is unsupported for generated vector fields. +// +// AISearch and its collections borrow credentials, generators, and transports +// and need no Close. Upserts and batch deletions can partially succeed; errors +// report failed documents. Generated keys are written to pointer and map inputs +// before upload. Plain structs are copied by value. Callers must not concurrently +// mutate inputs. +type AISearch struct { + *vectordata.Store +} + +type store struct { + client *serviceClient + generator *embedding.BoundGenerator +} + +// NewAISearch creates a store for an Azure AI Search service endpoint, such as +// https://example.search.windows.net. It uses the 2025-09-01 REST API and the +// Azure SDK HTTP pipeline. Construction does not contact the service. +// Retry, logging, and transport settings use config.ClientOptions. +func NewAISearch(endpoint string, config AISearchConfig) (*AISearch, error) { + client, err := newServiceClient(endpoint, config) + if err != nil { + return nil, err + } + provider := &store{client: client, generator: config.EmbeddingGenerator} + return &AISearch{Store: vectordata.NewStore(provider)}, nil +} + +func (s *store) OpenCollection(name string, keyType, recordType reflect.Type, definition *vectordata.Definition) (vectordata.Collection[any, any], error) { + return newCollection(s, name, keyType, recordType, definition) +} + +func (s *store) ListCollectionNames(ctx context.Context) iter.Seq2[string, error] { + return func(yield func(string, error) bool) { + endpoint := s.client.url("/indexes", url.Values{"$select": {"name"}}) + for { + var page struct { + Value []struct { + Name string `json:"name"` + } `json:"value"` + NextLink string `json:"@odata.nextLink"` + } + if err := s.client.request(ctx, http.MethodGet, endpoint, nil, &page, http.StatusOK); err != nil { + yield("", failure("", "GetIndexNames", err)) + return + } + if page.Value == nil { + yield("", failure("", "GetIndexNames", errors.New("response is missing index names"))) + return + } + for _, index := range page.Value { + if err := ctx.Err(); err != nil { + yield("", failure("", "GetIndexNames", err)) + return + } + if index.Name == "" { + yield("", failure("", "GetIndexNames", errors.New("response contains an index without a name"))) + return + } + if !yield(index.Name, nil) { + return + } + } + if page.NextLink == "" { + return + } + var err error + endpoint, err = s.client.nextURL(endpoint, page.NextLink) + if err != nil { + yield("", failure("", "GetIndexNames", err)) + return + } + } + } +} + +func (s *store) CollectionExists(ctx context.Context, name string) (bool, error) { + if strings.TrimSpace(name) == "" { + return false, failure(name, "GetIndex", errors.New("collection name is required")) + } + var index json.RawMessage + err := s.client.request(ctx, http.MethodGet, s.client.url(indexPath(name), nil), nil, &index, http.StatusOK) + if isStatus(err, http.StatusNotFound) { + return false, nil + } + if err != nil { + return false, failure(name, "GetIndex", err) + } + return true, nil +} + +func (s *store) EnsureCollectionDeleted(ctx context.Context, name string) error { + if strings.TrimSpace(name) == "" { + return failure(name, "DeleteIndex", errors.New("collection name is required")) + } + err := s.client.request(ctx, http.MethodDelete, s.client.url(indexPath(name), nil), nil, nil, http.StatusNoContent) + if err != nil && !isStatus(err, http.StatusNotFound) { + return failure(name, "DeleteIndex", err) + } + return nil +} + +func indexPath(name string) string { + return "/indexes('" + url.PathEscape(strings.ReplaceAll(name, "'", "''")) + "')" +} + +func isStatus(err error, status int) bool { + responseErr, ok := errors.AsType[*azcore.ResponseError](err) + return ok && responseErr.StatusCode == status +} + +func failure(name, operation string, err error) error { + return &vectordata.Error{ + SystemName: "azure.aisearch", CollectionName: name, OperationName: operation, Err: err, + } +} + +var _ vectordata.StoreProvider = (*store)(nil) diff --git a/provider/azureprovider/aisearch_e2e_test.go b/provider/azureprovider/aisearch_e2e_test.go new file mode 100644 index 000000000..6f7611daa --- /dev/null +++ b/provider/azureprovider/aisearch_e2e_test.go @@ -0,0 +1,392 @@ +// Copyright (c) Microsoft. All rights reserved. + +package azureprovider_test + +import ( + "bytes" + "cmp" + "context" + "log/slog" + "os" + "reflect" + "slices" + "strings" + "testing" + "time" + + "github.com/Azure/azure-sdk-for-go/sdk/azcore" + "github.com/Azure/azure-sdk-for-go/sdk/azidentity" + "github.com/google/uuid" + "github.com/microsoft/agent-framework-go/agent" + "github.com/microsoft/agent-framework-go/agent/memory" + "github.com/microsoft/agent-framework-go/embedding" + "github.com/microsoft/agent-framework-go/internal/agenttest" + "github.com/microsoft/agent-framework-go/message" + "github.com/microsoft/agent-framework-go/provider/azureprovider" + "github.com/microsoft/agent-framework-go/provider/foundryprovider" + "github.com/microsoft/agent-framework-go/vectordata" +) + +func newLiveSearch(t *testing.T, generateEmbeddings bool) *azureprovider.AISearch { + t.Helper() + if !strings.EqualFold(os.Getenv("RUN_AZURE_AI_SEARCH_INTEGRATION_TESTS"), "true") { + t.Skip("RUN_AZURE_AI_SEARCH_INTEGRATION_TESTS is not true") + } + endpoint := strings.TrimSpace(os.Getenv("AZURE_AI_SEARCH_ENDPOINT")) + if endpoint == "" { + t.Fatal("AZURE_AI_SEARCH_ENDPOINT is required") + } + credential, err := azidentity.NewAzureCLICredential(&azidentity.AzureCLICredentialOptions{ + Subscription: os.Getenv("AZURE_SUBSCRIPTION_ID"), + }) + if err != nil { + t.Fatal(err) + } + config := azureprovider.AISearchConfig{Credential: credential} + if key := os.Getenv("AZURE_AI_SEARCH_API_KEY"); key != "" { + config.Credential = nil + config.APIKey = azcore.NewKeyCredential(key) + } + if generateEmbeddings { + projectEndpoint := strings.TrimSpace(os.Getenv("FOUNDRY_PROJECT_ENDPOINT")) + if projectEndpoint == "" { + t.Skip("FOUNDRY_PROJECT_ENDPOINT is required for generated-vector and memory tests") + } + config.EmbeddingGenerator = embedding.Bind(foundryprovider.NewEmbeddingGenerator(projectEndpoint, credential, + foundryprovider.EmbeddingGeneratorConfig{ + Model: cmp.Or(os.Getenv("FOUNDRY_EMBEDDING_MODEL"), "text-embedding-3-large"), + })) + } + search, err := azureprovider.NewAISearch(endpoint, config) + if err != nil { + t.Fatal(err) + } + return search +} + +func liveIndex(t *testing.T, search *azureprovider.AISearch) string { + t.Helper() + name := "afgo-live-" + uuid.NewString() + t.Logf("Temporary Search index: %s", name) + t.Cleanup(func() { + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + if err := search.EnsureCollectionDeleted(ctx, name); err != nil { + t.Errorf("delete temporary index %s: %v", name, err) + } + }) + return name +} + +func waitForLiveIndex(t *testing.T, ready func(context.Context) (bool, error)) { + t.Helper() + ctx, cancel := context.WithTimeout(t.Context(), time.Minute) + defer cancel() + ticker := time.NewTicker(500 * time.Millisecond) + defer ticker.Stop() + for { + ok, err := ready(ctx) + if err != nil { + t.Fatal(err) + } + if ok { + return + } + select { + case <-ctx.Done(): + t.Fatalf("waiting for Search indexing: %v", ctx.Err()) + case <-ticker.C: + } + } +} + +type liveDocument struct { + ID uuid.UUID `vectordata:"key,storage=id"` + Text string `vectordata:"data,indexed,fulltext,storage=text"` + Bucket string `vectordata:"data,indexed,storage=bucket"` + Rank int32 `vectordata:"data,indexed,storage=rank"` + When time.Time `vectordata:"data,storage=when"` + Vector []float32 `vectordata:"vector,dimensions=3,storage=vector"` +} + +func TestE2E_AISearchNativeCRUD(t *testing.T) { + search := newLiveSearch(t, false) + name := liveIndex(t, search) + collection, err := vectordata.GetCollection[uuid.UUID, *liveDocument](search.Store, name, nil) + if err != nil { + t.Fatal(err) + } + for range 2 { + if err := collection.EnsureExists(t.Context()); err != nil { + t.Fatal(err) + } + } + names := collect(t, search.ListCollectionNames(t.Context())) + if !slices.Contains(names, name) { + t.Fatalf("created index %s was not listed", name) + } + now := time.Now().UTC().Truncate(time.Millisecond) + records := []*liveDocument{ + {Text: "A pirate joke about a ship.", Bucket: "alpha", Rank: 1, When: now, Vector: []float32{1, 0, 0}}, + {Text: "A pirate joke about a parrot.", Bucket: "alpha", Rank: 2, When: now, Vector: []float32{0.8, 0.2, 0}}, + {Text: "A private joke from another bucket.", Bucket: "beta", Rank: 3, When: now, Vector: []float32{1, 0, 0}}, + } + if err := collection.Upsert(t.Context(), records...); err != nil { + t.Fatal(err) + } + filter := vectordata.Filter{Operator: vectordata.Equal, Property: "Bucket", Value: "alpha"} + waitForLiveIndex(t, func(ctx context.Context) (bool, error) { + count := 0 + for _, err := range collection.GetFiltered(ctx, filter, 3, nil) { + if err != nil { + return false, err + } + count++ + } + return count == 2, nil + }) + got, found, err := collection.Get(t.Context(), records[0].ID, &vectordata.RecordRetrievalOptions{IncludeVectors: true}) + if err != nil || !found || !reflect.DeepEqual(got.Vector, records[0].Vector) || !got.When.Equal(now) { + t.Fatalf("native retrieval = (%+v, %v, %v)", got, found, err) + } + batch := collect(t, vectordata.GetBatch(t.Context(), collection, []uuid.UUID{records[0].ID, uuid.New(), records[1].ID}, nil)) + if len(batch) != 2 || batch[0].Vector != nil || batch[0].ID != records[0].ID || batch[1].ID != records[1].ID { + t.Fatalf("batch retrieval = %v", batch) + } + results := collect(t, collection.Search(t.Context(), []float32{1, 0, 0}, 3, &vectordata.SearchOptions{Filter: &filter})) + if len(results) != 2 || results[0].Record.ID != records[0].ID || results[0].Score == nil { + t.Fatalf("scoped vector results = %v", results) + } + ordered := collect(t, collection.GetFiltered(t.Context(), filter, 2, + &vectordata.FilteredRecordRetrievalOptions{OrderBy: []vectordata.Sort{{Property: "Rank", Descending: true}}})) + if len(ordered) != 2 || ordered[0].Rank != 2 { + t.Fatalf("filtered ordering = %v", ordered) + } + wantIDs := []string{records[0].ID.String(), records[1].ID.String()} + slices.Sort(wantIDs) + for _, descending := range []bool{false, true} { + if descending { + slices.Reverse(wantIDs) + } + orderedByKey := collect(t, collection.GetFiltered(t.Context(), filter, 2, + &vectordata.FilteredRecordRetrievalOptions{OrderBy: []vectordata.Sort{{Property: "ID", Descending: descending}}})) + if len(orderedByKey) != len(wantIDs) { + t.Fatalf("key ordering returned %d records, want %d", len(orderedByKey), len(wantIDs)) + } + for i, record := range orderedByKey { + if record.ID.String() != wantIDs[i] { + t.Fatalf("key ordering descending=%v: record %d = %s, want %s", descending, i, record.ID, wantIDs[i]) + } + } + } + hybrid := collection.(vectordata.KeywordHybridSearcher[*liveDocument]) + hybridResults := collect(t, hybrid.HybridSearch(t.Context(), []float32{1, 0, 0}, []string{"pirate", "ship"}, 2, + &vectordata.HybridSearchOptions{Filter: &filter, IncludeVectors: true})) + if len(hybridResults) != 2 || len(hybridResults[0].Record.Vector) != 3 { + t.Fatalf("hybrid results = %v", hybridResults) + } + dynamic, err := search.DynamicCollection(name, &vectordata.Definition{Properties: []vectordata.Property{ + &vectordata.KeyProperty{Name: "ID", StorageName: "id", Type: reflect.TypeFor[uuid.UUID]()}, + &vectordata.DataProperty{Name: "Text", StorageName: "text", Type: reflect.TypeFor[string]()}, + &vectordata.VectorProperty{Name: "Vector", StorageName: "vector", Type: reflect.TypeFor[[]float32](), Dimensions: 3}, + }}) + if err != nil { + t.Fatal(err) + } + mapped, found, err := dynamic.Get(t.Context(), records[0].ID, &vectordata.RecordRetrievalOptions{IncludeVectors: true}) + if err != nil || !found || mapped["ID"] != records[0].ID || !reflect.DeepEqual(mapped["Vector"], records[0].Vector) { + t.Fatalf("dynamic retrieval = (%v, %v, %v)", mapped, found, err) + } + if err := vectordata.DeleteBatch(t.Context(), collection, []uuid.UUID{records[0].ID, records[1].ID}); err != nil { + t.Fatal(err) + } + waitForLiveIndex(t, func(ctx context.Context) (bool, error) { + for _, err := range collection.GetFiltered(ctx, filter, 1, nil) { + return false, err + } + return true, nil + }) + if err := collection.EnsureDeleted(t.Context()); err != nil { + t.Fatal(err) + } + if exists, err := collection.Exists(t.Context()); err != nil || exists { + t.Fatalf("deleted index exists = (%v, %v)", exists, err) + } +} + +func TestE2E_AISearchEmptyCollectionsAndMembership(t *testing.T) { + type record struct { + ID string `vectordata:"key,storage=id"` + Text *string `vectordata:"data,indexed,storage=text"` + Tags []string `vectordata:"data,storage=tags"` + Vector []float32 `vectordata:"vector,dimensions=3,storage=vector"` + } + search := newLiveSearch(t, false) + collection, err := vectordata.GetCollection[string, *record](search.Store, liveIndex(t, search), nil) + if err != nil { + t.Fatal(err) + } + if err := collection.EnsureExists(t.Context()); err != nil { + t.Fatal(err) + } + records := []*record{ + {ID: "empty", Text: new("")}, + {ID: "word", Text: new("word"), Tags: []string{}, Vector: []float32{}}, + {ID: "null", Tags: []string{"tag"}, Vector: []float32{1, 0, 0}}, + } + if err := collection.Upsert(t.Context(), records...); err != nil { + t.Fatal(err) + } + waitForLiveIndex(t, func(ctx context.Context) (bool, error) { + count := 0 + for _, err := range collection.GetFiltered(ctx, + vectordata.Filter{Operator: vectordata.In, Property: "ID", Value: []string{"empty", "word", "null"}}, 3, nil) { + if err != nil { + return false, err + } + count++ + } + return count == len(records), nil + }) + got, found, err := collection.Get(t.Context(), "empty", &vectordata.RecordRetrievalOptions{IncludeVectors: true}) + if err != nil || !found || got.Tags == nil || len(got.Tags) != 0 || got.Vector == nil || len(got.Vector) != 0 { + t.Fatalf("empty collections = (%+v, %v, %v)", got, found, err) + } + if records[0].Tags != nil || records[0].Vector != nil { + t.Fatalf("upsert changed nil input fields: %+v", records[0]) + } + for _, test := range []struct { + name string + filter vectordata.Filter + want []string + }{ + {"empty set", vectordata.Filter{Operator: vectordata.In, Property: "Text", Value: []string{}}, nil}, + {"empty string", vectordata.Filter{Operator: vectordata.In, Property: "Text", Value: []string{""}}, []string{"empty"}}, + {"empty and nonempty strings", vectordata.Filter{Operator: vectordata.In, Property: "Text", Value: []string{"", "word"}}, []string{"empty", "word"}}, + {"nonempty string", vectordata.Filter{Operator: vectordata.In, Property: "Text", Value: []string{"word"}}, []string{"word"}}, + {"negated membership", vectordata.Filter{Operator: vectordata.Not, Operands: []vectordata.Filter{ + {Operator: vectordata.In, Property: "Text", Value: []string{"", "word"}}, + }}, []string{"null"}}, + } { + t.Run(test.name, func(t *testing.T) { + var ids []string + for record, err := range collection.GetFiltered(t.Context(), test.filter, 3, nil) { + if err != nil { + t.Fatal(err) + } + ids = append(ids, record.ID) + } + slices.Sort(ids) + if !slices.Equal(ids, test.want) { + t.Fatalf("matched IDs = %v, want %v", ids, test.want) + } + }) + } +} + +type liveGeneratedDocument struct { + ID uuid.UUID `vectordata:"key"` + Text string `vectordata:"data,fulltext"` + User string `vectordata:"data,indexed"` + Content string `vectordata:"vector,dimensions=3072"` +} + +func TestE2E_AISearchGeneratedVectors(t *testing.T) { + search := newLiveSearch(t, true) + name := liveIndex(t, search) + collection, err := vectordata.GetCollection[uuid.UUID, *liveGeneratedDocument](search.Store, name, nil) + if err != nil { + t.Fatal(err) + } + if err := collection.EnsureExists(t.Context()); err != nil { + t.Fatal(err) + } + first := &liveGeneratedDocument{Text: "I enjoy pirate jokes.", User: "alpha", Content: "I enjoy pirate jokes."} + second := &liveGeneratedDocument{Text: "I enjoy hiking in the mountains.", User: "beta", Content: "I enjoy hiking in the mountains."} + if err := collection.Upsert(t.Context(), first, second); err != nil { + t.Fatal(err) + } + filter := vectordata.Filter{Operator: vectordata.Equal, Property: "User", Value: "alpha"} + waitForLiveIndex(t, func(ctx context.Context) (bool, error) { + for _, err := range collection.GetFiltered(ctx, filter, 1, nil) { + return err == nil, err + } + return false, nil + }) + results := collect(t, collection.Search(t.Context(), "Tell me a pirate joke.", 2, &vectordata.SearchOptions{Filter: &filter})) + if len(results) != 1 || results[0].Record.ID != first.ID || results[0].Record.Content != "" { + t.Fatalf("generated-vector search = %v", results) + } + hybrid := collection.(vectordata.KeywordHybridSearcher[*liveGeneratedDocument]) + hybridResults := collect(t, hybrid.HybridSearch(t.Context(), "pirate jokes", []string{"pirate"}, 2, + &vectordata.HybridSearchOptions{Filter: &filter})) + if len(hybridResults) != 1 || hybridResults[0].Record.ID != first.ID { + t.Fatalf("generated hybrid search = %v", hybridResults) + } + type nativeView struct { + ID uuid.UUID `vectordata:"key"` + Content []float32 `vectordata:"vector,dimensions=3072"` + } + native, err := vectordata.GetCollection[uuid.UUID, nativeView](search.Store, name, nil) + if err != nil { + t.Fatal(err) + } + value, found, err := native.Get(t.Context(), first.ID, &vectordata.RecordRetrievalOptions{IncludeVectors: true}) + if err != nil || !found || len(value.Content) != 3072 { + t.Fatalf("stored embedding dimensions = %d, found=%v, error=%v", len(value.Content), found, err) + } +} + +func TestE2E_AISearchChatHistoryMemory(t *testing.T) { + search := newLiveSearch(t, true) + name := liveIndex(t, search) + var diagnostics bytes.Buffer + provider, err := memory.NewHistoryProvider(search.Store, name, 3072, func(*agent.Session) *memory.State { + return &memory.State{ + StorageScope: memory.Scope{UserID: new("alpha"), SessionID: new(uuid.NewString())}, + SearchScope: &memory.Scope{UserID: new("alpha")}, + } + }, memory.Config{Logger: slog.New(slog.NewJSONHandler(&diagnostics, nil))}) + if err != nil { + t.Fatal(err) + } + first := agenttest.CreateSession() + if err := provider.Invoked(t.Context(), agent.InvokedContext{ + RequestMessages: []*message.Message{message.NewText("I like pirate jokes.")}, + ResponseMessages: []*message.Message{message.NewText("Here is a pirate joke.")}, + Options: []agent.Option{agent.WithSession(first)}, + }); err != nil { + t.Fatal(err) + } + if strings.Contains(diagnostics.String(), `"level":"ERROR"`) { + t.Fatalf("memory storage diagnostics: %s", diagnostics.String()) + } + type storedMemory struct { + Key uuid.UUID `vectordata:"key"` + UserID string `vectordata:"data,indexed,storage=UserId"` + Content string `vectordata:"data,fulltext"` + } + view, err := vectordata.GetCollection[uuid.UUID, storedMemory](search.Store, name, nil) + if err != nil { + t.Fatal(err) + } + waitForLiveIndex(t, func(ctx context.Context) (bool, error) { + count := 0 + for _, err := range view.GetFiltered(ctx, vectordata.Filter{Operator: vectordata.Equal, Property: "UserID", Value: "alpha"}, 3, nil) { + if err != nil { + return false, err + } + count++ + } + return count == 2, nil + }) + reader := agenttest.CreateSession() + messages, _, err := provider.Invoking(t.Context(), agent.InvokingContext{ + Messages: []*message.Message{message.NewText("What kind of jokes do I like?")}, + Options: []agent.Option{agent.WithSession(reader)}, + }) + if err != nil || len(messages) != 2 || !strings.Contains(messages[1].String(), "I like pirate jokes.") || + strings.Contains(diagnostics.String(), `"level":"ERROR"`) { + t.Fatalf("cross-session memory = (%v, %v); diagnostics: %s", messages, err, diagnostics.String()) + } +} diff --git a/provider/azureprovider/aisearch_go127_test.go b/provider/azureprovider/aisearch_go127_test.go new file mode 100644 index 000000000..84dc848f5 --- /dev/null +++ b/provider/azureprovider/aisearch_go127_test.go @@ -0,0 +1,26 @@ +// Copyright (c) Microsoft. All rights reserved. + +//go:build go1.27 + +package azureprovider_test + +import ( + "net/http" + "testing" + + "github.com/microsoft/agent-framework-go/provider/azureprovider" +) + +func TestAISearch_GenericCollectionMethod(t *testing.T) { + store := newStore(t, func(*http.Request) (*http.Response, error) { + return response(http.StatusOK, `{"id":"one","text":"saved"}`), nil + }, azureprovider.AISearchConfig{}) + collection, err := store.Collection[string, document]("records", nil) + if err != nil { + t.Fatal(err) + } + value, found, err := collection.Get(t.Context(), "one", nil) + if err != nil || !found || value.Text != "saved" { + t.Fatalf("generic method retrieval = (%v, %v, %v)", value, found, err) + } +} diff --git a/provider/azureprovider/aisearch_test.go b/provider/azureprovider/aisearch_test.go new file mode 100644 index 000000000..1db674560 --- /dev/null +++ b/provider/azureprovider/aisearch_test.go @@ -0,0 +1,367 @@ +// Copyright (c) Microsoft. All rights reserved. + +package azureprovider_test + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "io" + "net/http" + "reflect" + "slices" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/Azure/azure-sdk-for-go/sdk/azcore" + "github.com/Azure/azure-sdk-for-go/sdk/azcore/policy" + "github.com/google/uuid" + "github.com/microsoft/agent-framework-go/provider/azureprovider" + "github.com/microsoft/agent-framework-go/vectordata" +) + +type transportFunc func(*http.Request) (*http.Response, error) + +func (f transportFunc) Do(request *http.Request) (*http.Response, error) { return f(request) } + +func response(status int, body string) *http.Response { + return &http.Response{ + StatusCode: status, Header: http.Header{"Content-Type": {"application/json"}}, + Body: io.NopCloser(strings.NewReader(body)), + } +} + +func jsonBody(t *testing.T, request *http.Request) map[string]any { + t.Helper() + var body map[string]any + if err := json.NewDecoder(request.Body).Decode(&body); err != nil { + t.Fatal(err) + } + return body +} + +func newStore(t *testing.T, transport transportFunc, config azureprovider.AISearchConfig) *azureprovider.AISearch { + t.Helper() + if config.Credential == nil && config.APIKey == nil { + config.APIKey = azcore.NewKeyCredential("test-key") + } + config.ClientOptions.Transport = transportFunc(func(request *http.Request) (*http.Response, error) { + if request.URL.Query().Get("api-version") != "2025-09-01" { + t.Errorf("API version = %q", request.URL.Query().Get("api-version")) + } + result, err := transport(request) + if result != nil { + result.Request = request + } + return result, err + }) + config.ClientOptions.Retry.MaxRetries = -1 + store, err := azureprovider.NewAISearch("https://example.search.windows.net", config) + if err != nil { + t.Fatal(err) + } + return store +} + +type document struct { + ID string `vectordata:"key,storage=id"` + Text string `vectordata:"data,indexed,fulltext,storage=text"` + Rank int32 `vectordata:"data,indexed,storage=rank"` + Tags []string `vectordata:"data,indexed,storage=tags"` + Vector []float32 `vectordata:"vector,dimensions=3,storage=vector"` + Other string +} + +func openCollection[K, R any](t *testing.T, store *azureprovider.AISearch, definition *vectordata.Definition) vectordata.Collection[K, R] { + t.Helper() + collection, err := vectordata.GetCollection[K, R](store.Store, "records", definition) + if err != nil { + t.Fatal(err) + } + return collection +} + +func collect[T any](t *testing.T, sequence func(func(T, error) bool)) []T { + t.Helper() + var values []T + for value, err := range sequence { + if err != nil { + t.Fatal(err) + } + values = append(values, value) + } + return values +} + +func sequenceError[T any](t *testing.T, sequence func(func(T, error) bool)) error { + t.Helper() + var failure error + for _, err := range sequence { + if failure != nil { + t.Fatal("iterator continued after an error") + } + failure = err + } + if failure == nil { + t.Fatal("iterator did not return an error") + } + return failure +} + +type credentialFunc func(context.Context, policy.TokenRequestOptions) (azcore.AccessToken, error) + +func (f credentialFunc) GetToken(ctx context.Context, options policy.TokenRequestOptions) (azcore.AccessToken, error) { + return f(ctx, options) +} + +func TestAISearch_AuthenticationAndValidation(t *testing.T) { + key := azcore.NewKeyCredential("initial") + for _, audience := range []string{"", "https://search.azure.us"} { + tokenCalls, requests := 0, 0 + credential := credentialFunc(func(_ context.Context, options policy.TokenRequestOptions) (azcore.AccessToken, error) { + tokenCalls++ + want := audience + if want == "" { + want = "https://search.azure.com" + } + if !slices.Equal(options.Scopes, []string{want + "/.default"}) { + t.Errorf("token scopes = %v", options.Scopes) + } + return azcore.AccessToken{Token: "token", ExpiresOn: time.Now().Add(time.Hour)}, nil + }) + store := newStore(t, func(request *http.Request) (*http.Response, error) { + requests++ + if request.Header.Get("Authorization") != "Bearer token" || request.Header.Get("api-key") != "" { + t.Errorf("token authentication headers = %v", request.Header) + } + return response(http.StatusOK, `{"name":"records"}`), nil + }, azureprovider.AISearchConfig{Credential: credential, Audience: audience}) + if tokenCalls != 0 || requests != 0 { + t.Fatal("constructor performed I/O") + } + if exists, err := store.CollectionExists(t.Context(), "records"); err != nil || !exists || tokenCalls != 1 || requests != 1 { + t.Fatalf("CollectionExists = (%v, %v), tokens=%d requests=%d", exists, err, tokenCalls, requests) + } + } + expectedKey := "initial" + store := newStore(t, func(request *http.Request) (*http.Response, error) { + if request.Header.Get("api-key") != expectedKey || request.Header.Get("Authorization") != "" { + t.Errorf("API key authentication headers = %v", request.Header) + } + return response(http.StatusOK, `{"name":"records"}`), nil + }, azureprovider.AISearchConfig{APIKey: key}) + for _, value := range []string{"initial", "rotated"} { + expectedKey = value + key.Update(value) + if _, err := store.CollectionExists(t.Context(), "records"); err != nil { + t.Fatal(err) + } + } + token := credentialFunc(func(context.Context, policy.TokenRequestOptions) (azcore.AccessToken, error) { + return azcore.AccessToken{}, errors.New("unused credential") + }) + for _, test := range []struct { + endpoint string + config azureprovider.AISearchConfig + }{ + {"", azureprovider.AISearchConfig{APIKey: key}}, + {"relative", azureprovider.AISearchConfig{APIKey: key}}, + {"https://user:secret@example.test", azureprovider.AISearchConfig{APIKey: key}}, + {"https://example.test/path", azureprovider.AISearchConfig{APIKey: key}}, + {"https://example.test?query=1", azureprovider.AISearchConfig{APIKey: key}}, + {"https://example.test#fragment", azureprovider.AISearchConfig{APIKey: key}}, + {"ftp://example.test", azureprovider.AISearchConfig{APIKey: key}}, + {"https://example.test", azureprovider.AISearchConfig{}}, + {"https://example.test", azureprovider.AISearchConfig{APIKey: key, Credential: token}}, + } { + if store, err := azureprovider.NewAISearch(test.endpoint, test.config); err == nil || store != nil { + t.Fatalf("accepted invalid configuration (%q, %#v)", test.endpoint, test.config) + } + } + if _, ok := any(store).(io.Closer); ok { + t.Fatal("store should not close borrowed resources") + } +} + +func TestAISearch_Lifecycle(t *testing.T) { + exists, gets, creates, deletes := false, 0, 0, 0 + var schema map[string]any + store := newStore(t, func(request *http.Request) (*http.Response, error) { + switch request.Method { + case http.MethodGet: + gets++ + if request.URL.Path != "/indexes('records')" { + t.Errorf("index path = %q", request.URL.Path) + } + if !exists { + return response(http.StatusNotFound, `{"error":{"code":"IndexNotFound","message":"missing"}}`), nil + } + return response(http.StatusOK, `{"name":"records"}`), nil + case http.MethodPost: + creates++ + if request.URL.Path != "/indexes" { + t.Errorf("create path = %q", request.URL.Path) + } + schema = jsonBody(t, request) + exists = true + return response(http.StatusCreated, `{"name":"records"}`), nil + case http.MethodDelete: + deletes++ + exists = false + return response(http.StatusNoContent, ""), nil + default: + t.Fatalf("unexpected method %s", request.Method) + return nil, nil + } + }, azureprovider.AISearchConfig{}) + collection := openCollection[string, document](t, store, nil) + if collection.Name() != "records" || gets != 0 || creates != 0 { + t.Fatal("opening a collection created it") + } + for range 2 { + if err := collection.EnsureExists(t.Context()); err != nil { + t.Fatal(err) + } + } + if gets != 2 || creates != 1 || schema["name"] != "records" { + t.Fatalf("gets=%d creates=%d schema=%v", gets, creates, schema) + } + fields := schema["fields"].([]any) + keyField := fields[0].(map[string]any) + vectorField := fields[4].(map[string]any) + if len(fields) != 5 || keyField["name"] != "id" || keyField["key"] != true || + keyField["type"] != "Edm.String" || keyField["filterable"] != true || keyField["sortable"] != true || + vectorField["name"] != "vector" || vectorField["dimensions"] != float64(3) || + vectorField["vectorSearchProfile"] != "vectorProfile" { + t.Fatalf("index fields = %v", fields) + } + vectorConfig := schema["vectorSearch"].(map[string]any) + algorithm := vectorConfig["algorithms"].([]any)[0].(map[string]any) + if algorithm["kind"] != "hnsw" || algorithm["name"] != "vectorAlgoConfig" || + algorithm["hnswParameters"].(map[string]any)["metric"] != "cosine" { + t.Fatalf("default vector index = %v", vectorConfig) + } + if err := collection.EnsureDeleted(t.Context()); err != nil { + t.Fatal(err) + } + if found, err := collection.Exists(t.Context()); err != nil || found || deletes != 1 { + t.Fatalf("deleted index exists = (%v, %v), deletes=%d", found, err, deletes) + } +} + +func TestAISearch_ListPaginationAndEarlyStop(t *testing.T) { + for _, stop := range []bool{false, true} { + calls := 0 + store := newStore(t, func(request *http.Request) (*http.Response, error) { + calls++ + if calls == 1 { + if request.URL.Query().Get("$select") != "name" { + t.Error("listing indexes did not project names") + } + return response(http.StatusOK, `{"value":[{"name":"one"}],"@odata.nextLink":"https://example.search.windows.net/indexes?api-version=2025-09-01&$skip=1"}`), nil + } + if request.URL.Query().Get("$skip") != "1" { + t.Error("continuation parameters were lost") + } + return response(http.StatusOK, `{"value":[{"name":"two"}]}`), nil + }, azureprovider.AISearchConfig{}) + sequence := store.ListCollectionNames(t.Context()) + if calls != 0 { + t.Fatal("listing started eagerly") + } + var names []string + for name, err := range sequence { + if err != nil { + t.Fatal(err) + } + names = append(names, name) + if stop { + break + } + } + want := []string{"one", "two"} + if stop { + want = want[:1] + } + if !slices.Equal(names, want) || calls != len(want) { + t.Fatalf("names=%v calls=%d", names, calls) + } + } +} + +func TestAISearch_ServiceErrorsAndCreateRace(t *testing.T) { + for _, code := range []string{"ResourceNameAlreadyInUse", "OtherConflict"} { + store := newStore(t, func(request *http.Request) (*http.Response, error) { + if request.Method == http.MethodGet { + return response(http.StatusNotFound, `{"error":{"code":"NotFound"}}`), nil + } + return response(http.StatusConflict, `{"error":{"code":"`+code+`","message":"conflict"}}`), nil + }, azureprovider.AISearchConfig{}) + err := openCollection[string, document](t, store, nil).EnsureExists(t.Context()) + if code == "ResourceNameAlreadyInUse" { + if err != nil { + t.Fatal(err) + } + } else { + var serviceErr *azcore.ResponseError + var vectorErr *vectordata.Error + if !errors.As(err, &serviceErr) || serviceErr.StatusCode != http.StatusConflict || + !errors.As(err, &vectorErr) || vectorErr.SystemName != "azure.aisearch" || + vectorErr.CollectionName != "records" || vectorErr.OperationName != "CreateIndex" { + t.Fatalf("service error was not preserved: %v", err) + } + } + } +} + +func TestAISearch_CollectionTypes(t *testing.T) { + calls := 0 + store := newStore(t, func(*http.Request) (*http.Response, error) { + calls++ + return nil, errors.New("unexpected HTTP request") + }, azureprovider.AISearchConfig{}) + definition := &vectordata.Definition{Properties: []vectordata.Property{ + &vectordata.KeyProperty{Name: "ID", Type: reflect.TypeFor[uuid.UUID]()}, + &vectordata.VectorProperty{Name: "Vector", Type: reflect.TypeFor[[]float32](), Dimensions: 3}, + }} + for _, recordType := range []func() error{ + func() error { + _, err := vectordata.GetCollection[string, document](store.Store, "records", nil) + return err + }, + func() error { + _, err := vectordata.GetCollection[string, *document](store.Store, "records", nil) + return err + }, + func() error { _, err := store.DynamicCollection("dynamic", definition); return err }, + } { + if err := recordType(); err != nil { + t.Fatal(err) + } + } + type unsupported struct { + ID int `vectordata:"key"` + } + if _, err := vectordata.GetCollection[int, unsupported](store.Store, "records", nil); err == nil { + t.Fatal("accepted integer key") + } + if _, err := vectordata.GetCollection[string, document](store.Store, " ", nil); err == nil { + t.Fatal("accepted blank collection name") + } + if _, err := store.DynamicCollection("dynamic", nil); err == nil { + t.Fatal("accepted dynamic collection without definition") + } + if calls != 0 || definition.Properties[0].(*vectordata.KeyProperty).StorageName != "" { + t.Fatal("collection construction performed I/O or mutated its definition") + } +} + +type trackingBody struct { + *bytes.Reader + closes atomic.Int32 +} + +func (b *trackingBody) Close() error { b.closes.Add(1); return nil } diff --git a/provider/azureprovider/client.go b/provider/azureprovider/client.go new file mode 100644 index 000000000..f709fa6b5 --- /dev/null +++ b/provider/azureprovider/client.go @@ -0,0 +1,155 @@ +// Copyright (c) Microsoft. All rights reserved. + +package azureprovider + +import ( + "context" + "errors" + "fmt" + "io" + "net/http" + "net/url" + "strings" + + "github.com/Azure/azure-sdk-for-go/sdk/azcore" + "github.com/Azure/azure-sdk-for-go/sdk/azcore/policy" + "github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime" +) + +const apiVersion = "2025-09-01" + +type serviceClient struct { + endpoint *url.URL + client *azcore.Client +} + +func newServiceClient(endpoint string, config AISearchConfig) (*serviceClient, error) { + parsed, err := url.Parse(strings.TrimSpace(endpoint)) + if err != nil { + return nil, fmt.Errorf("azureaisearch: invalid endpoint: %w", err) + } + if parsed.Host == "" || parsed.Scheme != "https" && parsed.Scheme != "http" || + parsed.User != nil || parsed.RawQuery != "" || parsed.Fragment != "" || + parsed.Path != "" && parsed.Path != "/" { + return nil, fmt.Errorf("azureaisearch: endpoint must be an absolute HTTP or HTTPS service URL without credentials, path, query, or fragment") + } + if (config.Credential == nil) == (config.APIKey == nil) { + return nil, fmt.Errorf("azureaisearch: set exactly one of Credential and APIKey") + } + var auth policy.Policy + if config.APIKey != nil { + auth = runtime.NewKeyCredentialPolicy(config.APIKey, "api-key", &runtime.KeyCredentialPolicyOptions{ + InsecureAllowCredentialWithHTTP: config.ClientOptions.InsecureAllowCredentialWithHTTP, + }) + } else { + audience := strings.TrimRight(config.Audience, "/") + if audience == "" { + audience = "https://search.azure.com" + } + auth = runtime.NewBearerTokenPolicy(config.Credential, []string{audience + "/.default"}, &policy.BearerTokenOptions{ + InsecureAllowCredentialWithHTTP: config.ClientOptions.InsecureAllowCredentialWithHTTP, + }) + } + options := config.ClientOptions + var transport *http.Client + switch configured := options.Transport.(type) { + case nil: + transport = &http.Client{} + case *http.Client: + clientCopy := *configured + transport = &clientCopy + } + if transport != nil { + checkRedirect := transport.CheckRedirect + transport.CheckRedirect = func(request *http.Request, via []*http.Request) error { + if !sameOrigin(parsed, request.URL) { + return fmt.Errorf("azureaisearch: redirects must remain on the configured service") + } + if checkRedirect != nil { + return checkRedirect(request, via) + } + if len(via) >= 10 { + return fmt.Errorf("azureaisearch: stopped after 10 redirects") + } + return nil + } + options.Transport = transport + } + client, err := azcore.NewClient( + "github.com/microsoft/agent-framework-go/provider/azureprovider", + "v0.1.0", + runtime.PipelineOptions{PerRetry: []policy.Policy{auth}}, + &options) + if err != nil { + return nil, err + } + return &serviceClient{endpoint: parsed, client: client}, nil +} + +func (c *serviceClient) url(path string, query url.Values) string { + if query == nil { + query = make(url.Values) + } + query.Set("api-version", apiVersion) + return strings.TrimRight(c.endpoint.String(), "/") + path + "?" + query.Encode() +} + +func (c *serviceClient) nextURL(current, next string) (string, error) { + base, err := url.Parse(current) + if err != nil { + return "", err + } + relative, err := url.Parse(next) + if err != nil { + return "", fmt.Errorf("azureaisearch: invalid continuation URL: %w", err) + } + resolved := base.ResolveReference(relative) + if !c.sameOrigin(resolved) { + return "", fmt.Errorf("azureaisearch: continuation URL must remain on the configured service") + } + return resolved.String(), nil +} + +func (c *serviceClient) sameOrigin(target *url.URL) bool { + return sameOrigin(c.endpoint, target) +} + +func sameOrigin(base, target *url.URL) bool { + return strings.EqualFold(target.Scheme, base.Scheme) && + strings.EqualFold(target.Host, base.Host) && target.User == nil && target.Fragment == "" +} + +func (c *serviceClient) request(ctx context.Context, method, endpoint string, body, result any, statusCodes ...int) (err error) { + if err := ctx.Err(); err != nil { + return err + } + target, err := url.Parse(endpoint) + if err != nil { + return err + } + if !c.sameOrigin(target) { + return fmt.Errorf("azureaisearch: request URL must remain on the configured service") + } + request, err := runtime.NewRequest(ctx, method, endpoint) + if err != nil { + return err + } + if body != nil { + if err := runtime.MarshalAsJSON(request, body); err != nil { + return err + } + } + response, err := c.client.Pipeline().Do(request) + if err != nil { + return err + } + defer func() { err = errors.Join(err, response.Body.Close()) }() + if !runtime.HasStatusCode(response, statusCodes...) { + return runtime.NewResponseError(response) + } + if result == nil { + _, err := io.Copy(io.Discard, response.Body) + return err + } + return runtime.UnmarshalAsJSON(response, result) +} diff --git a/provider/azureprovider/client_test.go b/provider/azureprovider/client_test.go new file mode 100644 index 000000000..a25c214c0 --- /dev/null +++ b/provider/azureprovider/client_test.go @@ -0,0 +1,140 @@ +// Copyright (c) Microsoft. All rights reserved. + +package azureprovider_test + +import ( + "context" + "errors" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + + "github.com/Azure/azure-sdk-for-go/sdk/azcore" + "github.com/microsoft/agent-framework-go/provider/azureprovider" +) + +func TestAISearch_RejectsCrossOriginRedirects(t *testing.T) { + var destinationCalls atomic.Int32 + destination := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { + destinationCalls.Add(1) + })) + defer destination.Close() + source := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + http.Redirect(writer, request, destination.URL, http.StatusTemporaryRedirect) + })) + defer source.Close() + for _, useCustomClient := range []bool{false, true} { + config := azureprovider.AISearchConfig{ + APIKey: azcore.NewKeyCredential("must-not-leak"), + ClientOptions: azcore.ClientOptions{ + InsecureAllowCredentialWithHTTP: true, + }, + } + config.ClientOptions.Retry.MaxRetries = -1 + client := &http.Client{} + if useCustomClient { + config.ClientOptions.Transport = client + } + search, err := azureprovider.NewAISearch(source.URL, config) + if err != nil { + t.Fatal(err) + } + if _, err := search.CollectionExists(t.Context(), "records"); err == nil || !strings.Contains(err.Error(), "redirects") { + t.Fatalf("cross-origin redirect = %v", err) + } + if destinationCalls.Load() != 0 || client.CheckRedirect != nil { + t.Fatal("redirect exposed credentials or modified the borrowed HTTP client") + } + } +} + +func TestAISearch_PreservesCustomRedirectPolicy(t *testing.T) { + var policyCalls atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + http.Redirect(writer, request, "/redirected", http.StatusTemporaryRedirect) + })) + defer server.Close() + wantErr := errors.New("custom redirect policy") + httpClient := &http.Client{CheckRedirect: func(*http.Request, []*http.Request) error { + policyCalls.Add(1) + return wantErr + }} + config := azureprovider.AISearchConfig{ + APIKey: azcore.NewKeyCredential("key"), + ClientOptions: azcore.ClientOptions{ + Transport: httpClient, InsecureAllowCredentialWithHTTP: true, + }, + } + config.ClientOptions.Retry.MaxRetries = -1 + search, err := azureprovider.NewAISearch(server.URL, config) + if err != nil { + t.Fatal(err) + } + _, err = search.CollectionExists(t.Context(), "records") + if !errors.Is(err, wantErr) || policyCalls.Load() != 1 { + t.Fatalf("custom redirect policy = %v, calls=%d", err, policyCalls.Load()) + } +} + +func TestAISearch_HTTPRequiresExplicitOptIn(t *testing.T) { + var calls atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { + calls.Add(1) + })) + defer server.Close() + search, err := azureprovider.NewAISearch(server.URL, azureprovider.AISearchConfig{APIKey: azcore.NewKeyCredential("key")}) + if err != nil { + t.Fatal(err) + } + if _, err := search.CollectionExists(t.Context(), "records"); err == nil { + t.Fatal("credentials were sent over HTTP without opting in") + } + if calls.Load() != 0 { + t.Fatal("insecure request reached the server") + } +} + +func TestAISearch_CancelInFlightRequest(t *testing.T) { + started := make(chan struct{}) + store := newStore(t, func(request *http.Request) (*http.Response, error) { + close(started) + <-request.Context().Done() + return nil, request.Context().Err() + }, azureprovider.AISearchConfig{}) + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + done := make(chan error, 1) + go func() { + _, err := store.CollectionExists(ctx, "records") + done <- err + }() + <-started + cancel() + if err := <-done; !errors.Is(err, context.Canceled) { + t.Fatalf("in-flight cancellation = %v", err) + } +} + +func TestAISearch_DeletionAndListFailures(t *testing.T) { + missing := newStore(t, func(*http.Request) (*http.Response, error) { + return response(http.StatusNotFound, `{"error":{"code":"NotFound"}}`), nil + }, azureprovider.AISearchConfig{}) + if err := missing.EnsureCollectionDeleted(t.Context(), "records"); err != nil { + t.Fatal(err) + } + for _, payload := range []string{ + `{}`, `{"value":[{}]}`, `{"value":[],"@odata.nextLink":"https://other.example/indexes"}`, + } { + calls := 0 + store := newStore(t, func(*http.Request) (*http.Response, error) { + calls++ + return response(http.StatusOK, payload), nil + }, azureprovider.AISearchConfig{}) + _ = sequenceError(t, store.ListCollectionNames(t.Context())) + if calls != 1 { + t.Fatalf("invalid index response triggered %d calls", calls) + } + } +} diff --git a/provider/azureprovider/collection.go b/provider/azureprovider/collection.go new file mode 100644 index 000000000..8b4f82623 --- /dev/null +++ b/provider/azureprovider/collection.go @@ -0,0 +1,196 @@ +// Copyright (c) Microsoft. All rights reserved. + +package azureprovider + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "net/url" + "strings" + + "github.com/Azure/azure-sdk-for-go/sdk/azcore" + "github.com/microsoft/agent-framework-go/vectordata" + "github.com/microsoft/agent-framework-go/vectordata/providerservices" +) + +func (c *collection) Name() string { return c.name } + +func (c *collection) Exists(ctx context.Context) (bool, error) { + return c.store.CollectionExists(ctx, c.name) +} + +func (c *collection) EnsureExists(ctx context.Context) error { + exists, err := c.Exists(ctx) + if err != nil || exists { + return err + } + definition, err := c.indexDefinition() + if err != nil { + return failure(c.name, "CreateIndex", err) + } + err = c.store.client.request(ctx, http.MethodPost, c.store.client.url("/indexes", nil), definition, nil, http.StatusCreated) + if responseErr, ok := errors.AsType[*azcore.ResponseError](err); ok && responseErr.ErrorCode == "ResourceNameAlreadyInUse" { + return nil + } + if err != nil { + return failure(c.name, "CreateIndex", err) + } + return nil +} + +func (c *collection) EnsureDeleted(ctx context.Context) error { + return c.store.EnsureCollectionDeleted(ctx, c.name) +} + +func (c *collection) Get(ctx context.Context, key any, options *vectordata.RecordRetrievalOptions) (any, bool, error) { + includeVectors := options != nil && options.IncludeVectors + if err := c.validateRetrieval(includeVectors); err != nil { + return nil, false, failure(c.name, "GetDocument", err) + } + stringKey, err := keyString(key) + if err != nil { + return nil, false, failure(c.name, "GetDocument", err) + } + query := make(url.Values) + if !includeVectors { + query.Set("$select", c.selectedFields()) + } + path := indexPath(c.name) + "/docs('" + url.PathEscape(strings.ReplaceAll(stringKey, "'", "''")) + "')" + var document map[string]json.RawMessage + if err := c.store.client.request(ctx, http.MethodGet, c.store.client.url(path, query), nil, &document, http.StatusOK); err != nil { + if isStatus(err, http.StatusNotFound) { + return nil, false, nil + } + return nil, false, failure(c.name, "GetDocument", err) + } + record, err := c.decodeRecord(document, includeVectors) + if err != nil { + return nil, false, failure(c.name, "GetDocument", err) + } + return record, true, nil +} + +func (c *collection) Delete(ctx context.Context, key any) error { + return c.DeleteBatch(ctx, []any{key}) +} + +func (c *collection) DeleteBatch(ctx context.Context, keys []any) error { + if err := ctx.Err(); err != nil { + return failure(c.name, "DeleteDocuments", err) + } + if len(keys) == 0 { + return nil + } + actions := make([]map[string]any, len(keys)) + for i, key := range keys { + value, err := keyString(key) + if err != nil { + return failure(c.name, "DeleteDocuments", err) + } + actions[i] = map[string]any{"@search.action": "delete", c.key.StorageName: value} + } + return c.indexDocuments(ctx, actions, "DeleteDocuments") +} + +func (c *collection) Upsert(ctx context.Context, records ...any) error { + if err := ctx.Err(); err != nil { + return failure(c.name, "UploadDocuments", err) + } + if len(records) == 0 { + return nil + } + owned := make([]any, len(records)) + for i, record := range records { + value, err := c.prepareRecord(record) + if err != nil { + return failure(c.name, "UploadDocuments", err) + } + owned[i] = value + } + generated := make(map[string][][]float32) + for _, property := range c.vectors { + if property.Type == property.EmbeddingType { + continue + } + result, err := providerservices.GenerateEmbeddings(ctx, property, owned) + if err != nil { + return failure(c.name, "UploadDocuments", err) + } + vectors := make([][]float32, len(result.Embeddings)) + for i, output := range result.Embeddings { + vector, ok := providerservices.AsVector[float32](output) + if !ok || vector == nil { + return failure(c.name, "UploadDocuments", fmt.Errorf("property %q: generator returned an invalid float32 embedding at index %d", property.Name, i)) + } + vectors[i] = vector + } + generated[property.Name] = vectors + } + actions := make([]map[string]any, len(records)) + for i, record := range owned { + document, err := c.encodeRecord(record, i, generated) + if err != nil { + return failure(c.name, "UploadDocuments", err) + } + document["@search.action"] = "upload" + actions[i] = document + } + return c.indexDocuments(ctx, actions, "UploadDocuments") +} + +func (c *collection) indexDocuments(ctx context.Context, actions []map[string]any, operation string) error { + body := struct { + Value []map[string]any `json:"value"` + }{Value: actions} + var response struct { + Value []struct { + Key string `json:"key"` + Status *bool `json:"status"` + StatusCode int `json:"statusCode"` + ErrorMessage string `json:"errorMessage"` + } `json:"value"` + } + if err := c.store.client.request(ctx, http.MethodPost, c.store.client.url(indexPath(c.name)+"/docs/index", nil), + body, &response, http.StatusOK, http.StatusMultiStatus); err != nil { + return failure(c.name, operation, err) + } + if len(response.Value) != len(actions) { + return failure(c.name, operation, fmt.Errorf("indexing response contains %d results for %d documents", len(response.Value), len(actions))) + } + var failures []error + for _, result := range response.Value { + if result.Status == nil { + failures = append(failures, fmt.Errorf("indexing result for document %q is missing its status", result.Key)) + } else if !*result.Status { + failures = append(failures, fmt.Errorf("document %q failed with status %d: %s", result.Key, result.StatusCode, result.ErrorMessage)) + } + } + if err := errors.Join(failures...); err != nil { + return failure(c.name, operation, err) + } + return nil +} + +func (c *collection) selectedFields() string { + names := make([]string, 0, len(c.data)+1) + names = append(names, c.key.StorageName) + for _, property := range c.data { + names = append(names, property.StorageName) + } + return strings.Join(names, ",") +} + +func (c *collection) validateRetrieval(includeVectors bool) error { + if includeVectors && c.model.EmbeddingGenerationRequired { + return fmt.Errorf("IncludeVectors is unsupported for collections with generated vector fields") + } + return nil +} + +var ( + _ vectordata.Collection[any, any] = (*collection)(nil) + _ vectordata.KeywordHybridSearcher[any] = (*collection)(nil) +) diff --git a/provider/azureprovider/collection_test.go b/provider/azureprovider/collection_test.go new file mode 100644 index 000000000..535ceefdd --- /dev/null +++ b/provider/azureprovider/collection_test.go @@ -0,0 +1,450 @@ +// Copyright (c) Microsoft. All rights reserved. + +package azureprovider_test + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "reflect" + "slices" + "strings" + "testing" + "time" + + "github.com/Azure/azure-sdk-for-go/sdk/azcore" + "github.com/google/uuid" + "github.com/microsoft/agent-framework-go/embedding" + "github.com/microsoft/agent-framework-go/provider/azureprovider" + "github.com/microsoft/agent-framework-go/vectordata" +) + +func indexedResponse(t *testing.T, actions []any, keyName string) *http.Response { + t.Helper() + results := make([]map[string]any, len(actions)) + for i, action := range actions { + results[i] = map[string]any{"key": action.(map[string]any)[keyName], "status": true, "statusCode": 200} + } + body, err := json.Marshal(map[string]any{"value": results}) + if err != nil { + t.Fatal(err) + } + return response(http.StatusOK, string(body)) +} + +func TestCollection_CRUDAndBatching(t *testing.T) { + for _, includeVectors := range []bool{false, true} { + t.Run(fmt.Sprintf("includeVectors=%t", includeVectors), func(t *testing.T) { + calls := 0 + var writes [][]any + store := newStore(t, func(request *http.Request) (*http.Response, error) { + calls++ + switch request.Method { + case http.MethodPost: + if request.URL.Path != "/indexes('records')/docs/index" { + t.Errorf("indexing path = %q", request.URL.Path) + } + actions := jsonBody(t, request)["value"].([]any) + writes = append(writes, actions) + return indexedResponse(t, actions, "id"), nil + case http.MethodGet: + selectFields := request.URL.Query().Get("$select") + if includeVectors && selectFields != "" || !includeVectors && selectFields != "id,text,rank,tags" { + t.Errorf("selected fields = %q", selectFields) + } + if strings.HasSuffix(request.URL.Path, "/docs('missing')") { + return response(http.StatusNotFound, `{"error":{"code":"NotFound"}}`), nil + } + return response(http.StatusOK, `{"id":"one","text":"saved","rank":42,"tags":["a","b"],"vector":[1,2,3],"unknown":"ignored"}`), nil + default: + t.Fatalf("unexpected method %s", request.Method) + return nil, nil + } + }, azureprovider.AISearchConfig{}) + collection := openCollection[string, document](t, store, nil) + input := []document{ + {ID: "one", Text: "first", Rank: 1, Tags: []string{"a"}, Vector: []float32{1, 2, 3}, Other: "not in schema"}, + {ID: "two", Text: "second", Vector: []float32{3, 2, 1}}, + } + before := slices.Clone(input) + if err := collection.Upsert(t.Context(), input...); err != nil { + t.Fatal(err) + } + if len(writes) != 1 || len(writes[0]) != 2 || calls != 1 || !reflect.DeepEqual(input, before) { + t.Fatalf("writes=%v calls=%d input=%v", writes, calls, input) + } + first := writes[0][0].(map[string]any) + if first["@search.action"] != "upload" || first["id"] != "one" || first["text"] != "first" || + !reflect.DeepEqual(first["vector"], []any{float64(1), float64(2), float64(3)}) { + t.Fatalf("upsert action = %v", first) + } + if _, found := first["Other"]; found { + t.Fatal("uploaded an unconfigured field") + } + options := &vectordata.RecordRetrievalOptions{IncludeVectors: includeVectors} + records := collect(t, vectordata.GetBatch(t.Context(), collection, []string{"one", "missing", "two"}, options)) + if len(records) != 2 || records[0].Text != "saved" || records[0].Rank != 42 || + !slices.Equal(records[0].Tags, []string{"a", "b"}) || records[0].Other != "" { + t.Fatalf("GetBatch = %v", records) + } + if includeVectors && !slices.Equal(records[0].Vector, []float32{1, 2, 3}) || + !includeVectors && records[0].Vector != nil { + t.Fatalf("retrieved vector = %v", records[0].Vector) + } + if err := vectordata.DeleteBatch(t.Context(), collection, []string{"one", "two"}); err != nil { + t.Fatal(err) + } + if len(writes) != 2 || len(writes[1]) != 2 { + t.Fatalf("deletion did not use a single batch: %v", writes) + } + for i, key := range []string{"one", "two"} { + if action := writes[1][i]; !reflect.DeepEqual(action, map[string]any{"@search.action": "delete", "id": key}) { + t.Fatalf("delete action = %v", action) + } + } + if err := collection.Delete(t.Context(), "one"); err != nil { + t.Fatal(err) + } + beforeCalls := calls + for _, batch := range [][]document{nil, {}} { + if err := collection.Upsert(t.Context(), batch...); err != nil { + t.Fatal(err) + } + } + for _, keys := range [][]string{nil, {}} { + if err := vectordata.DeleteBatch(t.Context(), collection, keys); err != nil { + t.Fatal(err) + } + collect(t, vectordata.GetBatch(t.Context(), collection, keys, options)) + } + if calls != beforeCalls { + t.Fatal("empty batches performed I/O") + } + }) + } +} + +func TestCollection_UUIDKeysAndDynamicRecords(t *testing.T) { + type entry struct { + ID uuid.UUID `vectordata:"key,storage=id"` + Text string `vectordata:"data,storage=text"` + } + var uploaded map[string]any + store := newStore(t, func(request *http.Request) (*http.Response, error) { + if request.Method == http.MethodPost { + actions := jsonBody(t, request)["value"].([]any) + uploaded = actions[0].(map[string]any) + return indexedResponse(t, actions, "id"), nil + } + body, err := json.Marshal(uploaded) + if err != nil { + t.Fatal(err) + } + return response(http.StatusOK, string(body)), nil + }, azureprovider.AISearchConfig{}) + pointers := openCollection[uuid.UUID, *entry](t, store, nil) + input := &entry{Text: "pointer"} + if err := pointers.Upsert(t.Context(), input); err != nil || input.ID == uuid.Nil || uploaded["id"] != input.ID.String() { + t.Fatalf("generated key = %v, document = %v, error = %v", input.ID, uploaded, err) + } + got, found, err := pointers.Get(t.Context(), input.ID, nil) + if err != nil || !found || !reflect.DeepEqual(got, input) || got == input { + t.Fatalf("pointer retrieval = (%v, %v, %v)", got, found, err) + } + values := openCollection[uuid.UUID, entry](t, store, nil) + value := entry{Text: "value"} + if err := values.Upsert(t.Context(), value); err != nil || value.ID != uuid.Nil { + t.Fatalf("value record mutation = (%v, %v)", value, err) + } + if key, err := uuid.Parse(uploaded["id"].(string)); err != nil || key == uuid.Nil { + t.Fatalf("stored struct key = (%v, %v)", key, err) + } + keyProperty := &vectordata.KeyProperty{Name: "ID", StorageName: "id", Type: reflect.TypeFor[uuid.UUID]()} + definition := &vectordata.Definition{Properties: []vectordata.Property{ + keyProperty, + &vectordata.DataProperty{Name: "Text", StorageName: "text", Type: reflect.TypeFor[string]()}, + &vectordata.DataProperty{Name: "When", StorageName: "when", Type: reflect.TypeFor[*time.Time]()}, + }} + dynamic, err := store.DynamicCollection("records", definition) + if err != nil { + t.Fatal(err) + } + instant := time.Date(2026, 1, 2, 3, 4, 5, 0, time.FixedZone("offset", 2*60*60)) + record := map[string]any{"Text": "dynamic", "When": &instant, "Other": "ignored"} + if err := dynamic.Upsert(t.Context(), record); err != nil { + t.Fatal(err) + } + key, ok := record["ID"].(uuid.UUID) + if !ok || key == uuid.Nil || uploaded["id"] != key.String() || uploaded["when"] != instant.Format(time.RFC3339) { + t.Fatalf("dynamic record = %v; uploaded = %v", record, uploaded) + } + gotMap, found, err := dynamic.Get(t.Context(), key, nil) + if err != nil || !found || gotMap["ID"] != key || gotMap["Text"] != "dynamic" || + !gotMap["When"].(*time.Time).Equal(instant) { + t.Fatalf("dynamic retrieval = (%v, %v, %v)", gotMap, found, err) + } + gotMap["Text"] = "changed" + if record["Text"] != "dynamic" || keyProperty.IsAutoGenerated != nil { + t.Fatal("retrieval or model resolution modified caller state") + } +} + +func TestCollection_NilCollectionFields(t *testing.T) { + type record struct { + ID string `vectordata:"key,storage=id"` + Tags []string `vectordata:"data,storage=tags"` + Numbers *[]int32 `vectordata:"data,storage=numbers"` + Text *string `vectordata:"data,storage=text"` + Vector []float32 `vectordata:"vector,dimensions=3,storage=vector"` + } + store := newStore(t, func(request *http.Request) (*http.Response, error) { + actions := jsonBody(t, request)["value"].([]any) + document := actions[0].(map[string]any) + for _, name := range []string{"tags", "numbers", "vector"} { + if value, ok := document[name].([]any); !ok || len(value) != 0 { + t.Errorf("field %s = %#v, want an empty array", name, document[name]) + } + } + if value, found := document["text"]; !found || value != nil { + t.Errorf("nullable scalar = %#v, present=%v", value, found) + } + return indexedResponse(t, actions, "id"), nil + }, azureprovider.AISearchConfig{}) + collection := openCollection[string, *record](t, store, nil) + input := &record{ID: "one", Numbers: new([]int32(nil))} + if err := collection.Upsert(t.Context(), input); err != nil { + t.Fatal(err) + } + if input.Tags != nil || *input.Numbers != nil || input.Text != nil || input.Vector != nil { + t.Fatalf("upsert changed the input: %+v", input) + } +} + +func TestCollection_CollectionFieldSerialization(t *testing.T) { + for _, test := range []struct { + name string + typ reflect.Type + value any + vector bool + want any + }{ + {"nil data", reflect.TypeFor[[]string](), nil, false, []any{}}, + {"nil slice", reflect.TypeFor[[]string](), []string(nil), false, []any{}}, + {"empty slice", reflect.TypeFor[[]string](), []string{}, false, []any{}}, + {"nonempty slice", reflect.TypeFor[[]string](), []string{"one"}, false, []any{"one"}}, + {"nil slice pointer", reflect.TypeFor[*[]string](), (*[]string)(nil), false, []any{}}, + {"pointer to nil slice", reflect.TypeFor[*[]string](), new([]string(nil)), false, []any{}}, + {"pointer to nonempty slice", reflect.TypeFor[*[]string](), new([]string{"one"}), false, []any{"one"}}, + {"nil array pointer", reflect.TypeFor[*[2]string](), (*[2]string)(nil), false, []any{}}, + {"nonempty array", reflect.TypeFor[[2]string](), [2]string{"one", "two"}, false, []any{"one", "two"}}, + {"untyped nil vector", reflect.TypeFor[[]float32](), nil, true, []any{}}, + {"nil vector", reflect.TypeFor[[]float32](), []float32(nil), true, []any{}}, + {"nil embedding vector", reflect.TypeFor[embedding.VectorOf[float32]](), embedding.VectorOf[float32](nil), true, []any{}}, + {"empty embedding", reflect.TypeFor[embedding.Of[float32]](), embedding.Of[float32]{}, true, []any{}}, + {"nil embedding pointer", reflect.TypeFor[*embedding.Of[float32]](), (*embedding.Of[float32])(nil), true, []any{}}, + {"pointer to empty embedding", reflect.TypeFor[*embedding.Of[float32]](), &embedding.Of[float32]{}, true, []any{}}, + {"nil scalar pointer", reflect.TypeFor[*string](), (*string)(nil), false, nil}, + } { + t.Run(test.name, func(t *testing.T) { + store := newStore(t, func(request *http.Request) (*http.Response, error) { + actions := jsonBody(t, request)["value"].([]any) + document := actions[0].(map[string]any) + if value, found := document["value"]; !found || !reflect.DeepEqual(value, test.want) { + t.Errorf("serialized value = %#v, present=%v; want %#v", value, found, test.want) + } + return indexedResponse(t, actions, "id"), nil + }, azureprovider.AISearchConfig{}) + var property vectordata.Property = &vectordata.DataProperty{Name: "Value", StorageName: "value", Type: test.typ} + if test.vector { + property = &vectordata.VectorProperty{Name: "Value", StorageName: "value", Type: test.typ, Dimensions: 3} + } + collection, err := store.DynamicCollection("records", &vectordata.Definition{Properties: []vectordata.Property{ + &vectordata.KeyProperty{Name: "ID", StorageName: "id", Type: reflect.TypeFor[string]()}, + property, + }}) + if err != nil { + t.Fatal(err) + } + input := map[string]any{"ID": "one", "Value": test.value} + before, err := json.Marshal(input) + if err != nil { + t.Fatal(err) + } + if err := collection.Upsert(t.Context(), input); err != nil { + t.Fatal(err) + } + after, err := json.Marshal(input) + if err != nil || string(after) != string(before) { + t.Fatalf("upsert changed the input: before=%s after=%s err=%v", before, after, err) + } + }) + } +} + +func TestCollection_NativeEmbeddingRepresentations(t *testing.T) { + for _, test := range []struct { + name string + typ reflect.Type + value any + }{ + {"slice", reflect.TypeFor[[]float32](), []float32{1, 2, 3}}, + {"vector", reflect.TypeFor[embedding.VectorOf[float32]](), embedding.VectorOf[float32]{1, 2, 3}}, + {"embedding", reflect.TypeFor[embedding.Of[float32]](), embedding.Of[float32]{Vector: embedding.VectorOf[float32]{1, 2, 3}}}, + {"pointer", reflect.TypeFor[*embedding.Of[float32]](), &embedding.Of[float32]{Vector: embedding.VectorOf[float32]{1, 2, 3}}}, + } { + t.Run(test.name, func(t *testing.T) { + store := newStore(t, func(request *http.Request) (*http.Response, error) { + if request.Method == http.MethodPost { + actions := jsonBody(t, request)["value"].([]any) + vector := actions[0].(map[string]any)["Vector"] + if !reflect.DeepEqual(vector, []any{float64(1), float64(2), float64(3)}) { + t.Fatalf("wire vector = %v", vector) + } + return indexedResponse(t, actions, "ID"), nil + } + return response(http.StatusOK, `{"ID":"one","Vector":[1,2,3]}`), nil + }, azureprovider.AISearchConfig{}) + collection, err := store.DynamicCollection("records", &vectordata.Definition{Properties: []vectordata.Property{ + &vectordata.KeyProperty{Name: "ID", Type: reflect.TypeFor[string]()}, + &vectordata.VectorProperty{Name: "Vector", Type: test.typ, Dimensions: 3}, + }}) + if err != nil { + t.Fatal(err) + } + input := map[string]any{"ID": "one", "Vector": test.value} + if err := collection.Upsert(t.Context(), input); err != nil { + t.Fatal(err) + } + output, found, err := collection.Get(t.Context(), "one", &vectordata.RecordRetrievalOptions{IncludeVectors: true}) + if err != nil || !found || !reflect.DeepEqual(output["Vector"], test.value) { + t.Fatalf("vector roundtrip = (%v, %v, %v)", output, found, err) + } + }) + } +} + +func TestCollection_PartialWriteFailures(t *testing.T) { + for _, operation := range []string{"upload", "delete"} { + for _, payload := range []string{ + `{"value":[{"key":"one","status":true,"statusCode":200},{"key":"two","status":false,"statusCode":400,"errorMessage":"bad document"}]}`, + `{"value":[{"key":"one","status":true},{"key":"two"}]}`, + `{"value":[]}`, + } { + t.Run(operation+"/"+payload, func(t *testing.T) { + calls := 0 + store := newStore(t, func(request *http.Request) (*http.Response, error) { + calls++ + return response(http.StatusMultiStatus, payload), nil + }, azureprovider.AISearchConfig{}) + collection := openCollection[string, document](t, store, nil) + var err error + if operation == "upload" { + err = collection.Upsert(t.Context(), document{ID: "one"}, document{ID: "two"}) + } else { + err = vectordata.DeleteBatch(t.Context(), collection, []string{"one", "two"}) + } + var vectorErr *vectordata.Error + if !errors.As(err, &vectorErr) || calls != 1 { + t.Fatalf("partial batch did not fail: %v; calls=%d", err, calls) + } + if strings.Contains(payload, "bad document") && !strings.Contains(err.Error(), "bad document") { + t.Fatalf("batch failure details were lost: %v", err) + } + }) + } + } +} + +func TestCollection_InvalidRecordsKeysAndResponses(t *testing.T) { + calls := 0 + var payload string + store := newStore(t, func(*http.Request) (*http.Response, error) { + calls++ + return response(http.StatusOK, payload), nil + }, azureprovider.AISearchConfig{}) + collection := openCollection[string, *document](t, store, nil) + if err := collection.Upsert(t.Context(), nil); err == nil { + t.Fatal("accepted nil record") + } + for _, key := range []string{"", " \t"} { + if _, _, err := collection.Get(t.Context(), key, nil); err == nil { + t.Fatal("accepted blank retrieval key") + } + if err := collection.Delete(t.Context(), key); err == nil { + t.Fatal("accepted blank deletion key") + } + } + if calls != 0 { + t.Fatal("invalid input performed I/O") + } + for _, invalid := range []string{ + `not json`, `null`, `{}`, `{"id":null}`, `{"id":7}`, `{"id":""}`, `{"id":"one","rank":"bad"}`, `{"id":"one","vector":"bad"}`, + } { + payload = invalid + if _, _, err := collection.Get(t.Context(), "one", &vectordata.RecordRetrievalOptions{IncludeVectors: true}); err == nil { + t.Fatalf("accepted malformed document %s", invalid) + } + } +} + +func TestCollection_CancellationAndProviderErrors(t *testing.T) { + wantErr := errors.New("transport failed") + calls := 0 + store := newStore(t, func(*http.Request) (*http.Response, error) { + calls++ + return nil, wantErr + }, azureprovider.AISearchConfig{}) + collection := openCollection[string, document](t, store, nil) + if _, _, err := collection.Get(t.Context(), "one", nil); !errors.Is(err, wantErr) { + t.Fatalf("transport error = %v", err) + } + ctx, cancel := context.WithCancel(t.Context()) + cancel() + before := calls + for _, err := range []error{ + collection.Upsert(ctx, document{ID: "one"}), + collection.Delete(ctx, "one"), + collection.EnsureExists(ctx), + collection.EnsureDeleted(ctx), + sequenceError(t, collection.Search(ctx, []float32{1, 2, 3}, 1, nil)), + } { + if !errors.Is(err, context.Canceled) { + t.Fatalf("cancellation = %v", err) + } + } + if calls != before { + t.Fatal("canceled operations reached the transport") + } + store = newStore(t, func(*http.Request) (*http.Response, error) { + return response(http.StatusForbidden, `{"error":{"code":"Forbidden","message":"denied"}}`), nil + }, azureprovider.AISearchConfig{}) + _, _, err := openCollection[string, document](t, store, nil).Get(t.Context(), "one", nil) + var serviceErr *azcore.ResponseError + if !errors.As(err, &serviceErr) || serviceErr.ErrorCode != "Forbidden" { + t.Fatalf("service failure was not preserved: %v", err) + } +} + +func TestCollection_BatchGetStopsWithConsumer(t *testing.T) { + calls := 0 + store := newStore(t, func(*http.Request) (*http.Response, error) { + calls++ + return response(http.StatusOK, `{"id":"one"}`), nil + }, azureprovider.AISearchConfig{}) + collection := openCollection[string, document](t, store, nil) + sequence := vectordata.GetBatch(t.Context(), collection, []string{"one", "two"}, nil) + if calls != 0 { + t.Fatal("batch retrieval started eagerly") + } + for _, err := range sequence { + if err != nil { + t.Fatal(err) + } + break + } + if calls != 1 { + t.Fatalf("consumer break retrieved %d documents", calls) + } +} diff --git a/provider/azureprovider/embedding_test.go b/provider/azureprovider/embedding_test.go new file mode 100644 index 000000000..1bff4e381 --- /dev/null +++ b/provider/azureprovider/embedding_test.go @@ -0,0 +1,303 @@ +// Copyright (c) Microsoft. All rights reserved. + +package azureprovider_test + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "log/slog" + "net/http" + "reflect" + "slices" + "strings" + "testing" + + "github.com/google/uuid" + "github.com/microsoft/agent-framework-go/agent" + "github.com/microsoft/agent-framework-go/agent/memory" + "github.com/microsoft/agent-framework-go/embedding" + "github.com/microsoft/agent-framework-go/internal/agenttest" + "github.com/microsoft/agent-framework-go/message" + "github.com/microsoft/agent-framework-go/provider/azureprovider" + "github.com/microsoft/agent-framework-go/vectordata" +) + +type generatorFunc func(context.Context, []string, *embedding.Options) (*embedding.Generated[*embedding.Of[float32]], error) + +func (f generatorFunc) Generate(ctx context.Context, values []string, options *embedding.Options) (*embedding.Generated[*embedding.Of[float32]], error) { + return f(ctx, values, options) +} + +func vectors(values []string) *embedding.Generated[*embedding.Of[float32]] { + result := &embedding.Generated[*embedding.Of[float32]]{} + for range values { + result.Embeddings = append(result.Embeddings, &embedding.Of[float32]{Vector: embedding.VectorOf[float32]{1, 2, 3}}) + } + return result +} + +type generatedDocument struct { + ID uuid.UUID `vectordata:"key,storage=id"` + Text string `vectordata:"data,fulltext,storage=text"` + Summary string `vectordata:"vector,dimensions=3,storage=summary"` + Details string `vectordata:"vector,dimensions=3,storage=details"` +} + +func TestCollection_BatchedEmbeddingGenerationAndQueries(t *testing.T) { + var batches [][]string + generator := embedding.Bind(generatorFunc(func(_ context.Context, values []string, options *embedding.Options) (*embedding.Generated[*embedding.Of[float32]], error) { + if options != nil { + t.Errorf("generation options = %+v; expected provider defaults", options) + } + batches = append(batches, slices.Clone(values)) + return vectors(values), nil + })) + var uploaded []any + var search map[string]any + calls := 0 + store := newStore(t, func(request *http.Request) (*http.Response, error) { + calls++ + if request.Method == http.MethodGet { + return response(http.StatusOK, `{"id":"af1b47a5-b046-49b5-aacd-53f219f9958d","text":"stored"}`), nil + } + body := jsonBody(t, request) + if strings.HasSuffix(request.URL.Path, "/docs/index") { + uploaded = body["value"].([]any) + return indexedResponse(t, uploaded, "id"), nil + } + search = body + return response(http.StatusOK, `{"value":[]}`), nil + }, azureprovider.AISearchConfig{EmbeddingGenerator: generator}) + collection := openCollection[uuid.UUID, *generatedDocument](t, store, nil) + inputs := []*generatedDocument{ + {Text: "one", Summary: "summary one", Details: "details one"}, + {Text: "two", Summary: "summary two", Details: "details two"}, + } + if err := collection.Upsert(t.Context(), inputs...); err != nil { + t.Fatal(err) + } + if !reflect.DeepEqual(batches, [][]string{{"summary one", "summary two"}, {"details one", "details two"}}) || + inputs[0].ID == uuid.Nil || inputs[0].Summary != "summary one" { + t.Fatalf("generation batches = %v; inputs = %v", batches, inputs) + } + first := uploaded[0].(map[string]any) + if !reflect.DeepEqual(first["summary"], []any{float64(1), float64(2), float64(3)}) || + !reflect.DeepEqual(first["details"], first["summary"]) || first["text"] != "one" { + t.Fatalf("generated wire record = %v", first) + } + sequence := collection.Search(t.Context(), "query", 3, &vectordata.SearchOptions{VectorProperty: "Summary"}) + if len(batches) != 2 { + t.Fatal("query generation started eagerly") + } + collect(t, sequence) + query := search["vectorQueries"].([]any)[0].(map[string]any) + if query["kind"] != "vector" || query["fields"] != "summary" || len(batches) != 3 || !slices.Equal(batches[2], []string{"query"}) { + t.Fatalf("query = %v; generation = %v", query, batches) + } + got, found, err := collection.Get(t.Context(), inputs[0].ID, nil) + if err != nil || !found || got.Summary != "" || got.Text != "stored" { + t.Fatalf("generated record retrieval = (%v, %v, %v)", got, found, err) + } + before := calls + if _, _, err := collection.Get(t.Context(), inputs[0].ID, &vectordata.RecordRetrievalOptions{IncludeVectors: true}); err == nil { + t.Fatal("generated-vector retrieval was accepted") + } + _ = sequenceError(t, collection.Search(t.Context(), "query", 1, &vectordata.SearchOptions{VectorProperty: "Summary", IncludeVectors: true})) + _ = sequenceError(t, collection.GetFiltered(t.Context(), vectordata.Filter{Operator: vectordata.Equal, Property: "ID", Value: inputs[0].ID}, 1, + &vectordata.FilteredRecordRetrievalOptions{IncludeVectors: true})) + hybrid := collection.(vectordata.KeywordHybridSearcher[*generatedDocument]) + _ = sequenceError(t, hybrid.HybridSearch(t.Context(), "query", []string{"query"}, 1, + &vectordata.HybridSearchOptions{VectorProperty: "Summary", IncludeVectors: true})) + if calls != before || len(batches) != 3 { + t.Fatal("rejected retrieval performed I/O or generated embeddings") + } +} + +func TestCollection_EmbeddingPrecedenceAndOwnership(t *testing.T) { + type entry struct { + ID string `vectordata:"key,storage=id"` + Text string `vectordata:"vector,dimensions=3"` + } + for level := range 3 { + t.Run(fmt.Sprint(level), func(t *testing.T) { + var calls []int + generators := make([]*embedding.BoundGenerator, 3) + for i := range generators { + generators[i] = embedding.Bind(generatorFunc(func(_ context.Context, values []string, _ *embedding.Options) (*embedding.Generated[*embedding.Of[float32]], error) { + calls = append(calls, i) + return vectors(values), nil + })) + } + property := &vectordata.VectorProperty{Name: "Text", Dimensions: 3} + definition := &vectordata.Definition{Properties: []vectordata.Property{property}} + if level >= 1 { + definition.EmbeddingGenerator = generators[1] + } + if level >= 2 { + property.EmbeddingGenerator = generators[2] + } + store := newStore(t, func(request *http.Request) (*http.Response, error) { + body := jsonBody(t, request) + if actions, found := body["value"]; found { + return indexedResponse(t, actions.([]any), "id"), nil + } + return response(http.StatusOK, `{"value":[]}`), nil + }, azureprovider.AISearchConfig{EmbeddingGenerator: generators[0]}) + collection := openCollection[string, entry](t, store, definition) + if err := collection.Upsert(t.Context(), entry{ID: "one", Text: "text"}); err != nil { + t.Fatal(err) + } + collect(t, collection.Search(t.Context(), "query", 1, nil)) + if !slices.Equal(calls, []int{level, level}) || property.Type != nil || property.EmbeddingType != nil || property.StorageName != "" { + t.Fatalf("generation=%v property=%+v", calls, property) + } + if level == 0 && definition.EmbeddingGenerator != nil || level >= 1 && definition.EmbeddingGenerator != generators[1] { + t.Fatal("store defaults changed the input definition") + } + }) + } +} + +func TestCollection_EmbeddingErrorsDoNotUploadOrGenerateKeys(t *testing.T) { + wantErr := errors.New("embedding failure") + for _, test := range []struct { + name string + result *embedding.Generated[*embedding.Of[float32]] + err error + }{ + {"provider error", nil, wantErr}, + {"nil batch", nil, nil}, + {"wrong count", vectors([]string{"one"}), nil}, + {"nil embedding", &embedding.Generated[*embedding.Of[float32]]{Embeddings: []*embedding.Of[float32]{nil, nil}}, nil}, + } { + t.Run(test.name, func(t *testing.T) { + calls := 0 + generator := embedding.Bind(generatorFunc(func(context.Context, []string, *embedding.Options) (*embedding.Generated[*embedding.Of[float32]], error) { + return test.result, test.err + })) + store := newStore(t, func(*http.Request) (*http.Response, error) { + calls++ + return nil, errors.New("unexpected request") + }, azureprovider.AISearchConfig{EmbeddingGenerator: generator}) + collection := openCollection[uuid.UUID, *generatedDocument](t, store, nil) + input := &generatedDocument{} + err := collection.Upsert(t.Context(), input, &generatedDocument{}) + if err == nil || test.err != nil && !errors.Is(err, wantErr) || calls != 0 || input.ID != uuid.Nil { + t.Fatalf("failure=%v calls=%d key=%v", err, calls, input.ID) + } + }) + } +} + +func TestAISearch_ChatHistoryMemory(t *testing.T) { + var stored []map[string]any + exists := false + var logs bytes.Buffer + store := newStore(t, func(request *http.Request) (*http.Response, error) { + switch { + case request.Method == http.MethodGet: + if !exists { + return response(http.StatusNotFound, `{"error":{"code":"NotFound"}}`), nil + } + return response(http.StatusOK, `{"name":"records"}`), nil + case request.URL.Path == "/indexes": + definition := jsonBody(t, request) + fields := definition["fields"].([]any) + if key := fields[0].(map[string]any); key["name"] != "Key" || key["type"] != "Edm.String" { + t.Errorf("memory key schema = %v", key) + } + exists = true + return response(http.StatusCreated, `{"name":"records"}`), nil + case strings.HasSuffix(request.URL.Path, "/docs/index"): + actions := jsonBody(t, request)["value"].([]any) + for _, action := range actions { + value := action.(map[string]any) + if _, err := uuid.Parse(value["Key"].(string)); err != nil { + t.Errorf("memory key was not encoded as a UUID string: %v", value["Key"]) + } + stored = append(stored, value) + } + return indexedResponse(t, actions, "Key"), nil + default: + body := jsonBody(t, request) + if body["filter"] != "(UserId eq 'u''1')" { + t.Errorf("memory search scope = %v", body["filter"]) + } + results := make([]map[string]any, len(stored)) + copy(results, stored) + payload, err := json.Marshal(map[string]any{"value": results}) + if err != nil { + t.Fatal(err) + } + return response(http.StatusOK, string(payload)), nil + } + }, azureprovider.AISearchConfig{EmbeddingGenerator: embedding.Bind(generatorFunc(func(_ context.Context, values []string, _ *embedding.Options) (*embedding.Generated[*embedding.Of[float32]], error) { + return vectors(values), nil + }))}) + memories, err := memory.NewHistoryProvider(store.Store, "records", 3, func(*agent.Session) *memory.State { + return &memory.State{ + StorageScope: memory.Scope{UserID: new("u'1"), SessionID: new(uuid.NewString())}, + SearchScope: &memory.Scope{UserID: new("u'1")}, + } + }, memory.Config{Logger: slog.New(slog.NewJSONHandler(&logs, nil))}) + if err != nil { + t.Fatal(err) + } + runner := &agenttest.Runner{Responses: agenttest.NewResponseBuilder(nil).AddText("A pirate joke."). + NewTurn(func(_ context.Context, messages []*message.Message, _ ...agent.Option) { + if len(messages) != 2 || !strings.Contains(messages[1].String(), "I like pirate jokes.") { + t.Fatalf("new session did not receive prior memories: %v", messages) + } + }).AddText("Another pirate joke.").Build()} + a := agent.New(agent.ProviderConfig{Run: runner.Run}, agent.Config{ContextProviders: []agent.ContextProvider{memories}}) + for _, query := range []string{"I like pirate jokes.", "Tell me a joke I might like."} { + session, err := a.CreateSession(t.Context()) + if err != nil { + t.Fatal(err) + } + if _, err := a.RunText(t.Context(), query, agent.WithSession(session)).Collect(); err != nil { + t.Fatal(err) + } + } + if len(stored) != 4 || strings.Contains(logs.String(), `"level":"ERROR"`) { + t.Fatalf("stored=%d diagnostics=%s", len(stored), logs.String()) + } +} + +func TestCollection_SearchEmbeddingFailures(t *testing.T) { + wantErr := errors.New("query embedding failed") + for _, test := range []struct { + name string + result *embedding.Generated[*embedding.Of[float32]] + err error + }{ + {"provider error", nil, wantErr}, + {"nil batch", nil, nil}, + {"nil embedding", &embedding.Generated[*embedding.Of[float32]]{Embeddings: []*embedding.Of[float32]{nil}}, nil}, + {"nil vector", &embedding.Generated[*embedding.Of[float32]]{Embeddings: []*embedding.Of[float32]{{}}}, nil}, + } { + t.Run(test.name, func(t *testing.T) { + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + generator := embedding.Bind(generatorFunc(func(context.Context, []string, *embedding.Options) (*embedding.Generated[*embedding.Of[float32]], error) { + if test.err != nil { + cancel() + } + return test.result, test.err + })) + calls := 0 + store := newStore(t, func(*http.Request) (*http.Response, error) { + calls++ + return response(http.StatusOK, `{"value":[]}`), nil + }, azureprovider.AISearchConfig{EmbeddingGenerator: generator}) + collection := openCollection[string, document](t, store, nil) + err := sequenceError(t, collection.Search(ctx, "query", 1, nil)) + if calls != 0 || test.err != nil && (!errors.Is(err, wantErr) || errors.Is(err, context.Canceled)) { + t.Fatalf("query embedding error = %v, HTTP calls = %d", err, calls) + } + }) + } +} diff --git a/provider/azureprovider/filter.go b/provider/azureprovider/filter.go new file mode 100644 index 000000000..ee1aa7f97 --- /dev/null +++ b/provider/azureprovider/filter.go @@ -0,0 +1,184 @@ +// Copyright (c) Microsoft. All rights reserved. + +package azureprovider + +import ( + "fmt" + "math" + "reflect" + "strconv" + "strings" + "time" + + "github.com/google/uuid" + "github.com/microsoft/agent-framework-go/vectordata" + "github.com/microsoft/agent-framework-go/vectordata/providerservices" +) + +func (c *collection) translateFilter(filter *vectordata.Filter) (string, error) { + if err := filter.Validate(); err != nil { + return "", err + } + if filter == nil { + return "", nil + } + return c.filterExpression(filter) +} + +func (c *collection) filterExpression(filter *vectordata.Filter) (string, error) { + switch filter.Operator { + case vectordata.And, vectordata.Or, vectordata.Not: + parts := make([]string, len(filter.Operands)) + for i := range filter.Operands { + part, err := c.filterExpression(&filter.Operands[i]) + if err != nil { + return "", err + } + parts[i] = part + } + if filter.Operator == vectordata.Not { + return "(not " + parts[0] + ")", nil + } + return "(" + strings.Join(parts, " "+string(filter.Operator)+" ") + ")", nil + } + name, typ, err := c.filterProperty(filter.Property) + if err != nil { + return "", err + } + if typ.Kind() == reflect.Pointer { + typ = typ.Elem() + } + switch filter.Operator { + case vectordata.In: + if typ.Kind() != reflect.String { + return "", fmt.Errorf("property %q: in filters require a string property", filter.Property) + } + value := reflect.ValueOf(filter.Value) + values := make([]string, 0, value.Len()) + hasEmpty := false + delimiter := ", " + for i := range value.Len() { + element := value.Index(i) + for element.Kind() == reflect.Interface && !element.IsNil() { + element = element.Elem() + } + if element.Kind() != reflect.String { + return "", fmt.Errorf("property %q: in filters require string values", filter.Property) + } + text := element.String() + // search.in ignores empty tokens, so match them separately. + if text == "" { + hasEmpty = true + continue + } + values = append(values, text) + if strings.ContainsAny(text, ", ") { + delimiter = "|" + } + } + if len(values) == 0 { + if hasEmpty { + return "(" + name + " eq '')", nil + } + return "false", nil + } + if delimiter == "|" { + for _, value := range values { + if strings.Contains(value, "|") { + return "", fmt.Errorf("property %q: in filter values cannot combine commas or spaces with pipes", filter.Property) + } + } + } + expression := "search.in(" + name + ", " + quote(strings.Join(values, delimiter)) + if delimiter == "|" { + expression += ", '|'" + } + expression += ")" + if hasEmpty { + expression = "(" + name + " eq '' or " + expression + ")" + } + return expression, nil + case vectordata.Contains: + if typ.Kind() != reflect.Slice && typ.Kind() != reflect.Array { + return "", fmt.Errorf("property %q: contains filters require a collection property", filter.Property) + } + value, err := odataLiteral(filter.Value) + if err != nil { + return "", err + } + return name + "/any(t: t eq " + value + ")", nil + default: + if typ.Kind() == reflect.Slice || typ.Kind() == reflect.Array && typ != reflect.TypeFor[uuid.UUID]() { + return "", fmt.Errorf("property %q: scalar comparison requires a scalar property", filter.Property) + } + value, err := odataLiteral(filter.Value) + if err != nil { + return "", err + } + return "(" + name + " " + string(filter.Operator) + " " + value + ")", nil + } +} + +func (c *collection) filterProperty(name string) (string, reflect.Type, error) { + switch property := c.properties[name].(type) { + case *providerservices.KeyPropertyModel: + return property.StorageName, property.Type, nil + case *providerservices.DataPropertyModel: + return property.StorageName, property.Type, nil + default: + return "", nil, fmt.Errorf("filter or ordering property %q must be a key or data property", name) + } +} + +func quote(value string) string { + return "'" + strings.ReplaceAll(value, "'", "''") + "'" +} + +func odataLiteral(value any) (string, error) { + reflected := reflect.ValueOf(value) + visited := make(map[uintptr]bool) + for reflected.IsValid() && (reflected.Kind() == reflect.Pointer || reflected.Kind() == reflect.Interface) { + if reflected.IsNil() { + return "null", nil + } + if reflected.Kind() == reflect.Pointer { + address := reflected.Pointer() + if visited[address] { + return "", fmt.Errorf("filter value contains a pointer cycle") + } + visited[address] = true + } + reflected = reflected.Elem() + } + if !reflected.IsValid() { + return "null", nil + } + if reflected.Type() == reflect.TypeFor[uuid.UUID]() { + return quote(reflected.Interface().(uuid.UUID).String()), nil + } + if reflected.Type() == reflect.TypeFor[time.Time]() { + return reflected.Interface().(time.Time).Format(time.RFC3339Nano), nil + } + switch reflected.Kind() { + case reflect.String: + return quote(reflected.String()), nil + case reflect.Bool: + return strconv.FormatBool(reflected.Bool()), nil + case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64: + return strconv.FormatInt(reflected.Int(), 10), nil + case reflect.Uint8: + return strconv.FormatUint(reflected.Uint(), 10), nil + case reflect.Float32, reflect.Float64: + number := reflected.Float() + switch { + case math.IsInf(number, 1): + return "INF", nil + case math.IsInf(number, -1): + return "-INF", nil + default: + return strconv.FormatFloat(number, 'g', -1, reflected.Type().Bits()), nil + } + default: + return "", fmt.Errorf("unsupported filter value type %v", reflected.Type()) + } +} diff --git a/provider/azureprovider/filter_test.go b/provider/azureprovider/filter_test.go new file mode 100644 index 000000000..519c67637 --- /dev/null +++ b/provider/azureprovider/filter_test.go @@ -0,0 +1,125 @@ +// Copyright (c) Microsoft. All rights reserved. + +package azureprovider_test + +import ( + "fmt" + "math" + "net/http" + "reflect" + "testing" + "time" + + "github.com/google/uuid" + "github.com/microsoft/agent-framework-go/provider/azureprovider" + "github.com/microsoft/agent-framework-go/vectordata" +) + +func TestCollection_FilterTranslation(t *testing.T) { + instant := time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC) + key := uuid.MustParse("d282e2e0-73ce-4b91-8973-b6bc418e31fa") + for _, test := range []struct { + name string + filter vectordata.Filter + want string + }{ + {"quoted string", vectordata.Filter{Operator: vectordata.Equal, Property: "Text", Value: "O'Reilly"}, "(text eq 'O''Reilly')"}, + {"not equal", vectordata.Filter{Operator: vectordata.NotEqual, Property: "Rank", Value: 3}, "(rank ne 3)"}, + {"greater", vectordata.Filter{Operator: vectordata.GreaterThan, Property: "Rank", Value: 3}, "(rank gt 3)"}, + {"greater equal", vectordata.Filter{Operator: vectordata.GreaterThanOrEqual, Property: "Rank", Value: 3}, "(rank ge 3)"}, + {"less", vectordata.Filter{Operator: vectordata.LessThan, Property: "Rank", Value: 3}, "(rank lt 3)"}, + {"less equal", vectordata.Filter{Operator: vectordata.LessThanOrEqual, Property: "Rank", Value: 3}, "(rank le 3)"}, + {"null", vectordata.Filter{Operator: vectordata.Equal, Property: "Text", Value: nil}, "(text eq null)"}, + {"typed nil", vectordata.Filter{Operator: vectordata.Equal, Property: "Text", Value: (*string)(nil)}, "(text eq null)"}, + {"pointer", vectordata.Filter{Operator: vectordata.Equal, Property: "Rank", Value: new(3)}, "(rank eq 3)"}, + {"contains", vectordata.Filter{Operator: vectordata.Contains, Property: "Tags", Value: "a'b"}, "tags/any(t: t eq 'a''b')"}, + {"in", vectordata.Filter{Operator: vectordata.In, Property: "Text", Value: []string{"one", "two"}}, "search.in(text, 'one, two')"}, + {"in pipe delimiter", vectordata.Filter{Operator: vectordata.In, Property: "Text", Value: []string{"one two", "O'Reilly"}}, "search.in(text, 'one two|O''Reilly', '|')"}, + {"in default pipe value", vectordata.Filter{Operator: vectordata.In, Property: "Text", Value: []string{"a|b", "c"}}, "search.in(text, 'a|b, c')"}, + {"empty in", vectordata.Filter{Operator: vectordata.In, Property: "Text", Value: []string{}}, "false"}, + {"nil in", vectordata.Filter{Operator: vectordata.In, Property: "Text", Value: []string(nil)}, "false"}, + {"in empty string", vectordata.Filter{Operator: vectordata.In, Property: "Text", Value: []string{""}}, "(text eq '')"}, + {"in repeated empty strings", vectordata.Filter{Operator: vectordata.In, Property: "Text", Value: []string{"", ""}}, "(text eq '')"}, + {"in empty and nonempty strings", vectordata.Filter{Operator: vectordata.In, Property: "Text", Value: []string{"", "one", ""}}, "(text eq '' or search.in(text, 'one'))"}, + {"in empty and delimited strings", vectordata.Filter{Operator: vectordata.In, Property: "Text", Value: []any{"one two", "", "O'Reilly"}}, "(text eq '' or search.in(text, 'one two|O''Reilly', '|'))"}, + {"not in empty and nonempty strings", vectordata.Filter{Operator: vectordata.Not, Operands: []vectordata.Filter{ + {Operator: vectordata.In, Property: "Text", Value: []string{"", "one"}}, + }}, "(not (text eq '' or search.in(text, 'one')))"}, + {"date", vectordata.Filter{Operator: vectordata.Equal, Property: "When", Value: instant}, "(when eq 2026-01-02T03:04:05Z)"}, + {"uuid", vectordata.Filter{Operator: vectordata.Equal, Property: "ID", Value: key}, "(id eq 'd282e2e0-73ce-4b91-8973-b6bc418e31fa')"}, + {"bool", vectordata.Filter{Operator: vectordata.Equal, Property: "Active", Value: true}, "(active eq true)"}, + {"float", vectordata.Filter{Operator: vectordata.Equal, Property: "Number", Value: float32(1.5)}, "(number eq 1.5)"}, + {"byte", vectordata.Filter{Operator: vectordata.Equal, Property: "Number", Value: byte(3)}, "(number eq 3)"}, + {"infinity", vectordata.Filter{Operator: vectordata.Equal, Property: "Number", Value: math.Inf(1)}, "(number eq INF)"}, + {"negative infinity", vectordata.Filter{Operator: vectordata.Equal, Property: "Number", Value: math.Inf(-1)}, "(number eq -INF)"}, + {"and", vectordata.Filter{Operator: vectordata.And, Operands: []vectordata.Filter{ + {Operator: vectordata.Equal, Property: "Text", Value: "a"}, + {Operator: vectordata.GreaterThan, Property: "Rank", Value: 2}, + }}, "((text eq 'a') and (rank gt 2))"}, + {"or", vectordata.Filter{Operator: vectordata.Or, Operands: []vectordata.Filter{ + {Operator: vectordata.Equal, Property: "Text", Value: "a"}, + {Operator: vectordata.Equal, Property: "Text", Value: "b"}, + }}, "((text eq 'a') or (text eq 'b'))"}, + {"not", vectordata.Filter{Operator: vectordata.Not, Operands: []vectordata.Filter{ + {Operator: vectordata.Equal, Property: "Text", Value: "a"}, + }}, "(not (text eq 'a'))"}, + } { + t.Run(test.name, func(t *testing.T) { + var body map[string]any + store := newStore(t, func(request *http.Request) (*http.Response, error) { + body = jsonBody(t, request) + return response(http.StatusOK, `{"value":[]}`), nil + }, azureprovider.AISearchConfig{}) + definition := &vectordata.Definition{Properties: []vectordata.Property{ + &vectordata.KeyProperty{Name: "ID", StorageName: "id", Type: reflect.TypeFor[string]()}, + &vectordata.DataProperty{Name: "Text", StorageName: "text", Type: reflect.TypeFor[string](), IsIndexed: true}, + &vectordata.DataProperty{Name: "Rank", StorageName: "rank", Type: reflect.TypeFor[int32](), IsIndexed: true}, + &vectordata.DataProperty{Name: "Tags", StorageName: "tags", Type: reflect.TypeFor[[]string](), IsIndexed: true}, + &vectordata.DataProperty{Name: "When", StorageName: "when", Type: reflect.TypeFor[time.Time](), IsIndexed: true}, + &vectordata.DataProperty{Name: "Active", StorageName: "active", Type: reflect.TypeFor[bool](), IsIndexed: true}, + &vectordata.DataProperty{Name: "Number", StorageName: "number", Type: reflect.TypeFor[float64](), IsIndexed: true}, + }} + collection := openCollection[string, map[string]any](t, store, definition) + collect(t, collection.GetFiltered(t.Context(), test.filter, 5, &vectordata.FilteredRecordRetrievalOptions{ + Skip: 2, OrderBy: []vectordata.Sort{{Property: "Rank", Descending: true}, {Property: "Text"}, {Property: "ID"}}, + })) + if body["filter"] != test.want || body["top"] != float64(5) || body["skip"] != float64(2) || + body["orderby"] != "rank desc,text,id" { + t.Fatalf("filtered search = %v; want filter %s", body, test.want) + } + }) + } +} + +func TestCollection_RejectsUnsupportedFilters(t *testing.T) { + var cycle any + cycle = &cycle + for _, filter := range []vectordata.Filter{ + {Operator: vectordata.Equal, Property: "Unknown", Value: "a"}, + {Operator: vectordata.Equal, Property: "Vector", Value: "a"}, + {Operator: vectordata.In, Property: "Rank", Value: []int{1, 2}}, + {Operator: vectordata.In, Property: "Text", Value: []any{"a", 2}}, + {Operator: vectordata.In, Property: "Text", Value: []string{"a b", "c|d"}}, + {Operator: vectordata.In, Property: "Text", Value: []any{"", 2}}, + {Operator: vectordata.In, Property: "Text", Value: []string{"", "a b", "c|d"}}, + {Operator: vectordata.Contains, Property: "Text", Value: "a"}, + {Operator: vectordata.Equal, Property: "Tags", Value: "a"}, + {Operator: vectordata.Equal, Property: "Text", Value: map[string]any{"bad": "value"}}, + {Operator: vectordata.Equal, Property: "Text", Value: cycle}, + {Operator: vectordata.And}, + {Operator: "unknown", Property: "Text", Value: "a"}, + } { + t.Run(fmt.Sprint(filter.Operator)+"/"+filter.Property, func(t *testing.T) { + calls := 0 + store := newStore(t, func(*http.Request) (*http.Response, error) { + calls++ + return response(http.StatusOK, `{"value":[]}`), nil + }, azureprovider.AISearchConfig{}) + collection := openCollection[string, document](t, store, nil) + _ = sequenceError(t, collection.GetFiltered(t.Context(), filter, 1, nil)) + if calls != 0 { + t.Fatal("unsupported filter reached the service") + } + }) + } +} diff --git a/provider/azureprovider/record.go b/provider/azureprovider/record.go new file mode 100644 index 000000000..e923fb443 --- /dev/null +++ b/provider/azureprovider/record.go @@ -0,0 +1,212 @@ +// Copyright (c) Microsoft. All rights reserved. + +package azureprovider + +import ( + "bytes" + "encoding/json" + "fmt" + "reflect" + "strings" + + "github.com/google/uuid" + "github.com/microsoft/agent-framework-go/embedding" + "github.com/microsoft/agent-framework-go/vectordata/providerservices" +) + +func keyString(key any) (string, error) { + if key == nil { + return "", fmt.Errorf("key is required") + } + if value, ok := key.(uuid.UUID); ok { + return value.String(), nil + } + value := reflect.ValueOf(key) + if value.Kind() != reflect.String { + return "", fmt.Errorf("key must be a string or uuid.UUID, got %T", key) + } + if strings.TrimSpace(value.String()) == "" { + return "", fmt.Errorf("key must not be blank") + } + return value.String(), nil +} + +func (c *collection) prepareRecord(record any) (any, error) { + value := reflect.ValueOf(record) + if !value.IsValid() || !value.Type().AssignableTo(c.recordType) { + return nil, fmt.Errorf("record must have type %v, got %T", c.recordType, record) + } + if value.Type() != c.recordType { + assigned := reflect.New(c.recordType).Elem() + assigned.Set(value) + value = assigned + } + if value.Kind() == reflect.Struct { + owned := reflect.New(c.recordType) + owned.Elem().Set(value) + return owned.Interface(), nil + } + if value.IsNil() || value.Kind() == reflect.Pointer && value.Elem().Kind() == reflect.Map && value.Elem().IsNil() { + return nil, fmt.Errorf("record must not be nil") + } + return value.Interface(), nil +} + +func (c *collection) encodeRecord(record any, index int, generated map[string][][]float32) (map[string]any, error) { + key, found, err := c.key.Value(record) + if err != nil { + return nil, err + } + if c.key.IsAutoGenerated && (!found || key == nil || key == uuid.Nil) { + key, err = uuid.NewRandom() + if err != nil { + return nil, err + } + if err := c.key.SetValue(record, key); err != nil { + return nil, err + } + } + stringKey, err := keyString(key) + if err != nil { + return nil, fmt.Errorf("property %q: %w", c.key.Name, err) + } + document := make(map[string]any, len(c.model.Properties)+1) + document[c.key.StorageName] = stringKey + for _, property := range c.data { + value, found, err := property.Value(record) + if err != nil { + return nil, err + } + if !found { + continue + } + typ := property.Type + if typ.Kind() == reflect.Pointer { + typ = typ.Elem() + } + if typ.Kind() == reflect.Slice || typ.Kind() == reflect.Array { + reflected := reflect.ValueOf(value) + if reflected.IsValid() && reflected.Kind() == reflect.Pointer { + reflected = reflected.Elem() + } + // Azure collection fields cannot be null. + if !reflected.IsValid() || reflected.Kind() == reflect.Slice && reflected.IsNil() { + value = []any{} + } + } + document[property.StorageName] = value + } + for _, property := range c.vectors { + value, found, err := property.Value(record) + if err != nil { + return nil, err + } + if !found { + continue + } + if outputs, wasGenerated := generated[property.Name]; wasGenerated { + document[property.StorageName] = outputs[index] + continue + } + vector, ok := providerservices.AsVector[float32](value) + if !ok && value != nil { + return nil, fmt.Errorf("property %q requires a float32 vector", property.Name) + } + if vector == nil { + vector = []float32{} + } + document[property.StorageName] = vector + } + return document, nil +} + +func (c *collection) decodeRecord(document map[string]json.RawMessage, includeVectors bool) (any, error) { + root := reflect.New(c.recordType) + target := root.Interface() + recordType := c.recordType + if recordType.Kind() == reflect.Pointer { + recordType = recordType.Elem() + root.Elem().Set(reflect.New(recordType)) + target = root.Elem().Interface() + } + if recordType.Kind() == reflect.Map { + value := reflect.ValueOf(target).Elem() + value.Set(reflect.MakeMap(recordType)) + } + key, ok := document[c.key.StorageName] + if !ok || bytes.Equal(bytes.TrimSpace(key), []byte("null")) { + return nil, fmt.Errorf("response is missing key property %q", c.key.StorageName) + } + var stringKey string + if err := json.Unmarshal(key, &stringKey); err != nil { + return nil, fmt.Errorf("property %q: %w", c.key.Name, err) + } + if _, err := keyString(stringKey); err != nil { + return nil, fmt.Errorf("property %q: %w", c.key.Name, err) + } + keyValue := reflect.New(c.key.Type) + if err := json.Unmarshal(key, keyValue.Interface()); err != nil { + return nil, fmt.Errorf("property %q: %w", c.key.Name, err) + } + if err := c.key.SetValue(target, keyValue.Elem().Interface()); err != nil { + return nil, err + } + for _, property := range c.data { + raw, found := document[property.StorageName] + if !found { + continue + } + value, err := decodeValue(raw, property.Type) + if err != nil { + return nil, fmt.Errorf("property %q: %w", property.Name, err) + } + if err := property.SetValue(target, value); err != nil { + return nil, err + } + } + if includeVectors { + for _, property := range c.vectors { + raw, found := document[property.StorageName] + if !found { + continue + } + value, err := decodeVector(raw, property) + if err != nil { + return nil, fmt.Errorf("property %q: %w", property.Name, err) + } + if err := property.SetValue(target, value); err != nil { + return nil, err + } + } + } + return root.Elem().Interface(), nil +} + +func decodeValue(raw json.RawMessage, typ reflect.Type) (any, error) { + if bytes.Equal(bytes.TrimSpace(raw), []byte("null")) { + return nil, nil + } + value := reflect.New(typ) + if err := json.Unmarshal(raw, value.Interface()); err != nil { + return nil, err + } + return value.Elem().Interface(), nil +} + +func decodeVector(raw json.RawMessage, property *providerservices.VectorPropertyModel) (any, error) { + if property.Type == reflect.TypeFor[embedding.Of[float32]]() || property.Type == reflect.TypeFor[*embedding.Of[float32]]() { + if bytes.Equal(bytes.TrimSpace(raw), []byte("null")) { + return nil, nil + } + var vector embedding.VectorOf[float32] + if err := json.Unmarshal(raw, &vector); err != nil { + return nil, err + } + result := embedding.Of[float32]{Vector: vector} + if property.Type.Kind() == reflect.Pointer { + return &result, nil + } + return result, nil + } + return decodeValue(raw, property.Type) +} diff --git a/provider/azureprovider/schema.go b/provider/azureprovider/schema.go new file mode 100644 index 000000000..7e2b2df17 --- /dev/null +++ b/provider/azureprovider/schema.go @@ -0,0 +1,234 @@ +// Copyright (c) Microsoft. All rights reserved. + +package azureprovider + +import ( + "fmt" + "reflect" + "strings" + "time" + "unicode" + + "github.com/google/uuid" + "github.com/microsoft/agent-framework-go/vectordata" + "github.com/microsoft/agent-framework-go/vectordata/providerservices" +) + +type collection struct { + store *store + name string + recordType reflect.Type + model *providerservices.CollectionModel + key *providerservices.KeyPropertyModel + data []*providerservices.DataPropertyModel + vectors []*providerservices.VectorPropertyModel + properties map[string]providerservices.PropertyModel +} + +func newCollection(s *store, name string, keyType, recordType reflect.Type, definition *vectordata.Definition) (*collection, error) { + if strings.TrimSpace(name) == "" { + return nil, failure(name, "Collection", fmt.Errorf("collection name is required")) + } + model, err := providerservices.Build(keyType, recordType, definition, providerservices.Config{ + SupportsMultipleVectors: true, + SupportsKeyType: keyTypeSupported, + SupportsDataType: dataTypeSupported, + SupportsVectorType: providerservices.SupportsVector[float32], + SupportsKeyAutoGeneration: func(typ reflect.Type) bool { return typ == reflect.TypeFor[uuid.UUID]() }, + DefaultEmbeddingGenerator: s.generator, + }) + if err != nil { + return nil, failure(name, "Collection", err) + } + c := &collection{ + store: s, name: name, recordType: recordType, model: model, + properties: make(map[string]providerservices.PropertyModel, len(model.Properties)), + } + for _, property := range model.Properties { + var propertyName, storageName string + switch p := property.(type) { + case *providerservices.KeyPropertyModel: + c.key = p + propertyName, storageName = p.Name, p.StorageName + case *providerservices.DataPropertyModel: + c.data = append(c.data, p) + propertyName, storageName = p.Name, p.StorageName + case *providerservices.VectorPropertyModel: + c.vectors = append(c.vectors, p) + propertyName, storageName = p.Name, p.StorageName + } + if !validIdentifier(storageName) { + return nil, failure(name, "Collection", fmt.Errorf("property %q has invalid OData storage name %q", propertyName, storageName)) + } + c.properties[propertyName] = property + } + return c, nil +} + +func keyTypeSupported(typ reflect.Type) bool { + return typ.Kind() == reflect.String || typ == reflect.TypeFor[uuid.UUID]() +} + +func dataTypeSupported(typ reflect.Type) bool { + if typ.Kind() == reflect.Pointer { + typ = typ.Elem() + } + if typ.Kind() == reflect.Slice || typ.Kind() == reflect.Array { + typ = typ.Elem() + } + if typ == reflect.TypeFor[time.Time]() { + return true + } + switch typ.Kind() { + case reflect.String, reflect.Bool, reflect.Int, reflect.Int32, reflect.Int64, reflect.Float32, reflect.Float64: + return true + default: + return false + } +} + +func validIdentifier(name string) bool { + if name == "" { + return false + } + for i, r := range name { + if r != '_' && !unicode.IsLetter(r) && (i == 0 || !unicode.IsDigit(r)) { + return false + } + } + return true +} + +type indexDefinition struct { + Name string `json:"name"` + Fields []indexField `json:"fields"` + VectorSearch vectorIndexDefinition `json:"vectorSearch"` +} + +type indexField struct { + Name string `json:"name"` + Type string `json:"type"` + Key bool `json:"key,omitempty"` + Searchable bool `json:"searchable"` + Filterable bool `json:"filterable"` + Sortable bool `json:"sortable"` + Facetable bool `json:"facetable"` + Retrievable bool `json:"retrievable"` + Dimensions int `json:"dimensions,omitempty"` + VectorSearchProfile string `json:"vectorSearchProfile,omitempty"` +} + +type vectorIndexDefinition struct { + Algorithms []vectorAlgorithm `json:"algorithms"` + Profiles []vectorProfile `json:"profiles"` +} + +type vectorAlgorithm struct { + Name string `json:"name"` + Kind string `json:"kind"` + HNSWParameters *vectorParameters `json:"hnswParameters,omitempty"` + ExhaustiveKNNParameters *vectorParameters `json:"exhaustiveKnnParameters,omitempty"` +} + +type vectorParameters struct { + Metric string `json:"metric"` +} + +type vectorProfile struct { + Name string `json:"name"` + Algorithm string `json:"algorithm"` +} + +func (c *collection) indexDefinition() (*indexDefinition, error) { + definition := &indexDefinition{ + Name: c.name, Fields: make([]indexField, 0, len(c.model.Properties)), + VectorSearch: vectorIndexDefinition{Algorithms: []vectorAlgorithm{}, Profiles: []vectorProfile{}}, + } + for _, property := range c.model.Properties { + field := indexField{Retrievable: true} + switch p := property.(type) { + case *providerservices.KeyPropertyModel: + field.Name, field.Type = p.StorageName, "Edm.String" + field.Key, field.Searchable, field.Filterable, field.Sortable = true, true, true, true + case *providerservices.DataPropertyModel: + field.Name = p.StorageName + var err error + field.Type, err = dataFieldType(p.Type) + if err != nil { + return nil, err + } + field.Filterable = p.IsIndexed + field.Sortable = p.IsIndexed && !strings.HasPrefix(field.Type, "Collection(") + if p.IsFullTextIndexed { + typ := p.Type + if typ.Kind() == reflect.Pointer { + typ = typ.Elem() + } + if typ.Kind() != reflect.String { + return nil, fmt.Errorf("property %q: full-text indexing requires a string", p.Name) + } + field.Searchable = true + } + case *providerservices.VectorPropertyModel: + field.Name, field.Type = p.StorageName, "Collection(Edm.Single)" + field.Searchable, field.Dimensions = true, p.Dimensions + field.VectorSearchProfile = p.StorageName + "Profile" + metric := "" + switch p.DistanceFunction { + case "", vectordata.CosineSimilarity: + metric = "cosine" + case vectordata.DotProductSimilarity: + metric = "dotProduct" + case vectordata.EuclideanDistance: + metric = "euclidean" + default: + return nil, fmt.Errorf("property %q: unsupported distance function %q", p.Name, p.DistanceFunction) + } + algorithm := vectorAlgorithm{Name: p.StorageName + "AlgoConfig"} + switch p.IndexKind { + case "", vectordata.HNSW: + algorithm.Kind, algorithm.HNSWParameters = "hnsw", &vectorParameters{Metric: metric} + case vectordata.Flat: + algorithm.Kind, algorithm.ExhaustiveKNNParameters = "exhaustiveKnn", &vectorParameters{Metric: metric} + default: + return nil, fmt.Errorf("property %q: unsupported index kind %q", p.Name, p.IndexKind) + } + definition.VectorSearch.Algorithms = append(definition.VectorSearch.Algorithms, algorithm) + definition.VectorSearch.Profiles = append(definition.VectorSearch.Profiles, vectorProfile{ + Name: field.VectorSearchProfile, Algorithm: algorithm.Name, + }) + } + definition.Fields = append(definition.Fields, field) + } + return definition, nil +} + +func dataFieldType(typ reflect.Type) (string, error) { + if typ.Kind() == reflect.Pointer { + return dataFieldType(typ.Elem()) + } + if typ.Kind() == reflect.Slice || typ.Kind() == reflect.Array { + element, err := dataFieldType(typ.Elem()) + if err != nil { + return "", err + } + return "Collection(" + element + ")", nil + } + if typ == reflect.TypeFor[time.Time]() { + return "Edm.DateTimeOffset", nil + } + switch typ.Kind() { + case reflect.String: + return "Edm.String", nil + case reflect.Bool: + return "Edm.Boolean", nil + case reflect.Int32: + return "Edm.Int32", nil + case reflect.Int, reflect.Int64: + return "Edm.Int64", nil + case reflect.Float32, reflect.Float64: + return "Edm.Double", nil + default: + return "", fmt.Errorf("unsupported data property type %v", typ) + } +} diff --git a/provider/azureprovider/schema_test.go b/provider/azureprovider/schema_test.go new file mode 100644 index 000000000..79fff32f9 --- /dev/null +++ b/provider/azureprovider/schema_test.go @@ -0,0 +1,176 @@ +// Copyright (c) Microsoft. All rights reserved. + +package azureprovider_test + +import ( + "encoding/json" + "net/http" + "reflect" + "testing" + "time" + + "github.com/microsoft/agent-framework-go/provider/azureprovider" + "github.com/microsoft/agent-framework-go/vectordata" +) + +func TestAISearch_IndexFieldTypes(t *testing.T) { + for _, test := range []struct { + typ reflect.Type + want string + }{ + {reflect.TypeFor[string](), "Edm.String"}, + {reflect.TypeFor[*string](), "Edm.String"}, + {reflect.TypeFor[bool](), "Edm.Boolean"}, + {reflect.TypeFor[int32](), "Edm.Int32"}, + {reflect.TypeFor[int](), "Edm.Int64"}, + {reflect.TypeFor[int64](), "Edm.Int64"}, + {reflect.TypeFor[float32](), "Edm.Double"}, + {reflect.TypeFor[float64](), "Edm.Double"}, + {reflect.TypeFor[time.Time](), "Edm.DateTimeOffset"}, + {reflect.TypeFor[*time.Time](), "Edm.DateTimeOffset"}, + {reflect.TypeFor[[]string](), "Collection(Edm.String)"}, + {reflect.TypeFor[[2]int32](), "Collection(Edm.Int32)"}, + {reflect.TypeFor[[]bool](), "Collection(Edm.Boolean)"}, + {reflect.TypeFor[[]float32](), "Collection(Edm.Double)"}, + {reflect.TypeFor[[]time.Time](), "Collection(Edm.DateTimeOffset)"}, + } { + t.Run(test.typ.String(), func(t *testing.T) { + var field map[string]any + store := newStore(t, func(request *http.Request) (*http.Response, error) { + if request.Method == http.MethodGet { + return response(http.StatusNotFound, `{"error":{"code":"NotFound"}}`), nil + } + fields := jsonBody(t, request)["fields"].([]any) + field = fields[1].(map[string]any) + return response(http.StatusCreated, `{}`), nil + }, azureprovider.AISearchConfig{}) + definition := &vectordata.Definition{Properties: []vectordata.Property{ + &vectordata.KeyProperty{Name: "ID", Type: reflect.TypeFor[string]()}, + &vectordata.DataProperty{Name: "Value", Type: test.typ, IsIndexed: true}, + }} + collection := openCollection[string, map[string]any](t, store, definition) + if err := collection.EnsureExists(t.Context()); err != nil { + t.Fatal(err) + } + collectionType := test.typ.Kind() == reflect.Array || test.typ.Kind() == reflect.Slice + if field["type"] != test.want || field["filterable"] != true || field["sortable"] != !collectionType { + t.Fatalf("mapped field = %v", field) + } + }) + } +} + +func TestAISearch_VectorIndexConfiguration(t *testing.T) { + for _, test := range []struct { + index vectordata.IndexKind + distance vectordata.DistanceFunction + kind string + parameters string + metric string + }{ + {vectordata.HNSW, vectordata.CosineSimilarity, "hnsw", "hnswParameters", "cosine"}, + {vectordata.Flat, vectordata.DotProductSimilarity, "exhaustiveKnn", "exhaustiveKnnParameters", "dotProduct"}, + {vectordata.Flat, vectordata.EuclideanDistance, "exhaustiveKnn", "exhaustiveKnnParameters", "euclidean"}, + } { + t.Run(string(test.index)+"/"+string(test.distance), func(t *testing.T) { + var schema map[string]any + store := newStore(t, func(request *http.Request) (*http.Response, error) { + if request.Method == http.MethodGet { + return response(http.StatusNotFound, `{"error":{"code":"NotFound"}}`), nil + } + schema = jsonBody(t, request) + return response(http.StatusCreated, `{}`), nil + }, azureprovider.AISearchConfig{}) + definition := &vectordata.Definition{Properties: []vectordata.Property{ + &vectordata.VectorProperty{Name: "Vector", StorageName: "embedding", Dimensions: 3, IndexKind: test.index, DistanceFunction: test.distance}, + }} + if err := openCollection[string, document](t, store, definition).EnsureExists(t.Context()); err != nil { + t.Fatal(err) + } + vectorSearch := schema["vectorSearch"].(map[string]any) + algorithm := vectorSearch["algorithms"].([]any)[0].(map[string]any) + profile := vectorSearch["profiles"].([]any)[0].(map[string]any) + if algorithm["kind"] != test.kind || algorithm[test.parameters].(map[string]any)["metric"] != test.metric || + profile["name"] != "embeddingProfile" || profile["algorithm"] != "embeddingAlgoConfig" { + t.Fatalf("vector index = %v", vectorSearch) + } + }) + } +} + +func TestAISearch_RejectsUnsupportedSchemas(t *testing.T) { + store := newStore(t, func(request *http.Request) (*http.Response, error) { + if request.Method != http.MethodGet { + t.Error("invalid index configuration was submitted") + } + return response(http.StatusNotFound, `{"error":{"code":"NotFound"}}`), nil + }, azureprovider.AISearchConfig{}) + for _, typ := range []reflect.Type{ + reflect.TypeFor[map[string]any](), reflect.TypeFor[uint64](), reflect.TypeFor[[][]string](), reflect.TypeFor[complex64](), + } { + definition := &vectordata.Definition{Properties: []vectordata.Property{ + &vectordata.KeyProperty{Name: "ID", Type: reflect.TypeFor[string]()}, + &vectordata.DataProperty{Name: "Data", Type: typ}, + }} + if _, err := store.DynamicCollection("records", definition); err == nil { + t.Fatalf("accepted unsupported data type %v", typ) + } + } + for _, storageName := range []string{"invalid-name", "value/child", "value eq 'x'", "1st", "@search.action"} { + if _, err := vectordata.GetCollection[string, document](store.Store, "records", &vectordata.Definition{ + Properties: []vectordata.Property{&vectordata.DataProperty{Name: "Text", StorageName: storageName}}, + }); err == nil { + t.Fatalf("accepted unsafe field name %q", storageName) + } + } + for _, property := range []vectordata.Property{ + &vectordata.VectorProperty{Name: "Vector", Dimensions: 3, IndexKind: vectordata.IVFFlat}, + &vectordata.VectorProperty{Name: "Vector", Dimensions: 3, DistanceFunction: vectordata.CosineDistance}, + &vectordata.DataProperty{Name: "Rank", IsFullTextIndexed: true}, + } { + collection := openCollection[string, document](t, store, &vectordata.Definition{Properties: []vectordata.Property{property}}) + if err := collection.EnsureExists(t.Context()); err == nil { + t.Fatalf("accepted unsupported index configuration %+v", property) + } + } +} + +func TestAISearch_NullableAndNamedDynamicRecords(t *testing.T) { + type named map[string]any + definition := &vectordata.Definition{Properties: []vectordata.Property{ + &vectordata.KeyProperty{Name: "ID", Type: reflect.TypeFor[string]()}, + &vectordata.DataProperty{Name: "Text", Type: reflect.TypeFor[*string](), IsFullTextIndexed: true}, + &vectordata.DataProperty{Name: "Number", Type: reflect.TypeFor[*int32]()}, + }} + store := newStore(t, func(request *http.Request) (*http.Response, error) { + switch request.Method { + case http.MethodGet: + if request.URL.Path == "/indexes('records')" { + return response(http.StatusNotFound, `{"error":{"code":"NotFound"}}`), nil + } + return response(http.StatusOK, `{"ID":"one","Text":null,"Number":null}`), nil + default: + body := jsonBody(t, request) + if request.URL.Path == "/indexes" { + if field := body["fields"].([]any)[1].(map[string]any); field["searchable"] != true || field["type"] != "Edm.String" { + t.Errorf("nullable full-text field = %v", field) + } + return response(http.StatusCreated, `{}`), nil + } + return indexedResponse(t, body["value"].([]any), "ID"), nil + } + }, azureprovider.AISearchConfig{}) + collection := openCollection[string, *named](t, store, definition) + if err := collection.EnsureExists(t.Context()); err != nil { + t.Fatal(err) + } + input := named{"ID": "one", "Text": (*string)(nil), "Number": (*int32)(nil)} + if err := collection.Upsert(t.Context(), &input); err != nil { + t.Fatal(err) + } + output, found, err := collection.Get(t.Context(), "one", nil) + if err != nil || !found || (*output)["Text"] != nil || (*output)["Number"] != nil { + data, _ := json.Marshal(output) + t.Fatalf("nullable dynamic record = (%s, %v, %v)", data, found, err) + } +} diff --git a/provider/azureprovider/search.go b/provider/azureprovider/search.go new file mode 100644 index 000000000..12d325784 --- /dev/null +++ b/provider/azureprovider/search.go @@ -0,0 +1,267 @@ +// Copyright (c) Microsoft. All rights reserved. + +package azureprovider + +import ( + "context" + "encoding/json" + "fmt" + "iter" + "net/http" + "strings" + + "github.com/microsoft/agent-framework-go/embedding" + "github.com/microsoft/agent-framework-go/vectordata" + "github.com/microsoft/agent-framework-go/vectordata/providerservices" +) + +type searchRequest struct { + Search *string `json:"search,omitempty"` + SearchFields string `json:"searchFields,omitempty"` + Filter string `json:"filter,omitempty"` + OrderBy string `json:"orderby,omitempty"` + Select string `json:"select,omitempty"` + Top int `json:"top"` + Skip int `json:"skip"` + VectorFilterMode string `json:"vectorFilterMode,omitempty"` + VectorQueries []vectorQuery `json:"vectorQueries,omitempty"` +} + +type vectorQuery struct { + Kind string `json:"kind"` + Fields string `json:"fields"` + K int `json:"k"` + Vector *[]float32 `json:"vector,omitempty"` + Text *string `json:"text,omitempty"` +} + +func (c *collection) GetFiltered(ctx context.Context, filter vectordata.Filter, top int, options *vectordata.FilteredRecordRetrievalOptions) iter.Seq2[any, error] { + return func(yield func(any, error) bool) { + if err := options.Validate(); err != nil { + yield(nil, failure(c.name, "Search", err)) + return + } + var settings vectordata.FilteredRecordRetrievalOptions + if options != nil { + settings = *options + } + request, err := c.baseSearchRequest(&filter, top, settings.Skip, settings.IncludeVectors) + if err != nil { + yield(nil, failure(c.name, "Search", err)) + return + } + ordering := make([]string, len(settings.OrderBy)) + for i, order := range settings.OrderBy { + name, _, err := c.filterProperty(order.Property) + if err != nil { + yield(nil, failure(c.name, "Search", err)) + return + } + ordering[i] = name + if order.Descending { + ordering[i] += " desc" + } + } + request.OrderBy = strings.Join(ordering, ",") + for result, err := range c.search(ctx, request, settings.IncludeVectors, nil) { + if err != nil { + yield(nil, err) + return + } + if !yield(result.Record, nil) { + return + } + } + } +} + +func (c *collection) Search(ctx context.Context, value any, top int, options *vectordata.SearchOptions) iter.Seq2[vectordata.SearchResult[any], error] { + return func(yield func(vectordata.SearchResult[any], error) bool) { + if err := options.Validate(); err != nil { + yield(vectordata.SearchResult[any]{}, failure(c.name, "Search", err)) + return + } + var settings vectordata.SearchOptions + if options != nil { + settings = *options + } + request, err := c.vectorSearchRequest(ctx, value, top, settings) + if err != nil { + yield(vectordata.SearchResult[any]{}, failure(c.name, "Search", err)) + return + } + for result, err := range c.search(ctx, request, settings.IncludeVectors, settings.ScoreThreshold) { + if !yield(result, err) { + return + } + } + } +} + +func (c *collection) HybridSearch(ctx context.Context, value any, keywords []string, top int, options *vectordata.HybridSearchOptions) iter.Seq2[vectordata.SearchResult[any], error] { + return func(yield func(vectordata.SearchResult[any], error) bool) { + if keywords == nil { + yield(vectordata.SearchResult[any]{}, failure(c.name, "Search", fmt.Errorf("keywords must not be nil"))) + return + } + if err := options.Validate(); err != nil { + yield(vectordata.SearchResult[any]{}, failure(c.name, "Search", err)) + return + } + var settings vectordata.HybridSearchOptions + if options != nil { + settings = *options + } + textProperty, err := c.model.FullTextProperty(settings.AdditionalProperty) + if err != nil { + yield(vectordata.SearchResult[any]{}, failure(c.name, "Search", err)) + return + } + request, err := c.vectorSearchRequest(ctx, value, top, vectordata.SearchOptions{ + Filter: settings.Filter, VectorProperty: settings.VectorProperty, + Skip: settings.Skip, IncludeVectors: settings.IncludeVectors, + }) + if err != nil { + yield(vectordata.SearchResult[any]{}, failure(c.name, "Search", err)) + return + } + request.Search = new(strings.Join(keywords, " ")) + request.SearchFields = textProperty.StorageName + for result, err := range c.search(ctx, request, settings.IncludeVectors, settings.ScoreThreshold) { + if !yield(result, err) { + return + } + } + } +} + +func (c *collection) baseSearchRequest(filter *vectordata.Filter, top, skip int, includeVectors bool) (*searchRequest, error) { + if top <= 0 { + return nil, fmt.Errorf("top must be positive") + } + if err := c.validateRetrieval(includeVectors); err != nil { + return nil, err + } + expression, err := c.translateFilter(filter) + if err != nil { + return nil, err + } + request := &searchRequest{Filter: expression, Top: top, Skip: skip} + if !includeVectors { + request.Select = c.selectedFields() + } + return request, nil +} + +func (c *collection) vectorSearchRequest(ctx context.Context, value any, top int, options vectordata.SearchOptions) (*searchRequest, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + request, err := c.baseSearchRequest(options.Filter, top, options.Skip, options.IncludeVectors) + if err != nil { + return nil, err + } + if top > int(^uint(0)>>1)-options.Skip { + return nil, fmt.Errorf("top plus skip exceeds the supported integer range") + } + property, err := c.model.VectorProperty(options.VectorProperty) + if err != nil { + return nil, err + } + query, err := searchVector(ctx, value, property) + if err != nil { + return nil, err + } + query.Fields, query.K = property.StorageName, top+options.Skip + request.VectorQueries = []vectorQuery{query} + request.VectorFilterMode = "preFilter" + return request, nil +} + +func searchVector(ctx context.Context, value any, property *providerservices.VectorPropertyModel) (vectorQuery, error) { + if vector, ok := providerservices.AsVector[float32](value); ok || value == nil { + if vector == nil { + return vectorQuery{}, fmt.Errorf("search requires a non-nil value") + } + return vectorQuery{Kind: "vector", Vector: &vector}, nil + } + if property.SearchEmbeddingGenerator != nil { + output, err := embedding.Generate(ctx, property.SearchEmbeddingGenerator, value, nil) + if err != nil { + return vectorQuery{}, err + } + vector, ok := providerservices.AsVector[float32](output) + if !ok || vector == nil { + return vectorQuery{}, fmt.Errorf("property %q: generator returned an invalid float32 embedding", property.Name) + } + return vectorQuery{Kind: "vector", Vector: &vector}, nil + } + if text, ok := value.(string); ok && property.EmbeddingGenerator == nil { + return vectorQuery{Kind: "text", Text: &text}, nil + } + return vectorQuery{}, fmt.Errorf("property %q has no compatible embedding generator for search input %T", property.Name, value) +} + +func (c *collection) search(ctx context.Context, request *searchRequest, includeVectors bool, threshold *float64) iter.Seq2[vectordata.SearchResult[any], error] { + return func(yield func(vectordata.SearchResult[any], error) bool) { + endpoint := c.store.client.url(indexPath(c.name)+"/docs/search", nil) + var body any = request + seen := 0 + for { + var page struct { + Value []map[string]json.RawMessage `json:"value"` + NextLink string `json:"@odata.nextLink"` + NextPageParameters json.RawMessage `json:"@search.nextPageParameters"` + } + if err := c.store.client.request(ctx, http.MethodPost, endpoint, body, &page, http.StatusOK); err != nil { + yield(vectordata.SearchResult[any]{}, failure(c.name, "Search", err)) + return + } + if page.Value == nil { + yield(vectordata.SearchResult[any]{}, failure(c.name, "Search", fmt.Errorf("response is missing search results"))) + return + } + for _, document := range page.Value { + if err := ctx.Err(); err != nil { + yield(vectordata.SearchResult[any]{}, failure(c.name, "Search", err)) + return + } + record, err := c.decodeRecord(document, includeVectors) + if err != nil { + yield(vectordata.SearchResult[any]{}, failure(c.name, "Search", err)) + return + } + var score *float64 + if raw, found := document["@search.score"]; found { + if err := json.Unmarshal(raw, &score); err != nil { + yield(vectordata.SearchResult[any]{}, failure(c.name, "Search", fmt.Errorf("invalid search score: %w", err))) + return + } + } + seen++ + if threshold == nil || score == nil || !(*score < *threshold) { + if !yield(vectordata.SearchResult[any]{Record: record, Score: score}, nil) { + return + } + } + if seen >= request.Top { + return + } + } + if page.NextLink == "" { + return + } + if len(page.NextPageParameters) == 0 || string(page.NextPageParameters) == "null" { + yield(vectordata.SearchResult[any]{}, failure(c.name, "Search", fmt.Errorf("search continuation is missing its request parameters"))) + return + } + var err error + endpoint, err = c.store.client.nextURL(endpoint, page.NextLink) + if err != nil { + yield(vectordata.SearchResult[any]{}, failure(c.name, "Search", err)) + return + } + body = page.NextPageParameters + } + } +} diff --git a/provider/azureprovider/search_test.go b/provider/azureprovider/search_test.go new file mode 100644 index 000000000..3dcf6a5b2 --- /dev/null +++ b/provider/azureprovider/search_test.go @@ -0,0 +1,269 @@ +// Copyright (c) Microsoft. All rights reserved. + +package azureprovider_test + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "reflect" + "slices" + "testing" + + "github.com/microsoft/agent-framework-go/embedding" + "github.com/microsoft/agent-framework-go/provider/azureprovider" + "github.com/microsoft/agent-framework-go/vectordata" +) + +func TestCollection_VectorHybridAndServerTextSearch(t *testing.T) { + for _, hybrid := range []bool{false, true} { + for _, test := range []struct { + name string + value any + kind string + }{ + {"vector", []float32{1, 2, 3}, "vector"}, + {"embedding", &embedding.Of[float32]{Vector: embedding.VectorOf[float32]{1, 2, 3}}, "vector"}, + {"server text", "search string", "text"}, + } { + t.Run(fmt.Sprintf("%s/hybrid=%t", test.name, hybrid), func(t *testing.T) { + calls := 0 + var body map[string]any + store := newStore(t, func(request *http.Request) (*http.Response, error) { + calls++ + if request.URL.Path != "/indexes('records')/docs/search" { + t.Errorf("search path = %q", request.URL.Path) + } + body = jsonBody(t, request) + return response(http.StatusOK, `{"value":[ + {"id":"one","text":"first","vector":[1,2,3],"@search.score":1}, + {"id":"two","text":"second","vector":[3,2,1],"@search.score":0.75}, + {"id":"three","@search.score":0.5} + ]}`), nil + }, azureprovider.AISearchConfig{}) + collection := openCollection[string, document](t, store, nil) + filter := &vectordata.Filter{Operator: vectordata.Equal, Property: "Text", Value: "filter"} + var results []vectordata.SearchResult[document] + if hybrid { + searcher := collection.(vectordata.KeywordHybridSearcher[document]) + sequence := searcher.HybridSearch(t.Context(), test.value, []string{"red", "blue"}, 5, + &vectordata.HybridSearchOptions{Filter: filter, Skip: 3, IncludeVectors: true, ScoreThreshold: new(0.75)}) + if calls != 0 { + t.Fatal("hybrid search started eagerly") + } + results = collect(t, sequence) + if body["search"] != "red blue" || body["searchFields"] != "text" { + t.Fatalf("hybrid text request = %v", body) + } + if _, found := body["select"]; found { + t.Fatal("hybrid IncludeVectors=true excluded vectors") + } + if !slices.Equal(results[0].Record.Vector, []float32{1, 2, 3}) { + t.Fatalf("hybrid vectors = %v", results[0].Record.Vector) + } + } else { + sequence := collection.Search(t.Context(), test.value, 5, + &vectordata.SearchOptions{Filter: filter, Skip: 3, ScoreThreshold: new(0.75)}) + if calls != 0 { + t.Fatal("vector search started eagerly") + } + results = collect(t, sequence) + if _, found := body["search"]; found { + t.Fatal("pure vector search included keyword text") + } + if body["select"] != "id,text,rank,tags" || results[0].Record.Vector != nil { + t.Fatalf("default vector projection = %v", body) + } + } + if len(results) != 2 || results[0].Record.ID != "one" || results[1].Record.ID != "two" || + results[1].Score == nil || *results[1].Score != 0.75 || calls != 1 { + t.Fatalf("search results = %v, calls = %d", results, calls) + } + query := body["vectorQueries"].([]any)[0].(map[string]any) + if query["kind"] != test.kind || query["fields"] != "vector" || query["k"] != float64(8) || + body["top"] != float64(5) || body["skip"] != float64(3) || + body["filter"] != "(text eq 'filter')" || body["vectorFilterMode"] != "preFilter" { + t.Fatalf("vector search request = %v", body) + } + if test.kind == "text" && query["text"] != "search string" { + t.Fatalf("server vectorization request = %v", query) + } + if filter.Property != "Text" { + t.Fatal("search modified the input filter") + } + }) + } + } +} + +func TestCollection_SearchPaginationAndEarlyStop(t *testing.T) { + for _, stop := range []bool{false, true} { + t.Run(fmt.Sprint(stop), func(t *testing.T) { + calls := 0 + var bodies []*trackingBody + store := newStore(t, func(request *http.Request) (*http.Response, error) { + calls++ + body := jsonBody(t, request) + payload := `{"value":[{"id":"two","@search.score":0.8}]}` + if calls == 1 { + payload = `{ + "value":[{"id":"one","@search.score":0.9}], + "@odata.nextLink":"/indexes('records')/docs/search?api-version=2025-09-01", + "@search.nextPageParameters":{"top":2,"skip":1,"filter":"rank gt 0"} + }` + } else if body["skip"] != float64(1) || body["filter"] != "rank gt 0" { + t.Errorf("continuation body = %v", body) + } + result := response(http.StatusOK, payload) + reader := &trackingBody{Reader: bytes.NewReader([]byte(payload))} + bodies = append(bodies, reader) + result.Body = reader + return result, nil + }, azureprovider.AISearchConfig{}) + collection := openCollection[string, document](t, store, nil) + var ids []string + for result, err := range collection.Search(t.Context(), []float32{1, 2, 3}, 2, nil) { + if err != nil { + t.Fatal(err) + } + for _, body := range bodies { + if body.closes.Load() == 0 { + t.Error("response body remained open while yielding a record") + } + } + ids = append(ids, result.Record.ID) + if stop { + break + } + } + want := []string{"one", "two"} + if stop { + want = want[:1] + } + if !slices.Equal(ids, want) || calls != len(want) { + t.Fatalf("records = %v, HTTP calls = %d", ids, calls) + } + }) + } +} + +func TestCollection_SearchFailuresAndContinuationOrigin(t *testing.T) { + for _, payload := range []string{ + `{}`, + `{"value":[{"id":"one","@search.score":"invalid"}]}`, + `{"value":[{"id":"one"},{"text":"missing key"}]}`, + `{"value":[],"@odata.nextLink":"https://other.example/indexes","@search.nextPageParameters":{"skip":1}}`, + `{"value":[],"@odata.nextLink":"https://user:secret@example.search.windows.net/indexes","@search.nextPageParameters":{"skip":1}}`, + `{"value":[],"@odata.nextLink":"/indexes('records')/docs/search"}`, + } { + t.Run(payload, func(t *testing.T) { + calls := 0 + store := newStore(t, func(*http.Request) (*http.Response, error) { + calls++ + return response(http.StatusOK, payload), nil + }, azureprovider.AISearchConfig{}) + collection := openCollection[string, document](t, store, nil) + _ = sequenceError(t, collection.Search(t.Context(), []float32{1, 2, 3}, 3, nil)) + if calls != 1 { + t.Fatalf("invalid response triggered %d requests", calls) + } + }) + } +} + +func TestCollection_SearchInputValidation(t *testing.T) { + calls := 0 + store := newStore(t, func(*http.Request) (*http.Response, error) { + calls++ + return response(http.StatusOK, `{"value":[]}`), nil + }, azureprovider.AISearchConfig{}) + collection := openCollection[string, document](t, store, nil) + for _, test := range []struct { + value any + top int + options *vectordata.SearchOptions + }{ + {nil, 1, nil}, + {[]float32(nil), 1, nil}, + {(*embedding.Of[float32])(nil), 1, nil}, + {[]float32{1, 2, 3}, 0, nil}, + {[]float32{1, 2, 3}, 1, &vectordata.SearchOptions{Skip: -1}}, + {[]float32{1, 2, 3}, 2, &vectordata.SearchOptions{Skip: int(^uint(0) >> 1)}}, + {[]float32{1, 2, 3}, 1, &vectordata.SearchOptions{VectorProperty: "Text"}}, + {42, 1, nil}, + } { + _ = sequenceError(t, collection.Search(t.Context(), test.value, test.top, test.options)) + } + hybrid := collection.(vectordata.KeywordHybridSearcher[document]) + _ = sequenceError(t, hybrid.HybridSearch(t.Context(), []float32{1, 2, 3}, nil, 1, nil)) + _ = sequenceError(t, hybrid.HybridSearch(t.Context(), []float32{1, 2, 3}, []string{"text"}, 1, + &vectordata.HybridSearchOptions{AdditionalProperty: "Rank"})) + if calls != 0 { + t.Fatal("invalid search input performed I/O") + } +} + +func TestCollection_CancelSearchBetweenPages(t *testing.T) { + calls := 0 + store := newStore(t, func(*http.Request) (*http.Response, error) { + calls++ + return response(http.StatusOK, `{ + "value":[{"id":"one"}], + "@odata.nextLink":"/indexes('records')/docs/search?api-version=2025-09-01", + "@search.nextPageParameters":{"top":2,"skip":1} + }`), nil + }, azureprovider.AISearchConfig{}) + collection := openCollection[string, document](t, store, nil) + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + count := 0 + var failure error + for _, err := range collection.Search(ctx, []float32{1, 2, 3}, 2, nil) { + if err != nil { + failure = err + continue + } + count++ + cancel() + } + if count != 1 || calls != 1 || !errors.Is(failure, context.Canceled) { + t.Fatalf("count=%d calls=%d failure=%v", count, calls, failure) + } +} + +func TestCollection_FilteredRetrievalProjection(t *testing.T) { + for _, include := range []bool{false, true} { + var requestBody map[string]any + store := newStore(t, func(request *http.Request) (*http.Response, error) { + requestBody = jsonBody(t, request) + return response(http.StatusOK, `{"value":[{"id":"one","vector":[1,2,3]}]}`), nil + }, azureprovider.AISearchConfig{}) + collection := openCollection[string, document](t, store, nil) + results := collect(t, collection.GetFiltered(t.Context(), vectordata.Filter{Operator: vectordata.Equal, Property: "ID", Value: "one"}, + 1, &vectordata.FilteredRecordRetrievalOptions{IncludeVectors: include})) + _, selected := requestBody["select"] + if len(results) != 1 || selected == include || (results[0].Vector != nil) != include { + t.Fatalf("projection=%v results=%v", requestBody, results) + } + if _, found := requestBody["vectorQueries"]; found { + t.Fatal("filtered retrieval included a vector query") + } + } +} + +func TestCollection_SearchScoreShapes(t *testing.T) { + store := newStore(t, func(*http.Request) (*http.Response, error) { + return response(http.StatusOK, `{"value":[{"id":"zero","@search.score":0},{"id":"null","@search.score":null},{"id":"missing"}]}`), nil + }, azureprovider.AISearchConfig{}) + results := collect(t, openCollection[string, document](t, store, nil).Search(t.Context(), []float32{1, 2, 3}, 3, nil)) + if len(results) != 3 || results[0].Score == nil || *results[0].Score != 0 || results[1].Score != nil || results[2].Score != nil { + data, _ := json.Marshal(results) + t.Fatalf("search scores = %s", data) + } + if !reflect.DeepEqual([]string{results[0].Record.ID, results[1].Record.ID, results[2].Record.ID}, []string{"zero", "null", "missing"}) { + t.Fatal("score handling changed result order") + } +} diff --git a/vectordata/collection.go b/vectordata/collection.go new file mode 100644 index 000000000..af154752c --- /dev/null +++ b/vectordata/collection.go @@ -0,0 +1,135 @@ +// Copyright (c) Microsoft. All rights reserved. + +package vectordata + +import ( + "context" + "fmt" + "iter" + "reflect" +) + +type typedCollection[K, R any] struct { + raw Collection[any, any] +} + +func (c *typedCollection[K, R]) Name() string { return c.raw.Name() } + +func (c *typedCollection[K, R]) Exists(ctx context.Context) (bool, error) { + return c.raw.Exists(ctx) +} + +func (c *typedCollection[K, R]) EnsureExists(ctx context.Context) error { + return c.raw.EnsureExists(ctx) +} + +func (c *typedCollection[K, R]) EnsureDeleted(ctx context.Context) error { + return c.raw.EnsureDeleted(ctx) +} + +func (c *typedCollection[K, R]) Get(ctx context.Context, key K, options *RecordRetrievalOptions) (R, bool, error) { + value, found, err := c.raw.Get(ctx, key, options) + var zero R + if err != nil || !found { + return zero, found, err + } + record, err := typedRecord[R](value) + return record, found, err +} + +func (c *typedCollection[K, R]) GetFiltered(ctx context.Context, filter Filter, top int, options *FilteredRecordRetrievalOptions) iter.Seq2[R, error] { + return func(yield func(R, error) bool) { + for value, err := range c.raw.GetFiltered(ctx, filter, top, options) { + var record R + if err == nil { + record, err = typedRecord[R](value) + } + if !yield(record, err) || err != nil { + return + } + } + } +} + +func (c *typedCollection[K, R]) Delete(ctx context.Context, key K) error { + return c.raw.Delete(ctx, key) +} + +func (c *typedCollection[K, R]) Upsert(ctx context.Context, records ...R) error { + return c.raw.Upsert(ctx, untypedValues(records)...) +} + +func (c *typedCollection[K, R]) Search(ctx context.Context, value any, top int, options *SearchOptions) iter.Seq2[SearchResult[R], error] { + return func(yield func(SearchResult[R], error) bool) { + yieldSearchResults(c.raw.Search(ctx, value, top, options), yield) + } +} + +func (c *typedCollection[K, R]) GetBatch(ctx context.Context, keys []K, options *RecordRetrievalOptions) iter.Seq2[R, error] { + return func(yield func(R, error) bool) { + for value, err := range GetBatch(ctx, c.raw, untypedValues(keys), options) { + var record R + if err == nil { + record, err = typedRecord[R](value) + } + if !yield(record, err) || err != nil { + return + } + } + } +} + +func (c *typedCollection[K, R]) DeleteBatch(ctx context.Context, keys []K) error { + return DeleteBatch(ctx, c.raw, untypedValues(keys)) +} + +type typedHybridCollection[K, R any] struct { + *typedCollection[K, R] + hybrid KeywordHybridSearcher[any] +} + +func (c *typedHybridCollection[K, R]) HybridSearch(ctx context.Context, value any, keywords []string, top int, options *HybridSearchOptions) iter.Seq2[SearchResult[R], error] { + return func(yield func(SearchResult[R], error) bool) { + yieldSearchResults(c.hybrid.HybridSearch(ctx, value, keywords, top, options), yield) + } +} + +func yieldSearchResults[R any](results iter.Seq2[SearchResult[any], error], yield func(SearchResult[R], error) bool) { + for result, err := range results { + var typed SearchResult[R] + if err == nil { + typed.Record, err = typedRecord[R](result.Record) + if err == nil { + typed.Score = result.Score + } + } + if !yield(typed, err) || err != nil { + return + } + } +} + +func typedRecord[R any](value any) (R, error) { + if record, ok := value.(R); ok { + return record, nil + } + var zero R + if value == nil { + switch reflect.TypeFor[R]().Kind() { + case reflect.Chan, reflect.Func, reflect.Interface, reflect.Map, reflect.Pointer, reflect.Slice, reflect.UnsafePointer: + return zero, nil + } + } + return zero, fmt.Errorf("vectordata: provider returned record type %T, want %v", value, reflect.TypeFor[R]()) +} + +func untypedValues[T any](values []T) []any { + if values == nil { + return nil + } + result := make([]any, len(values)) + for i, value := range values { + result[i] = value + } + return result +} diff --git a/vectordata/definition.go b/vectordata/definition.go index 810f10c7b..37f326053 100644 --- a/vectordata/definition.go +++ b/vectordata/definition.go @@ -118,8 +118,10 @@ type VectorProperty struct { // tags, while empty storage, index, and distance names retain tag settings. // Dynamic map[string]any records require explicit property types. // -// A property generator takes precedence over EmbeddingGenerator, which takes -// precedence over CollectionOptions.EmbeddingGenerator. +// EmbeddingGenerator supplies the default generator for vector properties, +// overriding any store default. A property's generator takes precedence. +// Use [embedding.Bind] to configure typed generators; input and output types +// may differ between vector properties. type Definition struct { Properties []Property EmbeddingGenerator *embedding.BoundGenerator diff --git a/vectordata/inmemory/collection.go b/vectordata/inmemory/collection.go index 654c5802e..64aca3158 100644 --- a/vectordata/inmemory/collection.go +++ b/vectordata/inmemory/collection.go @@ -16,46 +16,53 @@ import ( ) // Collection is a configured handle to an in-memory collection. Construct one -// with [Store.Collection] or [NewCollection]. Its schema is fixed at construction. +// with [NewCollection]. Its schema is fixed at construction. A shared [Store] +// also creates collections through its embedded [vectordata.Store]. type Collection[K comparable, R any] struct { - store *Store[K, R] - name string - model *providerservices.CollectionModel - schema recordSchema + store *store + name string + model *providerservices.CollectionModel + schema recordSchema + recordType reflect.Type } // NewCollection creates a handle backed by its own independent store. +// A nil definition uses the record's vectordata tags. // The collection does not exist until [Collection.EnsureExists] is called. -func NewCollection[K comparable, R any](name string, options vectordata.CollectionOptions) (*Collection[K, R], error) { - return new(Store[K, R]).collection(name, options) +func NewCollection[K comparable, R any](name string, definition *vectordata.Definition) (*Collection[K, R], error) { + return newCollection[K, R](new(store), name, reflect.TypeFor[K](), reflect.TypeFor[R](), definition) } -func (s *Store[K, R]) collection(name string, options vectordata.CollectionOptions) (*Collection[K, R], error) { +func newCollection[K comparable, R any](s *store, name string, keyType, recordType reflect.Type, definition *vectordata.Definition) (*Collection[K, R], error) { if strings.TrimSpace(name) == "" { return nil, failure(name, "Collection", fmt.Errorf("collection name is required")) } - if options.EmbeddingGenerator == nil { - options.EmbeddingGenerator = s.generator + s.mu.RLock() + data := s.collections[name] + s.mu.RUnlock() + if data != nil && data.recordType != recordType { + return nil, failure(name, "Collection", incompatibleCollection(name, data.recordType, recordType)) } - model, err := providerservices.Build[K, R](options, providerservices.Config{ + model, err := providerservices.Build(keyType, recordType, definition, providerservices.Config{ SupportsMultipleVectors: true, SupportsKeyType: reflect.Type.Comparable, SupportsDataType: func(reflect.Type) bool { return true }, - SupportsVectorType: vectorTypeSupported, + SupportsVectorType: providerservices.SupportsVector[float32], SupportsKeyAutoGeneration: func(typ reflect.Type) bool { return typ == reflect.TypeFor[uuid.UUID]() }, + DefaultEmbeddingGenerator: s.generator, }) if err != nil { return nil, failure(name, "Collection", err) } - schema := newRecordSchema(reflect.TypeFor[R](), model) + schema := newRecordSchema(recordType, model) for _, field := range schema.vectors { if field.vector.Type != field.vector.EmbeddingType && field.vector.EmbeddingGenerator == nil { return nil, failure(name, "Collection", fmt.Errorf("property %q requires an embedding generator; in-memory storage cannot generate embeddings itself", field.name)) } } - return &Collection[K, R]{store: s, name: name, model: model, schema: schema}, nil + return &Collection[K, R]{store: s, name: name, model: model, schema: schema, recordType: recordType}, nil } // Name returns the collection name. @@ -74,10 +81,14 @@ func (c *Collection[K, R]) EnsureExists(ctx context.Context) error { return failure(c.name, "EnsureExists", err) } if c.store.collections == nil { - c.store.collections = make(map[string]map[K]*storedRecord[R]) + c.store.collections = make(map[string]*collectionData) } - if _, exists := c.store.collections[c.name]; !exists { - c.store.collections[c.name] = make(map[K]*storedRecord[R]) + if data, exists := c.store.collections[c.name]; exists { + if data.recordType != c.recordType { + return failure(c.name, "EnsureExists", incompatibleCollection(c.name, data.recordType, c.recordType)) + } + } else { + c.store.collections[c.name] = &collectionData{recordType: c.recordType, records: make(map[any]*storedRecord)} } return nil } @@ -102,20 +113,31 @@ func (c *Collection[K, R]) Get(ctx context.Context, key K, options *vectordata.R if err := validateKey(reflect.ValueOf(key)); err != nil { return zero, false, failure(c.name, "Get", err) } + lookup := normalizeKey(reflect.ValueOf(key), c.schema.key.typ) c.store.mu.RLock() - records, exists := c.store.collections[c.name] - record, found := records[key] + data, exists := c.store.collections[c.name] + var record *storedRecord + if exists && data.recordType == c.recordType { + record = data.records[lookup] + } c.store.mu.RUnlock() if !exists { return zero, false, failure(c.name, "Get", missingCollection(c.name)) } - if !found { + if data.recordType != c.recordType { + return zero, false, failure(c.name, "Get", incompatibleCollection(c.name, data.recordType, c.recordType)) + } + if record == nil { return zero, false, nil } if err := ctx.Err(); err != nil { return zero, false, failure(c.name, "Get", err) } - return record.Record, true, nil + result, err := collectionRecord[R](record.Record) + if err != nil { + return zero, false, failure(c.name, "Get", err) + } + return result, true, nil } // Delete deletes a record if it exists. A missing collection is an error. @@ -126,16 +148,20 @@ func (c *Collection[K, R]) Delete(ctx context.Context, key K) error { if err := validateKey(reflect.ValueOf(key)); err != nil { return failure(c.name, "Delete", err) } + lookup := normalizeKey(reflect.ValueOf(key), c.schema.key.typ) c.store.mu.Lock() defer c.store.mu.Unlock() if err := ctx.Err(); err != nil { return failure(c.name, "Delete", err) } - records, exists := c.store.collections[c.name] + data, exists := c.store.collections[c.name] if !exists { return failure(c.name, "Delete", missingCollection(c.name)) } - delete(records, key) + if data.recordType != c.recordType { + return failure(c.name, "Delete", incompatibleCollection(c.name, data.recordType, c.recordType)) + } + delete(data.records, lookup) return nil } @@ -155,15 +181,24 @@ func (c *Collection[K, R]) Upsert(ctx context.Context, records ...R) error { if len(records) == 0 { return nil } - batch := make([]*storedRecord[R], len(records)) + batch := make([]*storedRecord, len(records)) for i, record := range records { if err := ctx.Err(); err != nil { return failure(c.name, "Upsert", err) } - if err := c.schema.validateRecord(reflect.ValueOf(record)); err != nil { + value := reflect.ValueOf(record) + if !value.IsValid() || !value.Type().AssignableTo(c.recordType) { + return failure(c.name, "Upsert", fmt.Errorf("record %d has type %T, want %v", i, record, c.recordType)) + } + if value.Type() != c.recordType { + typed := reflect.New(c.recordType).Elem() + typed.Set(value) + value = typed + } + if err := c.schema.validateRecord(value); err != nil { return failure(c.name, "Upsert", fmt.Errorf("record %d: %w", i, err)) } - batch[i] = &storedRecord[R]{Record: record} + batch[i] = &storedRecord{Record: value.Interface()} } for _, field := range c.schema.vectors { if field.vector.Type == field.vector.EmbeddingType { @@ -174,7 +209,7 @@ func (c *Collection[K, R]) Upsert(ctx context.Context, records ...R) error { return failure(c.name, "Upsert", err) } for i, output := range result.Embeddings { - vector, ok := nativeVector(output) + vector, ok := providerservices.AsVector[float32](output) if !ok || vector == nil { return failure(c.name, "Upsert", fmt.Errorf("property %q: generator returned an invalid embedding at index %d", field.name, i)) } @@ -194,12 +229,12 @@ func (c *Collection[K, R]) Upsert(ctx context.Context, records ...R) error { if !exists { return failure(c.name, "Upsert", missingCollection(c.name)) } - updates := make(map[K]*storedRecord[R], len(batch)) + updates := make(map[any]*storedRecord, len(batch)) for i, record := range batch { if err := ctx.Err(); err != nil { return failure(c.name, "Upsert", err) } - value := reflect.ValueOf(&record.Record).Elem() + value := reflect.ValueOf(record.Record) if c.schema.keyAutoGenerated { key, _, err := c.schema.key.read(value) if err != nil { @@ -212,14 +247,19 @@ func (c *Collection[K, R]) Upsert(ctx context.Context, records ...R) error { } target := value if target.Kind() == reflect.Struct { - target = target.Addr() + target = reflect.New(value.Type()) + target.Elem().Set(value) } if err := c.schema.key.property.SetValue(target.Interface(), id); err != nil { return failure(c.name, "Upsert", err) } + if value.Kind() == reflect.Struct { + value = target.Elem() + record.Record = value.Interface() + } } } - key, err := recordKey[K](c.schema.key, value) + key, err := recordKey(c.schema.key, value) if err != nil { return failure(c.name, "Upsert", fmt.Errorf("record %d: %w", i, err)) } @@ -228,42 +268,31 @@ func (c *Collection[K, R]) Upsert(ctx context.Context, records ...R) error { return c.commit(ctx, updates, false, "Upsert") } -func (c *Collection[K, R]) commit(ctx context.Context, updates map[K]*storedRecord[R], create bool, operation string) error { +func (c *Collection[K, R]) commit(ctx context.Context, updates map[any]*storedRecord, create bool, operation string) error { c.store.mu.Lock() defer c.store.mu.Unlock() if err := ctx.Err(); err != nil { return failure(c.name, operation, err) } - records, exists := c.store.collections[c.name] + data, exists := c.store.collections[c.name] if !exists { if !create { return failure(c.name, operation, missingCollection(c.name)) } if c.store.collections == nil { - c.store.collections = make(map[string]map[K]*storedRecord[R]) + c.store.collections = make(map[string]*collectionData) } - records = make(map[K]*storedRecord[R]) - c.store.collections[c.name] = records + data = &collectionData{recordType: c.recordType, records: make(map[any]*storedRecord)} + c.store.collections[c.name] = data + } else if data.recordType != c.recordType { + return failure(c.name, operation, incompatibleCollection(c.name, data.recordType, c.recordType)) } - maps.Copy(records, updates) + maps.Copy(data.records, updates) return nil } -func (c *Collection[K, R]) snapshot(ctx context.Context) ([]*storedRecord[R], error) { - if err := ctx.Err(); err != nil { - return nil, err - } - c.store.mu.RLock() - defer c.store.mu.RUnlock() - records, exists := c.store.collections[c.name] - if !exists { - return nil, missingCollection(c.name) - } - result := make([]*storedRecord[R], 0, len(records)) - for _, record := range records { - result = append(result, record) - } - return result, nil +func (c *Collection[K, R]) snapshot(ctx context.Context) ([]*storedRecord, error) { + return c.store.snapshot(ctx, c.name, c.recordType) } func (c *Collection[K, R]) validateRetrieval(includeVectors bool) error { diff --git a/vectordata/inmemory/collection_test.go b/vectordata/inmemory/collection_test.go index 875c691eb..524419ceb 100644 --- a/vectordata/inmemory/collection_test.go +++ b/vectordata/inmemory/collection_test.go @@ -28,7 +28,7 @@ type document struct { var allRecords = vectordata.Filter{Operator: vectordata.NotEqual, Property: "ID", Value: nil} -func createCollection[K comparable, R any](t *testing.T, options vectordata.CollectionOptions) *inmemory.Collection[K, R] { +func createCollection[K comparable, R any](t *testing.T, options *vectordata.Definition) *inmemory.Collection[K, R] { t.Helper() collection, err := inmemory.NewCollection[K, R]("records", options) if err != nil { @@ -70,18 +70,18 @@ func sequenceError[T any](t *testing.T, sequence iter.Seq2[T, error]) error { } func TestStore_LifecycleAndIsolation(t *testing.T) { - var store inmemory.Store[string, document] + store := inmemory.New(inmemory.Config{}) ctx := t.Context() for _, name := range []string{"", " \t\n"} { - if _, err := store.Collection(name, vectordata.CollectionOptions{}); err == nil { + if _, err := vectordata.GetCollection[string, document](store.Store, name, nil); err == nil { t.Fatalf("accepted collection name %q", name) } } - first, err := store.Collection("first", vectordata.CollectionOptions{}) + first, err := vectordata.GetCollection[string, document](store.Store, "first", nil) if err != nil { t.Fatal(err) } - same, err := store.Collection("first", vectordata.CollectionOptions{}) + same, err := vectordata.GetCollection[string, document](store.Store, "first", nil) if err != nil { t.Fatal(err) } @@ -108,7 +108,7 @@ func TestStore_LifecycleAndIsolation(t *testing.T) { if record, found, err := same.Get(ctx, "id", nil); err != nil || !found || record.Text != "original" { t.Fatalf("shared handle Get = %v, %v, %v", record, found, err) } - other := inmemory.New[string, document](inmemory.Config{}) + other := inmemory.New(inmemory.Config{}) if exists, err := other.CollectionExists(ctx, "first"); err != nil || exists { t.Fatalf("independent store shares data: %v, %v", exists, err) } @@ -135,8 +135,257 @@ func TestStore_LifecycleAndIsolation(t *testing.T) { } } +func TestStore_HeterogeneousCollections(t *testing.T) { + store := inmemory.New(inmemory.Config{}) + documents, err := vectordata.GetCollection[string, document](store.Store, "documents", nil) + if err != nil { + t.Fatal(err) + } + type counter struct { + ID int `vectordata:"key"` + Value int `vectordata:"data"` + } + counters, err := vectordata.GetCollection[int, *counter](store.Store, "counters", nil) + if err != nil { + t.Fatal(err) + } + options := dynamicDefinition(reflect.TypeFor[uuid.UUID]()) + dynamic, err := store.DynamicCollection("dynamic", options) + if err != nil { + t.Fatal(err) + } + for _, collection := range []interface{ EnsureExists(context.Context) error }{documents, counters, dynamic} { + if err := collection.EnsureExists(t.Context()); err != nil { + t.Fatal(err) + } + } + doc := document{ID: "one", Text: "document", Vector: []float32{1, 0, 0, 0}} + if err := documents.Upsert(t.Context(), doc); err != nil { + t.Fatal(err) + } + value := &counter{ID: 2, Value: 42} + if err := counters.Upsert(t.Context(), value); err != nil { + t.Fatal(err) + } + entry := map[string]any{"Text": "memory", "Vector": []float32{1, 0, 0, 0}} + if err := dynamic.Upsert(t.Context(), entry); err != nil { + t.Fatal(err) + } + key, ok := entry["ID"].(uuid.UUID) + if !ok || key == uuid.Nil { + t.Fatalf("dynamic collection did not generate a UUID: %v", entry) + } + otherHandle, err := store.DynamicCollection("dynamic", options) + if err != nil { + t.Fatal(err) + } + got, found, err := otherHandle.Get(t.Context(), key, nil) + if err != nil || !found || got["Text"] != "memory" { + t.Fatalf("dynamic Get = (%v, %v, %v)", got, found, err) + } + got["Text"] = "edited" + if entry["Text"] != "edited" { + t.Fatal("dynamic collection lost map identity") + } + if got, found, err := counters.Get(t.Context(), 2, nil); err != nil || !found || got != value { + t.Fatalf("typed pointer Get = (%v, %v, %v)", got, found, err) + } + if got, found, err := documents.Get(t.Context(), "one", nil); err != nil || !found || !reflect.DeepEqual(got, doc) { + t.Fatalf("typed value Get = (%v, %v, %v)", got, found, err) + } + results := collect(t, dynamic.Search(t.Context(), []float32{1, 0, 0, 0}, 1, nil)) + if len(results) != 1 || results[0].Record["ID"] != key { + t.Fatalf("dynamic search results = %v", results) + } + names := collect(t, store.ListCollectionNames(t.Context())) + slices.Sort(names) + if !slices.Equal(names, []string{"counters", "documents", "dynamic"}) { + t.Fatalf("shared namespace = %v", names) + } + if err := store.EnsureCollectionDeleted(t.Context(), "dynamic"); err != nil { + t.Fatal(err) + } + if exists, err := otherHandle.Exists(t.Context()); err != nil || exists { + t.Fatalf("deleted dynamic handle exists = (%v, %v)", exists, err) + } + if got, found, err := documents.Get(t.Context(), "one", nil); err != nil || !found || got.ID != "one" { + t.Fatal("deleting dynamic collection changed another collection") + } + if _, err := store.DynamicCollection("no-schema", nil); err == nil { + t.Fatal("dynamic collection accepted no schema") + } +} + +func TestStore_ConflictingRecordTypes(t *testing.T) { + store := inmemory.New(inmemory.Config{}) + first, err := vectordata.GetCollection[string, document](store.Store, "records", nil) + if err != nil { + t.Fatal(err) + } + second, err := vectordata.GetCollection[string, *document](store.Store, "records", nil) + if err != nil { + t.Fatal(err) + } + if err := first.EnsureExists(t.Context()); err != nil { + t.Fatal(err) + } + if err := second.EnsureExists(t.Context()); err == nil { + t.Fatal("incompatible handle recreated an existing collection") + } + if _, err := vectordata.GetCollection[string, *document](store.Store, "records", nil); err == nil { + t.Fatal("opened an existing collection with an incompatible record type") + } + if err := first.EnsureDeleted(t.Context()); err != nil { + t.Fatal(err) + } + if err := second.EnsureExists(t.Context()); err != nil { + t.Fatal(err) + } + input := &document{ID: "one", Text: "retained", Vector: []float32{1, 0, 0, 0}} + if err := second.Upsert(t.Context(), input); err != nil { + t.Fatal(err) + } + if _, _, err := first.Get(t.Context(), "one", nil); err == nil { + t.Fatal("stale handle read incompatible records") + } + for _, err := range []error{ + first.EnsureExists(t.Context()), + first.Upsert(t.Context(), document{ID: "one", Text: "changed"}), + first.Delete(t.Context(), "one"), + sequenceError(t, first.Search(t.Context(), []float32{1, 0, 0, 0}, 1, nil)), + sequenceError(t, first.GetFiltered(t.Context(), allRecords, 1, nil)), + } { + if err == nil { + t.Fatal("stale handle operated on an incompatible collection") + } + } + if got, found, err := second.Get(t.Context(), "one", nil); err != nil || !found || got != input || got.Text != "retained" { + t.Fatalf("stale operations changed the new collection: (%v, %v, %v)", got, found, err) + } +} + +func TestStore_ConcurrentRecordTypes(t *testing.T) { + store := inmemory.New(inmemory.Config{}) + first, err := vectordata.GetCollection[string, document](store.Store, "records", nil) + if err != nil { + t.Fatal(err) + } + second, err := vectordata.GetCollection[string, *document](store.Store, "records", nil) + if err != nil { + t.Fatal(err) + } + start := make(chan struct{}) + results := make(chan error, 2) + var wg sync.WaitGroup + for _, collection := range []interface{ EnsureExists(context.Context) error }{first, second} { + wg.Go(func() { + <-start + results <- collection.EnsureExists(t.Context()) + }) + } + close(start) + wg.Wait() + close(results) + successes := 0 + for err := range results { + if err == nil { + successes++ + } + } + if successes != 1 { + t.Fatalf("concurrent creation allowed %d record types", successes) + } +} + +func TestStore_GeneratedStructKeys(t *testing.T) { + type entry struct { + ID uuid.UUID `vectordata:"key"` + } + store := inmemory.New(inmemory.Config{}) + collection, err := vectordata.GetCollection[uuid.UUID, entry](store.Store, "records", nil) + if err != nil { + t.Fatal(err) + } + if err := collection.EnsureExists(t.Context()); err != nil { + t.Fatal(err) + } + input := entry{} + if err := collection.Upsert(t.Context(), input); err != nil { + t.Fatal(err) + } + records := collect(t, collection.GetFiltered(t.Context(), allRecords, 2, nil)) + if len(records) != 1 || records[0].ID == uuid.Nil || input.ID != uuid.Nil { + t.Fatalf("generated struct keys = %v, input = %v", records, input) + } + if got, found, err := collection.Get(t.Context(), records[0].ID, nil); err != nil || !found || got != records[0] { + t.Fatalf("generated key lookup = (%v, %v, %v)", got, found, err) + } +} + +func TestStore_AssignableDynamicKeys(t *testing.T) { + type key [2]int + type value [2]int + store := inmemory.New(inmemory.Config{}) + options := dynamicDefinition(reflect.TypeFor[[2]int]()) + collection, err := vectordata.GetCollection[key, map[string]any](store.Store, "records", options) + if err != nil { + t.Fatal(err) + } + dynamic, err := store.DynamicCollection("records", options) + if err != nil { + t.Fatal(err) + } + for _, collection := range []interface{ EnsureExists(context.Context) error }{collection, dynamic} { + if err := collection.EnsureExists(t.Context()); err != nil { + t.Fatal(err) + } + } + input := map[string]any{"ID": value{1, 2}, "Text": "typed"} + if err := collection.Upsert(t.Context(), input); err != nil { + t.Fatal(err) + } + if got, found, err := collection.Get(t.Context(), key{1, 2}, nil); err != nil || !found || got["ID"] != (value{1, 2}) { + t.Fatalf("assignable typed key Get = (%v, %v, %v)", got, found, err) + } + for _, lookup := range []any{[2]int{1, 2}, key{1, 2}, value{1, 2}} { + if got, found, err := dynamic.Get(t.Context(), lookup, nil); err != nil || !found || !reflect.DeepEqual(got, input) { + t.Fatalf("shared dynamic Get(%T) = (%v, %v, %v)", lookup, got, found, err) + } + } + replacement := map[string]any{"ID": [2]int{1, 2}, "Text": "dynamic"} + if err := dynamic.Upsert(t.Context(), replacement); err != nil { + t.Fatal(err) + } + if got, found, err := collection.Get(t.Context(), key{1, 2}, nil); err != nil || !found || !reflect.DeepEqual(got, replacement) { + t.Fatalf("shared typed Get = (%v, %v, %v)", got, found, err) + } + if records := collect(t, dynamic.GetFiltered(t.Context(), allRecords, 10, nil)); len(records) != 1 { + t.Fatalf("shared handles duplicated a logical key: %v", records) + } + if input["ID"] != (value{1, 2}) || replacement["ID"] != ([2]int{1, 2}) { + t.Fatal("key normalization changed the input records") + } + for _, lookup := range []any{[2]int{1, 2}, key{1, 2}, value{1, 2}} { + if err := dynamic.Delete(t.Context(), lookup); err != nil { + t.Fatal(err) + } + if _, found, err := collection.Get(t.Context(), key{1, 2}, nil); err != nil || found { + t.Fatalf("typed Get after dynamic Delete(%T) = (%v, %v)", lookup, found, err) + } + if err := collection.Upsert(t.Context(), input); err != nil { + t.Fatal(err) + } + } + if err := collection.Delete(t.Context(), key{1, 2}); err != nil { + t.Fatal(err) + } + if _, found, err := dynamic.Get(t.Context(), [2]int{1, 2}, nil); err != nil || found { + t.Fatalf("dynamic Get after typed Delete = (%v, %v)", found, err) + } +} + func TestCollection_CRUD(t *testing.T) { - c := createCollection[string, document](t, vectordata.CollectionOptions{}) + c := createCollection[string, document](t, nil) ctx := t.Context() input := document{ID: "first", Text: "initial", Vector: []float32{1, 2, 3, 4}} if err := c.Upsert(ctx, input); err != nil { @@ -178,7 +427,7 @@ func TestCollection_CRUD(t *testing.T) { } func TestCollection_MissingAndInvalidRecords(t *testing.T) { - c, err := inmemory.NewCollection[string, document]("missing", vectordata.CollectionOptions{}) + c, err := inmemory.NewCollection[string, document]("missing", nil) if err != nil { t.Fatal(err) } @@ -209,7 +458,7 @@ func TestCollection_MissingAndInvalidRecords(t *testing.T) { if _, found, err := c.Get(ctx, "valid", nil); err != nil || found { t.Fatalf("invalid batch partially committed: %v, %v", found, err) } - pointers := createCollection[string, *document](t, vectordata.CollectionOptions{}) + pointers := createCollection[string, *document](t, nil) if err := pointers.Upsert(ctx, nil); err == nil { t.Fatal("accepted nil record") } @@ -237,7 +486,7 @@ func TestCollection_RetainsRecordReferences(t *testing.T) { ID: "id", Root: root, Alias: root, Metadata: metadata, Array: [1][]string{{"original"}}, When: time.Now(), private: 42, } - c := createCollection[string, *record](t, vectordata.CollectionOptions{}) + c := createCollection[string, *record](t, nil) if err := c.Upsert(t.Context(), input); err != nil { t.Fatal(err) } @@ -266,7 +515,7 @@ func TestCollection_AcceptsReferenceAndPrivateState(t *testing.T) { Channel chan int hidden *int } - c := createCollection[string, record](t, vectordata.CollectionOptions{}) + c := createCollection[string, record](t, nil) called := false input := record{ ID: "id", Channel: make(chan int), hidden: new(42), @@ -283,7 +532,7 @@ func TestCollection_AcceptsReferenceAndPrivateState(t *testing.T) { } func TestCollection_ReturnsStoredRecordsWithVectors(t *testing.T) { - c := createCollection[string, *document](t, vectordata.CollectionOptions{}) + c := createCollection[string, *document](t, nil) input := &document{ID: "id", Vector: []float32{1, 0, 0, 0}} if err := c.Upsert(t.Context(), input); err != nil { t.Fatal(err) @@ -304,17 +553,17 @@ func TestCollection_ReturnsStoredRecordsWithVectors(t *testing.T) { } } -func dynamicOptions(keyType reflect.Type) vectordata.CollectionOptions { - return vectordata.CollectionOptions{Definition: &vectordata.Definition{Properties: []vectordata.Property{ +func dynamicDefinition(keyType reflect.Type) *vectordata.Definition { + return &vectordata.Definition{Properties: []vectordata.Property{ &vectordata.KeyProperty{Name: "ID", Type: keyType}, &vectordata.DataProperty{Name: "Text", Type: reflect.TypeFor[string]()}, &vectordata.VectorProperty{Name: "Vector", Type: reflect.TypeFor[[]float32](), Dimensions: 4}, - }}} + }} } func TestCollection_DynamicRecords(t *testing.T) { type dynamic map[string]any - c := createCollection[string, dynamic](t, dynamicOptions(reflect.TypeFor[string]())) + c := createCollection[string, dynamic](t, dynamicDefinition(reflect.TypeFor[string]())) input := dynamic{"ID": "id", "Text": "hello", "Vector": []float32{1, 2, 3, 4}, "Other": map[string]int{"value": 1}} if err := c.Upsert(t.Context(), input); err != nil { t.Fatal(err) @@ -343,7 +592,7 @@ func TestCollection_DynamicRecords(t *testing.T) { } }) } - pointers := createCollection[string, *dynamic](t, dynamicOptions(reflect.TypeFor[string]())) + pointers := createCollection[string, *dynamic](t, dynamicDefinition(reflect.TypeFor[string]())) if err := pointers.Upsert(t.Context(), nil); err == nil { t.Fatal("accepted nil pointer to map") } @@ -367,9 +616,9 @@ func TestCollection_UUIDKeys(t *testing.T) { name = fmt.Sprint(*auto) } t.Run(name, func(t *testing.T) { - options := vectordata.CollectionOptions{Definition: &vectordata.Definition{Properties: []vectordata.Property{ + options := &vectordata.Definition{Properties: []vectordata.Property{ &vectordata.KeyProperty{Name: "ID", IsAutoGenerated: auto}, - }}} + }} c := createCollection[uuid.UUID, record](t, options) input := record{Text: "generated"} explicit := record{ID: uuid.New(), Text: "explicit"} @@ -407,7 +656,7 @@ func TestCollection_UUIDKeys(t *testing.T) { } }) } - dynamic := createCollection[uuid.UUID, map[string]any](t, dynamicOptions(reflect.TypeFor[uuid.UUID]())) + dynamic := createCollection[uuid.UUID, map[string]any](t, dynamicDefinition(reflect.TypeFor[uuid.UUID]())) input := map[string]any{"Text": "missing key"} if err := dynamic.Upsert(t.Context(), input); err != nil { t.Fatal(err) @@ -417,7 +666,7 @@ func TestCollection_UUIDKeys(t *testing.T) { t.Fatalf("dynamic key was not generated: %v", got) } type dynamicRecord map[string]any - pointers := createCollection[uuid.UUID, *dynamicRecord](t, dynamicOptions(reflect.TypeFor[uuid.UUID]())) + pointers := createCollection[uuid.UUID, *dynamicRecord](t, dynamicDefinition(reflect.TypeFor[uuid.UUID]())) pointer := dynamicRecord{"Text": "pointer map"} if err := pointers.Upsert(t.Context(), &pointer); err != nil { t.Fatal(err) @@ -435,7 +684,7 @@ func TestCollection_KeyTypes(t *testing.T) { type record struct { ID any `vectordata:"key"` } - c := createCollection[any, record](t, vectordata.CollectionOptions{}) + c := createCollection[any, record](t, nil) valid := []any{ "string", 42, int64(43), uint64(44), true, [2]int{1, 2}, @@ -466,7 +715,7 @@ func TestCollection_KeyTypes(t *testing.T) { type pointerKey struct { ID *int `vectordata:"key"` } - pointers := createCollection[*int, pointerKey](t, vectordata.CollectionOptions{}) + pointers := createCollection[*int, pointerKey](t, nil) key := new(1) if err := pointers.Upsert(t.Context(), pointerKey{ID: key}); err != nil { t.Fatal(err) @@ -479,11 +728,11 @@ func TestCollection_KeyTypes(t *testing.T) { func checkAssignableDynamicKey[K comparable](t *testing.T, typ reflect.Type, key K, value any) { t.Helper() - c := createCollection[K, map[string]any](t, vectordata.CollectionOptions{ - Definition: &vectordata.Definition{Properties: []vectordata.Property{ + c := createCollection[K, map[string]any](t, + &vectordata.Definition{Properties: []vectordata.Property{ &vectordata.KeyProperty{Name: "ID", Type: typ}, }}, - }) + ) input := map[string]any{"ID": value} if err := c.Upsert(t.Context(), input); err != nil { t.Fatal(err) @@ -531,7 +780,7 @@ func TestCollection_NativeVectorRepresentations(t *testing.T) { Embedding embedding.Of[float32] `vectordata:"vector,dimensions=2"` Pointer *embedding.Of[float32] `vectordata:"vector,dimensions=2"` } - c := createCollection[int, record](t, vectordata.CollectionOptions{}) + c := createCollection[int, record](t, nil) input := record{ ID: 1, Slice: []float32{1, 0}, Vector: embedding.VectorOf[float32]{1, 0}, Embedding: embedding.Of[float32]{Vector: embedding.VectorOf[float32]{1, 0}}, @@ -558,7 +807,7 @@ func TestCollection_EmbeddedFields(t *testing.T) { fixed [2]int opaque any } - c := createCollection[uuid.UUID, record](t, vectordata.CollectionOptions{}) + c := createCollection[uuid.UUID, record](t, nil) input := record{when: time.Now(), fixed: [2]int{1, 2}, opaque: "immutable"} if err := c.Upsert(t.Context(), input); err != nil { t.Fatal(err) @@ -575,7 +824,7 @@ func TestCollection_EmbeddedFields(t *testing.T) { ID string `vectordata:"key"` *VectorFields } - native := createCollection[string, nullable](t, vectordata.CollectionOptions{}) + native := createCollection[string, nullable](t, nil) if err := native.Upsert(t.Context(), nullable{ID: "nil"}, nullable{ID: "value", VectorFields: &VectorFields{Vector: []float32{1, 0}}}); err != nil { t.Fatal(err) } @@ -591,14 +840,14 @@ func TestCollection_EmbeddedFields(t *testing.T) { ID uuid.UUID `vectordata:"key"` } type inaccessible struct{ *hiddenFields } - private := createCollection[uuid.UUID, inaccessible](t, vectordata.CollectionOptions{}) + private := createCollection[uuid.UUID, inaccessible](t, nil) if err := private.Upsert(t.Context(), inaccessible{}); err == nil { t.Fatal("accepted generation through an inaccessible embedded pointer") } } func TestCollection_ConcurrentCallsAndSnapshotIteration(t *testing.T) { - c := createCollection[string, document](t, vectordata.CollectionOptions{}) + c := createCollection[string, document](t, nil) ctx := t.Context() var work sync.WaitGroup for i := range 20 { @@ -642,8 +891,8 @@ func TestCollection_ConcurrentCallsAndSnapshotIteration(t *testing.T) { } func TestCollection_Cancellation(t *testing.T) { - c := createCollection[string, document](t, vectordata.CollectionOptions{}) - store := inmemory.New[string, document](inmemory.Config{}) + c := createCollection[string, document](t, nil) + store := inmemory.New(inmemory.Config{}) ctx, cancel := context.WithCancel(t.Context()) cancel() _, _, getErr := c.Get(ctx, "id", nil) diff --git a/vectordata/inmemory/embedding_test.go b/vectordata/inmemory/embedding_test.go index 49a0a1acf..7f979d08a 100644 --- a/vectordata/inmemory/embedding_test.go +++ b/vectordata/inmemory/embedding_test.go @@ -59,7 +59,7 @@ func TestEmbedding_BatchedPropertiesAndQueries(t *testing.T) { generated = append(generated, result.Embeddings...) return result, err }) - c := createCollection[string, record](t, vectordata.CollectionOptions{EmbeddingGenerator: embedding.Bind(generator)}) + c := createCollection[string, record](t, &vectordata.Definition{EmbeddingGenerator: embedding.Bind(generator)}) input := []record{ {ID: "one", Summary: "east", Details: "north", Native: []float32{1, 0}}, {ID: "two", Summary: "north", Details: "east", Native: []float32{0, 1}}, @@ -116,30 +116,41 @@ func TestEmbedding_BatchedPropertiesAndQueries(t *testing.T) { } func TestEmbedding_GeneratorPrecedence(t *testing.T) { - for level := range 4 { - t.Run([]string{"store", "collection", "definition", "property"}[level], func(t *testing.T) { + for _, test := range []struct { + name string + level int + explicit bool + }{ + {"store inferred", 0, false}, + {"store explicit", 0, true}, + {"definition", 1, true}, + {"property", 2, true}, + } { + t.Run(test.name, func(t *testing.T) { var calls []int - generators := make([]*embedding.BoundGenerator, 4) + generators := make([]*embedding.BoundGenerator, 3) for i := range generators { generators[i] = embedding.Bind(generatorFunc[string, *embedding.Of[float32]](func(ctx context.Context, values []string, options *embedding.Options) (*embedding.Generated[*embedding.Of[float32]], error) { calls = append(calls, i) return textGenerator().Generate(ctx, values, options) })) } - store := inmemory.New[string, embeddedDocument](inmemory.Config{EmbeddingGenerator: generators[0]}) - options := vectordata.CollectionOptions{} - if level >= 1 { - options.EmbeddingGenerator = generators[1] + store := inmemory.New(inmemory.Config{EmbeddingGenerator: generators[0]}) + var definition *vectordata.Definition + if test.explicit { + definition = &vectordata.Definition{} } - if level >= 2 { - options.Definition = &vectordata.Definition{EmbeddingGenerator: generators[2]} + var wantGenerator *embedding.BoundGenerator + if test.level >= 1 { + wantGenerator = generators[1] + definition.EmbeddingGenerator = wantGenerator } - if level >= 3 { - options.Definition.Properties = []vectordata.Property{ - &vectordata.VectorProperty{Name: "Text", Dimensions: 2, EmbeddingGenerator: generators[3]}, + if test.level >= 2 { + definition.Properties = []vectordata.Property{ + &vectordata.VectorProperty{Name: "Text", Dimensions: 2, EmbeddingGenerator: generators[2]}, } } - c, err := store.Collection("records", options) + c, err := vectordata.GetCollection[string, embeddedDocument](store.Store, "records", definition) if err != nil { t.Fatal(err) } @@ -150,13 +161,85 @@ func TestEmbedding_GeneratorPrecedence(t *testing.T) { t.Fatal(err) } collect(t, c.Search(t.Context(), "east", 1, nil)) - if !slices.Equal(calls, []int{level, level}) { + if !slices.Equal(calls, []int{test.level, test.level}) { t.Fatalf("generator selection = %v", calls) } + if definition != nil && definition.EmbeddingGenerator != wantGenerator { + t.Fatal("collection changed the definition's generator") + } + if test.level >= 2 { + property := definition.Properties[0].(*vectordata.VectorProperty) + if property.Type != nil || property.EmbeddingType != nil || property.StorageName != "" || + property.EmbeddingGenerator != generators[2] { + t.Fatal("collection changed the definition's vector property") + } + } }) } } +func TestEmbedding_SharedDefinitionDefaults(t *testing.T) { + property := &vectordata.VectorProperty{Name: "Text", Type: reflect.TypeFor[string](), Dimensions: 2} + definition := &vectordata.Definition{Properties: []vectordata.Property{ + &vectordata.KeyProperty{Name: "ID", Type: reflect.TypeFor[string]()}, + property, + }} + var calls []int + for i := range 2 { + generator := embedding.Bind(generatorFunc[string, *embedding.Of[float32]](func(ctx context.Context, values []string, options *embedding.Options) (*embedding.Generated[*embedding.Of[float32]], error) { + calls = append(calls, i) + return textGenerator().Generate(ctx, values, options) + })) + store := inmemory.New(inmemory.Config{EmbeddingGenerator: generator}) + collection, err := store.DynamicCollection("records", definition) + if err != nil { + t.Fatal(err) + } + if err := collection.EnsureExists(t.Context()); err != nil { + t.Fatal(err) + } + if err := collection.Upsert(t.Context(), map[string]any{"ID": "one", "Text": "east"}); err != nil { + t.Fatal(err) + } + results := collect(t, collection.Search(t.Context(), "east", 1, nil)) + if len(results) != 1 || results[0].Record["ID"] != "one" { + t.Fatalf("dynamic search results = %v", results) + } + if definition.EmbeddingGenerator != nil || property.EmbeddingGenerator != nil || + property.EmbeddingType != nil || property.StorageName != "" { + t.Fatal("store defaults modified the shared definition") + } + } + if !slices.Equal(calls, []int{0, 0, 1, 1}) { + t.Fatalf("generator defaults leaked between stores: %v", calls) + } + if _, err := inmemory.NewCollection[string, map[string]any]("standalone", definition); err == nil { + t.Fatal("standalone collection inherited a store's generator") + } +} + +func TestEmbedding_DefinitionOverridesInvalidStoreGenerator(t *testing.T) { + store := inmemory.New(inmemory.Config{EmbeddingGenerator: &embedding.BoundGenerator{}}) + if _, err := vectordata.GetCollection[string, embeddedDocument](store.Store, "invalid", nil); err == nil { + t.Fatal("invalid store generator was accepted") + } + definition := &vectordata.Definition{EmbeddingGenerator: embedding.Bind(textGenerator())} + collection, err := vectordata.GetCollection[string, embeddedDocument](store.Store, "records", definition) + if err != nil { + t.Fatalf("overridden store generator was selected: %v", err) + } + if err := collection.EnsureExists(t.Context()); err != nil { + t.Fatal(err) + } + if err := collection.Upsert(t.Context(), embeddedDocument{ID: "one", Text: "east"}); err != nil { + t.Fatal(err) + } + results := collect(t, collection.Search(t.Context(), "east", 1, nil)) + if len(results) != 1 || results[0].Record.ID != "one" { + t.Fatalf("definition generator search results = %v", results) + } +} + func TestEmbedding_InvalidOutputsAreAtomic(t *testing.T) { want := errors.New("generator failed") type output = embedding.Generated[*embedding.Of[float32]] @@ -186,7 +269,7 @@ func TestEmbedding_InvalidOutputsAreAtomic(t *testing.T) { } return textGenerator().Generate(ctx, values, options) }) - c := createCollection[string, embeddedDocument](t, vectordata.CollectionOptions{EmbeddingGenerator: embedding.Bind(generator)}) + c := createCollection[string, embeddedDocument](t, &vectordata.Definition{EmbeddingGenerator: embedding.Bind(generator)}) if err := c.Upsert(t.Context(), embeddedDocument{ID: "existing", Text: "east"}); err != nil { t.Fatal(err) } @@ -212,7 +295,7 @@ func TestEmbedding_SearchErrors(t *testing.T) { generator := generatorFunc[string, *embedding.Of[float32]](func(context.Context, []string, *embedding.Options) (*embedding.Generated[*embedding.Of[float32]], error) { return response, nil }) - c := createCollection[string, embeddedDocument](t, vectordata.CollectionOptions{EmbeddingGenerator: embedding.Bind(generator)}) + c := createCollection[string, embeddedDocument](t, &vectordata.Definition{EmbeddingGenerator: embedding.Bind(generator)}) _ = sequenceError(t, c.Search(t.Context(), "query", 1, nil)) }) } @@ -225,7 +308,7 @@ func TestEmbedding_SearchErrors(t *testing.T) { called = true return &embedding.Generated[[]float32]{Embeddings: [][]float32{{1, 0}}}, nil }) - c := createCollection[string, record](t, vectordata.CollectionOptions{EmbeddingGenerator: embedding.Bind(generator)}) + c := createCollection[string, record](t, &vectordata.Definition{EmbeddingGenerator: embedding.Bind(generator)}) _ = sequenceError(t, c.Search(t.Context(), (*string)(nil), 1, nil)) if called { t.Fatal("nil query reached the generator") @@ -240,7 +323,7 @@ func TestEmbedding_CancellationAndOriginalErrors(t *testing.T) { <-release return textGenerator().Generate(ctx, inputs, options) }) - c := createCollection[string, embeddedDocument](t, vectordata.CollectionOptions{EmbeddingGenerator: embedding.Bind(generator)}) + c := createCollection[string, embeddedDocument](t, &vectordata.Definition{EmbeddingGenerator: embedding.Bind(generator)}) ctx, cancel := context.WithCancel(t.Context()) defer cancel() done := make(chan error, 1) @@ -268,7 +351,7 @@ func TestEmbedding_CancellationAndOriginalErrors(t *testing.T) { cancel() return nil, want }) - c := createCollection[string, embeddedDocument](t, vectordata.CollectionOptions{EmbeddingGenerator: embedding.Bind(generator)}) + c := createCollection[string, embeddedDocument](t, &vectordata.Definition{EmbeddingGenerator: embedding.Bind(generator)}) var err error if search { err = sequenceError(t, c.Search(ctx, "query", 1, nil)) @@ -301,7 +384,7 @@ func TestEmbedding_CollectionDeletionDuringGeneration(t *testing.T) { } return textGenerator().Generate(ctx, values, options) }) - collection = createCollection[uuid.UUID, *record](t, vectordata.CollectionOptions{EmbeddingGenerator: embedding.Bind(generator)}) + collection = createCollection[uuid.UUID, *record](t, &vectordata.Definition{EmbeddingGenerator: embedding.Bind(generator)}) input := &record{Text: "east"} err := collection.Upsert(t.Context(), input) if recreate { @@ -343,7 +426,7 @@ func TestEmbedding_UUIDAssignmentFollowsGeneration(t *testing.T) { } return textGenerator().Generate(ctx, values, options) }) - c := createCollection[uuid.UUID, *record](t, vectordata.CollectionOptions{EmbeddingGenerator: embedding.Bind(generator)}) + c := createCollection[uuid.UUID, *record](t, &vectordata.Definition{EmbeddingGenerator: embedding.Bind(generator)}) if err := c.Upsert(t.Context(), input); !errors.Is(err, want) || input.ID != uuid.Nil { t.Fatalf("failed generation changed the key: %v, %v", input.ID, err) } @@ -357,9 +440,9 @@ func TestEmbedding_UUIDAssignmentFollowsGeneration(t *testing.T) { } func TestEmbedding_ProviderManagedGenerationUnsupported(t *testing.T) { - options := vectordata.CollectionOptions{Definition: &vectordata.Definition{Properties: []vectordata.Property{ + options := &vectordata.Definition{Properties: []vectordata.Property{ &vectordata.VectorProperty{Name: "Text", Dimensions: 2, EmbeddingType: reflect.TypeFor[*embedding.Of[float32]]()}, - }}} + }} if _, err := inmemory.NewCollection[string, embeddedDocument]("records", options); err == nil { t.Fatal("accepted generated field without a generator") } diff --git a/vectordata/inmemory/record.go b/vectordata/inmemory/record.go index b244716c4..b013db0a5 100644 --- a/vectordata/inmemory/record.go +++ b/vectordata/inmemory/record.go @@ -119,7 +119,7 @@ func (s recordSchema) validateRecord(record reflect.Value) error { if field.vector == nil || field.vector.Type != field.vector.EmbeddingType || !found || !value.IsValid() { continue } - vector, ok := nativeVector(value.Interface()) + vector, ok := providerservices.AsVector[float32](value.Interface()) if !ok { return fmt.Errorf("property %q is not a supported vector", field.name) } @@ -132,27 +132,35 @@ func (s recordSchema) validateRecord(record reflect.Value) error { return nil } -func recordKey[K comparable](field recordField, record reflect.Value) (K, error) { - var zero K +func recordKey(field recordField, record reflect.Value) (any, error) { value, found, err := field.read(record) if err != nil { - return zero, err + return nil, err } if !found { - return zero, fmt.Errorf("key property %q is missing", field.name) + return nil, fmt.Errorf("key property %q is missing", field.name) } if err := validateKey(value); err != nil { - return zero, fmt.Errorf("key property %q: %w", field.name, err) + return nil, fmt.Errorf("key property %q: %w", field.name, err) } - key := reflect.New(reflect.TypeFor[K]()).Elem() - if !value.Type().AssignableTo(key.Type()) { - // Assignability through the declared property type is not transitive. - declared := reflect.New(field.typ).Elem() - declared.Set(value) - value = declared + return normalizeKey(value, field.typ), nil +} + +func normalizeKey(value reflect.Value, typ reflect.Type) any { + if value.Type() == typ || !value.Type().AssignableTo(typ) { + return value.Interface() } + key := reflect.New(typ).Elem() key.Set(value) - return key.Interface().(K), nil + return key.Interface() +} + +func collectionRecord[R any](record any) (R, error) { + if typed, ok := record.(R); ok { + return typed, nil + } + var zero R + return zero, fmt.Errorf("stored record has type %T, want %v", record, reflect.TypeFor[R]()) } func validateKey(value reflect.Value) error { diff --git a/vectordata/inmemory/search.go b/vectordata/inmemory/search.go index 84fcbbc0c..19d886e02 100644 --- a/vectordata/inmemory/search.go +++ b/vectordata/inmemory/search.go @@ -35,19 +35,24 @@ func (c *Collection[K, R]) Search(ctx context.Context, searchValue any, top int, yield(vectordata.SearchResult[R]{}, failure(c.name, "Search", err)) return } - if !yield(vectordata.SearchResult[R]{Record: result.record.Record, Score: new(float64(result.score))}, nil) { + record, err := collectionRecord[R](result.record.Record) + if err != nil { + yield(vectordata.SearchResult[R]{}, failure(c.name, "Search", err)) + return + } + if !yield(vectordata.SearchResult[R]{Record: record, Score: new(float64(result.score))}, nil) { return } } } } -type scoredRecord[R any] struct { - record *storedRecord[R] +type scoredRecord struct { + record *storedRecord score float32 } -func (c *Collection[K, R]) search(ctx context.Context, searchValue any, top int, options *vectordata.SearchOptions) ([]scoredRecord[R], error) { +func (c *Collection[K, R]) search(ctx context.Context, searchValue any, top int, options *vectordata.SearchOptions) ([]scoredRecord, error) { if err := ctx.Err(); err != nil { return nil, err } @@ -75,7 +80,7 @@ func (c *Collection[K, R]) search(ctx context.Context, searchValue any, top int, if err != nil { return nil, err } - input, native := nativeVector(searchValue) + input, native := providerservices.AsVector[float32](searchValue) if !native { if property.SearchEmbeddingGenerator == nil { return nil, fmt.Errorf("property %q has no compatible generator for search input %T", property.Name, searchValue) @@ -84,7 +89,7 @@ func (c *Collection[K, R]) search(ctx context.Context, searchValue any, top int, if err != nil { return nil, err } - input, native = nativeVector(output) + input, native = providerservices.AsVector[float32](output) } if !native || input == nil { return nil, fmt.Errorf("search requires a non-nil float32 vector") @@ -97,7 +102,7 @@ func (c *Collection[K, R]) search(ctx context.Context, searchValue any, top int, return nil, err } field := c.schema.fields[property.Name] - results := make([]scoredRecord[R], 0, len(records)) + results := make([]scoredRecord, 0, len(records)) for _, record := range records { if err := ctx.Err(); err != nil { return nil, err @@ -118,7 +123,7 @@ func (c *Collection[K, R]) search(ctx context.Context, searchValue any, top int, if !found || !value.IsValid() { continue } - vector, _ = nativeVector(value.Interface()) + vector, _ = providerservices.AsVector[float32](value.Interface()) if vector == nil { continue } @@ -136,9 +141,9 @@ func (c *Collection[K, R]) search(ctx context.Context, searchValue any, top int, continue } } - results = append(results, scoredRecord[R]{record: record, score: score}) + results = append(results, scoredRecord{record: record, score: score}) } - slices.SortStableFunc(results, func(a, b scoredRecord[R]) int { + slices.SortStableFunc(results, func(a, b scoredRecord) int { if descending { return cmp.Compare(b.score, a.score) } @@ -213,8 +218,15 @@ func (c *Collection[K, R]) filtered(ctx context.Context, filter vectordata.Filte yield(zero, err) return } - if match && !yield(record.Record, nil) { - return + if match { + typed, err := collectionRecord[R](record.Record) + if err != nil { + yield(zero, err) + return + } + if !yield(typed, nil) { + return + } } } } diff --git a/vectordata/inmemory/search_test.go b/vectordata/inmemory/search_test.go index c75c6d56b..056f6df46 100644 --- a/vectordata/inmemory/search_test.go +++ b/vectordata/inmemory/search_test.go @@ -19,10 +19,10 @@ import ( "github.com/microsoft/agent-framework-go/vectordata/inmemory" ) -func vectorOptions(distance vectordata.DistanceFunction, index vectordata.IndexKind) vectordata.CollectionOptions { - return vectordata.CollectionOptions{Definition: &vectordata.Definition{Properties: []vectordata.Property{ +func vectorDefinition(distance vectordata.DistanceFunction, index vectordata.IndexKind) *vectordata.Definition { + return &vectordata.Definition{Properties: []vectordata.Property{ &vectordata.VectorProperty{Name: "Vector", Dimensions: 4, DistanceFunction: distance, IndexKind: index}, - }}} + }} } // Uses the exact, opposite, and orthogonal vectors from the upstream distance @@ -39,7 +39,7 @@ func TestSearch_DistanceFunctions(t *testing.T) { {vectordata.EuclideanDistance, []float64{0, math.Sqrt(3), 2}}, } { t.Run(string(test.distance), func(t *testing.T) { - c := createCollection[string, document](t, vectorOptions(test.distance, "")) + c := createCollection[string, document](t, vectorDefinition(test.distance, "")) base := []float32{1, 0, 0, 0} if err := c.Upsert(t.Context(), document{ID: "exact", Vector: base}, @@ -93,7 +93,7 @@ func TestSearch_UndefinedScores(t *testing.T) { vectordata.DotProductSimilarity, vectordata.EuclideanDistance, } { t.Run(string(distance), func(t *testing.T) { - c := createCollection[string, document](t, vectorOptions(distance, "")) + c := createCollection[string, document](t, vectorDefinition(distance, "")) base := []float32{1, 0, 0, 0} zero := []float32{0, 0, 0, 0} nan := []float32{float32(math.NaN()), 0, 0, 0} @@ -146,7 +146,7 @@ func TestSearch_UnsupportedMetricsAndIndexHints(t *testing.T) { vectordata.HammingDistance, vectordata.ManhattanDistance, "custom", } { t.Run(string(distance), func(t *testing.T) { - c := createCollection[string, document](t, vectorOptions(distance, "")) + c := createCollection[string, document](t, vectorDefinition(distance, "")) _ = sequenceError(t, c.Search(t.Context(), []float32{1, 0, 0, 0}, 1, nil)) }) } @@ -155,7 +155,7 @@ func TestSearch_UnsupportedMetricsAndIndexHints(t *testing.T) { vectordata.QuantizedFlat, vectordata.Dynamic, "custom", } { t.Run("index-"+string(index), func(t *testing.T) { - c := createCollection[string, document](t, vectorOptions("", index)) + c := createCollection[string, document](t, vectorDefinition("", index)) if err := c.Upsert(t.Context(), document{ID: "id", Vector: []float32{1, 0, 0, 0}}); err != nil { t.Fatal(err) } @@ -167,7 +167,7 @@ func TestSearch_UnsupportedMetricsAndIndexHints(t *testing.T) { } func TestSearch_Validation(t *testing.T) { - c := createCollection[string, document](t, vectordata.CollectionOptions{}) + c := createCollection[string, document](t, nil) vector := []float32{1, 0, 0, 0} for name, test := range map[string]struct { value any @@ -198,18 +198,18 @@ func TestSearch_Validation(t *testing.T) { First []float32 `vectordata:"vector,dimensions=4"` Second []float32 `vectordata:"vector,dimensions=4"` } - ambiguous := createCollection[string, multiple](t, vectordata.CollectionOptions{}) + ambiguous := createCollection[string, multiple](t, nil) _ = sequenceError(t, ambiguous.Search(t.Context(), vector, 1, nil)) type noVectors struct { ID string `vectordata:"key"` } - empty := createCollection[string, noVectors](t, vectordata.CollectionOptions{}) + empty := createCollection[string, noVectors](t, nil) _ = sequenceError(t, empty.Search(t.Context(), vector, 1, nil)) type unsupported struct { ID string `vectordata:"key"` Vector []float64 `vectordata:"vector,dimensions=4"` } - if _, err := inmemory.NewCollection[string, unsupported]("records", vectordata.CollectionOptions{}); err == nil { + if _, err := inmemory.NewCollection[string, unsupported]("records", nil); err == nil { t.Fatal("accepted unsupported float64 vectors") } } @@ -227,7 +227,7 @@ type filterRecord struct { func filterFixture(t *testing.T) (*inmemory.Collection[int, filterRecord], *inmemory.Collection[int, map[string]any]) { t.Helper() - typed := createCollection[int, filterRecord](t, vectordata.CollectionOptions{}) + typed := createCollection[int, filterRecord](t, nil) properties := []vectordata.Property{ &vectordata.KeyProperty{Name: "ID", Type: reflect.TypeFor[int]()}, &vectordata.VectorProperty{Name: "Vector", Type: reflect.TypeFor[[]float32](), Dimensions: 3}, @@ -237,7 +237,7 @@ func filterFixture(t *testing.T) (*inmemory.Collection[int, filterRecord], *inme properties = append(properties, &vectordata.DataProperty{Name: field.Name, Type: field.Type}) } } - dynamic := createCollection[int, map[string]any](t, vectordata.CollectionOptions{Definition: &vectordata.Definition{Properties: properties}}) + dynamic := createCollection[int, map[string]any](t, &vectordata.Definition{Properties: properties}) records := []filterRecord{ {ID: 0, Int: 8, String: new("foo"), Bool: true, Int2: 80, StringArray: []string{"x", "y"}, StringList: []string{"x", "y"}}, {ID: 1, Int: 9, String: new("bar"), Int2: 90, StringArray: []string{"a", "b"}, StringList: []string{"a", "b"}}, @@ -422,7 +422,7 @@ func TestGetFiltered_UnorderedPagination(t *testing.T) { {"ordered", 0, 1, []vectordata.Sort{{Property: "ID"}}, 1, 6}, } { t.Run(test.name, func(t *testing.T) { - c := createCollection[int, record](t, vectordata.CollectionOptions{}) + c := createCollection[int, record](t, nil) readings := 0 records := make([]record, 6) for i := range records { @@ -479,7 +479,7 @@ func TestGetFiltered_ValueTypes(t *testing.T) { "nested-null": {new((*int)(nil)), new(new(2))}, } { t.Run(name, func(t *testing.T) { - c := createCollection[int, record](t, vectordata.CollectionOptions{}) + c := createCollection[int, record](t, nil) if err := c.Upsert(t.Context(), record{ID: 2, Value: pair[1]}, record{ID: 1, Value: pair[0]}); err != nil { t.Fatal(err) } @@ -497,7 +497,7 @@ func TestGetFiltered_ValueTypes(t *testing.T) { } }) } - c := createCollection[int, record](t, vectordata.CollectionOptions{}) + c := createCollection[int, record](t, nil) if err := c.Upsert(t.Context(), record{ID: 1, Value: math.NaN()}, record{ID: 2, Value: float64(0)}); err != nil { t.Fatal(err) } @@ -519,7 +519,7 @@ func TestGetFiltered_CompareMethods(t *testing.T) { Value *comparisonValue `vectordata:"data"` Vector []float32 `vectordata:"vector,dimensions=1"` } - c := createCollection[int, record](t, vectordata.CollectionOptions{}) + c := createCollection[int, record](t, nil) if err := c.Upsert(t.Context(), record{ID: 3, Value: new(comparisonValue{rank: 3}), Vector: []float32{1}}, record{ID: 1, Value: new(comparisonValue{rank: 1}), Vector: []float32{1}}, @@ -606,12 +606,12 @@ func TestGetFiltered_InvalidCompareMethods(t *testing.T) { t.Run(reflect.TypeOf(value).Name(), func(t *testing.T) { for _, typ := range []reflect.Type{reflect.TypeOf(value), reflect.TypeFor[any]()} { t.Run(typ.String(), func(t *testing.T) { - c := createCollection[int, map[string]any](t, vectordata.CollectionOptions{Definition: &vectordata.Definition{ + c := createCollection[int, map[string]any](t, &vectordata.Definition{ Properties: []vectordata.Property{ &vectordata.KeyProperty{Name: "ID", Type: reflect.TypeFor[int]()}, &vectordata.DataProperty{Name: "Value", Type: typ}, }, - }}) + }) if err := c.Upsert(t.Context(), map[string]any{"ID": 1, "Value": value}); err != nil { t.Fatal(err) } @@ -670,7 +670,7 @@ func TestGetFiltered_EqualityMethods(t *testing.T) { {"invalid-result", invalidEqualResult{1}, invalidEqualResult{1}, true}, } { t.Run(test.name, func(t *testing.T) { - c := createCollection[int, record](t, vectordata.CollectionOptions{}) + c := createCollection[int, record](t, nil) if err := c.Upsert(t.Context(), record{ID: 1, Value: test.actual}); err != nil { t.Fatal(err) } @@ -695,7 +695,7 @@ func TestSearch_EqualityMethods(t *testing.T) { Values []equalityValue `vectordata:"data"` Vector []float32 `vectordata:"vector,dimensions=1"` } - c := createCollection[int, record](t, vectordata.CollectionOptions{}) + c := createCollection[int, record](t, nil) if err := c.Upsert(t.Context(), record{ID: 1, Value: new(equalityValue{"alpha"}), Values: []equalityValue{{"alpha"}}, Vector: []float32{1}}, record{ID: 2, Value: new(equalityValue{"beta"}), Values: []equalityValue{{"beta"}}, Vector: []float32{1}}, @@ -745,7 +745,7 @@ func TestGetFiltered_CyclicValues(t *testing.T) { } var cycle any cycle = &cycle - c := createCollection[int, record](t, vectordata.CollectionOptions{}) + c := createCollection[int, record](t, nil) if err := c.Upsert(t.Context(), record{ID: 1, Value: cycle}); err != nil { t.Fatal(err) } @@ -758,7 +758,7 @@ func TestGetFiltered_CyclicValues(t *testing.T) { ID int `vectordata:"key"` Value recursiveFirst } - recursiveCollection := createCollection[int, recursive](t, vectordata.CollectionOptions{}) + recursiveCollection := createCollection[int, recursive](t, nil) _ = sequenceError(t, recursiveCollection.GetFiltered(t.Context(), allRecords, 1, &vectordata.FilteredRecordRetrievalOptions{OrderBy: []vectordata.Sort{{Property: "Value"}}})) } @@ -769,7 +769,7 @@ func TestGetFiltered_NilCollectionsAndArrayMembership(t *testing.T) { Map map[string]int Array [2]string } - c := createCollection[int, record](t, vectordata.CollectionOptions{}) + c := createCollection[int, record](t, nil) if err := c.Upsert(t.Context(), record{ID: 1, Array: [2]string{"a", "b"}}, record{ID: 2, Slice: []string{}, Map: map[string]int{}}); err != nil { t.Fatal(err) } @@ -786,7 +786,7 @@ func TestGetFiltered_NilCollectionsAndArrayMembership(t *testing.T) { } func TestIterators_LazyAndEarlyTermination(t *testing.T) { - c, err := inmemory.NewCollection[string, document]("lazy", vectordata.CollectionOptions{}) + c, err := inmemory.NewCollection[string, document]("lazy", nil) if err != nil { t.Fatal(err) } @@ -846,7 +846,7 @@ func cancelAfterFirst[T any](t *testing.T, sequence func(context.Context) iter.S } func TestIterators_CancellationBetweenResults(t *testing.T) { - c := createCollection[string, document](t, vectordata.CollectionOptions{}) + c := createCollection[string, document](t, nil) if err := c.Upsert(t.Context(), document{ID: "one", Vector: []float32{1, 0, 0, 0}}, document{ID: "two", Vector: []float32{1, 0, 0, 0}}); err != nil { t.Fatal(err) } @@ -856,9 +856,9 @@ func TestIterators_CancellationBetweenResults(t *testing.T) { cancelAfterFirst(t, func(ctx context.Context) iter.Seq2[document, error] { return c.GetFiltered(ctx, allRecords, 2, nil) }) - store := inmemory.New[string, document](inmemory.Config{}) + store := inmemory.New(inmemory.Config{}) for _, name := range []string{"one", "two"} { - collection, err := store.Collection(name, vectordata.CollectionOptions{}) + collection, err := vectordata.GetCollection[string, document](store.Store, name, nil) if err != nil { t.Fatal(err) } diff --git a/vectordata/inmemory/snapshot.go b/vectordata/inmemory/snapshot.go index ea562fc0c..2e63641f7 100644 --- a/vectordata/inmemory/snapshot.go +++ b/vectordata/inmemory/snapshot.go @@ -19,17 +19,17 @@ import ( // generated vectors. The snapshot contains a name and an array of records, not // a schema or generator configuration. Records must support encoding/json. // It does not close the writer. Record order in the snapshot is unspecified. -func (s *Store[K, R]) Export(ctx context.Context, name string, writer io.Writer) error { +func (s *Store) Export(ctx context.Context, name string, writer io.Writer) error { if writer == nil { return failure(name, "Export", fmt.Errorf("writer is required")) } - records, err := (&Collection[K, R]{store: s, name: name}).snapshot(ctx) + records, err := s.provider.snapshot(ctx, name, nil) if err != nil { return failure(name, "Export", err) } snapshot := struct { - Name string `json:"name"` - Records []*storedRecord[R] `json:"records"` + Name string `json:"name"` + Records []*storedRecord `json:"records"` }{name, records} if err := json.NewEncoder(contextWriter{ctx, writer}).Encode(snapshot); err != nil { return failure(name, "Export", err) @@ -41,19 +41,23 @@ func (s *Store[K, R]) Export(ctx context.Context, name string, writer io.Writer) } // Import reads a JSON snapshot, creates its named collection if necessary, and -// merges records, replacing matching keys. Configuration is supplied separately. +// merges records, replacing matching keys. A nil definition uses the record's +// vectordata tags; otherwise the supplied definition configures the collection. // Generated vectors are restored without invoking a generator. Invalid snapshots // do not modify the store. The reader is not closed. // // Dynamic properties are decoded using their schema types. Other interface // values follow encoding/json's usual decoding rules. -func (s *Store[K, R]) Import(ctx context.Context, reader io.Reader, options vectordata.CollectionOptions) (vectordata.Collection[K, R], error) { +func Import[K comparable, R any](ctx context.Context, s *Store, reader io.Reader, definition *vectordata.Definition) (vectordata.Collection[K, R], error) { if err := ctx.Err(); err != nil { return nil, failure("", "Import", err) } if reader == nil { return nil, failure("", "Import", fmt.Errorf("reader is required")) } + if s == nil || s.provider == nil { + return nil, failure("", "Import", fmt.Errorf("store is required")) + } var snapshot struct { Name string `json:"name"` Records []struct { @@ -76,11 +80,11 @@ func (s *Store[K, R]) Import(ctx context.Context, reader io.Reader, options vect if snapshot.Records == nil { return nil, failure(snapshot.Name, "Import", fmt.Errorf("snapshot records must be an array")) } - collection, err := s.collection(snapshot.Name, options) + collection, err := newCollection[K, R](s.provider, snapshot.Name, reflect.TypeFor[K](), reflect.TypeFor[R](), definition) if err != nil { return nil, err } - updates := make(map[K]*storedRecord[R], len(snapshot.Records)) + updates := make(map[any]*storedRecord, len(snapshot.Records)) for i, source := range snapshot.Records { if err := ctx.Err(); err != nil { return nil, failure(snapshot.Name, "Import", err) @@ -93,7 +97,7 @@ func (s *Store[K, R]) Import(ctx context.Context, reader io.Reader, options vect if err := collection.schema.validateRecord(value); err != nil { return nil, failure(snapshot.Name, "Import", err) } - key, err := recordKey[K](collection.schema.key, value) + key, err := recordKey(collection.schema.key, value) if err != nil { return nil, failure(snapshot.Name, "Import", err) } @@ -114,7 +118,7 @@ func (s *Store[K, R]) Import(ctx context.Context, reader io.Reader, options vect if len(vectors) != len(source.Vectors) { return nil, failure(snapshot.Name, "Import", fmt.Errorf("record %d contains unknown generated vectors", i)) } - updates[key] = &storedRecord[R]{Record: record, Vectors: vectors} + updates[key] = &storedRecord{Record: record, Vectors: vectors} } if err := collection.commit(ctx, updates, true, "Import"); err != nil { return nil, err diff --git a/vectordata/inmemory/snapshot_test.go b/vectordata/inmemory/snapshot_test.go index fddba871d..39f212abd 100644 --- a/vectordata/inmemory/snapshot_test.go +++ b/vectordata/inmemory/snapshot_test.go @@ -30,8 +30,8 @@ func (s *snapshotStream) Close() error { func TestSnapshot_RoundTripAndMerge(t *testing.T) { ctx := t.Context() - source := inmemory.New[string, document](inmemory.Config{}) - c, err := source.Collection("records", vectordata.CollectionOptions{}) + source := inmemory.New(inmemory.Config{}) + c, err := vectordata.GetCollection[string, document](source.Store, "records", nil) if err != nil { t.Fatal(err) } @@ -49,8 +49,8 @@ func TestSnapshot_RoundTripAndMerge(t *testing.T) { if err := source.Export(ctx, "records", &stream); err != nil { t.Fatal(err) } - target := inmemory.New[string, document](inmemory.Config{}) - existing, err := target.Collection("records", vectordata.CollectionOptions{}) + target := inmemory.New(inmemory.Config{}) + existing, err := vectordata.GetCollection[string, document](target.Store, "records", nil) if err != nil { t.Fatal(err) } @@ -60,7 +60,7 @@ func TestSnapshot_RoundTripAndMerge(t *testing.T) { if err := existing.Upsert(ctx, document{ID: "kept"}, document{ID: "existing", Text: "old"}); err != nil { t.Fatal(err) } - imported, err := target.Import(ctx, &stream, vectordata.CollectionOptions{}) + imported, err := inmemory.Import[string, document](ctx, target, &stream, nil) if err != nil { t.Fatal(err) } @@ -76,8 +76,8 @@ func TestSnapshot_RoundTripAndMerge(t *testing.T) { if records := collect(t, imported.GetFiltered(ctx, allRecords, 10, nil)); len(records) != 3 { t.Fatalf("merge discarded records: %v", records) } - fresh := inmemory.New[string, document](inmemory.Config{}) - empty, err := fresh.Import(ctx, strings.NewReader(`{"name":"empty","records":[]}`), vectordata.CollectionOptions{}) + fresh := inmemory.New(inmemory.Config{}) + empty, err := inmemory.Import[string, document](ctx, fresh, strings.NewReader(`{"name":"empty","records":[]}`), nil) if err != nil { t.Fatal(err) } @@ -100,8 +100,8 @@ func TestSnapshot_GeneratedVectorsWithoutRegeneration(t *testing.T) { return textGenerator().Generate(ctx, values, options) }) config := inmemory.Config{EmbeddingGenerator: embedding.Bind(generator)} - source := inmemory.New[string, embeddedDocument](config) - c, err := source.Collection("generated", vectordata.CollectionOptions{}) + source := inmemory.New(config) + c, err := vectordata.GetCollection[string, embeddedDocument](source.Store, "generated", nil) if err != nil { t.Fatal(err) } @@ -115,8 +115,8 @@ func TestSnapshot_GeneratedVectorsWithoutRegeneration(t *testing.T) { if err := source.Export(t.Context(), "generated", &buffer); err != nil { t.Fatal(err) } - target := inmemory.New[string, embeddedDocument](config) - imported, err := target.Import(t.Context(), &buffer, vectordata.CollectionOptions{}) + target := inmemory.New(config) + imported, err := inmemory.Import[string, embeddedDocument](t.Context(), target, &buffer, nil) if err != nil { t.Fatal(err) } @@ -125,7 +125,7 @@ func TestSnapshot_GeneratedVectorsWithoutRegeneration(t *testing.T) { t.Fatalf("restored generated vectors = %+v, generator calls = %d", got, calls) } for _, vectors := range []string{`{}`, `{"Text":null}`, `{"Text":[1]}`, `{"Text":[1,0],"Other":[1,0]}`} { - _, err := target.Import(t.Context(), strings.NewReader(`{"name":"generated","records":[{"record":{"ID":"bad","Text":"east"},"vectors":`+vectors+`}]}`), vectordata.CollectionOptions{}) + _, err := inmemory.Import[string, embeddedDocument](t.Context(), target, strings.NewReader(`{"name":"generated","records":[{"record":{"ID":"bad","Text":"east"},"vectors":`+vectors+`}]}`), nil) if err == nil { t.Fatalf("accepted invalid sidecars %s", vectors) } @@ -138,9 +138,9 @@ func TestSnapshot_GeneratedVectorsWithoutRegeneration(t *testing.T) { func TestSnapshot_DynamicPropertyTypes(t *testing.T) { type dynamic map[string]any type pointer *dynamic - options := dynamicOptions(reflect.TypeFor[int]()) - source := inmemory.New[int, pointer](inmemory.Config{}) - c, err := source.Collection("dynamic", options) + options := dynamicDefinition(reflect.TypeFor[int]()) + source := inmemory.New(inmemory.Config{}) + c, err := vectordata.GetCollection[int, pointer](source.Store, "dynamic", options) if err != nil { t.Fatal(err) } @@ -156,8 +156,8 @@ func TestSnapshot_DynamicPropertyTypes(t *testing.T) { if err := source.Export(t.Context(), "dynamic", &snapshot); err != nil { t.Fatal(err) } - target := inmemory.New[int, pointer](inmemory.Config{}) - imported, err := target.Import(t.Context(), &snapshot, options) + target := inmemory.New(inmemory.Config{}) + imported, err := inmemory.Import[int, pointer](t.Context(), target, &snapshot, options) if err != nil { t.Fatal(err) } @@ -175,15 +175,63 @@ func TestSnapshot_DynamicPropertyTypes(t *testing.T) { t.Fatalf("typed nil vector was not preserved: %v, %v, %v", null, found, err) } for _, record := range []string{`{"ID":"wrong"}`, `{"ID":null}`, `{"Text":"missing key"}`, `{"ID":1,"Text":42}`} { - if _, err := target.Import(t.Context(), strings.NewReader(`{"name":"dynamic","records":[{"record":`+record+`}]}`), options); err == nil { + if _, err := inmemory.Import[int, pointer](t.Context(), target, strings.NewReader(`{"name":"dynamic","records":[{"record":`+record+`}]}`), options); err == nil { t.Fatalf("accepted invalid dynamic snapshot %s", record) } } } +func TestSnapshot_AssignableSharedKeys(t *testing.T) { + type key [2]int + for _, name := range []string{"typed import", "dynamic import"} { + t.Run(name, func(t *testing.T) { + store := inmemory.New(inmemory.Config{}) + definition := &vectordata.Definition{Properties: []vectordata.Property{ + &vectordata.KeyProperty{Name: "ID", Type: reflect.TypeFor[[2]int]()}, + }} + typed, err := vectordata.GetCollection[key, map[string]any](store.Store, "records", definition) + if err != nil { + t.Fatal(err) + } + dynamic, err := store.DynamicCollection("records", definition) + if err != nil { + t.Fatal(err) + } + if err := typed.EnsureExists(t.Context()); err != nil { + t.Fatal(err) + } + snapshot := strings.NewReader(`{"name":"records","records":[{"record":{"ID":[1,2],"Text":"imported"}}]}`) + if name == "typed import" { + if err := dynamic.Upsert(t.Context(), map[string]any{"ID": [2]int{1, 2}, "Text": "old"}); err != nil { + t.Fatal(err) + } + _, err = inmemory.Import[key, map[string]any](t.Context(), store, snapshot, definition) + } else { + if err := typed.Upsert(t.Context(), map[string]any{"ID": key{1, 2}, "Text": "old"}); err != nil { + t.Fatal(err) + } + _, err = inmemory.Import[any, map[string]any](t.Context(), store, snapshot, definition) + } + if err != nil { + t.Fatal(err) + } + want := map[string]any{"ID": [2]int{1, 2}, "Text": "imported"} + if got, found, err := typed.Get(t.Context(), key{1, 2}, nil); err != nil || !found || !reflect.DeepEqual(got, want) { + t.Fatalf("typed lookup after import = (%v, %v, %v)", got, found, err) + } + if got, found, err := dynamic.Get(t.Context(), [2]int{1, 2}, nil); err != nil || !found || !reflect.DeepEqual(got, want) { + t.Fatalf("dynamic lookup after import = (%v, %v, %v)", got, found, err) + } + if records := collect(t, dynamic.GetFiltered(t.Context(), allRecords, 10, nil)); len(records) != 1 { + t.Fatalf("import duplicated a logical key: %v", records) + } + }) + } +} + func TestSnapshot_InvalidInputsAreAtomic(t *testing.T) { - store := inmemory.New[string, document](inmemory.Config{}) - c, err := store.Collection("records", vectordata.CollectionOptions{}) + store := inmemory.New(inmemory.Config{}) + c, err := vectordata.GetCollection[string, document](store.Store, "records", nil) if err != nil { t.Fatal(err) } @@ -203,7 +251,7 @@ func TestSnapshot_InvalidInputsAreAtomic(t *testing.T) { `{"name":"records","records":[]} {}`, `{"name":"records","records":[],"unknown":true}`, `{"name":"records","records":[{"record":{"ID":"new"}},{"record":{"ID":"kept","Vector":[1]}}]}`, } { - if _, err := store.Import(t.Context(), strings.NewReader(snapshot), vectordata.CollectionOptions{}); err == nil { + if _, err := inmemory.Import[string, document](t.Context(), store, strings.NewReader(snapshot), nil); err == nil { t.Fatalf("accepted malformed snapshot %q", snapshot) } records := collect(t, c.GetFiltered(t.Context(), allRecords, 10, nil)) @@ -224,8 +272,8 @@ func (r *marshalingRecord) MarshalJSON() ([]byte, error) { } func TestSnapshot_ExportUsesStoredRecordsAndReportsEncodingErrors(t *testing.T) { - store := inmemory.New[string, *marshalingRecord](inmemory.Config{}) - c, err := store.Collection("records", vectordata.CollectionOptions{}) + store := inmemory.New(inmemory.Config{}) + c, err := vectordata.GetCollection[string, *marshalingRecord](store.Store, "records", nil) if err != nil { t.Fatal(err) } @@ -247,8 +295,8 @@ func TestSnapshot_ExportUsesStoredRecordsAndReportsEncodingErrors(t *testing.T) ID string `vectordata:"key"` Self *cyclicRecord } - cyclicStore := inmemory.New[string, *cyclicRecord](inmemory.Config{}) - cyclic, err := cyclicStore.Collection("cycles", vectordata.CollectionOptions{}) + cyclicStore := inmemory.New(inmemory.Config{}) + cyclic, err := vectordata.GetCollection[string, *cyclicRecord](cyclicStore.Store, "cycles", nil) if err != nil { t.Fatal(err) } @@ -276,8 +324,8 @@ type readerFunc func([]byte) (int, error) func (f readerFunc) Read(buffer []byte) (int, error) { return f(buffer) } func TestSnapshot_IOAndCancellation(t *testing.T) { - store := inmemory.New[string, document](inmemory.Config{}) - c, err := store.Collection("records", vectordata.CollectionOptions{}) + store := inmemory.New(inmemory.Config{}) + c, err := vectordata.GetCollection[string, document](store.Store, "records", nil) if err != nil { t.Fatal(err) } @@ -290,14 +338,14 @@ func TestSnapshot_IOAndCancellation(t *testing.T) { if err := store.Export(t.Context(), "records", nil); err == nil { t.Fatal("accepted nil writer") } - if _, err := store.Import(t.Context(), nil, vectordata.CollectionOptions{}); err == nil { + if _, err := inmemory.Import[string, document](t.Context(), store, nil, nil); err == nil { t.Fatal("accepted nil reader") } want := errors.New("IO failed") if err := store.Export(t.Context(), "records", failingIO{want}); !errors.Is(err, want) { t.Fatalf("export IO error = %v", err) } - if _, err := store.Import(t.Context(), failingIO{want}, vectordata.CollectionOptions{}); !errors.Is(err, want) { + if _, err := inmemory.Import[string, document](t.Context(), store, failingIO{want}, nil); !errors.Is(err, want) { t.Fatalf("import IO error = %v", err) } ctx, cancel := context.WithCancel(t.Context()) @@ -305,7 +353,7 @@ func TestSnapshot_IOAndCancellation(t *testing.T) { if err := store.Export(ctx, "records", io.Discard); !errors.Is(err, context.Canceled) { t.Fatalf("export cancellation = %v", err) } - if _, err := store.Import(ctx, failingIO{want}, vectordata.CollectionOptions{}); !errors.Is(err, context.Canceled) { + if _, err := inmemory.Import[string, document](ctx, store, failingIO{want}, nil); !errors.Is(err, context.Canceled) { t.Fatalf("import cancellation = %v", err) } ctx, cancel = context.WithCancel(t.Context()) @@ -317,7 +365,7 @@ func TestSnapshot_IOAndCancellation(t *testing.T) { cancel() return data.Read(buffer) }) - if _, err := store.Import(ctx, reader, vectordata.CollectionOptions{}); !errors.Is(err, context.Canceled) { + if _, err := inmemory.Import[string, document](ctx, store, reader, nil); !errors.Is(err, context.Canceled) { t.Fatalf("cancellation while reading = %v", err) } if reads != 1 { diff --git a/vectordata/inmemory/store.go b/vectordata/inmemory/store.go index 436ff3bb9..72ac38694 100644 --- a/vectordata/inmemory/store.go +++ b/vectordata/inmemory/store.go @@ -5,9 +5,12 @@ // unless explicitly exported. // // Collections support structs, pointers to structs, and map[string]any records. -// Keys must be non-nil comparable values without NaN. Records are retained and -// returned without cloning: pointer and map records, and reference-valued fields -// in struct records, share their underlying data with the caller. Plain structs +// Keys must be non-nil comparable values without NaN. Assignable keys are +// normalized to the schema's declared key type for storage and lookup, without +// changing records. Interface key properties retain concrete key types. +// Records are retained and returned without cloning: pointer and map records, +// and reference-valued fields in struct records, share their underlying data +// with the caller. Plain structs // are copied by value. Collection dictionaries support concurrent operations. // Shared record data is not synchronized: callers must prevent concurrent // mutation, including upserts that generate keys on the same record. Custom JSON @@ -36,6 +39,7 @@ import ( "context" "fmt" "iter" + "reflect" "sync" "github.com/microsoft/agent-framework-go/embedding" @@ -45,33 +49,43 @@ import ( // Config configures an in-memory store. type Config struct { // EmbeddingGenerator supplies the default generator for collection handles. - // Collection, definition, and property configuration may override it. + // Definition and property configuration may override it. EmbeddingGenerator *embedding.BoundGenerator } -// Store holds independently named collections with the same key and record -// types. The zero value is ready to use. A Store must not be copied after use. -// Its methods and collection handles support concurrent use. -type Store[K comparable, R any] struct { +// Store holds independently named collections with different key and record +// types. Create it with [New]. Its embedded [vectordata.Store] provides collection +// factories and lifecycle operations; pass that field to provider-independent +// APIs. A Store must not be copied after use. Its methods support concurrent use. +type Store struct { + *vectordata.Store + provider *store +} + +type store struct { mu sync.RWMutex - collections map[string]map[K]*storedRecord[R] + collections map[string]*collectionData generator *embedding.BoundGenerator } +type collectionData struct { + recordType reflect.Type + records map[any]*storedRecord +} + // New creates a store. The generator is borrowed, not closed by the store. -func New[K comparable, R any](config Config) *Store[K, R] { - return &Store[K, R]{generator: config.EmbeddingGenerator} +func New(config Config) *Store { + provider := &store{generator: config.EmbeddingGenerator} + return &Store{Store: vectordata.NewStore(provider), provider: provider} } -// Collection returns a configured handle without creating the collection. -// Handles with the same name share stored records, but not schema configuration. -func (s *Store[K, R]) Collection(name string, options vectordata.CollectionOptions) (vectordata.Collection[K, R], error) { - return s.collection(name, options) +func (s *store) OpenCollection(name string, keyType, recordType reflect.Type, definition *vectordata.Definition) (vectordata.Collection[any, any], error) { + return newCollection[any, any](s, name, keyType, recordType, definition) } // ListCollectionNames lazily returns a snapshot of the existing collection names. // The order is unspecified. -func (s *Store[K, R]) ListCollectionNames(ctx context.Context) iter.Seq2[string, error] { +func (s *store) ListCollectionNames(ctx context.Context) iter.Seq2[string, error] { return func(yield func(string, error) bool) { if err := ctx.Err(); err != nil { yield("", failure("", "ListCollectionNames", err)) @@ -96,7 +110,7 @@ func (s *Store[K, R]) ListCollectionNames(ctx context.Context) iter.Seq2[string, } // CollectionExists reports whether the named collection exists. -func (s *Store[K, R]) CollectionExists(ctx context.Context, name string) (bool, error) { +func (s *store) CollectionExists(ctx context.Context, name string) (bool, error) { if err := ctx.Err(); err != nil { return false, failure(name, "CollectionExists", err) } @@ -108,7 +122,7 @@ func (s *Store[K, R]) CollectionExists(ctx context.Context, name string) (bool, // EnsureCollectionDeleted removes the collection and all of its records. // Deleting a missing collection succeeds. Existing handles can recreate it. -func (s *Store[K, R]) EnsureCollectionDeleted(ctx context.Context, name string) error { +func (s *store) EnsureCollectionDeleted(ctx context.Context, name string) error { if err := ctx.Err(); err != nil { return failure(name, "EnsureCollectionDeleted", err) } @@ -121,8 +135,28 @@ func (s *Store[K, R]) EnsureCollectionDeleted(ctx context.Context, name string) return nil } -type storedRecord[R any] struct { - Record R `json:"record"` +func (s *store) snapshot(ctx context.Context, name string, recordType reflect.Type) ([]*storedRecord, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + s.mu.RLock() + defer s.mu.RUnlock() + data, exists := s.collections[name] + if !exists { + return nil, missingCollection(name) + } + if recordType != nil && data.recordType != recordType { + return nil, incompatibleCollection(name, data.recordType, recordType) + } + records := make([]*storedRecord, 0, len(data.records)) + for _, record := range data.records { + records = append(records, record) + } + return records, nil +} + +type storedRecord struct { + Record any `json:"record"` Vectors map[string][]float32 `json:"vectors,omitempty"` } @@ -136,4 +170,8 @@ func missingCollection(name string) error { return fmt.Errorf("collection %q does not exist", name) } -var _ vectordata.Store[string, any] = (*Store[string, any])(nil) +func incompatibleCollection(name string, existing, requested reflect.Type) error { + return fmt.Errorf("collection %q has record type %v, not %v", name, existing, requested) +} + +var _ vectordata.StoreProvider = (*store)(nil) diff --git a/vectordata/inmemory/vector.go b/vectordata/inmemory/vector.go index 576cd3e6f..937706bd9 100644 --- a/vectordata/inmemory/vector.go +++ b/vectordata/inmemory/vector.go @@ -5,40 +5,10 @@ package inmemory import ( "fmt" "math" - "reflect" - "github.com/microsoft/agent-framework-go/embedding" "github.com/microsoft/agent-framework-go/vectordata" ) -func vectorTypeSupported(typ reflect.Type) bool { - switch typ { - case reflect.TypeFor[[]float32](), reflect.TypeFor[embedding.VectorOf[float32]](), - reflect.TypeFor[embedding.Of[float32]](), reflect.TypeFor[*embedding.Of[float32]](): - return true - default: - return false - } -} - -func nativeVector(value any) ([]float32, bool) { - switch v := value.(type) { - case []float32: - return v, true - case embedding.VectorOf[float32]: - return v, true - case embedding.Of[float32]: - return v.Vector, true - case *embedding.Of[float32]: - if v == nil { - return nil, true - } - return v.Vector, true - default: - return nil, false - } -} - func validateVector(vector []float32, dimensions int) error { if len(vector) != dimensions { return fmt.Errorf("expected vector dimensions %d, got %d", dimensions, len(vector)) diff --git a/vectordata/providerservices/builder.go b/vectordata/providerservices/builder.go index 2f9367e24..d0e6b23de 100644 --- a/vectordata/providerservices/builder.go +++ b/vectordata/providerservices/builder.go @@ -1,10 +1,11 @@ // Copyright (c) Microsoft. All rights reserved. // Package providerservices provides experimental model-building, record-access, -// embedding-batching, and local filtering and sorting support for vector data -// providers, not applications. Providers build a model when creating a collection -// handle and retain it for subsequent operations. Local filtering and sorting -// are optional; remote providers translate them to their backend queries. +// vector-conversion, embedding-batching, and local filtering and sorting support +// for vector data providers, not applications. Providers build a model when +// creating a collection handle and retain it for subsequent operations. Local +// filtering and sorting are optional; remote providers translate them to their +// backend queries. // // This package does not implement a database connector, record serialization, // or vector storage. Its API may change independently of the core vector @@ -32,7 +33,7 @@ type Config struct { RequiresAtLeastOneVector bool // ReservedKeyStorageName overrides key storage names when nonempty. ReservedKeyStorageName string - // SupportsKeyType optionally restricts key types beyond TKey compatibility. + // SupportsKeyType optionally restricts supported key property types. SupportsKeyType func(reflect.Type) bool // SupportsDataType identifies supported stored data types. SupportsDataType func(reflect.Type) bool @@ -43,11 +44,16 @@ type Config struct { // SupportsKeyAutoGeneration controls both support and the default when // IsAutoGenerated is unset. Nil means auto-generation is unsupported. SupportsKeyAutoGeneration func(reflect.Type) bool + // DefaultEmbeddingGenerator supplies the fallback when the definition has + // no generator. Property and definition generators take precedence. + // The generator is borrowed, not closed by the model. + DefaultEmbeddingGenerator *embedding.BoundGenerator } -// Build resolves a collection's configuration into a provider-owned model. -// TRecord must be a struct, a pointer to a struct, or map[string]any. Dynamic -// records require a definition with explicit property types. +// Build resolves runtime key and record types into a provider-owned model. +// Both types must be non-nil. recordType must be a struct, map[string]any, or a +// pointer to either. Dynamic records require a definition with explicit property +// types. A nil definition uses the record's vectordata tags. // // Untagged fields are ignored unless included in Definition; "-" suppresses tag // inference. Explicit properties must have the same role as their tags. Definition @@ -56,24 +62,28 @@ type Config struct { // // The result owns its property models and annotation maps. Annotation values // and embedding generators are shared. Inputs are read-only. Exactly one key is -// required and its type must be assignable to TKey. +// required and its type must be assignable to keyType. // Property accessors are configured once for the resolved record type. They // read and assign values without cloning records or applying provider policies. // // Native vector fields retain their record type as their embedding type. Other // fields must match a generator's concrete input type or implement its input // interface, and the generator's output must be supported by the provider. +// Generators are resolved from the property, then the definition, then +// config.DefaultEmbeddingGenerator. // Without a generator, a non-native field requires an explicit supported // embedding type for provider-managed generation. // // Build does not invoke generators or create database collections. Providers // remain responsible for validating index choices, distance functions, and // representation-specific storage requirements. -func Build[TKey, TRecord any](options vectordata.CollectionOptions, config Config) (*CollectionModel, error) { +func Build(keyType, recordType reflect.Type, definition *vectordata.Definition, config Config) (*CollectionModel, error) { + if keyType == nil || recordType == nil { + return nil, fmt.Errorf("vectordata: key and record types are required") + } if config.SupportsDataType == nil || config.SupportsVectorType == nil { return nil, fmt.Errorf("vectordata: provider must specify supported data and vector types") } - recordType := reflect.TypeFor[TRecord]() if recordType.Kind() == reflect.Pointer { recordType = recordType.Elem() } @@ -82,17 +92,17 @@ func Build[TKey, TRecord any](options vectordata.CollectionOptions, config Confi if recordType.Kind() != reflect.Struct && !dynamic { return nil, fmt.Errorf("vectordata: record type %v must be a struct or map[string]any", recordType) } - if dynamic && options.Definition == nil { + if dynamic && definition == nil { return nil, fmt.Errorf("vectordata: dynamic records require a definition") } result := &CollectionModel{} - generator := options.EmbeddingGenerator + generator := config.DefaultEmbeddingGenerator definitions := make(map[string]PropertyModel) - if options.Definition != nil { - if options.Definition.EmbeddingGenerator != nil { - generator = options.Definition.EmbeddingGenerator + if definition != nil { + if definition.EmbeddingGenerator != nil { + generator = definition.EmbeddingGenerator } - for _, property := range options.Definition.Properties { + for _, property := range definition.Properties { clone, err := newPropertyModel(property) if err != nil { return nil, err @@ -176,8 +186,8 @@ func Build[TKey, TRecord any](options vectordata.CollectionOptions, config Confi storageNames[info.storageName] = info.name if _, ok := property.(*KeyPropertyModel); ok { keys++ - if !info.typ.AssignableTo(reflect.TypeFor[TKey]()) { - return nil, fmt.Errorf("vectordata: key property %q has type %v, not assignable to %v", info.name, info.typ, reflect.TypeFor[TKey]()) + if !info.typ.AssignableTo(keyType) { + return nil, fmt.Errorf("vectordata: key property %q has type %v, not assignable to %v", info.name, info.typ, keyType) } } if p, ok := property.(*VectorPropertyModel); ok { diff --git a/vectordata/providerservices/builder_test.go b/vectordata/providerservices/builder_test.go index 55c215351..b9fe7674f 100644 --- a/vectordata/providerservices/builder_test.go +++ b/vectordata/providerservices/builder_test.go @@ -59,20 +59,41 @@ func modelConfig() providerservices.Config { } } -func buildModel[K, R any](options vectordata.CollectionOptions) (*providerservices.CollectionModel, error) { - return providerservices.Build[K, R](options, modelConfig()) +func buildModel[K, R any](definition *vectordata.Definition) (*providerservices.CollectionModel, error) { + return providerservices.Build(reflect.TypeFor[K](), reflect.TypeFor[R](), definition, modelConfig()) +} + +func TestBuild_Validation(t *testing.T) { + for _, types := range [][2]reflect.Type{ + {nil, reflect.TypeFor[taggedRecord]()}, + {reflect.TypeFor[string](), nil}, + {reflect.TypeFor[int](), reflect.TypeFor[taggedRecord]()}, + {reflect.TypeFor[string](), reflect.TypeFor[any]()}, + } { + if _, err := providerservices.Build(types[0], types[1], nil, modelConfig()); err == nil { + t.Fatalf("accepted invalid runtime types %v", types) + } + } + model, err := providerservices.Build(reflect.TypeFor[string](), reflect.TypeFor[*taggedRecord](), nil, modelConfig()) + if err != nil { + t.Fatal(err) + } + input := &taggedRecord{ID: "one"} + if value, found, err := model.Properties[0].Value(input); err != nil || !found || value != "one" { + t.Fatalf("runtime model accessor = (%v, %v, %v)", value, found, err) + } } func TestPropertyModel_StructAccess(t *testing.T) { for _, test := range []struct { name string - build func(vectordata.CollectionOptions) (*providerservices.CollectionModel, error) + build func(*vectordata.Definition) (*providerservices.CollectionModel, error) }{ {"value model", buildModel[string, taggedRecord]}, {"pointer model", buildModel[string, *taggedRecord]}, } { t.Run(test.name, func(t *testing.T) { - model, err := test.build(vectordata.CollectionOptions{}) + model, err := test.build(nil) if err != nil { t.Fatal(err) } @@ -119,27 +140,27 @@ func TestPropertyModel_StructAccess(t *testing.T) { func TestPropertyModel_DynamicAccess(t *testing.T) { type namedRecord map[string]any - options := vectordata.CollectionOptions{Definition: &vectordata.Definition{Properties: []vectordata.Property{ + options := &vectordata.Definition{Properties: []vectordata.Property{ &vectordata.KeyProperty{Name: "ID", Type: reflect.TypeFor[string](), StorageName: "id"}, &vectordata.DataProperty{Name: "Text", Type: reflect.TypeFor[string](), StorageName: "text"}, &vectordata.VectorProperty{Name: "Vector", Type: reflect.TypeFor[[]float32](), Dimensions: 3}, &vectordata.DataProperty{Name: "Optional", Type: reflect.TypeFor[*string]()}, &vectordata.DataProperty{Name: "Label", Type: reflect.TypeFor[fmt.Stringer]()}, - }}} + }} config := modelConfig() config.SupportsDataType = func(reflect.Type) bool { return true } for _, test := range []struct { - name string - build func(vectordata.CollectionOptions, providerservices.Config) (*providerservices.CollectionModel, error) - record func(map[string]any) any + name string + recordType reflect.Type + record func(map[string]any) any }{ - {"map", providerservices.Build[string, map[string]any], func(record map[string]any) any { return record }}, - {"map pointer", providerservices.Build[string, *map[string]any], func(record map[string]any) any { return &record }}, - {"named map", providerservices.Build[string, namedRecord], func(record map[string]any) any { return namedRecord(record) }}, - {"named map pointer", providerservices.Build[string, *namedRecord], func(record map[string]any) any { return new(namedRecord(record)) }}, + {"map", reflect.TypeFor[map[string]any](), func(record map[string]any) any { return record }}, + {"map pointer", reflect.TypeFor[*map[string]any](), func(record map[string]any) any { return &record }}, + {"named map", reflect.TypeFor[namedRecord](), func(record map[string]any) any { return namedRecord(record) }}, + {"named map pointer", reflect.TypeFor[*namedRecord](), func(record map[string]any) any { return new(namedRecord(record)) }}, } { t.Run(test.name, func(t *testing.T) { - model, err := test.build(options, config) + model, err := providerservices.Build(reflect.TypeFor[string](), test.recordType, options, config) if err != nil { t.Fatal(err) } @@ -219,7 +240,7 @@ func TestPropertyModel_DynamicAccess(t *testing.T) { func TestPropertyModel_InvalidAccess(t *testing.T) { record := taggedRecord{ID: "one", Text: "before"} - model, err := buildModel[string, taggedRecord](vectordata.CollectionOptions{}) + model, err := buildModel[string, taggedRecord](nil) if err != nil { t.Fatal(err) } @@ -250,11 +271,11 @@ func TestPropertyModel_InvalidAccess(t *testing.T) { t.Fatalf("%T allowed writing without Build", unconfigured) } } - dynamic, err := buildModel[string, map[string]any](vectordata.CollectionOptions{ - Definition: &vectordata.Definition{Properties: []vectordata.Property{ + dynamic, err := buildModel[string, map[string]any]( + &vectordata.Definition{Properties: []vectordata.Property{ &vectordata.KeyProperty{Name: "ID", Type: reflect.TypeFor[string]()}, }}, - }) + ) if err != nil { t.Fatal(err) } @@ -282,7 +303,7 @@ func TestPropertyModel_NilAssignment(t *testing.T) { } config := modelConfig() config.SupportsDataType = func(reflect.Type) bool { return true } - model, err := providerservices.Build[string, record](vectordata.CollectionOptions{}, config) + model, err := providerservices.Build(reflect.TypeFor[string](), reflect.TypeFor[record](), nil, config) if err != nil { t.Fatal(err) } @@ -316,7 +337,7 @@ func TestPropertyModel_EmbeddedAccess(t *testing.T) { type Middle struct{ *Fields } type Outer struct{ Middle } type record struct{ *Outer } - model, err := buildModel[string, record](vectordata.CollectionOptions{}) + model, err := buildModel[string, record](nil) if err != nil { t.Fatal(err) } @@ -351,7 +372,7 @@ func TestPropertyModel_EmbeddedAccess(t *testing.T) { ID string `vectordata:"key"` } type hiddenRecord struct{ *hiddenFields } - hidden, err := buildModel[string, hiddenRecord](vectordata.CollectionOptions{}) + hidden, err := buildModel[string, hiddenRecord](nil) if err != nil { t.Fatal(err) } @@ -373,7 +394,7 @@ func TestPropertyModel_EmbeddedAccess(t *testing.T) { func TestBuild_Tags(t *testing.T) { generator := embedding.Bind(&taggedGenerator{name: "default"}) - definition, err := buildModel[string, taggedRecord](vectordata.CollectionOptions{EmbeddingGenerator: generator}) + definition, err := buildModel[string, taggedRecord](&vectordata.Definition{EmbeddingGenerator: generator}) if err != nil { t.Fatal(err) } @@ -398,14 +419,13 @@ func TestBuild_Tags(t *testing.T) { vector.EmbeddingType != vector.Type || vector.SearchEmbeddingGenerator != generator { t.Fatalf("vector = %+v", vector) } - fromPointer, err := buildModel[string, *taggedRecord](vectordata.CollectionOptions{}) + fromPointer, err := buildModel[string, *taggedRecord](nil) if err != nil || len(fromPointer.Properties) != 3 { t.Fatalf("pointer definition = %+v, %v", fromPointer, err) } } func TestBuild_DefinitionOverridesTags(t *testing.T) { - defaultGenerator := embedding.Bind(&taggedGenerator{name: "default"}) definitionGenerator := embedding.Bind(&taggedGenerator{name: "definition"}) propertyGenerator := embedding.Bind(&taggedGenerator{name: "property"}) input := &vectordata.Definition{ @@ -420,9 +440,7 @@ func TestBuild_DefinitionOverridesTags(t *testing.T) { &vectordata.DataProperty{Name: "Other"}, }, } - definition, err := buildModel[string, taggedRecord](vectordata.CollectionOptions{ - Definition: input, EmbeddingGenerator: defaultGenerator, - }) + definition, err := buildModel[string, taggedRecord](input) if err != nil { t.Fatal(err) } @@ -455,42 +473,78 @@ func TestBuild_EmbeddingGeneratorPrecedence(t *testing.T) { for _, dynamic := range []bool{false, true} { for _, test := range []struct { name string + fallback *embedding.BoundGenerator definition *embedding.BoundGenerator property *embedding.BoundGenerator want *embedding.BoundGenerator }{ - {"collection default", nil, nil, defaultGenerator}, - {"definition default", definitionGenerator, nil, definitionGenerator}, - {"property override", definitionGenerator, propertyGenerator, propertyGenerator}, + {"no generator", nil, nil, nil, nil}, + {"store default", defaultGenerator, nil, nil, defaultGenerator}, + {"definition default", nil, definitionGenerator, nil, definitionGenerator}, + {"definition override", defaultGenerator, definitionGenerator, nil, definitionGenerator}, + {"property override", defaultGenerator, definitionGenerator, propertyGenerator, propertyGenerator}, + {"property overrides store", defaultGenerator, nil, propertyGenerator, propertyGenerator}, + {"definition overrides invalid store", &embedding.BoundGenerator{}, definitionGenerator, nil, definitionGenerator}, } { t.Run(fmt.Sprintf("%s/dynamic=%t", test.name, dynamic), func(t *testing.T) { + source := &vectordata.VectorProperty{Name: "Vector", Type: reflect.TypeFor[embedding.VectorOf[float32]](), Dimensions: 3, EmbeddingGenerator: test.property} input := &vectordata.Definition{ EmbeddingGenerator: test.definition, Properties: []vectordata.Property{ &vectordata.KeyProperty{Name: "ID", Type: reflect.TypeFor[string]()}, - &vectordata.VectorProperty{Name: "Vector", Type: reflect.TypeFor[embedding.VectorOf[float32]](), Dimensions: 3, EmbeddingGenerator: test.property}, + source, }, } - options := vectordata.CollectionOptions{Definition: input, EmbeddingGenerator: defaultGenerator} - var definition *providerservices.CollectionModel - var err error + config := modelConfig() + config.DefaultEmbeddingGenerator = test.fallback + keyType, recordType := reflect.TypeFor[string](), reflect.TypeFor[taggedRecord]() if dynamic { - definition, err = buildModel[any, map[string]any](options) - } else { - definition, err = buildModel[string, taggedRecord](options) + keyType, recordType = reflect.TypeFor[any](), reflect.TypeFor[map[string]any]() } + definition, err := providerservices.Build(keyType, recordType, input, config) if err != nil { t.Fatal(err) } vector, err := definition.VectorProperty("") - if err != nil || vector.EmbeddingGenerator != test.want { + if err != nil || vector.EmbeddingGenerator != test.want || vector.SearchEmbeddingGenerator != test.want { t.Fatalf("generator = %+v, %v; want %+v", vector, err, test.want) } + if input.EmbeddingGenerator != test.definition || source.EmbeddingGenerator != test.property || + source.EmbeddingType != nil || source.StorageName != "" || config.DefaultEmbeddingGenerator != test.fallback { + t.Fatal("default resolution changed caller configuration") + } }) } } } +func TestBuild_DefaultEmbeddingGeneratorWithoutDefinition(t *testing.T) { + type record struct { + ID string `vectordata:"key"` + Content string `vectordata:"vector,dimensions=3"` + } + generator := embedding.Bind(&taggedGenerator{}) + config := modelConfig() + config.DefaultEmbeddingGenerator = generator + for _, input := range []*vectordata.Definition{nil, {}} { + model, err := providerservices.Build(reflect.TypeFor[string](), reflect.TypeFor[record](), input, config) + if err != nil { + t.Fatal(err) + } + property, err := model.VectorProperty("") + if err != nil || property.EmbeddingGenerator != generator || property.SearchEmbeddingGenerator != generator || + property.EmbeddingType != generator.EmbeddingType() || !model.EmbeddingGenerationRequired { + t.Fatalf("default generator model = %+v, %v", model, err) + } + if input != nil && input.EmbeddingGenerator != nil { + t.Fatal("default resolution changed the definition") + } + } + if _, err := providerservices.Build(reflect.TypeFor[string](), reflect.TypeFor[map[string]any](), nil, config); err == nil { + t.Fatal("a default generator replaced the required dynamic definition") + } +} + func TestBuild_TypedEmbeddingProperties(t *testing.T) { type image struct{ Name string } type mixedRecord struct { @@ -520,19 +574,21 @@ func TestBuild_TypedEmbeddingProperties(t *testing.T) { for _, dynamic := range []bool{false, true} { t.Run(fmt.Sprintf("dynamic=%t", dynamic), func(t *testing.T) { calls = 0 - source := &vectordata.Definition{Properties: []vectordata.Property{ - &vectordata.KeyProperty{Name: "ID", Type: reflect.TypeFor[string]()}, - &vectordata.VectorProperty{Name: "Text", Type: reflect.TypeFor[string](), Dimensions: 3}, - &vectordata.VectorProperty{Name: "Image", Type: reflect.TypeFor[image](), Dimensions: 2, EmbeddingGenerator: imageGenerator}, - &vectordata.VectorProperty{Name: "Dense", Type: reflect.TypeFor[[]float32](), Dimensions: 3}, - }} - options := vectordata.CollectionOptions{Definition: source, EmbeddingGenerator: textGenerator} + source := &vectordata.Definition{ + EmbeddingGenerator: textGenerator, + Properties: []vectordata.Property{ + &vectordata.KeyProperty{Name: "ID", Type: reflect.TypeFor[string]()}, + &vectordata.VectorProperty{Name: "Text", Type: reflect.TypeFor[string](), Dimensions: 3}, + &vectordata.VectorProperty{Name: "Image", Type: reflect.TypeFor[image](), Dimensions: 2, EmbeddingGenerator: imageGenerator}, + &vectordata.VectorProperty{Name: "Dense", Type: reflect.TypeFor[[]float32](), Dimensions: 3}, + }, + } var definition *providerservices.CollectionModel var err error if dynamic { - definition, err = buildModel[string, map[string]any](options) + definition, err = buildModel[string, map[string]any](source) } else { - definition, err = buildModel[string, mixedRecord](options) + definition, err = buildModel[string, mixedRecord](source) } if err != nil { t.Fatal(err) @@ -598,13 +654,12 @@ func TestBuild_PregeneratedVectorsWithSearchGenerator(t *testing.T) { config.SupportsVectorType = func(candidate reflect.Type) bool { return candidate == typ || candidate == generator.EmbeddingType() } - definition, err := providerservices.Build[string, map[string]any](vectordata.CollectionOptions{ - EmbeddingGenerator: generator, - Definition: &vectordata.Definition{Properties: []vectordata.Property{ - &vectordata.KeyProperty{Name: "ID", Type: reflect.TypeFor[string]()}, - &vectordata.VectorProperty{Name: "Vector", Type: typ, Dimensions: 3}, - }}, - }, config) + definition, err := providerservices.Build(reflect.TypeFor[string](), reflect.TypeFor[map[string]any](), &vectordata.Definition{EmbeddingGenerator: generator, Properties: []vectordata.Property{ + &vectordata.KeyProperty{Name: "ID", Type: reflect.TypeFor[string]()}, + &vectordata.VectorProperty{Name: "Vector", Type: typ, Dimensions: 3}, + }}, + + config) if err != nil { t.Fatalf("pregenerated type rejected: %v", err) } @@ -636,13 +691,10 @@ func TestBuild_GeneratorInputCompatibility(t *testing.T) { {"non-vector array", reflect.TypeFor[[3]string](), stringsGenerator, true}, } { t.Run(test.name, func(t *testing.T) { - definition, err := buildModel[string, map[string]any](vectordata.CollectionOptions{ - EmbeddingGenerator: test.generator, - Definition: &vectordata.Definition{Properties: []vectordata.Property{ - &vectordata.KeyProperty{Name: "ID", Type: reflect.TypeFor[string]()}, - &vectordata.VectorProperty{Name: "Content", Type: test.fieldType, Dimensions: 3}, - }}, - }) + definition, err := buildModel[string, map[string]any](&vectordata.Definition{EmbeddingGenerator: test.generator, Properties: []vectordata.Property{ + &vectordata.KeyProperty{Name: "ID", Type: reflect.TypeFor[string]()}, + &vectordata.VectorProperty{Name: "Content", Type: test.fieldType, Dimensions: 3}, + }}) if test.wantError { if err == nil || !strings.Contains(err.Error(), "Content") || definition != nil { t.Fatalf("expected incompatible input error: %+v, %v", definition, err) @@ -682,7 +734,7 @@ func testGeneratorInputIdentity[I, F any](t *testing.T, value F, wantError bool) t.Run(fmt.Sprintf("dynamic=%t", dynamic), func(t *testing.T) { calls := 0 want := &embedding.Of[float32]{Vector: embedding.VectorOf[float32]{1, 2, 3}} - options := vectordata.CollectionOptions{EmbeddingGenerator: embedding.Bind(&bindingGenerator[I, *embedding.Of[float32]]{ + options := &vectordata.Definition{EmbeddingGenerator: embedding.Bind(&bindingGenerator[I, *embedding.Of[float32]]{ generate: func(_ context.Context, inputs []I, _ *embedding.Options) (*embedding.GeneratedOf[float32], error) { calls++ if len(inputs) != 1 || !reflect.DeepEqual(inputs[0], value) { @@ -691,13 +743,14 @@ func testGeneratorInputIdentity[I, F any](t *testing.T, value F, wantError bool) return &embedding.GeneratedOf[float32]{Embeddings: []*embedding.Of[float32]{want}}, nil }, })} + var definition *providerservices.CollectionModel var err error if dynamic { - options.Definition = &vectordata.Definition{Properties: []vectordata.Property{ + options.Properties = []vectordata.Property{ &vectordata.KeyProperty{Name: "ID", Type: reflect.TypeFor[string]()}, &vectordata.VectorProperty{Name: "Content", Type: reflect.TypeFor[F](), Dimensions: 3}, - }} + } definition, err = buildModel[string, map[string]any](options) } else { definition, err = buildModel[string, model](options) @@ -736,12 +789,13 @@ func TestBuild_RejectsIncompatibleEmbeddingType(t *testing.T) { for _, output := range []reflect.Type{ nil, reflect.TypeFor[*embedding.Of[float32]](), reflect.TypeFor[*embedding.Of[float64]](), } { - input := &vectordata.Definition{Properties: []vectordata.Property{ - &vectordata.VectorProperty{Name: "Content", Dimensions: 3, EmbeddingType: output}, - }} - result, err := buildModel[string, model](vectordata.CollectionOptions{ - Definition: input, EmbeddingGenerator: generator, - }) + input := &vectordata.Definition{ + EmbeddingGenerator: generator, + Properties: []vectordata.Property{ + &vectordata.VectorProperty{Name: "Content", Dimensions: 3, EmbeddingType: output}, + }, + } + result, err := buildModel[string, model](input) if output == reflect.TypeFor[*embedding.Of[float64]]() { if err == nil || !strings.Contains(err.Error(), "Content") || result != nil { t.Fatalf("incompatible output was accepted: %+v, %v", result, err) @@ -767,23 +821,28 @@ func TestBuild_InvalidGeneratorBindings(t *testing.T) { Content string `vectordata:"vector,dimensions=3"` } invalid := &embedding.BoundGenerator{} - for _, options := range []vectordata.CollectionOptions{ + config := modelConfig() + config.DefaultEmbeddingGenerator = invalid + for _, definition := range []*vectordata.Definition{ + nil, + {Properties: []vectordata.Property{ + &vectordata.VectorProperty{Name: "Content", Dimensions: 3, EmbeddingGenerator: embedding.Bind(&taggedGenerator{})}, + }}, + } { + if result, err := providerservices.Build(reflect.TypeFor[string](), reflect.TypeFor[model](), definition, config); err == nil || result != nil { + t.Fatalf("uninitialized default binding accepted: %+v, %v", result, err) + } + } + for _, definition := range []*vectordata.Definition{ {EmbeddingGenerator: invalid}, - {Definition: &vectordata.Definition{EmbeddingGenerator: invalid}}, - {Definition: &vectordata.Definition{Properties: []vectordata.Property{ + {Properties: []vectordata.Property{ &vectordata.VectorProperty{Name: "Content", Dimensions: 3, EmbeddingGenerator: invalid}, - }}}, + }}, } { - if definition, err := buildModel[string, model](options); err == nil || definition != nil { - t.Fatalf("uninitialized binding accepted: %+v, %v", definition, err) + if model, err := buildModel[string, model](definition); err == nil || model != nil { + t.Fatalf("uninitialized binding accepted: %+v, %v", model, err) } } - valid := embedding.Bind(&taggedGenerator{}) - if _, err := buildModel[string, model](vectordata.CollectionOptions{ - EmbeddingGenerator: invalid, Definition: &vectordata.Definition{EmbeddingGenerator: valid}, - }); err != nil { - t.Fatalf("an overridden default should not be selected: %v", err) - } } func TestBuild_NoGeneratorRequiresSupportedEmbeddingType(t *testing.T) { @@ -792,11 +851,11 @@ func TestBuild_NoGeneratorRequiresSupportedEmbeddingType(t *testing.T) { Content string `vectordata:"vector,dimensions=3"` } for _, output := range []reflect.Type{nil, reflect.TypeFor[*embedding.Of[float32]](), reflect.TypeFor[*embedding.Of[uint8]]()} { - definition, err := buildModel[string, model](vectordata.CollectionOptions{ - Definition: &vectordata.Definition{Properties: []vectordata.Property{ + definition, err := buildModel[string, model]( + &vectordata.Definition{Properties: []vectordata.Property{ &vectordata.VectorProperty{Name: "Content", Dimensions: 3, EmbeddingType: output}, }}, - }) + ) if output != reflect.TypeFor[*embedding.Of[float32]]() { if err == nil || definition != nil || !strings.Contains(err.Error(), "Content") { t.Fatalf("unsupported provider configuration accepted: %+v, %v", definition, err) @@ -826,9 +885,9 @@ func TestBuild_DynamicAndOwnership(t *testing.T) { Name: "Vector", Type: reflect.TypeFor[[]float32](), Dimensions: 3, ProviderAnnotations: map[string]any{"index_setting": "value"}, } - definition, err := buildModel[string, dynamicRecord](vectordata.CollectionOptions{ - Definition: &vectordata.Definition{Properties: []vectordata.Property{key, data, vector}}, - }) + definition, err := buildModel[string, dynamicRecord]( + &vectordata.Definition{Properties: []vectordata.Property{key, data, vector}}, + ) if err != nil { t.Fatal(err) } @@ -865,13 +924,13 @@ func TestBuild_MissingDynamicTypes(t *testing.T) { case *vectordata.VectorProperty: p.Type = nil } - if _, err := buildModel[any, map[string]any](vectordata.CollectionOptions{ - Definition: &vectordata.Definition{Properties: properties}, - }); err == nil || !strings.Contains(err.Error(), "requires a type") { + if _, err := buildModel[any, map[string]any]( + &vectordata.Definition{Properties: properties}, + ); err == nil || !strings.Contains(err.Error(), "requires a type") { t.Fatalf("missing type at %d = %v", index, err) } } - if _, err := buildModel[any, map[string]any](vectordata.CollectionOptions{}); err == nil { + if _, err := buildModel[any, map[string]any](nil); err == nil { t.Fatal("dynamic records without a definition were accepted") } } @@ -896,14 +955,14 @@ func TestBuild_InvalidDefinitions(t *testing.T) { {"negative dimensions", []vectordata.Property{key, &vectordata.VectorProperty{Name: "Vector", Type: reflect.TypeFor[[]float32](), Dimensions: -1}}}, } { t.Run(test.name, func(t *testing.T) { - if _, err := buildModel[string, map[string]any](vectordata.CollectionOptions{ - Definition: &vectordata.Definition{Properties: test.properties}, - }); err == nil { + if _, err := buildModel[string, map[string]any]( + &vectordata.Definition{Properties: test.properties}, + ); err == nil { t.Fatal("invalid definition was accepted") } }) } - if _, err := buildModel[int, taggedRecord](vectordata.CollectionOptions{}); err == nil { + if _, err := buildModel[int, taggedRecord](nil); err == nil { t.Fatal("mismatched collection key type was accepted") } for _, property := range []vectordata.Property{ @@ -912,20 +971,20 @@ func TestBuild_InvalidDefinitions(t *testing.T) { &vectordata.DataProperty{Name: "ID"}, &vectordata.VectorProperty{Name: "Text", Dimensions: 3}, } { - if _, err := buildModel[string, taggedRecord](vectordata.CollectionOptions{ - Definition: &vectordata.Definition{Properties: []vectordata.Property{property}}, - }); err == nil { + if _, err := buildModel[string, taggedRecord]( + &vectordata.Definition{Properties: []vectordata.Property{property}}, + ); err == nil { t.Fatalf("invalid typed definition was accepted: %+v", property) } } - if _, err := buildModel[string, string](vectordata.CollectionOptions{}); err == nil { + if _, err := buildModel[string, string](nil); err == nil { t.Fatal("non-record type was accepted") } - if _, err := buildModel[string, map[int]string](vectordata.CollectionOptions{}); err == nil { + if _, err := buildModel[string, map[int]string](nil); err == nil { t.Fatal("unsupported map type was accepted") } type recursive *recursive - if _, err := buildModel[string, recursive](vectordata.CollectionOptions{}); err == nil { + if _, err := buildModel[string, recursive](nil); err == nil { t.Fatal("recursive pointer type was accepted") } } @@ -1036,7 +1095,7 @@ func TestBuild_InvalidTags(t *testing.T) { } func definitionError[TRecord any]() error { - _, err := buildModel[string, TRecord](vectordata.CollectionOptions{}) + _, err := buildModel[string, TRecord](nil) return err } @@ -1048,7 +1107,7 @@ func TestBuild_TagSettings(t *testing.T) { } config := modelConfig() config.SupportsKeyAutoGeneration = func(typ reflect.Type) bool { return typ == reflect.TypeFor[string]() } - definition, err := providerservices.Build[string, model](vectordata.CollectionOptions{}, config) + definition, err := providerservices.Build(reflect.TypeFor[string](), reflect.TypeFor[model](), nil, config) if err != nil { t.Fatal(err) } @@ -1104,12 +1163,12 @@ func TestBuild_BuiltinVectorSettingNames(t *testing.T) { } else { property.DistanceFunction = vectordata.DistanceFunction(value) } - definition, err := buildModel[string, map[string]any](vectordata.CollectionOptions{ - Definition: &vectordata.Definition{Properties: []vectordata.Property{ + definition, err := buildModel[string, map[string]any]( + &vectordata.Definition{Properties: []vectordata.Property{ &vectordata.KeyProperty{Name: "ID", Type: reflect.TypeFor[string]()}, property, }}, - }) + ) if value != test.name { if err == nil || definition != nil || !strings.Contains(err.Error(), fmt.Sprintf("must use %q", test.name)) { t.Fatalf("noncanonical setting accepted: %+v, %v", definition, err) @@ -1137,7 +1196,7 @@ func TestBuild_UnexportedAndEmbeddedFields(t *testing.T) { embedded Text string `vectordata:"data"` } - definition, err := buildModel[string, withEmbedded](vectordata.CollectionOptions{}) + definition, err := buildModel[string, withEmbedded](nil) if err != nil || len(definition.Properties) != 2 { t.Fatalf("embedded definition = %+v, %v", definition, err) } @@ -1145,7 +1204,7 @@ func TestBuild_UnexportedAndEmbeddedFields(t *testing.T) { id string `vectordata:"key"` } _ = private{}.id - if _, err := buildModel[string, private](vectordata.CollectionOptions{}); err == nil { + if _, err := buildModel[string, private](nil); err == nil { t.Fatal("unexported tagged field was accepted") } } @@ -1190,9 +1249,10 @@ func TestBuild_ProviderCapabilities(t *testing.T) { Name: fmt.Sprintf("Vector%d", i), Type: reflect.TypeFor[[]float32](), Dimensions: 3, }) } - model, err := providerservices.Build[string, map[string]any](vectordata.CollectionOptions{ - Definition: &vectordata.Definition{Properties: properties}, - }, config) + model, err := providerservices.Build(reflect.TypeFor[string](), reflect.TypeFor[map[string]any](), + &vectordata.Definition{Properties: properties}, + config) + if test.wantError != "" { if err == nil || model != nil || !strings.Contains(err.Error(), test.wantError) { t.Fatalf("model = %+v, %v; want %q", model, err, test.wantError) @@ -1211,20 +1271,20 @@ func TestBuild_ReservedKeyStorageName(t *testing.T) { config.ReservedKeyStorageName = "ProviderID" for _, explicit := range []bool{false, true} { t.Run(fmt.Sprintf("explicit=%t", explicit), func(t *testing.T) { - options := vectordata.CollectionOptions{} + var definition *vectordata.Definition if explicit { - options.Definition = &vectordata.Definition{Properties: []vectordata.Property{ + definition = &vectordata.Definition{Properties: []vectordata.Property{ &vectordata.KeyProperty{Name: "ID", StorageName: "CallerID"}, }} } - model, err := providerservices.Build[string, taggedRecord](options, config) + model, err := providerservices.Build(reflect.TypeFor[string](), reflect.TypeFor[taggedRecord](), definition, config) if err != nil { t.Fatal(err) } if key := model.Properties[0].(*providerservices.KeyPropertyModel); key.StorageName != "ProviderID" { t.Fatalf("reserved storage name = %q", key.StorageName) } - if explicit && options.Definition.Properties[0].(*vectordata.KeyProperty).StorageName != "CallerID" { + if explicit && definition.Properties[0].(*vectordata.KeyProperty).StorageName != "CallerID" { t.Fatal("provider policy modified caller configuration") } }) @@ -1246,15 +1306,15 @@ func TestBuild_KeyAutoGeneration(t *testing.T) { input := &vectordata.KeyProperty{ Name: "ID", Type: reflect.TypeFor[string](), IsAutoGenerated: requested, } - options := vectordata.CollectionOptions{Definition: &vectordata.Definition{ + options := &vectordata.Definition{ Properties: []vectordata.Property{input}, - }} + } var model *providerservices.CollectionModel var err error if dynamic { - model, err = providerservices.Build[string, map[string]any](options, config) + model, err = providerservices.Build(reflect.TypeFor[string](), reflect.TypeFor[map[string]any](), options, config) } else { - model, err = providerservices.Build[string, taggedRecord](options, config) + model, err = providerservices.Build(reflect.TypeFor[string](), reflect.TypeFor[taggedRecord](), options, config) } if !supported && requested != nil && *requested { if err == nil || model != nil || !strings.Contains(err.Error(), "does not support auto-generation") { @@ -1287,7 +1347,7 @@ func TestBuild_KeyAutoGeneration(t *testing.T) { for _, supported := range []bool{false, true} { config := modelConfig() config.SupportsKeyAutoGeneration = func(reflect.Type) bool { return supported } - model, err := providerservices.Build[string, record](vectordata.CollectionOptions{}, config) + model, err := providerservices.Build(reflect.TypeFor[string](), reflect.TypeFor[record](), nil, config) if err != nil || model.Properties[0].(*providerservices.KeyPropertyModel).IsAutoGenerated != supported { t.Fatalf("tag default = %+v, %v; want %t", model, err, supported) } @@ -1305,12 +1365,13 @@ func TestBuild_NativeVectorSupportComesFromProvider(t *testing.T) { config.SupportsVectorType = func(candidate reflect.Type) bool { return supported && candidate == typ } - model, err := providerservices.Build[string, map[string]any](vectordata.CollectionOptions{ - Definition: &vectordata.Definition{Properties: []vectordata.Property{ + model, err := providerservices.Build(reflect.TypeFor[string](), reflect.TypeFor[map[string]any](), + &vectordata.Definition{Properties: []vectordata.Property{ &vectordata.KeyProperty{Name: "ID", Type: reflect.TypeFor[string]()}, &vectordata.VectorProperty{Name: "Vector", Type: typ, Dimensions: 3}, }}, - }, config) + config) + if !supported { if err == nil || model != nil || !strings.Contains(err.Error(), "unsupported type") { t.Fatalf("unsupported native type accepted: %+v, %v", model, err) @@ -1336,10 +1397,7 @@ func TestBuild_NativeEmbeddingType(t *testing.T) { } for _, output := range []reflect.Type{nil, reflect.TypeFor[[]float32](), reflect.TypeFor[*embedding.Of[float32]]()} { source := &vectordata.VectorProperty{Name: "Vector", Dimensions: 3, EmbeddingType: output} - model, err := buildModel[string, record](vectordata.CollectionOptions{ - Definition: &vectordata.Definition{Properties: []vectordata.Property{source}}, - EmbeddingGenerator: embedding.Bind(&taggedGenerator{}), - }) + model, err := buildModel[string, record](&vectordata.Definition{Properties: []vectordata.Property{source}, EmbeddingGenerator: embedding.Bind(&taggedGenerator{})}) if output == reflect.TypeFor[*embedding.Of[float32]]() { if err == nil || model != nil || !strings.Contains(err.Error(), "must use its record type") { t.Fatalf("native embedding type override accepted: %+v, %v", model, err) @@ -1384,7 +1442,7 @@ func testSearchEmbeddingGenerator[E any](t *testing.T, output E, supported bool) return &embedding.Generated[E]{Embeddings: []E{output}}, nil }, }) - model, err := buildModel[string, record](vectordata.CollectionOptions{EmbeddingGenerator: generator}) + model, err := buildModel[string, record](&vectordata.Definition{EmbeddingGenerator: generator}) if err != nil { t.Fatal(err) } @@ -1401,7 +1459,7 @@ func testSearchEmbeddingGenerator[E any](t *testing.T, output E, supported bool) ID string `vectordata:"key"` Vector string `vectordata:"vector,dimensions=3"` } - model, err = buildModel[string, generatedRecord](vectordata.CollectionOptions{EmbeddingGenerator: generator}) + model, err = buildModel[string, generatedRecord](&vectordata.Definition{EmbeddingGenerator: generator}) if err == nil || model != nil || !strings.Contains(err.Error(), "unsupported embedding type") || calls != 0 { t.Fatalf("unsupported generated output accepted: %+v, %v", model, err) } @@ -1425,7 +1483,7 @@ func ExampleBuild() { Text string `vectordata:"data,fulltext"` Embedding []float32 `vectordata:"vector,dimensions=3,distance=cosine_similarity"` } - model, err := providerservices.Build[string, document](vectordata.CollectionOptions{}, providerservices.Config{ + model, err := providerservices.Build(reflect.TypeFor[string](), reflect.TypeFor[document](), nil, providerservices.Config{ RequiresAtLeastOneVector: true, SupportsDataType: func(typ reflect.Type) bool { return typ == reflect.TypeFor[string]() }, SupportsVectorType: func(typ reflect.Type) bool { return typ == reflect.TypeFor[[]float32]() }, diff --git a/vectordata/providerservices/embedding.go b/vectordata/providerservices/embedding.go index c4c76b1ef..c48162b8c 100644 --- a/vectordata/providerservices/embedding.go +++ b/vectordata/providerservices/embedding.go @@ -5,10 +5,47 @@ package providerservices import ( "context" "fmt" + "reflect" "github.com/microsoft/agent-framework-go/embedding" ) +// SupportsVector reports whether typ is []T, [embedding.VectorOf] with T +// elements, or [embedding.Of] with T elements (a value or pointer). +// A nil type is unsupported. Providers opt in by choosing T. +func SupportsVector[T any](typ reflect.Type) bool { + switch typ { + case reflect.TypeFor[[]T](), reflect.TypeFor[embedding.VectorOf[T]](), + reflect.TypeFor[embedding.Of[T]](), reflect.TypeFor[*embedding.Of[T]](): + return true + default: + return false + } +} + +// AsVector unwraps the representations supported by [SupportsVector] for T. +// The returned slice is borrowed, not cloned or normalized. Supported typed nils +// return (nil, true); untyped nil and unsupported values return (nil, false). +// Element types must match T exactly; no conversion is performed. +// Providers remain responsible for nil handling and dimension validation. +func AsVector[T any](value any) ([]T, bool) { + switch vector := value.(type) { + case []T: + return vector, true + case embedding.VectorOf[T]: + return vector, true + case embedding.Of[T]: + return vector.Vector, true + case *embedding.Of[T]: + if vector == nil { + return nil, true + } + return vector.Vector, true + default: + return nil, false + } +} + // GenerateEmbeddings reads a vector property's inputs from records and invokes // its configured generator once, with nil generation options. The property must // have a record accessor configured by [Build]. Missing values become nil; diff --git a/vectordata/providerservices/embedding_test.go b/vectordata/providerservices/embedding_test.go index 16fc41bee..c9b8eeb92 100644 --- a/vectordata/providerservices/embedding_test.go +++ b/vectordata/providerservices/embedding_test.go @@ -15,12 +15,71 @@ import ( "github.com/microsoft/agent-framework-go/vectordata/providerservices" ) +func TestVectorHelpers(t *testing.T) { + testVectorHelpers(t, float32(1), float32(2), []float64{1, 2}) + testVectorHelpers(t, float64(1), float64(2), []float32{1, 2}) + testVectorHelpers(t, uint8(1), uint8(2), []int8{1, 2}) + testVectorHelpers(t, int8(1), int8(2), []uint8{1, 2}) + type namedFloat float32 + testVectorHelpers(t, namedFloat(1), namedFloat(2), []float32{1, 2}) + testVectorHelpers(t, "one", "two", []byte{1, 2}) +} + +func testVectorHelpers[T, U any](t *testing.T, first, second T, other []U) { + t.Helper() + type namedSlice []T + t.Run(reflect.TypeFor[T]().String(), func(t *testing.T) { + for _, test := range []struct { + name string + value any + supported bool + want []T + }{ + {"slice", []T{first, second}, true, []T{first, second}}, + {"vector", embedding.VectorOf[T]{first, second}, true, []T{first, second}}, + {"embedding", embedding.Of[T]{Vector: embedding.VectorOf[T]{first, second}}, true, []T{first, second}}, + {"embedding pointer", &embedding.Of[T]{Vector: embedding.VectorOf[T]{first, second}}, true, []T{first, second}}, + {"nil slice", []T(nil), true, nil}, + {"nil vector", embedding.VectorOf[T](nil), true, nil}, + {"empty embedding", embedding.Of[T]{}, true, nil}, + {"nil embedding pointer", (*embedding.Of[T])(nil), true, nil}, + {"pointer to empty embedding", &embedding.Of[T]{}, true, nil}, + {"empty slice", []T{}, true, []T{}}, + {"untyped nil", nil, false, nil}, + {"other slice", other, false, nil}, + {"other vector", embedding.VectorOf[U](other), false, nil}, + {"other embedding", embedding.Of[U]{Vector: other}, false, nil}, + {"other embedding pointer", &embedding.Of[U]{Vector: other}, false, nil}, + {"array", [2]T{first, second}, false, nil}, + {"named slice", namedSlice{first, second}, false, nil}, + {"scalar", first, false, nil}, + } { + t.Run(test.name, func(t *testing.T) { + if supported := providerservices.SupportsVector[T](reflect.TypeOf(test.value)); supported != test.supported { + t.Fatalf("type support = %v, want %v", supported, test.supported) + } + vector, ok := providerservices.AsVector[T](test.value) + if ok != test.supported || !reflect.DeepEqual(vector, test.want) { + t.Fatalf("unwrapped vector = (%#v, %v), want (%#v, %v)", vector, ok, test.want, test.supported) + } + if len(vector) != 0 { + vector[0] = second + again, _ := providerservices.AsVector[T](test.value) + if !reflect.DeepEqual(again[0], second) { + t.Fatal("unwrapping cloned the vector") + } + } + }) + } + }) +} + type batchEmbeddingRecord struct { ID string `vectordata:"key"` Text string `vectordata:"vector,dimensions=2"` } -func buildEmbeddingProperty[R any](t *testing.T, options vectordata.CollectionOptions) *providerservices.VectorPropertyModel { +func buildEmbeddingProperty[R any](t *testing.T, options *vectordata.Definition) *providerservices.VectorPropertyModel { t.Helper() model, err := buildModel[string, R](options) if err != nil { @@ -36,7 +95,7 @@ func buildEmbeddingProperty[R any](t *testing.T, options vectordata.CollectionOp func TestGenerateEmbeddings_RecordsAndMetadata(t *testing.T) { for _, test := range []struct { name string - build func(vectordata.CollectionOptions) (*providerservices.CollectionModel, error) + build func(*vectordata.Definition) (*providerservices.CollectionModel, error) records []any }{ {"struct", buildModel[string, batchEmbeddingRecord], []any{ @@ -69,13 +128,10 @@ func TestGenerateEmbeddings_RecordsAndMetadata(t *testing.T) { }, nil }, } - model, err := test.build(vectordata.CollectionOptions{ - EmbeddingGenerator: embedding.Bind(generator), - Definition: &vectordata.Definition{Properties: []vectordata.Property{ - &vectordata.KeyProperty{Name: "ID", Type: reflect.TypeFor[string]()}, - &vectordata.VectorProperty{Name: "Text", Type: reflect.TypeFor[string](), Dimensions: 2, StorageName: "text"}, - }}, - }) + model, err := test.build(&vectordata.Definition{EmbeddingGenerator: embedding.Bind(generator), Properties: []vectordata.Property{ + &vectordata.KeyProperty{Name: "ID", Type: reflect.TypeFor[string]()}, + &vectordata.VectorProperty{Name: "Text", Type: reflect.TypeFor[string](), Dimensions: 2, StorageName: "text"}, + }}) if err != nil { t.Fatal(err) } @@ -117,7 +173,7 @@ func TestGenerateEmbeddings_EmptyBatches(t *testing.T) { return &embedding.Generated[[]float32]{}, nil }, } - property := buildEmbeddingProperty[batchEmbeddingRecord](t, vectordata.CollectionOptions{EmbeddingGenerator: embedding.Bind(generator)}) + property := buildEmbeddingProperty[batchEmbeddingRecord](t, &vectordata.Definition{EmbeddingGenerator: embedding.Bind(generator)}) result, err := providerservices.GenerateEmbeddings(t.Context(), property, records) if err != nil || result == nil || len(result.Embeddings) != 0 || calls != 1 { t.Fatalf("empty generation = %+v, %v; calls=%d", result, err, calls) @@ -141,7 +197,7 @@ func TestGenerateEmbeddings_TypedNilAndOutputPolicies(t *testing.T) { return &embedding.GeneratedOf[float64]{Embeddings: []*embedding.Of[float64]{nil, output}}, nil }, } - property := buildEmbeddingProperty[record](t, vectordata.CollectionOptions{EmbeddingGenerator: embedding.Bind(generator)}) + property := buildEmbeddingProperty[record](t, &vectordata.Definition{EmbeddingGenerator: embedding.Bind(generator)}) result, err := providerservices.GenerateEmbeddings(t.Context(), property, []record{{ID: "nil"}, {ID: "present", Text: text}}) if err != nil || result == nil || len(result.Embeddings) != 2 || result.Embeddings[0] != (*embedding.Of[float64])(nil) || result.Embeddings[1] != output { @@ -162,7 +218,7 @@ func TestGenerateEmbeddings_InvalidCounts(t *testing.T) { return output, nil }, } - property := buildEmbeddingProperty[batchEmbeddingRecord](t, vectordata.CollectionOptions{EmbeddingGenerator: embedding.Bind(generator)}) + property := buildEmbeddingProperty[batchEmbeddingRecord](t, &vectordata.Definition{EmbeddingGenerator: embedding.Bind(generator)}) result, err := providerservices.GenerateEmbeddings(t.Context(), property, []batchEmbeddingRecord{{Text: "first"}, {Text: "second"}}) if err == nil || result != nil { t.Fatalf("invalid count = %+v, %v", result, err) @@ -174,7 +230,7 @@ func TestGenerateEmbeddings_InvalidCounts(t *testing.T) { func TestGenerateEmbeddings_InvalidRecords(t *testing.T) { for _, test := range []struct { name string - build func(vectordata.CollectionOptions) (*providerservices.CollectionModel, error) + build func(*vectordata.Definition) (*providerservices.CollectionModel, error) records []any }{ {"wrong record type", buildModel[string, batchEmbeddingRecord], []any{batchEmbeddingRecord{Text: "first"}, struct{}{}}}, @@ -191,13 +247,10 @@ func TestGenerateEmbeddings_InvalidRecords(t *testing.T) { return nil, nil }, } - model, err := test.build(vectordata.CollectionOptions{ - EmbeddingGenerator: embedding.Bind(generator), - Definition: &vectordata.Definition{Properties: []vectordata.Property{ - &vectordata.KeyProperty{Name: "ID", Type: reflect.TypeFor[string]()}, - &vectordata.VectorProperty{Name: "Text", Type: reflect.TypeFor[string](), Dimensions: 2}, - }}, - }) + model, err := test.build(&vectordata.Definition{EmbeddingGenerator: embedding.Bind(generator), Properties: []vectordata.Property{ + &vectordata.KeyProperty{Name: "ID", Type: reflect.TypeFor[string]()}, + &vectordata.VectorProperty{Name: "Text", Type: reflect.TypeFor[string](), Dimensions: 2}, + }}) if err != nil { t.Fatal(err) } @@ -221,12 +274,12 @@ func TestGenerateEmbeddings_InvalidConfiguration(t *testing.T) { } }) } - property := buildEmbeddingProperty[batchEmbeddingRecord](t, vectordata.CollectionOptions{ - Definition: &vectordata.Definition{Properties: []vectordata.Property{ + property := buildEmbeddingProperty[batchEmbeddingRecord](t, + &vectordata.Definition{Properties: []vectordata.Property{ &vectordata.KeyProperty{Name: "ID", Type: reflect.TypeFor[string]()}, &vectordata.VectorProperty{Name: "Text", Type: reflect.TypeFor[string](), Dimensions: 2, EmbeddingType: reflect.TypeFor[[]float32]()}, }}, - }) + ) if result, err := providerservices.GenerateEmbeddings(t.Context(), property, records); err == nil || result != nil { t.Fatalf("missing generator = %+v, %v", result, err) } @@ -240,7 +293,7 @@ func TestGenerateEmbeddings_CancellationAndProviderErrors(t *testing.T) { return nil, nil }, } - property := buildEmbeddingProperty[batchEmbeddingRecord](t, vectordata.CollectionOptions{EmbeddingGenerator: embedding.Bind(generator)}) + property := buildEmbeddingProperty[batchEmbeddingRecord](t, &vectordata.Definition{EmbeddingGenerator: embedding.Bind(generator)}) ctx, cancel := context.WithCancel(t.Context()) cancel() if result, err := providerservices.GenerateEmbeddings(ctx, property, []batchEmbeddingRecord{{Text: "value"}}); !errors.Is(err, context.Canceled) || result != nil { @@ -263,7 +316,7 @@ func TestGenerateEmbeddings_CancellationAndProviderErrors(t *testing.T) { return nil, want }, } - property := buildEmbeddingProperty[batchEmbeddingRecord](t, vectordata.CollectionOptions{EmbeddingGenerator: embedding.Bind(generator)}) + property := buildEmbeddingProperty[batchEmbeddingRecord](t, &vectordata.Definition{EmbeddingGenerator: embedding.Bind(generator)}) result, err := providerservices.GenerateEmbeddings(ctx, property, []batchEmbeddingRecord{{Text: "first"}, {Text: "second"}}) if !errors.Is(err, want) || errors.Is(err, context.Canceled) { t.Fatalf("provider error was replaced: %v", err) diff --git a/vectordata/providerservices/filter_test.go b/vectordata/providerservices/filter_test.go index 2c5f4d150..6e028d9b2 100644 --- a/vectordata/providerservices/filter_test.go +++ b/vectordata/providerservices/filter_test.go @@ -233,13 +233,13 @@ func TestCompileFilter_Comparisons(t *testing.T) { } func TestCompileFilter_ModelAccessors(t *testing.T) { - options := vectordata.CollectionOptions{Definition: &vectordata.Definition{Properties: []vectordata.Property{ + options := &vectordata.Definition{Properties: []vectordata.Property{ &vectordata.KeyProperty{Name: "ID", Type: reflect.TypeFor[string]()}, &vectordata.DataProperty{Name: "Text", Type: reflect.TypeFor[string](), StorageName: "text"}, - }}} + }} for _, test := range []struct { name string - build func(vectordata.CollectionOptions) (*providerservices.CollectionModel, error) + build func(*vectordata.Definition) (*providerservices.CollectionModel, error) record any }{ {"struct", buildModel[string, taggedRecord], taggedRecord{Text: "value"}}, diff --git a/vectordata/store.go b/vectordata/store.go new file mode 100644 index 000000000..8aee9b8eb --- /dev/null +++ b/vectordata/store.go @@ -0,0 +1,109 @@ +// Copyright (c) Microsoft. All rights reserved. + +package vectordata + +import ( + "context" + "fmt" + "io" + "iter" + "reflect" +) + +// StoreProvider implements the database operations behind a [Store]. +// Implementations must support concurrent calls and treat inputs as read-only. +type StoreProvider interface { + // OpenCollection returns a handle without creating the collection. + // keyType and recordType describe the caller's Go types. The provider must + // validate their compatibility with the schema and retain the actual record + // types across operations; the any parameters are only a dispatch boundary. + // Optional batch, hybrid-search, and io.Closer capabilities are preserved + // by the store's typed collection adapters. + OpenCollection(name string, keyType, recordType reflect.Type, definition *Definition) (Collection[any, any], error) + + ListCollectionNames(ctx context.Context) iter.Seq2[string, error] + CollectionExists(ctx context.Context, name string) (bool, error) + EnsureCollectionDeleted(ctx context.Context, name string) error +} + +// Store manages collections that may have different key and record types. +// Create one with [NewStore] or a database provider's constructor. Use +// [Store.DynamicCollection] for map records or [GetCollection] for typed records. +// Go 1.27 also supports the generic Collection method. +type Store struct { + provider StoreProvider +} + +// NewStore wraps a database provider. It panics if provider is nil. +func NewStore(provider StoreProvider) *Store { + if provider == nil { + panic("vector store provider is required") + } + return &Store{provider: provider} +} + +// DynamicCollection returns a handle for map records without creating the +// collection. The definition must be non-nil and specify explicit property types. +// Keys use the schema's key type. This method is available on all supported Go +// versions. Call [Collection.EnsureExists] to create the collection. +func (s *Store) DynamicCollection(name string, definition *Definition) (Collection[any, map[string]any], error) { + if definition == nil { + return nil, fmt.Errorf("vectordata: dynamic collections require a definition") + } + return GetCollection[any, map[string]any](s, name, definition) +} + +// GetCollection returns a typed handle without creating the collection. +// A nil definition uses the record's vectordata tags. +// Providers validate the requested types and schema. The returned collection +// preserves record identity, provider errors, and optional provider capabilities. +// It reports an error if a provider returns a record of an incompatible type. +// +// This function supports Go 1.26 callers. On Go 1.27, s.Collection[TKey, TRecord] +// provides the same operation as a generic method. +func GetCollection[TKey, TRecord any](s *Store, name string, definition *Definition) (Collection[TKey, TRecord], error) { + if s == nil || s.provider == nil { + return nil, fmt.Errorf("vectordata: store is not initialized") + } + raw, err := s.provider.OpenCollection(name, reflect.TypeFor[TKey](), reflect.TypeFor[TRecord](), definition) + if err != nil { + return nil, err + } + if raw == nil { + return nil, fmt.Errorf("vectordata: provider returned a nil collection") + } + collection := &typedCollection[TKey, TRecord]{raw: raw} + closer, canClose := raw.(io.Closer) + if hybrid, ok := raw.(KeywordHybridSearcher[any]); ok { + collection := &typedHybridCollection[TKey, TRecord]{typedCollection: collection, hybrid: hybrid} + if canClose { + return &struct { + *typedHybridCollection[TKey, TRecord] + io.Closer + }{collection, closer}, nil + } + return collection, nil + } + if canClose { + return &struct { + *typedCollection[TKey, TRecord] + io.Closer + }{collection, closer}, nil + } + return collection, nil +} + +// ListCollectionNames returns the names of the collections in the store. +func (s *Store) ListCollectionNames(ctx context.Context) iter.Seq2[string, error] { + return s.provider.ListCollectionNames(ctx) +} + +// CollectionExists reports whether the named collection exists. +func (s *Store) CollectionExists(ctx context.Context, name string) (bool, error) { + return s.provider.CollectionExists(ctx, name) +} + +// EnsureCollectionDeleted deletes the named collection if it exists. +func (s *Store) EnsureCollectionDeleted(ctx context.Context, name string) error { + return s.provider.EnsureCollectionDeleted(ctx, name) +} diff --git a/vectordata/store_go127.go b/vectordata/store_go127.go new file mode 100644 index 000000000..cc943fc8c --- /dev/null +++ b/vectordata/store_go127.go @@ -0,0 +1,11 @@ +// Copyright (c) Microsoft. All rights reserved. + +//go:build go1.27 + +package vectordata + +// Collection returns a typed handle without creating the collection. +// It is equivalent to [GetCollection] and requires Go 1.27 or later. +func (s *Store) Collection[TKey, TRecord any](name string, definition *Definition) (Collection[TKey, TRecord], error) { + return GetCollection[TKey, TRecord](s, name, definition) +} diff --git a/vectordata/store_go127_test.go b/vectordata/store_go127_test.go new file mode 100644 index 000000000..dc1fd4fa9 --- /dev/null +++ b/vectordata/store_go127_test.go @@ -0,0 +1,64 @@ +// Copyright (c) Microsoft. All rights reserved. + +//go:build go1.27 + +package vectordata_test + +import ( + "context" + "reflect" + "testing" + + "github.com/microsoft/agent-framework-go/vectordata" + "github.com/microsoft/agent-framework-go/vectordata/inmemory" +) + +func TestStore_GenericMethod(t *testing.T) { + type document struct { + ID string `vectordata:"key"` + } + type counter struct { + ID int `vectordata:"key"` + } + store := inmemory.New(inmemory.Config{}) + documents, err := store.Collection[string, document]("documents", nil) + if err != nil { + t.Fatal(err) + } + facade := store.Store + counters, err := facade.Collection[int, *counter]("counters", nil) + if err != nil { + t.Fatal(err) + } + dynamic, err := store.DynamicCollection("dynamic", &vectordata.Definition{Properties: []vectordata.Property{ + &vectordata.KeyProperty{Name: "ID", Type: reflect.TypeFor[string]()}, + }}) + if err != nil { + t.Fatal(err) + } + for _, collection := range []interface{ EnsureExists(context.Context) error }{documents, counters, dynamic} { + if err := collection.EnsureExists(t.Context()); err != nil { + t.Fatal(err) + } + } + if err := documents.Upsert(t.Context(), document{ID: "one"}); err != nil { + t.Fatal(err) + } + value := &counter{ID: 2} + if err := counters.Upsert(t.Context(), value); err != nil { + t.Fatal(err) + } + if err := dynamic.Upsert(t.Context(), map[string]any{"ID": "three"}); err != nil { + t.Fatal(err) + } + if got, found, err := counters.Get(t.Context(), 2, nil); err != nil || !found || got != value { + t.Fatalf("generic method lost record identity: (%v, %v, %v)", got, found, err) + } + fallback, err := vectordata.GetCollection[string, document](store.Store, "documents", nil) + if err != nil { + t.Fatal(err) + } + if got, found, err := fallback.Get(t.Context(), "one", nil); err != nil || !found || got.ID != "one" { + t.Fatalf("method and function do not share storage: (%v, %v, %v)", got, found, err) + } +} diff --git a/vectordata/store_test.go b/vectordata/store_test.go new file mode 100644 index 000000000..c3f958d12 --- /dev/null +++ b/vectordata/store_test.go @@ -0,0 +1,605 @@ +// Copyright (c) Microsoft. All rights reserved. + +package vectordata_test + +import ( + "context" + "errors" + "fmt" + "io" + "iter" + "reflect" + "slices" + "testing" + + "github.com/microsoft/agent-framework-go/vectordata" +) + +type storeProvider struct { + vectordata.StoreProvider + open func(string, reflect.Type, reflect.Type, *vectordata.Definition) (vectordata.Collection[any, any], error) +} + +func (p *storeProvider) OpenCollection(name string, key, record reflect.Type, definition *vectordata.Definition) (vectordata.Collection[any, any], error) { + return p.open(name, key, record, definition) +} + +type rawCollection struct { + vectordata.Collection[any, any] + get func(context.Context, any, *vectordata.RecordRetrievalOptions) (any, bool, error) + upsert func(context.Context, ...any) error + delete func(context.Context, any) error + filtered func(context.Context, vectordata.Filter, int, *vectordata.FilteredRecordRetrievalOptions) iter.Seq2[any, error] + search func(context.Context, any, int, *vectordata.SearchOptions) iter.Seq2[vectordata.SearchResult[any], error] +} + +func (c *rawCollection) Get(ctx context.Context, key any, options *vectordata.RecordRetrievalOptions) (any, bool, error) { + return c.get(ctx, key, options) +} + +func (c *rawCollection) Upsert(ctx context.Context, records ...any) error { + return c.upsert(ctx, records...) +} + +func (c *rawCollection) Delete(ctx context.Context, key any) error { return c.delete(ctx, key) } + +func (c *rawCollection) GetFiltered(ctx context.Context, filter vectordata.Filter, top int, options *vectordata.FilteredRecordRetrievalOptions) iter.Seq2[any, error] { + return c.filtered(ctx, filter, top, options) +} + +func (c *rawCollection) Search(ctx context.Context, query any, top int, options *vectordata.SearchOptions) iter.Seq2[vectordata.SearchResult[any], error] { + return c.search(ctx, query, top, options) +} + +func collectionStore(raw vectordata.Collection[any, any]) *vectordata.Store { + return vectordata.NewStore(&storeProvider{open: func(string, reflect.Type, reflect.Type, *vectordata.Definition) (vectordata.Collection[any, any], error) { + return raw, nil + }}) +} + +func TestStore_FactoryTypesAndErrors(t *testing.T) { + definition := &vectordata.Definition{Properties: []vectordata.Property{ + &vectordata.KeyProperty{Name: "ID", Type: reflect.TypeFor[string]()}, + }} + var gotKey, gotRecord reflect.Type + calls := 0 + wantErr := errors.New("open failed") + var failure error + var raw vectordata.Collection[any, any] = &rawCollection{} + store := vectordata.NewStore(&storeProvider{open: func(name string, key, record reflect.Type, got *vectordata.Definition) (vectordata.Collection[any, any], error) { + calls++ + if name != "records" || got != definition { + t.Fatal("factory changed name or definition") + } + gotKey, gotRecord = key, record + return raw, failure + }}) + if _, err := vectordata.GetCollection[string, *record](store, "records", definition); err != nil || + gotKey != reflect.TypeFor[string]() || gotRecord != reflect.TypeFor[*record]() { + t.Fatalf("typed factory types = (%v, %v), error = %v", gotKey, gotRecord, err) + } + if _, err := store.DynamicCollection("records", definition); err != nil || + gotKey != reflect.TypeFor[any]() || gotRecord != reflect.TypeFor[map[string]any]() { + t.Fatalf("dynamic factory types = (%v, %v), error = %v", gotKey, gotRecord, err) + } + if _, err := store.DynamicCollection("records", nil); err == nil || calls != 2 { + t.Fatalf("dynamic collection without schema: calls=%d, error=%v", calls, err) + } + definition = nil + if _, err := vectordata.GetCollection[string, *record](store, "records", nil); err != nil || calls != 3 { + t.Fatalf("typed collection without explicit schema: calls=%d, error=%v", calls, err) + } + failure = wantErr + if _, err := vectordata.GetCollection[string, *record](store, "records", definition); err != wantErr { + t.Fatalf("provider failure = %v", err) + } + failure, raw = nil, nil + if _, err := vectordata.GetCollection[string, *record](store, "records", definition); err == nil { + t.Fatal("provider's nil collection was accepted") + } + for _, uninitialized := range []*vectordata.Store{nil, {}} { + if _, err := vectordata.GetCollection[string, record](uninitialized, "records", definition); err == nil { + t.Fatal("uninitialized store was accepted") + } + } + defer func() { + if recover() == nil { + t.Fatal("nil provider must panic") + } + }() + vectordata.NewStore(nil) +} + +func TestStore_TypedCRUD(t *testing.T) { + ctx := t.Context() + wantErr := errors.New("provider error") + input := &record{ID: "one", Text: "text"} + getOptions := &vectordata.RecordRetrievalOptions{IncludeVectors: true} + calls := 0 + raw := &rawCollection{ + get: func(got context.Context, key any, options *vectordata.RecordRetrievalOptions) (any, bool, error) { + if got != ctx || key != "one" || options != getOptions { + t.Fatal("Get arguments changed") + } + return input, true, nil + }, + upsert: func(got context.Context, values ...any) error { + if got != ctx { + t.Fatal("Upsert context changed") + } + switch calls { + case 0: + if values != nil { + t.Fatal("nil batch became non-nil") + } + case 1: + if values == nil || len(values) != 0 { + t.Fatal("empty batch lost its representation") + } + case 2: + if len(values) != 2 || values[0] != input || values[1].(*record) != nil { + t.Fatalf("upsert values = %v", values) + } + } + calls++ + return wantErr + }, + delete: func(got context.Context, key any) error { + if got != ctx || key != "one" { + t.Fatal("Delete arguments changed") + } + return wantErr + }, + } + collection, err := vectordata.GetCollection[string, *record](collectionStore(raw), "records", nil) + if err != nil { + t.Fatal(err) + } + got, found, err := collection.Get(ctx, "one", getOptions) + if err != nil || !found || got != input { + t.Fatalf("Get = (%v, %v, %v)", got, found, err) + } + for _, records := range [][]*record{nil, {}, {input, nil}} { + if err := collection.Upsert(ctx, records...); err != wantErr { + t.Fatalf("Upsert error = %v", err) + } + } + if calls != 3 { + t.Fatalf("Upsert calls = %d", calls) + } + if err := collection.Delete(ctx, "one"); err != wantErr { + t.Fatalf("Delete error = %v", err) + } + for _, tc := range []struct { + value any + found bool + err error + valid bool + }{ + {value: "wrong", found: true}, + {value: "wrong", found: false, valid: true}, + {value: "wrong", found: true, err: wantErr}, + {value: (*record)(nil), found: true, valid: true}, + {value: nil, found: true, valid: true}, + } { + raw.get = func(context.Context, any, *vectordata.RecordRetrievalOptions) (any, bool, error) { + return tc.value, tc.found, tc.err + } + got, found, err := collection.Get(ctx, "one", nil) + if got != nil || found != tc.found || tc.valid && err != nil || !tc.valid && err == nil || + tc.err != nil && err != tc.err { + t.Fatalf("Get with %#v = (%v, %v, %v)", tc, got, found, err) + } + } + valueCollection, err := vectordata.GetCollection[string, record](collectionStore(raw), "records", nil) + if err != nil { + t.Fatal(err) + } + if _, _, err := valueCollection.Get(ctx, "one", nil); err == nil { + t.Fatal("nil provider value was accepted as a struct") + } +} + +func TestStore_TypedIterators(t *testing.T) { + for _, operation := range []string{"filter", "search", "batch"} { + for _, outcome := range []string{"success", "bad type", "provider error", "consumer stop"} { + t.Run(operation+"/"+outcome, func(t *testing.T) { + wantErr := errors.New("provider failed") + input := &record{ID: "one"} + calls, yielded, released := 0, 0, false + values := func(yield func(any, error) bool) { + calls++ + defer func() { released = true }() + if !yield(input, nil) { + return + } + yielded++ + switch outcome { + case "bad type": + if yield("wrong", nil) { + t.Error("continued past a type error") + } + case "provider error": + if yield("wrong", wantErr) { + t.Error("continued past provider error") + } + } + } + options := &vectordata.SearchOptions{Skip: 2} + filterOptions := &vectordata.FilteredRecordRetrievalOptions{Skip: 2} + filter := vectordata.Filter{Operator: vectordata.Equal, Property: "ID", Value: "one"} + score := new(0.75) + raw := &rawCollection{ + filtered: func(ctx context.Context, got vectordata.Filter, top int, opts *vectordata.FilteredRecordRetrievalOptions) iter.Seq2[any, error] { + if ctx != t.Context() || !reflect.DeepEqual(got, filter) || top != 3 || opts != filterOptions { + t.Fatal("filter arguments changed") + } + return values + }, + search: func(ctx context.Context, query any, top int, opts *vectordata.SearchOptions) iter.Seq2[vectordata.SearchResult[any], error] { + if ctx != t.Context() || query != "query" || top != 3 || opts != options { + t.Fatal("search arguments changed") + } + return func(yield func(vectordata.SearchResult[any], error) bool) { + for value, err := range values { + if !yield(vectordata.SearchResult[any]{Record: value, Score: score}, err) { + return + } + } + } + }, + } + backend := &rawBatchCollection{rawCollection: raw, getBatch: func(context.Context, []any, *vectordata.RecordRetrievalOptions) iter.Seq2[any, error] { + return values + }} + collection, err := vectordata.GetCollection[string, *record](collectionStore(backend), "records", nil) + if err != nil { + t.Fatal(err) + } + var sequence iter.Seq2[*record, error] + switch operation { + case "filter": + sequence = collection.GetFiltered(t.Context(), filter, 3, filterOptions) + case "batch": + sequence = vectordata.GetBatch(t.Context(), collection, []string{"one", "two"}, nil) + case "search": + sequence = func(yield func(*record, error) bool) { + for result, err := range collection.Search(t.Context(), "query", 3, options) { + if err == nil && result.Score != score { + t.Error("search score identity changed") + } + if !yield(result.Record, err) { + return + } + } + } + } + if calls != 0 { + t.Fatal("iterator started eagerly") + } + count := 0 + var gotErr error + for value, err := range sequence { + if err != nil { + if value != nil { + t.Error("failed result contains a record") + } + gotErr = err + continue + } + count++ + if value != input { + t.Fatal("record identity changed") + } + if outcome == "consumer stop" { + break + } + } + if count != 1 || calls != 1 || !released || outcome == "consumer stop" && yielded != 0 || + outcome == "provider error" && gotErr != wantErr || outcome == "bad type" && gotErr == nil || + (outcome == "success" || outcome == "consumer stop") && gotErr != nil { + t.Fatalf("count=%d calls=%d released=%t yielded=%d error=%v", count, calls, released, yielded, gotErr) + } + }) + } + } +} + +type rawBatchCollection struct { + *rawCollection + getBatch func(context.Context, []any, *vectordata.RecordRetrievalOptions) iter.Seq2[any, error] + deleteBatch func(context.Context, []any) error +} + +func (c *rawBatchCollection) GetBatch(ctx context.Context, keys []any, options *vectordata.RecordRetrievalOptions) iter.Seq2[any, error] { + return c.getBatch(ctx, keys, options) +} + +func (c *rawBatchCollection) DeleteBatch(ctx context.Context, keys []any) error { + return c.deleteBatch(ctx, keys) +} + +func TestStore_BatchDispatch(t *testing.T) { + for _, native := range []bool{false, true} { + t.Run(fmt.Sprint(native), func(t *testing.T) { + var getKeys, deleteKeys []any + getCalls, deleteCalls := 0, 0 + wantErr := errors.New("delete failed") + base := &rawCollection{ + get: func(_ context.Context, key any, _ *vectordata.RecordRetrievalOptions) (any, bool, error) { + getCalls++ + getKeys = append(getKeys, key) + return &record{ID: key.(string)}, true, nil + }, + delete: func(_ context.Context, key any) error { + deleteCalls++ + deleteKeys = append(deleteKeys, key) + if key == "two" { + return wantErr + } + return nil + }, + } + var raw vectordata.Collection[any, any] = base + if native { + raw = &rawBatchCollection{ + rawCollection: base, + getBatch: func(_ context.Context, keys []any, _ *vectordata.RecordRetrievalOptions) iter.Seq2[any, error] { + getCalls++ + getKeys = keys + return func(yield func(any, error) bool) { + for _, key := range keys { + if !yield(&record{ID: key.(string)}, nil) { + return + } + } + } + }, + deleteBatch: func(_ context.Context, keys []any) error { + deleteCalls++ + deleteKeys = keys + return wantErr + }, + } + } + collection, err := vectordata.GetCollection[string, *record](collectionStore(raw), "records", nil) + if err != nil { + t.Fatal(err) + } + for _, keys := range [][]string{nil, {}} { + for _, err := range vectordata.GetBatch(t.Context(), collection, keys, nil) { + t.Fatalf("empty batch produced a result: %v", err) + } + if err := vectordata.DeleteBatch(t.Context(), collection, keys); err != nil { + t.Fatal(err) + } + } + keys := []string{"one", "two"} + var got []string + for record, err := range vectordata.GetBatch(t.Context(), collection, keys, nil) { + if err != nil { + t.Fatal(err) + } + got = append(got, record.ID) + } + if err := vectordata.DeleteBatch(t.Context(), collection, keys); err != wantErr { + t.Fatalf("delete error = %v", err) + } + wantCalls := 2 + if native { + wantCalls = 1 + } + if !slices.Equal(got, keys) || !reflect.DeepEqual(getKeys, []any{"one", "two"}) || + !reflect.DeepEqual(deleteKeys, getKeys) || getCalls != wantCalls || deleteCalls != wantCalls { + t.Fatalf("get=%v delete=%v calls=(%d,%d)", getKeys, deleteKeys, getCalls, deleteCalls) + } + }) + } +} + +type closeFunc func() error + +func (f closeFunc) Close() error { return f() } + +type hybridFunc func(context.Context, any, []string, int, *vectordata.HybridSearchOptions) iter.Seq2[vectordata.SearchResult[any], error] + +func (f hybridFunc) HybridSearch(ctx context.Context, value any, keywords []string, top int, options *vectordata.HybridSearchOptions) iter.Seq2[vectordata.SearchResult[any], error] { + return f(ctx, value, keywords, top, options) +} + +func TestStore_OptionalCapabilities(t *testing.T) { + for _, hybrid := range []bool{false, true} { + for _, closable := range []bool{false, true} { + t.Run(fmt.Sprintf("hybrid=%t/close=%t", hybrid, closable), func(t *testing.T) { + base := &rawCollection{} + input := &record{ID: "one"} + hybridCalls, closeCalls := 0, 0 + wantErr := errors.New("close failed") + keywords := []string{"one"} + options := &vectordata.HybridSearchOptions{Skip: 2} + search := hybridFunc(func(ctx context.Context, value any, got []string, top int, opts *vectordata.HybridSearchOptions) iter.Seq2[vectordata.SearchResult[any], error] { + hybridCalls++ + if ctx != t.Context() || value != "query" || &got[0] != &keywords[0] || top != 3 || opts != options { + t.Fatal("hybrid arguments changed") + } + return func(yield func(vectordata.SearchResult[any], error) bool) { + yield(vectordata.SearchResult[any]{Record: input}, nil) + } + }) + closeCollection := closeFunc(func() error { closeCalls++; return wantErr }) + var raw vectordata.Collection[any, any] = base + switch { + case hybrid && closable: + raw = &struct { + *rawCollection + vectordata.KeywordHybridSearcher[any] + io.Closer + }{base, search, closeCollection} + case hybrid: + raw = &struct { + *rawCollection + vectordata.KeywordHybridSearcher[any] + }{base, search} + case closable: + raw = &struct { + *rawCollection + io.Closer + }{base, closeCollection} + } + collection, err := vectordata.GetCollection[string, *record](collectionStore(raw), "records", nil) + if err != nil { + t.Fatal(err) + } + hybridCollection, hasHybrid := collection.(vectordata.KeywordHybridSearcher[*record]) + closer, hasClose := collection.(io.Closer) + if hasHybrid != hybrid || hasClose != closable { + t.Fatalf("capabilities = hybrid:%t close:%t", hasHybrid, hasClose) + } + if hybrid { + results := hybridCollection.HybridSearch(t.Context(), "query", keywords, 3, options) + if hybridCalls != 0 { + t.Fatal("hybrid iterator started eagerly") + } + for result, err := range results { + if err != nil || result.Record != input { + t.Fatalf("hybrid result = (%v, %v)", result, err) + } + } + if hybridCalls != 1 { + t.Fatalf("hybrid calls = %d", hybridCalls) + } + } + if closable { + if err := closer.Close(); err != wantErr || closeCalls != 1 { + t.Fatalf("close = %v, calls = %d", err, closeCalls) + } + } + }) + } + } +} + +func TestStore_CollectionCloseDoesNotCloseStore(t *testing.T) { + wantErr := errors.New("close failed") + calls, collectionCalls := 0, 0 + raw := &struct { + *rawCollection + io.Closer + }{&rawCollection{}, closeFunc(func() error { collectionCalls++; return nil })} + provider := &struct { + *storeProvider + io.Closer + }{&storeProvider{open: func(string, reflect.Type, reflect.Type, *vectordata.Definition) (vectordata.Collection[any, any], error) { + return raw, nil + }}, closeFunc(func() error { calls++; return wantErr })} + store := vectordata.NewStore(provider) + if _, ok := any(store).(io.Closer); ok { + t.Fatal("store facade advertises provider-specific cleanup") + } + collection, err := vectordata.GetCollection[string, record](store, "records", nil) + if err != nil { + t.Fatal(err) + } + if err := collection.(io.Closer).Close(); err != nil || collectionCalls != 1 || calls != 0 { + t.Fatalf("collection close = %v, store calls = %d, collection calls = %d", err, calls, collectionCalls) + } + if err := provider.Close(); err != wantErr || calls != 1 || collectionCalls != 1 { + t.Fatalf("provider close = %v, store calls = %d, collection calls = %d", err, calls, collectionCalls) + } +} + +type lifecycleBackend struct { + vectordata.Collection[any, any] + vectordata.StoreProvider + t *testing.T + err error + calls []string +} + +func (b *lifecycleBackend) check(ctx context.Context, operation, name string) { + b.t.Helper() + if ctx != b.t.Context() || name != "records" { + b.t.Fatal("lifecycle arguments changed") + } + b.calls = append(b.calls, operation) +} + +func (*lifecycleBackend) Name() string { return "records" } + +func (b *lifecycleBackend) Exists(ctx context.Context) (bool, error) { + b.check(ctx, "Exists", "records") + return true, b.err +} + +func (b *lifecycleBackend) EnsureExists(ctx context.Context) error { + b.check(ctx, "EnsureExists", "records") + return b.err +} + +func (b *lifecycleBackend) EnsureDeleted(ctx context.Context) error { + b.check(ctx, "EnsureDeleted", "records") + return b.err +} + +func (b *lifecycleBackend) ListCollectionNames(ctx context.Context) iter.Seq2[string, error] { + return func(yield func(string, error) bool) { + b.check(ctx, "ListCollectionNames", "records") + if yield("records", nil) { + yield("", b.err) + } + } +} + +func (b *lifecycleBackend) CollectionExists(ctx context.Context, name string) (bool, error) { + b.check(ctx, "CollectionExists", name) + return true, b.err +} + +func (b *lifecycleBackend) EnsureCollectionDeleted(ctx context.Context, name string) error { + b.check(ctx, "EnsureCollectionDeleted", name) + return b.err +} + +func TestStore_LifecycleForwarding(t *testing.T) { + backend := &lifecycleBackend{t: t, err: context.Canceled} + store := vectordata.NewStore(&storeProvider{StoreProvider: backend, open: func(string, reflect.Type, reflect.Type, *vectordata.Definition) (vectordata.Collection[any, any], error) { + return backend, nil + }}) + collection, err := vectordata.GetCollection[string, record](store, "records", nil) + if err != nil { + t.Fatal(err) + } + if collection.Name() != "records" { + t.Fatal("collection name changed") + } + if found, err := collection.Exists(t.Context()); !found || err != context.Canceled { + t.Fatalf("Exists = (%v, %v)", found, err) + } + if err := collection.EnsureExists(t.Context()); err != context.Canceled { + t.Fatalf("EnsureExists = %v", err) + } + if err := collection.EnsureDeleted(t.Context()); err != context.Canceled { + t.Fatalf("EnsureDeleted = %v", err) + } + if found, err := store.CollectionExists(t.Context(), "records"); !found || err != context.Canceled { + t.Fatalf("CollectionExists = (%v, %v)", found, err) + } + if err := store.EnsureCollectionDeleted(t.Context(), "records"); err != context.Canceled { + t.Fatalf("EnsureCollectionDeleted = %v", err) + } + var names []string + var listErr error + for name, err := range store.ListCollectionNames(t.Context()) { + if err != nil { + listErr = err + } else { + names = append(names, name) + } + } + if !slices.Equal(names, []string{"records"}) || listErr != context.Canceled { + t.Fatalf("ListCollectionNames = (%v, %v)", names, listErr) + } + if !slices.Equal(backend.calls, []string{"Exists", "EnsureExists", "EnsureDeleted", "CollectionExists", "EnsureCollectionDeleted", "ListCollectionNames"}) { + t.Fatalf("lifecycle calls = %v", backend.calls) + } +} diff --git a/vectordata/vectordata.go b/vectordata/vectordata.go index 6818958c9..b38916587 100644 --- a/vectordata/vectordata.go +++ b/vectordata/vectordata.go @@ -3,8 +3,8 @@ // Package vectordata defines provider-independent vector store contracts, // record schemas, retrieval options, and vector and keyword-hybrid searches. // -// Store and Collection are implemented by database providers. They must support -// concurrent calls, honor context cancellation, and treat caller-supplied schemas, +// Store wraps a database provider and creates typed or dynamic collections. +// Providers must support concurrent calls, honor context cancellation, and treat caller-supplied schemas, // filters, and options as read-only. Record ownership is provider-specific: // providers document whether records are copied or retained, whether returned // records share storage, and whether generated keys are written into inputs. @@ -24,27 +24,8 @@ package vectordata import ( "context" "iter" - - "github.com/microsoft/agent-framework-go/embedding" ) -// Store contains collections sharing the same key and record types. -// Providers may offer differently typed stores over the same client. -type Store[TKey, TRecord any] interface { - // Collection returns a configured handle for the named collection without - // creating it. Call [Collection.EnsureExists] to create the collection. - Collection(name string, options CollectionOptions) (Collection[TKey, TRecord], error) - - // ListCollectionNames returns the names of the collections in the store. - ListCollectionNames(ctx context.Context) iter.Seq2[string, error] - - // CollectionExists reports whether the named collection exists in the store. - CollectionExists(ctx context.Context, name string) (bool, error) - - // EnsureCollectionDeleted deletes the named collection if it exists. - EnsureCollectionDeleted(ctx context.Context, name string) error -} - // Collection manages a named collection and its records. Providers document // which key, data, vector, filter, and index types they support. type Collection[TKey, TRecord any] interface { @@ -113,20 +94,6 @@ type SearchResult[TRecord any] struct { Score *float64 } -// CollectionOptions configures a collection handle. Definition overrides schema -// tags. EmbeddingGenerator is the provider's default embedding generator; vector -// property and definition generators take precedence, in that order. -// -// Use [embedding.Bind] to configure a typed generator. The input and -// output types may differ between vector properties. Providers resolve schema -// configuration and validate input compatibility, supported representations, and -// storage requirements when constructing a collection handle. The supplied -// definition is configuration, not a resolved runtime model, and is read-only. -type CollectionOptions struct { - Definition *Definition - EmbeddingGenerator *embedding.BoundGenerator -} - // Error describes a vector store failure while preserving the original cause. // Provider errors may use it to attach database, collection, and operation names. type Error struct {