Python: Bambriz context provider

python/packages/azure-cosmos/AGENTS.md

@@ -1,30 +1,65 @@
 # Azure Cosmos DB Package (agent-framework-azure-cosmos)
 
-Azure Cosmos DB history provider integration for Agent Framework.
+Azure Cosmos DB history, retrieval, and workflow checkpointing integration for Agent Framework.
 
 ## Main Classes
 
 - **`CosmosHistoryProvider`** - Persistent conversation history storage backed by Azure Cosmos DB
+- **`CosmosContextProvider`** - Cosmos DB context provider for injecting relevant documents before a model run and writing request/response messages back into the same container after the run
+- **`CosmosCheckpointStorage`** - Cosmos DB-backed workflow checkpoint storage for durable workflow execution
 
 ## Usage
 
 ```python
-from agent_framework.azure import CosmosHistoryProvider
+from agent_framework_azure_cosmos import (
+    CosmosCheckpointStorage,
+    CosmosContextProvider,
+    CosmosContextSearchMode,
+    CosmosHistoryProvider,
+)
 
-provider = CosmosHistoryProvider(
+history_provider = CosmosHistoryProvider(
     endpoint="https://<account>.documents.azure.com:443/",
     credential="<key-or-token-credential>",
     database_name="agent-framework",
     container_name="chat-history",
 )
+
+context_provider = CosmosContextProvider(
+    endpoint="https://<account>.documents.azure.com:443/",
+    credential="<key-or-token-credential>",
+    database_name="agent-framework",
+    container_name="knowledge",
+    embedding_function=my_embedding_function,
+)
+
+checkpoint_storage = CosmosCheckpointStorage(
+    endpoint="https://<account>.documents.azure.com:443/",
+    credential="<key-or-token-credential>",
+    database_name="agent-framework",
+    container_name="workflow-checkpoints",
+)
 ```
 
-Container name is configured on the provider. `session_id` is used as the partition key.
+Container name is configured on each provider. `CosmosHistoryProvider` uses `session_id` as the partition key for reads/writes. `CosmosContextProvider` can optionally scope retrieval with `partition_key`.
+
+`CosmosContextProvider` joins the filtered `user` and `assistant` messages from the current run into one retrieval query string, and writes request/response messages back into the same Cosmos knowledge container after each run. All configuration — including search mode, weights, top_k, scan_limit, and partition key — is set on the constructor.
+
+The default search mode is `VECTOR`. Full-text and hybrid modes are also supported via the `search_mode` constructor parameter. Optional hybrid RRF weights can be provided through `weights=[...]` on the constructor.
+
+The application owner is responsible for making sure the Cosmos account, database, container, partitioning strategy, and any required full-text/vector/hybrid indexing configuration already exist. The provider does not create or manage Cosmos resources or search policies.
+
+`CosmosCheckpointStorage` creates the configured database and container on first use when needed, and stores workflow checkpoints using `/workflow_name` as the partition key.
 
 ## Import Path
 
 ```python
+from agent_framework_azure_cosmos import (
+    CosmosCheckpointStorage,
+    CosmosContextProvider,
+    CosmosHistoryProvider,
+)
+
+# `CosmosHistoryProvider` is also available from the Azure namespace:
 from agent_framework.azure import CosmosHistoryProvider
-# or directly:
-from agent_framework_azure_cosmos import CosmosHistoryProvider
 ```

python/packages/azure-cosmos/README.md

@@ -37,6 +37,73 @@ Container naming behavior:
 
 See `samples/02-agents/conversations/cosmos_history_provider.py` for a runnable example.
 
