diff --git a/examples/02-agents/providers/azure/cosmosdb/main.go b/examples/02-agents/providers/azure/cosmosdb/main.go new file mode 100644 index 00000000..9278047c --- /dev/null +++ b/examples/02-agents/providers/azure/cosmosdb/main.go @@ -0,0 +1,86 @@ +// Copyright (c) Microsoft. All rights reserved. + +// This sample stores generated embeddings in a temporary Cosmos DB container. +// It uses an existing database and deletes only its own container before exit. +package main + +import ( + "context" + "errors" + "fmt" + "os" + "strings" + "time" + + "github.com/Azure/azure-sdk-for-go/sdk/data/azcosmos" + "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"` + Topic string `vectordata:"data,indexed,storage=topic"` + Text string `vectordata:"data,storage=text"` + Content string `vectordata:"vector,dimensions=3072,index=disk_ann,storage=vector"` +} + +func main() { + ctx := context.Background() + demo.NewLogger("Azure Cosmos DB", "Demonstrates generated embeddings and scoped vector search.", + "Embedding Model", demo.FoundryEmbeddingModel) + endpoint := strings.TrimSpace(os.Getenv("AZURE_COSMOS_ENDPOINT")) + database := strings.TrimSpace(os.Getenv("AZURE_COSMOS_DATABASE")) + if endpoint == "" || database == "" { + demo.Panic("AZURE_COSMOS_ENDPOINT and AZURE_COSMOS_DATABASE must be set") + } + credential := demo.FoundryTokenCredential() + client, err := azcosmos.NewClient(endpoint, credential, nil) + if err != nil { + demo.Panic(err) + } + defer client.Close() + generator := foundryprovider.NewEmbeddingGenerator(demo.FoundryProjectEndpoint, credential, + foundryprovider.EmbeddingGeneratorConfig{Model: demo.FoundryEmbeddingModel}) + store, err := azureprovider.NewCosmosDB(client, database, azureprovider.CosmosDBConfig{ + PartitionKeyProperties: []string{"Topic"}, + EmbeddingGenerator: embedding.Bind(generator), + }) + if err != nil { + demo.Panic(err) + } + name := "afgo-example-" + uuid.NewString() + collection, err := vectordata.GetCollection[azureprovider.CosmosKey, *document](store.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 container: %s\n\n", name) + if err = collection.EnsureExists(ctx); err != nil { + return + } + if err = collection.Upsert(ctx, + &document{Topic: "humor", Text: "I enjoy pirate jokes about parrots and ships.", Content: "I enjoy pirate jokes about parrots and ships."}, + &document{Topic: "travel", Text: "I like hiking in the mountains.", Content: "I like hiking in the mountains."}, + ); err != nil { + return + } + filter := vectordata.Filter{Operator: vectordata.Equal, Property: "Topic", Value: "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) + } +} diff --git a/go.mod b/go.mod index 2c09bbc1..38c41a25 100644 --- a/go.mod +++ b/go.mod @@ -5,6 +5,7 @@ go 1.26.0 require ( github.com/Azure/azure-sdk-for-go/sdk/azcore v1.23.3 github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.15.0 + github.com/Azure/azure-sdk-for-go/sdk/data/azcosmos v1.5.0 github.com/a2aproject/a2a-go/v2 v2.6.0 github.com/ag-ui-protocol/ag-ui/sdks/community/go v0.0.0-20260312103001-8e7ab1df34c8 github.com/anthropics/anthropic-sdk-go v1.82.0 diff --git a/go.sum b/go.sum index f5ca2723..63a64ff7 100644 --- a/go.sum +++ b/go.sum @@ -12,6 +12,8 @@ github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.15.0 h1:w3gDUJlJ0BfBsHH02DCB github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.15.0/go.mod h1:7scl8mG/IFLVjJ9ef/ltdND2T1z0CkgVSaQoUxLvmeU= github.com/Azure/azure-sdk-for-go/sdk/azidentity/cache v0.4.0 h1:xFaZZ+IubdftrDHnGGwZ6QvQ3KHTtWl2MCK+GMt2vxs= github.com/Azure/azure-sdk-for-go/sdk/azidentity/cache v0.4.0/go.mod h1:mCBhUhlMjLLJKr5aqw2TNS/VqJOie8MzWq3DAMJeKso= +github.com/Azure/azure-sdk-for-go/sdk/data/azcosmos v1.5.0 h1:wtCn7MemMD9eo4/NdpJ6S/MFD2BV2CDwoEfvl5th2vM= +github.com/Azure/azure-sdk-for-go/sdk/data/azcosmos v1.5.0/go.mod h1:MIyTWizpwnsX4LS9/tW1II9JL+D25Ypzj6URaT9NcgQ= github.com/Azure/azure-sdk-for-go/sdk/internal v1.13.0 h1:Terj1CMW4+NY5CpWjSRkqLoxh2DkYDdv/FH2z0897Uc= github.com/Azure/azure-sdk-for-go/sdk/internal v1.13.0/go.mod h1:7Ut2Qw9vcwpLW5Q4MBy43uSxXS7tK2yjs94OJ6j5thE= github.com/AzureAD/microsoft-authentication-extensions-for-go/cache v0.1.1 h1:WJTmL004Abzc5wDB5VtZG2PJk5ndYDgVacGqfirKxjM= diff --git a/provider/azureprovider/cosmosdb.go b/provider/azureprovider/cosmosdb.go new file mode 100644 index 00000000..b3fe62bb --- /dev/null +++ b/provider/azureprovider/cosmosdb.go @@ -0,0 +1,321 @@ +// Copyright (c) Microsoft. All rights reserved. + +package azureprovider + +import ( + "context" + "errors" + "fmt" + "iter" + "net/http" + "reflect" + "strings" + + "github.com/Azure/azure-sdk-for-go/sdk/azcore" + "github.com/Azure/azure-sdk-for-go/sdk/data/azcosmos" + "github.com/google/uuid" + "github.com/microsoft/agent-framework-go/embedding" + "github.com/microsoft/agent-framework-go/vectordata" +) + +// CosmosDBConfig configures container defaults shared by collections created +// from a Cosmos DB store. Create separate stores over the same client and +// database when different collections need different partition-key paths or +// indexing settings. +type CosmosDBConfig struct { + // PartitionKeyProperties lists record property names to use as the + // collection's partition-key paths, in hierarchical order. An empty list + // uses the document ID (/id), so collection keys must be string or + // uuid.UUID. A non-empty list uses up to three key or scalar data + // properties and requires [CosmosKey] for point operations unless the + // only partition property is the document ID. Record ID properties remain + // strings or uuid.UUID values, independently of the collection's lookup + // key type. Supported partition values are strings, booleans, numbers, + // uuid.UUID values, and null, including nullable pointers. + // Property names refer to the VectorData schema, not serialized field + // names; their storage names become Cosmos DB partition paths. The list + // is copied when the store is created. + PartitionKeyProperties []string + + // IndexingMode selects the indexing mode for newly created containers. It + // accepts "consistent", "lazy", or "none", case-insensitively; the default + // is "consistent". "none" omits vector and full-text index policies. + IndexingMode string + + // Automatic controls automatic indexing for newly created containers. Nil + // defaults to true; a pointer preserves an explicit false. + Automatic *bool + + // EmbeddingGenerator supplies the default generator for vector properties + // in this store. A collection definition or property-specific generator + // takes precedence. + EmbeddingGenerator *embedding.BoundGenerator +} + +// CosmosKey identifies a Cosmos DB document and its scalar or hierarchical +// partition key components when a collection uses custom partition paths. +// DocumentID must be nonblank. PartitionKey must contain one to three values +// matching [CosmosDBConfig.PartitionKeyProperties] in order. Components may +// be strings, booleans, numbers, [uuid.UUID] values, or nil; UUIDs are sent as +// strings. The record's key property stores only the document ID as a string +// or UUID, not a CosmosKey. Use a plain string or [uuid.UUID] lookup key when +// the collection uses the default /id partition path. +type CosmosKey struct { + DocumentID string + PartitionKey []any +} + +// CosmosDB manages a Cosmos DB NoSQL database through its embedded +// [vectordata.Store]. Create it with [NewCosmosDB]. +// +// A Cosmos DB client and database must already exist. The store borrows the +// client and embedding generator and does not close either. The caller remains +// responsible for closing the client. +// +// Collections support string, [uuid.UUID], and [CosmosKey] lookup keys. Record +// ID properties are strings or UUIDs. With the default partition path, string +// and UUID lookup keys identify both the document and its partition; custom +// and hierarchical paths require CosmosKey unless the only partition property +// is the document ID. UUID IDs are generated when the key property enables +// auto-generation. +// +// Data properties support strings, booleans, signed integers, float32, +// float64, time.Time, nullable pointers to supported values, and slices or +// arrays of supported scalar values. Vector properties support float32, uint8, +// and int8 coordinates, including the corresponding [embedding.VectorOf] and +// [embedding.Of] types. Native vectors can be upserted directly; other vector +// properties require an embedding generator. Property and definition +// generators override the store default. +// +// Stored timestamps and timestamp filter values use UTC RFC3339 strings with +// nine fractional-second digits, including nullable and collection values. +// Existing timestamps in other formats must be rewritten to this format for +// chronological ordering and range filters. Storage names with special +// characters are quoted as Cosmos DB path segments. The storage name +// "__cosmos_score" is reserved for search results. +// +// Container policies include DiskANN, Flat, and QuantizedFlat vector indexes; +// DiskANN and cosine distance are defaults. Cosine, dot-product, and Euclidean +// distances are supported. Full-text indexing is available for string data +// properties with the en-US language policy. The account must have vector +// search enabled, and full-text queries require the corresponding account +// capability. Service constraints still apply; for example, small vector +// datasets may be scanned until an index can be built. +// +// Search supports vector and full-text/vector hybrid queries with portable +// filters. The native Go client does not support cross-partition ranking or +// ordering: search and ordered retrieval require equality filters for every +// partition-key property, optionally nested in conjunctions. Unsupported +// queries return an error before calling the service or embedding generator. +// Choose partition paths that group records searched together; with the +// default /id path, ranked queries must filter to a single document ID. +// Unordered filtered retrieval supports cross-partition queries and applies +// Skip and top while paging. Queries project only schema properties, omitting +// vector fields unless IncludeVectors is true. Hybrid results are ranked with +// reciprocal-rank fusion and expose the raw vector distance as Score. +// Score thresholds compare against that distance (similarities use >=; +// Euclidean distance uses <=). +// +// IncludeVectors is unsupported on collections with generated vector fields, +// including missing records and empty query results. Nil native vectors are +// stored as null; non-nil vectors must have the configured dimensions. +// +// Item batches execute sequentially and may partially succeed. Generated +// input keys are not rolled back on failure. Per-collection partition and +// indexing settings are not available on the shared VectorData collection +// options; this provider configures those settings at store creation. +type CosmosDB struct { + *vectordata.Store +} + +type cosmosDBStore struct { + database *azcosmos.DatabaseClient + databaseName string + generator *embedding.BoundGenerator + partitionProperties []string + indexingMode azcosmos.IndexingMode + automatic bool +} + +// NewCosmosDB creates a store for databaseName using a caller-owned Cosmos DB +// client. The database must already exist; construction does not contact the +// service. The returned store borrows client and does not close it. +// +// Create and configure the client separately, so applications control its +// authentication and transport: +// +// client, err := azcosmos.NewClient(endpoint, credential, nil) +// if err != nil { +// return err +// } +// defer client.Close() +// +// store, err := azureprovider.NewCosmosDB(client, "vector-database", azureprovider.CosmosDBConfig{}) +// if err != nil { +// return err +// } +// collection, err := vectordata.GetCollection[string, Record](store.Store, "records", nil) +// +// Use [vectordata.Collection.EnsureExists] to create the container and its +// indexing policies. It does not enable account-level vector or full-text +// search capabilities. The config's partition and indexing settings apply to +// all collections created from this store. +func NewCosmosDB(client *azcosmos.Client, databaseName string, config CosmosDBConfig) (*CosmosDB, error) { + if client == nil { + return nil, errors.New("azurecosmosdb: client is required") + } + if strings.TrimSpace(databaseName) == "" { + return nil, errors.New("azurecosmosdb: database name is required") + } + if len(config.PartitionKeyProperties) > 3 { + return nil, errors.New("azurecosmosdb: Cosmos DB supports at most 3 hierarchical partition key properties") + } + mode := strings.ToLower(strings.TrimSpace(config.IndexingMode)) + if mode == "" { + mode = "consistent" + } + switch mode { + case "consistent", "lazy", "none": + default: + return nil, fmt.Errorf("azurecosmosdb: unsupported indexing mode %q", config.IndexingMode) + } + database, err := client.NewDatabase(databaseName) + if err != nil { + return nil, err + } + automatic := true + if config.Automatic != nil { + automatic = *config.Automatic + } + indexingMode := azcosmos.IndexingModeConsistent + switch mode { + case "lazy": + indexingMode = azcosmos.IndexingMode("Lazy") + case "none": + indexingMode = azcosmos.IndexingModeNone + } + provider := &cosmosDBStore{ + database: database, databaseName: databaseName, + generator: config.EmbeddingGenerator, + partitionProperties: append([]string(nil), config.PartitionKeyProperties...), + indexingMode: indexingMode, automatic: automatic, + } + return &CosmosDB{Store: vectordata.NewStore(provider)}, nil +} + +func (s *cosmosDBStore) OpenCollection(name string, keyType, recordType reflect.Type, definition *vectordata.Definition) (vectordata.Collection[any, any], error) { + return newCosmosCollection(s, name, keyType, recordType, definition) +} + +func (s *cosmosDBStore) ListCollectionNames(ctx context.Context) iter.Seq2[string, error] { + return func(yield func(string, error) bool) { + pager := s.database.NewQueryContainersPager("SELECT c.id FROM c", nil) + for pager.More() { + page, err := pager.NextPage(ctx) + if err != nil { + yield("", cosmosFailure(s.databaseName, "", "ListCollectionNames", err)) + return + } + for _, container := range page.Containers { + if err := ctx.Err(); err != nil { + yield("", cosmosFailure(s.databaseName, "", "ListCollectionNames", err)) + return + } + if container.ID == "" { + yield("", cosmosFailure(s.databaseName, "", "ListCollectionNames", errors.New("response contains a collection without an ID"))) + return + } + if !yield(container.ID, nil) { + return + } + } + } + } +} + +func (s *cosmosDBStore) CollectionExists(ctx context.Context, name string) (bool, error) { + if err := validateCosmosName(name); err != nil { + return false, cosmosFailure(s.databaseName, name, "CollectionExists", err) + } + container, err := s.database.NewContainer(name) + if err != nil { + return false, cosmosFailure(s.databaseName, name, "CollectionExists", err) + } + _, err = container.Read(ctx, nil) + if cosmosStatus(err, http.StatusNotFound) { + return false, nil + } + if err != nil { + return false, cosmosFailure(s.databaseName, name, "CollectionExists", err) + } + return true, nil +} + +func (s *cosmosDBStore) EnsureCollectionDeleted(ctx context.Context, name string) error { + if err := validateCosmosName(name); err != nil { + return cosmosFailure(s.databaseName, name, "DeleteCollection", err) + } + container, err := s.database.NewContainer(name) + if err != nil { + return cosmosFailure(s.databaseName, name, "DeleteCollection", err) + } + _, err = container.Delete(ctx, nil) + if err != nil && !cosmosStatus(err, http.StatusNotFound) { + return cosmosFailure(s.databaseName, name, "DeleteCollection", err) + } + return nil +} + +func validateCosmosName(name string) error { + if strings.TrimSpace(name) == "" { + return errors.New("collection name is required") + } + if strings.ContainsAny(name, "/\\?#") { + return fmt.Errorf("invalid Cosmos DB collection name %q", name) + } + return nil +} + +func cosmosFailure(database any, collection, operation string, err error) error { + databaseName, _ := database.(string) + if value, ok := database.(*azcosmos.DatabaseClient); ok && value != nil { + databaseName = value.ID() + } + return &vectordata.Error{ + SystemName: "azure.cosmosdb", StoreName: databaseName, CollectionName: collection, + OperationName: operation, Err: err, + } +} + +func cosmosStatus(err error, status int) bool { + responseErr, ok := errors.AsType[*azcore.ResponseError](err) + return ok && responseErr.StatusCode == status +} + +func cosmosID(key any) (string, error) { + var id string + switch value := key.(type) { + case uuid.UUID: + if value == uuid.Nil { + return "", errors.New("key must not be nil") + } + id = value.String() + case CosmosKey: + if strings.TrimSpace(value.DocumentID) == "" { + return "", errors.New("CosmosKey.DocumentID must not be blank") + } + id = value.DocumentID + default: + keyValue := reflect.ValueOf(key) + if !keyValue.IsValid() || keyValue.Kind() != reflect.String { + return "", fmt.Errorf("key must be a string, uuid.UUID, or CosmosKey, got %T", key) + } + id = keyValue.String() + } + if strings.TrimSpace(id) == "" { + return "", errors.New("key must not be blank") + } + return id, nil +} + +var _ vectordata.StoreProvider = (*cosmosDBStore)(nil) diff --git a/provider/azureprovider/cosmosdb_collection.go b/provider/azureprovider/cosmosdb_collection.go new file mode 100644 index 00000000..f39613d3 --- /dev/null +++ b/provider/azureprovider/cosmosdb_collection.go @@ -0,0 +1,1032 @@ +// Copyright (c) Microsoft. All rights reserved. + +package azureprovider + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "iter" + "net/http" + "reflect" + "strings" + "time" + + "github.com/Azure/azure-sdk-for-go/sdk/data/azcosmos" + "github.com/google/uuid" + "github.com/microsoft/agent-framework-go/embedding" + "github.com/microsoft/agent-framework-go/vectordata" + "github.com/microsoft/agent-framework-go/vectordata/providerservices" +) + +const cosmosScoreField = "__cosmos_score" + +type cosmosCollection struct { + store *cosmosDBStore + container *azcosmos.ContainerClient + name string + recordType reflect.Type + model *providerservices.CollectionModel + key *providerservices.KeyPropertyModel + data []*providerservices.DataPropertyModel + vectors []*providerservices.VectorPropertyModel + properties map[string]providerservices.PropertyModel + partitionProperties []providerservices.PropertyModel +} + +func newCosmosCollection(store *cosmosDBStore, name string, keyType, recordType reflect.Type, definition *vectordata.Definition) (*cosmosCollection, error) { + if err := validateCosmosName(name); err != nil { + return nil, cosmosFailure(store.databaseName, name, "Collection", err) + } + modelKeyType := keyType + if keyType == reflect.TypeFor[CosmosKey]() { + // The composite lookup key is not the record's document ID. + modelKeyType = reflect.TypeFor[any]() + } + model, err := providerservices.Build(modelKeyType, recordType, definition, providerservices.Config{ + SupportsMultipleVectors: true, + SupportsKeyType: keyTypeSupported, + SupportsDataType: dataTypeSupported, + SupportsVectorType: cosmosVectorTypeSupported, + SupportsKeyAutoGeneration: func(typ reflect.Type) bool { return typ == reflect.TypeFor[uuid.UUID]() }, + ReservedKeyStorageName: "id", + DefaultEmbeddingGenerator: store.generator, + }) + if err != nil { + return nil, cosmosFailure(store.databaseName, name, "Collection", err) + } + c := &cosmosCollection{ + store: store, name: name, recordType: recordType, model: model, + properties: make(map[string]providerservices.PropertyModel, len(model.Properties)), + } + c.container, err = store.database.NewContainer(name) + if err != nil { + return nil, cosmosFailure(store.databaseName, name, "Collection", err) + } + for _, property := range model.Properties { + if storageName := propertyStorageName(property); storageName == cosmosScoreField { + return nil, cosmosFailure(store.databaseName, name, "Collection", fmt.Errorf("storage name %q is reserved for search scores", storageName)) + } + switch p := property.(type) { + case *providerservices.KeyPropertyModel: + c.key = p + case *providerservices.DataPropertyModel: + c.data = append(c.data, p) + case *providerservices.VectorPropertyModel: + c.vectors = append(c.vectors, p) + } + c.properties[propertyName(property)] = property + } + if len(store.partitionProperties) != 0 { + seen := make(map[string]struct{}, len(store.partitionProperties)) + for _, propertyName := range store.partitionProperties { + property, found := c.properties[propertyName] + if !found { + return nil, cosmosFailure(store.databaseName, name, "Collection", fmt.Errorf("partition key property %q is not part of the record definition", propertyName)) + } + storageName := propertyStorageName(property) + if _, found := seen[storageName]; found { + return nil, cosmosFailure(store.databaseName, name, "Collection", fmt.Errorf("partition key property %q is duplicated", propertyName)) + } + seen[storageName] = struct{}{} + switch p := property.(type) { + case *providerservices.KeyPropertyModel: + c.partitionProperties = append(c.partitionProperties, p) + case *providerservices.DataPropertyModel: + if err := validatePartitionType(p.Type); err != nil { + return nil, cosmosFailure(store.databaseName, name, "Collection", fmt.Errorf("partition key property %q: %w", p.Name, err)) + } + c.partitionProperties = append(c.partitionProperties, p) + default: + return nil, cosmosFailure(store.databaseName, name, "Collection", fmt.Errorf("partition key property %q must be a key or scalar data property", propertyName)) + } + } + } + if len(c.partitionProperties) == 0 { + c.partitionProperties = []providerservices.PropertyModel{c.key} + } + if keyType != reflect.TypeFor[CosmosKey]() && keyType != reflect.TypeFor[any]() && !c.partitionedByID() { + return nil, cosmosFailure(store.databaseName, name, "Collection", errors.New("custom or hierarchical partition keys require CosmosKey as the collection lookup key type")) + } + for _, vector := range c.vectors { + if _, _, err := c.vectorPolicy(vector); err != nil { + return nil, cosmosFailure(store.database, name, "Collection", err) + } + } + for _, property := range c.data { + if property.IsFullTextIndexed { + typ := property.Type + if typ.Kind() == reflect.Pointer { + typ = typ.Elem() + } + if typ.Kind() != reflect.String { + return nil, cosmosFailure(store.databaseName, name, "Collection", fmt.Errorf("full-text indexing requires string property %q", property.Name)) + } + } + } + return c, nil +} + +func validatePartitionType(typ reflect.Type) error { + if typ.Kind() == reflect.Pointer { + typ = typ.Elem() + } + if typ == reflect.TypeFor[uuid.UUID]() { + return nil + } + switch typ.Kind() { + case reflect.String, reflect.Bool, reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64, + reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64, + reflect.Float32, reflect.Float64: + return nil + default: + return fmt.Errorf("unsupported Cosmos DB partition key type %v", typ) + } +} + +func propertyName(property providerservices.PropertyModel) string { + switch p := property.(type) { + case *providerservices.KeyPropertyModel: + return p.Name + case *providerservices.DataPropertyModel: + return p.Name + case *providerservices.VectorPropertyModel: + return p.Name + default: + return "" + } +} + +func (c *cosmosCollection) Name() string { return c.name } + +func (c *cosmosCollection) Exists(ctx context.Context) (bool, error) { + _, err := c.container.Read(ctx, nil) + if cosmosStatus(err, http.StatusNotFound) { + return false, nil + } + if err != nil { + return false, cosmosFailure(c.store.databaseName, c.name, "CollectionExists", err) + } + return true, nil +} + +func (c *cosmosCollection) EnsureExists(ctx context.Context) error { + properties, err := c.containerDefinition() + if err != nil { + return cosmosFailure(c.store.databaseName, c.name, "CreateCollection", err) + } + _, err = c.store.database.CreateContainer(ctx, properties, nil) + if err != nil && !cosmosStatus(err, http.StatusConflict) { + return cosmosFailure(c.store.databaseName, c.name, "CreateCollection", err) + } + return nil +} + +func (c *cosmosCollection) EnsureDeleted(ctx context.Context) error { + _, err := c.container.Delete(ctx, nil) + if err != nil && !cosmosStatus(err, http.StatusNotFound) { + return cosmosFailure(c.store.databaseName, c.name, "DeleteCollection", err) + } + return nil +} + +func (c *cosmosCollection) containerDefinition() (azcosmos.ContainerProperties, error) { + paths := make([]string, len(c.partitionProperties)) + for i, property := range c.partitionProperties { + paths[i] = "/" + escapeCosmosPath(propertyStorageName(property)) + } + kind := azcosmos.PartitionKeyKindHash + if len(paths) > 1 { + kind = azcosmos.PartitionKeyKindMultiHash + } + indexingPolicy := &azcosmos.IndexingPolicy{ + Automatic: c.store.automatic, IndexingMode: c.store.indexingMode, + IncludedPaths: []azcosmos.IncludedPath{}, ExcludedPaths: []azcosmos.ExcludedPath{}, + VectorIndexes: []azcosmos.VectorIndex{}, FullTextIndexes: []azcosmos.FullTextIndex{}, + } + properties := azcosmos.ContainerProperties{ + ID: c.name, + PartitionKeyDefinition: azcosmos.PartitionKeyDefinition{ + Kind: kind, Paths: paths, Version: 2, + }, + IndexingPolicy: indexingPolicy, + } + if c.store.indexingMode == azcosmos.IndexingModeNone { + return properties, nil + } + fullTextPaths := []azcosmos.FullTextPath{} + embeddings := make([]azcosmos.VectorEmbedding, 0, len(c.vectors)) + for _, property := range c.vectors { + index, embedding, err := c.vectorPolicy(property) + if err != nil { + return azcosmos.ContainerProperties{}, err + } + indexingPolicy.VectorIndexes = append(indexingPolicy.VectorIndexes, index) + embeddings = append(embeddings, embedding) + indexingPolicy.ExcludedPaths = append(indexingPolicy.ExcludedPaths, azcosmos.ExcludedPath{Path: "/" + escapeCosmosPath(property.StorageName) + "/*"}) + } + for _, property := range c.data { + if property.IsIndexed || property.IsFullTextIndexed { + indexingPolicy.IncludedPaths = append(indexingPolicy.IncludedPaths, azcosmos.IncludedPath{Path: "/" + escapeCosmosPath(property.StorageName) + "/?"}) + } + if property.IsFullTextIndexed { + path := "/" + escapeCosmosPath(property.StorageName) + indexingPolicy.FullTextIndexes = append(indexingPolicy.FullTextIndexes, azcosmos.FullTextIndex{Path: path}) + fullTextPaths = append(fullTextPaths, azcosmos.FullTextPath{Path: path, Language: "en-US"}) + } + } + indexingPolicy.IncludedPaths = append(indexingPolicy.IncludedPaths, azcosmos.IncludedPath{Path: "/*"}) + properties.VectorEmbeddingPolicy = &azcosmos.VectorEmbeddingPolicy{VectorEmbeddings: embeddings} + properties.FullTextPolicy = &azcosmos.FullTextPolicy{DefaultLanguage: "en-US", FullTextPaths: fullTextPaths} + return properties, nil +} + +func (c *cosmosCollection) vectorPolicy(property *providerservices.VectorPropertyModel) (azcosmos.VectorIndex, azcosmos.VectorEmbedding, error) { + path := "/" + escapeCosmosPath(property.StorageName) + dataType, ok := cosmosVectorDataType(property.EmbeddingType) + if !ok { + return azcosmos.VectorIndex{}, azcosmos.VectorEmbedding{}, fmt.Errorf("property %q: unsupported Cosmos DB vector data type %v", property.Name, property.EmbeddingType) + } + indexType := azcosmos.VectorIndexTypeDiskANN + switch strings.ToLower(string(property.IndexKind)) { + case "", "diskann", "disk_ann": + case "flat": + indexType = azcosmos.VectorIndexTypeFlat + case "quantizedflat", "quantized_flat": + indexType = azcosmos.VectorIndexTypeQuantizedFlat + default: + return azcosmos.VectorIndex{}, azcosmos.VectorEmbedding{}, fmt.Errorf("property %q: unsupported Cosmos DB vector index kind %q; supported kinds are disk_ann, flat, and quantized_flat", property.Name, property.IndexKind) + } + distance := azcosmos.VectorDistanceFunctionCosine + switch property.DistanceFunction { + case "", vectordata.CosineSimilarity: + case vectordata.DotProductSimilarity: + distance = azcosmos.VectorDistanceFunctionDotProduct + case vectordata.EuclideanDistance: + distance = azcosmos.VectorDistanceFunctionEuclidean + default: + return azcosmos.VectorIndex{}, azcosmos.VectorEmbedding{}, fmt.Errorf("property %q: unsupported Cosmos DB distance function %q", property.Name, property.DistanceFunction) + } + return azcosmos.VectorIndex{Path: path, Type: indexType}, + azcosmos.VectorEmbedding{Path: path, DataType: dataType, Dimensions: int32(property.Dimensions), DistanceFunction: distance}, nil +} + +func cosmosVectorDataType(typ reflect.Type) (azcosmos.VectorDataType, bool) { + element, ok := cosmosVectorElementType(typ) + if !ok { + return "", false + } + switch element { + case reflect.TypeFor[float32](): + return azcosmos.VectorDataTypeFloat32, true + case reflect.TypeFor[uint8](): + return azcosmos.VectorDataTypeUint8, true + case reflect.TypeFor[int8](): + return azcosmos.VectorDataTypeInt8, true + default: + return "", false + } +} + +func escapeCosmosPath(value string) string { + if validIdentifier(value) { + return value + } + return cosmosJSONString(value) +} + +func propertyStorageName(property providerservices.PropertyModel) string { + switch p := property.(type) { + case *providerservices.KeyPropertyModel: + return p.StorageName + case *providerservices.DataPropertyModel: + return p.StorageName + case *providerservices.VectorPropertyModel: + return p.StorageName + default: + return "" + } +} + +func (c *cosmosCollection) 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, cosmosFailure(c.store.databaseName, c.name, "Get", err) + } + id, partitionValues, err := c.partitionForKey(key) + if err != nil { + return nil, false, cosmosFailure(c.store.databaseName, c.name, "Get", err) + } + partition, err := cosmosPartitionKey(partitionValues) + if err != nil { + return nil, false, cosmosFailure(c.store.databaseName, c.name, "Get", err) + } + response, err := c.container.ReadItem(ctx, partition, id, nil) + if cosmosStatus(err, http.StatusNotFound) { + return nil, false, nil + } + if err != nil { + return nil, false, cosmosFailure(c.store.databaseName, c.name, "Get", err) + } + document, err := rawDocument(response.Value) + if err != nil { + return nil, false, cosmosFailure(c.store.databaseName, c.name, "Get", err) + } + record, err := c.decodeRecord(document, includeVectors) + if err != nil { + return nil, false, cosmosFailure(c.store.databaseName, c.name, "Get", err) + } + return record, true, nil +} + +func (c *cosmosCollection) GetFiltered(ctx context.Context, filter vectordata.Filter, top int, options *vectordata.FilteredRecordRetrievalOptions) iter.Seq2[any, error] { + return func(yield func(any, error) bool) { + var zero any + if err := ctx.Err(); err != nil { + yield(zero, cosmosFailure(c.store.databaseName, c.name, "GetFiltered", err)) + return + } + var opts vectordata.FilteredRecordRetrievalOptions + if options != nil { + opts = *options + if err := opts.Validate(); err != nil { + yield(zero, cosmosFailure(c.store.databaseName, c.name, "GetFiltered", err)) + return + } + if err := c.validateRetrieval(opts.IncludeVectors); err != nil { + yield(zero, cosmosFailure(c.store.databaseName, c.name, "GetFiltered", err)) + return + } + } + if top < 0 { + yield(zero, cosmosFailure(c.store.databaseName, c.name, "GetFiltered", errors.New("top must be nonnegative"))) + return + } + if top == 0 { + return + } + filterSQL, parameters, err := c.translateFilter(&filter) + if err != nil { + yield(zero, cosmosFailure(c.store.databaseName, c.name, "GetFiltered", err)) + return + } + query := "SELECT " + strings.Join(c.selectFields(opts.IncludeVectors), ", ") + " FROM c" + if filterSQL != "" { + query += " WHERE " + filterSQL + } + orderBy, err := c.orderBy(opts.OrderBy) + if err != nil { + yield(zero, cosmosFailure(c.store.databaseName, c.name, "GetFiltered", err)) + return + } + partition, err := c.queryPartition(&filter, orderBy != "") + if err != nil { + yield(zero, cosmosFailure(c.store.databaseName, c.name, "GetFiltered", err)) + return + } + if orderBy != "" { + query += " ORDER BY " + orderBy + } + skip, remaining := opts.Skip, top + err = c.eachQueryDocument(ctx, query, parameters, partition, func(raw json.RawMessage) bool { + if skip > 0 { + skip-- + return true + } + document, err := rawDocument(raw) + if err != nil { + yield(zero, cosmosFailure(c.store.databaseName, c.name, "GetFiltered", err)) + return false + } + record, err := c.decodeRecord(document, opts.IncludeVectors) + if err != nil { + yield(zero, cosmosFailure(c.store.databaseName, c.name, "GetFiltered", err)) + return false + } + remaining-- + return yield(record, nil) && remaining > 0 + }) + if err != nil { + yield(zero, cosmosFailure(c.store.databaseName, c.name, "GetFiltered", err)) + } + } +} + +func (c *cosmosCollection) Delete(ctx context.Context, key any) error { + id, partitionValues, err := c.partitionForKey(key) + if err != nil { + return cosmosFailure(c.store.databaseName, c.name, "Delete", err) + } + partition, err := cosmosPartitionKey(partitionValues) + if err != nil { + return cosmosFailure(c.store.databaseName, c.name, "Delete", err) + } + _, err = c.container.DeleteItem(ctx, partition, id, nil) + if err != nil && !cosmosStatus(err, http.StatusNotFound) { + return cosmosFailure(c.store.databaseName, c.name, "Delete", err) + } + return nil +} + +func (c *cosmosCollection) Upsert(ctx context.Context, records ...any) error { + if len(records) == 0 { + return nil + } + owned := make([]any, len(records)) + for i, record := range records { + prepared, err := c.prepareRecord(record) + if err != nil { + return cosmosFailure(c.store.databaseName, c.name, "Upsert", fmt.Errorf("record %d: %w", i, err)) + } + owned[i] = prepared + } + generated := make(map[string][]any) + for _, property := range c.vectors { + if property.EmbeddingType == property.Type { + continue + } + result, err := providerservices.GenerateEmbeddings(ctx, property, owned) + if err != nil { + return cosmosFailure(c.store.databaseName, c.name, "Upsert", fmt.Errorf("property %q: %w", property.Name, err)) + } + vectors := make([]any, len(result.Embeddings)) + for i, value := range result.Embeddings { + vector, dimensions, ok := cosmosVectorSlice(value, property.EmbeddingType) + if !ok || dimensions != property.Dimensions { + return cosmosFailure(c.store.databaseName, c.name, "Upsert", fmt.Errorf("property %q: generated embedding %d must have dimension %d and a supported Cosmos DB vector type", property.Name, i, property.Dimensions)) + } + vectors[i] = vector + } + generated[property.Name] = vectors + } + for i, record := range owned { + if err := ctx.Err(); err != nil { + return cosmosFailure(c.store.databaseName, c.name, "Upsert", err) + } + document, err := c.encodeRecord(record, i, generated) + if err != nil { + return cosmosFailure(c.store.databaseName, c.name, "Upsert", fmt.Errorf("record %d: %w", i, err)) + } + id, _ := document["id"].(string) + partition, err := c.partitionForRecord(record, id) + if err != nil { + return cosmosFailure(c.store.databaseName, c.name, "Upsert", fmt.Errorf("record %d: %w", i, err)) + } + partitionKey, err := cosmosPartitionKey(partition) + if err == nil { + var body []byte + body, err = json.Marshal(document) + if err == nil { + _, err = c.container.UpsertItem(ctx, partitionKey, body, nil) + } + } + if err != nil { + return cosmosFailure(c.store.databaseName, c.name, "Upsert", fmt.Errorf("record %d: %w", i, err)) + } + } + return nil +} + +func (c *cosmosCollection) GetBatch(ctx context.Context, keys []any, options *vectordata.RecordRetrievalOptions) iter.Seq2[any, error] { + return func(yield func(any, error) bool) { + if err := c.validateRetrieval(options != nil && options.IncludeVectors); err != nil { + yield(nil, cosmosFailure(c.store.databaseName, c.name, "GetBatch", err)) + return + } + for _, key := range keys { + if err := ctx.Err(); err != nil { + yield(nil, cosmosFailure(c.store.databaseName, c.name, "GetBatch", err)) + return + } + record, found, err := c.Get(ctx, key, options) + if err != nil { + yield(nil, cosmosFailure(c.store.databaseName, c.name, "GetBatch", err)) + return + } + if found && !yield(record, nil) { + return + } + } + } +} + +func (c *cosmosCollection) DeleteBatch(ctx context.Context, keys []any) error { + for i, key := range keys { + if err := ctx.Err(); err != nil { + return cosmosFailure(c.store.databaseName, c.name, "DeleteBatch", err) + } + if err := c.Delete(ctx, key); err != nil { + return cosmosFailure(c.store.databaseName, c.name, "DeleteBatch", fmt.Errorf("key %d: %w", i, err)) + } + } + return nil +} + +func (c *cosmosCollection) 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.Kind() == reflect.Struct { + owned := reflect.New(c.recordType) + owned.Elem().Set(value) + return owned.Interface(), nil + } + if value.IsNil() { + return nil, errors.New("record must not be nil") + } + if value.Kind() == reflect.Pointer && value.Elem().Kind() == reflect.Map && value.Elem().IsNil() { + return nil, errors.New("record map must not be nil") + } + return value.Interface(), nil +} + +func (c *cosmosCollection) encodeRecord(record any, index int, generated map[string][]any) (map[string]any, error) { + keyValue, found, err := c.key.Value(record) + if err != nil { + return nil, err + } + if c.key.IsAutoGenerated && (!found || keyValue == nil || keyValue == uuid.Nil) { + keyValue, err = uuid.NewRandom() + if err != nil { + return nil, err + } + if err := c.key.SetValue(record, keyValue); err != nil { + return nil, err + } + } + id, err := cosmosID(keyValue) + 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["id"] = id + for _, property := range c.data { + value, found, err := property.Value(record) + if err != nil { + return nil, err + } + if found { + document[property.StorageName], err = cosmosDataJSON(value) + if err != nil { + return nil, fmt.Errorf("property %q: %w", property.Name, err) + } + } + } + for _, property := range c.vectors { + value, found, err := property.Value(record) + if err != nil { + return nil, err + } + if !found { + continue + } + if outputs, ok := generated[property.Name]; ok { + value = outputs[index] + } + vector, dimensions, ok := cosmosVectorSlice(value, property.EmbeddingType) + if !ok { + return nil, fmt.Errorf("property %q requires a supported Cosmos DB vector type", property.Name) + } + if vector != nil && dimensions != property.Dimensions { + return nil, fmt.Errorf("property %q requires %d dimensions, got %d", property.Name, property.Dimensions, dimensions) + } + document[property.StorageName] = cosmosVectorJSON(vector) + } + return document, nil +} + +func cosmosDataJSON(value any) (any, error) { + reflected := reflect.ValueOf(value) + for reflected.IsValid() && reflected.Kind() == reflect.Pointer { + if reflected.IsNil() { + return nil, nil + } + reflected = reflected.Elem() + } + if !reflected.IsValid() { + return nil, nil + } + if instant, ok := reflected.Interface().(time.Time); ok { + instant = instant.UTC() + if _, err := instant.MarshalJSON(); err != nil { + return nil, err + } + return instant.Format("2006-01-02T15:04:05.000000000Z"), nil + } + if reflected.Kind() == reflect.Slice || reflected.Kind() == reflect.Array { + element := reflected.Type().Elem() + if element == reflect.TypeFor[time.Time]() || element.Kind() == reflect.Interface { + if reflected.Kind() == reflect.Slice && reflected.IsNil() { + return nil, nil + } + values := make([]any, reflected.Len()) + for i := range reflected.Len() { + var err error + values[i], err = cosmosDataJSON(reflected.Index(i).Interface()) + if err != nil { + return nil, fmt.Errorf("element %d: %w", i, err) + } + } + return values, nil + } + } + return value, nil +} + +func (c *cosmosCollection) decodeRecord(document map[string]json.RawMessage, includeVectors bool) (any, error) { + if document == nil { + return nil, errors.New("response document is missing") + } + 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)) + } + rawID, found := document["id"] + if !found || bytes.Equal(bytes.TrimSpace(rawID), []byte("null")) { + return nil, errors.New("response is missing document id") + } + var id string + if err := json.Unmarshal(rawID, &id); err != nil || strings.TrimSpace(id) == "" { + if err == nil { + err = errors.New("document id must not be blank") + } + return nil, fmt.Errorf("property %q: %w", c.key.Name, err) + } + keyValue, err := c.decodedKey(id) + if err != nil { + return nil, err + } + if err := c.key.SetValue(target, keyValue); err != nil { + return nil, err + } + for _, property := range c.data { + if raw, found := document[property.StorageName]; found { + value, err := decodeCosmosValue(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 { + if raw, found := document[property.StorageName]; found { + value, err := decodeCosmosVector(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 (c *cosmosCollection) decodedKey(id string) (any, error) { + switch c.key.Type { + case reflect.TypeFor[uuid.UUID](): + value, err := uuid.Parse(id) + if err != nil { + return nil, fmt.Errorf("property %q: invalid UUID document ID: %w", c.key.Name, err) + } + return value, nil + default: + return reflect.ValueOf(id).Convert(c.key.Type).Interface(), nil + } +} + +func decodeCosmosValue(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 decodeCosmosVector(raw json.RawMessage, property *providerservices.VectorPropertyModel) (any, error) { + if bytes.Equal(bytes.TrimSpace(raw), []byte("null")) { + if property.Type.Kind() == reflect.Pointer { + return nil, nil + } + if isCosmosEmbeddingType(property.Type) { + return reflect.Zero(property.Type).Interface(), nil + } + return nil, nil + } + if property.EmbeddingType != property.Type { + return nil, errors.New("generated vector cannot be assigned to its non-vector input property") + } + if _, ok := cosmosVectorElementType(property.Type); !ok { + return nil, fmt.Errorf("unsupported vector type %v", property.Type) + } + var decoded any + typ := property.Type + embeddingType := isCosmosEmbeddingType(typ) + if typ.Kind() == reflect.Pointer { + typ = typ.Elem() + } + if embeddingType { + value := reflect.New(typ) + vector := value.Elem().FieldByName("Vector") + if !vector.IsValid() || !vector.CanAddr() { + return nil, fmt.Errorf("unsupported embedding type %v", property.Type) + } + if err := json.Unmarshal(raw, vector.Addr().Interface()); err != nil { + return nil, err + } + if property.Type.Kind() == reflect.Pointer { + decoded = value.Interface() + } else { + decoded = value.Elem().Interface() + } + } else { + value := reflect.New(typ) + if err := json.Unmarshal(raw, value.Interface()); err != nil { + return nil, err + } + decoded = value.Elem().Interface() + } + _, dimensions, ok := cosmosVectorSlice(decoded, property.Type) + if !ok || dimensions != property.Dimensions { + return nil, fmt.Errorf("expected a supported vector of dimension %d", property.Dimensions) + } + return decoded, nil +} + +func cosmosVectorTypeSupported(typ reflect.Type) bool { + return providerservices.SupportsVector[float32](typ) || + providerservices.SupportsVector[uint8](typ) || + providerservices.SupportsVector[int8](typ) +} + +func cosmosVectorElementType(typ reflect.Type) (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 reflect.TypeFor[float32](), true + case reflect.TypeFor[[]uint8](), reflect.TypeFor[embedding.VectorOf[uint8]](), + reflect.TypeFor[embedding.Of[uint8]](), reflect.TypeFor[*embedding.Of[uint8]](): + return reflect.TypeFor[uint8](), true + case reflect.TypeFor[[]int8](), reflect.TypeFor[embedding.VectorOf[int8]](), + reflect.TypeFor[embedding.Of[int8]](), reflect.TypeFor[*embedding.Of[int8]](): + return reflect.TypeFor[int8](), true + default: + return nil, false + } +} + +func isCosmosEmbeddingType(typ reflect.Type) bool { + for typ.Kind() == reflect.Pointer { + typ = typ.Elem() + } + return typ == reflect.TypeFor[embedding.Of[float32]]() || + typ == reflect.TypeFor[embedding.Of[uint8]]() || + typ == reflect.TypeFor[embedding.Of[int8]]() +} + +func cosmosVectorSlice(value any, typ reflect.Type) (any, int, bool) { + if value == nil { + return nil, 0, true + } + switch elementType, ok := cosmosVectorElementType(typ); { + case ok && elementType == reflect.TypeFor[float32](): + vector, supported := providerservices.AsVector[float32](value) + if vector == nil { + return nil, 0, supported + } + return vector, len(vector), supported + case ok && elementType == reflect.TypeFor[uint8](): + vector, supported := providerservices.AsVector[uint8](value) + if vector == nil { + return nil, 0, supported + } + return vector, len(vector), supported + case ok && elementType == reflect.TypeFor[int8](): + vector, supported := providerservices.AsVector[int8](value) + if vector == nil { + return nil, 0, supported + } + return vector, len(vector), supported + default: + return nil, 0, false + } +} + +func cosmosVectorJSON(vector any) any { + if bytes, ok := vector.([]uint8); ok { + values := make([]int, len(bytes)) + for i, value := range bytes { + values[i] = int(value) + } + return values + } + return vector +} + +func (c *cosmosCollection) partitionedByID() bool { + return len(c.partitionProperties) == 1 && c.partitionProperties[0] == c.key +} + +func (c *cosmosCollection) partitionForRecord(record any, id string) ([]any, error) { + if c.partitionedByID() { + return []any{id}, nil + } + values := make([]any, 0, len(c.partitionProperties)) + for _, property := range c.partitionProperties { + switch property.(type) { + case *providerservices.KeyPropertyModel: + values = append(values, id) + default: + value, found, err := property.Value(record) + if err != nil { + return nil, err + } + if !found { + return nil, fmt.Errorf("partition key property %q is missing", propertyName(property)) + } + normalized, err := normalizePartitionValue(value) + if err != nil { + return nil, fmt.Errorf("partition key property %q: %w", propertyName(property), err) + } + values = append(values, normalized) + } + } + return values, nil +} + +func normalizePartitionValue(value any) (any, error) { + if value == nil { + return nil, nil + } + valueOf := reflect.ValueOf(value) + for valueOf.Kind() == reflect.Pointer { + if valueOf.IsNil() { + return nil, nil + } + valueOf = valueOf.Elem() + } + value = valueOf.Interface() + if id, ok := value.(uuid.UUID); ok { + return id.String(), nil + } + switch valueOf.Kind() { + case reflect.String, reflect.Bool, + reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64, + reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64, + reflect.Float32, reflect.Float64: + default: + return nil, fmt.Errorf("unsupported partition key value type %T", value) + } + // Preserve the document's JSON number instead of widening float32 first. + encoded, err := json.Marshal(value) + if err != nil { + return nil, err + } + var normalized any + if err := json.Unmarshal(encoded, &normalized); err != nil { + return nil, err + } + switch normalized.(type) { + case nil, string, bool, float64: + return normalized, nil + default: + return nil, fmt.Errorf("partition key value of type %T must serialize to a JSON scalar", value) + } +} + +func (c *cosmosCollection) partitionForKey(key any) (string, []any, error) { + id, err := cosmosID(key) + if err != nil { + return "", nil, err + } + cosmosKey, composite := key.(CosmosKey) + if !composite { + if c.partitionedByID() { + return id, []any{id}, nil + } + return "", nil, errors.New("custom or hierarchical partition keys require a CosmosKey lookup key") + } + parts := cosmosKey.PartitionKey + if len(parts) != len(c.partitionProperties) { + return "", nil, fmt.Errorf("CosmosKey requires %d partition key values, got %d", len(c.partitionProperties), len(parts)) + } + values := make([]any, len(parts)) + for i, part := range parts { + normalized, err := normalizePartitionValue(part) + if err != nil { + return "", nil, fmt.Errorf("CosmosKey partition key value %d: %w", i, err) + } + values[i] = normalized + } + return id, values, nil +} + +func cosmosPartitionKey(values []any) (azcosmos.PartitionKey, error) { + key := azcosmos.NewPartitionKey() + for _, value := range values { + switch v := value.(type) { + case nil: + key = key.AppendNull() + case string: + key = key.AppendString(v) + case bool: + key = key.AppendBool(v) + case float64: + key = key.AppendNumber(v) + default: + return azcosmos.PartitionKey{}, fmt.Errorf("unsupported normalized partition key value %T", value) + } + } + return key, nil +} + +func (c *cosmosCollection) vectorForSearch(ctx context.Context, property *providerservices.VectorPropertyModel, value any) (any, error) { + vector, dimensions, ok := cosmosVectorSlice(value, property.EmbeddingType) + if !ok { + generator := property.SearchEmbeddingGenerator + if generator == nil { + return nil, fmt.Errorf("property %q does not have a generator for search value type %T", property.Name, value) + } + result, err := generator.Generate(ctx, []any{value}, nil) + if err != nil { + return nil, err + } + if result == nil || len(result.Embeddings) != 1 { + return nil, fmt.Errorf("property %q: search generator must return exactly one embedding", property.Name) + } + vector, dimensions, ok = cosmosVectorSlice(result.Embeddings[0], property.EmbeddingType) + if !ok { + return nil, fmt.Errorf("property %q: search generator must return a supported Cosmos DB vector type", property.Name) + } + } + if dimensions != property.Dimensions { + return nil, fmt.Errorf("property %q: search vector must have dimension %d, got %d", property.Name, property.Dimensions, dimensions) + } + return cosmosVectorJSON(vector), nil +} + +func (c *cosmosCollection) validateRetrieval(includeVectors bool) error { + if includeVectors && c.model.EmbeddingGenerationRequired { + return errors.New("IncludeVectors is unsupported for collections with generated vector fields") + } + return nil +} + +func (c *cosmosCollection) eachQueryDocument(ctx context.Context, query string, parameters []cosmosParameter, partition azcosmos.PartitionKey, yield func(json.RawMessage) bool) error { + queryParameters := make([]azcosmos.QueryParameter, len(parameters)) + for i, parameter := range parameters { + queryParameters[i] = azcosmos.QueryParameter{Name: parameter.Name, Value: parameter.Value} + } + pager := c.container.NewQueryItemsPager(query, partition, &azcosmos.QueryOptions{ + QueryParameters: queryParameters, + }) + for pager.More() { + if err := ctx.Err(); err != nil { + return err + } + page, err := pager.NextPage(ctx) + if err != nil { + return err + } + for _, raw := range page.Items { + if err := ctx.Err(); err != nil { + return err + } + if !yield(raw) { + return nil + } + } + } + return nil +} + +func rawDocument(raw json.RawMessage) (map[string]json.RawMessage, error) { + var document map[string]json.RawMessage + if err := json.Unmarshal(raw, &document); err != nil { + return nil, err + } + return document, nil +} + +var ( + _ vectordata.Collection[any, any] = (*cosmosCollection)(nil) + _ vectordata.KeywordHybridSearcher[any] = (*cosmosCollection)(nil) +) diff --git a/provider/azureprovider/cosmosdb_e2e_test.go b/provider/azureprovider/cosmosdb_e2e_test.go new file mode 100644 index 00000000..3701ccf6 --- /dev/null +++ b/provider/azureprovider/cosmosdb_e2e_test.go @@ -0,0 +1,442 @@ +// Copyright (c) Microsoft. All rights reserved. + +package azureprovider_test + +import ( + "cmp" + "context" + "errors" + "fmt" + "math" + "os" + "reflect" + "slices" + "strings" + "testing" + "time" + + "github.com/Azure/azure-sdk-for-go/sdk/azidentity" + "github.com/Azure/azure-sdk-for-go/sdk/data/azcosmos" + "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/provider/foundryprovider" + "github.com/microsoft/agent-framework-go/vectordata" +) + +// Live tests require RUN_AZURE_COSMOS_INTEGRATION_TESTS=true, +// AZURE_COSMOS_ENDPOINT, AZURE_COSMOS_DATABASE, and an Azure CLI login. +// AZURE_SUBSCRIPTION_ID selects the CLI subscription. +// The account must support vector and full-text search, and the database must +// already exist. The principal needs Cosmos DB data contributor access. +// Only UUID-named test containers are created and deleted. +// +// Generated-vector tests also require FOUNDRY_PROJECT_ENDPOINT and access to +// FOUNDRY_EMBEDDING_MODEL, which defaults to text-embedding-3-large (3072 dimensions). +func newLiveCosmosDB(t *testing.T, config azureprovider.CosmosDBConfig, generateEmbeddings bool) *azureprovider.CosmosDB { + t.Helper() + if !strings.EqualFold(os.Getenv("RUN_AZURE_COSMOS_INTEGRATION_TESTS"), "true") { + t.Skip("RUN_AZURE_COSMOS_INTEGRATION_TESTS is not true") + } + endpoint := strings.TrimSpace(os.Getenv("AZURE_COSMOS_ENDPOINT")) + database := strings.TrimSpace(os.Getenv("AZURE_COSMOS_DATABASE")) + if endpoint == "" || database == "" { + t.Fatal("AZURE_COSMOS_ENDPOINT and AZURE_COSMOS_DATABASE are required") + } + credential, err := azidentity.NewAzureCLICredential(&azidentity.AzureCLICredentialOptions{ + Subscription: os.Getenv("AZURE_SUBSCRIPTION_ID"), + }) + if err != nil { + t.Fatal(err) + } + if generateEmbeddings { + projectEndpoint := strings.TrimSpace(os.Getenv("FOUNDRY_PROJECT_ENDPOINT")) + if projectEndpoint == "" { + t.Skip("FOUNDRY_PROJECT_ENDPOINT is required for generated-vector tests") + } + config.EmbeddingGenerator = embedding.Bind(foundryprovider.NewEmbeddingGenerator(projectEndpoint, credential, + foundryprovider.EmbeddingGeneratorConfig{ + Model: cmp.Or(os.Getenv("FOUNDRY_EMBEDDING_MODEL"), "text-embedding-3-large"), + })) + } + client, err := azcosmos.NewClient(endpoint, credential, nil) + if err != nil { + t.Fatal(err) + } + t.Cleanup(client.Close) + store, err := azureprovider.NewCosmosDB(client, database, config) + if err != nil { + t.Fatal(err) + } + return store +} + +func liveCosmosContainer(t *testing.T, store *azureprovider.CosmosDB) string { + t.Helper() + name := "afgo-live-" + uuid.NewString() + t.Logf("Temporary Cosmos DB container: %s", name) + t.Cleanup(func() { + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + if err := store.EnsureCollectionDeleted(ctx, name); err != nil { + t.Errorf("delete temporary container %s: %v", name, err) + } + }) + return name +} + +type liveCosmosDocument struct { + ID uuid.UUID `vectordata:"key"` + Text string `vectordata:"data,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,index=flat,storage=vector"` + Bytes []uint8 `vectordata:"vector,dimensions=3,index=quantized_flat,storage=bytes"` + Signed []int8 `vectordata:"vector,dimensions=3,index=disk_ann,storage=signed"` +} + +func TestE2E_CosmosDBNativeCRUDAndSearch(t *testing.T) { + store := newLiveCosmosDB(t, azureprovider.CosmosDBConfig{PartitionKeyProperties: []string{"Bucket"}}, false) + name := liveCosmosContainer(t, store) + collection, err := vectordata.GetCollection[azureprovider.CosmosKey, *liveCosmosDocument](store.Store, name, nil) + if err != nil { + t.Fatal(err) + } + if exists, err := collection.Exists(t.Context()); err != nil || exists { + t.Fatalf("absent container exists = (%v, %v)", exists, err) + } + for range 2 { + if err := collection.EnsureExists(t.Context()); err != nil { + t.Fatal(err) + } + } + if exists, err := collection.Exists(t.Context()); err != nil || !exists { + t.Fatalf("created container exists = (%v, %v)", exists, err) + } + if !slices.Contains(collect(t, store.ListCollectionNames(t.Context())), name) { + t.Fatalf("created container %s was not listed", name) + } + now := time.Now().UTC().Truncate(time.Millisecond) + records := []*liveCosmosDocument{ + {Text: "A pirate joke about a ship.", Bucket: "alpha", Rank: 1, When: now, Vector: []float32{1, 0, 0}, Bytes: []uint8{255, 0, 0}, Signed: []int8{127, 0, 0}}, + {Text: "A pirate joke about a parrot.", Bucket: "alpha", Rank: 2, When: now, Vector: []float32{0.8, 0.2, 0}, Bytes: []uint8{200, 50, 0}, Signed: []int8{100, 25, 0}}, + {Text: "A private joke from another bucket.", Bucket: "beta", Rank: 3, When: now, Vector: []float32{1, 0, 0}, Bytes: []uint8{255, 0, 0}, Signed: []int8{127, 0, 0}}, + } + if err := collection.Upsert(t.Context(), records...); err != nil { + t.Fatal(err) + } + key := func(record *liveCosmosDocument) azureprovider.CosmosKey { + return azureprovider.CosmosKey{DocumentID: record.ID.String(), PartitionKey: []any{record.Bucket}} + } + if records[0].ID == uuid.Nil || records[0].ID == records[1].ID { + t.Fatalf("generated record IDs = %s, %s", records[0].ID, records[1].ID) + } + got, found, err := collection.Get(t.Context(), key(records[0]), &vectordata.RecordRetrievalOptions{IncludeVectors: true}) + if err != nil || !found || !reflect.DeepEqual(got, records[0]) { + t.Fatalf("native retrieval = (%+v, %v, %v)", got, found, err) + } + wrongPartition := azureprovider.CosmosKey{DocumentID: records[0].ID.String(), PartitionKey: []any{"beta"}} + if _, found, err := collection.Get(t.Context(), wrongPartition, nil); err != nil || found { + t.Fatalf("wrong-partition retrieval = (%v, %v)", found, err) + } + batch := collect(t, vectordata.GetBatch(t.Context(), collection, []azureprovider.CosmosKey{ + key(records[0]), wrongPartition, key(records[1]), + }, nil)) + if len(batch) != 2 || batch[0].ID != records[0].ID || batch[1].ID != records[1].ID || + batch[0].Vector != nil || batch[0].Bytes != nil || batch[0].Signed != nil { + t.Fatalf("batch retrieval = %+v", batch) + } + records[0].Text = "An updated pirate joke about a ship." + if err := collection.Upsert(t.Context(), records[0]); err != nil { + t.Fatal(err) + } + got, found, err = collection.Get(t.Context(), key(records[0]), nil) + if err != nil || !found || got.Text != records[0].Text { + t.Fatalf("updated retrieval = (%+v, %v, %v)", got, found, err) + } + filter := vectordata.Filter{Operator: vectordata.Equal, Property: "Bucket", Value: "alpha"} + for _, vector := range []struct { + property string + value any + }{ + {"Vector", []float32{1, 0, 0}}, + {"Bytes", []uint8{255, 0, 0}}, + {"Signed", []int8{127, 0, 0}}, + } { + t.Run(vector.property, func(t *testing.T) { + results := collect(t, collection.Search(t.Context(), vector.value, 3, + &vectordata.SearchOptions{Filter: &filter, VectorProperty: vector.property, IncludeVectors: true})) + if len(results) != 2 || results[0].Record.ID != records[0].ID || results[0].Score == nil || + !reflect.DeepEqual(results[0].Record, records[0]) { + t.Fatalf("scoped %s search = %+v", vector.property, results) + } + if score := *results[0].Score; math.IsNaN(score) || math.IsInf(score, 0) { + t.Fatalf("invalid search score %v", score) + } + thresholded := collect(t, collection.Search(t.Context(), vector.value, 3, + &vectordata.SearchOptions{Filter: &filter, VectorProperty: vector.property, ScoreThreshold: new(0.99)})) + if len(thresholded) != 1 || thresholded[0].Record.ID != records[0].ID || + thresholded[0].Score == nil || *thresholded[0].Score < 0.99 { + t.Fatalf("score-threshold search = %+v", thresholded) + } + }) + } + results := collect(t, collection.Search(t.Context(), []float32{1, 0, 0}, 1, + &vectordata.SearchOptions{Filter: &filter, VectorProperty: "Vector", Skip: 1})) + if len(results) != 1 || results[0].Record.ID != records[1].ID || results[0].Record.Vector != nil { + t.Fatalf("vector Skip/top = %+v", results) + } + hybrid := collection.(vectordata.KeywordHybridSearcher[*liveCosmosDocument]) + hybridResults := collect(t, hybrid.HybridSearch(t.Context(), []float32{1, 0, 0}, []string{"pirate", "ship"}, 2, + &vectordata.HybridSearchOptions{Filter: &filter, VectorProperty: "Vector", IncludeVectors: true})) + if len(hybridResults) != 2 || hybridResults[0].Record.ID != records[0].ID || hybridResults[0].Score == nil || + !reflect.DeepEqual(hybridResults[0].Record.Vector, records[0].Vector) { + t.Fatalf("hybrid search = %+v", hybridResults) + } + ordered := collect(t, collection.GetFiltered(t.Context(), filter, 1, + &vectordata.FilteredRecordRetrievalOptions{Skip: 1, OrderBy: []vectordata.Sort{{Property: "Rank", Descending: true}}})) + if len(ordered) != 1 || ordered[0].ID != records[0].ID { + t.Fatalf("ordered Skip/top = %+v", ordered) + } + membership := collect(t, collection.GetFiltered(t.Context(), + vectordata.Filter{Operator: vectordata.In, Property: "ID", Value: []azureprovider.CosmosKey{key(records[0]), key(records[2])}}, 3, nil)) + if len(membership) != 2 { + t.Fatalf("composite-key membership = %+v", membership) + } + if err := vectordata.DeleteBatch(t.Context(), collection, []azureprovider.CosmosKey{key(records[0]), key(records[1])}); err != nil { + t.Fatal(err) + } + if _, found, err := collection.Get(t.Context(), key(records[0]), nil); err != nil || found { + t.Fatalf("deleted document retrieval = (%v, %v)", found, err) + } + if err := collection.Delete(t.Context(), key(records[0])); err != nil { + t.Fatalf("idempotent delete: %v", err) + } + for range 2 { + if err := collection.EnsureDeleted(t.Context()); err != nil { + t.Fatal(err) + } + } + if exists, err := collection.Exists(t.Context()); err != nil || exists { + t.Fatalf("deleted container exists = (%v, %v)", exists, err) + } +} + +func TestE2E_CosmosDBHierarchicalNullablePartitions(t *testing.T) { + type record struct { + ID string `vectordata:"key"` + Tenant *float32 `vectordata:"data,indexed,storage=tenant"` + Region *string `vectordata:"data,indexed,storage=region"` + Active *bool `vectordata:"data,indexed,storage=active"` + } + store := newLiveCosmosDB(t, azureprovider.CosmosDBConfig{PartitionKeyProperties: []string{"Tenant", "Region", "Active"}}, false) + name := liveCosmosContainer(t, store) + collection, err := vectordata.GetCollection[azureprovider.CosmosKey, *record](store.Store, name, nil) + if err != nil { + t.Fatal(err) + } + if err := collection.EnsureExists(t.Context()); err != nil { + t.Fatal(err) + } + records := []*record{ + {ID: "numeric", Tenant: new(float32(0.1)), Region: new("north"), Active: new(true)}, + {ID: "null"}, + } + if err := collection.Upsert(t.Context(), records...); err != nil { + t.Fatal(err) + } + dynamic, err := store.DynamicCollection(name, &vectordata.Definition{Properties: []vectordata.Property{ + &vectordata.KeyProperty{Name: "ID", Type: reflect.TypeFor[string]()}, + &vectordata.DataProperty{Name: "Tenant", StorageName: "tenant", Type: reflect.TypeFor[*float32](), IsIndexed: true}, + &vectordata.DataProperty{Name: "Region", StorageName: "region", Type: reflect.TypeFor[*string](), IsIndexed: true}, + &vectordata.DataProperty{Name: "Active", StorageName: "active", Type: reflect.TypeFor[*bool](), IsIndexed: true}, + }}) + if err != nil { + t.Fatal(err) + } + for _, want := range records { + t.Run(want.ID, func(t *testing.T) { + key := azureprovider.CosmosKey{DocumentID: want.ID, PartitionKey: []any{want.Tenant, want.Region, want.Active}} + got, found, err := collection.Get(t.Context(), key, nil) + if err != nil || !found || !reflect.DeepEqual(got, want) { + t.Fatalf("hierarchical retrieval = (%+v, %v, %v), want %+v", got, found, err, want) + } + mapped, found, err := dynamic.Get(t.Context(), key, nil) + var wantTenant any + if want.Tenant != nil { + wantTenant = want.Tenant + } + if err != nil || !found || mapped["ID"] != want.ID || !reflect.DeepEqual(mapped["Tenant"], wantTenant) { + t.Fatalf("dynamic retrieval = (%#v, %v, %v)", mapped, found, err) + } + if err := dynamic.Upsert(t.Context(), mapped); err != nil { + t.Fatal(err) + } + filter := vectordata.Filter{Operator: vectordata.And, Operands: []vectordata.Filter{ + {Operator: vectordata.Equal, Property: "Tenant", Value: want.Tenant}, + {Operator: vectordata.Equal, Property: "Region", Value: want.Region}, + {Operator: vectordata.Equal, Property: "Active", Value: want.Active}, + }} + matches := collect(t, collection.GetFiltered(t.Context(), filter, 2, + &vectordata.FilteredRecordRetrievalOptions{OrderBy: []vectordata.Sort{{Property: "ID"}}})) + if len(matches) != 1 || matches[0].ID != want.ID { + t.Fatalf("hierarchical filtered retrieval = %+v", matches) + } + if err := dynamic.Delete(t.Context(), key); err != nil { + t.Fatal(err) + } + }) + } +} + +func TestE2E_CosmosDBCrossPartitionPagingAndNamedKeys(t *testing.T) { + type key string + type record struct { + ID key `vectordata:"key"` + Bucket string `vectordata:"data,indexed,storage=bucket"` + Vector []float32 `vectordata:"vector,dimensions=3,index=flat,storage=vector"` + } + store := newLiveCosmosDB(t, azureprovider.CosmosDBConfig{}, false) + name := liveCosmosContainer(t, store) + collection, err := vectordata.GetCollection[key, *record](store.Store, name, nil) + if err != nil { + t.Fatal(err) + } + if err := collection.EnsureExists(t.Context()); err != nil { + t.Fatal(err) + } + const count = 125 + records := make([]*record, count) + for i := range records { + records[i] = &record{ID: key(fmt.Sprintf("item-%03d", i)), Bucket: "alpha"} + } + if err := collection.Upsert(t.Context(), records...); err != nil { + t.Fatal(err) + } + got, found, err := collection.Get(t.Context(), records[0].ID, &vectordata.RecordRetrievalOptions{IncludeVectors: true}) + if err != nil || !found || got.ID != records[0].ID || got.Vector != nil { + t.Fatalf("named key and nil vector retrieval = (%+v, %v, %v)", got, found, err) + } + filter := vectordata.Filter{Operator: vectordata.Equal, Property: "Bucket", Value: "alpha"} + matches := collect(t, collection.GetFiltered(t.Context(), filter, count+1, nil)) + seen := make(map[key]bool, count) + for _, match := range matches { + if match.Bucket != "alpha" || seen[match.ID] { + t.Fatalf("unexpected or duplicate result: %+v", match) + } + seen[match.ID] = true + } + if len(seen) != count { + t.Fatalf("cross-partition retrieval returned %d distinct records, want %d", len(seen), count) + } + for _, record := range records { + if !seen[record.ID] { + t.Fatalf("cross-partition retrieval omitted %s", record.ID) + } + } + page := collect(t, collection.GetFiltered(t.Context(), filter, 10, &vectordata.FilteredRecordRetrievalOptions{Skip: 110})) + if len(page) != 10 { + t.Fatalf("cross-partition Skip/top returned %d records, want 10", len(page)) + } + exhausted := collect(t, collection.GetFiltered(t.Context(), filter, 1, &vectordata.FilteredRecordRetrievalOptions{Skip: count})) + if len(exhausted) != 0 { + t.Fatalf("Skip past final record = %+v", exhausted) + } + yielded := 0 + for _, err := range collection.GetFiltered(t.Context(), filter, count, nil) { + if err != nil { + t.Fatal(err) + } + yielded++ + break + } + if yielded != 1 { + t.Fatalf("early-stopped retrieval yielded %d records", yielded) + } + ctx, cancel := context.WithCancel(t.Context()) + yielded = 0 + var canceled bool + for _, err := range collection.GetFiltered(ctx, filter, count, nil) { + if err != nil { + if !errors.Is(err, context.Canceled) { + t.Fatal(err) + } + canceled = true + break + } + yielded++ + cancel() + } + cancel() + if !canceled || yielded != 1 { + t.Fatalf("canceled retrieval yielded %d records, cancellation error=%v", yielded, canceled) + } + var rejected bool + for _, err := range collection.Search(t.Context(), []float32{1, 0, 0}, 1, nil) { + if err == nil || !strings.Contains(err.Error(), "partition") { + t.Fatalf("unscoped search error = %v", err) + } + rejected = true + } + if !rejected { + t.Fatal("unscoped ranking was not rejected") + } +} + +func TestE2E_CosmosDBGeneratedVectors(t *testing.T) { + type record struct { + ID uuid.UUID `vectordata:"key"` + Bucket string `vectordata:"data,indexed,storage=bucket"` + Text string `vectordata:"data,fulltext,storage=text"` + Content string `vectordata:"vector,dimensions=3072,storage=vector"` + } + store := newLiveCosmosDB(t, azureprovider.CosmosDBConfig{PartitionKeyProperties: []string{"Bucket"}}, true) + name := liveCosmosContainer(t, store) + collection, err := vectordata.GetCollection[azureprovider.CosmosKey, *record](store.Store, name, nil) + if err != nil { + t.Fatal(err) + } + if err := collection.EnsureExists(t.Context()); err != nil { + t.Fatal(err) + } + records := []*record{ + {Bucket: "alpha", Text: "I enjoy pirate jokes about parrots and ships.", Content: "I enjoy pirate jokes about parrots and ships."}, + {Bucket: "alpha", Text: "I like hiking in the mountains.", Content: "I like hiking in the mountains."}, + {Bucket: "beta", Text: "I enjoy pirate jokes.", Content: "I enjoy pirate jokes."}, + } + if err := collection.Upsert(t.Context(), records...); err != nil { + t.Fatal(err) + } + filter := vectordata.Filter{Operator: vectordata.Equal, Property: "Bucket", Value: "alpha"} + results := collect(t, collection.Search(t.Context(), "Tell me a pirate joke.", 2, &vectordata.SearchOptions{Filter: &filter})) + if len(results) != 2 || results[0].Record.ID != records[0].ID || results[0].Record.Content != "" || results[0].Score == nil { + t.Fatalf("generated-vector search = %+v", results) + } + hybrid := collection.(vectordata.KeywordHybridSearcher[*record]) + hybridResults := collect(t, hybrid.HybridSearch(t.Context(), "pirate jokes", []string{"pirate"}, 2, + &vectordata.HybridSearchOptions{Filter: &filter})) + if len(hybridResults) != 2 || hybridResults[0].Record.ID != records[0].ID { + t.Fatalf("generated hybrid search = %+v", hybridResults) + } + type nativeView struct { + ID uuid.UUID `vectordata:"key"` + Bucket string `vectordata:"data,indexed,storage=bucket"` + Content []float32 `vectordata:"vector,dimensions=3072,storage=vector"` + } + native, err := vectordata.GetCollection[azureprovider.CosmosKey, nativeView](store.Store, name, nil) + if err != nil { + t.Fatal(err) + } + key := azureprovider.CosmosKey{DocumentID: records[0].ID.String(), PartitionKey: []any{"alpha"}} + value, found, err := native.Get(t.Context(), key, &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) + } + if _, _, err := collection.Get(t.Context(), key, &vectordata.RecordRetrievalOptions{IncludeVectors: true}); err == nil { + t.Fatal("generated-vector IncludeVectors unexpectedly succeeded") + } + if records[0].Content != records[0].Text { + t.Fatal("embedding generation changed the source record") + } +} diff --git a/provider/azureprovider/cosmosdb_query.go b/provider/azureprovider/cosmosdb_query.go new file mode 100644 index 00000000..5b3dcb7f --- /dev/null +++ b/provider/azureprovider/cosmosdb_query.go @@ -0,0 +1,498 @@ +// Copyright (c) Microsoft. All rights reserved. + +package azureprovider + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "iter" + "reflect" + "strconv" + "strings" + + "github.com/Azure/azure-sdk-for-go/sdk/data/azcosmos" + "github.com/google/uuid" + "github.com/microsoft/agent-framework-go/vectordata" + "github.com/microsoft/agent-framework-go/vectordata/providerservices" +) + +type cosmosParameter struct { + Name string + Value any +} + +func (c *cosmosCollection) translateFilter(filter *vectordata.Filter) (string, []cosmosParameter, error) { + if err := filter.Validate(); err != nil { + return "", nil, err + } + if filter == nil { + return "", nil, nil + } + var parameters []cosmosParameter + next := 0 + var translate func(*vectordata.Filter) (string, error) + translate = func(node *vectordata.Filter) (string, error) { + switch node.Operator { + case vectordata.And, vectordata.Or: + parts := make([]string, len(node.Operands)) + for i := range node.Operands { + part, err := translate(&node.Operands[i]) + if err != nil { + return "", err + } + parts[i] = part + } + return "(" + strings.Join(parts, " "+strings.ToUpper(string(node.Operator))+" ") + ")", nil + case vectordata.Not: + part, err := translate(&node.Operands[0]) + if err != nil { + return "", err + } + return "(NOT " + part + ")", nil + } + property, found := c.properties[node.Property] + if !found { + return "", fmt.Errorf("filter property %q was not found", node.Property) + } + if _, ok := property.(*providerservices.VectorPropertyModel); ok { + return "", fmt.Errorf("filter property %q is a vector, not a data property", node.Property) + } + typ := cosmosFilterType(property) + expression := "c[" + cosmosJSONString(propertyStorageName(property)) + "]" + parameterValue := func(value any) any { + if _, ok := property.(*providerservices.KeyPropertyModel); ok { + switch key := value.(type) { + case uuid.UUID: + return key.String() + case CosmosKey: + return key.DocumentID + } + } + return value + } + valueName := func(value any) string { + name := "@p" + strconv.Itoa(next) + next++ + parameters = append(parameters, cosmosParameter{Name: name, Value: parameterValue(value)}) + return name + } + switch node.Operator { + case vectordata.Equal, vectordata.NotEqual: + if node.Value == nil { + if node.Operator == vectordata.Equal { + return "IS_NULL(" + expression + ")", nil + } + return "NOT IS_NULL(" + expression + ")", nil + } + if typ.Kind() == reflect.Slice || typ.Kind() == reflect.Array { + return "", fmt.Errorf("filter property %q is a collection; scalar comparison is unsupported", node.Property) + } + if !cosmosFilterValueCompatible(typ, node.Value) { + return "", fmt.Errorf("filter value for property %q has incompatible type %T", node.Property, node.Value) + } + operator := "=" + if node.Operator == vectordata.NotEqual { + operator = "<>" + } + return expression + " " + operator + " " + valueName(node.Value), nil + case vectordata.GreaterThan: + if typ.Kind() == reflect.Slice || typ.Kind() == reflect.Array || !cosmosFilterValueCompatible(typ, node.Value) { + return "", fmt.Errorf("filter value for property %q has incompatible type %T", node.Property, node.Value) + } + return expression + " > " + valueName(node.Value), nil + case vectordata.GreaterThanOrEqual: + if typ.Kind() == reflect.Slice || typ.Kind() == reflect.Array || !cosmosFilterValueCompatible(typ, node.Value) { + return "", fmt.Errorf("filter value for property %q has incompatible type %T", node.Property, node.Value) + } + return expression + " >= " + valueName(node.Value), nil + case vectordata.LessThan: + if typ.Kind() == reflect.Slice || typ.Kind() == reflect.Array || !cosmosFilterValueCompatible(typ, node.Value) { + return "", fmt.Errorf("filter value for property %q has incompatible type %T", node.Property, node.Value) + } + return expression + " < " + valueName(node.Value), nil + case vectordata.LessThanOrEqual: + if typ.Kind() == reflect.Slice || typ.Kind() == reflect.Array || !cosmosFilterValueCompatible(typ, node.Value) { + return "", fmt.Errorf("filter value for property %q has incompatible type %T", node.Property, node.Value) + } + return expression + " <= " + valueName(node.Value), nil + case vectordata.Contains: + if typ.Kind() != reflect.Slice && typ.Kind() != reflect.Array { + return "", fmt.Errorf("filter property %q is not a collection", node.Property) + } + if !cosmosFilterValueCompatible(typ.Elem(), node.Value) { + return "", fmt.Errorf("filter value for property %q has incompatible type %T", node.Property, node.Value) + } + return "ARRAY_CONTAINS(" + expression + ", " + valueName(node.Value) + ")", nil + case vectordata.In: + if typ.Kind() == reflect.Slice || typ.Kind() == reflect.Array { + return "", fmt.Errorf("filter property %q is a collection; in filters require a scalar property", node.Property) + } + values := reflect.ValueOf(node.Value) + normalized := make([]any, values.Len()) + for i := range values.Len() { + value := values.Index(i).Interface() + if !cosmosFilterValueCompatible(typ, value) { + return "", fmt.Errorf("in filter value %d for property %q has incompatible type %T", i, node.Property, values.Index(i).Interface()) + } + normalized[i] = parameterValue(value) + } + return "ARRAY_CONTAINS(" + valueName(normalized) + ", " + expression + ")", nil + default: + return "", fmt.Errorf("unsupported filter operator %q", node.Operator) + } + } + sql, err := translate(filter) + if err != nil { + return "", nil, err + } + for i := range parameters { + parameters[i].Value, err = cosmosDataJSON(parameters[i].Value) + if err != nil { + return "", nil, fmt.Errorf("filter parameter %q: %w", parameters[i].Name, err) + } + } + return sql, parameters, nil +} + +func cosmosFilterType(property providerservices.PropertyModel) reflect.Type { + var typ reflect.Type + switch p := property.(type) { + case *providerservices.KeyPropertyModel: + typ = p.Type + if typ == reflect.TypeFor[uuid.UUID]() || typ == reflect.TypeFor[CosmosKey]() { + typ = reflect.TypeFor[string]() + } + case *providerservices.DataPropertyModel: + typ = p.Type + } + for typ != nil && typ.Kind() == reflect.Pointer { + typ = typ.Elem() + } + return typ +} + +func cosmosFilterValueCompatible(target reflect.Type, value any) bool { + if target == nil || value == nil { + return value == nil + } + if target.Kind() == reflect.String { + switch value.(type) { + case uuid.UUID, CosmosKey: + return true + } + } + for target.Kind() == reflect.Pointer { + target = target.Elem() + } + valueType := reflect.TypeOf(value) + for valueType.Kind() == reflect.Pointer { + valueType = valueType.Elem() + } + if target == valueType || valueType.AssignableTo(target) { + return true + } + numeric := func(kind reflect.Kind) bool { + return kind >= reflect.Int && kind <= reflect.Float64 + } + return numeric(target.Kind()) && numeric(valueType.Kind()) +} + +func cosmosJSONString(value string) string { + encoded, _ := json.Marshal(value) + return string(encoded) +} + +func (c *cosmosCollection) orderBy(orderBy []vectordata.Sort) (string, error) { + parts := make([]string, len(orderBy)) + for i, order := range orderBy { + property, found := c.properties[order.Property] + if !found { + return "", fmt.Errorf("sort property %q was not found", order.Property) + } + if _, ok := property.(*providerservices.VectorPropertyModel); ok { + return "", fmt.Errorf("sort property %q is a vector", order.Property) + } + direction := "ASC" + if order.Descending { + direction = "DESC" + } + parts[i] = "c[" + cosmosJSONString(propertyStorageName(property)) + "] " + direction + } + return strings.Join(parts, ", "), nil +} + +func (c *cosmosCollection) queryPartition(filter *vectordata.Filter, required bool) (azcosmos.PartitionKey, error) { + values := make([]any, len(c.partitionProperties)) + for i, property := range c.partitionProperties { + name := propertyName(property) + value, found := cosmosPartitionFilterValue(filter, name) + if !found { + if required { + return azcosmos.PartitionKey{}, fmt.Errorf("azcosmos requires equality filters for every partition key property on ranked or ordered queries; missing property %q", name) + } + return azcosmos.NewPartitionKey(), nil + } + if property == c.key && value != nil { + id, err := cosmosID(value) + if err != nil { + return azcosmos.PartitionKey{}, err + } + value = id + } + normalized, err := normalizePartitionValue(value) + if err != nil { + return azcosmos.PartitionKey{}, fmt.Errorf("partition key filter %q: %w", name, err) + } + values[i] = normalized + } + return cosmosPartitionKey(values) +} + +func cosmosPartitionFilterValue(filter *vectordata.Filter, property string) (any, bool) { + if filter == nil { + return nil, false + } + switch filter.Operator { + case vectordata.Equal: + if filter.Property == property { + return filter.Value, true + } + case vectordata.And: + for i := range filter.Operands { + if value, found := cosmosPartitionFilterValue(&filter.Operands[i], property); found { + return value, true + } + } + } + return nil, false +} + +func (c *cosmosCollection) 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 := ctx.Err(); err != nil { + yield(vectordata.SearchResult[any]{}, cosmosFailure(c.store.database, c.name, "Search", err)) + return + } + var opts vectordata.SearchOptions + if options != nil { + opts = *options + if err := opts.Validate(); err != nil { + yield(vectordata.SearchResult[any]{}, cosmosFailure(c.store.database, c.name, "Search", err)) + return + } + if err := c.validateRetrieval(opts.IncludeVectors); err != nil { + yield(vectordata.SearchResult[any]{}, cosmosFailure(c.store.database, c.name, "Search", err)) + return + } + } + if top <= 0 { + yield(vectordata.SearchResult[any]{}, cosmosFailure(c.store.database, c.name, "Search", errors.New("top must be positive"))) + return + } + property, err := c.model.VectorProperty(opts.VectorProperty) + if err != nil { + yield(vectordata.SearchResult[any]{}, cosmosFailure(c.store.database, c.name, "Search", err)) + return + } + filter, parameters, err := c.translateFilter(opts.Filter) + if err != nil { + yield(vectordata.SearchResult[any]{}, cosmosFailure(c.store.database, c.name, "Search", err)) + return + } + partition, err := c.queryPartition(opts.Filter, true) + if err != nil { + yield(vectordata.SearchResult[any]{}, cosmosFailure(c.store.database, c.name, "Search", err)) + return + } + vector, err := c.vectorForSearch(ctx, property, value) + if err != nil { + yield(vectordata.SearchResult[any]{}, cosmosFailure(c.store.database, c.name, "Search", err)) + return + } + query, parameters, err := c.buildSearchQuery(property, vector, nil, nil, filter, parameters, top, opts.Skip, opts.ScoreThreshold, opts.IncludeVectors) + if err != nil { + yield(vectordata.SearchResult[any]{}, cosmosFailure(c.store.database, c.name, "Search", err)) + return + } + err = c.eachQueryDocument(ctx, query, parameters, partition, func(raw json.RawMessage) bool { + result, err := c.decodeSearchResult(raw, opts.IncludeVectors) + if err != nil { + yield(vectordata.SearchResult[any]{}, cosmosFailure(c.store.database, c.name, "Search", err)) + return false + } + return yield(result, nil) + }) + if err != nil { + yield(vectordata.SearchResult[any]{}, cosmosFailure(c.store.database, c.name, "Search", err)) + } + } +} + +func (c *cosmosCollection) 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 err := ctx.Err(); err != nil { + yield(vectordata.SearchResult[any]{}, cosmosFailure(c.store.database, c.name, "HybridSearch", err)) + return + } + var opts vectordata.HybridSearchOptions + if options != nil { + opts = *options + if err := opts.Validate(); err != nil { + yield(vectordata.SearchResult[any]{}, cosmosFailure(c.store.database, c.name, "HybridSearch", err)) + return + } + if err := c.validateRetrieval(opts.IncludeVectors); err != nil { + yield(vectordata.SearchResult[any]{}, cosmosFailure(c.store.database, c.name, "HybridSearch", err)) + return + } + } + if top <= 0 { + yield(vectordata.SearchResult[any]{}, cosmosFailure(c.store.database, c.name, "HybridSearch", errors.New("top must be positive"))) + return + } + vectorProperty, err := c.model.VectorProperty(opts.VectorProperty) + if err != nil { + yield(vectordata.SearchResult[any]{}, cosmosFailure(c.store.database, c.name, "HybridSearch", err)) + return + } + textProperty, err := c.model.FullTextProperty(opts.AdditionalProperty) + if err != nil { + yield(vectordata.SearchResult[any]{}, cosmosFailure(c.store.database, c.name, "HybridSearch", err)) + return + } + filter, parameters, err := c.translateFilter(opts.Filter) + if err != nil { + yield(vectordata.SearchResult[any]{}, cosmosFailure(c.store.database, c.name, "HybridSearch", err)) + return + } + partition, err := c.queryPartition(opts.Filter, true) + if err != nil { + yield(vectordata.SearchResult[any]{}, cosmosFailure(c.store.database, c.name, "HybridSearch", err)) + return + } + vector, err := c.vectorForSearch(ctx, vectorProperty, value) + if err != nil { + yield(vectordata.SearchResult[any]{}, cosmosFailure(c.store.database, c.name, "HybridSearch", err)) + return + } + query, parameters, err := c.buildSearchQuery(vectorProperty, vector, textProperty, keywords, filter, parameters, top, opts.Skip, opts.ScoreThreshold, opts.IncludeVectors) + if err != nil { + yield(vectordata.SearchResult[any]{}, cosmosFailure(c.store.database, c.name, "HybridSearch", err)) + return + } + err = c.eachQueryDocument(ctx, query, parameters, partition, func(raw json.RawMessage) bool { + result, err := c.decodeSearchResult(raw, opts.IncludeVectors) + if err != nil { + yield(vectordata.SearchResult[any]{}, cosmosFailure(c.store.database, c.name, "HybridSearch", err)) + return false + } + return yield(result, nil) + }) + if err != nil { + yield(vectordata.SearchResult[any]{}, cosmosFailure(c.store.database, c.name, "HybridSearch", err)) + } + } +} + +func (c *cosmosCollection) selectFields(includeVectors bool) []string { + fields := make([]string, 0, len(c.model.Properties)+1) + for _, property := range c.model.Properties { + if _, isVector := property.(*providerservices.VectorPropertyModel); isVector && !includeVectors { + continue + } + fields = append(fields, "c["+cosmosJSONString(propertyStorageName(property))+"]") + } + return fields +} + +func (c *cosmosCollection) buildSearchQuery( + vectorProperty *providerservices.VectorPropertyModel, + vector any, + textProperty *providerservices.DataPropertyModel, + keywords []string, + filter string, + parameters []cosmosParameter, + top, skip int, + scoreThreshold *float64, + includeVectors bool, +) (string, []cosmosParameter, error) { + vectorExpression := "VectorDistance(c[" + cosmosJSONString(vectorProperty.StorageName) + "], @vector)" + parameters = append(parameters, cosmosParameter{Name: "@vector", Value: vector}) + selectFields := c.selectFields(includeVectors) + selectFields = append(selectFields, vectorExpression+" AS "+cosmosScoreField) + rankingExpression := vectorExpression + if textProperty != nil && keywords != nil { + fullTextArguments := []string{"c[" + cosmosJSONString(textProperty.StorageName) + "]"} + for i, keyword := range keywords { + name := "@keyword" + strconv.Itoa(i) + fullTextArguments = append(fullTextArguments, name) + parameters = append(parameters, cosmosParameter{Name: name, Value: keyword}) + } + rankingExpression = "RANK RRF(" + vectorExpression + ", FullTextScore(" + strings.Join(fullTextArguments, ", ") + "))" + } + conditions := make([]string, 0, 2) + if filter != "" { + conditions = append(conditions, filter) + } + if scoreThreshold != nil { + operator := ">=" + switch vectorProperty.DistanceFunction { + case "", vectordata.CosineSimilarity, vectordata.DotProductSimilarity: + case vectordata.EuclideanDistance: + operator = "<=" + default: + return "", nil, fmt.Errorf("score threshold is unsupported for distance function %q", vectorProperty.DistanceFunction) + } + conditions = append(conditions, vectorExpression+" "+operator+" @scoreThreshold") + parameters = append(parameters, cosmosParameter{Name: "@scoreThreshold", Value: *scoreThreshold}) + } + topClause := fmt.Sprintf("TOP %d ", top) + var query strings.Builder + query.WriteString("SELECT ") + query.WriteString(topClause) + query.WriteString(strings.Join(selectFields, ", ")) + query.WriteString(" FROM c") + if len(conditions) > 0 { + query.WriteString(" WHERE ") + query.WriteString(strings.Join(conditions, " AND ")) + } + query.WriteString(" ORDER BY ") + query.WriteString(rankingExpression) + if skip > 0 { + query.Reset() + query.WriteString("SELECT ") + query.WriteString(strings.Join(selectFields, ", ")) + query.WriteString(" FROM c") + if len(conditions) > 0 { + query.WriteString(" WHERE ") + query.WriteString(strings.Join(conditions, " AND ")) + } + query.WriteString(" ORDER BY ") + query.WriteString(rankingExpression) + fmt.Fprintf(&query, " OFFSET %d LIMIT %d", skip, top) + } + return query.String(), parameters, nil +} + +func (c *cosmosCollection) decodeSearchResult(raw json.RawMessage, includeVectors bool) (vectordata.SearchResult[any], error) { + document, err := rawDocument(raw) + if err != nil { + return vectordata.SearchResult[any]{}, err + } + var score *float64 + if rawScore, found := document[cosmosScoreField]; found && !bytes.Equal(bytes.TrimSpace(rawScore), []byte("null")) { + var value float64 + if err := json.Unmarshal(rawScore, &value); err != nil { + return vectordata.SearchResult[any]{}, fmt.Errorf("invalid Cosmos DB score: %w", err) + } + score = &value + } + record, err := c.decodeRecord(document, includeVectors) + if err != nil { + return vectordata.SearchResult[any]{}, err + } + return vectordata.SearchResult[any]{Record: record, Score: score}, nil +} diff --git a/provider/azureprovider/cosmosdb_test.go b/provider/azureprovider/cosmosdb_test.go new file mode 100644 index 00000000..bd236380 --- /dev/null +++ b/provider/azureprovider/cosmosdb_test.go @@ -0,0 +1,1890 @@ +// Copyright (c) Microsoft. All rights reserved. + +package azureprovider_test + +import ( + "context" + "encoding/base64" + "encoding/json" + "errors" + "io" + "net/http" + "net/http/httptest" + "reflect" + "strings" + "sync" + "testing" + "time" + + "github.com/Azure/azure-sdk-for-go/sdk/azcore" + "github.com/Azure/azure-sdk-for-go/sdk/data/azcosmos" + "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" +) + +type cosmosProtocolRecord struct { + ID uuid.UUID `vectordata:"key"` + Category string `vectordata:"data,indexed"` + Text string `vectordata:"data,fulltext"` + Tags []string `vectordata:"data"` + Vector []float32 `vectordata:"vector,dimensions=2,index=disk_ann,distance=cosine_similarity"` + Bytes []uint8 `vectordata:"vector,dimensions=2,index=flat"` + Signed []int8 `vectordata:"vector,dimensions=2,index=quantized_flat"` +} + +type cosmosExampleShapedRecord struct { + ID uuid.UUID `vectordata:"key"` + Topic string `vectordata:"data,indexed,storage=topic"` + Text string `vectordata:"data,storage=text"` + Content string `vectordata:"vector,dimensions=3072,index=disk_ann,storage=vector"` +} + +type cosmosTestEmbeddingGenerator struct { + calls int + inputs [][]string +} + +func (g *cosmosTestEmbeddingGenerator) Generate(_ context.Context, values []string, _ *embedding.Options) (*embedding.Generated[*embedding.Of[float32]], error) { + g.calls++ + g.inputs = append(g.inputs, append([]string(nil), values...)) + result := &embedding.Generated[*embedding.Of[float32]]{ + Embeddings: make([]*embedding.Of[float32], len(values)), + } + for i, value := range values { + result.Embeddings[i] = &embedding.Of[float32]{Vector: []float32{float32(len(value)), float32(i)}} + } + return result, nil +} + +func TestCosmosDBNativeSDKLifecycleCRUDAndSearch(t *testing.T) { + var mu sync.Mutex + var stored map[string]any + var createBody map[string]any + var queries []string + var server *httptest.Server + server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + switch { + case r.Method == http.MethodGet && r.URL.Path == "/": + _ = json.NewEncoder(w).Encode(map[string]any{ + "id": "mock", "enableMultipleWriteLocations": false, + "readableLocations": []any{map[string]any{"name": "local", "databaseAccountEndpoint": server.URL + "/"}}, + "writableLocations": []any{map[string]any{"name": "local", "databaseAccountEndpoint": server.URL + "/"}}, + }) + case r.Method == http.MethodPost && r.URL.Path == "/dbs/sample/colls" && strings.Contains(r.Header.Get("Content-Type"), "application/query+json"): + _, _ = io.WriteString(w, `{"DocumentCollections":[{"id":"vectors"}]}`) + case r.Method == http.MethodPost && r.URL.Path == "/dbs/sample/colls": + if err := json.NewDecoder(r.Body).Decode(&createBody); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + w.WriteHeader(http.StatusCreated) + _, _ = io.WriteString(w, `{"id":"vectors","partitionKey":{"kind":"Hash","paths":["/id"],"version":2}}`) + case r.Method == http.MethodGet && r.URL.Path == "/dbs/sample/colls/vectors": + _, _ = io.WriteString(w, `{"id":"vectors","partitionKey":{"kind":"Hash","paths":["/id"],"version":2}}`) + case r.Method == http.MethodPost && r.URL.Path == "/dbs/sample/colls/vectors/docs" && strings.Contains(r.Header.Get("Content-Type"), "application/query+json"): + var request struct { + Query string `json:"query"` + } + if err := json.NewDecoder(r.Body).Decode(&request); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + mu.Lock() + queries = append(queries, request.Query) + document := map[string]any{ + "id": stored["id"], "Category": stored["Category"], "Text": stored["Text"], + "Tags": stored["Tags"], + } + document["__cosmos_score"] = 0.75 + if strings.Contains(request.Query, `c["Vector"]`) { + document["Vector"] = stored["Vector"] + } + mu.Unlock() + _ = json.NewEncoder(w).Encode(map[string]any{"Documents": []any{document}}) + case r.Method == http.MethodPost && r.URL.Path == "/dbs/sample/colls/vectors/docs": + mu.Lock() + if err := json.NewDecoder(r.Body).Decode(&stored); err != nil { + mu.Unlock() + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + mu.Unlock() + w.WriteHeader(http.StatusOK) + _ = json.NewEncoder(w).Encode(stored) + case r.Method == http.MethodGet && strings.HasPrefix(r.URL.Path, "/dbs/sample/colls/vectors/docs/"): + _ = json.NewEncoder(w).Encode(stored) + case r.Method == http.MethodDelete && strings.HasPrefix(r.URL.Path, "/dbs/sample/colls/vectors/docs/"): + w.WriteHeader(http.StatusNoContent) + case r.Method == http.MethodDelete && r.URL.Path == "/dbs/sample/colls/vectors": + w.WriteHeader(http.StatusNoContent) + default: + http.Error(w, r.Method+" "+r.URL.Path, http.StatusNotFound) + } + })) + defer server.Close() + + credential, err := azcosmos.NewKeyCredential(base64.StdEncoding.EncodeToString([]byte("local-test-key"))) + if err != nil { + t.Fatal(err) + } + client, err := azcosmos.NewClientWithKey(server.URL, credential, &azcosmos.ClientOptions{ + ClientOptions: azcore.ClientOptions{ + Transport: server.Client(), + InsecureAllowCredentialWithHTTP: true, + }, + }) + if err != nil { + t.Fatal(err) + } + db, err := azureprovider.NewCosmosDB(client, "sample", azureprovider.CosmosDBConfig{}) + if err != nil { + t.Fatal(err) + } + collection, err := vectordata.GetCollection[uuid.UUID, *cosmosProtocolRecord](db.Store, "vectors", nil) + if err != nil { + t.Fatal(err) + } + if err := collection.EnsureExists(t.Context()); err != nil { + t.Fatal(err) + } + if createBody["id"] != "vectors" { + t.Fatalf("container create payload = %#v", createBody) + } + embeddingPolicy, ok := createBody["vectorEmbeddingPolicy"].(map[string]any) + if !ok { + t.Fatalf("container missing vector embedding policy: %#v", createBody) + } + embeddings, ok := embeddingPolicy["vectorEmbeddings"].([]any) + if !ok || len(embeddings) != 3 { + t.Fatalf("vector embedding policy = %#v", embeddingPolicy) + } + if embeddings[0].(map[string]any)["dataType"] != "float32" || + embeddings[1].(map[string]any)["dataType"] != "uint8" || + embeddings[2].(map[string]any)["dataType"] != "int8" { + t.Fatalf("vector embedding data types = %#v", embeddings) + } + fullTextPolicy, ok := createBody["fullTextPolicy"].(map[string]any) + if !ok || fullTextPolicy["defaultLanguage"] != "en-US" { + t.Fatalf("full-text policy = %#v", createBody["fullTextPolicy"]) + } + + record := &cosmosProtocolRecord{ + Category: "engineering", Text: "vector search", Tags: []string{"go", "cosmos"}, + Vector: []float32{0.1, 0.2}, Bytes: []uint8{1, 2}, Signed: []int8{-1, 2}, + } + if err := collection.Upsert(t.Context(), record); err != nil { + t.Fatal(err) + } + if record.ID == uuid.Nil { + t.Fatal("upsert did not generate a UUID key") + } + mu.Lock() + gotStored := stored + mu.Unlock() + if gotStored["id"] != record.ID.String() || gotStored["Vector"] == nil { + t.Fatalf("stored document = %#v", gotStored) + } + if vector, ok := gotStored["Bytes"].([]any); !ok || !reflect.DeepEqual(vector, []any{float64(1), float64(2)}) { + t.Fatalf("byte vector must be a numeric JSON array: %#v", gotStored["Bytes"]) + } + + got, found, err := collection.Get(t.Context(), record.ID, nil) + if err != nil || !found { + t.Fatalf("Get() = (%v, %v, %v)", got, found, err) + } + if got.ID != record.ID || got.Vector != nil { + t.Fatalf("Get() did not return key and omit vectors by default: %#v", got) + } + gotWithVectors, found, err := collection.Get(t.Context(), record.ID, &vectordata.RecordRetrievalOptions{IncludeVectors: true}) + if err != nil || !found || gotWithVectors.Vector[0] != 0.1 || + gotWithVectors.Bytes[0] != 1 || gotWithVectors.Signed[0] != -1 { + t.Fatalf("Get() with vectors = (%#v, %v, %v)", gotWithVectors, found, err) + } + + idFilter := vectordata.Filter{Operator: vectordata.Equal, Property: "ID", Value: record.ID} + filtered := make([]*cosmosProtocolRecord, 0, 1) + for item, err := range collection.GetFiltered(t.Context(), vectordata.Filter{ + Operator: vectordata.And, Operands: []vectordata.Filter{ + {Operator: vectordata.Equal, Property: "Category", Value: "engineering"}, + idFilter, + }, + }, 5, nil) { + if err != nil { + t.Fatal(err) + } + filtered = append(filtered, item) + } + if len(filtered) != 1 || filtered[0].Category != "engineering" { + t.Fatalf("filtered records = %#v", filtered) + } + for item, err := range collection.GetFiltered(t.Context(), vectordata.Filter{ + Operator: vectordata.Contains, Property: "Tags", Value: "go", + }, 5, nil) { + if err != nil { + t.Fatal(err) + } + if item.Category != "engineering" { + t.Fatalf("contains result = %#v", item) + } + } + + searchResults := make([]vectordata.SearchResult[*cosmosProtocolRecord], 0, 1) + threshold := 0.5 + for result, err := range collection.Search(t.Context(), []float32{0.1, 0.2}, 3, &vectordata.SearchOptions{ + VectorProperty: "Vector", ScoreThreshold: &threshold, Filter: &idFilter, + }) { + if err != nil { + t.Fatal(err) + } + searchResults = append(searchResults, result) + } + if len(searchResults) != 1 || searchResults[0].Score == nil || *searchResults[0].Score != 0.75 || searchResults[0].Record.Vector != nil { + t.Fatalf("search results = %#v", searchResults) + } + + hybridResults := make([]vectordata.SearchResult[*cosmosProtocolRecord], 0, 1) + hybrid, ok := collection.(vectordata.KeywordHybridSearcher[*cosmosProtocolRecord]) + if !ok { + t.Fatal("collection does not expose hybrid search") + } + for result, err := range hybrid.HybridSearch(t.Context(), []float32{0.1, 0.2}, []string{"cosmos", "vectors"}, 3, + &vectordata.HybridSearchOptions{VectorProperty: "Vector", ScoreThreshold: &threshold, Filter: &idFilter}) { + if err != nil { + t.Fatal(err) + } + hybridResults = append(hybridResults, result) + } + if len(hybridResults) != 1 { + t.Fatalf("hybrid results = %#v", hybridResults) + } + if hybridResults[0].Score == nil || *hybridResults[0].Score != 0.75 { + t.Fatalf("hybrid results = %#v", hybridResults) + } + + names := make([]string, 0, 1) + for name, err := range db.ListCollectionNames(t.Context()) { + if err != nil { + t.Fatal(err) + } + names = append(names, name) + } + if len(names) != 1 || names[0] != "vectors" { + t.Fatalf("collection names = %#v", names) + } + mu.Lock() + gotQueries := append([]string(nil), queries...) + mu.Unlock() + if len(gotQueries) != 4 || + !strings.Contains(gotQueries[0], "@p0") || + !strings.Contains(gotQueries[1], `ARRAY_CONTAINS(c["Tags"]`) || + !strings.Contains(gotQueries[2], "VectorDistance") || + !strings.Contains(gotQueries[2], ">= @scoreThreshold") || + !strings.Contains(gotQueries[3], "FullTextScore") || + !strings.Contains(gotQueries[3], "RANK RRF") || + !strings.Contains(gotQueries[3], ">= @scoreThreshold") { + t.Fatalf("query SQL = %#v", gotQueries) + } + + var invalidVectorErr error + for _, err := range collection.Search(t.Context(), []float32{0.1}, 1, &vectordata.SearchOptions{ + VectorProperty: "Vector", Filter: &idFilter, + }) { + invalidVectorErr = err + } + if invalidVectorErr == nil || !strings.Contains(invalidVectorErr.Error(), "dimension 2") { + t.Fatalf("invalid vector search error = %v", invalidVectorErr) + } + var invalidFilterErr error + for _, err := range collection.GetFiltered(t.Context(), vectordata.Filter{ + Operator: vectordata.Contains, Property: "Category", Value: "engineering", + }, 1, nil) { + invalidFilterErr = err + } + if invalidFilterErr == nil || !strings.Contains(invalidFilterErr.Error(), "not a collection") { + t.Fatalf("invalid filter error = %v", invalidFilterErr) + } + if err := collection.Delete(t.Context(), record.ID); err != nil { + t.Fatal(err) + } + if err := collection.EnsureDeleted(t.Context()); err != nil { + t.Fatal(err) + } +} + +func TestCosmosDBExampleShapedGeneratedVectorSchema(t *testing.T) { + var createBody map[string]any + var server *httptest.Server + server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + if r.Method == http.MethodGet && r.URL.Path == "/" { + _ = json.NewEncoder(w).Encode(map[string]any{ + "id": "mock", "enableMultipleWriteLocations": false, + "readableLocations": []any{map[string]any{"name": "local", "databaseAccountEndpoint": server.URL + "/"}}, + "writableLocations": []any{map[string]any{"name": "local", "databaseAccountEndpoint": server.URL + "/"}}, + }) + return + } + if r.Method == http.MethodPost && r.URL.Path == "/dbs/sample/colls" { + if err := json.NewDecoder(r.Body).Decode(&createBody); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + w.WriteHeader(http.StatusCreated) + _, _ = io.WriteString(w, `{"id":"example-shaped","partitionKey":{"kind":"Hash","paths":["/id"],"version":2}}`) + return + } + http.Error(w, r.Method+" "+r.URL.Path, http.StatusNotFound) + })) + defer server.Close() + + credential, err := azcosmos.NewKeyCredential(base64.StdEncoding.EncodeToString([]byte("local-test-key"))) + if err != nil { + t.Fatal(err) + } + client, err := azcosmos.NewClientWithKey(server.URL, credential, &azcosmos.ClientOptions{ + ClientOptions: azcore.ClientOptions{ + Transport: server.Client(), + InsecureAllowCredentialWithHTTP: true, + }, + }) + if err != nil { + t.Fatal(err) + } + store, err := azureprovider.NewCosmosDB(client, "sample", azureprovider.CosmosDBConfig{ + PartitionKeyProperties: []string{"Topic"}, + EmbeddingGenerator: embedding.Bind[string, *embedding.Of[float32]](&cosmosTestEmbeddingGenerator{}), + }) + if err != nil { + t.Fatal(err) + } + collection, err := vectordata.GetCollection[azureprovider.CosmosKey, *cosmosExampleShapedRecord](store.Store, "example-shaped", nil) + if err != nil { + t.Fatalf("open collection with index=diskann: %v", err) + } + if err := collection.EnsureExists(t.Context()); err != nil { + t.Fatal(err) + } + partition, ok := createBody["partitionKey"].(map[string]any) + paths, pathsOK := partition["paths"].([]any) + if !ok || !pathsOK || len(paths) != 1 || paths[0] != "/topic" { + t.Fatalf("UUID collection partition key = %#v", createBody["partitionKey"]) + } + vectorPolicy, ok := createBody["vectorEmbeddingPolicy"].(map[string]any) + if !ok { + t.Fatalf("vector embedding policy = %#v", createBody["vectorEmbeddingPolicy"]) + } + embeddings, ok := vectorPolicy["vectorEmbeddings"].([]any) + if !ok || len(embeddings) != 1 { + t.Fatalf("vector embeddings = %#v", vectorPolicy["vectorEmbeddings"]) + } + vectorEmbedding := embeddings[0].(map[string]any) + if vectorEmbedding["path"] != "/vector" || vectorEmbedding["dimensions"] != float64(3072) || + vectorEmbedding["dataType"] != "float32" || vectorEmbedding["distanceFunction"] != "cosine" { + t.Fatalf("example-shaped vector embedding policy = %#v", vectorEmbedding) + } + indexingPolicy, ok := createBody["indexingPolicy"].(map[string]any) + if !ok { + t.Fatalf("indexing policy = %#v", createBody["indexingPolicy"]) + } + indexes, ok := indexingPolicy["vectorIndexes"].([]any) + if !ok || len(indexes) != 1 { + t.Fatalf("vector indexes = %#v", indexingPolicy["vectorIndexes"]) + } + if index := indexes[0].(map[string]any); index["path"] != "/vector" || index["type"] != "diskANN" { + t.Fatalf("example-shaped vector index = %#v", index) + } +} + +func TestCosmosDBIteratorStopsBeforeFetchingContinuationPage(t *testing.T) { + pageRequests := 0 + var server *httptest.Server + server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + if r.Method == http.MethodGet && r.URL.Path == "/" { + _ = json.NewEncoder(w).Encode(map[string]any{ + "id": "mock", "enableMultipleWriteLocations": false, + "readableLocations": []any{map[string]any{"name": "local", "databaseAccountEndpoint": server.URL + "/"}}, + "writableLocations": []any{map[string]any{"name": "local", "databaseAccountEndpoint": server.URL + "/"}}, + }) + return + } + if r.Method == http.MethodPost && r.URL.Path == "/dbs/sample/colls" && + strings.Contains(r.Header.Get("Content-Type"), "application/query+json") { + pageRequests++ + w.Header().Set("x-ms-continuation", "next-page") + _, _ = io.WriteString(w, `{"DocumentCollections":[{"id":"first"}]}`) + return + } + http.Error(w, r.Method+" "+r.URL.Path, http.StatusNotFound) + })) + defer server.Close() + credential, err := azcosmos.NewKeyCredential(base64.StdEncoding.EncodeToString([]byte("local-test-key"))) + if err != nil { + t.Fatal(err) + } + client, err := azcosmos.NewClientWithKey(server.URL, credential, &azcosmos.ClientOptions{ + ClientOptions: azcore.ClientOptions{Transport: server.Client(), InsecureAllowCredentialWithHTTP: true}, + }) + if err != nil { + t.Fatal(err) + } + store, err := azureprovider.NewCosmosDB(client, "sample", azureprovider.CosmosDBConfig{}) + if err != nil { + t.Fatal(err) + } + for name, err := range store.ListCollectionNames(t.Context()) { + if err != nil { + t.Fatal(err) + } + if name != "first" { + t.Fatalf("first collection name = %q", name) + } + break + } + if pageRequests != 1 { + t.Fatalf("query pages requested after early stop = %d, want 1", pageRequests) + } +} + +func TestCosmosDBRejectsInvalidConfigAndPartitionModel(t *testing.T) { + if _, err := azureprovider.NewCosmosDB(nil, "sample", azureprovider.CosmosDBConfig{}); err == nil { + t.Fatal("expected nil client error") + } + server := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {})) + defer server.Close() + credential, err := azcosmos.NewKeyCredential(base64.StdEncoding.EncodeToString([]byte("test"))) + if err != nil { + t.Fatal(err) + } + client, err := azcosmos.NewClientWithKey(server.URL, credential, &azcosmos.ClientOptions{ + ClientOptions: azcore.ClientOptions{ + Transport: server.Client(), + InsecureAllowCredentialWithHTTP: true, + }, + }) + if err != nil { + t.Fatal(err) + } + store, err := azureprovider.NewCosmosDB(client, "sample", azureprovider.CosmosDBConfig{ + PartitionKeyProperties: []string{"Category"}, + }) + if err != nil { + t.Fatal(err) + } + if _, err := vectordata.GetCollection[uuid.UUID, *cosmosProtocolRecord](store.Store, "vectors", nil); err == nil || !strings.Contains(err.Error(), "CosmosKey") { + t.Fatalf("expected custom partition key type error, got %v", err) + } +} + +func TestCosmosDBCompositePartitionKeyDefinitionAndOperations(t *testing.T) { + var createBody map[string]any + var server *httptest.Server + server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodGet && r.URL.Path == "/" { + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]any{ + "id": "mock", "enableMultipleWriteLocations": false, + "readableLocations": []any{map[string]any{"name": "local", "databaseAccountEndpoint": server.URL + "/"}}, + "writableLocations": []any{map[string]any{"name": "local", "databaseAccountEndpoint": server.URL + "/"}}, + }) + return + } + + if r.Method == http.MethodPost && r.URL.Path == "/dbs/sample/colls" { + _ = json.NewDecoder(r.Body).Decode(&createBody) + w.WriteHeader(http.StatusCreated) + _, _ = io.WriteString(w, `{"id":"vectors","partitionKey":{"kind":"MultiHash","paths":["/Category","/Tenant"],"version":2}}`) + return + } + if r.Method == http.MethodPost && r.URL.Path == "/dbs/sample/colls/vectors/docs" && strings.Contains(r.Header.Get("Content-Type"), "application/json") { + var document map[string]any + _ = json.NewDecoder(r.Body).Decode(&document) + if got := r.Header.Get("x-ms-documentdb-partitionkey"); got != `["engineering","tenant-1"]` { + http.Error(w, "partition key "+got, http.StatusBadRequest) + return + } + w.WriteHeader(http.StatusOK) + _ = json.NewEncoder(w).Encode(document) + return + } + http.Error(w, r.Method+" "+r.URL.Path, http.StatusNotFound) + })) + defer server.Close() + credential, err := azcosmos.NewKeyCredential(base64.StdEncoding.EncodeToString([]byte("local-test-key"))) + if err != nil { + t.Fatal(err) + } + client, err := azcosmos.NewClientWithKey(server.URL, credential, &azcosmos.ClientOptions{ + ClientOptions: azcore.ClientOptions{Transport: server.Client(), InsecureAllowCredentialWithHTTP: true}, + }) + if err != nil { + t.Fatal(err) + } + store, err := azureprovider.NewCosmosDB(client, "sample", azureprovider.CosmosDBConfig{ + PartitionKeyProperties: []string{"Category", "Tenant"}, + }) + if err != nil { + t.Fatal(err) + } + type record struct { + ID string `vectordata:"key"` + Category string `vectordata:"data"` + Tenant string `vectordata:"data"` + Vector []float32 `vectordata:"vector,dimensions=2"` + } + collection, err := vectordata.GetCollection[azureprovider.CosmosKey, record](store.Store, "vectors", nil) + if err != nil { + t.Fatal(err) + } + if err := collection.EnsureExists(t.Context()); err != nil { + t.Fatal(err) + } + partitionDefinition, ok := createBody["partitionKey"].(map[string]any) + if !ok || partitionDefinition["kind"] != "MultiHash" { + t.Fatalf("partition key definition = %#v", createBody["partitionKey"]) + } + if err := collection.Upsert(t.Context(), record{ + ID: "item-1", Category: "engineering", Tenant: "tenant-1", Vector: []float32{0.2, 0.3}, + }); err != nil { + t.Fatal(err) + } +} + +func TestCosmosDBGeneratesEmbeddingsOncePerBatch(t *testing.T) { + var stored []map[string]any + var server *httptest.Server + server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + if r.Method == http.MethodGet && r.URL.Path == "/" { + _ = json.NewEncoder(w).Encode(map[string]any{ + "id": "mock", "enableMultipleWriteLocations": false, + "readableLocations": []any{map[string]any{"name": "local", "databaseAccountEndpoint": server.URL + "/"}}, + "writableLocations": []any{map[string]any{"name": "local", "databaseAccountEndpoint": server.URL + "/"}}, + }) + return + } + if r.Method == http.MethodPost && r.URL.Path == "/dbs/sample/colls/vectors/docs" && + strings.Contains(r.Header.Get("Content-Type"), "application/query+json") { + _, _ = io.WriteString(w, `{"Documents":[]}`) + return + } + if r.Method == http.MethodPost && r.URL.Path == "/dbs/sample/colls/vectors/docs" { + var document map[string]any + if err := json.NewDecoder(r.Body).Decode(&document); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + stored = append(stored, document) + w.WriteHeader(http.StatusOK) + _, _ = io.WriteString(w, `{}`) + return + } + http.Error(w, r.Method+" "+r.URL.Path, http.StatusNotFound) + })) + defer server.Close() + credential, err := azcosmos.NewKeyCredential(base64.StdEncoding.EncodeToString([]byte("local-test-key"))) + if err != nil { + t.Fatal(err) + } + client, err := azcosmos.NewClientWithKey(server.URL, credential, &azcosmos.ClientOptions{ + ClientOptions: azcore.ClientOptions{Transport: server.Client(), InsecureAllowCredentialWithHTTP: true}, + }) + if err != nil { + t.Fatal(err) + } + generator := &cosmosTestEmbeddingGenerator{} + store, err := azureprovider.NewCosmosDB(client, "sample", azureprovider.CosmosDBConfig{ + EmbeddingGenerator: embedding.Bind(generator), + }) + if err != nil { + t.Fatal(err) + } + type generatedRecord struct { + ID uuid.UUID `vectordata:"key"` + Content string `vectordata:"vector,dimensions=2"` + } + collection, err := vectordata.GetCollection[uuid.UUID, *generatedRecord](store.Store, "vectors", nil) + if err != nil { + t.Fatal(err) + } + records := []*generatedRecord{{Content: "first"}, {Content: "second"}} + if err := collection.Upsert(t.Context(), records...); err != nil { + t.Fatal(err) + } + if generator.calls != 1 || len(generator.inputs) != 1 || + len(generator.inputs[0]) != 2 || generator.inputs[0][0] != "first" || generator.inputs[0][1] != "second" { + t.Fatalf("generator calls/inputs = %d, %#v", generator.calls, generator.inputs) + } + if len(stored) != 2 { + t.Fatalf("stored items = %d, want 2", len(stored)) + } + for i, document := range stored { + vector, ok := document["Content"].([]any) + if !ok || len(vector) != 2 || document["id"] != records[i].ID.String() { + t.Fatalf("stored generated document = %#v", document) + } + } + for _, err := range collection.Search(t.Context(), "query", 1, &vectordata.SearchOptions{ + Filter: &vectordata.Filter{Operator: vectordata.Equal, Property: "ID", Value: records[0].ID}, + }) { + if err != nil { + t.Fatal(err) + } + } + if generator.calls != 2 || len(generator.inputs) != 2 || + len(generator.inputs[1]) != 1 || generator.inputs[1][0] != "query" { + t.Fatalf("search generator calls/inputs = %d, %#v", generator.calls, generator.inputs) + } +} + +func TestCosmosDBDynamicRecordMapping(t *testing.T) { + var stored map[string]any + var server *httptest.Server + server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + if r.Method == http.MethodGet && r.URL.Path == "/" { + _ = json.NewEncoder(w).Encode(map[string]any{ + "id": "mock", "enableMultipleWriteLocations": false, + "readableLocations": []any{map[string]any{"name": "local", "databaseAccountEndpoint": server.URL + "/"}}, + "writableLocations": []any{map[string]any{"name": "local", "databaseAccountEndpoint": server.URL + "/"}}, + }) + return + } + if r.Method == http.MethodPost && r.URL.Path == "/dbs/sample/colls/dynamic/docs" { + if err := json.NewDecoder(r.Body).Decode(&stored); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + w.WriteHeader(http.StatusOK) + _ = json.NewEncoder(w).Encode(stored) + return + } + if r.Method == http.MethodGet && strings.HasPrefix(r.URL.Path, "/dbs/sample/colls/dynamic/docs/") { + _ = json.NewEncoder(w).Encode(stored) + return + } + http.Error(w, r.Method+" "+r.URL.Path, http.StatusNotFound) + })) + defer server.Close() + credential, err := azcosmos.NewKeyCredential(base64.StdEncoding.EncodeToString([]byte("local-test-key"))) + if err != nil { + t.Fatal(err) + } + client, err := azcosmos.NewClientWithKey(server.URL, credential, &azcosmos.ClientOptions{ + ClientOptions: azcore.ClientOptions{Transport: server.Client(), InsecureAllowCredentialWithHTTP: true}, + }) + if err != nil { + t.Fatal(err) + } + store, err := azureprovider.NewCosmosDB(client, "sample", azureprovider.CosmosDBConfig{}) + if err != nil { + t.Fatal(err) + } + definition := &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]()}, + }} + collection, err := store.DynamicCollection("dynamic", definition) + if err != nil { + t.Fatal(err) + } + record := map[string]any{"Text": "dynamic"} + if err := collection.Upsert(t.Context(), record); err != nil { + t.Fatal(err) + } + id, ok := record["ID"].(uuid.UUID) + if !ok || id == uuid.Nil || stored["id"] != id.String() || stored["text"] != "dynamic" { + t.Fatalf("dynamic record after upsert = %#v; stored document = %#v", record, stored) + } + result, found, err := collection.Get(t.Context(), id, nil) + if err != nil || !found { + t.Fatalf("dynamic Get() = (%#v, %v, %v)", result, found, err) + } + if result["ID"] != id || result["Text"] != "dynamic" { + t.Fatalf("dynamic Get() record = %#v", result) + } +} + +type cosmosTestQuery struct { + Query string `json:"query"` + Parameters []struct { + Name string `json:"name"` + Value json.RawMessage `json:"value"` + } `json:"parameters"` + partition string +} + +func (q cosmosTestQuery) parameter(name string) string { + for _, parameter := range q.Parameters { + if parameter.Name == name { + return string(parameter.Value) + } + } + return "" +} + +type cosmosTestProtocol struct { + mu sync.Mutex + requests int + containers []azcosmos.ContainerProperties + documents map[string]map[string]json.RawMessage + partitions []string + queries []cosmosTestQuery +} + +func newCosmosTestStore(t *testing.T, config azureprovider.CosmosDBConfig, queryHandler http.HandlerFunc) (*azureprovider.CosmosDB, *cosmosTestProtocol) { + t.Helper() + protocol := &cosmosTestProtocol{documents: make(map[string]map[string]json.RawMessage)} + var server *httptest.Server + server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + protocol.mu.Lock() + protocol.requests++ + protocol.mu.Unlock() + w.Header().Set("Content-Type", "application/json") + writeJSON := func(value any) { + if err := json.NewEncoder(w).Encode(value); err != nil { + t.Errorf("encode mock response: %v", err) + } + } + switch { + case r.Method == http.MethodGet && r.URL.Path == "/": + writeJSON(map[string]any{ + "id": "mock", "enableMultipleWriteLocations": false, + "readableLocations": []any{map[string]any{"name": "local", "databaseAccountEndpoint": server.URL + "/"}}, + "writableLocations": []any{map[string]any{"name": "local", "databaseAccountEndpoint": server.URL + "/"}}, + }) + case r.Method == http.MethodPost && strings.Contains(r.Header.Get("Content-Type"), "application/query+json"): + var query cosmosTestQuery + if err := json.NewDecoder(r.Body).Decode(&query); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + query.partition = r.Header.Get("x-ms-documentdb-partitionkey") + protocol.mu.Lock() + protocol.queries = append(protocol.queries, query) + protocol.mu.Unlock() + if queryHandler != nil { + queryHandler(w, r) + } else { + writeJSON(map[string]any{"Documents": []any{}}) + } + case r.Method == http.MethodPost && r.URL.Path == "/dbs/sample/colls": + var properties azcosmos.ContainerProperties + if err := json.NewDecoder(r.Body).Decode(&properties); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + protocol.mu.Lock() + protocol.containers = append(protocol.containers, properties) + protocol.mu.Unlock() + w.WriteHeader(http.StatusCreated) + writeJSON(properties) + case r.Method == http.MethodPost && strings.HasSuffix(r.URL.Path, "/docs"): + var document map[string]json.RawMessage + if err := json.NewDecoder(r.Body).Decode(&document); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + var id string + if err := json.Unmarshal(document["id"], &id); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + protocol.mu.Lock() + protocol.documents[id] = document + protocol.partitions = append(protocol.partitions, r.Header.Get("x-ms-documentdb-partitionkey")) + protocol.mu.Unlock() + writeJSON(document) + case strings.Contains(r.URL.Path, "/docs/") && (r.Method == http.MethodGet || r.Method == http.MethodDelete): + id := r.URL.Path[strings.LastIndex(r.URL.Path, "/")+1:] + protocol.mu.Lock() + document, found := protocol.documents[id] + protocol.partitions = append(protocol.partitions, r.Header.Get("x-ms-documentdb-partitionkey")) + if r.Method == http.MethodDelete { + delete(protocol.documents, id) + } + protocol.mu.Unlock() + if !found { + w.WriteHeader(http.StatusNotFound) + writeJSON(map[string]string{"code": "NotFound", "message": "missing"}) + } else if r.Method == http.MethodDelete { + w.WriteHeader(http.StatusNoContent) + } else { + writeJSON(document) + } + default: + http.Error(w, r.Method+" "+r.URL.Path, http.StatusNotFound) + } + })) + t.Cleanup(server.Close) + credential, err := azcosmos.NewKeyCredential(base64.StdEncoding.EncodeToString([]byte("local-test-key"))) + if err != nil { + t.Fatal(err) + } + client, err := azcosmos.NewClientWithKey(server.URL, credential, &azcosmos.ClientOptions{ + ClientOptions: azcore.ClientOptions{Transport: server.Client(), InsecureAllowCredentialWithHTTP: true}, + }) + if err != nil { + t.Fatal(err) + } + t.Cleanup(client.Close) + store, err := azureprovider.NewCosmosDB(client, "sample", config) + if err != nil { + t.Fatal(err) + } + return store, protocol +} + +func TestCosmosDBByteVectorsUseNumericJSONArrays(t *testing.T) { + for _, value := range []any{ + []uint8{0, 255}, + embedding.VectorOf[uint8]{0, 255}, + embedding.Of[uint8]{Vector: []uint8{0, 255}}, + &embedding.Of[uint8]{Vector: []uint8{0, 255}}, + } { + t.Run(reflect.TypeOf(value).String(), func(t *testing.T) { + store, protocol := newCosmosTestStore(t, azureprovider.CosmosDBConfig{}, nil) + definition := &vectordata.Definition{Properties: []vectordata.Property{ + &vectordata.KeyProperty{Name: "ID", Type: reflect.TypeFor[string]()}, + &vectordata.VectorProperty{Name: "Vector", Type: reflect.TypeOf(value), Dimensions: 2}, + }} + collection, err := store.DynamicCollection("records", definition) + if err != nil { + t.Fatal(err) + } + record := map[string]any{"ID": "item-1", "Vector": value} + if err := collection.Upsert(t.Context(), record); err != nil { + t.Fatal(err) + } + protocol.mu.Lock() + stored := string(protocol.documents["item-1"]["Vector"]) + protocol.mu.Unlock() + if stored != "[0,255]" { + t.Fatalf("stored byte vector = %s, want [0,255]", stored) + } + got, found, err := collection.Get(t.Context(), "item-1", &vectordata.RecordRetrievalOptions{IncludeVectors: true}) + if err != nil || !found || !reflect.DeepEqual(got["Vector"], value) { + t.Fatalf("byte-vector round trip = (%#v, %v, %v)", got, found, err) + } + for _, err := range collection.Search(t.Context(), value, 1, &vectordata.SearchOptions{ + Filter: &vectordata.Filter{Operator: vectordata.Equal, Property: "ID", Value: "item-1"}, + }) { + if err != nil { + t.Fatal(err) + } + } + protocol.mu.Lock() + query := protocol.queries[0] + protocol.mu.Unlock() + if query.parameter("@vector") != "[0,255]" || query.partition != `["item-1"]` { + t.Fatalf("byte-vector query = %#v; @vector = %s", query, query.parameter("@vector")) + } + if !reflect.DeepEqual(record["Vector"], value) { + t.Fatal("upsert changed the input vector") + } + }) + } +} + +type cosmosTestByteEmbeddingGenerator struct { + calls int +} + +func (g *cosmosTestByteEmbeddingGenerator) Generate(_ context.Context, values []string, _ *embedding.Options) (*embedding.Generated[*embedding.Of[uint8]], error) { + g.calls++ + result := &embedding.Generated[*embedding.Of[uint8]]{Embeddings: make([]*embedding.Of[uint8], len(values))} + for i := range values { + result.Embeddings[i] = &embedding.Of[uint8]{Vector: []uint8{7, 255}} + } + return result, nil +} + +func TestCosmosDBGeneratedByteVectorsUseNumericJSONArrays(t *testing.T) { + generator := &cosmosTestByteEmbeddingGenerator{} + store, protocol := newCosmosTestStore(t, azureprovider.CosmosDBConfig{EmbeddingGenerator: embedding.Bind(generator)}, nil) + type record struct { + ID string `vectordata:"key"` + Content string `vectordata:"vector,dimensions=2"` + } + collection, err := vectordata.GetCollection[string, record](store.Store, "records", nil) + if err != nil { + t.Fatal(err) + } + if err := collection.Upsert(t.Context(), record{ID: "item-1", Content: "hello"}); err != nil { + t.Fatal(err) + } + protocol.mu.Lock() + vector := string(protocol.documents["item-1"]["Content"]) + protocol.mu.Unlock() + if vector != "[7,255]" { + t.Fatalf("generated byte-vector JSON = %s", vector) + } + for _, err := range collection.Search(t.Context(), "query", 1, &vectordata.SearchOptions{ + Filter: &vectordata.Filter{Operator: vectordata.Equal, Property: "ID", Value: "item-1"}, + }) { + if err != nil { + t.Fatal(err) + } + } + protocol.mu.Lock() + query := protocol.queries[0] + protocol.mu.Unlock() + if query.parameter("@vector") != "[7,255]" || generator.calls != 2 { + t.Fatalf("generated search: vector=%s, generator calls=%d", query.parameter("@vector"), generator.calls) + } +} + +func TestCosmosDBCompositeLookupKeysPreserveScalarRecordIDs(t *testing.T) { + for _, keyType := range []reflect.Type{reflect.TypeFor[string](), reflect.TypeFor[uuid.UUID]()} { + t.Run(keyType.String(), func(t *testing.T) { + store, protocol := newCosmosTestStore(t, azureprovider.CosmosDBConfig{PartitionKeyProperties: []string{"Tenant"}}, nil) + definition := &vectordata.Definition{Properties: []vectordata.Property{ + &vectordata.KeyProperty{Name: "ID", Type: keyType}, + &vectordata.DataProperty{Name: "Tenant", Type: reflect.TypeFor[string]()}, + }} + collection, err := vectordata.GetCollection[azureprovider.CosmosKey, map[string]any](store.Store, "records", definition) + if err != nil { + t.Fatal(err) + } + record := map[string]any{"Tenant": "tenant-1"} + if keyType == reflect.TypeFor[string]() { + record["ID"] = "item-1" + } + if err := collection.Upsert(t.Context(), record); err != nil { + t.Fatal(err) + } + id := "item-1" + if keyType == reflect.TypeFor[uuid.UUID]() { + generated, ok := record["ID"].(uuid.UUID) + if !ok || generated == uuid.Nil { + t.Fatalf("generated record ID = %#v", record["ID"]) + } + id = generated.String() + } + key := azureprovider.CosmosKey{DocumentID: id, PartitionKey: []any{"tenant-1"}} + got, found, err := collection.Get(t.Context(), key, nil) + if err != nil || !found || got["ID"] != record["ID"] || got["Tenant"] != "tenant-1" { + t.Fatalf("composite-key Get = (%#v, %v, %v)", got, found, err) + } + if err := collection.Delete(t.Context(), key); err != nil { + t.Fatal(err) + } + _, found, err = collection.Get(t.Context(), key, nil) + if err != nil || found { + t.Fatalf("Get after Delete = (found=%v, err=%v)", found, err) + } + protocol.mu.Lock() + partitions := append([]string(nil), protocol.partitions...) + protocol.mu.Unlock() + if len(partitions) != 4 { + t.Fatalf("point-operation partition headers = %#v", partitions) + } + for _, partition := range partitions { + if partition != `["tenant-1"]` { + t.Fatalf("point operation used wrong partition: %s", partition) + } + } + }) + } +} + +func TestCosmosDBPartitionValuesMatchDocumentJSONAndRoundTrip(t *testing.T) { + text, number, boolean := "tenant-1", float32(0.1), true + tests := []struct { + name string + typ reflect.Type + value any + partition string + }{ + {"float32", reflect.TypeFor[float32](), float32(0.1), "[0.1]"}, + {"float64", reflect.TypeFor[float64](), 0.1, "[0.1]"}, + {"integer", reflect.TypeFor[int64](), int64(42), "[42]"}, + {"string-pointer", reflect.TypeFor[*string](), &text, `["tenant-1"]`}, + {"float32-pointer", reflect.TypeFor[*float32](), &number, "[0.1]"}, + {"bool-pointer", reflect.TypeFor[*bool](), &boolean, "[true]"}, + {"nil-pointer", reflect.TypeFor[*string](), (*string)(nil), "[null]"}, + {"dynamic-null", reflect.TypeFor[string](), nil, "[null]"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + store, protocol := newCosmosTestStore(t, azureprovider.CosmosDBConfig{PartitionKeyProperties: []string{"Tenant"}}, nil) + definition := &vectordata.Definition{Properties: []vectordata.Property{ + &vectordata.KeyProperty{Name: "ID", Type: reflect.TypeFor[string]()}, + &vectordata.DataProperty{Name: "Tenant", StorageName: "tenant", Type: tt.typ}, + }} + collection, err := store.DynamicCollection("records", definition) + if err != nil { + t.Fatal(err) + } + if err := collection.Upsert(t.Context(), map[string]any{"ID": "item-1", "Tenant": tt.value}); err != nil { + t.Fatal(err) + } + protocol.mu.Lock() + partition := protocol.partitions[0] + stored := append(json.RawMessage(nil), protocol.documents["item-1"]["tenant"]...) + protocol.mu.Unlock() + if partition != tt.partition || "["+string(stored)+"]" != tt.partition { + t.Fatalf("partition header=%s, document tenant=%s, want %s", partition, stored, tt.partition) + } + key := azureprovider.CosmosKey{DocumentID: "item-1", PartitionKey: []any{tt.value}} + got, found, err := collection.Get(t.Context(), key, nil) + if err != nil || !found || got["ID"] != "item-1" { + t.Fatalf("partitioned Get = (%#v, %v, %v)", got, found, err) + } + tenantJSON, err := json.Marshal(got["Tenant"]) + if err != nil || "["+string(tenantJSON)+"]" != tt.partition { + t.Fatalf("retrieved tenant JSON=%s, err=%v", tenantJSON, err) + } + for _, err := range collection.GetFiltered(t.Context(), vectordata.Filter{ + Operator: vectordata.Equal, Property: "Tenant", Value: tt.value, + }, 1, &vectordata.FilteredRecordRetrievalOptions{OrderBy: []vectordata.Sort{{Property: "ID"}}}) { + if err != nil { + t.Fatal(err) + } + } + protocol.mu.Lock() + query := protocol.queries[0] + protocol.mu.Unlock() + if query.partition != tt.partition { + t.Fatalf("query partition=%s, want %s", query.partition, tt.partition) + } + }) + } +} + +func TestCosmosDBNilNativeVectorsAreNull(t *testing.T) { + values := []any{ + []float32(nil), []uint8(nil), []int8(nil), + embedding.VectorOf[float32](nil), embedding.VectorOf[uint8](nil), embedding.VectorOf[int8](nil), + embedding.Of[float32]{}, + embedding.Of[uint8]{}, + embedding.Of[int8]{}, + (*embedding.Of[float32])(nil), (*embedding.Of[uint8])(nil), (*embedding.Of[int8])(nil), + } + for _, value := range values { + t.Run(reflect.TypeOf(value).String(), func(t *testing.T) { + store, protocol := newCosmosTestStore(t, azureprovider.CosmosDBConfig{}, nil) + definition := &vectordata.Definition{Properties: []vectordata.Property{ + &vectordata.KeyProperty{Name: "ID", Type: reflect.TypeFor[string]()}, + &vectordata.VectorProperty{Name: "Vector", Type: reflect.TypeOf(value), Dimensions: 2}, + }} + collection, err := store.DynamicCollection("records", definition) + if err != nil { + t.Fatal(err) + } + if err := collection.Upsert(t.Context(), map[string]any{"ID": "item-1", "Vector": value}); err != nil { + t.Fatal(err) + } + protocol.mu.Lock() + vector := string(protocol.documents["item-1"]["Vector"]) + protocol.mu.Unlock() + if vector != "null" { + t.Fatalf("nil native vector JSON = %s", vector) + } + if _, found, err := collection.Get(t.Context(), "item-1", &vectordata.RecordRetrievalOptions{IncludeVectors: true}); err != nil || !found { + t.Fatalf("null-vector retrieval = (found=%v, err=%v)", found, err) + } + }) + } +} + +func TestCosmosDBEmptyNativeVectorsRemainInvalid(t *testing.T) { + for _, value := range []any{[]float32{}, []uint8{}, []int8{}} { + t.Run(reflect.TypeOf(value).String(), func(t *testing.T) { + store, protocol := newCosmosTestStore(t, azureprovider.CosmosDBConfig{}, nil) + definition := &vectordata.Definition{Properties: []vectordata.Property{ + &vectordata.KeyProperty{Name: "ID", Type: reflect.TypeFor[string]()}, + &vectordata.VectorProperty{Name: "Vector", Type: reflect.TypeOf(value), Dimensions: 2}, + }} + collection, err := store.DynamicCollection("records", definition) + if err != nil { + t.Fatal(err) + } + err = collection.Upsert(t.Context(), map[string]any{"ID": "item-1", "Vector": value}) + if err == nil || !strings.Contains(err.Error(), "requires 2 dimensions") { + t.Fatalf("empty vector upsert error = %v", err) + } + var searchErr error + for _, err := range collection.Search(t.Context(), value, 1, &vectordata.SearchOptions{ + Filter: &vectordata.Filter{Operator: vectordata.Equal, Property: "ID", Value: "item-1"}, + }) { + searchErr = err + } + if searchErr == nil || !strings.Contains(searchErr.Error(), "dimension 2") { + t.Fatalf("empty search vector error = %v", searchErr) + } + protocol.mu.Lock() + requests := protocol.requests + protocol.mu.Unlock() + if requests != 0 { + t.Fatalf("invalid vectors performed %d HTTP requests", requests) + } + }) + } +} + +func TestCosmosDBNamedStringKeys(t *testing.T) { + type recordID string + type record struct { + ID recordID `vectordata:"key"` + } + for _, partitionProperties := range [][]string{nil, {"ID"}} { + t.Run(strings.Join(partitionProperties, ","), func(t *testing.T) { + store, protocol := newCosmosTestStore(t, azureprovider.CosmosDBConfig{PartitionKeyProperties: partitionProperties}, nil) + collection, err := vectordata.GetCollection[recordID, record](store.Store, "records", nil) + if err != nil { + t.Fatal(err) + } + if err := collection.Upsert(t.Context(), record{ID: "item-1"}); err != nil { + t.Fatal(err) + } + got, found, err := collection.Get(t.Context(), recordID("item-1"), nil) + if err != nil || !found || got.ID != recordID("item-1") { + t.Fatalf("named-key Get = (%#v, %v, %v)", got, found, err) + } + if err := collection.Delete(t.Context(), recordID("item-1")); err != nil { + t.Fatal(err) + } + protocol.mu.Lock() + partitions := append([]string(nil), protocol.partitions...) + protocol.mu.Unlock() + for _, partition := range partitions { + if partition != `["item-1"]` { + t.Fatalf("named-key partition = %s", partition) + } + } + }) + } +} + +func TestCosmosDBKeyMembershipFiltersNormalizeElements(t *testing.T) { + id := uuid.MustParse("aa11aa11-aa11-4a11-8a11-aa11aa11aa11") + tests := []struct { + name string + keys any + want string + }{ + {"composite-slice", []azureprovider.CosmosKey{{DocumentID: "item-1", PartitionKey: []any{"tenant-1"}}}, `["item-1"]`}, + {"composite-array", [1]azureprovider.CosmosKey{{DocumentID: "item-1", PartitionKey: []any{"tenant-1"}}}, `["item-1"]`}, + {"uuid", []uuid.UUID{id}, `["aa11aa11-aa11-4a11-8a11-aa11aa11aa11"]`}, + {"empty", []string{}, `[]`}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + store, protocol := newCosmosTestStore(t, azureprovider.CosmosDBConfig{}, nil) + type record struct { + ID string `vectordata:"key"` + } + collection, err := vectordata.GetCollection[string, record](store.Store, "records", nil) + if err != nil { + t.Fatal(err) + } + filter := vectordata.Filter{Operator: vectordata.In, Property: "ID", Value: tt.keys} + before, err := json.Marshal(filter) + if err != nil { + t.Fatal(err) + } + for _, err := range collection.GetFiltered(t.Context(), filter, 1, nil) { + if err != nil { + t.Fatal(err) + } + } + + protocol.mu.Lock() + query := protocol.queries[0] + protocol.mu.Unlock() + if query.parameter("@p0") != tt.want || !strings.Contains(query.Query, `ARRAY_CONTAINS(@p0, c["id"])`) { + t.Fatalf("membership query = %#v, parameter = %s", query, query.parameter("@p0")) + } + after, err := json.Marshal(filter) + if err != nil || string(after) != string(before) { + t.Fatalf("filter was changed: before=%s, after=%s, err=%v", before, after, err) + } + }) + } +} + +func TestCosmosDBGeneratedVectorsRejectRetrievalBeforeIO(t *testing.T) { + generator := &cosmosTestEmbeddingGenerator{} + store, protocol := newCosmosTestStore(t, azureprovider.CosmosDBConfig{EmbeddingGenerator: embedding.Bind(generator)}, nil) + type record struct { + ID string `vectordata:"key"` + Text string `vectordata:"data,fulltext"` + Content string `vectordata:"vector,dimensions=2"` + } + collection, err := vectordata.GetCollection[string, record](store.Store, "records", nil) + if err != nil { + t.Fatal(err) + } + hybrid, ok := collection.(vectordata.KeywordHybridSearcher[record]) + if !ok { + t.Fatal("missing hybrid-search capability") + } + tests := []struct { + name string + run func() error + }{ + {"get-missing", func() error { + _, _, err := collection.Get(t.Context(), "missing", &vectordata.RecordRetrievalOptions{IncludeVectors: true}) + return err + }}, + {"get-batch-missing", func() error { + for _, err := range vectordata.GetBatch(t.Context(), collection, []string{"missing"}, &vectordata.RecordRetrievalOptions{IncludeVectors: true}) { + return err + } + return nil + }}, + {"filtered", func() error { + for _, err := range collection.GetFiltered(t.Context(), vectordata.Filter{ + Operator: vectordata.Equal, Property: "ID", Value: "missing", + }, 1, &vectordata.FilteredRecordRetrievalOptions{IncludeVectors: true}) { + return err + } + return nil + }}, + {"filtered-zero-top", func() error { + for _, err := range collection.GetFiltered(t.Context(), vectordata.Filter{ + Operator: vectordata.Equal, Property: "ID", Value: "missing", + }, 0, &vectordata.FilteredRecordRetrievalOptions{IncludeVectors: true}) { + return err + } + return nil + }}, + {"search", func() error { + for _, err := range collection.Search(t.Context(), "query", 1, &vectordata.SearchOptions{IncludeVectors: true}) { + return err + } + return nil + }}, + {"hybrid-search", func() error { + for _, err := range hybrid.HybridSearch(t.Context(), "query", []string{"keyword"}, 1, &vectordata.HybridSearchOptions{IncludeVectors: true}) { + return err + } + return nil + }}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := tt.run() + if err == nil || !strings.Contains(err.Error(), "IncludeVectors") { + t.Fatalf("error = %v, want unsupported IncludeVectors", err) + } + }) + } + protocol.mu.Lock() + requests := protocol.requests + protocol.mu.Unlock() + if requests != 0 || generator.calls != 0 { + t.Fatalf("unsupported retrieval performed I/O: HTTP requests=%d, embedding calls=%d", requests, generator.calls) + } +} + +func TestCosmosDBRejectsCrossPartitionRankingBeforeIO(t *testing.T) { + generator := &cosmosTestEmbeddingGenerator{} + store, protocol := newCosmosTestStore(t, azureprovider.CosmosDBConfig{ + PartitionKeyProperties: []string{"Tenant", "Region"}, + EmbeddingGenerator: embedding.Bind(generator), + }, nil) + type record struct { + ID string `vectordata:"key"` + Tenant string `vectordata:"data"` + Region string `vectordata:"data"` + Text string `vectordata:"data,fulltext"` + Content string `vectordata:"vector,dimensions=2"` + } + collection, err := vectordata.GetCollection[azureprovider.CosmosKey, record](store.Store, "records", nil) + if err != nil { + t.Fatal(err) + } + hybrid, ok := collection.(vectordata.KeywordHybridSearcher[record]) + if !ok { + t.Fatal("missing hybrid-search capability") + } + filters := []struct { + name string + filter *vectordata.Filter + }{ + {"none", nil}, + {"partial-hierarchy", &vectordata.Filter{Operator: vectordata.Equal, Property: "Tenant", Value: "tenant-1"}}, + {"disjunction", &vectordata.Filter{Operator: vectordata.Or, Operands: []vectordata.Filter{ + {Operator: vectordata.Equal, Property: "Tenant", Value: "tenant-1"}, + {Operator: vectordata.Equal, Property: "Region", Value: "region-1"}, + }}}, + {"negation", &vectordata.Filter{Operator: vectordata.Not, Operands: []vectordata.Filter{ + {Operator: vectordata.Equal, Property: "Tenant", Value: "tenant-1"}, + }}}, + } + for _, tt := range filters { + t.Run(tt.name, func(t *testing.T) { + var searchErr error + for _, err := range collection.Search(t.Context(), "query", 1, &vectordata.SearchOptions{Filter: tt.filter}) { + searchErr = err + } + var hybridErr error + for _, err := range hybrid.HybridSearch(t.Context(), "query", []string{"keyword"}, 1, &vectordata.HybridSearchOptions{Filter: tt.filter}) { + hybridErr = err + } + for _, err := range []error{searchErr, hybridErr} { + if err == nil || !strings.Contains(err.Error(), "equality filters for every partition key") { + t.Errorf("cross-partition ranking error = %v", err) + } + } + if tt.filter != nil { + var retrievalErr error + for _, err := range collection.GetFiltered(t.Context(), *tt.filter, 1, &vectordata.FilteredRecordRetrievalOptions{ + OrderBy: []vectordata.Sort{{Property: "ID"}}, + }) { + retrievalErr = err + } + if retrievalErr == nil || !strings.Contains(retrievalErr.Error(), "equality filters for every partition key") { + t.Errorf("cross-partition ordering error = %v", retrievalErr) + } + } + }) + } + protocol.mu.Lock() + requests := protocol.requests + protocol.mu.Unlock() + if requests != 0 || generator.calls != 0 { + t.Fatalf("unsupported ranking performed I/O: HTTP requests=%d, embedding calls=%d", requests, generator.calls) + } +} + +func TestCosmosDBDefaultPartitionRejectsUnscopedSearch(t *testing.T) { + type record struct { + ID string `vectordata:"key"` + Vector []float32 `vectordata:"vector,dimensions=2"` + } + store, protocol := newCosmosTestStore(t, azureprovider.CosmosDBConfig{}, nil) + collection, err := vectordata.GetCollection[string, record](store.Store, "records", nil) + if err != nil { + t.Fatal(err) + } + var searchErr error + for _, err := range collection.Search(t.Context(), []float32{1, 2}, 5, nil) { + searchErr = err + } + if searchErr == nil || !strings.Contains(searchErr.Error(), `missing property "ID"`) { + t.Fatalf("unscoped default-partition search error = %v", searchErr) + } + protocol.mu.Lock() + requests := protocol.requests + protocol.mu.Unlock() + if requests != 0 { + t.Fatalf("unsupported search performed %d HTTP requests", requests) + } +} + +func TestCosmosDBHierarchicalQueriesUseCompletePartition(t *testing.T) { + store, protocol := newCosmosTestStore(t, azureprovider.CosmosDBConfig{PartitionKeyProperties: []string{"Tenant", "Region"}}, nil) + type record struct { + ID string `vectordata:"key"` + Tenant string `vectordata:"data,storage=tenant"` + Region int32 `vectordata:"data,storage=region"` + Text string `vectordata:"data,fulltext"` + Vector []float32 `vectordata:"vector,dimensions=2"` + } + collection, err := vectordata.GetCollection[azureprovider.CosmosKey, record](store.Store, "records", nil) + if err != nil { + t.Fatal(err) + } + filter := vectordata.Filter{Operator: vectordata.And, Operands: []vectordata.Filter{ + {Operator: vectordata.Equal, Property: "Tenant", Value: "tenant-1"}, + {Operator: vectordata.And, Operands: []vectordata.Filter{ + {Operator: vectordata.Equal, Property: "Region", Value: int32(3)}, + {Operator: vectordata.Or, Operands: []vectordata.Filter{ + {Operator: vectordata.Equal, Property: "ID", Value: "item-1"}, + {Operator: vectordata.Equal, Property: "ID", Value: "item-2"}, + }}, + }}, + }} + for _, err := range collection.Search(t.Context(), []float32{1, 2}, 3, &vectordata.SearchOptions{Filter: &filter, Skip: 1}) { + if err != nil { + t.Fatal(err) + } + } + hybrid, ok := collection.(vectordata.KeywordHybridSearcher[record]) + if !ok { + t.Fatal("missing hybrid-search capability") + } + for _, err := range hybrid.HybridSearch(t.Context(), []float32{1, 2}, []string{"keyword"}, 3, &vectordata.HybridSearchOptions{Filter: &filter}) { + if err != nil { + t.Fatal(err) + } + } + for _, err := range collection.GetFiltered(t.Context(), filter, 2, &vectordata.FilteredRecordRetrievalOptions{ + OrderBy: []vectordata.Sort{{Property: "ID", Descending: true}}, Skip: 1, + }) { + if err != nil { + t.Fatal(err) + } + } + protocol.mu.Lock() + queries := append([]cosmosTestQuery(nil), protocol.queries...) + protocol.mu.Unlock() + if len(queries) != 3 { + t.Fatalf("queries = %#v", queries) + } + for _, query := range queries { + if query.partition != `["tenant-1",3]` || !strings.Contains(query.Query, `c["tenant"]`) || !strings.Contains(query.Query, `c["region"]`) { + t.Fatalf("query lost partition routing or filter: %#v", query) + } + } + if !strings.Contains(queries[0].Query, "OFFSET 1 LIMIT 3") || + !strings.Contains(queries[1].Query, "RANK RRF") || + !strings.Contains(queries[2].Query, `ORDER BY c["id"] DESC`) { + t.Fatalf("partition-scoped query shapes = %#v", queries) + } +} + +func TestCosmosDBCrossPartitionFilteredPaging(t *testing.T) { + for _, stopEarly := range []bool{false, true} { + name := "skip-and-top" + if stopEarly { + name = "consumer-stop" + } + t.Run(name, func(t *testing.T) { + pages := 0 + store, protocol := newCosmosTestStore(t, azureprovider.CosmosDBConfig{}, func(w http.ResponseWriter, r *http.Request) { + pages++ + var body string + switch pages { + case 1: + if continuation := r.Header.Get("x-ms-continuation"); continuation != "" { + t.Errorf("first page continuation = %q", continuation) + } + w.Header().Set("x-ms-continuation", "one") + body = `{"Documents":[]}` + case 2: + if continuation := r.Header.Get("x-ms-continuation"); continuation != "one" { + t.Errorf("second page continuation = %q", continuation) + } + w.Header().Set("x-ms-continuation", "two") + body = `{"Documents":[{"id":"a","Category":"match"},{"id":"b","Category":"match"}]}` + case 3: + if continuation := r.Header.Get("x-ms-continuation"); continuation != "two" { + t.Errorf("third page continuation = %q", continuation) + } + w.Header().Set("x-ms-continuation", "three") + body = `{"Documents":[{"id":"c","Category":"match"},{"id":"d","Category":"match"}]}` + default: + t.Error("fetched an unnecessary continuation page") + body = `{"Documents":[]}` + } + if _, err := io.WriteString(w, body); err != nil { + t.Errorf("write query response: %v", err) + } + }) + type record struct { + ID string `vectordata:"key"` + Category string `vectordata:"data"` + } + collection, err := vectordata.GetCollection[string, record](store.Store, "records", nil) + if err != nil { + t.Fatal(err) + } + var ids []string + for item, err := range collection.GetFiltered(t.Context(), vectordata.Filter{ + Operator: vectordata.Equal, Property: "Category", Value: "match", + }, 2, &vectordata.FilteredRecordRetrievalOptions{Skip: 1}) { + if err != nil { + t.Fatal(err) + } + ids = append(ids, item.ID) + if stopEarly { + break + } + } + wantIDs, wantPages := []string{"b", "c"}, 3 + if stopEarly { + wantIDs, wantPages = []string{"b"}, 2 + } + if !reflect.DeepEqual(ids, wantIDs) || pages != wantPages { + t.Fatalf("filtered paging: IDs=%v, pages=%d, want %v/%d", ids, pages, wantIDs, wantPages) + } + protocol.mu.Lock() + queries := append([]cosmosTestQuery(nil), protocol.queries...) + protocol.mu.Unlock() + for _, query := range queries { + if query.partition != "" || strings.Contains(query.Query, "TOP") || strings.Contains(query.Query, "OFFSET") || + strings.Contains(query.Query, "LIMIT") || strings.Contains(query.Query, "ORDER BY") || + !strings.Contains(query.Query, `WHERE c["Category"] = @p0`) || query.parameter("@p0") != `"match"` { + t.Fatalf("unsupported or incomplete cross-partition query: %#v", query) + } + } + }) + } +} + +func TestCosmosDBFilteredPagingPropagatesFailuresAndCancellation(t *testing.T) { + for _, cancelAfterFirst := range []bool{false, true} { + name := "service-error" + if cancelAfterFirst { + name = "cancellation" + } + t.Run(name, func(t *testing.T) { + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + pages := 0 + store, _ := newCosmosTestStore(t, azureprovider.CosmosDBConfig{}, func(w http.ResponseWriter, _ *http.Request) { + pages++ + body := `{"Documents":[{"id":"item-1"}]}` + if pages == 1 { + w.Header().Set("x-ms-continuation", "next") + } else { + w.WriteHeader(http.StatusBadRequest) + body = `{"code":"BadRequest","message":"query failed"}` + } + if _, err := io.WriteString(w, body); err != nil { + t.Errorf("write response: %v", err) + } + }) + type record struct { + ID string `vectordata:"key"` + } + collection, err := vectordata.GetCollection[string, record](store.Store, "records", nil) + if err != nil { + t.Fatal(err) + } + var terminalErr error + results := 0 + for _, err := range collection.GetFiltered(ctx, vectordata.Filter{ + Operator: vectordata.NotEqual, Property: "ID", Value: "other", + }, 3, nil) { + if err != nil { + terminalErr = err + continue + } + results++ + if cancelAfterFirst { + cancel() + } + } + + if results != 1 || terminalErr == nil { + t.Fatalf("query yielded %d results with error %v", results, terminalErr) + } + if cancelAfterFirst { + if !errors.Is(terminalErr, context.Canceled) || pages != 1 { + t.Fatalf("canceled query: pages=%d, err=%v", pages, terminalErr) + } + } else { + responseErr, ok := errors.AsType[*azcore.ResponseError](terminalErr) + if !ok || responseErr.StatusCode != http.StatusBadRequest || pages != 2 { + t.Fatalf("failed query: pages=%d, err=%v", pages, terminalErr) + } + } + }) + } +} + +func TestCosmosDBRejectsReservedScoreStorageName(t *testing.T) { + type record struct { + ID string + Value float64 + Vector []float32 + } + for _, test := range []struct { + name string + property vectordata.Property + }{ + {"data", &vectordata.DataProperty{Name: "Value", StorageName: "__cosmos_score", Type: reflect.TypeFor[float64]()}}, + {"vector", &vectordata.VectorProperty{Name: "Vector", StorageName: "__cosmos_score", Type: reflect.TypeFor[[]float32](), Dimensions: 2}}, + } { + t.Run(test.name, func(t *testing.T) { + store, protocol := newCosmosTestStore(t, azureprovider.CosmosDBConfig{}, nil) + definition := &vectordata.Definition{Properties: []vectordata.Property{ + &vectordata.KeyProperty{Name: "ID", Type: reflect.TypeFor[string]()}, + test.property, + }} + _, err := vectordata.GetCollection[string, record](store.Store, "records", definition) + if err == nil || !strings.Contains(err.Error(), "__cosmos_score") { + t.Fatalf("typed reserved storage name error = %v", err) + } + if _, err := store.DynamicCollection("records", definition); err == nil || !strings.Contains(err.Error(), "__cosmos_score") { + t.Fatalf("dynamic reserved storage name error = %v", err) + } + protocol.mu.Lock() + defer protocol.mu.Unlock() + if protocol.requests != 0 { + t.Fatalf("invalid schema made %d service requests", protocol.requests) + } + }) + } +} + +func TestCosmosDBSpecialCharacterStoragePaths(t *testing.T) { + for _, test := range []struct { + name string + storage string + path string + vectorPath string + }{ + {"tilde", "field~name", `/"field~name"`, `/"field~nameVector"`}, + {"slash", "field/name", `/"field/name"`, `/"field/nameVector"`}, + {"hyphen", "field-name", `/"field-name"`, `/"field-nameVector"`}, + {"space", "field name", `/"field name"`, `/"field nameVector"`}, + {"wildcards", "field[?*]", `/"field[?*]"`, `/"field[?*]Vector"`}, + } { + t.Run(test.name, func(t *testing.T) { + store, protocol := newCosmosTestStore(t, azureprovider.CosmosDBConfig{PartitionKeyProperties: []string{"Text"}}, nil) + collection, err := store.DynamicCollection("records", &vectordata.Definition{Properties: []vectordata.Property{ + &vectordata.KeyProperty{Name: "ID", Type: reflect.TypeFor[string]()}, + &vectordata.DataProperty{Name: "Text", StorageName: test.storage, Type: reflect.TypeFor[string](), IsIndexed: true, IsFullTextIndexed: true}, + &vectordata.VectorProperty{Name: "Vector", StorageName: test.storage + "Vector", Type: reflect.TypeFor[[]float32](), Dimensions: 2}, + }}) + if err != nil { + t.Fatal(err) + } + if err := collection.EnsureExists(t.Context()); err != nil { + t.Fatal(err) + } + protocol.mu.Lock() + properties := protocol.containers[0] + protocol.mu.Unlock() + if !reflect.DeepEqual(properties.PartitionKeyDefinition.Paths, []string{test.path}) { + t.Fatalf("partition paths = %v, want %s", properties.PartitionKeyDefinition.Paths, test.path) + } + if properties.FullTextPolicy == nil || len(properties.FullTextPolicy.FullTextPaths) != 1 || + properties.FullTextPolicy.FullTextPaths[0].Path != test.path { + t.Fatalf("full-text policy = %+v", properties.FullTextPolicy) + } + if properties.VectorEmbeddingPolicy == nil || len(properties.VectorEmbeddingPolicy.VectorEmbeddings) != 1 || + properties.VectorEmbeddingPolicy.VectorEmbeddings[0].Path != test.vectorPath { + t.Fatalf("vector policy = %+v", properties.VectorEmbeddingPolicy) + } + indexes := properties.IndexingPolicy + if indexes == nil || !reflect.DeepEqual(indexes.IncludedPaths, []azcosmos.IncludedPath{{Path: test.path + "/?"}, {Path: "/*"}}) || + !reflect.DeepEqual(indexes.ExcludedPaths, []azcosmos.ExcludedPath{{Path: test.vectorPath + "/*"}}) || + len(indexes.FullTextIndexes) != 1 || indexes.FullTextIndexes[0].Path != test.path || + len(indexes.VectorIndexes) != 1 || indexes.VectorIndexes[0].Path != test.vectorPath { + t.Fatalf("index paths = %+v", indexes) + } + record := map[string]any{"ID": "item-1", "Text": "tenant-1", "Vector": []float32{1, 2}} + if err := collection.Upsert(t.Context(), record); err != nil { + t.Fatal(err) + } + key := azureprovider.CosmosKey{DocumentID: "item-1", PartitionKey: []any{"tenant-1"}} + got, found, err := collection.Get(t.Context(), key, &vectordata.RecordRetrievalOptions{IncludeVectors: true}) + if err != nil || !found || !reflect.DeepEqual(got, record) { + t.Fatalf("special-name retrieval = (%#v, %v, %v)", got, found, err) + } + protocol.mu.Lock() + defer protocol.mu.Unlock() + document := protocol.documents["item-1"] + if string(document[test.storage]) != `"tenant-1"` || string(document[test.storage+"Vector"]) != "[1,2]" || + protocol.partitions[0] != `["tenant-1"]` { + t.Fatalf("special-name document = %s, partitions = %v", document, protocol.partitions) + } + }) + } +} + +func TestCosmosDBTimestampsUseFixedWidthUTC(t *testing.T) { + type record struct { + ID string `vectordata:"key"` + Category string `vectordata:"data,indexed,fulltext"` + When time.Time `vectordata:"data,indexed"` + Nano time.Time `vectordata:"data"` + Optional *time.Time `vectordata:"data,indexed"` + Missing *time.Time `vectordata:"data"` + Times []time.Time `vectordata:"data"` + Array [2]time.Time `vectordata:"data"` + Pointer *[]time.Time `vectordata:"data"` + NilTimes []time.Time `vectordata:"data"` + Empty []time.Time `vectordata:"data"` + Vector []float32 `vectordata:"vector,dimensions=2"` + } + whole := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + nano := whole.Add(time.Nanosecond) + east, west := time.FixedZone("east", 3*60*60), time.FixedZone("west", -5*60*60) + fraction := whole.Add(500 * time.Millisecond).In(east) + const wholeJSON = `"2026-01-01T00:00:00.000000000Z"` + const fractionJSON = `"2026-01-01T00:00:00.500000000Z"` + const timesJSON = "[" + wholeJSON + "," + fractionJSON + "]" + times := []time.Time{whole.In(west), fraction} + input := &record{ + ID: "item-1", Category: "group", When: whole, Nano: nano, Optional: &fraction, + Times: times, Array: [2]time.Time{times[0], times[1]}, Pointer: ×, + Empty: []time.Time{}, Vector: []float32{1, 2}, + } + store, protocol := newCosmosTestStore(t, azureprovider.CosmosDBConfig{PartitionKeyProperties: []string{"Category"}}, nil) + collection, err := vectordata.GetCollection[azureprovider.CosmosKey, *record](store.Store, "records", nil) + if err != nil { + t.Fatal(err) + } + if err := collection.Upsert(t.Context(), input); err != nil { + t.Fatal(err) + } + protocol.mu.Lock() + document := protocol.documents[input.ID] + protocol.mu.Unlock() + for property, want := range map[string]string{ + "When": wholeJSON, "Nano": `"2026-01-01T00:00:00.000000001Z"`, + "Optional": fractionJSON, "Missing": "null", "Times": timesJSON, "Array": timesJSON, + "Pointer": timesJSON, "NilTimes": "null", "Empty": "[]", + } { + if got := string(document[property]); got != want { + t.Errorf("stored %s = %s, want %s", property, got, want) + } + } + if string(document["When"]) >= string(document["Nano"]) || string(document["Nano"]) >= string(document["Optional"]) { + t.Fatal("stored timestamps do not sort chronologically") + } + if input.Optional.Location() != east || input.Times[0].Location() != west || + input.Array[1].Location() != east || input.Pointer != × || (*input.Pointer)[1].Location() != east { + t.Fatal("timestamp normalization modified the input record") + } + key := azureprovider.CosmosKey{DocumentID: input.ID, PartitionKey: []any{"group"}} + got, found, err := collection.Get(t.Context(), key, nil) + if err != nil || !found || !got.When.Equal(whole) || !got.Nano.Equal(nano) || + got.Optional == nil || !got.Optional.Equal(fraction) || got.Missing != nil || + len(got.Times) != 2 || !got.Times[0].Equal(times[0]) || !got.Times[1].Equal(fraction) || + !got.Array[0].Equal(times[0]) || !got.Array[1].Equal(fraction) || got.Pointer == nil || + len(*got.Pointer) != 2 || !(*got.Pointer)[1].Equal(fraction) || got.NilTimes != nil || got.Empty == nil { + t.Fatalf("timestamp round trip = (%+v, %v, %v)", got, found, err) + } + for _, test := range []struct { + name string + filter vectordata.Filter + expected string + }{ + {"equal", vectordata.Filter{Operator: vectordata.Equal, Property: "When", Value: fraction}, fractionJSON}, + {"not-equal", vectordata.Filter{Operator: vectordata.NotEqual, Property: "When", Value: fraction}, fractionJSON}, + {"greater", vectordata.Filter{Operator: vectordata.GreaterThan, Property: "When", Value: fraction}, fractionJSON}, + {"greater-equal", vectordata.Filter{Operator: vectordata.GreaterThanOrEqual, Property: "When", Value: fraction}, fractionJSON}, + {"less", vectordata.Filter{Operator: vectordata.LessThan, Property: "When", Value: fraction}, fractionJSON}, + {"less-equal", vectordata.Filter{Operator: vectordata.LessThanOrEqual, Property: "When", Value: fraction}, fractionJSON}, + {"nullable", vectordata.Filter{Operator: vectordata.Equal, Property: "Optional", Value: &fraction}, fractionJSON}, + {"contains", vectordata.Filter{Operator: vectordata.Contains, Property: "Times", Value: fraction}, fractionJSON}, + {"membership", vectordata.Filter{Operator: vectordata.In, Property: "When", Value: times}, timesJSON}, + {"nullable-membership", vectordata.Filter{Operator: vectordata.In, Property: "Optional", Value: []*time.Time{&whole, nil, &fraction}}, "[" + wholeJSON + ",null," + fractionJSON + "]"}, + } { + t.Run(test.name, func(t *testing.T) { + filter := vectordata.Filter{Operator: vectordata.And, Operands: []vectordata.Filter{ + {Operator: vectordata.Equal, Property: "Category", Value: "group"}, + test.filter, + }} + protocol.mu.Lock() + start := len(protocol.queries) + protocol.mu.Unlock() + collect(t, collection.GetFiltered(t.Context(), filter, 1, + &vectordata.FilteredRecordRetrievalOptions{OrderBy: []vectordata.Sort{{Property: "When"}}})) + collect(t, collection.Search(t.Context(), []float32{1, 2}, 1, &vectordata.SearchOptions{Filter: &filter})) + hybrid := collection.(vectordata.KeywordHybridSearcher[*record]) + collect(t, hybrid.HybridSearch(t.Context(), []float32{1, 2}, []string{"group"}, 1, + &vectordata.HybridSearchOptions{Filter: &filter})) + protocol.mu.Lock() + defer protocol.mu.Unlock() + if len(protocol.queries)-start != 3 { + t.Fatalf("sent %d timestamp queries, want 3", len(protocol.queries)-start) + } + for _, query := range protocol.queries[start:] { + if query.parameter("@p1") != test.expected { + t.Errorf("timestamp filter parameter = %s, want %s", query.parameter("@p1"), test.expected) + } + } + }) + } +} + +func TestCosmosDBInvalidTimestampsFailBeforeIO(t *testing.T) { + type record struct { + ID string `vectordata:"key"` + When time.Time `vectordata:"data"` + } + store, protocol := newCosmosTestStore(t, azureprovider.CosmosDBConfig{}, nil) + collection, err := vectordata.GetCollection[string, record](store.Store, "records", nil) + if err != nil { + t.Fatal(err) + } + invalid := time.Date(10000, 1, 1, 0, 0, 0, 0, time.UTC) + if err := collection.Upsert(t.Context(), record{ID: "item-1", When: invalid}); err == nil { + t.Fatal("invalid stored timestamp was accepted") + } + errorsSeen := 0 + for _, err := range collection.GetFiltered(t.Context(), vectordata.Filter{ + Operator: vectordata.Equal, Property: "When", Value: invalid, + }, 1, nil) { + if err == nil { + t.Fatal("invalid timestamp filter yielded a record") + } + errorsSeen++ + } + protocol.mu.Lock() + defer protocol.mu.Unlock() + if errorsSeen != 1 || protocol.requests != 0 { + t.Fatalf("invalid timestamp: errors=%d, service requests=%d", errorsSeen, protocol.requests) + } +} + +func TestCosmosDBFilteredRetrievalProjectsRequestedFields(t *testing.T) { + for _, generated := range []bool{false, true} { + name := "native" + if generated { + name = "generated" + } + t.Run(name, func(t *testing.T) { + for _, options := range []*vectordata.FilteredRecordRetrievalOptions{ + nil, {Skip: 1}, {Skip: 1, IncludeVectors: true}, + } { + if generated && options != nil && options.IncludeVectors { + continue + } + generator := &cosmosTestEmbeddingGenerator{} + store, protocol := newCosmosTestStore(t, azureprovider.CosmosDBConfig{ + EmbeddingGenerator: embedding.Bind(generator), + }, func(w http.ResponseWriter, _ *http.Request) { + body := `{"Documents":[{"id":"first","text":"one"},{"id":"second","text":"two"}]}` + if options != nil && options.IncludeVectors { + body = `{"Documents":[{"id":"first","text":"one","vector":[1,2]},{"id":"second","text":"two","vector":[3,4]}]}` + } + if _, err := io.WriteString(w, body); err != nil { + t.Errorf("write projection response: %v", err) + } + }) + vectorType := reflect.TypeFor[[]float32]() + if generated { + vectorType = reflect.TypeFor[string]() + } + collection, err := store.DynamicCollection("records", &vectordata.Definition{Properties: []vectordata.Property{ + &vectordata.KeyProperty{Name: "ID", Type: reflect.TypeFor[string]()}, + &vectordata.DataProperty{Name: "Text", StorageName: "text", Type: reflect.TypeFor[string]()}, + &vectordata.VectorProperty{Name: "Vector", StorageName: "vector", Type: vectorType, Dimensions: 2}, + }}) + if err != nil { + t.Fatal(err) + } + results := collect(t, collection.GetFiltered(t.Context(), vectordata.Filter{ + Operator: vectordata.NotEqual, Property: "ID", Value: "other", + }, 1, options)) + wantID, wantText := "first", "one" + includeVectors := options != nil && options.IncludeVectors + if options != nil { + wantID, wantText = "second", "two" + } + if len(results) != 1 || results[0]["ID"] != wantID || results[0]["Text"] != wantText { + t.Fatalf("projected result = %#v", results) + } + if vector, found := results[0]["Vector"]; found != includeVectors || + (includeVectors && !reflect.DeepEqual(vector, []float32{3, 4})) { + t.Fatalf("projected vector = %#v, present=%v, include=%v", vector, found, includeVectors) + } + protocol.mu.Lock() + query := protocol.queries[0] + protocol.mu.Unlock() + projection, _, found := strings.Cut(query.Query, " FROM c") + if !found || !strings.Contains(projection, `c["id"]`) || !strings.Contains(projection, `c["text"]`) || + strings.Contains(projection, `c["vector"]`) != includeVectors || strings.Contains(projection, "VALUE c") { + t.Fatalf("filtered projection = %s, include=%v", projection, includeVectors) + } + if generator.calls != 0 { + t.Fatalf("filtered retrieval generated embeddings %d times", generator.calls) + } + } + }) + } +}