Python: Add Azure Cosmos DB NoSQL vector store connector

python/AGENTS.md

@@ -112,7 +112,7 @@ python/
 - [foundry](packages/foundry/README.md) - Microsoft Foundry chat, agent, memory, and embedding integrations
 - [azure-contentunderstanding](packages/azure-contentunderstanding/AGENTS.md) - Azure Content Understanding context provider
 - [azure-ai-search](packages/azure-ai-search/AGENTS.md) - Azure AI Search RAG
-- [azure-cosmos](packages/azure-cosmos/AGENTS.md) - Azure Cosmos DB-backed history provider
+- [azure-cosmos](packages/azure-cosmos/AGENTS.md) - Azure Cosmos DB for NoSQL vector, history, and checkpoint integrations
 
 Durable Task and Azure Functions integrations are maintained in the [Durable Agent Framework extension](https://github.com/microsoft/agent-framework-durable-extension).
 

python/packages/azure-cosmos/AGENTS.md

@@ -1,10 +1,14 @@
 # Azure Cosmos DB Package (agent-framework-azure-cosmos)
 
-Azure Cosmos DB history provider integration for Agent Framework.
+Azure Cosmos DB for NoSQL integrations for Agent Framework.
 
 ## Main Classes
 
 - **`CosmosHistoryProvider`** - Persistent conversation history storage backed by Azure Cosmos DB
+- **`CosmosCheckpointStorage`** - Workflow checkpoint storage backed by Azure Cosmos DB
+- **`CosmosCollection`** - Vector collection using NoSQL `VectorDistance` queries
+- **`CosmosStore`** - Factory and database administration for vector collections
+- **`AzureCosmosSettings`** - Shared vector connector settings shape
 
 ## Usage
 
@@ -21,10 +25,17 @@ provider = CosmosHistoryProvider(
 
 Container name is configured on the provider. `session_id` is used as the partition key.
 
+Vector collections require an application-provided string key stored as `id` and a
+container partition key of `/id`. Container vector and indexing policies are created
+only by `ensure_collection_exists()` and validated before use. Existing incompatible
+containers are never updated or recreated.
+Supported vector element types are float32, int8, and uint8. Euclidean search is supported,
+but Euclidean score thresholds fail before service I/O.
+
 ## Import Path
 
 ```python
-from agent_framework.azure import CosmosHistoryProvider
+from agent_framework.azure import CosmosCollection, CosmosHistoryProvider, CosmosStore
 # or directly:
-from agent_framework_azure_cosmos import CosmosHistoryProvider
+from agent_framework_azure_cosmos import CosmosCollection, CosmosHistoryProvider, CosmosStore
 ```

python/packages/azure-cosmos/README.md

@@ -1,126 +1,171 @@
-# Get Started with Microsoft Agent Framework Azure Cosmos DB
+# Microsoft Agent Framework integrations for Azure Cosmos DB
 
-Please install this package via pip:
+Use Azure Cosmos DB for NoSQL as an Agent Framework vector store, conversation history
+provider, or workflow checkpoint store.
+
+## Contents
+
+- `CosmosCollection` and `CosmosStore` for vector CRUD and `VectorDistance` search
+- `CosmosHistoryProvider` for persistent conversation history
+- `CosmosCheckpointStorage` for durable workflow checkpoints
+
+## Install
 
 ```bash
 pip install agent-framework-azure-cosmos --pre
 ```
 
-## Azure Cosmos DB History Provider
+The package requires Python 3.10 or later, an Azure Cosmos DB for NoSQL account,
+and `azure-cosmos` 4.7 or later. Vector collections also require the account's
+[NoSQL vector search capability](https://learn.microsoft.com/azure/cosmos-db/vector-search)
+to be enabled before use.
 
-The Azure Cosmos DB integration provides `CosmosHistoryProvider` for persistent conversation history storage.
+## Authentication and settings
 
-### Basic Usage Example
+Pass a caller-owned Azure credential or account key. If constructor values are omitted,
+the package reads an explicitly selected `.env` file and then these environment variables:
 
-```python
-from azure.identity.aio import DefaultAzureCredential
-from agent_framework_azure_cosmos import CosmosHistoryProvider
+| Variable | Purpose |
+| --- | --- |
+| `AZURE_COSMOS_ENDPOINT` | Azure Cosmos DB account endpoint |
+| `AZURE_COSMOS_DATABASE_NAME` | Database name |
+| `AZURE_COSMOS_CONTAINER_NAME` | Container name for direct collection/provider use |
+| `AZURE_COSMOS_KEY` | Account key; omit when passing an Azure credential |
 
-provider = CosmosHistoryProvider(
+Explicit constructor values take precedence over a selected `.env` file, which takes
+precedence over process environment variables. Injected asynchronous `CosmosClient`,
+`DatabaseProxy`, and `ContainerProxy` objects bypass connection settings and remain
+caller-owned.
+
+## Vector store
+
+Vector collections require:
+
+- one application-provided string key with storage name `id`;
+- a single Hash partition key path `/id`;
+- top-level dense vector fields; and
+- the container vector and indexing policies derived from the collection definition.
+
+This makes `get()` and `delete()` unambiguous point operations because the item ID is also
+its partition key. IDs must contain 1-1,023 UTF-8 bytes and cannot contain `/`, `\`,
+`?`, or `#`; the connector never encodes them. Custom and hierarchical partition keys
+are not supported.
+
+```python
+from dataclasses import dataclass
+from typing import Annotated
+
+from agent_framework import VectorStoreField, vectorstoremodel
+from agent_framework.azure import CosmosStore
+from azure.identity.aio import AzureCliCredential
+
+
+@vectorstoremodel(collection_name="documents")
+@dataclass
+class Document:
+    id: Annotated[str, VectorStoreField("key", storage_name="id")]
+    text: Annotated[str, VectorStoreField("data", is_indexed=True)]
+    embedding: Annotated[
+        list[float],
+        VectorStoreField(
+            "vector",
+            dimensions=1536,
+            distance_function="cosine_similarity",
+        ),
+    ]
+
+
+async with CosmosStore(
     endpoint="https://<account>.documents.azure.com:443/",
-    credential=DefaultAzureCredential(),
     database_name="agent-framework",
-    container_name="chat-history",
-)
+    credential=AzureCliCredential(),
+) as store:
+    collection = store.get_collection(Document)
+    await collection.ensure_collection_exists()
+    await collection.upsert(
+        [Document(id="doc-1", text="Vector search", embedding=[1.0] + [0.0] * 1535)],
+        generate_vectors=False,
+    )
+    results = await collection.search(vector=[1.0] + [0.0] * 1535, top=3)
+    async for result in results:
+        print(result["record"].text, result["score"])
 ```
 
-Credentials follow the same pattern used by other Azure connectors in the repository:
-
-- Pass a credential object (for example `DefaultAzureCredential`)
-- Or pass a key string directly
-- Or set `AZURE_COSMOS_KEY` in the environment
+The default vector index is `quantizedFlat`, which supports up to 4,096 dimensions and
+uses quantized rather than exact-float recall. Below 1,000 indexed vectors, Azure Cosmos DB
+falls back to a full scan. For exact search, explicitly set `index_kind="flat"` on vector
+fields with at most 505 dimensions.
 
-Container naming behavior:
+Supported vector element types are `float32`, `int8`, and `uint8`. Supported metrics are
+cosine similarity, dot product, and Euclidean distance. Scores are returned unchanged:
+higher is better for cosine and dot product, while lower is better for Euclidean distance.
+Score thresholds are supported for cosine and dot product; Euclidean search is supported
+without `score_threshold` because direct `VectorDistance` predicates do not reliably apply
+Euclidean cutoffs.
+After results are consumed, `SearchResults.metadata` contains the bounded request charge,
+last activity ID, and whether the last response page advertised more results; continuation
+tokens are not exposed.
 
-- Container name is configured on the provider (`container_name` or `AZURE_COSMOS_CONTAINER_NAME`)
-- `session_id` is used as the Cosmos partition key for reads/writes
+Vector index tuning is available through the field's `azure_cosmos` provider annotations:
 
-See `samples/02-agents/conversations/cosmos_history_provider.py` for a runnable example.
+```python
+VectorStoreField(
+    "vector",
+    dimensions=1536,
+    index_kind="disk_ann",
+    provider_annotations={
+        "azure_cosmos": {
+            "quantizer_type": "spherical",
+            "quantization_byte_size": 256,
+            "indexing_search_list_size": 200,
+        }
+    },
+)
+```
 
-## Cosmos DB Workflow Checkpoint Storage
+Search operation options support `search_list_size_multiplier`,
+`quantized_vector_list_multiplier`, `filter_priority`, and `brute_force`.
+Values and vectors are always sent as query parameters.
 
-`CosmosCheckpointStorage` implements the `CheckpointStorage` protocol, enabling
-durable workflow checkpointing backed by Azure Cosmos DB NoSQL. Workflows can be
-paused and resumed across process restarts by persisting checkpoint state in Cosmos DB.
+Writes and deletes span `/id` partitions and are not multi-item transactions. The connector
+validates the complete input batch before the first request, but a service failure can still
+leave an operation partially applied. Retry with stable application keys for idempotent
+upserts and deletes.
 
-### Basic Usage
+See the [Azure Cosmos DB vector search documentation](https://learn.microsoft.com/azure/cosmos-db/vector-search)
+and the [Agent Framework vector-store samples](https://github.com/microsoft/agent-framework/tree/main/python/samples/02-agents/vector_stores).
 
-#### Managed Identity / RBAC (recommended for production)
+## Conversation history
 
 ```python
-from azure.identity.aio import DefaultAzureCredential
-from agent_framework import WorkflowBuilder
-from agent_framework_azure_cosmos import CosmosCheckpointStorage
+from agent_framework.azure import CosmosHistoryProvider
+from azure.identity.aio import AzureCliCredential
 
-checkpoint_storage = CosmosCheckpointStorage(
+provider = CosmosHistoryProvider(
     endpoint="https://<account>.documents.azure.com:443/",
-    credential=DefaultAzureCredential(),
+    credential=AzureCliCredential(),
     database_name="agent-framework",
-    container_name="workflow-checkpoints",
+    container_name="chat-history",
 )
 ```
 
-#### Account Key
+`CosmosHistoryProvider` stores each conversation under its `session_id` partition key.
+See the [conversation sample](https://github.com/microsoft/agent-framework/blob/main/python/samples/02-agents/conversations/cosmos_history_provider.py).
+
+## Workflow checkpoints
 
 ```python
 from agent_framework_azure_cosmos import CosmosCheckpointStorage
+from azure.identity.aio import AzureCliCredential
 
-checkpoint_storage = CosmosCheckpointStorage(
+storage = CosmosCheckpointStorage(
     endpoint="https://<account>.documents.azure.com:443/",
-    credential="<your-account-key>",
+    credential=AzureCliCredential(),
     database_name="agent-framework",
     container_name="workflow-checkpoints",
 )
 ```
 
-#### Then use with a workflow
-
-```python
-from agent_framework import WorkflowBuilder
-
-# Build a workflow with checkpointing enabled
-workflow = WorkflowBuilder(
-    start_executor=start,
-    checkpoint_storage=checkpoint_storage,
-).build()
-
-# Run the workflow — checkpoints are automatically saved after each superstep
-result = await workflow.run(message="input data")
-
-# Resume from a checkpoint
-latest = await checkpoint_storage.get_latest(workflow_name=workflow.name)
-if latest:
-    resumed = await workflow.run(checkpoint_id=latest.checkpoint_id)
-```
-
-### Authentication Options
-
-`CosmosCheckpointStorage` supports the same authentication modes as `CosmosHistoryProvider`:
-
-- **Managed identity / RBAC** (recommended): Pass `DefaultAzureCredential()`,
-  `ManagedIdentityCredential()`, or any Azure `TokenCredential`
-- **Account key**: Pass a key string via `credential` parameter
-- **Environment variables**: Set `AZURE_COSMOS_ENDPOINT`, `AZURE_COSMOS_DATABASE_NAME`,
-  `AZURE_COSMOS_CONTAINER_NAME`, and `AZURE_COSMOS_KEY` (key not required when using
-  Azure credentials)
-- **Pre-created client**: Pass an existing `CosmosClient` or `ContainerProxy`
-
-### Database and Container Setup
-
-The database and container are created automatically on first use (via
-`create_database_if_not_exists` and `create_container_if_not_exists`). The container
-uses `/workflow_name` as the partition key. You can also pre-create them in the Azure
-portal with this partition key configuration.
-
-### Environment Variables
-
-| Variable | Description |
-|---|---|
-| `AZURE_COSMOS_ENDPOINT` | Cosmos DB account endpoint |
-| `AZURE_COSMOS_DATABASE_NAME` | Database name |
-| `AZURE_COSMOS_CONTAINER_NAME` | Container name |
-| `AZURE_COSMOS_KEY` | Account key (optional if using Azure credentials) |
-
-See `samples/03-workflows/checkpoint/cosmos_workflow_checkpointing.py` for a standalone example,
-or `samples/03-workflows/checkpoint/cosmos_workflow_checkpointing_foundry.py` for an end-to-end
-example with Microsoft Foundry agents.
+`CosmosCheckpointStorage` uses `/workflow_name` as its partition key and creates its
+database and container on first use. See the
+[checkpoint sample](https://github.com/microsoft/agent-framework/blob/main/python/samples/03-workflows/checkpoint/cosmos_workflow_checkpointing.py).

python/packages/azure-cosmos/agent_framework_azure_cosmos/__init__.py

@@ -1,17 +1,55 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+import importlib
 import importlib.metadata
+from typing import TYPE_CHECKING, Any
+
+import agent_framework
 
 from ._checkpoint_storage import CosmosCheckpointStorage
 from ._history_provider import CosmosHistoryProvider
 
+if TYPE_CHECKING:
+    from ._vector_store import AzureCosmosSettings, CosmosCollection, CosmosStore  # pyright: ignore[reportUnusedImport]
+
+_VECTOR_EXPORTS = frozenset({"AzureCosmosSettings", "CosmosCollection", "CosmosStore"})
+_HAS_VECTOR_CORE = all(
+    hasattr(agent_framework, name) for name in ("BaseVectorCollection", "BaseVectorSearch", "BaseVectorStore")
+)
+
 try:
     __version__ = importlib.metadata.version(__name__)
 except importlib.metadata.PackageNotFoundError:
     __version__ = "0.0.0"  # Fallback for development mode
 
-__all__ = [
-    "CosmosCheckpointStorage",
-    "CosmosHistoryProvider",
-    "__version__",
-]
+
+def __getattr__(name: str) -> Any:
+    if name not in _VECTOR_EXPORTS:
+        raise AttributeError(f"Module {__name__!r} has no attribute {name!r}.")
+    try:
+        return getattr(importlib.import_module("._vector_store", __name__), name)
+    except ImportError as exc:
+        raise ImportError(
+            "Azure Cosmos DB vector APIs require agent-framework-core with vector-store support."
+        ) from exc
+
+
+def __dir__() -> list[str]:
+    return sorted((*globals(), *_VECTOR_EXPORTS))
+
+
+if _HAS_VECTOR_CORE:
+    __all__ = [
+        "AzureCosmosSettings",
+        "CosmosCheckpointStorage",
+        "CosmosCollection",
+        "CosmosHistoryProvider",
+        "CosmosStore",
+        "__version__",
+    ]
+else:
+    __all__ = [
+        "CosmosCheckpointStorage",
+        "CosmosHistoryProvider",
+        "__version__",
+    ]