feat(azureprovider): add Cosmos DB vector store

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)
+	}
+}

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

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=