From a88e2ea20b512a4aeb46e6d985f5d9aec0abaaef Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Wed, 16 Sep 2026 10:58:06 +0200 Subject: [PATCH] Python: Add generic vector store providers --- .../vector-stores-and-embeddings/README.md | 25 +- docs/specs/feature-usage-bit-registry.md | 4 +- python/packages/core/AGENTS.md | 6 + .../packages/core/agent_framework/__init__.py | 10 + .../core/agent_framework/__init__.pyi | 10 + .../core/agent_framework/_telemetry.py | 2 + .../packages/core/agent_framework/_vectors.py | 855 +++++++++++++++++- .../packages/core/tests/core/test_vectors.py | 345 +++++++ .../samples/02-agents/conversations/README.md | 5 + .../vector_store_history_provider.py | 80 ++ .../samples/02-agents/vector_stores/README.md | 7 +- .../vector_collection_context_provider.py | 98 ++ ..._context_provider_multiple_search_tools.py | 152 ++++ 13 files changed, 1590 insertions(+), 9 deletions(-) create mode 100644 python/samples/02-agents/conversations/vector_store_history_provider.py create mode 100644 python/samples/02-agents/vector_stores/vector_collection_context_provider.py create mode 100644 python/samples/02-agents/vector_stores/vector_collection_context_provider_multiple_search_tools.py diff --git a/docs/features/vector-stores-and-embeddings/README.md b/docs/features/vector-stores-and-embeddings/README.md index ecdac79117a..c509aa73444 100644 --- a/docs/features/vector-stores-and-embeddings/README.md +++ b/docs/features/vector-stores-and-embeddings/README.md @@ -11,6 +11,7 @@ This feature ports the vector store abstractions, embedding generator abstractio | Vector search | Unified search interface with `search_type` parameter (`"vector"`, `"keyword_hybrid"`) | | Data model decorator | `@vectorstoremodel` decorator for defining vector store data models (supports Pydantic, dataclasses, plain classes, dicts) | | Agent tools | `create_vector_search_tool`, `create_upsert_tool`, `create_get_tool`, `create_delete_tool` for agent-usable vector store operations | +| Agent providers | `VectorStoreHistoryProvider` for provider-owned scoped history and `VectorCollectionContextProvider` for caller-owned collections | | In-memory store | Zero-dependency vector store for testing and development | | 13+ connectors | Azure AI Search, Qdrant, Redis, PostgreSQL, MongoDB, Cosmos DB, Pinecone, Chroma, Weaviate, Oracle, SQL Server, FAISS | @@ -130,6 +131,11 @@ Options considered: | `SupportsVectorUpsert` | `_vectors.py` | Protocol for collection CRUD | | `SupportsVectorSearch` | `_vectors.py` | Protocol for vector search | | `create_vector_search_tool` | `_vectors.py` | Creates AF `FunctionTool` from vector search | +| `create_upsert_tool` | `_vectors.py` | Creates an agent tool for record upserts | +| `create_get_tool` | `_vectors.py` | Creates an agent tool for key-based retrieval | +| `create_delete_tool` | `_vectors.py` | Creates an agent tool for key-based deletion | +| `VectorStoreHistoryProvider` | `_vectors.py` | Stores scoped full history with optional compacted loading and search | +| `VectorCollectionContextProvider` | `_vectors.py` | Adds CRUD and one or more search tools for a caller-owned collection | ## Source Files Reference (SK → AF mapping) @@ -437,7 +443,7 @@ Each connector follows the AF package structure: --- ### Phase 8: Vector Store CRUD Tools -**Goal:** Provide a full set of agent-usable tools for CRUD operations on vector store collections. +**Goal:** Provide agent-usable CRUD tools plus generic history and context providers for vector stores. **Mergeable:** Yes — adds tools without changing existing APIs. #### 8.1 — `create_upsert_tool` — tool for upserting records into a collection @@ -447,6 +453,8 @@ Each connector follows the AF package structure: - Consider if this overlaps with filtered search and document when to use which #### 8.3 — `create_delete_tool` — tool for deleting records by key #### 8.4 — Tests and samples for CRUD tools +#### 8.5 — `VectorStoreHistoryProvider` — provider-owned scoped history with optional embeddings, compaction, and search +#### 8.6 — `VectorCollectionContextProvider` — configurable tools and instructions for a caller-owned collection --- @@ -502,9 +510,18 @@ Each connector follows the AF package structure: - `create_upsert_tool(...)` → tool for upserting records - `create_get_tool(...)` → tool for retrieving records by key - `create_delete_tool(...)` → tool for deleting records - - These are separate from search and are placed in a later phase - -10. **Score threshold filtering**: Scoring, filter execution, score thresholds, and paging belong to the connector + - Upsert and delete require approval by default; get is read-only by default + +10. **Generic vector providers**: + - `VectorStoreHistoryProvider` takes a store because it owns the history model and collection. Every operation + is filtered by application, optional tenant/agent, provider source, and session. Compaction changes only the + messages loaded into model context; optional search still queries the full scoped transcript. Content arrays + can be stored as JSON text or msgspec-encoded MessagePack bytes to match the backing store. + - `VectorCollectionContextProvider` takes a caller-owned collection and adds generated instructions plus + configurable CRUD/search tools. Additional `create_vector_search_tool` instances can expose different filters + and result mappings. + +11. **Score threshold filtering**: Scoring, filter execution, score thresholds, and paging belong to the connector and backing store (ref: [SK .NET PR #13501](https://github.com/microsoft/semantic-kernel/pull/13501)). Core passes `score_threshold` through without requiring a known distance function or an explicit metric and does not post-filter returned results, including results without scores. Each connector defines its score units, diff --git a/docs/specs/feature-usage-bit-registry.md b/docs/specs/feature-usage-bit-registry.md index 1e955610acf..4f7ee4cd937 100644 --- a/docs/specs/feature-usage-bit-registry.md +++ b/docs/specs/feature-usage-bit-registry.md @@ -138,7 +138,9 @@ only to approved first-party endpoints. | 17 | `core.session_store` | Agent session store | `agent_framework.SessionStore` / `FileSessionStore` | | 18 | `core.agent_hooks` | Agent Hooks middleware | `agent_framework.create_agent_hooks_middleware` | | 19 | `core.vector_stores` | Vector store abstractions | `BaseVectorCollection` / `BaseVectorSearch` operations | -| 20–31 | _reserved_ | core growth | — | +| 20 | `core.vector_store_history_provider` | Vector-store history provider | `agent_framework.VectorStoreHistoryProvider` | +| 21 | `core.vector_collection_context_provider` | Vector-collection context provider | `agent_framework.VectorCollectionContextProvider` | +| 22–31 | _reserved_ | core growth | — | | 32 | `orchestration.sequential` | Sequential orchestration | `agent_framework_orchestrations.SequentialBuilder` | | 33 | `orchestration.concurrent` | Concurrent orchestration | `agent_framework_orchestrations.ConcurrentBuilder` | | 34 | `orchestration.group_chat` | Group-chat orchestration | `agent_framework_orchestrations.GroupChatBuilder` | diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index 7311bc3942d..278159e29f5 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -112,6 +112,12 @@ The vector store API is experimental under the shared `VECTOR_STORES` feature ID filter execution, score thresholds (including provider-defined/default metrics), and paging. Use native backend execution where available, otherwise an explicit connector-local fallback or reject unsupported options - **`create_vector_search_tool`** - Creates an agent tool from any `SupportsVectorSearch` implementation +- **`create_upsert_tool` / `create_get_tool` / `create_delete_tool`** - Create agent tools for collection CRUD; + upsert and delete require approval by default, while get does not +- **`VectorStoreHistoryProvider`** - Stores full scoped conversation history in a provider-owned collection; + optional embeddings enable session-scoped history search, while optional compaction affects only loaded context +- **`VectorCollectionContextProvider`** - Adds instructions and configurable CRUD/search tools for a caller-owned + collection, including independently configured additional search tools - **`InMemoryCollection` / `InMemoryStore`** - Dependency-free, process-local development and test implementation; cosine scoring scales finite inputs, all metrics reject non-finite scores, and unsupported distance functions fail before record scanning. Hamming scores/thresholds use the fraction of unequal dimensions, not a count. diff --git a/python/packages/core/agent_framework/__init__.py b/python/packages/core/agent_framework/__init__.py index ea9e4577a22..2039452197e 100644 --- a/python/packages/core/agent_framework/__init__.py +++ b/python/packages/core/agent_framework/__init__.py @@ -310,8 +310,13 @@ "SearchType", "SupportsVectorSearch", "SupportsVectorUpsert", + "VectorCollectionContextProvider", "VectorStoreCollectionDefinition", "VectorStoreField", + "VectorStoreHistoryProvider", + "create_delete_tool", + "create_get_tool", + "create_upsert_tool", "create_vector_search_tool", "register_vectorstoremodel", "vectorstoremodel", @@ -633,8 +638,10 @@ "UsageDetails", "UserInputRequiredException", "ValidationTypeEnum", + "VectorCollectionContextProvider", "VectorStoreCollectionDefinition", "VectorStoreField", + "VectorStoreHistoryProvider", "Workflow", "WorkflowAgent", "WorkflowBuilder", @@ -667,8 +674,11 @@ "create_agent_hooks_middleware_from_emitter", "create_always_approve_tool_response", "create_always_approve_tool_with_arguments_response", + "create_delete_tool", "create_edge_runner", + "create_get_tool", "create_harness_agent", + "create_upsert_tool", "create_vector_search_tool", "detect_media_type_from_base64", "enqueue_messages", diff --git a/python/packages/core/agent_framework/__init__.pyi b/python/packages/core/agent_framework/__init__.pyi index 350a458a389..e35b6dc843a 100644 --- a/python/packages/core/agent_framework/__init__.pyi +++ b/python/packages/core/agent_framework/__init__.pyi @@ -266,8 +266,13 @@ from ._vectors import ( SearchType, SupportsVectorSearch, SupportsVectorUpsert, + VectorCollectionContextProvider, VectorStoreCollectionDefinition, VectorStoreField, + VectorStoreHistoryProvider, + create_delete_tool, + create_get_tool, + create_upsert_tool, create_vector_search_tool, register_vectorstoremodel, vectorstoremodel, @@ -588,8 +593,10 @@ __all__ = [ "UsageDetails", "UserInputRequiredException", "ValidationTypeEnum", + "VectorCollectionContextProvider", "VectorStoreCollectionDefinition", "VectorStoreField", + "VectorStoreHistoryProvider", "Workflow", "WorkflowAgent", "WorkflowBuilder", @@ -622,8 +629,11 @@ __all__ = [ "create_agent_hooks_middleware_from_emitter", "create_always_approve_tool_response", "create_always_approve_tool_with_arguments_response", + "create_delete_tool", "create_edge_runner", + "create_get_tool", "create_harness_agent", + "create_upsert_tool", "create_vector_search_tool", "detect_media_type_from_base64", "enqueue_messages", diff --git a/python/packages/core/agent_framework/_telemetry.py b/python/packages/core/agent_framework/_telemetry.py index d20f31b86cf..af6309e5d32 100644 --- a/python/packages/core/agent_framework/_telemetry.py +++ b/python/packages/core/agent_framework/_telemetry.py @@ -56,6 +56,8 @@ class FeatureIndex(IntEnum): CORE_SESSION_STORE = 17 CORE_AGENT_HOOKS = 18 CORE_VECTOR_STORES = 19 + CORE_VECTOR_STORE_HISTORY_PROVIDER = 20 + CORE_VECTOR_COLLECTION_CONTEXT_PROVIDER = 21 # This environment variable is reserved by the Foundry hosting environment to diff --git a/python/packages/core/agent_framework/_vectors.py b/python/packages/core/agent_framework/_vectors.py index 3ed815e0911..1d12e6cf8eb 100644 --- a/python/packages/core/agent_framework/_vectors.py +++ b/python/packages/core/agent_framework/_vectors.py @@ -4,15 +4,19 @@ from __future__ import annotations +import asyncio import operator +import time +import uuid from abc import ABC, abstractmethod from collections.abc import AsyncIterable, AsyncIterator, Callable, Mapping, Sequence from copy import deepcopy from dataclasses import dataclass, is_dataclass, replace from dataclasses import field as dataclass_field from inspect import Parameter, signature -from types import UnionType +from types import MappingProxyType, UnionType from typing import ( + TYPE_CHECKING, Annotated, Any, ClassVar, @@ -36,12 +40,16 @@ from typing_extensions import Self, TypedDict, TypeVar from ._clients import SupportsGetEmbeddings +from ._compaction import CompactionStrategy, TokenizerProtocol, apply_compaction from ._feature_stage import ExperimentalFeature, experimental +from ._sessions import AgentSession, ContextProvider, HistoryProvider, SessionContext, filter_new_messages from ._telemetry import FeatureIndex, mark_feature_used -from ._tools import FunctionTool -from ._types import Content, EmbeddingGenerationOptions +from ._tools import ApprovalMode, FunctionTool +from ._types import Content, EmbeddingGenerationOptions, Message from ._vector_filters import ( + Filter, FilterExpression, + FilterGroup, Param, iter_filter_params, param_schema, @@ -52,6 +60,9 @@ ) from .exceptions import IntegrationException, IntegrationInvalidResponseException +if TYPE_CHECKING: + from ._agents import SupportsAgentRun + ModelT = TypeVar("ModelT", default=Any) KeyT = TypeVar("KeyT", default=Any) ResultT = TypeVar("ResultT") @@ -84,6 +95,21 @@ _DEFAULT_SEARCH_TOOL_DESCRIPTION: Final[str] = ( "Perform a vector search for data in a vector store using the provided query." ) +_DEFAULT_UPSERT_TOOL_NAME: Final[str] = "upsert" +_DEFAULT_UPSERT_TOOL_DESCRIPTION: Final[str] = "Upsert records into the vector collection." +_DEFAULT_GET_TOOL_NAME: Final[str] = "get" +_DEFAULT_GET_TOOL_DESCRIPTION: Final[str] = "Get vector collection records by their keys." +_DEFAULT_DELETE_TOOL_NAME: Final[str] = "delete" +_DEFAULT_DELETE_TOOL_DESCRIPTION: Final[str] = "Delete vector collection records by their keys." +_VECTOR_HISTORY_PAGE_SIZE: Final[int] = 1000 +_VECTOR_HISTORY_SEARCH_TOP: Final[int] = 5 +_VectorCollectionOperation: TypeAlias = Literal["get", "delete", "upsert", "search"] +_DEFAULT_VECTOR_COLLECTION_APPROVAL_MODES: Final[Mapping[_VectorCollectionOperation, ApprovalMode]] = MappingProxyType({ + "get": "never_require", + "delete": "always_require", + "upsert": "always_require", + "search": "never_require", +}) DISTANCE_FUNCTION_DIRECTION_HELPER: Final[Mapping[DistanceFunction, Callable[[float | int, float | int], bool]]] = { "cosine_similarity": operator.ge, "cosine_distance": operator.le, @@ -1842,6 +1868,267 @@ async def search( ... +def _vector_tool_field_schema(field: VectorStoreField) -> dict[str, Any]: + if field.field_type == "vector": + return { + "anyOf": [ + {"type": "array", "items": {"type": "number"}}, + {"type": "string"}, + {"type": "null"}, + ] + } + json_type = { + "str": "string", + "int": "integer", + "float": "number", + "bool": "boolean", + "bytes": "string", + "bytearray": "string", + "list": "array", + "tuple": "array", + "set": "array", + "dict": "object", + }.get(field.type_ or "") + return {"type": json_type} if json_type is not None else {} + + +def _vector_tool_record_schema(collection: BaseVectorCollection[Any, Any]) -> dict[str, Any]: + properties = {field.name: _vector_tool_field_schema(field) for field in collection.definition.fields} + required = [ + field.name + for field in collection.definition.fields + if not (field.field_type == "key" and field.is_auto_generated) + and (collection.record_type is dict or not _has_default(collection.record_type, field.name)) + ] + return { + "type": "object", + "properties": properties, + "required": required, + "additionalProperties": False, + } + + +def _vector_tool_key_schema(collection: BaseVectorCollection[Any, Any]) -> dict[str, Any]: + return _vector_tool_field_schema(collection.definition.key_field) + + +def _validate_vector_tool_sequence(value: Any, *, name: str) -> list[Any]: + if not _is_non_string_sequence(value): + raise TypeError(f"{name} must be a sequence.") + values = list(value) + if not values: + raise ValueError(f"{name} must not be empty.") + return values + + +def _decode_vector_tool_records( + collection: BaseVectorCollection[KeyT, ModelT], + records: Any, +) -> list[ModelT]: + raw_records = _validate_vector_tool_sequence(records, name="records") + decoded: list[ModelT] = [] + registration = _VECTOR_MODEL_REGISTRY.get(collection.record_type) + for record in raw_records: + if not isinstance(record, Mapping): + raise TypeError("Each record must be a mapping.") + logical_record = dict(cast(Mapping[str, Any], record)) + if collection.record_type is dict: + decoded.append(cast(ModelT, logical_record)) + continue + if registration is None: + raise RuntimeError(f"Vector model {collection.record_type.__name__!r} is not registered.") + decoded.append(cast(ModelT, registration.decoder(logical_record))) + return decoded + + +def _encode_vector_tool_record( + collection: BaseVectorCollection[Any, ModelT], + record: ModelT, + *, + include_vectors: bool, +) -> dict[str, Any]: + if collection.record_type is dict: + source = _VectorStoreRecordHandler._to_builtin_mapping( # pyright: ignore[reportPrivateUsage] + record + ) + else: + registration = _VECTOR_MODEL_REGISTRY.get(collection.record_type) + if registration is None: + raise RuntimeError(f"Vector model {collection.record_type.__name__!r} is not registered.") + source = _VectorStoreRecordHandler._to_builtin_mapping( # pyright: ignore[reportPrivateUsage] + registration.encoder(record) + ) + + result: dict[str, Any] = {} + for field in collection.definition.fields: + if field.field_type == "vector" and not include_vectors: + continue + if field.name in source: + result[field.name] = source[field.name] + continue + storage_name = field.storage_name or field.name + if storage_name not in source: + raise IntegrationInvalidResponseException(f"Vector model is missing field '{field.name}' after retrieval.") + result[field.name] = source[storage_name] + return result + + +@experimental(feature_id=ExperimentalFeature.VECTOR_STORES) +def create_upsert_tool( + collection: BaseVectorCollection[KeyT, ModelT], + *, + name: str = _DEFAULT_UPSERT_TOOL_NAME, + description: str = _DEFAULT_UPSERT_TOOL_DESCRIPTION, + approval_mode: Literal["always_require", "never_require"] = "always_require", + generate_vectors: GenerateVectors = True, +) -> FunctionTool: + """Create an agent-usable tool that upserts vector collection records. + + Args: + collection: The vector collection CRUD capability invoked by the tool. + name: The tool name. + description: The tool description shown to the model. + approval_mode: Whether the tool requires approval before invocation. + generate_vectors: Which vector fields the collection generates during upsert. + + Returns: + A function tool accepting a non-empty ``records`` array. + """ + + async def upsert_tool(records: Any) -> dict[str, Any]: + decoded = _decode_vector_tool_records(collection, records) + keys = await collection.upsert(decoded, generate_vectors=generate_vectors) + return {"keys": list(keys)} + + return FunctionTool( + name=name, + description=description, + approval_mode=approval_mode, + func=upsert_tool, + input_model={ + "type": "object", + "properties": { + "records": { + "type": "array", + "items": _vector_tool_record_schema(collection), + "minItems": 1, + } + }, + "required": ["records"], + "additionalProperties": False, + }, + ) + + +@experimental(feature_id=ExperimentalFeature.VECTOR_STORES) +def create_get_tool( + collection: BaseVectorCollection[KeyT, ModelT], + *, + name: str = _DEFAULT_GET_TOOL_NAME, + description: str = _DEFAULT_GET_TOOL_DESCRIPTION, + approval_mode: Literal["always_require", "never_require"] = "never_require", + include_vectors: bool = False, + result_mapper: Callable[[ModelT], str | Content | Sequence[Content]] | None = None, +) -> FunctionTool: + """Create an agent-usable tool that gets vector collection records by key. + + Args: + collection: The vector collection CRUD capability invoked by the tool. + name: The tool name. + description: The tool description shown to the model. + approval_mode: Whether the tool requires approval before invocation. + include_vectors: Whether returned records include vector fields. + result_mapper: Optional model-specific projection for each retrieved record. + + Returns: + A function tool accepting a non-empty ``keys`` array. + """ + + async def get_tool(keys: Any) -> dict[str, Any] | list[Content]: + validated_keys = cast(list[KeyT], _validate_vector_tool_sequence(keys, name="keys")) + records = await collection.get(validated_keys, include_vectors=include_vectors) + if result_mapper is not None: + mapped_results: list[Content] = [] + for record in records: + mapped = result_mapper(record) + if isinstance(mapped, str): + mapped_results.append(Content.from_text(mapped)) + elif isinstance(mapped, Content): + mapped_results.append(mapped) + else: + mapped_results.extend(mapped) + return mapped_results + return { + "records": [ + _encode_vector_tool_record(collection, record, include_vectors=include_vectors) for record in records + ] + } + + return FunctionTool( + name=name, + description=description, + approval_mode=approval_mode, + func=get_tool, + input_model={ + "type": "object", + "properties": { + "keys": { + "type": "array", + "items": _vector_tool_key_schema(collection), + "minItems": 1, + } + }, + "required": ["keys"], + "additionalProperties": False, + }, + ) + + +@experimental(feature_id=ExperimentalFeature.VECTOR_STORES) +def create_delete_tool( + collection: BaseVectorCollection[KeyT, Any], + *, + name: str = _DEFAULT_DELETE_TOOL_NAME, + description: str = _DEFAULT_DELETE_TOOL_DESCRIPTION, + approval_mode: Literal["always_require", "never_require"] = "always_require", +) -> FunctionTool: + """Create an agent-usable tool that deletes vector collection records by key. + + Args: + collection: The vector collection CRUD capability invoked by the tool. + name: The tool name. + description: The tool description shown to the model. + approval_mode: Whether the tool requires approval before invocation. + + Returns: + A function tool accepting a non-empty ``keys`` array. + """ + + async def delete_tool(keys: Any) -> dict[str, Any]: + validated_keys = cast(list[KeyT], _validate_vector_tool_sequence(keys, name="keys")) + await collection.delete(validated_keys) + return {"processed_keys": validated_keys} + + return FunctionTool( + name=name, + description=description, + approval_mode=approval_mode, + func=delete_tool, + input_model={ + "type": "object", + "properties": { + "keys": { + "type": "array", + "items": _vector_tool_key_schema(collection), + "minItems": 1, + } + }, + "required": ["keys"], + "additionalProperties": False, + }, + ) + + @experimental(feature_id=ExperimentalFeature.VECTOR_STORES) def create_vector_search_tool( search: SupportsVectorSearch[ModelT], @@ -2063,3 +2350,565 @@ def _default_search_result_mapper(response: SearchResponse[Any]) -> str: response, enc_hook=_msgspec_enc_hook, ).decode() + + +def _vector_history_definition( + *, + dimensions: int | None, + contents_format: Literal["json", "msgpack"], +) -> VectorStoreCollectionDefinition: + fields = [ + VectorStoreField("key", name="id", type_="str"), + VectorStoreField("data", name="application_id", type_="str", is_indexed=True), + VectorStoreField("data", name="tenant_id", type_="str", is_indexed=True), + VectorStoreField("data", name="agent_id", type_="str", is_indexed=True), + VectorStoreField("data", name="source_id", type_="str", is_indexed=True), + VectorStoreField("data", name="session_id", type_="str", is_indexed=True), + VectorStoreField("data", name="timestamp", type_="int"), + VectorStoreField("data", name="message", type_="str"), + VectorStoreField( + "data", + name="contents", + type_="str" if contents_format == "json" else "bytes", + is_full_text_indexed=contents_format == "json", + ), + ] + if dimensions is not None: + fields.append( + VectorStoreField( + "vector", + name="embedding", + type_="float", + dimensions=dimensions, + distance_function="cosine_similarity", + ) + ) + return VectorStoreCollectionDefinition(fields) + + +@experimental(feature_id=ExperimentalFeature.VECTOR_STORES) +class VectorStoreHistoryProvider(HistoryProvider): + """Store full conversation history in a provider-owned vector collection. + + Use this provider when the records are Agent Framework :class:`Message` + instances and the agent should automatically load and save its conversation + through the normal :class:`HistoryProvider` lifecycle. The provider owns the + collection definition because chat-history records have a framework-defined + shape. + + Use :class:`VectorCollectionContextProvider` instead when the application + already owns a vector collection and data model and wants to expose that + domain data through CRUD or search tools. That provider does not participate + in conversation-history loading or persistence. + + History selection is scoped by ``application_id``, optional tenant and agent + identifiers, ``source_id``, and the current session. These identifiers prevent + accidental overlap but are not authorization boundaries; applications must + still use appropriately scoped store credentials and namespaces. + + ``get_messages`` returns the full scoped transcript. When a compaction strategy + is configured, only the projection produced after loading is added to model + context. The optional search tool always searches the full scoped transcript. + """ + + DEFAULT_SOURCE_ID: ClassVar[str] = "vector_store_history" + DEFAULT_COLLECTION_NAME: ClassVar[str] = "agent_framework_history" + SEARCH_TOOL_NAME: ClassVar[str] = "search_history" + SEARCH_TOOL_DESCRIPTION: ClassVar[str] = ( + "Search the full conversation history for information that may not be present in the loaded context." + ) + + def __init__( + self, + vector_store: BaseVectorStore, + source_id: str = DEFAULT_SOURCE_ID, + *, + application_id: str, + tenant_id: str | None = None, + agent_id: str | None = None, + collection_name: str = DEFAULT_COLLECTION_NAME, + contents_format: Literal["json", "msgpack"] = "json", + embedding_generator: EmbeddingClient | None = None, + embedding_options: Mapping[str, Any] | None = None, + compaction_strategy: CompactionStrategy | None = None, + compaction_tokenizer: TokenizerProtocol | None = None, + include_search_tool: bool = False, + search_approval_mode: Literal["always_require", "never_require"] = "never_require", + load_messages: bool = True, + store_inputs: bool = True, + store_context_messages: bool = False, + store_context_from: set[str] | None = None, + store_outputs: bool = True, + ) -> None: + """Initialize a vector-store-backed history provider. + + This provider is for framework-owned chat history. It creates the + collection client and translates messages to its fixed history schema. + For tools over a caller-supplied vector collection with a user-defined + data model, use :class:`VectorCollectionContextProvider`. + + Args: + vector_store: Store used to create the provider-owned history collection. + source_id: Provider identifier included in every history scope. + application_id: Required application isolation identifier. + tenant_id: Optional tenant isolation identifier. + agent_id: Optional agent isolation identifier. + collection_name: Name of the provider-owned history collection. + contents_format: Encoding used for the stored content array. JSON + stores text; MessagePack uses msgspec and stores bytes. + embedding_generator: Optional client used to embed each message's serialized contents. + embedding_options: Options passed to every embedding request. When an + embedding generator is supplied, this mapping must include positive + integer ``dimensions`` for the collection definition. + compaction_strategy: Optional strategy applied after loading history and + before adding it to model context. + compaction_tokenizer: Optional tokenizer used by compaction. + include_search_tool: Whether to add a scoped full-history search tool. + search_approval_mode: Approval mode for the optional history search tool. + load_messages: Whether to load messages before invocation. + store_inputs: Whether to store input messages. + store_context_messages: Whether to store context from other providers. + store_context_from: If set, only store context from these source IDs. + store_outputs: Whether to store response messages. + + Raises: + TypeError: If an isolation identifier, option mapping, or dimensions value has an invalid type. + ValueError: If required scope or embedding configuration is missing. + """ + super().__init__( + source_id, + load_messages=load_messages, + store_inputs=store_inputs, + store_context_messages=store_context_messages, + store_context_from=store_context_from, + store_outputs=store_outputs, + ) + for name, value in ( + ("source_id", source_id), + ("application_id", application_id), + ("collection_name", collection_name), + ): + if not isinstance(value, str): + raise TypeError(f"{name} must be a string.") + if not value: + raise ValueError(f"{name} must be a non-empty string.") + for name, value in (("tenant_id", tenant_id), ("agent_id", agent_id)): + if value is not None and not isinstance(value, str): + raise TypeError(f"{name} must be a string when supplied.") + if value == "": + raise ValueError(f"{name} must be non-empty when supplied.") + if not isinstance(include_search_tool, bool): + raise TypeError("include_search_tool must be a boolean.") + if contents_format not in ("json", "msgpack"): + raise ValueError("contents_format must be 'json' or 'msgpack'.") + if embedding_generator is None and embedding_options is not None: + raise ValueError("embedding_options requires embedding_generator.") + if embedding_generator is not None and embedding_options is None: + raise ValueError("embedding_generator requires embedding_options with dimensions.") + if include_search_tool and embedding_generator is None: + raise ValueError("include_search_tool requires embedding_generator.") + if include_search_tool and not load_messages: + raise ValueError("include_search_tool requires load_messages=True.") + + resolved_embedding_options: dict[str, Any] | None = None + dimensions: int | None = None + if embedding_options is not None: + if any(not isinstance(key, str) for key in embedding_options): + raise TypeError("embedding_options keys must be strings.") + resolved_embedding_options = deepcopy(dict(embedding_options)) + dimensions_value = resolved_embedding_options.get("dimensions") + if not isinstance(dimensions_value, int) or isinstance(dimensions_value, bool): + raise TypeError("embedding_options['dimensions'] must be an integer.") + if dimensions_value <= 0: + raise ValueError("embedding_options['dimensions'] must be positive.") + dimensions = dimensions_value + + self.vector_store = vector_store + self.application_id = application_id + self.tenant_id = tenant_id + self.agent_id = agent_id + self.collection_name = collection_name + self.contents_format = contents_format + self.embedding_generator = embedding_generator + self.embedding_options = resolved_embedding_options + self.compaction_strategy = compaction_strategy + self.compaction_tokenizer = compaction_tokenizer + self.include_search_tool = include_search_tool + self.search_approval_mode: ApprovalMode = search_approval_mode + self._collection: BaseVectorCollection[Any, dict[str, Any]] = vector_store.get_collection( + cast(type[dict[str, Any]], dict), + definition=_vector_history_definition( + dimensions=dimensions, + contents_format=contents_format, + ), + collection_name=collection_name, + embedding_generator=embedding_generator, + ) + if include_search_tool and not isinstance(self._collection, SupportsVectorSearch): + raise ValueError("The vector store collection does not support search.") + self._collection_ready = False + self._collection_lock = asyncio.Lock() + + def _scope_filter(self, session_id: str | None) -> FilterGroup: + if not isinstance(session_id, str) or not session_id: + raise ValueError("session_id must be a non-empty string.") + return FilterGroup( + "and", + ( + Filter("application_id", "eq", self.application_id), + Filter("tenant_id", "eq", self.tenant_id or ""), + Filter("agent_id", "eq", self.agent_id or ""), + Filter("source_id", "eq", self.source_id), + Filter("session_id", "eq", session_id), + ), + ) + + async def _ensure_collection(self) -> None: + if self._collection_ready: + return + async with self._collection_lock: + if self._collection_ready: + return + await self._collection.ensure_collection_exists() + self._collection_ready = True + + async def _get_records(self, session_id: str | None) -> list[dict[str, Any]]: + await self._ensure_collection() + records: list[dict[str, Any]] = [] + seen_keys: set[str] = set() + skip = 0 + while True: + page = await self._collection.get( + filter=self._scope_filter(session_id), + top=_VECTOR_HISTORY_PAGE_SIZE, + skip=skip, + ) + if not page: + break + for record in page: + key = record.get("id") + timestamp = record.get("timestamp") + if not isinstance(key, str): + raise IntegrationInvalidResponseException("Vector history record has a non-string id.") + if not isinstance(timestamp, int) or isinstance(timestamp, bool): + raise IntegrationInvalidResponseException("Vector history record has a non-integer timestamp.") + if key in seen_keys: + raise IntegrationInvalidResponseException( + "Vector history collection returned a duplicate record while paging." + ) + seen_keys.add(key) + records.append(record) + skip += len(page) + records.sort(key=lambda record: (record["timestamp"], record["id"])) + return records + + @staticmethod + def _messages_from_records(records: Sequence[Mapping[str, Any]]) -> list[Message]: + messages: list[Message] = [] + for record in records: + payload = record.get("message") + if not isinstance(payload, str): + raise IntegrationInvalidResponseException("Vector history record has a non-string message payload.") + try: + messages.append(Message.from_json(payload)) + except (TypeError, ValueError) as exc: + raise IntegrationInvalidResponseException( + "Vector history record has an invalid message payload." + ) from exc + return messages + + def _serialize_contents(self, message: Message) -> str | bytes: + contents = [content.to_dict() for content in message.contents] + if self.contents_format == "msgpack": + return msgspec.msgpack.encode(contents, enc_hook=_msgspec_enc_hook) + return msgspec.json.encode( + contents, + enc_hook=_msgspec_enc_hook, + ).decode() + + @staticmethod + def _serialize_contents_for_embedding(message: Message) -> str: + return msgspec.json.encode( + [content.to_dict() for content in message.contents], + enc_hook=_msgspec_enc_hook, + ).decode() + + async def _generate_embeddings(self, values: Sequence[str]) -> list[Vector]: + if self.embedding_generator is None or self.embedding_options is None: + raise RuntimeError("Embedding generation is not configured.") + try: + embeddings = await self.embedding_generator.get_embeddings( + values, + options=cast(Any, deepcopy(self.embedding_options)), + ) + except (TypeError, ValueError, NotImplementedError, IntegrationException): + raise + except Exception as exc: + raise IntegrationException(f"Error generating vector history embeddings: {exc}") from exc + if len(embeddings) != len(values): + raise IntegrationInvalidResponseException( + f"Embedding client returned {len(embeddings)} vectors for {len(values)} values." + ) + return [_normalize_vector(embedding.vector) for embedding in embeddings] + + async def get_messages( + self, + session_id: str | None, + *, + state: dict[str, Any] | None = None, + **kwargs: Any, + ) -> list[Message]: + """Return the full transcript for one isolated history scope.""" + mark_feature_used(FeatureIndex.CORE_VECTOR_STORE_HISTORY_PROVIDER) + del state, kwargs + return self._messages_from_records(await self._get_records(session_id)) + + async def save_messages( + self, + session_id: str | None, + messages: Sequence[Message], + *, + state: dict[str, Any] | None = None, + **kwargs: Any, + ) -> None: + """Persist new messages for one isolated history scope.""" + mark_feature_used(FeatureIndex.CORE_VECTOR_STORE_HISTORY_PROVIDER) + del state, kwargs + if not messages: + return + existing = self._messages_from_records(await self._get_records(session_id)) + new_messages = filter_new_messages(existing, messages) + if not new_messages: + return + + base_timestamp = time.time_ns() + records: list[dict[str, Any]] = [ + { + "id": str(uuid.uuid4()), + "application_id": self.application_id, + "tenant_id": self.tenant_id or "", + "agent_id": self.agent_id or "", + "source_id": self.source_id, + "session_id": cast(str, session_id), + "timestamp": base_timestamp + index, + "message": message.to_json(), + "contents": self._serialize_contents(message), + } + for index, message in enumerate(new_messages) + ] + if self.embedding_generator is not None: + vectors = await self._generate_embeddings([ + self._serialize_contents_for_embedding(message) for message in new_messages + ]) + for record, vector in zip(records, vectors, strict=True): + record["embedding"] = vector + await self._collection.upsert(records, generate_vectors=False) + + async def clear(self, session_id: str | None) -> None: + """Delete all messages for one isolated history scope.""" + mark_feature_used(FeatureIndex.CORE_VECTOR_STORE_HISTORY_PROVIDER) + records = await self._get_records(session_id) + if records: + await self._collection.delete([record["id"] for record in records]) + + def _create_search_tool(self, session_id: str | None) -> FunctionTool: + async def search_history(query: str) -> list[Content]: + if not isinstance(query, str): + raise TypeError("query must be a string.") + vector = (await self._generate_embeddings([query]))[0] + search = cast(SupportsVectorSearch[dict[str, Any]], self._collection) + results = await search.search( + vector=vector, + filter=self._scope_filter(session_id), + top=_VECTOR_HISTORY_SEARCH_TOP, + ) + mapped: list[Content] = [] + async for result in results: + message = self._messages_from_records([result["record"]])[0] + mapped.append(Content.from_text(f"[{message.role}] {message.text or message.to_json()}")) + return mapped + + return FunctionTool( + name=self.SEARCH_TOOL_NAME, + description=self.SEARCH_TOOL_DESCRIPTION, + approval_mode=self.search_approval_mode, + additional_properties={}, + func=search_history, + ) + + async def before_run( + self, + *, + agent: SupportsAgentRun, + session: AgentSession, + context: SessionContext, + state: dict[str, Any], + ) -> None: + """Load history, compact its context projection, and add scoped search.""" + await super().before_run(agent=agent, session=session, context=context, state=state) + if self.compaction_strategy is not None: + loaded = context.context_messages.get(self.source_id, []) + context.context_messages[self.source_id] = await apply_compaction( + loaded, + strategy=self.compaction_strategy, + tokenizer=self.compaction_tokenizer, + ) + if self.include_search_tool: + context.extend_instructions( + self.source_id, + f"Use {self.SEARCH_TOOL_NAME} when relevant information is not present in the loaded history.", + ) + context.extend_tools(self.source_id, [self._create_search_tool(context.session_id)]) + + +@experimental(feature_id=ExperimentalFeature.VECTOR_STORES) +class VectorCollectionContextProvider(ContextProvider, Generic[KeyT, ModelT]): + """Add agent-usable CRUD and search tools for a caller-owned vector collection. + + Use this provider when the application already owns the collection and its + domain model, such as products, documents, or hotel records, and wants the + agent to interact with those records through tools. The provider leaves + collection creation, schema design, and record lifecycle with the caller. + + Use :class:`VectorStoreHistoryProvider` instead to persist the agent's own + conversation. That provider owns a chat-history schema and automatically + loads and stores :class:`Message` objects through the + :class:`HistoryProvider` lifecycle. + """ + + DEFAULT_SOURCE_ID: ClassVar[str] = "vector_collection" + + def __init__( + self, + collection: BaseVectorCollection[KeyT, ModelT], + source_id: str = DEFAULT_SOURCE_ID, + *, + instructions: str | Sequence[str] | None = None, + include_upsert_tool: bool = True, + include_get_tool: bool = True, + include_delete_tool: bool = True, + include_search_tool: bool = True, + approval_mode: ( + Literal["always_require", "never_require"] + | Mapping[ + Literal["get", "delete", "upsert", "search"], + Literal["always_require", "never_require"], + ] + ) = _DEFAULT_VECTOR_COLLECTION_APPROVAL_MODES, + additional_search_tools: Sequence[FunctionTool] | None = None, + ) -> None: + """Initialize a vector collection context provider. + + This provider exposes a caller-supplied vector collection with a + user-defined data model as tools; it does not store or load conversation history. Use + :class:`VectorStoreHistoryProvider` for that purpose. + + Args: + collection: Caller-owned vector collection exposed through tools. + source_id: Provider identifier used for instruction and tool attribution. + instructions: Instructions added before each run. ``None`` uses generated defaults. + include_upsert_tool: Whether to add the default upsert tool. + include_get_tool: Whether to add the default get-by-key tool. + include_delete_tool: Whether to add the default delete-by-key tool. + include_search_tool: Whether to add the default vector search tool. + approval_mode: One mode for every generated tool, or per-tool overrides + merged over the safe defaults. + additional_search_tools: Additional caller-configured search tools. + + Raises: + TypeError: If a flag, instruction, approval setting, or additional tool is invalid. + ValueError: If a tool name is duplicated or default search is unsupported. + """ + super().__init__(source_id) + for name, value in ( + ("include_upsert_tool", include_upsert_tool), + ("include_get_tool", include_get_tool), + ("include_delete_tool", include_delete_tool), + ("include_search_tool", include_search_tool), + ): + if not isinstance(value, bool): + raise TypeError(f"{name} must be a boolean.") + + approval_modes = self._resolve_approval_modes(approval_mode) + tools: list[FunctionTool] = [] + if include_upsert_tool: + tools.append(create_upsert_tool(collection, approval_mode=approval_modes["upsert"])) + if include_get_tool: + tools.append(create_get_tool(collection, approval_mode=approval_modes["get"])) + if include_delete_tool: + tools.append(create_delete_tool(collection, approval_mode=approval_modes["delete"])) + if include_search_tool: + if not isinstance(collection, SupportsVectorSearch): + raise ValueError("The vector collection does not support search.") + tools.append( + create_vector_search_tool( + cast(SupportsVectorSearch[ModelT], collection), + approval_mode=approval_modes["search"], + ) + ) + + for tool in additional_search_tools or (): + if not isinstance(tool, FunctionTool): + raise TypeError("additional_search_tools must contain FunctionTool instances.") + tools.append(tool) + names = [tool.name for tool in tools] + if len(names) != len(set(names)): + raise ValueError("Vector collection tool names must be unique.") + for tool in tools: + tool.additional_properties = dict(tool.additional_properties or {}) + + self.collection = collection + self.tools = tuple(tools) + self.instructions = self._resolve_instructions(instructions) + + @staticmethod + def _resolve_approval_modes( + approval_mode: ApprovalMode | Mapping[_VectorCollectionOperation, ApprovalMode], + ) -> dict[_VectorCollectionOperation, ApprovalMode]: + resolved = dict(_DEFAULT_VECTOR_COLLECTION_APPROVAL_MODES) + if isinstance(approval_mode, str): + if approval_mode not in ("always_require", "never_require"): + raise ValueError("approval_mode must be 'always_require' or 'never_require'.") + return {name: approval_mode for name in resolved} + if not isinstance(approval_mode, Mapping): + raise TypeError("approval_mode must be a string or mapping.") + unknown = sorted(set(approval_mode) - set(resolved)) + if unknown: + raise ValueError(f"Unknown approval_mode tool(s): {', '.join(unknown)}.") + for name, mode in approval_mode.items(): + if mode not in ("always_require", "never_require"): + raise ValueError(f"Invalid approval mode for '{name}'.") + resolved[name] = mode + return resolved + + def _resolve_instructions(self, instructions: str | Sequence[str] | None) -> tuple[str, ...]: + if instructions is None: + generated = ["Use the available vector collection tools as the source of truth for stored records."] + tool_names = {tool.name for tool in self.tools} + if "get" in tool_names and "search" in tool_names: + generated.append("Use get when record keys are known; use search to find records by similarity.") + if "upsert" in tool_names: + generated.append("Use upsert to create or update collection records.") + if "delete" in tool_names: + generated.append("Use delete only when collection records should be removed.") + return tuple(generated) + if isinstance(instructions, str): + return (instructions,) if instructions else () + if not _is_non_string_sequence(instructions) or any(not isinstance(item, str) for item in instructions): + raise TypeError("instructions must be a string or sequence of strings.") + return tuple(instructions) + + async def before_run( + self, + *, + agent: SupportsAgentRun, + session: AgentSession, + context: SessionContext, + state: dict[str, Any], + ) -> None: + """Add configured instructions and tools to the current invocation.""" + mark_feature_used(FeatureIndex.CORE_VECTOR_COLLECTION_CONTEXT_PROVIDER) + if self.instructions: + context.extend_instructions(self.source_id, self.instructions) + if self.tools: + context.extend_tools(self.source_id, self.tools) diff --git a/python/packages/core/tests/core/test_vectors.py b/python/packages/core/tests/core/test_vectors.py index 0378804b180..5c1f32f9fe5 100644 --- a/python/packages/core/tests/core/test_vectors.py +++ b/python/packages/core/tests/core/test_vectors.py @@ -18,6 +18,7 @@ from agent_framework import ( DISTANCE_FUNCTION_DIRECTION_HELPER, + AgentSession, BaseEmbeddingClient, BaseVectorCollection, BaseVectorSearch, @@ -32,14 +33,22 @@ FilterGroup, GeneratedEmbeddings, IndexKind, + InMemoryStore, + Message, Param, SearchResponse, SearchResults, SearchType, + SessionContext, SupportsVectorSearch, SupportsVectorUpsert, + VectorCollectionContextProvider, VectorStoreCollectionDefinition, VectorStoreField, + VectorStoreHistoryProvider, + create_delete_tool, + create_get_tool, + create_upsert_tool, create_vector_search_tool, register_vectorstoremodel, vectorstoremodel, @@ -283,6 +292,11 @@ def test_vector_apis_are_marked_experimental() -> None: BaseVectorStore, BaseVectorSearch, register_vectorstoremodel, + create_upsert_tool, + create_get_tool, + create_delete_tool, + VectorStoreHistoryProvider, + VectorCollectionContextProvider, ) for api in staged_apis: assert getattr(api, "__feature_stage__", None) == "experimental" @@ -2529,3 +2543,334 @@ async def test_runtime_operations_mark_vector_store_feature_usage() -> None: mark_feature_used_mock.reset_mock() await collection.search(vector=[1.0, 0.0]) mark_feature_used_mock.assert_called_once_with(FeatureIndex.CORE_VECTOR_STORES) + + +async def test_vector_crud_tools_round_trip_records() -> None: + collection = MockCollection() + upsert_tool = create_upsert_tool(collection, generate_vectors=False) + get_tool = create_get_tool(collection) + delete_tool = create_delete_tool(collection) + + assert upsert_tool.approval_mode == "always_require" + assert get_tool.approval_mode == "never_require" + assert delete_tool.approval_mode == "always_require" + record_schema = upsert_tool.parameters()["properties"]["records"]["items"] + assert record_schema["required"] == ["id", "text"] + assert set(record_schema["properties"]) == {"id", "text", "vector"} + + upserted = await upsert_tool.invoke( + arguments={"records": [{"id": "one", "text": "first", "vector": [1.0, 0.0]}]}, + skip_parsing=True, + ) + assert upserted == {"keys": ["one"]} + + fetched = await get_tool.invoke(arguments={"keys": ["one"]}, skip_parsing=True) + assert fetched == {"records": [{"id": "one", "text": "first"}]} + + mapped_get_tool = create_get_tool(collection, result_mapper=lambda record: f"{record.id}: {record.text}") + mapped = await mapped_get_tool.invoke(arguments={"keys": ["one"]}, skip_parsing=True) + assert [content.text for content in mapped] == ["one: first"] + + deleted = await delete_tool.invoke(arguments={"keys": ["one"]}, skip_parsing=True) + assert deleted == {"processed_keys": ["one"]} + assert await get_tool.invoke(arguments={"keys": ["one"]}, skip_parsing=True) == {"records": []} + + +@pytest.mark.parametrize("factory", [create_upsert_tool, create_get_tool, create_delete_tool]) +async def test_vector_crud_tools_reject_empty_batches(factory: Callable[..., Any]) -> None: + collection = MockCollection() + tool = factory(collection) + argument_name = "records" if factory is create_upsert_tool else "keys" + + with pytest.raises(ValueError, match="must not be empty"): + await tool.invoke(arguments={argument_name: []}, skip_parsing=True) + + +def test_vector_collection_context_provider_configures_tools_and_approvals() -> None: + collection = MockCollection(embedding_generator=MockEmbeddingClient()) + provider = VectorCollectionContextProvider( + collection, + approval_mode={"upsert": "never_require"}, + ) + + assert [tool.name for tool in provider.tools] == ["upsert", "get", "delete", "search"] + assert {tool.name: tool.approval_mode for tool in provider.tools} == { + "upsert": "never_require", + "get": "never_require", + "delete": "always_require", + "search": "never_require", + } + assert any("get when record keys are known" in instruction for instruction in provider.instructions) + + require_all = VectorCollectionContextProvider(collection, approval_mode="always_require") + assert all(tool.approval_mode == "always_require" for tool in require_all.tools) + + +async def test_vector_collection_context_provider_adds_attributed_context() -> None: + collection = MockCollection(embedding_generator=MockEmbeddingClient()) + details_tool = create_vector_search_tool(collection, name="search_details") + provider = VectorCollectionContextProvider( + collection, + include_upsert_tool=False, + include_get_tool=False, + include_delete_tool=False, + include_search_tool=False, + additional_search_tools=[details_tool], + ) + session = AgentSession(session_id="session") + context = SessionContext(input_messages=[], session_id=session.session_id) + + with patch("agent_framework._vectors.mark_feature_used") as mark_feature_used_mock: + await provider.before_run(agent=cast(Any, None), session=session, context=context, state={}) + + assert [tool.name for tool in context.tools] == ["search_details"] + assert context.tools[0].additional_properties["context_source"] == provider.source_id + assert context.instructions == [ + "Use the available vector collection tools as the source of truth for stored records." + ] + mark_feature_used_mock.assert_called_once_with(FeatureIndex.CORE_VECTOR_COLLECTION_CONTEXT_PROVIDER) + + +def test_vector_collection_context_provider_validates_configuration() -> None: + collection = MockCollection(embedding_generator=MockEmbeddingClient()) + + with pytest.raises(ValueError, match="Unknown approval_mode"): + VectorCollectionContextProvider(collection, approval_mode=cast(Any, {"unknown": "never_require"})) + with pytest.raises(ValueError, match="Invalid approval mode"): + VectorCollectionContextProvider(collection, approval_mode=cast(Any, {"get": "sometimes"})) + with pytest.raises(ValueError, match="tool names must be unique"): + VectorCollectionContextProvider( + collection, + additional_search_tools=[create_vector_search_tool(collection)], + ) + with pytest.raises(TypeError, match="FunctionTool"): + VectorCollectionContextProvider( + collection, + include_search_tool=False, + additional_search_tools=cast(Any, [object()]), + ) + + +async def test_vector_store_history_provider_round_trips_and_deduplicates_messages() -> None: + generator = MockEmbeddingClient() + embedding_options: dict[str, Any] = {"model": "embedding-model", "dimensions": 2} + provider = VectorStoreHistoryProvider( + InMemoryStore(), + application_id="app", + tenant_id="tenant", + agent_id="agent", + embedding_generator=generator, + embedding_options=embedding_options, + ) + embedding_options["model"] = "changed-after-construction" + messages = [ + Message( + role="user", + contents=[ + "hello", + Content.from_uri("https://example.com/image.png", media_type="image/png"), + ], + ), + Message(role="assistant", contents=["hi"]), + ] + + await provider.save_messages("session", messages) + embedded_contents = [msgspec.json.decode(value) for value in generator.values] + assert embedded_contents == [ + [ + {"type": "text", "text": "hello", "additional_properties": {}}, + { + "type": "uri", + "uri": "https://example.com/image.png", + "media_type": "image/png", + "additional_properties": {}, + }, + ], + [{"type": "text", "text": "hi", "additional_properties": {}}], + ] + assert generator.options == {"model": "embedding-model", "dimensions": 2} + assert [message.to_dict() for message in await provider.get_messages("session")] == [ + message.to_dict() for message in messages + ] + + await provider.save_messages("session", messages) + assert len(await provider.get_messages("session")) == 2 + + await provider.clear("session") + assert await provider.get_messages("session") == [] + + +async def test_vector_store_history_provider_isolates_every_scope_dimension() -> None: + store = InMemoryStore() + provider = VectorStoreHistoryProvider( + store, + source_id="history", + application_id="app", + tenant_id="tenant", + agent_id="agent", + embedding_generator=MockEmbeddingClient(), + embedding_options={"dimensions": 2}, + ) + await provider.save_messages("session", [Message(role="user", contents=["visible"])]) + + other_providers = [ + VectorStoreHistoryProvider( + store, + source_id="other-history", + application_id="app", + tenant_id="tenant", + agent_id="agent", + embedding_generator=MockEmbeddingClient(), + embedding_options={"dimensions": 2}, + ), + VectorStoreHistoryProvider( + store, + source_id="history", + application_id="other-app", + tenant_id="tenant", + agent_id="agent", + embedding_generator=MockEmbeddingClient(), + embedding_options={"dimensions": 2}, + ), + VectorStoreHistoryProvider( + store, + source_id="history", + application_id="app", + tenant_id="other-tenant", + agent_id="agent", + embedding_generator=MockEmbeddingClient(), + embedding_options={"dimensions": 2}, + ), + VectorStoreHistoryProvider( + store, + source_id="history", + application_id="app", + tenant_id="tenant", + agent_id="other-agent", + embedding_generator=MockEmbeddingClient(), + embedding_options={"dimensions": 2}, + ), + ] + for index, other_provider in enumerate(other_providers): + await other_provider.save_messages("session", [Message(role="user", contents=[f"secret-{index}"])]) + + assert [message.text for message in await provider.get_messages("session")] == ["visible"] + for index, other_provider in enumerate(other_providers): + assert [message.text for message in await other_provider.get_messages("session")] == [f"secret-{index}"] + assert await provider.get_messages("other-session") == [] + + +async def test_vector_store_history_provider_works_without_embeddings() -> None: + provider = VectorStoreHistoryProvider(InMemoryStore(), application_id="app") + + await provider.save_messages("session", [Message(role="user", contents=["plain history"])]) + + assert provider._collection.definition.vector_fields == [] + assert [message.text for message in await provider.get_messages("session")] == ["plain history"] + + +async def test_vector_store_history_provider_supports_msgpack_contents() -> None: + provider = VectorStoreHistoryProvider( + InMemoryStore(), + application_id="app", + contents_format="msgpack", + ) + message = Message( + role="user", + contents=[ + "hello", + Content.from_uri("https://example.com/image.png", media_type="image/png"), + ], + ) + + await provider.save_messages("session", [message]) + + records = await provider._collection.get(filter=provider._scope_filter("session")) + contents_field = provider._collection.definition.try_get_field("contents") + assert contents_field is not None + assert contents_field.type_ == "bytes" + assert msgspec.msgpack.decode(records[0]["contents"]) == [content.to_dict() for content in message.contents] + assert [stored.to_dict() for stored in await provider.get_messages("session")] == [message.to_dict()] + + +async def test_vector_store_history_provider_compacts_loaded_context_but_searches_full_history() -> None: + async def keep_last_message(messages: list[Message]) -> bool: + del messages[:-1] + return True + + store = InMemoryStore() + provider = VectorStoreHistoryProvider( + store, + source_id="history", + application_id="app", + embedding_generator=MockEmbeddingClient(), + embedding_options={"dimensions": 2}, + compaction_strategy=keep_last_message, + include_search_tool=True, + ) + other_provider = VectorStoreHistoryProvider( + store, + source_id="other-history", + application_id="app", + embedding_generator=MockEmbeddingClient(), + embedding_options={"dimensions": 2}, + ) + await provider.save_messages( + "session", + [ + Message(role="user", contents=["old visible detail"]), + Message(role="assistant", contents=["recent visible detail"]), + ], + ) + await other_provider.save_messages("session", [Message(role="user", contents=["other source secret"])]) + + session = AgentSession(session_id="session") + context = SessionContext(input_messages=[], session_id=session.session_id) + with patch("agent_framework._vectors.mark_feature_used") as mark_feature_used_mock: + await provider.before_run(agent=cast(Any, None), session=session, context=context, state={}) + + assert [message.text for message in context.context_messages[provider.source_id]] == ["recent visible detail"] + assert context.instructions == [ + "Use search_history when relevant information is not present in the loaded history." + ] + assert [tool.name for tool in context.tools] == ["search_history"] + search_result = await context.tools[0].invoke( + arguments={"query": "visible detail"}, + skip_parsing=True, + ) + result_text = "\n".join(content.text or "" for content in search_result) + assert "old visible detail" in result_text + assert "recent visible detail" in result_text + assert "other source secret" not in result_text + mark_feature_used_mock.assert_any_call(FeatureIndex.CORE_VECTOR_STORE_HISTORY_PROVIDER) + + +@pytest.mark.parametrize( + ("kwargs", "error_type", "message"), + [ + ({"application_id": ""}, ValueError, "application_id"), + ({"application_id": "app", "embedding_options": {"dimensions": 2}}, ValueError, "requires"), + ( + {"application_id": "app", "embedding_generator": MockEmbeddingClient()}, + ValueError, + "requires embedding_options", + ), + ( + { + "application_id": "app", + "embedding_generator": MockEmbeddingClient(), + "embedding_options": {"dimensions": 0}, + }, + ValueError, + "must be positive", + ), + ({"application_id": "app", "contents_format": "yaml"}, ValueError, "contents_format"), + ], +) +def test_vector_store_history_provider_validates_configuration( + kwargs: dict[str, Any], + error_type: type[Exception], + message: str, +) -> None: + with pytest.raises(error_type, match=message): + VectorStoreHistoryProvider(InMemoryStore(), **kwargs) diff --git a/python/samples/02-agents/conversations/README.md b/python/samples/02-agents/conversations/README.md index 52845236c00..5d3c2e77294 100644 --- a/python/samples/02-agents/conversations/README.md +++ b/python/samples/02-agents/conversations/README.md @@ -10,6 +10,7 @@ These samples demonstrate different approaches to managing conversation history | [`custom_history_provider.py`](custom_history_provider.py) | Implement a custom history provider by extending `HistoryProvider`, enabling conversation persistence in your preferred storage backend. | | [`file_history_provider.py`](file_history_provider.py) | Use the experimental `FileHistoryProvider` with `FoundryChatClient` and a function tool so the local JSON Lines file shows the full tool-calling loop. | | [`file_history_provider_conversation_persistence.py`](file_history_provider_conversation_persistence.py) | Persist a tool-driven weather conversation with `FileHistoryProvider`, inspect the stored JSONL records, and continue with another city. | +| [`vector_store_history_provider.py`](vector_store_history_provider.py) | Store JSON- or MessagePack-encoded conversation contents in a provider-owned vector collection while loading compacted context and exposing scoped full-history search. | | [`cosmos_history_provider.py`](cosmos_history_provider.py) | Use Azure Cosmos DB as a history provider for durable conversation storage with `CosmosHistoryProvider`. | | [`cosmos_history_provider_conversation_persistence.py`](cosmos_history_provider_conversation_persistence.py) | Persist and resume conversations across application restarts using `CosmosHistoryProvider` — serialize session state, restore it, and continue with full Cosmos DB history. | | [`cosmos_history_provider_messages.py`](cosmos_history_provider_messages.py) | Direct message history operations — retrieve stored messages as a transcript, clear session history, and verify data deletion. | @@ -41,6 +42,10 @@ These samples demonstrate different approaches to managing conversation history - The sample writes plaintext JSONL conversation logs to disk; use a trusted local directory and avoid treating the history files as secure secret storage +**For `vector_store_history_provider.py`:** +- `OPENAI_API_KEY`: Used by the OpenAI chat and embedding clients +- The sample uses `InMemoryStore`; persisted history lasts only for the current process + **For Cosmos DB samples (`cosmos_history_provider*.py`):** - `FOUNDRY_PROJECT_ENDPOINT`: Your Microsoft Foundry project endpoint - `FOUNDRY_MODEL`: The Foundry model deployment name diff --git a/python/samples/02-agents/conversations/vector_store_history_provider.py b/python/samples/02-agents/conversations/vector_store_history_provider.py new file mode 100644 index 00000000000..5ab3bcb5c45 --- /dev/null +++ b/python/samples/02-agents/conversations/vector_store_history_provider.py @@ -0,0 +1,80 @@ +# Copyright (c) Microsoft. All rights reserved. + +import asyncio + +from agent_framework import Agent, InMemoryStore, SlidingWindowStrategy, VectorStoreHistoryProvider +from agent_framework.openai import OpenAIChatClient, OpenAIEmbeddingClient +from dotenv import load_dotenv + +load_dotenv() + +""" +This sample demonstrates a vector-store-backed conversation history provider. + +The provider owns its history collection and data model. It stores the full +conversation, loads a compacted window into each model invocation, and exposes +a search tool over the full transcript. Every storage and search operation is +scoped to the application, tenant, agent, provider source, and current session. + +Set `OPENAI_API_KEY` before running this sample. +""" + + +async def main() -> None: + """Run a compacted conversation backed by an in-memory vector store.""" + # 1. The history provider takes a store rather than a collection because it + # owns the history record model and collection definition. + history = VectorStoreHistoryProvider( + InMemoryStore(), + application_id="release-planning", + tenant_id="contoso", + agent_id="release-assistant", + contents_format="json", + embedding_generator=OpenAIEmbeddingClient( + model="text-embedding-3-small", + ), + embedding_options={ + "dimensions": 1536, + "encoding_format": "float", + }, + compaction_strategy=SlidingWindowStrategy( + keep_last_groups=2, + preserve_system=True, + ), + include_search_tool=True, + ) + + # 2. Only the compacted projection is loaded into the model context. The + # provider-owned search tool can still retrieve older scoped messages. + async with Agent( + client=OpenAIChatClient(model="gpt-5.4-nano"), + name="ReleaseAssistant", + instructions=( + "Help with release planning. Use the history search tool when an " + "older detail is not present in the loaded conversation." + ), + context_providers=[history], + ) as agent: + session = agent.create_session() + + for prompt in ( + "Remember that the deployment region is westus3.", + "Remember that the release ring is canary.", + "Remember that the rollback owner is Mira.", + "Remember that release alerts go to the operations channel.", + "Which deployment region did I choose? Search the full history if needed.", + ): + response = await agent.run(prompt, session=session) + print(f"User: {prompt}") + print(f"Assistant: {response.text}\n") + + +if __name__ == "__main__": + asyncio.run(main()) + + +""" +Sample output: +User: Which deployment region did I choose? Search the full history if needed. +Assistant: You chose westus3 as the deployment region. +""" diff --git a/python/samples/02-agents/vector_stores/README.md b/python/samples/02-agents/vector_stores/README.md index 89d2529fa9f..d8b855c3223 100644 --- a/python/samples/02-agents/vector_stores/README.md +++ b/python/samples/02-agents/vector_stores/README.md @@ -11,7 +11,8 @@ before calling the batch API. No database is needed for the in-memory examples. The model, format, and direct in-memory filter samples need no credentials. The search-tool sample loads the existing Azure AI Search hotel dataset and uses OpenAI for embeddings and the -agent; set `OPENAI_API_KEY` before running it. +agent; set `OPENAI_API_KEY` before running it. The context-provider examples +also use an in-memory collection with OpenAI for embeddings and agent responses. | File | Demonstrates | |------|--------------| @@ -19,6 +20,8 @@ agent; set `OPENAI_API_KEY` before running it. | [`optimized_data_formats.py`](optimized_data_formats.py) | Keeping NumPy vector fields and adapting pandas DataFrames to the batch record API. | | [`in_memory_filters.py`](in_memory_filters.py) | Direct vector search with `Filter` and `FilterGroup`. | | [`in_memory_search_tool.py`](in_memory_search_tool.py) | Model-set filter values with native typed `Param` declarations. | +| [`vector_collection_context_provider.py`](vector_collection_context_provider.py) | Adding default upsert, get, delete, and search tools for a user-defined collection. | +| [`vector_collection_context_provider_multiple_search_tools.py`](vector_collection_context_provider_multiple_search_tools.py) | Replacing default tools with multiple searches that expose different inputs and result detail. | | [`azure_ai_search.py`](azure_ai_search.py) | Native Azure vector/hybrid search with deterministic vectors and a disposable index. | | [`redis_store.py`](redis_store.py) | Native HASH and JSON storage, vector search, filtering, and lifecycle with a disposable Redis server. | @@ -78,4 +81,6 @@ uv run samples/02-agents/vector_stores/vector_store_models.py uv run samples/02-agents/vector_stores/optimized_data_formats.py uv run samples/02-agents/vector_stores/in_memory_filters.py uv run samples/02-agents/vector_stores/in_memory_search_tool.py +uv run samples/02-agents/vector_stores/vector_collection_context_provider.py +uv run samples/02-agents/vector_stores/vector_collection_context_provider_multiple_search_tools.py ``` diff --git a/python/samples/02-agents/vector_stores/vector_collection_context_provider.py b/python/samples/02-agents/vector_stores/vector_collection_context_provider.py new file mode 100644 index 00000000000..8ca18d526f3 --- /dev/null +++ b/python/samples/02-agents/vector_stores/vector_collection_context_provider.py @@ -0,0 +1,98 @@ +# Copyright (c) Microsoft. All rights reserved. + +import asyncio +from typing import Annotated, Any + +from agent_framework import ( + Agent, + InMemoryCollection, + VectorCollectionContextProvider, + VectorStoreField, + vectorstoremodel, +) +from agent_framework.openai import OpenAIChatClient, OpenAIEmbeddingClient +from dotenv import load_dotenv +from pydantic import BaseModel + +load_dotenv() + +""" +This sample demonstrates adding agent-usable vector collection tools through +VectorCollectionContextProvider. + +The provider takes a collection with the application's own data model. By +default it adds instructions plus upsert, get, delete, and search tools. +Approval behavior can be configured for all tools or selected tools through +one `approval_mode` argument. + +Set `OPENAI_API_KEY` before running this sample. +""" + + +@vectorstoremodel(collection_name="project-notes") +class ProjectNote(BaseModel): + """A project note managed by the agent.""" + + note_id: Annotated[str, VectorStoreField("key")] + title: Annotated[str, VectorStoreField("data", is_full_text_indexed=True)] + body: Annotated[str, VectorStoreField("data", is_full_text_indexed=True)] + body_vector: Annotated[ + list[float] | str | None, + VectorStoreField("vector", dimensions=1536, distance_function="cosine_similarity"), + ] = None + + def model_post_init(self, context: Any) -> None: + """Use the note body as the embedding input when no vector is supplied.""" + if self.body_vector is None: + self.body_vector = self.body + + +async def main() -> None: + """Give an agent CRUD and search access to a project-note collection.""" + collection: InMemoryCollection[str, ProjectNote] = InMemoryCollection( + ProjectNote, + embedding_generator=OpenAIEmbeddingClient( + model="text-embedding-3-small", + ), + ) + await collection.ensure_collection_exists() + + # Omitted mapping entries keep their safe defaults. This sample disables + # approval for upsert so the scripted interaction can run unattended; + # delete still requires approval, while get and search remain read-only. + collection_context = VectorCollectionContextProvider( + collection, + approval_mode={"upsert": "never_require"}, + ) + + async with Agent( + client=OpenAIChatClient(model="gpt-5.4-nano"), + name="ProjectNotesAssistant", + instructions="Use the collection tools to manage project notes. Do not invent stored notes.", + context_providers=[collection_context], + ) as agent: + session = agent.create_session() + + for prompt in ( + ( + "Save a note with id release-checklist, title Release checklist, " + "and body Verify rollback, monitoring, and owner sign-off." + ), + "Search the project notes for release readiness checks.", + "Get the note with id release-checklist.", + ): + response = await agent.run(prompt, session=session) + print(f"User: {prompt}") + print(f"Assistant: {response.text}\n") + + +if __name__ == "__main__": + asyncio.run(main()) + + +""" +Sample output: +Assistant: I saved the "Release checklist" note. +Assistant: The release checklist covers rollback, monitoring, and owner sign-off. +Assistant: The note release-checklist is titled "Release checklist". +""" diff --git a/python/samples/02-agents/vector_stores/vector_collection_context_provider_multiple_search_tools.py b/python/samples/02-agents/vector_stores/vector_collection_context_provider_multiple_search_tools.py new file mode 100644 index 00000000000..3b0312ec587 --- /dev/null +++ b/python/samples/02-agents/vector_stores/vector_collection_context_provider_multiple_search_tools.py @@ -0,0 +1,152 @@ +# Copyright (c) Microsoft. All rights reserved. + +import asyncio +from typing import Annotated, Any + +from agent_framework import ( + Agent, + Filter, + InMemoryCollection, + Param, + VectorCollectionContextProvider, + VectorStoreField, + create_vector_search_tool, + vectorstoremodel, +) +from agent_framework.openai import OpenAIChatClient, OpenAIEmbeddingClient +from dotenv import load_dotenv +from pydantic import BaseModel + +load_dotenv() + +""" +This sample demonstrates replacing the default collection search tool with +multiple searches that expose different inputs and result detail. + +The context provider disables its default CRUD and search tools, then adds a +generic hotel discovery tool and a hotel-detail tool. Both tools use the same +user-defined collection while keeping their own descriptions, filters, result +mappers, and approval settings. + +Set `OPENAI_API_KEY` before running this sample. +""" + + +@vectorstoremodel(collection_name="hotels") +class Hotel(BaseModel): + """A hotel record with summary and detail fields.""" + + hotel_id: Annotated[str, VectorStoreField("key")] + name: Annotated[str, VectorStoreField("data", is_full_text_indexed=True)] + city: Annotated[str, VectorStoreField("data", is_indexed=True)] + summary: Annotated[str, VectorStoreField("data", is_full_text_indexed=True)] + details: Annotated[str, VectorStoreField("data")] + summary_vector: Annotated[ + list[float] | str | None, + VectorStoreField("vector", dimensions=1536, distance_function="cosine_similarity"), + ] = None + + def model_post_init(self, context: Any) -> None: + """Use the summary as the embedding input when no vector is supplied.""" + if self.summary_vector is None: + self.summary_vector = self.summary + + +async def main() -> None: + """Expose two differently shaped search tools for one collection.""" + collection: InMemoryCollection[str, Hotel] = InMemoryCollection( + Hotel, + embedding_generator=OpenAIEmbeddingClient( + model="text-embedding-3-small", + ), + ) + await collection.ensure_collection_exists() + await collection.upsert([ + Hotel( + hotel_id="hotel-1", + name="Harbor View", + city="Lisbon", + summary="A waterfront Lisbon hotel with breakfast and a pool.", + details="Includes airport transfers, late checkout, a rooftop pool, and breakfast from 7:00.", + ), + Hotel( + hotel_id="hotel-2", + name="Old Town Rooms", + city="Lisbon", + summary="A quiet guesthouse in Lisbon's historic center.", + details="Includes breakfast, luggage storage, and self-service check-in. There is no pool.", + ), + Hotel( + hotel_id="hotel-3", + name="Market Square", + city="Seattle", + summary="A central Seattle hotel for short business stays.", + details="Includes a workspace, gym access, and express checkout.", + ), + ]) + + # 1. The discovery tool returns a compact result suitable for broad searches. + discovery_tool = create_vector_search_tool( + collection, + name="search_hotels", + description="Find hotels from a natural-language request and return a short summary.", + top=3, + result_mapper=lambda result: ( + f"{result['record'].hotel_id}: {result['record'].name} " + f"({result['record'].city}) - {result['record'].summary}" + ), + ) + + # 2. The detail tool exposes a required hotel_id filter and returns fields + # intentionally omitted from the discovery result. + details_tool = create_vector_search_tool( + collection, + name="search_hotel_details", + description="Get detailed information for a known hotel_id.", + filter=Filter( + "hotel_id", + "eq", + Param("hotel_id", str, description="The hotel_id returned by search_hotels.", required=True), + ), + top=1, + result_mapper=lambda result: ( + f"{result['record'].name} ({result['record'].hotel_id}): {result['record'].details}" + ), + ) + + # 3. Disable every generated tool and provide only the two tailored searches. + collection_context = VectorCollectionContextProvider( + collection, + include_upsert_tool=False, + include_get_tool=False, + include_delete_tool=False, + include_search_tool=False, + additional_search_tools=[discovery_tool, details_tool], + ) + + async with Agent( + client=OpenAIChatClient(model="gpt-5.4-nano"), + name="HotelAssistant", + instructions=("Use search_hotels for discovery. Use search_hotel_details only after you know a hotel_id."), + context_providers=[collection_context], + ) as agent: + session = agent.create_session() + + for prompt in ( + "Find a Lisbon hotel with a pool.", + "Give me the full details for that hotel.", + ): + response = await agent.run(prompt, session=session) + print(f"User: {prompt}") + print(f"Assistant: {response.text}\n") + + +if __name__ == "__main__": + asyncio.run(main()) + + +""" +Sample output: +Assistant: Harbor View (hotel-1) is a waterfront Lisbon hotel with breakfast and a pool. +Assistant: Harbor View includes airport transfers, late checkout, a rooftop pool, and breakfast from 7:00. +"""