+## Azure Cosmos DB Context Provider
+
+The Azure Cosmos DB integration also provides `CosmosContextProvider` for context injection before model invocation. It also writes input and response messages back into the same Cosmos container after each run so the knowledge container can accumulate additional context over time.
+
+### Basic Usage Example
+
+```python
+from azure.identity.aio import DefaultAzureCredential
+from agent_framework_azure_cosmos import CosmosContextProvider, CosmosContextSearchMode
+
+provider = CosmosContextProvider(
+    endpoint="https://<account>.documents.azure.com:443/",
+    credential=DefaultAzureCredential(),
+    database_name="agent-framework",
+    container_name="knowledge",
+    embedding_function=my_embedding_function,
+    content_field_names=("content", "text"),
+)
+```
+
+Supported retrieval configuration includes:
+
+- `search_mode`: `CosmosContextSearchMode.VECTOR` (default), `.FULL_TEXT`, or `.HYBRID`
+- `weights` for hybrid RRF runs (optional, omitted by default)
+- `top_k` for controlling the number of context messages injected
+- `scan_limit` for controlling the number of Cosmos candidate items scanned
+- `partition_key` for scoping Cosmos retrieval
+
+All configuration is set on the constructor. The default search mode is `VECTOR`, which requires an `embedding_function`. For full-text mode, set `search_mode=CosmosContextSearchMode.FULL_TEXT`:
+
+```python
+provider = CosmosContextProvider(
+    endpoint="https://<account>.documents.azure.com:443/",
+    credential=DefaultAzureCredential(),
+    database_name="agent-framework",
+    container_name="knowledge",
+    search_mode=CosmosContextSearchMode.FULL_TEXT,
+)
+```
+
+For hybrid retrieval with optional weights:
+
+```python
+provider = CosmosContextProvider(
+    endpoint="https://<account>.documents.azure.com:443/",
+    credential=DefaultAzureCredential(),
+    database_name="agent-framework",
+    container_name="knowledge",
+    embedding_function=my_embedding_function,
+    search_mode=CosmosContextSearchMode.HYBRID,
+    weights=[2.0, 1.0],
+    top_k=3,
+    scan_limit=10,
+    partition_key="tenant-a",
+)
+```
+
+`CosmosContextProvider` contributes retrieval context in `before_run(...)` and persists input/response messages in `after_run(...)`.
+
+The provider builds retrieval input by joining the filtered `user` and `assistant` messages from the current run into a single query string. That joined query text is then used for full-text tokenization, vector embedding generation, or hybrid retrieval depending on the configured search mode.
+
+The provider writes the request/response messages back into the same knowledge container configured by `container_name`.
+
+The provider assumes the Cosmos account, database, container, partitioning strategy, and any required Cosmos full-text/vector/hybrid indexing policies already exist and are correctly configured by the application owner. It does not create or manage Cosmos resources, schema, or search policies for you.
+
+See `packages/azure-cosmos/samples/cosmos_context_provider.py` for a package-local context provider example.
+
 ## Cosmos DB Workflow Checkpoint Storage
 
 `CosmosCheckpointStorage` implements the `CheckpointStorage` protocol, enabling
@@ -84,7 +151,7 @@ workflow = WorkflowBuilder(
     checkpoint_storage=checkpoint_storage,
 ).build()
 
-# Run the workflow — checkpoints are automatically saved after each superstep
+# Run the workflow - checkpoints are automatically saved after each superstep
 result = await workflow.run(message="input data")
 
 # Resume from a checkpoint
@@ -124,3 +191,4 @@ portal with this partition key configuration.
 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 Azure AI Foundry agents.
+

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

@@ -1,8 +1,18 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+"""Azure Cosmos DB provider exports.
+
+Supported classes:
+    - ``CosmosContextProvider``
+    - ``CosmosCheckpointStorage``
+    - ``CosmosContextSearchMode``
+    - ``CosmosHistoryProvider``
+"""
+
 import importlib.metadata
 
 from ._checkpoint_storage import CosmosCheckpointStorage
+from ._context_provider import CosmosContextProvider, CosmosContextSearchMode
 from ._history_provider import CosmosHistoryProvider
 
 try:
@@ -12,6 +22,8 @@
 
 __all__ = [
     "CosmosCheckpointStorage",
+    "CosmosContextProvider",
+    "CosmosContextSearchMode",
     "CosmosHistoryProvider",
     "__version__",
 ]