diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 3b04ba06be6..c2b687ce266 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -61,6 +61,8 @@ /python/packages/claude/ @chetantoshniwal @eavanvalkenburg @giles17 @moonbox3 /python/packages/copilotstudio/ @chetantoshniwal @eavanvalkenburg @giles17 @moonbox3 /python/packages/core/ @chetantoshniwal @eavanvalkenburg @moonbox3 @TaoChenOSU @jpalvarezl @giles17 +/python/packages/core/agent_framework/_in_memory.py @chetantoshniwal @westey-m @eavanvalkenburg @giles17 @moonbox3 @TaoChenOSU @jpalvarezl @peibekwe @baywet @rogerbarreto @SergeyMenshykh +/python/packages/core/agent_framework/_vector_filters.py @chetantoshniwal @westey-m @eavanvalkenburg @giles17 @moonbox3 @TaoChenOSU @jpalvarezl @peibekwe @baywet @rogerbarreto @SergeyMenshykh /python/packages/core/agent_framework/_vectors.py @chetantoshniwal @westey-m @eavanvalkenburg @giles17 @moonbox3 @TaoChenOSU @jpalvarezl @peibekwe @baywet @rogerbarreto @SergeyMenshykh /python/packages/core/agent_framework/_workflows/ @chetantoshniwal @eavanvalkenburg @moonbox3 @TaoChenOSU @jpalvarezl /python/packages/core/agent_framework/_harness/ @chetantoshniwal @westey-m @eavanvalkenburg @moonbox3 diff --git a/docs/features/vector-stores-and-embeddings/README.md b/docs/features/vector-stores-and-embeddings/README.md index 15e058fd58a..ecdac79117a 100644 --- a/docs/features/vector-stores-and-embeddings/README.md +++ b/docs/features/vector-stores-and-embeddings/README.md @@ -36,17 +36,40 @@ This feature ports the vector store abstractions, embedding generator abstractio - `SupportsVectorUpsert` / `SupportsVectorSearch` — Protocols for duck-typing (follows `Supports` naming convention) - `BaseVectorCollection` / `BaseVectorSearch` — ABC base classes for implementations - `BaseVectorStore` — ABC base class for store operations (factory for collections, no protocol needed) -- **TypeVar naming convention**: `ModelT`, `KeyT`, `FilterT` (suffix T, per AF standard) +- **TypeVar naming convention**: `ModelT`, `KeyT` (suffix T, per AF standard) - **Support Pydantic for user-facing data models** — the `@vectorstoremodel` decorator and `VectorStoreCollectionDefinition` should work with Pydantic models, dataclasses, plain classes, and dicts - **Remove SK-specific dependencies** — no `KernelBaseModel`, `KernelFunction`, `KernelParameterMetadata`, `kernel_function`, `PromptExecutionSettings` - **Embedding types in `_types.py`**, embedding protocol/base class in `_clients.py` -- **All vector store specific types, enums, protocols, base classes** in `_vectors.py` +- **Portable filters** are data-only operation inputs in `_vector_filters.py`; no Python source or AST translation +- **Dependency-free local storage** is isolated in `_in_memory.py` +- **Vector store definitions, protocols, and base classes** remain in `_vectors.py` - **Error handling** uses AF's exception hierarchy (e.g., `IntegrationException` variants) +### Vector Filter Representation + +The original Phase 3 design accepted callable or string lambdas, recovered +their source with `inspect`, parsed the source into an AST, and delegated +translation to each connector. Phase 4 replaces that experimental model before +connector implementations depend on it. + +Options considered: + +- **Lambda source and connector-specific AST translation** — concise authoring, + but brittle across Python execution contexts, vulnerable to semantic drift, + and carries security concerns when evaluated locally. +- **A closed class hierarchy with one type per operation** — strongly typed, + but every provider-specific capability would require another core type. +- **A small data-only tree with namespaced provider extensions** — chosen. + `Filter` represents a field, operator, and value; `FilterGroup` provides + explicit AND, OR, and NOT composition; `Param` marks model-set values when + creating a search tool. Common operators have shared semantics, while + namespaced operators let connectors add structured provider capabilities + without accepting raw query source. + ### Package Structure - **Embedding types** (`Embedding`, `GeneratedEmbeddings`, `EmbeddingGenerationOptions`) in `agent_framework/_types.py` - **Embedding protocol + base class** (`SupportsGetEmbeddings`, `BaseEmbeddingClient`) in `agent_framework/_clients.py` -- **All vector store specific code** in a new `agent_framework/_vectors.py` module — this includes: +- **Vector store abstractions** in `agent_framework/_vectors.py` — this includes: - String literal aliases: `FieldTypes`, `IndexKind`, `DistanceFunction` - `VectorStoreField`, `VectorStoreCollectionDefinition` - `SearchResponse`, `SearchResults`, and explicit CRUD/search keyword arguments @@ -57,8 +80,8 @@ This feature ports the vector store abstractions, embedding generator abstractio - **OpenAI embeddings** in `agent_framework/openai/` (built into core, like OpenAI chat) - **Azure OpenAI embeddings** in `agent_framework/azure/` (built into core, follows `AzureOpenAIChatClient` pattern) - **Each vector store connector** in its own AF package under `packages/` -- **In-memory store** in core (no external deps) -- **TextSearch and its implementations** (Brave, Google) — last phase, separate work +- **Portable filters** (`Filter`, `FilterGroup`, `Param`) in `agent_framework/_vector_filters.py` +- **In-memory store** in `agent_framework/_in_memory.py` ## Naming: SK → AF @@ -190,12 +213,14 @@ This feature ports the vector store abstractions, embedding generator abstractio #### 3.1 — Vector store literal aliases and field types in `_vectors.py` - `FieldTypes`: `Literal["key", "vector", "data"]` -- `IndexKind`: literal alias covering `hnsw`, `flat`, `ivf_flat`, `disk_ann`, `quantized_flat`, `dynamic`, and `default` -- `DistanceFunction`: literal alias covering the supported similarity and distance functions +- `IndexKind`: common literal values for IDE guidance plus open provider-defined strings +- `DistanceFunction`: common literal values, including negative dot product, plus open provider-defined strings - `SearchType`: `Literal["vector", "keyword_hybrid"]` - `VectorStoreField` plain class (not Pydantic) + - Key fields can opt into store-generated keys + - Provider annotations are copied on construction and carry mutable connector-specific field configuration - `VectorStoreCollectionDefinition` class (not Pydantic internally, but supports Pydantic models as input) -- `SearchResponse` generic class +- `SearchResponse` generic `TypedDict` - `SearchResults` generic result container - Explicit keyword arguments on `get()` and `search()` instead of options classes - `DISTANCE_FUNCTION_DIRECTION_HELPER` dict @@ -223,7 +248,15 @@ This feature ports the vector store abstractions, embedding generator abstractio - Batch-oriented `upsert`, `get`, and `delete` - `upsert()` generates vector values by default and requires an embedding generator for every vector field; pass `generate_vectors=False` to preserve supplied vector values + - `generate_vectors` can also take a list or tuple of selected vector field names, allowing one model to combine + locally generated, precomputed, and provider-vectorized values + - After optional generation, check materialized dense sequence lengths against each field's `dimensions` for the + entire batch before connector conversion or writes. A mismatch raises `ValueError` with the zero-based record + index, logical field name, and expected/actual lengths; this rejection performs no writes + - Batch upsert does not promise atomicity; stable application keys make retries safer, while store-generated keys + may produce duplicates after a partial failure - CRUD `get()` excludes vectors by default; pass `include_vectors=True` when stored embeddings are needed + - CRUD `get()` accepts either keys or a portable filter with paging and ordering - `ensure_collection_exists`, `collection_exists`, `ensure_collection_deleted` - Async context manager support - `BaseVectorStore` — base for stores @@ -234,8 +267,13 @@ This feature ports the vector store abstractions, embedding generator abstractio - `BaseVectorSearch` — base for vector search - Single `search(search_type=...)` method with `search_type: Literal["vector", "keyword_hybrid"]` parameter — no enum, just a literal - `_inner_search` abstract method for implementations - - Filter building with lambda parser (AST-based) + - Portable `Filter` and `FilterGroup` trees passed unchanged to connector implementations + - Core validates portable request structure and deserializes returned records, but does not compute scores, + interpret score thresholds, or re-filter connector results. Connectors own execution and paging - Vector generation from values using embedding generator + - Check supplied or locally generated dense query length against the selected field's `dimensions` before + connector dispatch, including empty collections. In-memory search also checks array-like queries after its + numeric normalization; existing query/stored length checks remain in scoring #### 3.6 — Protocols for type checking - `SupportsVectorUpsert` — Protocol for upsert/get/delete operations @@ -251,9 +289,10 @@ This feature ports the vector store abstractions, embedding generator abstractio #### 3.8 — `create_vector_search_tool` - Standalone factory that creates an AF `FunctionTool` from any `SupportsVectorSearch` implementation - Wraps the single `search()` method, passing `search_type` parameter -- Accepts: `name`, `description`, `approval_mode`, `search_type`, `parameters`, `top`, `skip`, `filter`, `filter_mapper`, `result_mapper` -- Defaults to `query`; a custom Pydantic model or JSON schema can expose `top`, `skip`, and additional filter fields -- Custom schemas must require a string `query`; exposed `top` and `skip` fields must declare finite maximum values +- Accepts: `name`, `description`, `approval_mode`, `search_type`, `top`, `skip`, `filter`, `result_mapper` +- Defaults to a required string `query` +- Discovers `Param` values in filters and paging options, generating a closed JSON Schema without Pydantic +- Validates model-set values against native Python types and inline constraints before resolving the filter - The tool vectorizes the query, searches, and maps results to text or multimodal `Content` - Can also be a standalone factory function in `_vectors.py` @@ -265,19 +304,84 @@ This feature ports the vector store abstractions, embedding generator abstractio --- -### Phase 4: In-Memory Vector Store -**Goal:** Provide a zero-dependency vector store for testing and development. -**Mergeable:** Yes — first usable vector store. - -#### 4.1 — Port `InMemoryCollection` and `InMemoryStore` into core -- Place in `agent_framework/_vectors.py` (alongside the abstractions) -- Supports vector search (cosine similarity, etc.) -- No external dependencies - -#### 4.2 — Port FAISS extension (optional, can be separate package) -- Extends InMemory with FAISS indexing - -#### 4.3 — Tests and sample code +### Phase 4: Portable Filters and In-Memory Vector Store +**Goal:** Provide safe cross-store filters and a zero-dependency vector store for testing and development. +**Mergeable:** Yes — the filter contract and in-memory implementation can be reviewed independently. + +#### 4.1 — Replace source filters with portable data +- `Filter(field_name, operator, value)` for leaf conditions +- `FilterGroup(operator, filters)` for explicit AND, OR, and NOT composition +- Common operators plus namespaced provider extensions +- No callable inspection, source strings, AST parsing, `eval`, `exec`, or `compile` + +#### 4.1.1 — Connector extensibility aligned with Microsoft.Extensions.VectorData +- Index kinds and distance functions provide common literal hints but remain open to provider-defined strings +- `VectorStoreField.provider_annotations` is copied when the field is created and carries mutable provider-specific + configuration; the frozen field protects core schema attributes, not nested provider values +- Key fields can declare `is_auto_generated=True`; connectors decide which generated key types they support +- Search already supports provider-side query vectorization: without a local generator, connectors receive the + original `values` and `vector=None` +- Server-side write vectorization needs no core flag; leave that field out of `generate_vectors` so its source value + reaches a connector that supports it +- Dense vectors support numeric sequences and binary `bytes`; supplied and generated mutable `bytearray` values + normalize to `bytes` +- Float16, float32, float64, and integer element support remains connector-specific through `type_` and + `supported_vector_types` +- Sparse vectors remain provider-native values supplied through model codecs or `search(values=...)`; core does not + define a sparse representation or dense+sparse fusion mode + +#### 4.1.2 — Dense vector dimension checks +- Enforce declared dimensions at the shared write and search boundaries by default. This replaces the earlier + decision to leave dense length enforcement entirely to providers +- Check sequence length only, without copying, converting, or scanning elements solely for validation. Resolve + vector fields and storage names once per write batch, and validate final values after any local generation +- Null vectors remain allowed. Source text and non-sequence provider-native values are not treated as dense + vectors; `bytes`/`bytearray` length is not assumed to equal dimensionality. Connectors validate these representations +- The contract covers materialized non-string, non-binary sequences, not arbitrary provider-specific encodings, + numeric element validity, or revalidation on retrieval. Provider-side vectorization remains unchanged +- Local benchmarking of 1,000-record batches found length checking inexpensive relative to serialization and + in-memory copying. Use a straightforward pass without an opt-out flag or a more complex serialization path + +#### 4.2 — Derive search-tool parameters from filters +- A `Param` used as a complete filter value defines its model-visible name, native type, default, and constraints +- The tool factory emits a closed JSON Schema and resolves parameters before search +- An absent optional parameter without a default removes its containing filter; fixed filters remain unchanged +- `omit_if_none=True` also removes the leaf when its resolved argument is `None` (JSON `null`). It requires a + nullable type and an explicit `default=None`, for example + `Param("text", str | None, default=None, omit_if_none=True)`. Absent/null arguments omit the leaf; non-null + arguments retain normal type, constraint, and operator validation. This policy is not supported for paging. +- AND/OR groups evaluate their remaining children, not a `True` replacement for an omitted leaf. Groups left + empty, including NOT groups whose child is removed, are removed recursively. If the whole tree is removed, + search receives no filter; other search options still apply. +- Without the opt-in, explicit null values retain normal validation and provider semantics. Strings such as + `"*"` are literal filter values, not omission markers. +- String operators reject non-string operands centrally, after substitution for parameterized leaves +- Defaults and supplied mutable parameter values are copied per invocation, including nested containers +- Bound structural inspection before copying: filter depth/node limits apply to the tree and collection members, + including non-sequence collections such as sets; mapping keys cannot hide a `Param`. Limits also apply to + search tools whose search implementation has no collection definition. Unknown field names remain + connector-owned in that case +- Structural budgets do not sandbox arbitrary provider-native objects or trusted Python hooks + +#### 4.3 — Add `InMemoryCollection` and `InMemoryStore` +- Dedicated `_in_memory.py` module +- Shared process-local collection state, full CRUD/listing/order behavior, and flat vector search +- Pure-Python distance functions with no NumPy or SciPy dependency +- Scoring, filters, and thresholds execute locally before paging. `DEFAULT` resolves to cosine distance and + therefore accepts scores at or below the threshold, including zero for identical vectors +- Cosine calculations scale each vector independently to avoid overflow/underflow from finite magnitudes. + Non-finite scores are rejected for every metric, and unsupported distance functions fail before scanning records +- Hamming distance is the proportion of unequal dimensions, not a mismatch count; scores and thresholds use + the range zero to one, consistent with `scipy.spatial.distance.hamming` +- Strict filter evaluator over serialized mappings with the shared conservative resource limits +- Dictionary inputs and custom encoder outputs both pass through `msgspec.to_builtins` before storage, so + ordinary filtering operates on normalized data rather than original object comparison methods. Custom codecs + and connector overrides are trusted Python code, not a sandbox + +#### 4.4 — Tests and samples +- Direct filter composition and model-set search-tool filter parameters +- Security regressions for fail-closed behavior, scope preservation, and the SK exploit class +- FAISS remains deferred to its own optional connector phase --- @@ -390,8 +494,8 @@ Each connector follows the AF package structure: 8. **`create_vector_search_tool`**: The AF-native equivalent of SK's `create_search_function`. Instead of creating a `KernelFunction`, this creates an AF `FunctionTool` from any `SupportsVectorSearch` implementation. This allows agents to use vector search as a tool during conversations. Design: - `create_vector_search_tool(search, name, description, search_type, ...)` returns a `FunctionTool` - - The tool accepts declared parameters, performs embedding + vector search, and returns text or multimodal content - - Defaults to `query`; custom parameters can expose `top`, `skip`, and additional fields for the filter mapper + - The tool accepts `query` plus `Param` values discovered in filters or paging options + - It generates a closed native JSON Schema, performs embedding + vector search, and returns text or multimodal content - Lives in `_vectors.py` without expanding the structural search protocol 9. **CRUD tools**: A full set of create/read/update/delete tools for vector store collections, allowing agents to manage data in vector stores. Design: @@ -400,4 +504,11 @@ Each connector follows the AF package structure: - `create_delete_tool(...)` → tool for deleting records - These are separate from search and are placed in a later phase -10. **Score threshold filtering**: `search(score_threshold=...)` filters results by relevance score (ref: [SK .NET PR #13501](https://github.com/microsoft/semantic-kernel/pull/13501)). The semantics depend on the distance function: for similarity functions (cosine similarity, dot product), results *below* the threshold are filtered out; for distance functions (cosine distance, euclidean), results *above* the threshold are filtered out. Use `DISTANCE_FUNCTION_DIRECTION_HELPER` to determine direction. Connectors should implement this natively where the database supports it, falling back to client-side post-filtering otherwise. +10. **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, + threshold direction, and default metric. Execute filtering and thresholding natively where supported; + otherwise implement an explicit connector-local fallback coordinated with paging, or reject the unsupported + option rather than silently ignoring it. `DISTANCE_FUNCTION_DIRECTION_HELPER` remains available for + connectors implementing comparisons for common metrics locally; it is not a core capability gate. diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index c6a4affe0eb..75cb4f2e496 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -65,17 +65,50 @@ agent_framework/ - **`@tool`** decorator - Converts functions to tools - **`use_function_invocation()`** - Decorator to add automatic function calling to chat clients -### Vector stores (`_vectors.py`) +### Vector stores The vector store API is experimental under the shared `VECTOR_STORES` feature ID. - **`@vectorstoremodel`** - Declares key, data, and vector fields on dataclasses, Pydantic models, and plain classes - **`register_vectorstoremodel`** - Registers one definition and msgspec-backed codec pair per model type +- **`VectorStoreField`** - Frozen core key/data/vector metadata; common index and distance values remain open to + provider-defined strings, key fields can be store-generated, and copied `provider_annotations` remains mutable for + connector-specific configuration +- **`Filter` / `FilterGroup`** - Mutable data-only filter inputs shared by local and remote vector stores; collection + operations bound structural traversal before copying and pass an independent snapshot to connectors. Parameter + detection includes collection members and mapping keys; string operators require string operands after resolution +- **`Param`** - Native typed search-tool parameter reference embedded in filter values or paging options + - Defaults and supplied mutable values are copied per filter invocation, including for definition-less search tools + - Filter parameters may opt into null omission with a nullable type and explicit `default=None`, such as + `Param("text", str | None, default=None, omit_if_none=True)`; absent/null arguments remove that leaf, while + remaining AND/OR children still apply. Empty groups (including NOT with an omitted child) are removed + recursively; removing the whole tree means no filter. Paging parameters do not support null omission. - **`BaseVectorCollection`** - Base class for collection lifecycle and msgspec-backed record CRUD operations; - upserts generate embeddings by default and retrieval excludes vectors by default + upserts generate embeddings by default, retrieval excludes vectors by default, and filtered retrieval is an + alternate mode to key lookup +- **Embedding generation selection** - `generate_vectors=True` regenerates every vector field, `False` preserves all + values, and a list or tuple of logical vector field names generates only those fields so connectors can combine + local, precomputed, and provider-side vectorization +- **Vector payloads** - Shared dense query/generated vectors accept numeric sequences and binary bytes; connectors + declare supported element/representation types. Sparse vectors remain provider-native through codecs or + `search(values=...)`, not a core sparse type. +- **Vector dimensions** - Final dense sequence lengths are checked after optional generation for the whole write + batch before connector conversion or writes, and for the selected query field before search dispatch. In-memory + queries also check their normalized numeric sequence, including array-like inputs and empty collections. These + are length checks, not element validation; null vectors, source text, binary payloads, and non-sequence + provider-native representations remain connector-owned. - **`BaseVectorStore`** - Base class for stores that create collection clients -- **`BaseVectorSearch`** - Base class for vector and keyword-hybrid search +- **`BaseVectorSearch`** - Base class for vector and keyword-hybrid search; core validates portable requests and + deserializes results without interpreting thresholds or re-filtering returned scores. Connectors own scoring, + 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 +- **`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. + Scoring, filtering, and thresholds run locally before paging; `DEFAULT` means cosine distance and uses a maximum + distance threshold. + Shared serialization normalizes stored data; codecs and connector overrides remain trusted Python code - **`SupportsVectorUpsert`** / **`SupportsVectorSearch`** - Structural protocols for vector store capabilities ### Middleware (`_middleware.py`) diff --git a/python/packages/core/agent_framework/__init__.py b/python/packages/core/agent_framework/__init__.py index 241bcda3aae..c80b5025eb0 100644 --- a/python/packages/core/agent_framework/__init__.py +++ b/python/packages/core/agent_framework/__init__.py @@ -284,6 +284,17 @@ "validate_tool_mode", "validate_tools", ), + "._in_memory": ( + "InMemoryCollection", + "InMemoryStore", + ), + "._vector_filters": ( + "Filter", + "FilterGroup", + "FilterGroupOperator", + "FilterOperator", + "Param", + ), "._vectors": ( "DISTANCE_FUNCTION_DIRECTION_HELPER", "BaseVectorCollection", @@ -291,6 +302,7 @@ "BaseVectorStore", "DistanceFunction", "FieldTypes", + "GenerateVectors", "IndexKind", "SearchResponse", "SearchResults", @@ -493,6 +505,10 @@ "FileSkillsSource", "FileStoreEntry", "FileSystemAgentFileStore", + "Filter", + "FilterGroup", + "FilterGroupOperator", + "FilterOperator", "FilteringSkillsSource", "FinalT", "FinishReason", @@ -507,13 +523,16 @@ "FunctionalWorkflow", "FunctionalWorkflowAgent", "FunctionalWorkflowDefinition", + "GenerateVectors", "GeneratedEmbeddings", "GraphConnectivityError", "HistoryProvider", "InMemoryAgentFileStore", "InMemoryCheckpointStorage", + "InMemoryCollection", "InMemoryHistoryProvider", "InMemorySkillsSource", + "InMemoryStore", "InProcRunnerContext", "IndexKind", "InlineSkill", @@ -543,6 +562,7 @@ "MiddlewareTypes", "OuterFinalT", "OuterUpdateT", + "Param", "RawAgent", "ReleaseCandidateFeature", "ResponseStream", diff --git a/python/packages/core/agent_framework/__init__.pyi b/python/packages/core/agent_framework/__init__.pyi index 5816c90600f..a18848f56f3 100644 --- a/python/packages/core/agent_framework/__init__.pyi +++ b/python/packages/core/agent_framework/__init__.pyi @@ -123,6 +123,7 @@ from ._harness._tool_approval import ( create_always_approve_tool_response, create_always_approve_tool_with_arguments_response, ) +from ._in_memory import InMemoryCollection, InMemoryStore from ._mcp import ( MCPStdioTool, MCPStreamableHTTPTool, @@ -250,6 +251,7 @@ from ._types import ( validate_tool_mode, validate_tools, ) +from ._vector_filters import Filter, FilterGroup, FilterGroupOperator, FilterOperator, Param from ._vectors import ( DISTANCE_FUNCTION_DIRECTION_HELPER, BaseVectorCollection, @@ -257,6 +259,7 @@ from ._vectors import ( BaseVectorStore, DistanceFunction, FieldTypes, + GenerateVectors, IndexKind, SearchResponse, SearchResults, @@ -457,6 +460,10 @@ __all__ = [ "FileSkillsSource", "FileStoreEntry", "FileSystemAgentFileStore", + "Filter", + "FilterGroup", + "FilterGroupOperator", + "FilterOperator", "FilteringSkillsSource", "FinalT", "FinishReason", @@ -471,13 +478,16 @@ __all__ = [ "FunctionalWorkflow", "FunctionalWorkflowAgent", "FunctionalWorkflowDefinition", + "GenerateVectors", "GeneratedEmbeddings", "GraphConnectivityError", "HistoryProvider", "InMemoryAgentFileStore", "InMemoryCheckpointStorage", + "InMemoryCollection", "InMemoryHistoryProvider", "InMemorySkillsSource", + "InMemoryStore", "InProcRunnerContext", "IndexKind", "InlineSkill", @@ -507,6 +517,7 @@ __all__ = [ "MiddlewareTypes", "OuterFinalT", "OuterUpdateT", + "Param", "RawAgent", "ReleaseCandidateFeature", "ResponseStream", diff --git a/python/packages/core/agent_framework/_in_memory.py b/python/packages/core/agent_framework/_in_memory.py new file mode 100644 index 00000000000..b816aeabccf --- /dev/null +++ b/python/packages/core/agent_framework/_in_memory.py @@ -0,0 +1,604 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Dependency-free in-memory vector store.""" + +from __future__ import annotations + +import math +from collections.abc import Collection, Mapping, Sequence +from copy import deepcopy +from dataclasses import dataclass +from datetime import date, datetime, time, timedelta +from decimal import Decimal +from typing import Any, ClassVar, Generic, cast +from uuid import UUID, uuid4 + +from typing_extensions import TypeVar + +from ._feature_stage import ExperimentalFeature, experimental +from ._telemetry import FeatureIndex, mark_feature_used +from ._vector_filters import ( + Filter, + FilterExpression, + FilterGroup, + filter_values_equal, + require_filter_collection, + require_filter_string, + validate_filter, +) +from ._vectors import ( + DISTANCE_FUNCTION_DIRECTION_HELPER, + BaseVectorCollection, + BaseVectorSearch, + BaseVectorStore, + DistanceFunction, + EmbeddingClient, + SearchResults, + SearchType, + Vector, + VectorStoreCollectionDefinition, + _validate_vector_dimensions, # pyright: ignore[reportPrivateUsage] +) +from .exceptions import IntegrationException + +KeyT = TypeVar("KeyT", default=Any) +ModelT = TypeVar("ModelT", default=Any) + +_IN_MEMORY_FILTER_OPERATORS = frozenset({ + "eq", + "ne", + "gt", + "gte", + "lt", + "lte", + "between", + "in", + "not_in", + "is_null", + "is_not_null", + "exists", + "contains", + "contains_any", + "contains_all", + "starts_with", + "ends_with", + "contains_text", +}) +_SCALAR_FILTER_TYPES = (str, int, float, bool, bytes, date, datetime, time, timedelta, Decimal, UUID) +_DESCENDING_DISTANCE_FUNCTIONS = frozenset({"cosine_similarity", "dot_prod"}) +_IN_MEMORY_DISTANCE_FUNCTIONS = frozenset({ + "cosine_similarity", + "cosine_distance", + "dot_prod", + "negative_dot_prod", + "euclidean_distance", + "euclidean_squared_distance", + "manhattan", + "hamming", + "DEFAULT", +}) + + +@dataclass(slots=True) +class _InMemoryCollectionState: + definition: VectorStoreCollectionDefinition + records: dict[Any, dict[str, Any]] + exists: bool = False + + +def _numeric_vector(value: Any, *, field_name: str) -> tuple[float, ...]: + to_list = getattr(value, "tolist", None) + if callable(to_list): + value = to_list() + if not isinstance(value, Sequence) or isinstance(value, (str, bytes, bytearray)): + raise TypeError(f"Vector field '{field_name}' must contain a numeric sequence.") + vector: list[float] = [] + for item in cast(Sequence[Any], value): + if not isinstance(item, int | float) or isinstance(item, bool): + raise TypeError(f"Vector field '{field_name}' must contain only numbers.") + number = float(item) + if not math.isfinite(number): + raise ValueError(f"Vector field '{field_name}' must contain only finite numbers.") + vector.append(number) + if not vector: + raise ValueError(f"Vector field '{field_name}' cannot be empty.") + return tuple(vector) + + +def _paired_vectors(left: Vector, right: Vector) -> tuple[tuple[float, ...], tuple[float, ...]]: + normalized_left = _numeric_vector(left, field_name="query") + normalized_right = _numeric_vector(right, field_name="stored") + if len(normalized_left) != len(normalized_right): + raise ValueError( + f"Query and stored vectors must have the same length; " + f"got {len(normalized_left)} and {len(normalized_right)}." + ) + return normalized_left, normalized_right + + +def _calculate_score(left: Vector, right: Vector, distance_function: DistanceFunction) -> float: + left_values, right_values = _paired_vectors(left, right) + if distance_function in ("cosine_similarity", "cosine_distance", "DEFAULT"): + left_scale = max(abs(value) for value in left_values) + right_scale = max(abs(value) for value in right_values) + if left_scale == 0 or right_scale == 0: + raise ValueError("Cosine distance is undefined for zero-magnitude vectors.") + left_scaled = tuple(value / left_scale for value in left_values) + right_scaled = tuple(value / right_scale for value in right_values) + dot_product = math.fsum(a * b for a, b in zip(left_scaled, right_scaled, strict=True)) + left_norm = math.sqrt(math.fsum(value * value for value in left_scaled)) + right_norm = math.sqrt(math.fsum(value * value for value in right_scaled)) + similarity = dot_product / (left_norm * right_norm) + if not math.isfinite(similarity): + raise ValueError("Cosine similarity must be finite.") + similarity = max(-1.0, min(1.0, similarity)) + score = similarity if distance_function == "cosine_similarity" else 1 - similarity + elif distance_function == "dot_prod": + score = sum(a * b for a, b in zip(left_values, right_values, strict=True)) + elif distance_function == "negative_dot_prod": + score = -sum(a * b for a, b in zip(left_values, right_values, strict=True)) + elif distance_function == "euclidean_distance": + score = math.dist(left_values, right_values) + elif distance_function == "euclidean_squared_distance": + score = sum((a - b) * (a - b) for a, b in zip(left_values, right_values, strict=True)) + elif distance_function == "manhattan": + score = sum(abs(a - b) for a, b in zip(left_values, right_values, strict=True)) + elif distance_function == "hamming": + score = sum(a != b for a, b in zip(left_values, right_values, strict=True)) / len(left_values) + else: + raise NotImplementedError(f"Distance function '{distance_function}' is not supported by InMemoryCollection.") + if not math.isfinite(score): + raise ValueError(f"Distance function '{distance_function}' produced a non-finite score.") + return score + + +def _validate_in_memory_filter_value(value: Any) -> None: + if value is None or isinstance(value, _SCALAR_FILTER_TYPES): + if isinstance(value, float) and not math.isfinite(value): + raise ValueError("In-memory filter values must be finite.") + if isinstance(value, Decimal) and not value.is_finite(): + raise ValueError("In-memory filter values must be finite.") + return + if isinstance(value, Mapping): + raise TypeError("InMemoryCollection does not support mapping filter values.") + if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): + for item in cast(Sequence[Any], value): + _validate_in_memory_filter_value(item) + return + raise TypeError(f"InMemoryCollection does not support filter values of type '{type(value).__name__}'.") + + +def _resolve_record_value( + record: Mapping[str, Any], + filter_: Filter, + definition: VectorStoreCollectionDefinition, +) -> tuple[bool, Any]: + if "." in filter_.field_name: + raise NotImplementedError("InMemoryCollection does not support nested filter field paths.") + field = definition.try_get_field(filter_.field_name) + if field is None: + raise ValueError(f"Filter field '{filter_.field_name}' is not part of the vector store definition.") + storage_name = field.storage_name or field.name + return storage_name in record, record.get(storage_name) + + +def _evaluate_filter( + expression: FilterExpression, + record: Mapping[str, Any], + definition: VectorStoreCollectionDefinition, +) -> bool: + if isinstance(expression, FilterGroup): + values = (_evaluate_filter(item, record, definition) for item in expression.filters) + match expression.operator: + case "and": + return all(values) + case "or": + return any(values) + case "not": + return not next(values) + case _: + raise ValueError(f"Unknown filter group operator '{expression.operator}'.") + + exists, actual = _resolve_record_value(record, expression, definition) + operator = expression.operator + expected = expression.value + if operator == "exists": + return exists + if operator == "is_null": + return exists and actual is None + if operator == "is_not_null": + return exists and actual is not None + if not exists: + return False + if actual is None and operator not in ("eq", "ne"): + return False + + try: + match operator: + case "eq": + return filter_values_equal(actual, expected) + case "ne": + return not filter_values_equal(actual, expected) + case "gt": + return actual > expected + case "gte": + return actual >= expected + case "lt": + return actual < expected + case "lte": + return actual <= expected + case "between": + lower, upper = cast(Sequence[Any], expected) + return lower <= actual <= upper + case "in": + return any(filter_values_equal(actual, item) for item in cast(Collection[Any], expected)) + case "not_in": + return all(not filter_values_equal(actual, item) for item in cast(Collection[Any], expected)) + case "contains": + return any(filter_values_equal(item, expected) for item in require_filter_collection(actual)) + case "contains_any": + collection = require_filter_collection(actual) + return any( + filter_values_equal(item, value) for value in cast(Sequence[Any], expected) for item in collection + ) + case "contains_all": + collection = require_filter_collection(actual) + return all( + any(filter_values_equal(item, value) for item in collection) + for value in cast(Sequence[Any], expected) + ) + case "starts_with": + return require_filter_string(actual).startswith(require_filter_string(expected)) + case "ends_with": + return require_filter_string(actual).endswith(require_filter_string(expected)) + case "contains_text": + return require_filter_string(expected) in require_filter_string(actual) + case _: + raise NotImplementedError(f"Filter operator '{operator}' is not supported by InMemoryCollection.") + except TypeError as exc: + raise ValueError( + f"Filter operator '{operator}' cannot compare field '{expression.field_name}' with value {expected!r}." + ) from exc + + +@experimental(feature_id=ExperimentalFeature.VECTOR_STORES) +class InMemoryCollection( + BaseVectorCollection[KeyT, ModelT], + BaseVectorSearch[KeyT, ModelT], + Generic[KeyT, ModelT], +): + """Store and search vector records in process memory. + + This implementation is intended for tests and development. It is + nonpersistent, uses linear scans, and is not thread-safe. + + Scoring, filtering, and score thresholds are applied locally before paging. + The default metric is cosine distance, so its threshold is a maximum distance. + Hamming scores are the proportion of unequal dimensions, between zero and + one, not a mismatch count. Non-finite scores are rejected. Records are + normalized by the shared serializer before storage; custom codecs and + connector overrides remain trusted Python code, not sandboxed execution. + """ + + supported_search_types: ClassVar[set[SearchType]] = {"vector"} + supported_vector_types: ClassVar[set[str] | None] = { + "float", + "float16", + "float32", + "float64", + "int", + "int8", + "int16", + "int32", + "int64", + "uint8", + "uint16", + "uint32", + "uint64", + } + + def __init__( + self, + record_type: type[ModelT], + *, + definition: VectorStoreCollectionDefinition | None = None, + collection_name: str | None = None, + embedding_generator: EmbeddingClient | None = None, + ) -> None: + """Initialize an in-memory collection. + + Args: + record_type: The application record type. + definition: An explicit definition for dictionary or externally owned records. + collection_name: The collection name, overriding the model definition. + embedding_generator: The default client used to generate vectors. + """ + super().__init__( + record_type, + definition=definition, + collection_name=collection_name, + embedding_generator=embedding_generator, + ) + self._state = _InMemoryCollectionState(self.definition, {}) + + def _require_collection(self) -> None: + if not self._state.exists: + raise IntegrationException(f"Collection '{self.collection_name}' does not exist.") + + async def ensure_collection_exists( + self, + *, + operation_options: Mapping[str, Any] | None = None, + ) -> None: + """Create the collection when it does not exist.""" + mark_feature_used(FeatureIndex.CORE_VECTOR_STORES) + self._state.exists = True + + async def collection_exists( + self, + *, + operation_options: Mapping[str, Any] | None = None, + ) -> bool: + """Return whether the collection exists.""" + mark_feature_used(FeatureIndex.CORE_VECTOR_STORES) + return self._state.exists + + async def ensure_collection_deleted( + self, + *, + operation_options: Mapping[str, Any] | None = None, + ) -> None: + """Delete all records and mark the collection as absent.""" + mark_feature_used(FeatureIndex.CORE_VECTOR_STORES) + self._state.records.clear() + self._state.exists = False + + async def _inner_upsert( + self, + records: Sequence[Any], + *, + operation_options: Mapping[str, Any] | None = None, + ) -> Sequence[KeyT]: + self._require_collection() + keys: list[KeyT] = [] + for record in records: + if not isinstance(record, dict): + raise TypeError("In-memory records must serialize to dictionaries.") + stored_record = deepcopy(cast(dict[str, Any], record)) + key_field = self.definition.key_field + key_storage_name = self.definition.key_field_storage_name + if key_storage_name not in stored_record: + if not key_field.is_auto_generated: + raise ValueError(f"Record is missing key field '{key_field.name}'.") + if key_field.type_ == "str": + stored_record[key_storage_name] = str(uuid4()) + else: + raise NotImplementedError("InMemoryCollection can auto-generate only string keys.") + key = stored_record[key_storage_name] + try: + hash(key) + except TypeError as exc: + raise TypeError("In-memory record keys must be hashable.") from exc + self._state.records[key] = stored_record + keys.append(cast(KeyT, key)) + return keys + + async def _inner_get( + self, + *, + keys: Sequence[KeyT] | None = None, + filter: FilterExpression | None = None, + top: int = 10, + skip: int = 0, + order_by: Mapping[str, bool] | None = None, + include_vectors: bool = False, + operation_options: Mapping[str, Any] | None = None, + ) -> Sequence[Any] | None: + self._require_collection() + if keys is not None: + return [deepcopy(self._state.records[key]) for key in keys if key in self._state.records] + records = list(self._state.records.values()) + if filter is not None: + _validate_in_memory_filter(filter, self.definition) + records = [record for record in records if _evaluate_filter(filter, record, self.definition)] + if order_by: + for field_name, ascending in reversed(tuple(order_by.items())): + if not isinstance(ascending, bool): + raise TypeError(f"Order direction for field '{field_name}' must be a boolean.") + field = self.definition.try_get_field(field_name) + if field is None: + raise ValueError(f"Order field '{field_name}' is not part of the vector store definition.") + storage_name = field.storage_name or field.name + try: + records_with_values = [record for record in records if record[storage_name] is not None] + records_without_values = [record for record in records if record[storage_name] is None] + records_with_values.sort(key=lambda record: record[storage_name], reverse=not ascending) + records[:] = [*records_with_values, *records_without_values] + except (KeyError, TypeError) as exc: + raise ValueError(f"Records cannot be ordered by field '{field_name}'.") from exc + return deepcopy(records[skip : skip + top]) + + async def _inner_delete( + self, + keys: Sequence[KeyT], + *, + operation_options: Mapping[str, Any] | None = None, + ) -> None: + self._require_collection() + for key in keys: + self._state.records.pop(key, None) + + async def _inner_search( + self, + *, + search_type: SearchType, + filter: FilterExpression | None = None, + values: Any | None = None, + vector: Vector | None = None, + top: int = 3, + skip: int = 0, + include_vectors: bool = False, + vector_property_name: str | None = None, + additional_property_name: str | None = None, + score_threshold: float | None = None, + operation_options: Mapping[str, Any] | None = None, + ) -> SearchResults[Any]: + self._require_collection() + if search_type != "vector": + raise NotImplementedError("InMemoryCollection supports only vector search.") + if vector is None: + raise ValueError("InMemoryCollection vector search requires a query vector or embedding generator.") + query_vector = _numeric_vector(vector, field_name="query") + vector_field = self.definition.try_get_vector_field(vector_property_name) + if vector_field is None: + raise ValueError("InMemoryCollection vector search requires a vector field.") + _validate_vector_dimensions(query_vector, vector_field) + if filter is not None: + _validate_in_memory_filter(filter, self.definition) + + distance_function = vector_field.distance_function or "DEFAULT" + if distance_function == "DEFAULT": + distance_function = "cosine_distance" + if distance_function not in _IN_MEMORY_DISTANCE_FUNCTIONS: + raise NotImplementedError( + f"Distance function '{distance_function}' is not supported by InMemoryCollection." + ) + comparison = DISTANCE_FUNCTION_DIRECTION_HELPER[distance_function] + storage_name = vector_field.storage_name or vector_field.name + results: list[dict[str, Any]] = [] + key_storage_name = self.definition.key_field_storage_name + for record in self._state.records.values(): + if filter is not None and not _evaluate_filter(filter, record, self.definition): + continue + stored_vector = record.get(storage_name) + if stored_vector is None: + continue + try: + score = _calculate_score(query_vector, stored_vector, distance_function) + except TypeError as exc: + raise TypeError( + f"Record {record.get(key_storage_name)!r} has an invalid vector in field '{storage_name}': {exc}" + ) from exc + except ValueError as exc: + raise ValueError( + f"Record {record.get(key_storage_name)!r} has an invalid vector in field '{storage_name}': {exc}" + ) from exc + if score_threshold is not None and not comparison(score, score_threshold): + continue + results.append({"record": deepcopy(record), "score": score}) + results.sort( + key=lambda result: cast(float, result["score"]), + reverse=distance_function in _DESCENDING_DISTANCE_FUNCTIONS, + ) + total_count = len(results) + return SearchResults( + results[skip : skip + top], + metadata={"in_memory_total_count": total_count}, + ) + + def _get_record_from_result(self, result: Any) -> Any: + if not isinstance(result, Mapping): + raise TypeError("In-memory search results must be mappings.") + return cast(Mapping[str, Any], result)["record"] + + def _get_score_from_result(self, result: Any) -> float | None: + if not isinstance(result, Mapping): + raise TypeError("In-memory search results must be mappings.") + score = cast(Mapping[str, Any], result).get("score") + return cast(float | None, score) + + +def _walk_filters(expression: FilterExpression) -> tuple[Filter, ...]: + if isinstance(expression, Filter): + return (expression,) + return tuple(item for child in expression.filters for item in _walk_filters(child)) + + +def _validate_in_memory_filter( + expression: FilterExpression, + definition: VectorStoreCollectionDefinition, +) -> None: + validate_filter(expression, field_names=definition.names) + for item in _walk_filters(expression): + if item.operator not in _IN_MEMORY_FILTER_OPERATORS: + raise NotImplementedError(f"Filter operator '{item.operator}' is not supported by InMemoryCollection.") + if "." in item.field_name: + raise NotImplementedError("InMemoryCollection does not support nested filter field paths.") + _validate_in_memory_filter_value(item.value) + + +@experimental(feature_id=ExperimentalFeature.VECTOR_STORES) +class InMemoryStore(BaseVectorStore): + """Create in-memory collection clients that share process-local state.""" + + def __init__( + self, + *, + embedding_generator: EmbeddingClient | None = None, + ) -> None: + """Initialize an in-memory vector store. + + Args: + embedding_generator: The default client used by collection search and upsert operations. + """ + super().__init__(embedding_generator=embedding_generator) + self._collections: dict[str, _InMemoryCollectionState] = {} + + def get_collection( + self, + record_type: type[ModelT], + *, + definition: VectorStoreCollectionDefinition | None = None, + collection_name: str | None = None, + embedding_generator: EmbeddingClient | None = None, + ) -> InMemoryCollection[Any, ModelT]: + """Create a collection client tied to shared in-memory state. + + Args: + record_type: The application record type. + definition: An explicit definition for dictionary or externally owned records. + collection_name: The collection name, overriding the model definition. + embedding_generator: A collection-specific embedding client. + + Returns: + A collection client sharing state with other clients for the same name. + + Raises: + ValueError: If the collection name is already associated with another definition. + """ + mark_feature_used(FeatureIndex.CORE_VECTOR_STORES) + collection = InMemoryCollection( + record_type, + definition=definition, + collection_name=collection_name, + embedding_generator=embedding_generator if embedding_generator is not None else self.embedding_generator, + ) + state = self._collections.get(collection.collection_name) + if state is None: + self._collections[collection.collection_name] = collection._state # pyright: ignore[reportPrivateUsage] + return collection + if state.definition != collection.definition: + raise ValueError( + f"Collection '{collection.collection_name}' is already registered with another definition." + ) + collection._state = state # pyright: ignore[reportPrivateUsage] + return collection + + async def list_collection_names( + self, + *, + operation_options: Mapping[str, Any] | None = None, + ) -> Sequence[str]: + """List existing in-memory collection names.""" + mark_feature_used(FeatureIndex.CORE_VECTOR_STORES) + return sorted(name for name, state in self._collections.items() if state.exists) + + async def _inner_ensure_collection_deleted( + self, + collection_name: str, + *, + operation_options: Mapping[str, Any] | None = None, + ) -> None: + state = self._collections[collection_name] + state.records.clear() + state.exists = False diff --git a/python/packages/core/agent_framework/_vector_filters.py b/python/packages/core/agent_framework/_vector_filters.py new file mode 100644 index 00000000000..c4488c715c1 --- /dev/null +++ b/python/packages/core/agent_framework/_vector_filters.py @@ -0,0 +1,874 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Portable vector store filter expressions. + +Common operator semantics: + +- ``eq`` / ``ne`` compare one field value and treat booleans as distinct from numbers. +- ``gt`` / ``gte`` / ``lt`` / ``lte`` perform ordered scalar comparisons. +- ``between`` is inclusive and accepts exactly ``(lower, upper)``. +- ``in`` / ``not_in`` test whether the field value occurs in the supplied sequence. +- ``contains`` tests whether a non-mapping collection field contains one supplied value. +- ``contains_any`` / ``contains_all`` test collection fields against supplied sequences. +- ``is_null`` / ``is_not_null`` require the field to exist; use ``exists`` to test presence alone. +- ``starts_with`` / ``ends_with`` / ``contains_text`` require string operands. + +A missing field returns ``False`` for every operator except ``exists``. Use +``FilterGroup`` for explicit AND, OR, and NOT composition. Provider-specific +operators must be namespaced and are interpreted only by their connector. +""" + +from __future__ import annotations + +import keyword +import math +import re +from collections.abc import Collection, Mapping, Sequence +from copy import deepcopy +from dataclasses import dataclass +from decimal import Decimal +from itertools import islice +from types import UnionType +from typing import Any, Final, Literal, TypeAlias, Union, cast, get_args, get_origin + +from typing_extensions import Sentinel + +from ._feature_stage import ExperimentalFeature, experimental + +FilterOperator: TypeAlias = ( + Literal[ + "eq", + "ne", + "gt", + "gte", + "lt", + "lte", + "between", + "in", + "not_in", + "is_null", + "is_not_null", + "exists", + "contains", + "contains_any", + "contains_all", + "starts_with", + "ends_with", + "contains_text", + ] + | str +) +FilterGroupOperator: TypeAlias = Literal["and", "or", "not"] + +_STANDARD_FILTER_OPERATORS: Final[frozenset[str]] = frozenset({ + "eq", + "ne", + "gt", + "gte", + "lt", + "lte", + "between", + "in", + "not_in", + "is_null", + "is_not_null", + "exists", + "contains", + "contains_any", + "contains_all", + "starts_with", + "ends_with", + "contains_text", +}) +_NO_VALUE_OPERATORS: Final[frozenset[str]] = frozenset({"is_null", "is_not_null", "exists"}) +_TEXT_VALUE_OPERATORS: Final[frozenset[str]] = frozenset({"starts_with", "ends_with", "contains_text"}) +_SEQUENCE_VALUE_OPERATORS: Final[frozenset[str]] = frozenset({ + "between", + "in", + "not_in", + "contains_any", + "contains_all", +}) +_GROUP_OPERATORS: Final[frozenset[str]] = frozenset({"and", "or", "not"}) +_PROVIDER_OPERATOR_PATTERN: Final[re.Pattern[str]] = re.compile(r"[a-z][a-z0-9_]*(?:\.[a-z][a-z0-9_]*)+") +_PARAM_UNSET = Sentinel("_PARAM_UNSET") +_OMIT_FILTER = Sentinel("_OMIT_FILTER") +_MAX_FILTER_DEPTH: Final[int] = 8 +_MAX_FILTER_NODES: Final[int] = 64 + + +def _is_non_string_sequence(value: Any) -> bool: + return isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)) + + +def require_filter_collection(value: Any) -> Collection[Any]: + """Return a non-string, non-mapping collection filter value.""" + if not isinstance(value, Collection) or isinstance(value, (str, bytes, bytearray, Mapping)): + raise TypeError("Filter value must be a non-string, non-mapping collection.") + return cast(Collection[Any], value) + + +def require_filter_string(value: Any) -> str: + """Return a string filter value.""" + if not isinstance(value, str): + raise TypeError("Filter value must be a string.") + return value + + +def filter_values_equal(left: Any, right: Any) -> bool: + """Compare validated scalars, sequences, and mappings without equating booleans to numbers.""" + if isinstance(left, bool) or isinstance(right, bool): + return isinstance(left, bool) and isinstance(right, bool) and left == right + if _is_non_string_sequence(left) and _is_non_string_sequence(right): + # Preserve native container semantics, such as lists not equaling tuples. + return left == right and all( + filter_values_equal(left_item, right_item) + for left_item, right_item in zip(cast(Sequence[Any], left), cast(Sequence[Any], right), strict=True) + ) + if isinstance(left, Mapping) and isinstance(right, Mapping): + left_mapping = cast(Mapping[str, Any], left) + right_mapping = cast(Mapping[str, Any], right) + return left_mapping == right_mapping and all( + filter_values_equal(item, right_mapping[key]) for key, item in left_mapping.items() + ) + return left == right + + +def _validate_name(name: str, *, kind: str, allow_path: bool) -> None: + if not name: + raise ValueError(f"{kind} cannot be empty.") + segments = name.split(".") if allow_path else [name] + for segment in segments: + if ( + not segment.isidentifier() + or keyword.iskeyword(segment) + or (segment.startswith("__") and segment.endswith("__")) + ): + raise ValueError(f"Invalid {kind.lower()} '{name}'.") + + +@experimental(feature_id=ExperimentalFeature.VECTOR_STORES) +@dataclass(frozen=True, slots=True, init=False) +class Param: + """Reference one model-set search-tool parameter. + + An optional parameter without a default removes the containing filter when + omitted. Otherwise, the supplied value or declared default is substituted + before the filter reaches the vector store. A missing required parameter + raises an error. + + Set ``omit_if_none=True`` to remove the containing filter when the resolved + value is ``None`` (JSON ``null``). This requires a nullable ``value_type`` + and an explicit ``default=None``, for example + ``Param("text", str | None, default=None, omit_if_none=True)``. Both an + absent argument and an explicit ``None`` then omit the filter. Non-None + values still undergo the declared type and constraint checks. + + Omission removes the entire leaf, not the field value and not a boolean + ``True`` substitute. AND/OR groups evaluate their remaining children; + empty groups, including NOT groups whose child was removed, are removed + recursively. If the whole tree is removed, search receives no filter. + Fixed filters are retained. See ``FilterGroup`` for composition examples. + + With ``omit_if_none=False`` (the default), ``None`` is an ordinary supplied + value subject to type and operator validation, not an omission request. + Null omission is only supported for filter parameters, not paging options. + + Defaults and supplied mutable values are copied for each filter invocation. + Parameter data is bounded before copying or type validation. + """ + + name: str + value_type: Any + required: bool + _default: Any + omit_if_none: bool + description: str | None + minimum: float | int | None + maximum: float | int | None + min_length: int | None + max_length: int | None + + def __init__( + self, + name: str, + value_type: Any, + *, + required: bool = False, + default: Any = _PARAM_UNSET, + omit_if_none: bool = False, + description: str | None = None, + minimum: float | int | None = None, + maximum: float | int | None = None, + min_length: int | None = None, + max_length: int | None = None, + ) -> None: + """Initialize a parameter reference. + + Args: + name: The tool parameter name exposed to the model. + value_type: The Python type used for schema generation and validation. + Must accept ``None``, such as ``str | None``, when ``omit_if_none=True``. + required: Whether the model must supply the parameter. + default: The value used when an optional parameter is omitted. + Must be explicitly set to ``None`` when ``omit_if_none=True``. + omit_if_none: Whether a resolved ``None`` removes the containing filter. + Requires a nullable type and ``default=None``; not supported for paging. + description: The parameter description shown to the model. + minimum: The inclusive minimum for numeric values. + maximum: The inclusive maximum for numeric values. + min_length: The minimum length for string, array, or object values. + max_length: The maximum length for string, array, or object values. + + Raises: + TypeError: If the annotation is unsupported or ``omit_if_none`` is not a boolean. + ValueError: If the name or constraints are invalid, a required parameter declares a default, + or ``omit_if_none=True`` is used without a nullable type and explicit ``default=None``. + """ + _validate_name(name, kind="Parameter name", allow_path=False) + if name == "query": + raise ValueError("'query' is reserved by vector search tools.") + if required and default is not _PARAM_UNSET: + raise ValueError("A required parameter cannot declare a default.") + if not isinstance(omit_if_none, bool): + raise TypeError("Param omit_if_none must be a boolean.") + if omit_if_none and default is not None: + raise ValueError("Param omit_if_none=True requires an explicit default=None.") + if minimum is not None and (isinstance(minimum, bool) or not math.isfinite(minimum)): + raise ValueError("Param minimum must be a finite number.") + if maximum is not None and (isinstance(maximum, bool) or not math.isfinite(maximum)): + raise ValueError("Param maximum must be a finite number.") + if minimum is not None and maximum is not None and minimum > maximum: + raise ValueError("Param minimum cannot exceed maximum.") + if min_length is not None and ( + not isinstance(min_length, int) or isinstance(min_length, bool) or min_length < 0 + ): + raise ValueError("Param min_length must be a non-negative integer.") + if max_length is not None and ( + not isinstance(max_length, int) or isinstance(max_length, bool) or max_length < 0 + ): + raise ValueError("Param max_length must be a non-negative integer.") + if min_length is not None and max_length is not None and min_length > max_length: + raise ValueError("Param min_length cannot exceed max_length.") + if default is not _PARAM_UNSET: + _validate_param_data(default) + object.__setattr__(self, "name", name) + object.__setattr__(self, "value_type", value_type) + object.__setattr__(self, "required", required) + object.__setattr__(self, "_default", deepcopy(default)) + object.__setattr__(self, "omit_if_none", omit_if_none) + object.__setattr__(self, "description", description) + object.__setattr__(self, "minimum", minimum) + object.__setattr__(self, "maximum", maximum) + object.__setattr__(self, "min_length", min_length) + object.__setattr__(self, "max_length", max_length) + param_schema(self) + + @property + def has_default(self) -> bool: + """Return whether the parameter declares a default value.""" + return self._default is not _PARAM_UNSET + + @property + def default(self) -> Any: + """Return an independent copy of the default value.""" + return deepcopy(self._default) + + def __deepcopy__(self, memo: dict[int, Any]) -> Param: + return self + + +def _json_type(value_type: Any) -> str | None: + if value_type is str: + return "string" + if value_type is int: + return "integer" + if value_type is float: + return "number" + if value_type is bool: + return "boolean" + if value_type is type(None): + return "null" + return None + + +def _schema_for_type(value_type: Any) -> dict[str, Any]: + origin = get_origin(value_type) + if origin is Literal: + values = list(get_args(value_type)) + json_types: list[str] = [] + for value in values: + json_type = _json_type(type(value)) + if json_type is None: + raise TypeError("Param Literal values must be JSON-compatible scalar values.") + if isinstance(value, float) and not math.isfinite(value): + raise ValueError("Param Literal numbers must be finite.") + if json_type not in json_types: + json_types.append(json_type) + schema: dict[str, Any] = {"enum": values} + if len(json_types) == 1: + schema["type"] = json_types[0] + else: + schema["anyOf"] = [{"type": json_type} for json_type in json_types] + return schema + if origin in (Union, UnionType): + return {"anyOf": [_schema_for_type(item) for item in get_args(value_type)]} + if origin is tuple: + args = get_args(value_type) + if len(args) != 2 or args[1] is not Ellipsis: + raise TypeError("Param tuple types must use the homogeneous tuple[T, ...] form.") + item_type = args[0] + return {"type": "array", "items": _schema_for_type(item_type)} + if origin in (list, Sequence): + args = get_args(value_type) + item_type = args[0] if args else Any + return {"type": "array", "items": _schema_for_type(item_type)} + if origin in (dict, Mapping): + args = get_args(value_type) + key_annotation, value_annotation = args if len(args) == 2 else (Any, Any) + if key_annotation is not str: + raise TypeError("Param mapping keys must use the str type.") + return {"type": "object", "additionalProperties": _schema_for_type(value_annotation)} + if value_type is Any: + raise TypeError("Param requires an explicit JSON-compatible value type.") + schema_type = _json_type(value_type) + if schema_type is not None: + return {"type": schema_type} + raise TypeError(f"Param type '{value_type}' cannot be represented as JSON Schema.") + + +def param_schema(param: Param) -> dict[str, Any]: + """Build a JSON Schema property for a parameter without Pydantic.""" + schema = _schema_for_type(param.value_type) + schema_types = _schema_types(schema) + if param.omit_if_none and "null" not in schema_types: + raise ValueError("Param omit_if_none=True requires a value type that accepts None.") + non_null_schema_types = schema_types - {"null"} + if (param.minimum is not None or param.maximum is not None) and ( + not non_null_schema_types or not non_null_schema_types <= {"integer", "number"} + ): + raise ValueError("Param numeric constraints require numeric value types.") + if param.description is not None: + schema["description"] = param.description + if param.has_default: + schema["default"] = param.default + if param.minimum is not None: + schema["minimum"] = param.minimum + if param.maximum is not None: + schema["maximum"] = param.maximum + if (param.min_length is not None or param.max_length is not None) and len(non_null_schema_types) != 1: + raise ValueError("Param length constraints require one string, array, or object type.") + schema_type = next(iter(non_null_schema_types), None) + if schema_type not in (None, "string", "array", "object") and ( + param.min_length is not None or param.max_length is not None + ): + raise ValueError("Param length constraints require a string, array, or object type.") + if param.min_length is not None: + length_key = ( + "minLength" if schema_type == "string" else "minProperties" if schema_type == "object" else "minItems" + ) + _set_schema_constraint(schema, schema_type, length_key, param.min_length) + if param.max_length is not None: + length_key = ( + "maxLength" if schema_type == "string" else "maxProperties" if schema_type == "object" else "maxItems" + ) + _set_schema_constraint(schema, schema_type, length_key, param.max_length) + return schema + + +def _schema_types(schema: Mapping[str, Any]) -> set[str]: + schema_type = schema.get("type") + if isinstance(schema_type, str): + return {schema_type} + variants = schema.get("anyOf") + if not isinstance(variants, Sequence): + return set() + return { + item_type + for variant in cast(Sequence[Any], variants) + if isinstance(variant, Mapping) + for item_type in _schema_types(cast(Mapping[str, Any], variant)) + } + + +def _set_schema_constraint( + schema: dict[str, Any], + schema_type: str | None, + key: str, + value: Any, +) -> None: + if schema.get("type") == schema_type: + schema[key] = value + return + for variant in cast(Sequence[Any], schema.get("anyOf", ())): + if isinstance(variant, dict): + typed_variant = cast(dict[str, Any], variant) + if typed_variant.get("type") == schema_type: + typed_variant[key] = value + + +def _matches_param_type(value: Any, value_type: Any) -> bool: + origin = get_origin(value_type) + if origin is Literal: + return any(filter_values_equal(value, item) for item in get_args(value_type)) + if origin in (Union, UnionType): + return any(_matches_param_type(value, item) for item in get_args(value_type)) + if origin is tuple: + if not _is_non_string_sequence(value): + return False + args = get_args(value_type) + if len(args) != 2 or args[1] is not Ellipsis: + return False + return all(_matches_param_type(item, args[0]) for item in cast(Sequence[Any], value)) + if origin in (list, Sequence): + if not _is_non_string_sequence(value): + return False + args = get_args(value_type) + item_type = args[0] if args else Any + return all(_matches_param_type(item, item_type) for item in cast(Sequence[Any], value)) + if origin in (dict, Mapping): + if not isinstance(value, Mapping): + return False + args = get_args(value_type) + key_type, item_type = args if len(args) == 2 else (Any, Any) + mapping = cast(Mapping[Any, Any], value) + return all( + _matches_param_type(key, key_type) and _matches_param_type(item, item_type) for key, item in mapping.items() + ) + if value_type is Any: + return True + if value_type is int: + return isinstance(value, int) and not isinstance(value, bool) + if value_type is float: + return isinstance(value, int | float) and not isinstance(value, bool) + return isinstance(value, value_type) + + +@dataclass(slots=True) +class _ParamDataValidator: + max_depth: int = 16 + max_nodes: int = 256 + node_count: int = 0 + + def validate(self, value: Any, *, seen: set[int] | None = None, depth: int = 0) -> None: + self.node_count += 1 + if self.node_count > self.max_nodes: + raise ValueError(f"Search parameter values cannot contain more than {self.max_nodes} nodes.") + if depth > self.max_depth: + raise ValueError(f"Search parameter values cannot exceed a depth of {self.max_depth}.") + if isinstance(value, float) and not math.isfinite(value): + raise ValueError("Filter and search parameter numbers must be finite.") + if isinstance(value, Decimal) and not value.is_finite(): + raise ValueError("Filter and search parameter numbers must be finite.") + if seen is None: + seen = set() + if isinstance(value, Mapping): + mapping = cast(Mapping[Any, Any], value) + if id(mapping) in seen: + raise ValueError("Search parameter values cannot contain cycles.") + seen.add(id(mapping)) + try: + for item in mapping.values(): + self.validate(item, seen=seen, depth=depth + 1) + finally: + seen.remove(id(mapping)) + elif _is_non_string_sequence(value): + sequence = cast(Sequence[Any], value) + if id(sequence) in seen: + raise ValueError("Search parameter values cannot contain cycles.") + seen.add(id(sequence)) + try: + for item in sequence: + self.validate(item, seen=seen, depth=depth + 1) + finally: + seen.remove(id(sequence)) + + +def _validate_param_data(value: Any) -> None: + _ParamDataValidator().validate(value) + + +def validate_param_value(param: Param, value: Any) -> Any: + """Validate one parameter value with native type and constraint checks.""" + _validate_param_data(value) + if not _matches_param_type(value, param.value_type): + raise TypeError(f"Search parameter '{param.name}' does not match {param.value_type}.") + if isinstance(value, int | float) and not isinstance(value, bool): + if param.minimum is not None and value < param.minimum: + raise ValueError(f"Search parameter '{param.name}' must be at least {param.minimum}.") + if param.maximum is not None and value > param.maximum: + raise ValueError(f"Search parameter '{param.name}' must be at most {param.maximum}.") + if isinstance(value, str | Sequence | Mapping): + sized_value = cast(str | Sequence[Any] | Mapping[Any, Any], value) + if param.min_length is not None and len(sized_value) < param.min_length: + raise ValueError(f"Search parameter '{param.name}' is shorter than {param.min_length}.") + if param.max_length is not None and len(sized_value) > param.max_length: + raise ValueError(f"Search parameter '{param.name}' is longer than {param.max_length}.") + return cast(Any, value) + + +@experimental(feature_id=ExperimentalFeature.VECTOR_STORES) +@dataclass(slots=True, init=False) +class Filter: + """Describe one data-only vector store filter. + + A ``Param`` must be the complete value of a leaf. In a search tool, + ``Filter("description", "contains_text", + Param("text", str | None, default=None, omit_if_none=True))`` is omitted + when ``text`` is absent or explicitly ``None`` (JSON ``null``). + ``omit_if_none=True`` requires both a nullable parameter type and an + explicit ``default=None``. Non-None arguments use normal operator + semantics; strings such as ``"*"`` are not special omission values. + + Omission removes this leaf from its group rather than making it match + every record. Remaining AND/OR children still apply; a NOT group is + removed if its child is removed. Empty groups are removed recursively, + and removing the whole tree means search receives no filter. + + With ``omit_if_none=False`` (the default), a supplied ``None`` is validated + as a value, not omitted. A literal ``Filter(..., value=None)`` does not + opt into omission either; use ``is_null`` to test for a null field. + + Nested parameters are forbidden in collection members and mapping keys as + well as mapping values. Structural depth/node limits apply during parameter + inspection and before operation snapshots are copied; oversized inputs may + therefore fail at construction as well as at operation time. + """ + + field_name: str + operator: FilterOperator + value: Any + + def __init__(self, field_name: str, operator: FilterOperator, value: Any = None) -> None: + """Initialize a filter. + + Args: + field_name: The logical model field, optionally followed by a provider-supported path. + operator: A standard operator or a namespaced provider operator. + value: The structured value consumed by the operator. + + Raises: + TypeError: If the field name or operator is not a string, or an operand has an invalid shape. + ValueError: If the field name or operator is invalid, an operator that takes no value receives one, + an operator that requires a value receives ``None``, ``between`` does not receive two boundaries, + a ``Param`` is nested inside a larger value, or structural limits are exceeded. + """ + if not isinstance(field_name, str): + raise TypeError("Filter field_name must be a string.") + if not isinstance(operator, str): + raise TypeError("Filter operator must be a string.") + _validate_name(field_name, kind="Filter field name", allow_path=True) + if operator not in _STANDARD_FILTER_OPERATORS and not _PROVIDER_OPERATOR_PATTERN.fullmatch(operator): + raise ValueError( + f"Unknown filter operator '{operator}'. Provider-specific operators must use a namespaced name." + ) + self.field_name = field_name + self.operator = operator + self.value = value + _validate_filter_value_shape(self, allow_params=True) + + +@experimental(feature_id=ExperimentalFeature.VECTOR_STORES) +@dataclass(slots=True, init=False) +class FilterGroup: + """Combine vector store filters with explicit boolean semantics. + + Search tools resolve parameters before evaluating the group. An absent + optional parameter without a default removes its leaf. A parameter such + as ``Param("text", str | None, default=None, omit_if_none=True)`` also + removes its leaf for an explicit ``None`` (JSON ``null``). This opt-in + requires a nullable type and an explicit ``default=None``. + + After omission, ``"and"`` requires all remaining children to match, and + ``"or"`` requires any remaining child to match. For example, combining + ``Filter("rating", "gte", 4)`` with an omitted text filter leaves only + the rating condition in either group. The omitted leaf is not replaced + with ``True``, which would make an OR group match every record. + + A ``"not"`` group negates its remaining child; if that child is removed, + the NOT group is removed too. Any group left with no children is removed + recursively. If the whole tree disappears, search receives no filter; + other search options still apply. Fixed children are never omitted. + + With ``omit_if_none=False`` (the default), explicit ``None`` values retain + normal type and operator validation rather than removing a child. + """ + + operator: FilterGroupOperator + filters: tuple[Filter | FilterGroup, ...] + + def __init__(self, operator: FilterGroupOperator, filters: Sequence[Filter | FilterGroup]) -> None: + """Initialize a filter group. + + Args: + operator: ``"and"``, ``"or"``, or unary ``"not"``. + filters: The filters combined by the operator. + + Raises: + TypeError: If the operator or filters have unsupported types. + ValueError: If the operator or number of filters is invalid, including structural limits. + """ + if not isinstance(operator, str): + raise TypeError("Filter group operator must be a string.") + if operator not in _GROUP_OPERATORS: + raise ValueError(f"Unknown filter group operator '{operator}'.") + if not _is_non_string_sequence(filters): + raise TypeError("FilterGroup filters must be a sequence.") + resolved_filters = tuple(islice(filters, _MAX_FILTER_NODES)) + if len(resolved_filters) >= _MAX_FILTER_NODES: + raise ValueError(f"Filters cannot contain more than {_MAX_FILTER_NODES} nodes.") + if not resolved_filters: + raise ValueError("FilterGroup requires at least one filter.") + if operator == "not" and len(resolved_filters) != 1: + raise ValueError("A 'not' FilterGroup requires exactly one filter.") + if any(not isinstance(item, Filter | FilterGroup) for item in resolved_filters): + raise TypeError("FilterGroup entries must be Filter or FilterGroup instances.") + self.operator = operator + self.filters = resolved_filters + + +FilterExpression: TypeAlias = Filter | FilterGroup + + +def _snapshot_filter(filter_: FilterExpression, *, active: set[int]) -> FilterExpression: + if id(filter_) in active: + raise ValueError("Filter expressions cannot contain cycles.") + active.add(id(filter_)) + try: + if isinstance(filter_, Filter): + value = ( + _snapshot_filter(filter_.value, active=active) + if isinstance(filter_.value, Filter | FilterGroup) + else deepcopy(filter_.value) + ) + return Filter(filter_.field_name, filter_.operator, value) + if isinstance(filter_, FilterGroup): + return FilterGroup( + filter_.operator, + tuple(_snapshot_filter(item, active=active) for item in filter_.filters), + ) + finally: + active.remove(id(filter_)) + raise TypeError("filter must be a Filter or FilterGroup.") + + +def snapshot_filter(filter_: FilterExpression) -> FilterExpression: + """Bound and validate the filter structure before copying it for one operation.""" + validate_filter(filter_, allow_params=True) + return _snapshot_filter(filter_, active=set()) + + +def _value_contains_param(value: Any) -> bool: + active: set[int] = set() + node_count = 0 + + def visit(value: Any, depth: int) -> bool: + nonlocal node_count + node_count += 1 + if node_count > _MAX_FILTER_NODES: + raise ValueError(f"Filters cannot contain more than {_MAX_FILTER_NODES} nodes.") + if depth > _MAX_FILTER_DEPTH: + raise ValueError(f"Filter values cannot exceed a depth of {_MAX_FILTER_DEPTH}.") + if isinstance(value, Param): + return True + if isinstance(value, (str, bytes, bytearray)) or not isinstance(value, Filter | FilterGroup | Collection): + return False + identity = id(cast(object, value)) + if identity in active: + raise ValueError("Filter values cannot contain cycles.") + active.add(identity) + try: + if isinstance(value, Filter): + return visit(value.value, depth + 1) + if isinstance(value, FilterGroup): + return any(visit(item, depth + 1) for item in value.filters) + if isinstance(value, Mapping): + return any( + (not isinstance(key, str) and visit(key, depth + 1)) or visit(item, depth + 1) + for key, item in cast(Mapping[Any, Any], value).items() + ) + return any(visit(item, depth + 1) for item in cast(Collection[Any], value)) + finally: + active.remove(identity) + + return visit(value, 1) + + +def _validate_filter_value_shape(filter_: Filter, *, allow_params: bool) -> None: + value = filter_.value + has_param = _value_contains_param(value) + if has_param and not isinstance(value, Param): + raise ValueError("Param must be the entire Filter value, not nested inside a collection or mapping.") + if has_param and not allow_params: + raise ValueError("Param references must be resolved before searching.") + if filter_.operator not in _STANDARD_FILTER_OPERATORS: + return + if filter_.operator in _NO_VALUE_OPERATORS: + if value is not None: + raise ValueError(f"Filter operator '{filter_.operator}' does not accept a value.") + return + if value is None: + raise ValueError(f"Filter operator '{filter_.operator}' requires a value.") + if filter_.operator in _TEXT_VALUE_OPERATORS and not has_param: + require_filter_string(value) + if filter_.operator not in _SEQUENCE_VALUE_OPERATORS or has_param: + return + if not _is_non_string_sequence(value): + raise TypeError(f"Filter operator '{filter_.operator}' requires a sequence value.") + if filter_.operator == "between" and len(cast(Sequence[Any], value)) != 2: + raise ValueError("'between' requires exactly two boundary values.") + + +def iter_filter_params(value: Any) -> tuple[Param, ...]: + if isinstance(value, Param): + return (value,) + if isinstance(value, Filter): + filter_value: Any = value.value + return (filter_value,) if isinstance(filter_value, Param) else () + if isinstance(value, FilterGroup): + return tuple(param for item in value.filters for param in iter_filter_params(item)) + return () + + +def _resolve_param_value( + value: Any, + arguments: Mapping[str, Any], +) -> Any: + if isinstance(value, Param): + if value.name in arguments: + resolved = arguments[value.name] + elif value.has_default: + resolved = value.default + elif value.required: + raise TypeError(f"Missing required search parameter '{value.name}'.") + else: + return _OMIT_FILTER + return _OMIT_FILTER if value.omit_if_none and resolved is None else deepcopy(resolved) + return deepcopy(value) + + +def resolve_filter_params( + filter_: FilterExpression, + arguments: Mapping[str, Any], +) -> FilterExpression | None: + if isinstance(filter_, Filter): + value = _resolve_param_value(filter_.value, arguments) + if value is _OMIT_FILTER: + return None + return Filter(filter_.field_name, filter_.operator, value) + resolved_filters = tuple( + resolved for item in filter_.filters if (resolved := resolve_filter_params(item, arguments)) is not None + ) + if not resolved_filters: + return None + if filter_.operator == "not" and len(resolved_filters) != 1: + return None + return FilterGroup(filter_.operator, resolved_filters) + + +@dataclass(slots=True) +class _FilterValidator: + field_names: Collection[str] | None + allow_params: bool + max_depth: int + max_nodes: int + node_count: int = 0 + + def validate_value(self, value: Any, *, depth: int, active: set[int]) -> None: + if depth > self.max_depth: + raise ValueError(f"Filter values cannot exceed a depth of {self.max_depth}.") + if isinstance(value, float) and not math.isfinite(value): + raise ValueError("Filter numbers must be finite.") + if isinstance(value, Decimal) and not value.is_finite(): + raise ValueError("Filter numbers must be finite.") + if isinstance(value, Param): + raise ValueError("Param must be the entire Filter value, not nested inside a collection or mapping.") + if isinstance(value, Filter | FilterGroup): + self.validate_expression(value, depth=depth, relative_fields=True, active=active) + return + if isinstance(value, Mapping): + mapping = cast(Mapping[Any, Any], value) + if id(mapping) in active: + raise ValueError("Filter values cannot contain cycles.") + active.add(id(mapping)) + try: + for key, item in mapping.items(): + if not isinstance(key, str): + raise TypeError("Filter mapping keys must be strings.") + self.count_node() + self.validate_value(item, depth=depth + 1, active=active) + finally: + active.remove(id(mapping)) + return + if isinstance(value, Collection) and not isinstance(value, (str, bytes, bytearray)): + collection = cast(Collection[Any], value) + if id(collection) in active: + raise ValueError("Filter values cannot contain cycles.") + active.add(id(collection)) + try: + for item in collection: + self.count_node() + self.validate_value(item, depth=depth + 1, active=active) + finally: + active.remove(id(collection)) + + def validate_expression( + self, + expression: FilterExpression, + *, + depth: int, + relative_fields: bool, + active: set[int], + ) -> None: + if depth > self.max_depth: + raise ValueError(f"Filters cannot exceed a depth of {self.max_depth}.") + if id(expression) in active: + raise ValueError("Filter expressions cannot contain cycles.") + active.add(id(expression)) + try: + self.count_node() + if isinstance(expression, Filter): + if ( + not relative_fields + and self.field_names is not None + and expression.field_name.split(".", maxsplit=1)[0] not in self.field_names + ): + raise ValueError( + f"Filter field '{expression.field_name}' is not part of the vector store definition." + ) + if not isinstance(expression.value, Param): + self.validate_value(expression.value, depth=depth + 1, active=active) + _validate_filter_value_shape(expression, allow_params=self.allow_params) + return + if not isinstance(expression, FilterGroup): + raise TypeError("filter must be a Filter or FilterGroup.") + for item in expression.filters: + self.validate_expression( + item, + depth=depth + 1, + relative_fields=relative_fields, + active=active, + ) + finally: + active.remove(id(expression)) + + def count_node(self) -> None: + self.node_count += 1 + if self.node_count > self.max_nodes: + raise ValueError(f"Filters cannot contain more than {self.max_nodes} nodes.") + + +def validate_filter( + filter_: FilterExpression, + *, + field_names: Collection[str] | None = None, + allow_params: bool = False, + max_depth: int = _MAX_FILTER_DEPTH, + max_nodes: int = _MAX_FILTER_NODES, +) -> None: + _FilterValidator( + field_names=field_names, + allow_params=allow_params, + max_depth=max_depth, + max_nodes=max_nodes, + ).validate_expression(filter_, depth=1, relative_fields=False, active=set()) diff --git a/python/packages/core/agent_framework/_vectors.py b/python/packages/core/agent_framework/_vectors.py index 761cc3a5bb2..3ed815e0911 100644 --- a/python/packages/core/agent_framework/_vectors.py +++ b/python/packages/core/agent_framework/_vectors.py @@ -6,10 +6,11 @@ import operator from abc import ABC, abstractmethod -from ast import AST, Lambda, NodeVisitor, expr, parse from collections.abc import AsyncIterable, AsyncIterator, Callable, Mapping, Sequence +from copy import deepcopy from dataclasses import dataclass, is_dataclass, replace -from inspect import Parameter, getsource, signature +from dataclasses import field as dataclass_field +from inspect import Parameter, signature from types import UnionType from typing import ( Annotated, @@ -39,30 +40,42 @@ from ._telemetry import FeatureIndex, mark_feature_used from ._tools import FunctionTool from ._types import Content, EmbeddingGenerationOptions +from ._vector_filters import ( + FilterExpression, + Param, + iter_filter_params, + param_schema, + resolve_filter_params, + snapshot_filter, + validate_filter, + validate_param_value, +) from .exceptions import IntegrationException, IntegrationInvalidResponseException ModelT = TypeVar("ModelT", default=Any) KeyT = TypeVar("KeyT", default=Any) -FilterT = TypeVar("FilterT") ResultT = TypeVar("ResultT") DecoratedModelT = TypeVar("DecoratedModelT") SearchType: TypeAlias = Literal["vector", "keyword_hybrid"] FieldTypes: TypeAlias = Literal["key", "vector", "data"] -IndexKind: TypeAlias = Literal["hnsw", "flat", "ivf_flat", "disk_ann", "quantized_flat", "dynamic", "default"] -DistanceFunction: TypeAlias = Literal[ - "cosine_similarity", - "cosine_distance", - "dot_prod", - "euclidean_distance", - "euclidean_squared_distance", - "manhattan", - "hamming", - "DEFAULT", -] -Vector: TypeAlias = Sequence[float | int] -RecordFilter: TypeAlias = Callable[[Any], bool] | str -RecordFilters: TypeAlias = RecordFilter | Sequence[RecordFilter] +IndexKind: TypeAlias = Literal["hnsw", "flat", "ivf_flat", "disk_ann", "quantized_flat", "dynamic", "default"] | str +DistanceFunction: TypeAlias = ( + Literal[ + "cosine_similarity", + "cosine_distance", + "dot_prod", + "negative_dot_prod", + "euclidean_distance", + "euclidean_squared_distance", + "manhattan", + "hamming", + "DEFAULT", + ] + | str +) +Vector: TypeAlias = Sequence[float | int] | bytes | bytearray +GenerateVectors: TypeAlias = bool | list[str] | tuple[str, ...] EmbeddingClient: TypeAlias = SupportsGetEmbeddings[Any, Any, Any] VectorModelEncoder: TypeAlias = Callable[[Any], Mapping[str, Any]] VectorModelDecoder: TypeAlias = Callable[[Mapping[str, Any]], Any] @@ -71,31 +84,11 @@ _DEFAULT_SEARCH_TOOL_DESCRIPTION: Final[str] = ( "Perform a vector search for data in a vector store using the provided query." ) -_INDEX_KINDS: Final[tuple[str, ...]] = ( - "hnsw", - "flat", - "ivf_flat", - "disk_ann", - "quantized_flat", - "dynamic", - "default", -) -_DISTANCE_FUNCTIONS: Final[tuple[str, ...]] = ( - "cosine_similarity", - "cosine_distance", - "dot_prod", - "euclidean_distance", - "euclidean_squared_distance", - "manhattan", - "hamming", - "DEFAULT", -) - - DISTANCE_FUNCTION_DIRECTION_HELPER: Final[Mapping[DistanceFunction, Callable[[float | int, float | int], bool]]] = { "cosine_similarity": operator.ge, "cosine_distance": operator.le, "dot_prod": operator.ge, + "negative_dot_prod": operator.le, "euclidean_distance": operator.le, "euclidean_squared_distance": operator.le, "manhattan": operator.le, @@ -103,9 +96,18 @@ } +def _copy_provider_annotations(value: Mapping[str, Any] | None) -> dict[str, Any]: + annotations = dict(value or {}) + if any(not isinstance(key, str) for key in annotations): + raise TypeError("Provider annotation keys must be strings.") + return deepcopy(annotations) + + def _msgspec_enc_hook(value: Any) -> Any: if isinstance(value, BaseModel): return value.model_dump() + if isinstance(value, Mapping): + return dict(cast(Mapping[Any, Any], value)) to_list = getattr(value, "tolist", None) if callable(to_list): return to_list() @@ -115,20 +117,48 @@ def _msgspec_enc_hook(value: Any) -> Any: def _normalize_vector(value: Any) -> Vector: - if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): + if isinstance(value, bytes): + return value + if isinstance(value, bytearray): + return bytes(value) + if isinstance(value, Sequence) and not isinstance(value, str): return cast(Vector, value) to_list = getattr(value, "tolist", None) if callable(to_list): converted = to_list() - if isinstance(converted, Sequence) and not isinstance(converted, (str, bytes, bytearray)): + if isinstance(converted, bytearray): + return bytes(converted) + if isinstance(converted, Sequence) and not isinstance(converted, str): return cast(Vector, converted) raise TypeError("The embedding client returned an unsupported vector type.") +def _validate_vector_dimensions( + vector: Any, + field: VectorStoreField, + *, + record_index: int | None = None, +) -> None: + """Check dense sequence length without inspecting elements or interpreting native encodings.""" + # Normalized vectors do not need the more expensive generic sequence check. + if type(vector) not in (list, tuple) and not _is_non_string_sequence(vector): + return + actual_dimensions = len(vector) + if actual_dimensions != field.dimensions: + context = "Query" if record_index is None else f"Record at index {record_index}," + raise ValueError( + f"{context} vector field '{field.name}' expects {field.dimensions} dimensions; got {actual_dimensions}." + ) + + @experimental(feature_id=ExperimentalFeature.VECTOR_STORES) @dataclass(frozen=True, slots=True, init=False) class VectorStoreField: - """Describe one field in a vector store model.""" + """Describe one field in a vector store model. + + Vector ``dimensions`` is the expected length of materialized dense sequences. + Binary bytes and non-sequence provider-native representations remain connector-validated. + """ field_type: FieldTypes name: str @@ -140,6 +170,8 @@ class VectorStoreField: index_kind: IndexKind | None distance_function: DistanceFunction | None embedding_generator: EmbeddingClient | None + is_auto_generated: bool + provider_annotations: dict[str, Any] = dataclass_field(hash=False) @overload def __init__( @@ -149,6 +181,8 @@ def __init__( name: str | None = None, type_: str | None = None, storage_name: str | None = None, + is_auto_generated: bool = False, + provider_annotations: Mapping[str, Any] | None = None, ) -> None: """Initialize a key field. @@ -157,6 +191,8 @@ def __init__( name: The model field name. The decorator supplies this when omitted. type_: The scalar type name used by the backing store. storage_name: The field name used by the backing store. + is_auto_generated: Whether the backing store generates missing key values. + provider_annotations: Mutable provider-specific configuration, copied when the field is created. """ ... @@ -170,6 +206,7 @@ def __init__( storage_name: str | None = None, is_indexed: bool | None = None, is_full_text_indexed: bool | None = None, + provider_annotations: Mapping[str, Any] | None = None, ) -> None: """Initialize a data field with optional indexing. @@ -180,6 +217,7 @@ def __init__( storage_name: The field name used by the backing store. is_indexed: Whether the field should be indexed. is_full_text_indexed: Whether the field should have a full-text index. + provider_annotations: Mutable provider-specific configuration, copied when the field is created. """ ... @@ -195,6 +233,7 @@ def __init__( index_kind: IndexKind | None = None, distance_function: DistanceFunction | None = None, embedding_generator: EmbeddingClient | None = None, + provider_annotations: Mapping[str, Any] | None = None, ) -> None: """Initialize a vector field with required dimensions. @@ -207,6 +246,7 @@ def __init__( index_kind: The vector index kind. distance_function: The vector distance function. embedding_generator: An optional client used to generate this field's embeddings. + provider_annotations: Mutable provider-specific configuration, copied when the field is created. Raises: ValueError: If dimensions or vector options are invalid. @@ -226,6 +266,8 @@ def __init__( index_kind: IndexKind | None = None, distance_function: DistanceFunction | None = None, embedding_generator: EmbeddingClient | None = None, + is_auto_generated: bool = False, + provider_annotations: Mapping[str, Any] | None = None, ) -> None: """Initialize a vector store field. @@ -240,12 +282,17 @@ def __init__( index_kind: The vector index kind. distance_function: The vector distance function. embedding_generator: An optional client used to generate this field's embeddings. + is_auto_generated: Whether a key field is generated by the backing store when missing. + provider_annotations: Mutable provider-specific configuration, copied when the field is created. Raises: + TypeError: If ``is_auto_generated`` is not a boolean. ValueError: If field options are invalid. """ if field_type not in ("key", "vector", "data"): raise ValueError(f"Unknown vector store field type '{field_type}'.") + if not isinstance(is_auto_generated, bool): + raise TypeError("Vector is_auto_generated must be a boolean.") resolved_dimensions: int | None = None resolved_index_kind: IndexKind | None = None resolved_distance_function: DistanceFunction | None = None @@ -253,16 +300,18 @@ def __init__( if field_type == "vector": if dimensions is None or dimensions <= 0: raise ValueError("Vector fields must specify a positive number of dimensions.") - if index_kind is not None and index_kind not in _INDEX_KINDS: - raise ValueError(f"Unknown vector index kind '{index_kind}'.") - if distance_function is not None and distance_function not in _DISTANCE_FUNCTIONS: - raise ValueError(f"Unknown vector distance function '{distance_function}'.") + if index_kind is not None and not isinstance(index_kind, str): + raise TypeError("Vector index_kind must be a string.") + if distance_function is not None and not isinstance(distance_function, str): + raise TypeError("Vector distance_function must be a string.") resolved_dimensions = dimensions resolved_index_kind = index_kind or "default" resolved_distance_function = distance_function or "DEFAULT" resolved_embedding_generator = embedding_generator elif any(value is not None for value in (dimensions, index_kind, distance_function, embedding_generator)): raise ValueError("Vector-only options can only be set on vector fields.") + if field_type != "key" and is_auto_generated: + raise ValueError("Only key fields can be auto-generated.") object.__setattr__(self, "field_type", field_type) object.__setattr__(self, "name", name or "") @@ -274,6 +323,8 @@ def __init__( object.__setattr__(self, "index_kind", resolved_index_kind) object.__setattr__(self, "distance_function", resolved_distance_function) object.__setattr__(self, "embedding_generator", resolved_embedding_generator) + object.__setattr__(self, "is_auto_generated", is_auto_generated) + object.__setattr__(self, "provider_annotations", _copy_provider_annotations(provider_annotations)) @experimental(feature_id=ExperimentalFeature.VECTOR_STORES) @@ -325,6 +376,12 @@ def _validate(self) -> str: storage_names = [field.storage_name or field.name for field in self.fields] if len(storage_names) != len(set(storage_names)): raise ValueError("Vector store field storage names must be unique.") + name_set = set(names) + if any( + field.storage_name is not None and field.storage_name != field.name and field.storage_name in name_set + for field in self.fields + ): + raise ValueError("A vector store field storage name cannot match another field's model name.") key_fields = [field for field in self.fields if field.field_type == "key"] if len(key_fields) != 1: @@ -371,14 +428,19 @@ def data_field_names(self) -> list[str]: """Get the data field names.""" return [field.name for field in self.data_fields] + def try_get_field(self, field_name: str) -> VectorStoreField | None: + """Get a field by model or storage name.""" + model_field = next((field for field in self.fields if field.name == field_name), None) + if model_field is not None: + return model_field + return next((field for field in self.fields if field.storage_name == field_name), None) + def try_get_vector_field(self, field_name: str | None = None) -> VectorStoreField | None: """Get a vector field by model or storage name, defaulting to the first vector field.""" if field_name is None: return self.vector_fields[0] if self.vector_fields else None - return next( - (field for field in self.vector_fields if field.name == field_name or field.storage_name == field_name), - None, - ) + field = self.try_get_field(field_name) + return field if field is not None and field.field_type == "vector" else None def get_names(self, *, include_vector_fields: bool = True, include_key_field: bool = True) -> list[str]: """Get selected model field names.""" @@ -416,7 +478,12 @@ def _default_vector_model_encoder(record_type: type[Any]) -> VectorModelEncoder: def encode(value: Any) -> Mapping[str, Any]: if not isinstance(value, record_type): raise TypeError(f"Expected {record_type.__name__}, got {type(value).__name__}.") - converted = msgspec.to_builtins(value, str_keys=True, enc_hook=_msgspec_enc_hook) + converted = msgspec.to_builtins( + value, + str_keys=True, + builtin_types=(bytes, bytearray), + enc_hook=_msgspec_enc_hook, + ) if not isinstance(converted, Mapping): raise TypeError(f"Vector model {record_type.__name__!r} must serialize to a mapping.") return cast(Mapping[str, Any], converted) @@ -484,6 +551,14 @@ def register_vectorstoremodel( "Vector fields omitted by include_vectors=False must declare defaults when using the default decoder. " f"Add defaults or supply a custom decoder for: {', '.join(required_vector_fields)}." ) + if is_dataclass(record_type) or issubclass(record_type, msgspec.Struct): + try: + msgspec.inspect.type_info(record_type) + except TypeError as exc: + raise ValueError( + f"Vector model {record_type.__name__!r} is not supported by the default msgspec decoder: {exc}. " + "Supply a custom decoder." + ) from exc resolved_encoder = ( cast(VectorModelEncoder, encoder) if encoder is not None else _default_vector_model_encoder(record_type) ) @@ -533,6 +608,9 @@ def _infer_type_name(annotation: Any, *, vector: bool) -> str | None: if origin is not None and args: candidate = next((arg for arg in args if arg is not Ellipsis), candidate) return getattr(candidate, "__name__", str(candidate)) + binary_candidate = next((candidate for candidate in candidates if candidate in (bytes, bytearray)), None) + if binary_candidate is not None: + return "bytes" candidate = candidates[0] if candidates else annotation origin = get_origin(candidate) return getattr(origin or candidate, "__name__", None) @@ -806,15 +884,21 @@ async def serialize( self, records: ModelT | Sequence[ModelT], *, - generate_vectors: bool = True, + generate_vectors: GenerateVectors = True, context: Mapping[str, Any] | None = None, ) -> Any: """Serialize one or more application records for the backing store. + After optional embedding generation, materialized dense sequence lengths + are checked against their fields' dimensions for the entire batch before + connector conversion. This does not inspect vector elements. Null vectors, + source text, binary payloads, and non-sequence provider-native values are + left to the connector. + Args: records: One application record or a sequence of records. - generate_vectors: Whether to generate vector values, overwriting any supplied values. When ``False``, - supplied values are preserved. + generate_vectors: Whether to generate all vector fields, preserve all supplied values, or generate only + the vector fields named in a sequence. Generated values overwrite supplied values. context: Connector-specific serialization context. Raises: @@ -827,8 +911,13 @@ async def serialize( input_records = list(cast(Sequence[ModelT], records)) if is_batch else [cast(ModelT, records)] dict_records = [self._serialize_record_to_dict(record) for record in input_records] - if generate_vectors: - await self._add_vectors_to_records(dict_records) + vector_fields = self._resolve_vector_fields_to_generate(generate_vectors) + if vector_fields: + await self._add_vectors_to_records(dict_records, vector_fields=vector_fields) + dimension_fields = tuple((field.storage_name or field.name, field) for field in self.definition.vector_fields) + for record_index, record in enumerate(dict_records): + for storage_name, field in dimension_fields: + _validate_vector_dimensions(record.get(storage_name), field, record_index=record_index) store_models = list(self._serialize_dicts_to_store_models(dict_records, context=context)) if len(store_models) != len(dict_records): @@ -852,7 +941,12 @@ def _serialize_record_to_dict(self, record: ModelT) -> dict[str, Any]: @staticmethod def _to_builtin_mapping(record: Any) -> Mapping[str, Any]: - converted = msgspec.to_builtins(record, str_keys=True, enc_hook=_msgspec_enc_hook) + converted = msgspec.to_builtins( + record, + str_keys=True, + builtin_types=(bytes, bytearray), + enc_hook=_msgspec_enc_hook, + ) if not isinstance(converted, Mapping): raise TypeError("Vector records must serialize to mappings.") return cast(Mapping[str, Any], converted) @@ -864,14 +958,47 @@ def _serialize_mapping_to_store(self, source: Mapping[str, Any]) -> dict[str, An value = source[field.name] elif field.storage_name is not None and field.storage_name in source: value = source[field.storage_name] + elif field.field_type == "key" and field.is_auto_generated: + continue else: raise ValueError(f"Record is missing vector store field '{field.name}'.") + if field.field_type == "key" and field.is_auto_generated and value is None: + continue + if field.field_type == "vector" and isinstance(value, bytearray): + value = bytes(value) serialized[field.storage_name or field.name] = value return serialized - async def _add_vectors_to_records(self, records: Sequence[dict[str, Any]]) -> None: + def _resolve_vector_fields_to_generate( + self, + generate_vectors: GenerateVectors, + ) -> tuple[VectorStoreField, ...]: + if isinstance(generate_vectors, bool): + return tuple(self.definition.vector_fields) if generate_vectors else () + if not _is_non_string_sequence(generate_vectors) or any( + not isinstance(field_name, str) for field_name in generate_vectors + ): + raise TypeError("generate_vectors must be a boolean or a sequence of vector field names.") + field_names = list(cast(Sequence[str], generate_vectors)) + if len(field_names) != len(set(field_names)): + raise ValueError("generate_vectors field names must be unique.") + vector_fields = {field.name: field for field in self.definition.vector_fields} + unknown = sorted(set(field_names) - set(vector_fields)) + if unknown: + raise ValueError(f"Unknown vector field(s) in generate_vectors: {', '.join(unknown)}.") + selected = set(field_names) + return tuple(field for field in self.definition.vector_fields if field.name in selected) + + async def _add_vectors_to_records( + self, + records: Sequence[dict[str, Any]], + *, + vector_fields: Sequence[VectorStoreField], + ) -> None: + if not records: + return field_generators: list[tuple[VectorStoreField, EmbeddingClient]] = [] - for field in self.definition.vector_fields: + for field in vector_fields: embedding_generator = field.embedding_generator or self.embedding_generator if embedding_generator is None: raise ValueError( @@ -1029,6 +1156,7 @@ async def _inner_get( self, *, keys: Sequence[KeyT] | None = None, + filter: FilterExpression | None = None, top: int = 10, skip: int = 0, order_by: Mapping[str, bool] | None = None, @@ -1052,15 +1180,26 @@ async def upsert( self, records: Sequence[ModelT], *, - generate_vectors: bool = True, + generate_vectors: GenerateVectors = True, operation_options: Mapping[str, Any] | None = None, ) -> Sequence[KeyT]: """Upsert a batch of records. + Dense sequence lengths are checked after optional embedding generation, + before connector conversion or writes. A dimension mismatch rejects the + whole batch at this boundary. Binary and non-sequence provider-native + representations remain connector-validated. + + A connector may partially persist a batch before reporting an error; the + abstraction does not guarantee rollback or atomicity. Retrying records + with stable application-provided keys should be idempotent when the + backing store supports ordinary upsert semantics. Retrying records whose + keys are generated by the store may create duplicates. + Args: records: A sequence of models. - generate_vectors: Whether to generate vector values, overwriting any supplied values. When ``False``, - supplied values are preserved. + generate_vectors: Whether to generate all vector fields, preserve all supplied values, or generate only + the vector fields named in a sequence. Generated values overwrite supplied values. operation_options: Store-specific operation options. Returns: @@ -1079,7 +1218,7 @@ async def upsert( serialized = await self.serialize(records, generate_vectors=generate_vectors) store_records = list(serialized) if _is_non_string_sequence(serialized) else [serialized] keys = list(await self._inner_upsert(store_records, operation_options=operation_options)) - except (TypeError, ValueError): + except (TypeError, ValueError, NotImplementedError): raise except IntegrationException: raise @@ -1097,6 +1236,7 @@ async def get( self, keys: Sequence[KeyT] | None = None, *, + filter: FilterExpression | None = None, top: int = 10, skip: int = 0, order_by: Mapping[str, bool] | None = None, @@ -1107,6 +1247,7 @@ async def get( Args: keys: A sequence of keys, or ``None`` to list a page of records. + filter: A portable data-only filter used when listing records. top: The maximum number of records returned when listing. skip: The number of records skipped when listing. order_by: Field names mapped to ascending (``True``) or descending (``False``) order. @@ -1117,7 +1258,7 @@ async def get( A sequence of models. Keys that do not exist are omitted. Raises: - ValueError: If paging arguments are invalid. + ValueError: If paging arguments or filters are invalid, or keys and a filter are supplied together. TypeError: If keys or a returned record has an unsupported type. IntegrationException: If retrieval fails. """ @@ -1125,15 +1266,23 @@ async def get( _validate_paging(top=top, skip=skip) if keys is not None and not _is_non_string_sequence(keys): raise TypeError("keys must be a sequence.") + if keys is not None and filter is not None: + raise ValueError("keys and filter are alternate retrieval modes and cannot be combined.") + operation_filter = snapshot_filter(filter) if filter is not None else None + if operation_filter is not None: + validate_filter(operation_filter, field_names=self.definition.names) try: records = await self._inner_get( keys=keys, + filter=operation_filter, top=top, skip=skip, order_by=order_by, include_vectors=include_vectors, operation_options=operation_options, ) + except (TypeError, ValueError, NotImplementedError): + raise except IntegrationException: raise except Exception as exc: @@ -1251,18 +1400,14 @@ async def _inner_ensure_collection_deleted( ... -class _LambdaVisitor(NodeVisitor, Generic[FilterT]): - def __init__(self, lambda_parser: Callable[[expr], FilterT]) -> None: - self.lambda_parser = lambda_parser - self.output_filters: list[FilterT] = [] - - def visit_Lambda(self, node: Lambda) -> None: - self.output_filters.append(self.lambda_parser(node.body)) - - @experimental(feature_id=ExperimentalFeature.VECTOR_STORES) class BaseVectorSearch(_VectorStoreRecordHandler[KeyT, ModelT], ABC): - """Base class for vector and keyword-hybrid search.""" + """Base class for vector and keyword-hybrid search. + + Core validates portable request structure and deserializes connector results. + Connectors own scoring, filter execution, score thresholds, and paging. + Returned scores and results are not re-filtered by core. + """ supported_search_types: ClassVar[set[SearchType]] = {"vector"} @@ -1271,7 +1416,7 @@ async def _inner_search( self, *, search_type: SearchType, - filter: Any | list[Any] | None = None, + filter: FilterExpression | None = None, values: Any | None = None, vector: Vector | None = None, top: int = 3, @@ -1282,7 +1427,14 @@ async def _inner_search( score_threshold: float | None = None, operation_options: Mapping[str, Any] | None = None, ) -> SearchResults[Any]: - """Execute a search and return raw connector results.""" + """Execute a search and return raw connector results. + + Apply filters and score thresholds natively in the backing store where + supported. Otherwise implement an explicit connector-local fallback or + raise ``NotImplementedError``; do not silently ignore these options. + The connector owns score units, comparison direction, default metrics, + and coordinating filtering with paging. Core does not post-filter results. + """ ... @abstractmethod @@ -1295,11 +1447,6 @@ def _get_score_from_result(self, result: Any) -> float | None: """Extract a score from one raw search result.""" ... - @abstractmethod - def _lambda_parser(self, node: AST) -> Any: - """Translate one lambda expression body into a store filter.""" - ... - @overload async def search( self, @@ -1307,7 +1454,7 @@ async def search( *, search_type: SearchType = "vector", vector: Vector | None = None, - filter: RecordFilters | None = None, + filter: FilterExpression | None = None, top: int = 3, skip: int = 0, include_vectors: bool = False, @@ -1318,18 +1465,22 @@ async def search( ) -> SearchResults[SearchResponse[ModelT]]: """Search from a value, optionally with a precomputed vector. + Materialized dense sequence length must match the selected vector field's + dimensions before connector dispatch. Binary bytes and non-sequence + provider-native formats remain connector-validated. + Args: values: The value to search for or vectorize. search_type: Whether to perform vector or keyword-hybrid search. vector: An optional precomputed query vector. - filter: One or more lambda filters. + filter: A portable data-only filter. top: The maximum number of results. skip: The number of results to skip. include_vectors: Whether returned records include vector fields. vector_property_name: The vector field used for search. additional_property_name: The data field used for keyword-hybrid search. - score_threshold: The minimum similarity or maximum distance accepted. - Results without scores remain included. + score_threshold: An optional cutoff interpreted and enforced by the connector. + Score units, comparison direction, and default metrics are connector-specific. operation_options: Store-specific operation options. Returns: @@ -1348,7 +1499,7 @@ async def search( *, search_type: Literal["vector"] = "vector", vector: Vector, - filter: RecordFilters | None = None, + filter: FilterExpression | None = None, top: int = 3, skip: int = 0, include_vectors: bool = False, @@ -1359,17 +1510,21 @@ async def search( ) -> SearchResults[SearchResponse[ModelT]]: """Search from a required precomputed vector. + Dense sequence length must match the selected vector field's dimensions + before connector dispatch. Binary bytes and non-sequence provider-native + formats remain connector-validated. + Args: search_type: The vector search type. vector: The precomputed query vector. - filter: One or more lambda filters. + filter: A portable data-only filter. top: The maximum number of results. skip: The number of results to skip. include_vectors: Whether returned records include vector fields. vector_property_name: The vector field used for search. additional_property_name: The data field used for keyword-hybrid search. - score_threshold: The minimum similarity or maximum distance accepted. - Results without scores remain included. + score_threshold: An optional cutoff interpreted and enforced by the connector. + Score units, comparison direction, and default metrics are connector-specific. operation_options: Store-specific operation options. Returns: @@ -1388,7 +1543,7 @@ async def search( *, search_type: SearchType = "vector", vector: Vector | None = None, - filter: RecordFilters | None = None, + filter: FilterExpression | None = None, top: int = 3, skip: int = 0, include_vectors: bool = False, @@ -1399,18 +1554,29 @@ async def search( ) -> SearchResults[SearchResponse[ModelT]]: """Search the vector store. + Supplied or locally generated dense sequence length is checked against the + selected vector field's dimensions before connector dispatch, even for + empty collections. This does not inspect elements or convert native + payloads. Binary and non-sequence provider-native formats remain + connector-validated; provider-side vectorization still receives ``values`` + and no vector. + + Filters and score thresholds are passed to the connector for execution, + normally in the backing store. Core deserializes returned records without + applying another threshold comparison or changing returned scores. + Args: values: The value to search for or vectorize. search_type: Whether to perform vector or keyword-hybrid search. vector: A precomputed query vector. - filter: One or more lambda filters. + filter: A portable data-only filter. top: The maximum number of results. skip: The number of results to skip. include_vectors: Whether returned records include vector fields. vector_property_name: The vector field used for search. additional_property_name: The data field used for keyword-hybrid search. - score_threshold: The minimum similarity or maximum distance accepted. - Results without scores remain included. + score_threshold: An optional cutoff interpreted and enforced by the connector. + Score units, comparison direction, and default metrics are connector-specific. operation_options: Store-specific operation options. Returns: @@ -1432,21 +1598,27 @@ async def search( raise ValueError("Keyword-hybrid search requires values.") _validate_paging(top=top, skip=skip) + operation_filter = snapshot_filter(filter) if filter is not None else None + if operation_filter is not None: + validate_filter(operation_filter, field_names=self.definition.names) try: - self._validate_score_threshold( - score_threshold=score_threshold, - vector_property_name=vector_property_name, - ) resolved_vector = vector if resolved_vector is None and values is not None: resolved_vector = await self._generate_vector_from_values( values, vector_property_name=vector_property_name, ) - translated_filter = self._build_filter(filter) + if resolved_vector is not None: + vector_field = self.definition.try_get_vector_field(vector_property_name) + if vector_field is None and vector_property_name is not None: + raise ValueError( + f"Vector field '{vector_property_name}' was not found in the collection definition." + ) + if vector_field is not None: + _validate_vector_dimensions(resolved_vector, vector_field) raw_results = await self._inner_search( search_type=search_type, - filter=translated_filter, + filter=operation_filter, values=values, vector=resolved_vector, top=top, @@ -1461,32 +1633,16 @@ async def search( self._get_search_results_from_results( raw_results.results, include_vectors=include_vectors, - vector_property_name=vector_property_name, - score_threshold=score_threshold, ), metadata=raw_results.metadata, ) - except (TypeError, ValueError): + except (TypeError, ValueError, NotImplementedError): raise except IntegrationException: raise except Exception as exc: raise IntegrationException(f"Vector search failed: {exc}") from exc - def _validate_score_threshold( - self, - *, - score_threshold: float | None, - vector_property_name: str | None, - ) -> None: - if score_threshold is None: - return - vector_field = self.definition.try_get_vector_field(vector_property_name) - if vector_field is None: - raise ValueError("A score threshold requires a vector field.") - if vector_field.distance_function == "DEFAULT": - raise ValueError("A score threshold requires an explicit distance function on the vector field.") - async def _generate_vector_from_values( self, values: Any, @@ -1512,37 +1668,11 @@ async def _generate_vector_from_values( generated_vector = embeddings[0].vector return _normalize_vector(generated_vector) - def _build_filter(self, search_filter: RecordFilters | None) -> Any | list[Any] | None: - """Translate lambda filters with the connector's AST parser.""" - if not search_filter: - return None - filters: list[RecordFilter] - if _is_non_string_sequence(search_filter) and not callable(search_filter): - filters = cast(list[RecordFilter], list(search_filter)) - else: - filters = [cast(RecordFilter, search_filter)] - visitor = _LambdaVisitor(self._lambda_parser) - try: - for filter_item in filters: - source = ( - filter_item - if isinstance(filter_item, str) - else getsource(cast(Callable[..., Any], filter_item)).strip() - ) - visitor.visit(parse(source)) - except (OSError, SyntaxError, TypeError) as exc: - raise ValueError(f"Unable to parse vector search filter: {exc}") from exc - if not visitor.output_filters: - raise ValueError("No lambda expression was found in the vector search filter.") - return visitor.output_filters[0] if len(visitor.output_filters) == 1 else visitor.output_filters - def _get_search_results_from_results( self, results: AsyncIterable[Any] | Sequence[Any], *, include_vectors: bool, - vector_property_name: str | None, - score_threshold: float | None, ) -> AsyncIterable[SearchResponse[ModelT]]: """Convert raw connector results into deserialized search responses.""" @@ -1561,12 +1691,6 @@ async def generate() -> AsyncIterator[SearchResponse[ModelT]]: "A search result must deserialize to exactly one record." ) score = self._get_score_from_result(result) - if not self._meets_score_threshold( - score, - score_threshold=score_threshold, - vector_property_name=vector_property_name, - ): - continue yield SearchResponse(record=cast(ModelT, record), score=score) except IntegrationException: raise @@ -1581,26 +1705,6 @@ async def generate() -> AsyncIterator[SearchResponse[ModelT]]: return generate() - def _meets_score_threshold( - self, - score: float | None, - *, - score_threshold: float | None, - vector_property_name: str | None, - ) -> bool: - """Apply a threshold when a result includes a comparable score. - - Results without scores remain included because the threshold cannot be - evaluated for them. - """ - if score_threshold is None or score is None: - return True - vector_field = self.definition.try_get_vector_field(vector_property_name) - if vector_field is None or vector_field.distance_function is None: - return True - comparison = DISTANCE_FUNCTION_DIRECTION_HELPER.get(vector_field.distance_function) - return comparison(score, score_threshold) if comparison is not None else True - @runtime_checkable @experimental(feature_id=ExperimentalFeature.VECTOR_STORES) @@ -1615,23 +1719,24 @@ async def upsert( self, records: Sequence[ModelT], *, - generate_vectors: bool = True, + generate_vectors: GenerateVectors = True, operation_options: Mapping[str, Any] | None = None, ) -> Sequence[KeyT]: - """Upsert a batch of records, generating embeddings by default.""" + """Upsert a batch, which may partially succeed, generating embeddings by default.""" ... async def get( self, keys: Sequence[KeyT] | None = None, *, + filter: FilterExpression | None = None, top: int = 10, skip: int = 0, order_by: Mapping[str, bool] | None = None, include_vectors: bool = False, operation_options: Mapping[str, Any] | None = None, ) -> Sequence[ModelT]: - """Get records by keys or list a page of records, excluding vectors by default.""" + """Get records by keys or filter, or list a page of records, excluding vectors by default.""" ... async def delete( @@ -1647,7 +1752,12 @@ async def delete( @runtime_checkable @experimental(feature_id=ExperimentalFeature.VECTOR_STORES) class SupportsVectorSearch(Protocol[ModelT]): - """Protocol for vector and keyword-hybrid search.""" + """Protocol for vector and keyword-hybrid search. + + Implementations own scoring, filter execution, score thresholds, and paging. + Execute these in the backing store where supported, otherwise use an explicit + local fallback or reject unsupported options. + """ @overload async def search( @@ -1656,7 +1766,7 @@ async def search( *, search_type: SearchType = "vector", vector: Vector | None = None, - filter: RecordFilters | None = None, + filter: FilterExpression | None = None, top: int = 3, skip: int = 0, include_vectors: bool = False, @@ -1671,14 +1781,14 @@ async def search( values: The value to search for or vectorize. search_type: Whether to perform vector or keyword-hybrid search. vector: An optional precomputed query vector. - filter: One or more lambda filters. + filter: A portable data-only filter. top: The maximum number of results. skip: The number of results to skip. include_vectors: Whether returned records include vector fields. vector_property_name: The vector field used for search. additional_property_name: The data field used for keyword-hybrid search. - score_threshold: The minimum similarity or maximum distance accepted. - Results without scores remain included. + score_threshold: An optional cutoff interpreted and enforced by the connector. + Score units, comparison direction, and default metrics are connector-specific. operation_options: Store-specific operation options. Returns: @@ -1697,7 +1807,7 @@ async def search( *, search_type: Literal["vector"] = "vector", vector: Vector, - filter: RecordFilters | None = None, + filter: FilterExpression | None = None, top: int = 3, skip: int = 0, include_vectors: bool = False, @@ -1711,14 +1821,14 @@ async def search( Args: search_type: The vector search type. vector: The precomputed query vector. - filter: One or more lambda filters. + filter: A portable data-only filter. top: The maximum number of results. skip: The number of results to skip. include_vectors: Whether returned records include vector fields. vector_property_name: The vector field used for search. additional_property_name: The data field used for keyword-hybrid search. - score_threshold: The minimum similarity or maximum distance accepted. - Results without scores remain included. + score_threshold: An optional cutoff interpreted and enforced by the connector. + Score units, comparison direction, and default metrics are connector-specific. operation_options: Store-specific operation options. Returns: @@ -1740,11 +1850,9 @@ def create_vector_search_tool( description: str = _DEFAULT_SEARCH_TOOL_DESCRIPTION, approval_mode: Literal["always_require", "never_require"] = "never_require", search_type: SearchType = "vector", - parameters: type[BaseModel] | Mapping[str, Any] | None = None, - top: int = 5, - skip: int = 0, - filter: RecordFilters | None = None, - filter_mapper: Callable[[RecordFilters | None, Mapping[str, Any]], RecordFilters | None] | None = None, + top: int | Param = 5, + skip: int | Param = 0, + filter: FilterExpression | None = None, result_mapper: Callable[[SearchResponse[ModelT]], str | Content | Sequence[Content]] | None = None, ) -> FunctionTool: """Create an agent-usable tool backed by vector search. @@ -1755,51 +1863,82 @@ def create_vector_search_tool( description: The tool description shown to the model. approval_mode: Whether the tool requires approval before invocation. search_type: Whether the tool performs vector or keyword-hybrid search. - parameters: A Pydantic model or JSON schema declaring the tool parameters. - It must declare ``query`` as a required string. A custom schema can - expose ``top`` and ``skip`` as integers with finite ``maximum`` values; - additional fields are passed to ``filter_mapper``. - top: The default result limit and the maximum when ``parameters`` does not expose ``top``. - skip: The default offset and the maximum when ``parameters`` does not expose ``skip``. - filter: A fixed filter applied to each tool invocation. - filter_mapper: Maps additional declared tool arguments to search filters. - The default creates equality filters for each additional argument. + top: A fixed result limit or a bounded model-set parameter. + skip: A fixed result offset or a bounded model-set parameter. + filter: A fixed filter that may contain model-set ``Param`` values. + A nullable ``Param`` with ``default=None`` and ``omit_if_none=True`` removes + its leaf for an absent or null argument. Remaining group children still apply; + empty groups are removed recursively. See ``FilterGroup`` for details. result_mapper: Maps each search response to text or one or more multimodal content items. Returns: - A function tool with only a ``query`` parameter by default. Custom parameters can expose - ``top``, ``skip``, and fields mapped into filters by ``filter_mapper``. + A function tool with ``query`` and any parameters discovered in ``filter``, ``top``, or ``skip``. Raises: - ValueError: If parameters or paging limits are invalid. + TypeError: If a parameter annotation or default is invalid. + ValueError: If filters, parameters, or paging limits are invalid. NotImplementedError: If the search type is unsupported. """ - _validate_paging(top=top, skip=skip) - map_filter = filter_mapper or _default_search_filter_mapper + if isinstance(top, bool) or isinstance(skip, bool): + raise TypeError("top and skip must be integers or Param instances.") + if not isinstance(top, int | Param) or not isinstance(skip, int | Param): + raise TypeError("top and skip must be integers or Param instances.") + if isinstance(top, int): + _validate_paging(top=top, skip=0) + if isinstance(skip, int): + _validate_paging(top=1, skip=skip) + map_result = result_mapper or _default_search_result_mapper - input_model = parameters if parameters is not None else _default_search_tool_parameters() - max_top, max_skip = _validate_search_tool_parameters( - input_model, - default_top=top, - default_skip=skip, + configured_filter = snapshot_filter(filter) if filter is not None else None + input_schema, param_definitions = _create_search_tool_input_schema( + filter=configured_filter, + top=top, + skip=skip, ) + _validate_search_tool_paging_param("top", top, input_schema) + _validate_search_tool_paging_param("skip", skip, input_schema) + definition = getattr(search, "definition", None) + if configured_filter is not None and isinstance(definition, VectorStoreCollectionDefinition): + validate_filter(configured_filter, field_names=definition.names, allow_params=True) async def search_tool(**arguments: Any) -> list[Content]: - query = arguments.pop("query") + unexpected = sorted(set(arguments) - set(cast(Mapping[str, Any], input_schema["properties"]))) + if unexpected: + raise TypeError(f"Unexpected argument(s) for '{name}': {', '.join(unexpected)}") + missing = sorted(set(cast(Sequence[str], input_schema["required"])) - set(arguments)) + if missing: + raise TypeError(f"Missing required argument(s) for '{name}': {', '.join(missing)}") + query = arguments.get("query") if not isinstance(query, str): raise TypeError("The search tool 'query' argument must be a string.") - invocation_top = arguments.pop("top", top) - invocation_skip = arguments.pop("skip", skip) + validated_arguments = { + parameter_name: validate_param_value(param, arguments[parameter_name]) + for parameter_name, param in param_definitions.items() + if parameter_name in arguments + } + resolved_arguments = { + **{ + parameter_name: param.default + for parameter_name, param in param_definitions.items() + if param.has_default + }, + **validated_arguments, + } + invocation_top = _resolve_search_tool_option("top", top, resolved_arguments) + invocation_skip = _resolve_search_tool_option("skip", skip, resolved_arguments) _validate_paging(top=invocation_top, skip=invocation_skip) - if invocation_top > max_top: - raise ValueError(f"top must not exceed the configured maximum of {max_top}.") - if invocation_skip > max_skip: - raise ValueError(f"skip must not exceed the configured maximum of {max_skip}.") - dynamic_filter = map_filter(filter, arguments) + resolved_filter = ( + resolve_filter_params(configured_filter, resolved_arguments) if configured_filter is not None else None + ) + if resolved_filter is not None: + validate_filter( + resolved_filter, + field_names=definition.names if isinstance(definition, VectorStoreCollectionDefinition) else None, + ) results = await search.search( query, search_type=search_type, - filter=dynamic_filter, + filter=resolved_filter, top=invocation_top, skip=invocation_skip, ) @@ -1823,7 +1962,7 @@ async def search_tool(**arguments: Any) -> list[Content]: description=description, approval_mode=approval_mode, func=search_tool, - input_model=input_model, + input_model=input_schema, ) @@ -1842,78 +1981,81 @@ async def _as_async_iterable( yield value -def _default_search_tool_parameters() -> dict[str, Any]: - return { - "type": "object", - "properties": { - "query": { - "type": "string", - "description": "The query to search for.", - }, +def _create_search_tool_input_schema( + *, + filter: FilterExpression | None, + top: int | Param, + skip: int | Param, +) -> tuple[dict[str, Any], dict[str, Param]]: + params = [*iter_filter_params(filter)] if filter is not None else [] + params.extend(option for option in (top, skip) if isinstance(option, Param)) + definitions: dict[str, Param] = {} + + for param in params: + existing = definitions.get(param.name) + if existing is not None: + if existing != param: + raise ValueError(f"Search parameter '{param.name}' has conflicting declarations.") + continue + definitions[param.name] = param + if param.has_default: + validate_param_value(param, param.default) + + properties = { + "query": { + "type": "string", + "description": "The query to search for.", }, - "required": ["query"], - "additionalProperties": False, + **{name: param_schema(param) for name, param in definitions.items()}, } + required = ["query", *(name for name, param in definitions.items() if param.required)] + return ( + { + "type": "object", + "properties": properties, + "required": required, + "additionalProperties": False, + }, + definitions, + ) -def _validate_search_tool_parameters( - parameters: type[BaseModel] | Mapping[str, Any], - *, - default_top: int, - default_skip: int, -) -> tuple[int, int]: - schema: Mapping[str, Any] = parameters.model_json_schema() if isinstance(parameters, type) else parameters - raw_properties = schema.get("properties") - if not isinstance(raw_properties, Mapping): - raise ValueError("Search tool parameters must define object properties.") - properties = cast(Mapping[str, Any], raw_properties) - query_schema = properties.get("query") - required = schema.get("required") - query_type = cast(Mapping[str, Any], query_schema).get("type") if isinstance(query_schema, Mapping) else None - if ( - not isinstance(query_schema, Mapping) - or query_type != "string" - or not _is_non_string_sequence(required) - or "query" not in required - ): - raise ValueError("Search tool parameters must define 'query' as a required string.") - - limits = {"top": default_top, "skip": default_skip} - for name, minimum in (("top", 1), ("skip", 0)): - parameter_schema = properties.get(name) - if parameter_schema is None: - continue - if not isinstance(parameter_schema, Mapping): - raise ValueError(f"Search tool parameter '{name}' must be an integer.") - typed_parameter_schema = cast(Mapping[str, Any], parameter_schema) - if typed_parameter_schema.get("type") != "integer": - raise ValueError(f"Search tool parameter '{name}' must be an integer.") - maximum = typed_parameter_schema.get("maximum") - if not isinstance(maximum, int) or isinstance(maximum, bool) or maximum < minimum: - raise ValueError(f"Search tool parameter '{name}' must declare an integer maximum of at least {minimum}.") - configured_default = default_top if name == "top" else default_skip - if configured_default > maximum: - raise ValueError(f"Configured {name}={configured_default} exceeds the parameter maximum of {maximum}.") - limits[name] = maximum - return limits["top"], limits["skip"] - - -def _default_search_filter_mapper( - search_filter: RecordFilters | None, +def _validate_search_tool_paging_param( + option_name: Literal["top", "skip"], + option: int | Param, + input_schema: Mapping[str, Any], +) -> None: + if isinstance(option, int): + return + if option.omit_if_none: + raise ValueError(f"The {option_name} Param does not support omit_if_none.") + if not option.required and not option.has_default: + raise ValueError(f"A model-set {option_name} Param must be required or declare a default.") + parameter_schema = cast(Mapping[str, Any], cast(Mapping[str, Any], input_schema["properties"])[option.name]) + minimum = parameter_schema.get("minimum") + maximum = parameter_schema.get("maximum") + required_minimum = 1 if option_name == "top" else 0 + if parameter_schema.get("type") != "integer": + raise ValueError(f"The {option_name} Param must use an integer annotation.") + if not isinstance(minimum, int | float) or isinstance(minimum, bool) or minimum < required_minimum: + raise ValueError(f"The {option_name} Param must declare a minimum of at least {required_minimum}.") + if not isinstance(maximum, int) or isinstance(maximum, bool) or maximum < required_minimum: + raise ValueError(f"The {option_name} Param must declare a finite integer maximum.") + + +def _resolve_search_tool_option( + option_name: Literal["top", "skip"], + option: int | Param, arguments: Mapping[str, Any], -) -> RecordFilters | None: - dynamic_filters: list[RecordFilter] = [] - for name, value in arguments.items(): - if not name.isidentifier(): - raise ValueError(f"Search tool parameter '{name}' cannot be mapped to a model field.") - dynamic_filters.append(f"lambda record: record.{name} == {value!r}") - if not dynamic_filters: - return search_filter - if search_filter is None: - return dynamic_filters - if _is_non_string_sequence(search_filter) and not callable(search_filter): - return [*cast(Sequence[RecordFilter], search_filter), *dynamic_filters] - return [cast(RecordFilter, search_filter), *dynamic_filters] +) -> int: + if isinstance(option, int): + return option + if option.name not in arguments: + raise TypeError(f"Missing search parameter '{option.name}' for {option_name}.") + value = arguments[option.name] + if not isinstance(value, int) or isinstance(value, bool): + raise TypeError(f"Search parameter '{option.name}' for {option_name} must be an integer.") + return value def _default_search_result_mapper(response: SearchResponse[Any]) -> str: diff --git a/python/packages/core/tests/core/test_in_memory.py b/python/packages/core/tests/core/test_in_memory.py new file mode 100644 index 00000000000..ca72fdcef61 --- /dev/null +++ b/python/packages/core/tests/core/test_in_memory.py @@ -0,0 +1,811 @@ +# Copyright (c) Microsoft. All rights reserved. + +from __future__ import annotations + +import ast +import inspect +import math +import warnings +from collections.abc import Iterator, Sequence +from dataclasses import dataclass +from typing import Annotated, Any, cast + +import pytest + +from agent_framework import ( + BaseEmbeddingClient, + DistanceFunction, + Embedding, + EmbeddingGenerationOptions, + Filter, + FilterGroup, + GeneratedEmbeddings, + InMemoryCollection, + InMemoryStore, + Param, + VectorStoreCollectionDefinition, + VectorStoreField, + create_vector_search_tool, + register_vectorstoremodel, + vectorstoremodel, +) +from agent_framework import _in_memory as in_memory_module +from agent_framework._feature_stage import ExperimentalWarning +from agent_framework.exceptions import IntegrationException + +pytestmark = pytest.mark.filterwarnings("ignore::agent_framework._feature_stage.ExperimentalWarning") + +with warnings.catch_warnings(): + warnings.simplefilter("ignore", ExperimentalWarning) + + @vectorstoremodel(collection_name="documents") + @dataclass + class Document: + id: Annotated[str, VectorStoreField("key")] + text: Annotated[str, VectorStoreField("data")] + category: Annotated[str, VectorStoreField("data")] + rating: Annotated[int, VectorStoreField("data")] + tags: Annotated[list[str], VectorStoreField("data")] + optional: Annotated[str | None, VectorStoreField("data")] + vector: Annotated[ + list[float] | None, + VectorStoreField("vector", dimensions=2, distance_function="cosine_similarity"), + ] = None + + +DOCUMENTS = ( + Document("one", "Luxury hotel", "travel", 5, ["wifi", "pool"], "featured", [1.0, 0.0]), + Document("two", "Budget hostel", "work", 3, ["wifi"], None, [0.0, 1.0]), +) + + +class QueryEmbeddingClient(BaseEmbeddingClient[str, list[float], EmbeddingGenerationOptions]): + async def get_embeddings( + self, + values: Sequence[str], + *, + options: EmbeddingGenerationOptions | None = None, + ) -> GeneratedEmbeddings[list[float], EmbeddingGenerationOptions]: + return GeneratedEmbeddings([Embedding(vector=[1.0, 0.0]) for _ in values], options=options) + + +async def _create_collection() -> InMemoryCollection[str, Document]: + collection = InMemoryCollection(Document) + await collection.ensure_collection_exists() + await collection.upsert(DOCUMENTS, generate_vectors=False) + return collection + + +async def _result_ids( + collection: InMemoryCollection[str, Document], + filter_: Filter | FilterGroup, +) -> list[str]: + results = await collection.search(vector=[1.0, 0.0], filter=filter_, top=10) + return [result["record"].id async for result in results] + + +async def test_in_memory_store_shares_collection_state_and_lifecycle() -> None: + store = InMemoryStore() + first = store.get_collection(Document) + second = store.get_collection(Document) + + assert not await first.collection_exists() + assert await store.list_collection_names() == [] + + await first.ensure_collection_exists() + await first.upsert([DOCUMENTS[0]], generate_vectors=False) + + assert await store.list_collection_names() == ["documents"] + assert await second.get(["one"], include_vectors=True) == [DOCUMENTS[0]] + + await store.ensure_collection_deleted("documents") + assert not await first.collection_exists() + with pytest.raises(IntegrationException, match="does not exist"): + await second.get(["one"]) + + +def test_in_memory_store_rejects_conflicting_collection_definitions() -> None: + store = InMemoryStore() + store.get_collection(Document) + other_definition = VectorStoreCollectionDefinition( + [ + VectorStoreField("key", name="other_id"), + VectorStoreField("vector", name="vector", dimensions=2), + ], + collection_name="documents", + ) + + with pytest.raises(ValueError, match="another definition"): + store.get_collection(dict, definition=other_definition) + + +async def test_in_memory_collection_requires_creation() -> None: + collection: InMemoryCollection[str, Document] = InMemoryCollection(Document) + + with pytest.raises(IntegrationException, match="does not exist"): + await collection.get() + + +async def test_in_memory_crud_listing_ordering_and_defensive_copies() -> None: + collection = await _create_collection() + + assert await collection.get(["one"]) == [ + Document("one", "Luxury hotel", "travel", 5, ["wifi", "pool"], "featured"), + ] + assert await collection.get(order_by={"rating": False}, top=1) == [ + Document("one", "Luxury hotel", "travel", 5, ["wifi", "pool"], "featured"), + ] + assert await collection.get(order_by={"rating": True}, skip=1, top=1) == [ + Document("one", "Luxury hotel", "travel", 5, ["wifi", "pool"], "featured"), + ] + assert await collection.get(order_by={"optional": True}) == [ + Document("one", "Luxury hotel", "travel", 5, ["wifi", "pool"], "featured"), + Document("two", "Budget hostel", "work", 3, ["wifi"], None), + ] + with pytest.raises(ValueError, match="not part of the vector store definition"): + await collection.get(order_by={"missing": True}) + with pytest.raises(TypeError, match="must be a boolean"): + await collection.get(order_by=cast(Any, {"rating": "descending"})) + + fetched = (await collection.get(["one"], include_vectors=True))[0] + cast(list[float], fetched.vector)[0] = 99.0 + assert (await collection.get(["one"], include_vectors=True))[0].vector == [1.0, 0.0] + + await collection.delete(["one", "missing"]) + assert await collection.get(["one"]) == [] + + +async def test_in_memory_get_filters_before_ordering_and_paging() -> None: + collection = await _create_collection() + + results = await collection.get( + filter=Filter("category", "eq", "travel"), + order_by={"rating": False}, + top=1, + ) + + assert results == [Document("one", "Luxury hotel", "travel", 5, ["wifi", "pool"], "featured")] + + +async def test_in_memory_generates_missing_string_keys() -> None: + definition = VectorStoreCollectionDefinition( + [ + VectorStoreField("key", name="id", type_="str", is_auto_generated=True), + VectorStoreField("data", name="text"), + VectorStoreField("vector", name="vector", dimensions=2), + ], + collection_name="generated-keys", + ) + collection: InMemoryCollection[str, dict[str, Any]] = InMemoryCollection(dict, definition=definition) + await collection.ensure_collection_exists() + + keys = await collection.upsert( + [{"text": "generated", "vector": [1.0, 0.0]}], + generate_vectors=False, + ) + + assert len(keys) == 1 + assert isinstance(keys[0], str) + assert await collection.get(keys, include_vectors=True) == [ + {"id": keys[0], "text": "generated", "vector": [1.0, 0.0]} + ] + + +@pytest.mark.parametrize( + ("filter_", "expected"), + [ + (Filter("category", "eq", "travel"), ["one"]), + (Filter("category", "ne", "travel"), ["two"]), + (Filter("rating", "gt", 4), ["one"]), + (Filter("rating", "gte", 5), ["one"]), + (Filter("rating", "lt", 4), ["two"]), + (Filter("rating", "lte", 3), ["two"]), + (Filter("rating", "between", (4, 5)), ["one"]), + (Filter("category", "in", ("travel",)), ["one"]), + (Filter("category", "not_in", ("travel",)), ["two"]), + (Filter("optional", "is_null"), ["two"]), + (Filter("optional", "is_not_null"), ["one"]), + (Filter("optional", "exists"), ["one", "two"]), + (Filter("optional", "starts_with", "feat"), ["one"]), + (Filter("tags", "contains", "pool"), ["one"]), + (Filter("tags", "contains_any", ("pool", "missing")), ["one"]), + (Filter("tags", "contains_all", ("wifi", "pool")), ["one"]), + (Filter("text", "starts_with", "Luxury"), ["one"]), + (Filter("text", "ends_with", "hostel"), ["two"]), + (Filter("text", "contains_text", "hotel"), ["one"]), + (Filter("rating", "eq", True), []), + (Filter("rating", "in", (True,)), []), + (FilterGroup("not", (Filter("category", "eq", "travel"),)), ["two"]), + ], +) +async def test_in_memory_filter_operators( + filter_: Filter | FilterGroup, + expected: list[str], +) -> None: + assert await _result_ids(await _create_collection(), filter_) == expected + + +@pytest.mark.parametrize( + ("operator", "operand", "expected"), + [ + ("eq", [[1]], {"integer", "float"}), + ("eq", [[True]], {"boolean"}), + ("ne", [[1]], {"boolean", "zero"}), + ("ne", [[True]], {"integer", "float", "zero"}), + ("in", ([[1]],), {"integer", "float"}), + ("in", ([[True]],), {"boolean"}), + ("not_in", ([[1]],), {"boolean", "zero"}), + ("not_in", ([[True]],), {"integer", "float", "zero"}), + ("contains", [1], {"integer", "float"}), + ("contains", [True], {"boolean"}), + ("contains_any", ([1],), {"integer", "float"}), + ("contains_any", ([True],), {"boolean"}), + ("contains_all", ([1],), {"integer", "float"}), + ("contains_all", ([True],), {"boolean"}), + ], +) +async def test_in_memory_filters_distinguish_nested_booleans_from_numbers( + operator: str, + operand: Any, + expected: set[str], +) -> None: + definition = VectorStoreCollectionDefinition( + [ + VectorStoreField("key", name="id"), + VectorStoreField("data", name="value"), + VectorStoreField("vector", name="vector", dimensions=2), + ], + collection_name="nested-values", + ) + collection: InMemoryCollection[str, dict[str, Any]] = InMemoryCollection(dict, definition=definition) + await collection.ensure_collection_exists() + await collection.upsert( + [ + {"id": "integer", "value": [[1]], "vector": [1.0, 0.0]}, + {"id": "float", "value": [[1.0]], "vector": [1.0, 0.0]}, + {"id": "boolean", "value": [[True]], "vector": [1.0, 0.0]}, + {"id": "zero", "value": [[0]], "vector": [1.0, 0.0]}, + ], + generate_vectors=False, + ) + filter_ = Filter("value", operator, operand) + + records = await collection.get(filter=filter_) + assert {record["id"] for record in records} == expected + results = await collection.search(vector=[1.0, 0.0], filter=filter_) + assert {result["record"]["id"] async for result in results} == expected + + +async def test_in_memory_filter_groups_use_explicit_and_or_semantics() -> None: + collection = await _create_collection() + tenant_scope = Filter("category", "eq", "travel") + model_filter = Filter("id", "eq", "two") + + assert await _result_ids(collection, FilterGroup("and", (tenant_scope, model_filter))) == [] + assert await _result_ids(collection, FilterGroup("or", (tenant_scope, model_filter))) == ["one", "two"] + + +async def test_in_memory_search_tool_resolves_params_without_weakening_fixed_filter() -> None: + collection: InMemoryCollection[str, Document] = InMemoryCollection( + Document, + embedding_generator=QueryEmbeddingClient(), + ) + await collection.ensure_collection_exists() + await collection.upsert(DOCUMENTS, generate_vectors=False) + record_id = Param("record_id", str) + tool = create_vector_search_tool( + collection, + filter=FilterGroup( + "and", + ( + Filter("category", "eq", "travel"), + Filter("id", "eq", record_id), + ), + ), + ) + + assert len(await tool(query="hotel", record_id="one")) == 1 + assert await tool(query="hotel", record_id="two") == [] + assert len(await tool(query="hotel")) == 1 + + +async def test_in_memory_search_tool_omits_null_text_condition_in_and_group() -> None: + definition = VectorStoreCollectionDefinition( + [ + VectorStoreField("key", name="id"), + VectorStoreField("data", name="description"), + VectorStoreField("data", name="rating"), + VectorStoreField("vector", name="vector", dimensions=2), + ], + collection_name="optional-text", + ) + collection: InMemoryCollection[str, dict[str, Any]] = InMemoryCollection( + dict, definition=definition, embedding_generator=QueryEmbeddingClient() + ) + await collection.ensure_collection_exists() + records = [ + {"id": "text", "rating": 5, "description": "Pool hotel"}, + {"id": "empty", "rating": 5, "description": ""}, + {"id": "star", "rating": 5, "description": "5* hotel"}, + {"id": "null", "rating": 5, "description": None}, + {"id": "low", "rating": 3, "description": "Pool hotel"}, + ] + await collection.upsert( + [{**record, "vector": [1.0, 0.0]} for record in records], + generate_vectors=False, + ) + tool = create_vector_search_tool( + collection, + top=10, + filter=FilterGroup( + "and", + ( + Filter("rating", "gte", 4), + Filter( + "description", + "contains_text", + Param("text", str | None, default=None, omit_if_none=True), + ), + ), + ), + result_mapper=lambda response: response["record"]["id"], + ) + + assert {item.text for item in await tool(query="hotel")} == {"text", "empty", "star", "null"} + assert {item.text for item in await tool(query="hotel", text=None)} == {"text", "empty", "star", "null"} + assert {item.text for item in await tool(query="hotel", text="Pool")} == {"text"} + assert {item.text for item in await tool(query="hotel", text="")} == {"text", "empty", "star"} + assert {item.text for item in await tool(query="hotel", text="*")} == {"star"} + + +@pytest.mark.parametrize( + ("distance_function", "expected_scores"), + [ + ("cosine_similarity", (1.0, 0.0)), + ("cosine_distance", (0.0, 1.0)), + ("dot_prod", (1.0, 0.0)), + ("negative_dot_prod", (-1.0, 0.0)), + ("euclidean_distance", (0.0, math.sqrt(2))), + ("euclidean_squared_distance", (0.0, 2.0)), + ("manhattan", (0.0, 2.0)), + ("hamming", (0.0, 1.0)), + ("DEFAULT", (0.0, 1.0)), + ], +) +async def test_in_memory_distance_functions( + distance_function: DistanceFunction, + expected_scores: tuple[float, float], +) -> None: + definition = VectorStoreCollectionDefinition( + [ + VectorStoreField("key", name="id"), + VectorStoreField( + "vector", + name="vector", + dimensions=2, + distance_function=distance_function, + ), + ], + collection_name=f"distance-{distance_function}", + ) + collection: InMemoryCollection[str, dict[str, Any]] = InMemoryCollection(dict, definition=definition) + await collection.ensure_collection_exists() + await collection.upsert( + [ + {"id": "same", "vector": [1.0, 0.0]}, + {"id": "other", "vector": [0.0, 1.0]}, + ], + generate_vectors=False, + ) + + results = await collection.search(vector=[1.0, 0.0], top=2) + responses = [result async for result in results] + + assert [response["record"]["id"] for response in responses] == ["same", "other"] + assert [cast(float, response["score"]) for response in responses] == pytest.approx(expected_scores) + assert results.metadata == {"in_memory_total_count": 2} + + +@pytest.mark.parametrize("scale", [1e308, 1e-308, 5e-324]) +@pytest.mark.parametrize("similarity", [1.0, 0.0, -1.0]) +@pytest.mark.parametrize("distance_function", ["cosine_similarity", "cosine_distance", "DEFAULT"]) +def test_cosine_scoring_handles_extreme_finite_magnitudes( + scale: float, similarity: float, distance_function: DistanceFunction +) -> None: + left = [scale, 0.0] + right = [scale, 0.0] if similarity == 1.0 else [-scale, 0.0] if similarity == -1.0 else [0.0, scale] + expected = similarity if distance_function == "cosine_similarity" else 1 - similarity + + assert in_memory_module._calculate_score(left, right, distance_function) == pytest.approx(expected) + + +async def test_large_opposing_vectors_are_not_misranked() -> None: + collection = InMemoryCollection(Document) + await collection.ensure_collection_exists() + await collection.upsert( + [ + Document("opposed", "text", "travel", 5, [], None, [-1e308, 0.0]), + Document("aligned", "text", "travel", 5, [], None, [1e308, 0.0]), + ], + generate_vectors=False, + ) + results = await collection.search(vector=[1e308, 0.0], score_threshold=0.5) + + assert [result["record"].id async for result in results] == ["aligned"] + + +@pytest.mark.parametrize( + ("right", "expected"), + [([1e-308, 1e-308], 1.0), ([-1e-308, -1e-308], -1.0), ([1e-308, -1e-308], 0.0)], +) +def test_cosine_scales_each_vector_independently(right: list[float], expected: float) -> None: + assert in_memory_module._calculate_score([1e308, 1e308], right, "cosine_similarity") == pytest.approx(expected) + + +async def test_non_finite_score_reports_the_record_context() -> None: + definition = VectorStoreCollectionDefinition( + [ + VectorStoreField("key", name="id"), + VectorStoreField("vector", name="vector", dimensions=2, distance_function="dot_prod"), + ], + collection_name="non-finite-score", + ) + collection: InMemoryCollection[str, dict[str, Any]] = InMemoryCollection(dict, definition=definition) + await collection.ensure_collection_exists() + await collection.upsert([{"id": "overflow", "vector": [1e308, 0.0]}], generate_vectors=False) + + with pytest.raises(ValueError, match="Record 'overflow'.*non-finite score"): + await collection.search(vector=[1e308, 0.0]) + + +@pytest.mark.parametrize( + ("distance_function", "left", "right"), + [ + ("dot_prod", [1e308, 0.0], [1e308, 0.0]), + ("dot_prod", [1e308, 1e308], [1e308, -1e308]), + ("negative_dot_prod", [1e308, 0.0], [1e308, 0.0]), + ("euclidean_distance", [1e308, 0.0], [-1e308, 0.0]), + ("euclidean_squared_distance", [1e200, 0.0], [0.0, 0.0]), + ("manhattan", [1e308, 0.0], [-1e308, 0.0]), + ], +) +def test_distance_functions_reject_non_finite_scores( + distance_function: DistanceFunction, left: list[float], right: list[float] +) -> None: + with pytest.raises(ValueError, match="non-finite score"): + in_memory_module._calculate_score(left, right, distance_function) + + +def test_distance_functions_do_not_compute_unnecessary_squares() -> None: + assert in_memory_module._calculate_score([1e200, 0], [0, 0], "euclidean_distance") == 1e200 + assert in_memory_module._calculate_score([1e308, 0], [-1e308, 0], "hamming") == 0.5 + + +async def test_hamming_scores_and_thresholds_use_mismatch_proportions() -> None: + definition = VectorStoreCollectionDefinition( + [ + VectorStoreField("key", name="id"), + VectorStoreField("vector", name="vector", dimensions=2, distance_function="hamming"), + ], + collection_name="normalized-hamming", + ) + collection: InMemoryCollection[str, dict[str, Any]] = InMemoryCollection(dict, definition=definition) + await collection.ensure_collection_exists() + await collection.upsert( + [ + {"id": "same", "vector": [0, 0]}, + {"id": "partial", "vector": [1, 0]}, + {"id": "different", "vector": [1, 1]}, + ], + generate_vectors=False, + ) + results = await collection.search(vector=[0, 0], score_threshold=0.5) + + assert [(result["record"]["id"], result["score"]) async for result in results] == [ + ("same", 0.0), + ("partial", 0.5), + ] + + +@pytest.mark.parametrize("contents", ["empty", "missing_vector", "filtered_out"]) +@pytest.mark.parametrize("score_threshold", [None, 0.5]) +async def test_unsupported_distance_is_rejected_independently_of_records( + contents: str, score_threshold: float | None +) -> None: + definition = VectorStoreCollectionDefinition( + [ + VectorStoreField("key", name="id"), + VectorStoreField("vector", name="vector", dimensions=2, distance_function="provider.unsupported"), + ], + collection_name="unsupported-distance", + ) + collection: InMemoryCollection[str, dict[str, Any]] = InMemoryCollection(dict, definition=definition) + await collection.ensure_collection_exists() + if contents != "empty": + await collection.upsert( + [{"id": "one", "vector": None if contents == "missing_vector" else [1.0, 0.0]}], + generate_vectors=False, + ) + with pytest.raises(NotImplementedError, match="provider.unsupported"): + await collection.search( + vector=[1.0, 0.0], + filter=Filter("id", "eq", "absent") if contents == "filtered_out" else None, + score_threshold=score_threshold, + ) + + +@pytest.mark.parametrize( + ("distance_function", "threshold"), + [("DEFAULT", 1.0), ("cosine_distance", 1.0), ("cosine_similarity", 0.0)], +) +async def test_in_memory_owns_filtering_and_thresholds_before_paging( + distance_function: DistanceFunction, threshold: float +) -> None: + definition = VectorStoreCollectionDefinition( + [ + VectorStoreField("key", name="id"), + VectorStoreField("data", name="category"), + VectorStoreField("vector", name="vector", dimensions=2, distance_function=distance_function), + ], + collection_name="threshold-paging", + ) + collection: InMemoryCollection[str, dict[str, Any]] = InMemoryCollection(dict, definition=definition) + await collection.ensure_collection_exists() + await collection.upsert( + [ + {"id": "threshold-excluded", "category": "keep", "vector": [-1.0, 0.0]}, + {"id": "filter-excluded", "category": "exclude", "vector": [1.0, 0.0]}, + {"id": "best", "category": "keep", "vector": [1.0, 0.0]}, + {"id": "boundary", "category": "keep", "vector": [0.0, 1.0]}, + ], + generate_vectors=False, + ) + + results = await collection.search( + vector=[1.0, 0.0], + filter=Filter("category", "eq", "keep"), + score_threshold=threshold, + skip=1, + top=1, + ) + + assert [(result["record"]["id"], result["score"]) async for result in results] == [("boundary", threshold)] + assert results.metadata == {"in_memory_total_count": 2} + + +@pytest.mark.parametrize("has_records", [False, True]) +async def test_in_memory_default_threshold_accepts_zero_distance(has_records: bool) -> None: + definition = VectorStoreCollectionDefinition( + [ + VectorStoreField("key", name="id"), + VectorStoreField("vector", name="vector", dimensions=2), + ], + collection_name="default-threshold", + ) + collection: InMemoryCollection[str, dict[str, Any]] = InMemoryCollection(dict, definition=definition) + await collection.ensure_collection_exists() + if has_records: + await collection.upsert( + [ + {"id": "same", "vector": [1.0, 0.0]}, + {"id": "different", "vector": [0.0, 1.0]}, + ], + generate_vectors=False, + ) + + results = await collection.search(vector=[1.0, 0.0], score_threshold=0.0) + + assert [(result["record"]["id"], result["score"]) async for result in results] == ( + [("same", 0.0)] if has_records else [] + ) + + +@pytest.mark.parametrize("operator", ["starts_with", "ends_with", "contains_text"]) +async def test_invalid_string_operand_is_rejected_for_empty_collection(operator: str) -> None: + collection = InMemoryCollection(Document) + await collection.ensure_collection_exists() + expression = Filter("text", operator, "valid") + expression.value = 1 + + with pytest.raises(TypeError, match="must be a string"): + await collection.search(vector=[1.0, 0.0], filter=expression) + + +@pytest.mark.parametrize("custom_codec", [False, True]) +async def test_stored_values_are_normalized_before_filter_comparisons(custom_codec: bool) -> None: + class ComparisonProbe: + def __init__(self) -> None: + self.value = 1 + + def __eq__(self, other: object) -> bool: + raise AssertionError("Custom equality must not execute during filtering.") + + class ComparisonList(list[Any]): + def __eq__(self, other: object) -> bool: + raise AssertionError("Custom list equality must not execute during filtering.") + + def __iter__(self) -> Iterator[Any]: + raise AssertionError("Custom iteration must not execute during filtering.") + + @dataclass + class EncodedRecord: + id: str + payload: Any + vector: list[float] | None = None + + definition = VectorStoreCollectionDefinition( + [ + VectorStoreField("key", name="id"), + VectorStoreField("data", name="payload"), + VectorStoreField("vector", name="vector", dimensions=2), + ], + collection_name="normalized-records", + ) + payload = ComparisonList([ComparisonProbe()]) + collection: InMemoryCollection[str, Any] + if custom_codec: + register_vectorstoremodel( + EncodedRecord, + definition=definition, + encoder=lambda record: {"id": record.id, "payload": record.payload, "vector": record.vector}, + ) + collection = InMemoryCollection(EncodedRecord) + await collection.ensure_collection_exists() + await collection.upsert([EncodedRecord("one", payload, [1.0, 0.0])], generate_vectors=False) + else: + collection = InMemoryCollection(dict, definition=definition) + await collection.ensure_collection_exists() + await collection.upsert([{"id": "one", "payload": payload, "vector": [1.0, 0.0]}], generate_vectors=False) + fetched = (await collection.get(["one"]))[0] + stored = fetched.payload if isinstance(fetched, EncodedRecord) else fetched["payload"] + + assert type(stored) is list + assert stored == [{"value": 1}] + assert await collection.get(filter=Filter("payload", "eq", [1])) == [] + assert await collection.get(filter=Filter("payload", "contains", 1)) == [] + + +async def test_in_memory_rejects_unsafe_or_unsupported_filters() -> None: + collection = await _create_collection() + + with pytest.raises(TypeError, match="Filter or FilterGroup"): + await collection.search(vector=[1.0, 0.0], filter=cast(Any, "lambda record: True")) + with pytest.raises(NotImplementedError, match="python.eval"): + await collection.search( + vector=[1.0, 0.0], + filter=Filter("text", "python.eval", "__import__('os').system('echo unsafe')"), + ) + with pytest.raises(NotImplementedError, match="nested"): + await collection.search(vector=[1.0, 0.0], filter=Filter("text.value", "eq", "unsafe")) + with pytest.raises(ValueError, match="more than 64 nodes"): + await collection.search(vector=[1.0, 0.0], filter=Filter("id", "in", tuple(str(i) for i in range(65)))) + with pytest.raises(ValueError, match="must be resolved"): + await collection.search(vector=[1.0, 0.0], filter=Filter("category", "eq", Param("category", str))) + with pytest.raises(ValueError, match="must be finite"): + await collection.search(vector=[1.0, 0.0], filter=Filter("rating", "ne", float("nan"))) + + +@pytest.mark.parametrize( + "payload", + [ + "x.__class__.__mro__[-1].__subclasses__()", + "x.__init__.__globals__", + "__import__('os').system('echo unsafe')", + "x.clear()", + "x['a']['b']()", + ], +) +async def test_in_memory_treats_filter_payloads_only_as_data(payload: str) -> None: + collection = await _create_collection() + + assert await _result_ids(collection, Filter("text", "eq", payload)) == [] + + +async def test_in_memory_rejects_invalid_vectors() -> None: + collection = await _create_collection() + + with pytest.raises(TypeError, match="field 'query'.*numeric sequence"): + await collection.search(vector=b"\x01\x02") + with pytest.raises(ValueError, match="Query vector field 'vector' expects 2 dimensions; got 1"): + await collection.search(vector=[1.0]) + with pytest.raises(ValueError, match="zero-magnitude"): + await collection.search(vector=[0.0, 0.0]) + + await collection.upsert( + [Document("no-vector", "No vector", "travel", 5, [], None)], + generate_vectors=False, + ) + results = await collection.search(vector=[1.0, 0.0], top=10) + assert [result["record"].id async for result in results] == ["one", "two"] + + await collection.upsert( + [Document("bad-vector", "Bad vector", "travel", 5, [], None, cast(Any, ["bad", 0.0]))], + generate_vectors=False, + ) + with pytest.raises(TypeError, match="Record 'bad-vector'.*field 'vector'"): + await collection.search(vector=[1.0, 0.0], top=10) + + binary_definition = VectorStoreCollectionDefinition( + [ + VectorStoreField("key", name="id"), + VectorStoreField("vector", name="vector", dimensions=8, type_="bytes"), + ], + collection_name="binary", + ) + with pytest.raises(ValueError, match="type must be one of"): + InMemoryCollection(dict, definition=binary_definition) + + +@pytest.mark.parametrize("vector", [[], [1.0], [1.0, 0.0, 0.0]]) +async def test_in_memory_dimension_mismatch_rejects_batch_without_writes(vector: list[float]) -> None: + collection = InMemoryCollection(Document) + await collection.ensure_collection_exists() + + with pytest.raises( + ValueError, match=f"Record at index 1, vector field 'vector' expects 2 dimensions; got {len(vector)}" + ): + await collection.upsert( + [DOCUMENTS[0], Document("invalid", "text", "travel", 5, [], None, vector)], + generate_vectors=False, + ) + + assert await collection.get() == [] + + +@pytest.mark.parametrize("has_records", [False, True]) +@pytest.mark.parametrize("dimensions", [1, 2, 3]) +async def test_in_memory_checks_normalized_array_query_dimensions(has_records: bool, dimensions: int) -> None: + class ArrayLike: + def tolist(self) -> list[float]: + return [1.0] * dimensions + + collection = InMemoryCollection(Document) + await collection.ensure_collection_exists() + if has_records: + await collection.upsert(DOCUMENTS, generate_vectors=False) + + if dimensions != 2: + with pytest.raises(ValueError, match=f"Query vector field 'vector' expects 2 dimensions; got {dimensions}"): + await collection.search(vector=cast(Any, ArrayLike())) + else: + results = await collection.search(vector=cast(Any, ArrayLike())) + assert len([result async for result in results]) == (2 if has_records else 0) + + +def test_in_memory_retains_pairwise_vector_length_check() -> None: + with pytest.raises(ValueError, match="Query and stored vectors must have the same length"): + in_memory_module._calculate_score([1.0, 0.0], [1.0], "cosine_similarity") + + +async def test_in_memory_contains_rejects_mapping_fields() -> None: + definition = VectorStoreCollectionDefinition( + [ + VectorStoreField("key", name="id"), + VectorStoreField("data", name="metadata"), + VectorStoreField("vector", name="vector", dimensions=2), + ], + collection_name="mapping-values", + ) + collection: InMemoryCollection[str, dict[str, Any]] = InMemoryCollection(dict, definition=definition) + await collection.ensure_collection_exists() + await collection.upsert( + [{"id": "one", "metadata": {"wifi": True}, "vector": [1.0, 0.0]}], + generate_vectors=False, + ) + + with pytest.raises(ValueError, match="cannot compare field 'metadata'"): + await collection.search(vector=[1.0, 0.0], filter=Filter("metadata", "contains", "wifi")) + + +def test_in_memory_filter_implementation_does_not_compile_or_evaluate_source() -> None: + tree = ast.parse(inspect.getsource(in_memory_module)) + called_names = { + node.func.id for node in ast.walk(tree) if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) + } + + assert called_names.isdisjoint({"compile", "eval", "exec"}) + + +def test_in_memory_and_filter_apis_are_experimental() -> None: + for api in (Filter, FilterGroup, Param, InMemoryCollection, InMemoryStore): + assert getattr(api, "__feature_stage__", None) == "experimental" diff --git a/python/packages/core/tests/core/test_vectors.py b/python/packages/core/tests/core/test_vectors.py index 9bb045e61d7..ac56df4e048 100644 --- a/python/packages/core/tests/core/test_vectors.py +++ b/python/packages/core/tests/core/test_vectors.py @@ -3,10 +3,10 @@ from __future__ import annotations import warnings -from ast import AST, unparse -from collections.abc import AsyncIterable, Mapping, Sequence +from collections.abc import AsyncIterable, Callable, Iterator, Mapping, Sequence from dataclasses import FrozenInstanceError, dataclass, field -from typing import Annotated, Any, ClassVar, cast +from decimal import Decimal +from typing import Annotated, Any, ClassVar, Literal, cast from unittest.mock import patch import msgspec @@ -27,8 +27,11 @@ EmbeddingGenerationOptions, ExperimentalFeature, FieldTypes, + Filter, + FilterGroup, GeneratedEmbeddings, IndexKind, + Param, SearchResponse, SearchResults, SearchType, @@ -42,6 +45,7 @@ ) from agent_framework._feature_stage import ExperimentalWarning from agent_framework._telemetry import FeatureIndex +from agent_framework._vector_filters import filter_values_equal from agent_framework._vectors import _VectorStoreRecordHandler as VectorStoreRecordHandler from agent_framework.exceptions import IntegrationException, IntegrationInvalidResponseException @@ -56,6 +60,7 @@ dimensions=2, index_kind="hnsw", distance_function="cosine_similarity", + provider_annotations={"index": {"ef_construction": 200}}, ), ] @@ -95,15 +100,19 @@ def __init__(self, *, embedding_generator: MockEmbeddingClient | None = None) -> self.created = False self.records: dict[str, dict[str, Any]] = {} self.last_search_type: str | None = None + self.last_search_values: Any | None = None self.last_search_vector: Sequence[float | int] | None = None - self.last_search_filter: Any | list[Any] | None = None + self.last_search_filter: Filter | FilterGroup | None = None + self.last_get_filter: Filter | FilterGroup | None = None self.last_search_top = 0 self.last_search_skip = 0 + self.last_search_score_threshold: float | None = None self.fail_upsert = False self.upsert_error: Exception | None = None self.get_error: Exception | None = None self.delete_error: Exception | None = None self.search_error: Exception | None = None + self.mutate_search_filter = False self.upsert_keys: Sequence[str] | None = None self.raw_search_results: AsyncIterable[Any] | Sequence[Any] | None = None @@ -151,6 +160,7 @@ async def _inner_get( self, *, keys: Sequence[str] | None = None, + filter: Filter | FilterGroup | None = None, top: int = 10, skip: int = 0, order_by: Mapping[str, bool] | None = None, @@ -159,6 +169,7 @@ async def _inner_get( ) -> Sequence[Any] | None: if self.get_error is not None: raise self.get_error + self.last_get_filter = filter if keys is not None: return [self.records[key] for key in keys if key in self.records] return list(self.records.values())[skip : skip + top] @@ -178,7 +189,7 @@ async def _inner_search( self, *, search_type: SearchType, - filter: Any | list[Any] | None = None, + filter: Filter | FilterGroup | None = None, values: Any | None = None, vector: Sequence[float | int] | None = None, top: int = 3, @@ -192,10 +203,14 @@ async def _inner_search( if self.search_error is not None: raise self.search_error self.last_search_type = search_type + self.last_search_values = values self.last_search_vector = vector self.last_search_filter = filter self.last_search_top = top self.last_search_skip = skip + self.last_search_score_threshold = score_threshold + if self.mutate_search_filter and isinstance(filter, Filter) and isinstance(filter.value, list): + filter.value.append("connector mutation") raw_results = self.raw_search_results or [ {"record": record, "score": score} for record, score in zip(self.records.values(), (0.9, 0.4), strict=False) ] @@ -207,9 +222,6 @@ def _get_record_from_result(self, result: Any) -> Any: def _get_score_from_result(self, result: Any) -> float | None: return cast(float | None, result["score"]) - def _lambda_parser(self, node: AST) -> str: - return unparse(node) - StoreModelT = TypeVar("StoreModelT") @@ -289,10 +301,30 @@ def test_vector_field_validates_vector_options() -> None: cast(Any, VectorStoreField)("vector") with pytest.raises(ValueError, match="Vector-only"): cast(Any, VectorStoreField)("data", dimensions=3) - with pytest.raises(ValueError, match="index kind"): - cast(Any, VectorStoreField)("vector", dimensions=3, index_kind="unknown") - with pytest.raises(ValueError, match="distance function"): - cast(Any, VectorStoreField)("vector", dimensions=3, distance_function="unknown") + with pytest.raises(ValueError, match="Only key fields"): + cast(Any, VectorStoreField)("data", is_auto_generated=True) + with pytest.raises(TypeError, match="index_kind must be a string"): + cast(Any, VectorStoreField)("vector", dimensions=3, index_kind=1) + with pytest.raises(TypeError, match="distance_function must be a string"): + cast(Any, VectorStoreField)("vector", dimensions=3, distance_function=object()) + + annotations = {"index": {"ef_construction": 200}} + provider_field = VectorStoreField( + "vector", + dimensions=3, + index_kind="provider.custom_index", + distance_function="provider.custom_distance", + provider_annotations=annotations, + ) + annotations["index"]["ef_construction"] = 100 + assert provider_field.index_kind == "provider.custom_index" + assert provider_field.distance_function == "provider.custom_distance" + assert provider_field.provider_annotations["index"]["ef_construction"] == 200 + provider_field.provider_annotations["index"]["ef_construction"] = 100 + assert provider_field.provider_annotations["index"]["ef_construction"] == 100 + assert hash(provider_field) + with pytest.raises(TypeError, match="keys must be strings"): + VectorStoreField("data", provider_annotations=cast(Any, {1: "value"})) def test_collection_definition_exposes_fields() -> None: @@ -309,8 +341,11 @@ def test_collection_definition_exposes_fields() -> None: assert definition.get_storage_names(include_key_field=False) == ["body", "vector"] assert isinstance(definition.fields, tuple) assert definition.vector_fields[0].dimensions == 2 + assert definition.vector_fields[0].type_ == "float" assert definition.vector_fields[0].index_kind == "hnsw" assert definition.vector_fields[0].distance_function == "cosine_similarity" + assert definition.vector_fields[0].provider_annotations["index"]["ef_construction"] == 200 + assert not definition.key_field.is_auto_generated frozen_field = cast(Any, definition.fields[0]) with pytest.raises(FrozenInstanceError): @@ -320,6 +355,33 @@ def test_collection_definition_exposes_fields() -> None: frozen_definition.fields = () +@pytest.mark.parametrize("value", [0, 1, "false", None]) +@pytest.mark.parametrize("field_type", ["key", "data", "vector"]) +def test_auto_generated_flag_requires_a_boolean(value: Any, field_type: FieldTypes) -> None: + with pytest.raises(TypeError, match="is_auto_generated must be a boolean"): + cast(Any, VectorStoreField)( + field_type, is_auto_generated=value, dimensions=2 if field_type == "vector" else None + ) + + +async def test_bytearray_annotation_uses_normalized_binary_metadata() -> None: + @vectorstoremodel + @dataclass + class BinaryRecord: + id: Annotated[str, VectorStoreField("key")] + vector: Annotated[bytearray | None, VectorStoreField("vector", dimensions=24)] = None + + class BinaryHandler(VectorStoreRecordHandler[str, BinaryRecord]): + supported_vector_types: ClassVar[set[str] | None] = {"bytes"} + + handler = BinaryHandler(BinaryRecord) + assert handler.definition.vector_fields[0].type_ == "bytes" + assert await handler.serialize(BinaryRecord("one", bytearray((1, 2, 3))), generate_vectors=False) == { + "id": "one", + "vector": b"\x01\x02\x03", + } + + @pytest.mark.parametrize( "fields, message", [ @@ -339,6 +401,13 @@ def test_collection_definition_exposes_fields() -> None: ], "must be unique", ), + ( + [ + VectorStoreField("key", name="id", storage_name="record_id"), + VectorStoreField("data", name="record_id", storage_name="body"), + ], + "storage name cannot match another field's model name", + ), ], ) def test_collection_definition_rejects_invalid_fields( @@ -585,6 +654,9 @@ class ArrayRecord: assert isinstance(restored, ArrayRecord) assert restored.vector.values == [0.1, 0.2, 0.3] + with pytest.raises(ValueError, match="vector field 'vector' expects 3 dimensions; got 1"): + await handler.serialize(ArrayRecord("bad", ArrayLike([0.1])), generate_vectors=False) + async def test_custom_encoder_normalizes_array_like_vectors() -> None: class ArrayLike: @@ -646,6 +718,15 @@ async def test_collection_serializes_records_and_generates_vectors() -> None: assert embedding_client.options == {"dimensions": 2} +async def test_collection_empty_upsert_skips_embedding_generation() -> None: + embedding_client = MockEmbeddingClient() + collection = MockCollection(embedding_generator=embedding_client) + + assert await collection.upsert([]) == [] + assert embedding_client.values == [] + assert await MockCollection().upsert([]) == [] + + async def test_upsert_controls_embedding_generation() -> None: embedding_client = MockEmbeddingClient() collection = MockCollection(embedding_generator=embedding_client) @@ -668,6 +749,206 @@ async def test_upsert_controls_embedding_generation() -> None: await MockCollection().upsert([Record("missing-generator", "text", [1.0, 0.0])]) +@pytest.mark.parametrize("vector", [[], [1.0], [1.0, 0.0, 0.0]]) +async def test_upsert_rejects_dimension_mismatches_before_connector_conversion(vector: list[float]) -> None: + collection = MockCollection() + records = [Record("valid", "text", [1.0, 0.0]), Record("invalid", "text", vector)] + + with patch.object( + collection, "_serialize_dicts_to_store_models", wraps=collection._serialize_dicts_to_store_models + ) as convert: + with pytest.raises( + ValueError, match=f"Record at index 1, vector field 'vector' expects 2 dimensions; got {len(vector)}" + ): + await collection.upsert(records, generate_vectors=False) + convert.assert_not_called() + + assert collection.records == {} + + +async def test_upsert_checks_generated_dimensions_instead_of_replaced_input() -> None: + collection = MockCollection(embedding_generator=MockEmbeddingClient()) + record = Record("replaced", "text", [1.0]) + + await collection.upsert([record]) + + assert len(collection.records["replaced"]["vector"]) == 2 + assert record.vector == [1.0] + + +async def test_upsert_rejects_generated_dimension_mismatch_before_writing() -> None: + class WrongDimensionEmbeddingClient(MockEmbeddingClient): + async def get_embeddings( + self, + values: Sequence[Any], + *, + options: EmbeddingGenerationOptions | None = None, + ) -> GeneratedEmbeddings[list[float]]: + return GeneratedEmbeddings([ + Embedding(vector=[1.0, 0.0] if index == 0 else [1.0], dimensions=2) for index, _ in enumerate(values) + ]) + + collection = MockCollection(embedding_generator=WrongDimensionEmbeddingClient()) + + with pytest.raises(ValueError, match="Record at index 1, vector field 'vector' expects 2 dimensions; got 1"): + await collection.upsert([Record("valid", "text", "source"), Record("invalid", "text", "source")]) + + assert collection.records == {} + + +@pytest.mark.parametrize(("field_name", "dimensions"), [("primary", 2), ("secondary", 3)]) +async def test_serialization_validates_all_vector_fields_and_storage_aliases(field_name: str, dimensions: int) -> None: + definition = VectorStoreCollectionDefinition([ + VectorStoreField("key", name="id"), + VectorStoreField("vector", name="primary", storage_name="primary_vector", dimensions=2), + VectorStoreField("vector", name="secondary", storage_name="secondary_vector", dimensions=3), + ]) + handler = VectorStoreRecordHandler(dict, definition=definition) + records = [ + {"id": "one", "primary": [1.0, 0.0], "secondary": [1.0, 0.0, 0.0]}, + {"id": "two", "primary_vector": [1.0, 0.0], "secondary_vector": [1.0, 0.0, 0.0]}, + ] + records[1][f"{field_name}_vector"] = [1.0] + + with pytest.raises( + ValueError, match=f"Record at index 1, vector field '{field_name}' expects {dimensions} dimensions; got 1" + ): + await handler.serialize(records, generate_vectors=False) + + +@pytest.mark.parametrize( + ("payload", "expected"), + [ + (None, None), + ("provider-side source", "provider-side source"), + (b"\x01\x02", b"\x01\x02"), + (bytearray((1, 2)), b"\x01\x02"), + ({"indices": [4], "values": [0.5]}, {"indices": [4], "values": [0.5]}), + ], +) +async def test_serialization_leaves_non_dense_dimensions_to_connector(payload: Any, expected: Any) -> None: + definition = VectorStoreCollectionDefinition([ + VectorStoreField("key", name="id"), + VectorStoreField("vector", name="vector", dimensions=1536), + ]) + handler = VectorStoreRecordHandler(dict, definition=definition) + + assert await handler.serialize({"id": "one", "vector": payload}, generate_vectors=False) == { + "id": "one", + "vector": expected, + } + + +async def test_serialization_selects_vector_fields_for_generation() -> None: + embedding_client = MockEmbeddingClient() + + @dataclass + class MixedVectorRecord: + id: str + local_vector: str | list[float] | None = None + provider_vector: str | list[float] | None = None + + definition = VectorStoreCollectionDefinition([ + VectorStoreField("key", name="id"), + VectorStoreField( + "vector", + name="local_vector", + dimensions=2, + embedding_generator=embedding_client, + ), + VectorStoreField("vector", name="provider_vector", dimensions=2), + ]) + register_vectorstoremodel(MixedVectorRecord, definition=definition) + handler = VectorStoreRecordHandler(MixedVectorRecord) + serialized = await handler.serialize( + MixedVectorRecord("one", "embed locally", "send to provider"), + generate_vectors=["local_vector"], + ) + + assert serialized == { + "id": "one", + "local_vector": [13.0, 0.5], + "provider_vector": "send to provider", + } + assert embedding_client.values == ["embed locally"] + + with pytest.raises(ValueError, match="vector field 'provider_vector' expects 2 dimensions; got 1"): + await handler.serialize( + MixedVectorRecord("bad", "embed locally", [1.0]), + generate_vectors=["local_vector"], + ) + + with pytest.raises(ValueError, match="Unknown vector field"): + await handler.serialize( + MixedVectorRecord("one", "local", "provider"), + generate_vectors=["missing"], + ) + with pytest.raises(ValueError, match="must be unique"): + await handler.serialize( + MixedVectorRecord("one", "local", "provider"), + generate_vectors=["local_vector", "local_vector"], + ) + with pytest.raises(TypeError, match="boolean or a sequence"): + await handler.serialize( + MixedVectorRecord("one", "local", "provider"), + generate_vectors=cast(Any, "local_vector"), + ) + + +async def test_generated_binary_vectors_are_preserved() -> None: + class ByteArrayEmbeddingClient(BaseEmbeddingClient[Any, bytearray, EmbeddingGenerationOptions]): + async def get_embeddings( + self, + values: Sequence[Any], + *, + options: EmbeddingGenerationOptions | None = None, + ) -> GeneratedEmbeddings[bytearray]: + return GeneratedEmbeddings([Embedding(vector=bytearray((1, 2, 3))) for _ in values]) + + @vectorstoremodel + class BinaryRecord(BaseModel): + id: Annotated[str, VectorStoreField("key")] + vector: Annotated[ + str | bytes | None, + VectorStoreField("vector", dimensions=24), + ] = None + + handler = VectorStoreRecordHandler( + BinaryRecord, + embedding_generator=ByteArrayEmbeddingClient(), + ) + serialized = await handler.serialize(BinaryRecord(id="one", vector="source")) + + assert serialized == {"id": "one", "vector": b"\x01\x02\x03"} + assert handler.definition.vector_fields[0].type_ == "bytes" + assert await handler.serialize( + BinaryRecord(id="two", vector=b"\x04\x05\x06"), + generate_vectors=False, + ) == {"id": "two", "vector": b"\x04\x05\x06"} + + @dataclass + class UnsupportedBinaryRecord: + id: Annotated[str, VectorStoreField("key")] + vector: Annotated[str | bytes | None, VectorStoreField("vector", dimensions=24)] = None + + with pytest.raises(ValueError, match="default msgspec decoder.*custom decoder"): + vectorstoremodel(UnsupportedBinaryRecord) + + @vectorstoremodel + class BytesOnlyRecord(BaseModel): + id: Annotated[str, VectorStoreField("key")] + vector: Annotated[bytes | None, VectorStoreField("vector", dimensions=24)] = None + + pydantic_handler = VectorStoreRecordHandler(BytesOnlyRecord) + supplied = BytesOnlyRecord(id="three", vector=b"\x07\x08\x09") + serialized_supplied = await pydantic_handler.serialize(supplied, generate_vectors=False) + restored = pydantic_handler.deserialize(serialized_supplied) + + assert serialized_supplied == {"id": "three", "vector": b"\x07\x08\x09"} + assert isinstance(restored, BytesOnlyRecord) + assert restored.vector == b"\x07\x08\x09" + + async def test_collection_crud_preserves_single_and_batch_shapes() -> None: collection = MockCollection(embedding_generator=MockEmbeddingClient()) await collection.ensure_collection_exists() @@ -706,6 +987,18 @@ async def test_collection_get_without_keys_lists_records() -> None: assert await MockCollection().get() == [] +async def test_collection_get_accepts_filter_as_alternate_retrieval_mode() -> None: + collection = MockCollection() + filter_ = Filter("text", "eq", "hello") + + await collection.get(filter=filter_) + + assert collection.last_get_filter == filter_ + assert collection.last_get_filter is not filter_ + with pytest.raises(ValueError, match="alternate retrieval modes"): + await collection.get(["one"], filter=filter_) + + async def test_collection_crud_rejects_singular_ordinary_inputs() -> None: collection = MockCollection() @@ -717,7 +1010,7 @@ async def test_collection_crud_rejects_singular_ordinary_inputs() -> None: await collection.delete("one") -async def test_vector_search_generates_query_vector_and_filters_threshold() -> None: +async def test_vector_search_generates_query_vector_and_forwards_threshold() -> None: embedding_client = MockEmbeddingClient() collection = MockCollection(embedding_generator=embedding_client) await collection.upsert([ @@ -734,9 +1027,11 @@ async def test_vector_search_generates_query_vector_and_filters_threshold() -> N assert results.metadata == {"mock_count": 2} assert embedding_client.values == ["find this"] assert collection.last_search_vector == [9.0, 0.5] + assert collection.last_search_score_threshold == 0.5 assert responses[0]["record"].id == "one" assert responses[0]["score"] == 0.9 - assert len(responses) == 1 + assert [response["record"].id for response in responses] == ["one", "two"] + assert responses[1]["score"] == 0.4 async def test_keyword_hybrid_search_uses_single_search_method() -> None: @@ -747,11 +1042,110 @@ async def test_keyword_hybrid_search_uses_single_search_method() -> None: assert collection.last_search_type == "keyword_hybrid" +@pytest.mark.parametrize("values", ["vectorize this on the provider", {"indices": [4], "values": [0.5]}, [1.0]]) +async def test_search_passes_values_to_provider_when_no_generator_is_configured(values: Any) -> None: + collection = MockCollection() + + await collection.search(values) + + assert collection.last_search_values is values + assert collection.last_search_vector is None + + +@pytest.mark.parametrize("vector", [[], [1.0], (1.0,), range(1), [1.0, 0.0, 0.0]]) +@pytest.mark.parametrize("search_type", ["vector", "keyword_hybrid"]) +async def test_search_rejects_dense_dimension_mismatches_before_dispatch( + vector: Sequence[float], search_type: SearchType +) -> None: + collection = MockCollection() + + with pytest.raises(ValueError, match=f"Query vector field 'vector' expects 2 dimensions; got {len(vector)}"): + await collection.search("query", vector=vector, search_type=search_type) + + assert collection.last_search_type is None + + +@pytest.mark.parametrize( + ("selected_field", "logical_name", "dimensions"), + [(None, "vector", 2), ("secondary", "secondary", 3), ("secondary_vector", "secondary", 3)], +) +async def test_search_checks_dimensions_of_selected_vector_field( + selected_field: str | None, logical_name: str, dimensions: int +) -> None: + collection = MockCollection() + collection.definition = VectorStoreCollectionDefinition([ + VectorStoreField("key", name="id"), + VectorStoreField("vector", name="vector", dimensions=2), + VectorStoreField("vector", name="secondary", storage_name="secondary_vector", dimensions=3), + ]) + vector = [1.0] * dimensions + + await collection.search(vector=vector, vector_property_name=selected_field) + assert collection.last_search_vector is vector + + with pytest.raises( + ValueError, match=f"Query vector field '{logical_name}' expects {dimensions} dimensions; got {dimensions + 1}" + ): + await collection.search(vector=[1.0] * (dimensions + 1), vector_property_name=selected_field) + assert collection.last_search_vector is vector + + +async def test_search_rejects_generated_dimension_mismatch() -> None: + class WrongDimensionEmbeddingClient(MockEmbeddingClient): + async def get_embeddings( + self, + values: Sequence[Any], + *, + options: EmbeddingGenerationOptions | None = None, + ) -> GeneratedEmbeddings[list[float]]: + return GeneratedEmbeddings([Embedding(vector=[1.0], dimensions=2)]) + + collection = MockCollection(embedding_generator=WrongDimensionEmbeddingClient()) + + with pytest.raises(ValueError, match="Query vector field 'vector' expects 2 dimensions; got 1"): + await collection.search("source") + + assert collection.last_search_type is None + + +async def test_search_dimension_check_does_not_iterate_or_copy_vectors() -> None: + class LengthOnlyVector(list[float]): + def __iter__(self) -> Iterator[float]: + raise AssertionError("Dimension checks must not iterate vector elements.") + + vector = LengthOnlyVector([1.0, 0.0]) + collection = MockCollection() + + await collection.search(vector=vector) + + assert collection.last_search_vector is vector + + +@pytest.mark.parametrize("vector", [(1.0, 0.0), range(2)]) +async def test_search_accepts_matching_dense_sequence_dimensions(vector: Sequence[float | int]) -> None: + collection = MockCollection() + + await collection.search(vector=vector) + + assert collection.last_search_vector is vector + + +@pytest.mark.parametrize("vector", [b"\x01", bytearray((1,))]) +async def test_search_leaves_binary_dimensions_to_connector(vector: bytes | bytearray) -> None: + collection = MockCollection() + + await collection.search(vector=vector) + + assert collection.last_search_vector is vector + + async def test_vector_search_validates_inputs_and_supported_type() -> None: collection = MockCollection() with pytest.raises(ValueError, match="requires values"): await cast(Any, collection.search)() + with pytest.raises(ValueError, match="Vector field 'missing' was not found"): + await collection.search(vector=[1.0, 0.0], vector_property_name="missing") class VectorOnlyCollection(MockCollection): supported_search_types: ClassVar[set[SearchType]] = {"vector"} @@ -760,17 +1154,107 @@ class VectorOnlyCollection(MockCollection): await VectorOnlyCollection().search("words", search_type="keyword_hybrid") -async def test_vector_search_requires_explicit_distance_for_score_threshold() -> None: +@pytest.mark.parametrize( + "distance_function", [None, "DEFAULT", "provider.custom_distance", "cosine_similarity", "cosine_distance"] +) +@pytest.mark.parametrize("precomputed", [False, True]) +@pytest.mark.parametrize("threshold", [0.0, 0.5]) +async def test_core_preserves_connector_scores_without_interpreting_thresholds( + distance_function: DistanceFunction | None, precomputed: bool, threshold: float +) -> None: collection = MockCollection() + fields = [ + VectorStoreField("key", name="id", storage_name="record_id"), + VectorStoreField("data", name="text", storage_name="body"), + ] + if distance_function is not None: + fields.append( + VectorStoreField("vector", name="vector", type_="float", dimensions=2, distance_function=distance_function) + ) + collection.definition = VectorStoreCollectionDefinition( + fields, + collection_name="records", + ) + collection.raw_search_results = [ + {"record": {"record_id": "low", "body": "low"}, "score": -0.2}, + {"record": {"record_id": "equal", "body": "equal"}, "score": threshold}, + {"record": {"record_id": "high", "body": "high"}, "score": 0.9}, + {"record": {"record_id": "scoreless", "body": "scoreless"}, "score": None}, + ] + filter_ = Filter("text", "eq", "provider interprets this") + + results = await collection.search( + "query", + vector=[1.0, 0.0] if precomputed else None, + filter=filter_, + score_threshold=threshold, + top=4, + skip=2, + ) + responses = [response async for response in results] + + assert collection.last_search_score_threshold == threshold + assert collection.last_search_filter == filter_ + assert collection.last_search_top == 4 + assert collection.last_search_skip == 2 + assert [response["record"].id for response in responses] == ["low", "equal", "high", "scoreless"] + assert [response["score"] for response in responses] == [-0.2, threshold, 0.9, None] + + +async def test_connector_can_enforce_provider_threshold_before_core_deserialization() -> None: + class ProviderThresholdCollection(MockCollection): + async def _inner_search( + self, + *, + search_type: SearchType, + filter: Filter | FilterGroup | None = None, + values: Any | None = None, + vector: Sequence[float | int] | bytes | bytearray | None = None, + top: int = 3, + skip: int = 0, + include_vectors: bool = False, + vector_property_name: str | None = None, + additional_property_name: str | None = None, + score_threshold: float | None = None, + operation_options: Mapping[str, Any] | None = None, + ) -> SearchResults[Any]: + self.last_search_score_threshold = score_threshold + assert self.raw_search_results is not None + results = SearchResults(self.raw_search_results) + return SearchResults([ + result async for result in results if score_threshold is None or result["score"] <= score_threshold + ]) + + collection = ProviderThresholdCollection() collection.definition = VectorStoreCollectionDefinition( [ - VectorStoreField("key", name="id", type_="str"), - VectorStoreField("vector", name="vector", type_="float", dimensions=2), + VectorStoreField("key", name="id", storage_name="record_id"), + VectorStoreField("data", name="text", storage_name="body"), + VectorStoreField( + "vector", + name="vector", + dimensions=2, + distance_function="provider.custom_distance", + ), ], collection_name="records", ) + collection.raw_search_results = [ + {"record": {"record_id": "rejected", "body": "far"}, "score": 0.9}, + {"record": {"record_id": "accepted", "body": "near"}, "score": 0.2}, + ] + + results = await collection.search(vector=[1.0, 0.0], score_threshold=0.5) - with pytest.raises(ValueError, match="explicit distance"): + assert collection.last_search_score_threshold == 0.5 + assert [(result["record"].id, result["score"]) async for result in results] == [("accepted", 0.2)] + + +async def test_connector_can_reject_unsupported_threshold_execution() -> None: + collection = MockCollection() + collection.search_error = NotImplementedError("Connector does not support score thresholds.") + + with pytest.raises(NotImplementedError, match="Connector does not support score thresholds"): await collection.search(vector=[1.0, 0.0], score_threshold=0.5) @@ -790,30 +1274,25 @@ async def get_embeddings( await collection.search("query") -def test_vector_search_builds_connector_filter() -> None: - collection = MockCollection() - - assert collection._build_filter("lambda record: record.category == 'travel'") == "record.category == 'travel'" - assert collection._build_filter([ - "lambda record: record.category == 'travel'", - "lambda record: record.id != 'ignored'", - ]) == ["record.category == 'travel'", "record.id != 'ignored'"] - - -async def test_vector_search_passes_translated_filter_to_connector() -> None: +async def test_vector_search_passes_filter_to_connector() -> None: collection = MockCollection() + search_filter = FilterGroup("and", (Filter("text", "eq", "travel"), Filter("id", "ne", "ignored"))) await collection.search( "query", - filter="lambda record: record.category == 'travel'", + filter=search_filter, ) - assert collection.last_search_filter == "record.category == 'travel'" + assert collection.last_search_filter == search_filter + assert collection.last_search_filter is not search_filter -def test_vector_search_rejects_filter_without_lambda() -> None: - with pytest.raises(ValueError, match="No lambda"): - MockCollection()._build_filter("record.category == 'travel'") +async def test_vector_search_rejects_invalid_or_unresolved_filters() -> None: + collection = MockCollection() + with pytest.raises(ValueError, match="not part of the vector store definition"): + await collection.search("query", filter=Filter("missing", "eq", "value")) + with pytest.raises(ValueError, match="must be resolved"): + await collection.search("query", filter=Filter("text", "eq", Param("text", str))) async def test_create_search_tool_returns_mapped_results() -> None: @@ -838,62 +1317,54 @@ async def test_create_search_tool_returns_mapped_results() -> None: async def test_create_search_tool_supports_declared_filter_parameters() -> None: collection = MockCollection() collection.records["one"] = {"record_id": "one", "body": "first", "vector": [1.0, 0.0]} + category = Param( + "category", + Literal["travel", "work"], + required=True, + description="The category to match.", + ) tool = create_vector_search_tool( collection, - parameters={ - "type": "object", - "properties": { - "query": {"type": "string", "description": "The search query."}, - "category": {"type": "string", "description": "The category to match."}, - "top": { - "type": "integer", - "description": "The maximum number of results.", - "maximum": 5, - }, - "skip": { - "type": "integer", - "description": "The number of results to skip.", - "maximum": 10, - }, - }, - "required": ["query", "category"], - "additionalProperties": False, - }, + filter=Filter("text", "eq", category), + top=Param("top", int, default=1, minimum=1, maximum=5), + skip=Param("skip", int, default=0, minimum=0, maximum=10), ) await tool(query="first", category="travel", top=1, skip=2) assert set(tool.parameters()["properties"]) == {"query", "category", "top", "skip"} - assert collection.last_search_filter == "record.category == 'travel'" + assert tool.parameters()["additionalProperties"] is False + assert tool.parameters()["properties"]["category"]["enum"] == ["travel", "work"] + assert tool.parameters()["properties"]["top"]["minimum"] == 1 + assert tool.parameters()["properties"]["top"]["maximum"] == 5 + assert collection.last_search_filter == Filter("text", "eq", "travel") assert collection.last_search_top == 1 assert collection.last_search_skip == 2 -def test_create_search_tool_validates_custom_schema() -> None: +def test_create_search_tool_validates_params() -> None: collection = MockCollection() - with pytest.raises(ValueError, match="required string"): - create_vector_search_tool(collection, parameters={"type": "object", "properties": {}}) - with pytest.raises(ValueError, match="required string"): + with pytest.raises(ValueError, match="minimum"): create_vector_search_tool( collection, - parameters={ - "type": "object", - "properties": {"query": {"type": "integer"}}, - "required": ["query"], - }, + top=Param("top", int, default=1, maximum=5), ) - with pytest.raises(ValueError, match="declare an integer maximum"): + with pytest.raises(ValueError, match="finite integer maximum"): create_vector_search_tool( collection, - parameters={ - "type": "object", - "properties": { - "query": {"type": "string"}, - "top": {"type": "integer"}, - }, - "required": ["query"], - }, + top=Param("top", int, default=1, minimum=1), + ) + with pytest.raises(ValueError, match="conflicting declarations"): + create_vector_search_tool( + collection, + filter=FilterGroup( + "and", + ( + Filter("id", "eq", Param("record_id", str)), + Filter("text", "eq", Param("record_id", int)), + ), + ), ) @@ -904,25 +1375,22 @@ async def test_create_search_tool_enforces_paging_limits_and_result_cap() -> Non ] tool = create_vector_search_tool( collection, - top=2, - parameters={ - "type": "object", - "properties": { - "query": {"type": "string"}, - "top": {"type": "integer", "maximum": 2}, - "skip": {"type": "integer", "maximum": 4}, - }, - "required": ["query"], - }, + top=Param("top", int, default=2, minimum=1, maximum=2), + skip=Param("skip", int, default=0, minimum=0, maximum=4), ) results = await tool(query="records", top=2, skip=4) assert len(results) == 2 - with pytest.raises(ValueError, match="top must not exceed"): + with pytest.raises(ValueError, match="must be at most 2"): await tool(query="records", top=3) - with pytest.raises(ValueError, match="skip must not exceed"): + with pytest.raises(ValueError, match="must be at most 4"): await tool(query="records", skip=5) + with pytest.raises(TypeError, match="does not match"): + await create_vector_search_tool( + collection, + filter=Filter("text", "eq", Param("category", Literal["travel", "work"], required=True)), + )(query="records", category="other") async def test_create_search_tool_supports_multimodal_results() -> None: @@ -1301,43 +1769,684 @@ async def test_scoreless_results_remain_when_threshold_cannot_be_applied() -> No assert responses[0]["score"] is None -def test_filter_parser_and_default_mapper_edge_paths() -> None: - collection = MockCollection() - assert collection._build_filter(lambda record: record.id == "one") == "record.id == 'one'" - with pytest.raises(ValueError, match="Unable to parse"): - collection._build_filter("lambda record:") - +@pytest.mark.parametrize( + ("left", "right", "expected"), + [ + (True, 1, False), + (False, 0, False), + (True, True, True), + (False, False, True), + (1, 1.0, True), + ([1], [True], False), + ([0], [False], False), + ((1,), (True,), False), + ([[1]], [[True]], False), + ([(1, [0])], [(True, [False])], False), + ([True, [False]], [True, [False]], True), + ([1, [0]], [1.0, [0.0]], True), + ((1, (0,)), (1.0, (0.0,)), True), + ([1], (1,), False), + ([], (), False), + ([1, 2], [2, 1], False), + ([1], [1, 2], False), + ([], [], True), + ({"value": [1]}, {"value": [True]}, False), + ([{"value": (0,)}], [{"value": (False,)}], False), + ({"value": [1]}, {"value": [1.0]}, True), + ({"one": True, "two": False}, {"two": False, "one": True}, True), + ({"one": True}, {"two": True}, False), + ({"value": True}, {}, False), + ("1", 1, False), + ("text", "text", True), + (b"\x01", b"\x01", True), + (None, None, True), + ], +) +def test_filter_values_equal_preserves_nested_types(left: Any, right: Any, expected: bool) -> None: + assert filter_values_equal(left, right) is expected + assert filter_values_equal(right, left) is expected + + +def test_filter_model_accepts_namespaced_provider_operators() -> None: + assert Filter("text", "azure_ai_search.full_text", "query").operator == "azure_ai_search.full_text" + with pytest.raises(ValueError, match="namespaced"): + Filter("text", "full_text", "query") + + +def test_filter_model_rejects_invalid_shapes() -> None: + with pytest.raises(ValueError, match="requires a value"): + Filter("text", "eq") + with pytest.raises(TypeError, match="sequence value"): + Filter("text", "in", "value") + with pytest.raises(ValueError, match="two boundary"): + Filter("text", "between", (1,)) + with pytest.raises(ValueError, match="exactly one"): + FilterGroup("not", (Filter("text", "eq", "one"), Filter("text", "eq", "two"))) + with pytest.raises(ValueError, match="Invalid filter field name"): + Filter("__class__", "eq", "unsafe") + + cyclic: list[Any] = [] + cyclic.append(cyclic) + with pytest.raises(ValueError, match="cycles"): + Filter("text", "in", cyclic) + with pytest.raises(ValueError, match="entire Filter value"): + Filter("text", "in", ("fixed", Param("dynamic", str))) + + +def test_param_generates_native_schema_and_validates_constraints() -> None: + param = Param( + "category", + Literal["travel", "work"], + description="The category.", + required=True, + ) + tool = create_vector_search_tool(MockCollection(), filter=Filter("text", "eq", param)) -async def test_search_tool_filter_mapper_edge_paths() -> None: - collection = MockCollection() - parameters = { + assert tool.parameters() == { "type": "object", "properties": { - "query": {"type": "string"}, - "category": {"type": "string"}, + "query": {"type": "string", "description": "The query to search for."}, + "category": { + "enum": ["travel", "work"], + "type": "string", + "description": "The category.", + }, }, "required": ["query", "category"], + "additionalProperties": False, } + + with pytest.raises(TypeError, match="cannot be represented"): + Param("unsupported", object) + with pytest.raises(ValueError, match="minimum cannot exceed"): + Param("invalid_range", float, minimum=2, maximum=1) + + optional_code = Param("code", str | None, max_length=8) + schema = create_vector_search_tool(MockCollection(), filter=Filter("text", "eq", optional_code)).parameters() + assert {"type": "string", "maxLength": 8} in schema["properties"]["code"]["anyOf"] + with pytest.raises(TypeError, match="homogeneous"): + Param("pair", tuple[int, str]) + with pytest.raises(TypeError, match="JSON-compatible"): + Param("binary_literal", Literal[b"value"]) + non_finite_literal = cast(Any, Literal)[float("inf")] + with pytest.raises(ValueError, match="Literal numbers must be finite"): + Param("non_finite_literal", non_finite_literal) + with pytest.raises(TypeError, match="mapping keys must use the str type"): + Param("integer_keys", dict[int, str]) + with pytest.raises(ValueError, match="numeric constraints require numeric"): + Param("category_with_minimum", str, minimum=1) + with pytest.raises(ValueError, match="numeric constraints require numeric"): + Param("mixed_with_minimum", int | str, minimum=1) + + tuple_param = Param("numbers", tuple[int, ...]) + tuple_schema = create_vector_search_tool( + MockCollection(), + filter=Filter("text", "provider.numbers", tuple_param), + ).parameters() + assert tuple_schema["properties"]["numbers"] == {"type": "array", "items": {"type": "integer"}} + + nullable_number = Param("score", float | None, minimum=0, maximum=1) + nullable_number_schema = create_vector_search_tool( + MockCollection(), + filter=Filter("text", "provider.score", nullable_number), + ).parameters() + assert nullable_number_schema["properties"]["score"] == { + "anyOf": [{"type": "number"}, {"type": "null"}], + "minimum": 0, + "maximum": 1, + } + assert Param("level", Literal[1, None], minimum=1).minimum == 1 + + source_default = ["travel"] + frozen_default = Param("categories", list[str], default=source_default) + source_default.append("work") + returned_default = frozen_default.default + returned_default.append("personal") + assert frozen_default.default == ["travel"] + + +def test_filter_and_param_constructor_boundaries() -> None: + with pytest.raises(ValueError, match="cannot be empty"): + Param("", str) + with pytest.raises(ValueError, match="reserved"): + Param("query", str) + with pytest.raises(ValueError, match="required parameter cannot declare a default"): + Param("required", str, required=True, default="value") + with pytest.raises(ValueError, match="finite number"): + Param("minimum", float, minimum=float("inf")) + with pytest.raises(ValueError, match="non-negative integer"): + Param("length", str, min_length=-1) + with pytest.raises(ValueError, match="cannot exceed max_length"): + Param("length", str, min_length=2, max_length=1) + with pytest.raises(TypeError, match="field_name must be a string"): + Filter(cast(Any, 1), "eq", "value") + with pytest.raises(TypeError, match="operator must be a string"): + Filter("text", cast(Any, 1), "value") + with pytest.raises(ValueError, match="does not accept a value"): + Filter("text", "is_null", "value") + with pytest.raises(TypeError, match="operator must be a string"): + FilterGroup(cast(Any, 1), (Filter("text", "eq", "value"),)) + with pytest.raises(ValueError, match="Unknown filter group"): + FilterGroup(cast(Any, "xor"), (Filter("text", "eq", "value"),)) + with pytest.raises(TypeError, match="must be a sequence"): + FilterGroup("and", cast(Any, "not-a-sequence")) + with pytest.raises(ValueError, match="at least one"): + FilterGroup("and", ()) + with pytest.raises(TypeError, match="Filter or FilterGroup"): + FilterGroup("and", cast(Any, (object(),))) + + +async def test_param_native_schema_and_runtime_validation_for_container_types() -> None: + items = Param("items", list[str], min_length=1, max_length=3) + flags = Param("flags", dict[str, int], max_length=2) + enabled = Param("enabled", bool, required=True) + collection = MockCollection() tool = create_vector_search_tool( collection, - parameters=parameters, - filter=["lambda record: record.id != 'ignored'"], + filter=FilterGroup( + "and", + ( + Filter("text", "eq", items), + Filter("text", "eq", flags), + Filter("text", "eq", enabled), + ), + ), ) - await tool(query="query", category="travel") - assert collection.last_search_filter == ["record.id != 'ignored'", "record.category == 'travel'"] - invalid_tool = create_vector_search_tool( + schema = tool.parameters() + assert schema["properties"]["items"] == { + "type": "array", + "items": {"type": "string"}, + "minItems": 1, + "maxItems": 3, + } + assert schema["properties"]["flags"] == { + "type": "object", + "additionalProperties": {"type": "integer"}, + "maxProperties": 2, + } + assert schema["required"] == ["query", "enabled"] + + await tool(query="query", items=["a"], flags={"one": 1}, enabled=True) + with pytest.raises(TypeError, match="does not match"): + await tool(query="query", items=[1], enabled=True) + with pytest.raises(ValueError, match="longer than 3"): + await tool(query="query", items=["a", "b", "c", "d"], enabled=True) + with pytest.raises(TypeError, match="Missing required"): + await tool(query="query") + + +async def test_filter_resource_validation_boundaries() -> None: + collection = MockCollection() + invalid_mapping = Filter("text", "provider.native", cast(Any, {1: "value"})) + with pytest.raises(TypeError, match="mapping keys must be strings"): + await collection.search("query", filter=invalid_mapping) + + mutable_value: list[Any] = ["fixed"] + mutated_filter = Filter("text", "in", mutable_value) + mutable_value.append(Param("dynamic", str)) + with pytest.raises(ValueError, match="entire Filter value"): + await collection.search("query", filter=mutated_filter) + + too_many = FilterGroup("and", (Filter("text", "eq", "value"),)) + too_many.filters = tuple(Filter("text", "eq", str(index)) for index in range(65)) + with pytest.raises(ValueError, match="more than 64 nodes"): + await collection.search("query", filter=too_many) + + cyclic_group = FilterGroup("and", (Filter("text", "eq", "value"),)) + cyclic_group.filters = (cyclic_group,) + with pytest.raises(ValueError, match="cannot contain cycles"): + await collection.search("query", filter=cyclic_group) + + cyclic_expression_value = Filter("text", "provider.nested", "value") + cyclic_expression_value.value = cyclic_expression_value + with pytest.raises(ValueError, match="cannot contain cycles"): + await collection.search("query", filter=cyclic_expression_value) + + nested: Filter | FilterGroup = Filter("relative_name", "eq", "value") + assert Filter("text", "provider.nested", nested).value is nested + + +@pytest.mark.parametrize("operation", ["get", "search", "tool"]) +async def test_filter_limits_are_checked_before_snapshot_copy(operation: str) -> None: + collection = MockCollection() + expression: Filter | FilterGroup = Filter("text", "eq", "value") + for _ in range(1100): + expression = FilterGroup("and", (expression,)) + + with patch("agent_framework._vector_filters.deepcopy") as copy: + with pytest.raises(ValueError, match="depth of 8"): + if operation == "get": + await collection.get(filter=expression) + elif operation == "search": + await collection.search("query", filter=expression) + else: + create_vector_search_tool(collection, filter=expression) + copy.assert_not_called() + + +def test_filter_group_constructor_bounds_child_copying() -> None: + child = Filter("text", "eq", "value") + + class LargeFilterSequence(Sequence[Filter]): + def __init__(self) -> None: + self.visited = 0 + + def __len__(self) -> int: + return 1_000_000 + + def __getitem__(self, index: Any) -> Any: + self.visited += 1 + assert self.visited <= 64 + return child + + filters = LargeFilterSequence() + with pytest.raises(ValueError, match="more than 64 nodes"): + FilterGroup("and", filters) + assert filters.visited == 64 + assert len(FilterGroup("and", [child] * 63).filters) == 63 + + +async def test_filter_value_traversal_stops_at_node_budget_before_copy() -> None: + class LargeSequence(Sequence[int]): + def __init__(self) -> None: + self.visited = 0 + + def __len__(self) -> int: + return 1_000_000 + + def __getitem__(self, index: Any) -> Any: + self.visited += 1 + assert self.visited <= 64 + return index + + def __deepcopy__(self, memo: dict[int, Any]) -> Any: + raise AssertionError("Oversized values must not be copied.") + + values = LargeSequence() + expression = Filter("text", "provider.values", []) + expression.value = values + with pytest.raises(ValueError, match="more than 64 nodes"): + await MockCollection().search("query", filter=expression) + assert values.visited == 64 + + +async def test_filter_collection_values_keep_existing_node_budget_and_snapshot_semantics() -> None: + values = set(range(63)) + expression = Filter("text", "provider.native", values) + collection = MockCollection() + tool = create_vector_search_tool(collection, filter=expression) + values.add(63) + + await tool(query="query") + assert collection.last_search_filter == Filter("text", "provider.native", set(range(63))) + with pytest.raises(ValueError, match="more than 64 nodes"): + await collection.search("query", filter=expression) + + +@pytest.mark.parametrize( + "wrap", + [ + lambda param: {param}, + lambda param: frozenset((param,)), + lambda param: {"key": param}.values(), + lambda param: {param: "value"}, + lambda param: {(param,): "value"}, + lambda param: {"nested": {param}}, + ], +) +def test_filters_reject_params_in_collection_values_and_mapping_keys(wrap: Callable[[Param], Any]) -> None: + with pytest.raises(ValueError, match="entire Filter value"): + Filter("text", "provider.native", wrap(Param("hidden", str))) + + +def test_filter_constructor_bounds_nested_value_inspection() -> None: + value: Any = "value" + for _ in range(1100): + value = [value] + with pytest.raises(ValueError, match="depth of 8"): + Filter("text", "provider.nested", value) + + +def test_param_default_is_bounded_before_copying() -> None: + with patch("agent_framework._vector_filters.deepcopy") as copy: + with pytest.raises(ValueError, match="more than 256 nodes"): + Param("values", list[int], default=list(range(300))) + copy.assert_not_called() + + +@pytest.mark.parametrize("operator", ["starts_with", "ends_with", "contains_text"]) +@pytest.mark.parametrize("value", [1, False, []]) +def test_string_filter_operators_require_string_operands(operator: str, value: Any) -> None: + with pytest.raises(TypeError, match="must be a string"): + Filter("text", operator, value) + + +@pytest.mark.parametrize("case", ["nested_param", "string_operand", "combined_budget"]) +async def test_definitionless_search_still_validates_portable_filter_structure(case: str) -> None: + class SearchOnly: + async def search(self, values: Any, **kwargs: Any) -> SearchResults[SearchResponse[Record]]: + raise AssertionError("Invalid filter reached the connector.") + + search = cast(SupportsVectorSearch[Record], SearchOnly()) + if case == "nested_param": + expression = Filter("text", "provider.native", set()) + expression.value.add(Param("hidden", str)) + with pytest.raises(ValueError, match="entire Filter value"): + create_vector_search_tool(search, filter=expression) + elif case == "string_operand": + tool = create_vector_search_tool(search, filter=Filter("text", "contains_text", Param("text", int))) + with pytest.raises(TypeError, match="must be a string"): + await tool(query="query", text=1) + else: + tool = create_vector_search_tool( + search, + filter=FilterGroup( + "and", + ( + Filter("first", "in", Param("first", list[str])), + Filter("second", "in", Param("second", list[str])), + ), + ), + ) + with pytest.raises(ValueError, match="more than 64 nodes"): + await tool(query="query", first=["one"] * 40, second=["two"] * 40) + + +async def test_search_tool_param_edge_paths() -> None: + collection = MockCollection() + category = Param("category", str) + tool = create_vector_search_tool( collection, - parameters={ - "type": "object", - "properties": {"query": {"type": "string"}, "bad-name": {"type": "string"}}, - "required": ["query"], - }, + filter=FilterGroup( + "and", + ( + Filter("id", "ne", "ignored"), + Filter("text", "eq", category), + ), + ), ) - with pytest.raises(ValueError, match="cannot be mapped"): - await invalid_tool(query="query", **{"bad-name": "value"}) + await tool(query="query", category="travel") + assert collection.last_search_filter == FilterGroup( + "and", + ( + Filter("id", "ne", "ignored"), + Filter("text", "eq", "travel"), + ), + ) + + await tool(query="query") + assert collection.last_search_filter == FilterGroup("and", (Filter("id", "ne", "ignored"),)) + + with pytest.raises(ValueError, match="Invalid parameter name"): + Param("bad-name", str) with pytest.raises(TypeError, match="'query'.*string"): await cast(Any, create_vector_search_tool(collection))(query=1) + with pytest.raises(TypeError, match="Unexpected argument"): + await tool(query="query", undeclared="value") + with pytest.raises(ValueError, match="must be finite"): + await create_vector_search_tool( + collection, + filter=Filter("text", "eq", Param("score", float)), + )(query="query", score=float("nan")) + with pytest.raises(TypeError, match="does not match"): + await create_vector_search_tool( + collection, + filter=Filter("text", "eq", Param("level", Literal[1, 2], required=True)), + )(query="query", level=True) + + nullable_tool = create_vector_search_tool( + collection, + filter=Filter("text", "provider.nullable", Param("nullable", str | None)), + ) + await nullable_tool(query="query", nullable=None) + assert collection.last_search_filter == Filter("text", "provider.nullable", None) + + nested_values_tool = create_vector_search_tool( + collection, + filter=Filter("text", "provider.values", Param("values", list[list[int]])), + ) + with pytest.raises(ValueError, match="more than 256 nodes"): + await nested_values_tool(query="query", values=[[index for index in range(20)] for _ in range(20)]) + + with pytest.raises(ValueError, match="must be finite"): + await collection.search("query", filter=Filter("text", "provider.number", Decimal("NaN"))) + + +@pytest.mark.parametrize("value_type", [str, int, list[str | None]]) +def test_param_omit_if_none_requires_nullable_type(value_type: Any) -> None: + with pytest.raises(ValueError, match="value type that accepts None"): + Param("value", value_type, default=None, omit_if_none=True) + + +def test_param_omit_if_none_requires_explicit_null_default() -> None: + with pytest.raises(ValueError, match="explicit default=None"): + Param("text", str | None, omit_if_none=True) + with pytest.raises(ValueError, match="explicit default=None"): + Param("text", str | None, default="hotel", omit_if_none=True) + with pytest.raises(ValueError, match="required parameter cannot declare a default"): + Param("text", str | None, required=True, default=None, omit_if_none=True) + with pytest.raises(TypeError, match="omit_if_none must be a boolean"): + Param("text", str | None, default=None, omit_if_none=cast(Any, "true")) + + +async def test_param_omit_if_none_exposes_nullable_schema_and_validates_values() -> None: + param = Param("text", str | None, default=None, omit_if_none=True, max_length=8) + tool = create_vector_search_tool(MockCollection(), filter=Filter("text", "contains_text", param)) + + assert param.omit_if_none + assert param.has_default + assert param.default is None + assert tool.parameters()["properties"]["text"] == { + "anyOf": [{"type": "string", "maxLength": 8}, {"type": "null"}], + "default": None, + } + assert tool.parameters()["required"] == ["query"] + with pytest.raises(TypeError, match="does not match"): + await tool(query="query", text=123) + with pytest.raises(ValueError, match="longer than 8"): + await tool(query="query", text="too long a value") + + +@pytest.mark.parametrize( + ("operator", "value_type", "supplied"), + [ + ("contains_text", str | None, "hotel"), + ("contains_text", str | None, ""), + ("contains_text", str | None, "*"), + ("eq", Literal["hotel", None], "hotel"), + ("gte", float | None, 0), + ("in", list[str] | None, []), + ("provider.enabled", bool | None, False), + ], +) +async def test_search_tool_omits_only_null_param_values(operator: str, value_type: Any, supplied: Any) -> None: + collection = MockCollection() + param = Param("value", value_type, default=None, omit_if_none=True) + search_filter = Filter("text", operator, param) + tool = create_vector_search_tool(collection, filter=search_filter) + + await tool(query="query") + assert collection.last_search_filter is None + await tool(query="query", value=None) + assert collection.last_search_filter is None + await tool(query="query", value=supplied) + assert collection.last_search_filter == Filter("text", operator, supplied) + await tool(query="query", value=None) + assert collection.last_search_filter is None + assert search_filter.value is param + + +@pytest.mark.parametrize("operator", ["and", "or"]) +async def test_search_tool_null_omission_keeps_remaining_group_children(operator: Literal["and", "or"]) -> None: + collection = MockCollection() + param = Param("text", str | None, default=None, omit_if_none=True) + fixed = Filter("id", "eq", "one") + tool = create_vector_search_tool( + collection, + filter=FilterGroup(operator, (fixed, Filter("text", "contains_text", param))), + ) + + await tool(query="query") + assert collection.last_search_filter == FilterGroup(operator, (fixed,)) + await tool(query="query", text=None) + assert collection.last_search_filter == FilterGroup(operator, (fixed,)) + await tool(query="query", text="hotel") + assert collection.last_search_filter == FilterGroup(operator, (fixed, Filter("text", "contains_text", "hotel"))) + + +@pytest.mark.parametrize("operator", ["and", "or", "not"]) +@pytest.mark.parametrize("nested", [False, True]) +async def test_search_tool_null_omission_prunes_empty_groups( + operator: Literal["and", "or", "not"], nested: bool +) -> None: + collection = MockCollection() + param = Param("text", str | None, default=None, omit_if_none=True) + fixed = Filter("id", "eq", "one") + group = FilterGroup(operator, (Filter("text", "contains_text", param),)) + tool = create_vector_search_tool( + collection, + filter=FilterGroup("and", (fixed, group)) if nested else group, + ) + expected = FilterGroup("and", (fixed,)) if nested else None + + await tool(query="query") + assert collection.last_search_filter == expected + await tool(query="query", text=None) + assert collection.last_search_filter == expected + await tool(query="query", text="hotel") + resolved_group = FilterGroup(operator, (Filter("text", "contains_text", "hotel"),)) + assert collection.last_search_filter == (FilterGroup("and", (fixed, resolved_group)) if nested else resolved_group) + + +async def test_search_tool_null_omission_is_opt_in_per_parameter() -> None: + collection = MockCollection() + omitted = Param("text", str | None, default=None, omit_if_none=True) + retained = Param("native", str | None, default=None) + tool = create_vector_search_tool( + collection, + filter=FilterGroup( + "and", + (Filter("text", "contains_text", omitted), Filter("text", "provider.nullable", retained)), + ), + ) + + assert not retained.omit_if_none + await tool(query="query", text=None, native=None) + assert collection.last_search_filter == FilterGroup("and", (Filter("text", "provider.nullable", None),)) + await tool(query="query") + assert collection.last_search_filter == FilterGroup("and", (Filter("text", "provider.nullable", None),)) + + strict_tool = create_vector_search_tool(collection, filter=Filter("text", "contains_text", retained)) + with pytest.raises(ValueError, match="requires a value"): + await strict_tool(query="query", native=None) + with pytest.raises(ValueError, match="requires a value"): + await strict_tool(query="query") + with pytest.raises(ValueError, match="requires a value"): + Filter("text", "contains_text", None) + + +def test_search_tool_rejects_conflicting_null_omission_policies() -> None: + with pytest.raises(ValueError, match="conflicting declarations"): + create_vector_search_tool( + MockCollection(), + filter=FilterGroup( + "and", + ( + Filter("text", "contains_text", Param("text", str | None, default=None, omit_if_none=True)), + Filter("text", "provider.nullable", Param("text", str | None, default=None)), + ), + ), + ) + + +@pytest.mark.parametrize("option_name", ["top", "skip"]) +def test_search_tool_rejects_null_omission_for_paging(option_name: Literal["top", "skip"]) -> None: + param = Param(option_name, int | None, default=None, omit_if_none=True, minimum=1, maximum=10) + with pytest.raises(ValueError, match=f"{option_name} Param does not support omit_if_none"): + create_vector_search_tool( + MockCollection(), + top=param if option_name == "top" else 5, + skip=param if option_name == "skip" else 0, + ) + + +async def test_search_operations_snapshot_mutable_filters() -> None: + collection = MockCollection() + collection.mutate_search_filter = True + values = ["one"] + search_filter = Filter("id", "in", values) + + await collection.search("query", filter=search_filter) + + assert values == ["one"] + assert search_filter.value == ["one"] + + tool = create_vector_search_tool(collection, filter=search_filter) + search_filter.field_name = "text" + values.append("two") + await tool(query="query") + + assert isinstance(collection.last_search_filter, Filter) + assert collection.last_search_filter.field_name == "id" + assert collection.last_search_filter.value == ["one", "connector mutation"] + + +@pytest.mark.parametrize("supply_argument", [False, True]) +async def test_search_tool_copies_mutable_param_values_per_invocation(supply_argument: bool) -> None: + class MutatingSearch: + def __init__(self) -> None: + self.seen_values: list[list[str]] = [] + + async def search( + self, + values: Any, + *, + search_type: SearchType = "vector", + filter: Filter | FilterGroup | None = None, + top: int = 3, + skip: int = 0, + **kwargs: Any, + ) -> SearchResults[SearchResponse[Record]]: + assert isinstance(filter, Filter) + received = cast(list[str], filter.value) + self.seen_values.append(list(received)) + received.append("mutated") + return SearchResults([]) + + search = MutatingSearch() + tool = create_vector_search_tool( + cast(SupportsVectorSearch[Record], search), + filter=Filter("tenant_id", "in", Param("tenant_ids", list[str], default=["acme"])), + ) + + supplied = ["acme"] + arguments = {"tenant_ids": supplied} if supply_argument else {} + await tool(query="first", **arguments) + await tool(query="second", **arguments) + + assert search.seen_values == [["acme"], ["acme"]] + assert supplied == ["acme"] + + +async def test_search_tool_deep_copies_supplied_mapping_values() -> None: + class MutatingSearch: + async def search( + self, values: Any, *, filter: Filter | FilterGroup | None = None, **kwargs: Any + ) -> SearchResults[SearchResponse[Record]]: + assert isinstance(filter, Filter) + filter.value["tags"].append("mutated") + return SearchResults([]) + + tool = create_vector_search_tool( + cast(SupportsVectorSearch[Record], MutatingSearch()), + filter=Filter("metadata", "provider.native", Param("metadata", dict[str, list[str]])), + ) + supplied = {"tags": ["original"]} + + await tool(query="query", metadata=supplied) + + assert supplied == {"tags": ["original"]} async def test_runtime_operations_mark_vector_store_feature_usage() -> None: diff --git a/python/samples/02-agents/vector_stores/README.md b/python/samples/02-agents/vector_stores/README.md index 92227f0b1f0..d59041810fd 100644 --- a/python/samples/02-agents/vector_stores/README.md +++ b/python/samples/02-agents/vector_stores/README.md @@ -8,12 +8,17 @@ explicit definition and codecs. Dictionaries use a collection-specific definition; DataFrames and other containers can convert to row dictionaries before calling the batch API. -No database or credentials are needed for these examples. +No database is needed for these 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. | File | Demonstrates | |------|--------------| | [`vector_store_models.py`](vector_store_models.py) | Choosing among owned models, third-party model registration, and loose dictionary definitions. | | [`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. | The first section shows the two equivalent custom-codec registration forms. `@vectorstoremodel` derives the definition from annotations and registers it; @@ -53,4 +58,6 @@ Run the sample from the `python` directory: ```bash 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 ``` diff --git a/python/samples/02-agents/vector_stores/in_memory_filters.py b/python/samples/02-agents/vector_stores/in_memory_filters.py new file mode 100644 index 00000000000..193584cc41c --- /dev/null +++ b/python/samples/02-agents/vector_stores/in_memory_filters.py @@ -0,0 +1,73 @@ +# Copyright (c) Microsoft. All rights reserved. + +import asyncio +from dataclasses import dataclass +from typing import Annotated + +from agent_framework import Filter, FilterGroup, InMemoryCollection, VectorStoreField, vectorstoremodel + +""" +This sample demonstrates direct vector searches with portable, data-only filters. + +The in-memory store is process-local and uses a linear scan. It is intended for +tests and development, not as a production vector database. +""" + + +@vectorstoremodel(collection_name="hotels") +@dataclass +class Hotel: + hotel_id: Annotated[str, VectorStoreField("key")] + name: Annotated[str, VectorStoreField("data")] + city: Annotated[str, VectorStoreField("data")] + rating: Annotated[float, VectorStoreField("data")] + amenities: Annotated[list[str], VectorStoreField("data")] + vector: Annotated[ + list[float] | None, + VectorStoreField("vector", dimensions=2, distance_function="cosine_similarity"), + ] = None + + +async def main() -> None: + """Store precomputed vectors and search them with direct filters.""" + collection: InMemoryCollection[str, Hotel] = InMemoryCollection(Hotel) + await collection.ensure_collection_exists() + + # 1. The sample already has vectors, so generation is disabled explicitly. + await collection.upsert( + [ + Hotel("hotel-1", "Harbor View", "Lisbon", 4.8, ["wifi", "pool"], [1.0, 0.1]), + Hotel("hotel-2", "Old Town Rooms", "Lisbon", 4.1, ["wifi"], [0.8, 0.2]), + Hotel("hotel-3", "City Center", "Seattle", 4.7, ["wifi", "gym"], [0.1, 1.0]), + ], + generate_vectors=False, + ) + + # 2. Filter values are ordinary data. No Python source is parsed or executed. + search_filter = FilterGroup( + "and", + ( + Filter("city", "eq", "Lisbon"), + Filter("rating", "between", (4.5, 5.0)), + Filter("amenities", "contains", "pool"), + ), + ) + results = await collection.search( + vector=[1.0, 0.0], + filter=search_filter, + top=5, + ) + + # 3. Search results are consumed asynchronously. + async for result in results: + print(f"{result['record'].name}: {result['score']:.3f}") + + +if __name__ == "__main__": + asyncio.run(main()) + + +""" +Sample output: +Harbor View: 0.995 +""" diff --git a/python/samples/02-agents/vector_stores/in_memory_search_tool.py b/python/samples/02-agents/vector_stores/in_memory_search_tool.py new file mode 100644 index 00000000000..c8072f88cc9 --- /dev/null +++ b/python/samples/02-agents/vector_stores/in_memory_search_tool.py @@ -0,0 +1,189 @@ +# Copyright (c) Microsoft. All rights reserved. + +import asyncio +import json +import os +from typing import Annotated, Any, Literal +from urllib.request import urlopen + +from agent_framework import ( + Agent, + Filter, + FilterGroup, + InMemoryCollection, + Param, + VectorStoreField, + create_vector_search_tool, + vectorstoremodel, +) +from agent_framework.openai import OpenAIChatClient, OpenAIEmbeddingClient +from dotenv import load_dotenv +from pydantic import BaseModel, ConfigDict, Field + +load_dotenv() + +""" +This sample demonstrates an agent using model-set filter values in a search tool. + +This uses the hotel model and dataset based on the Azure AI Search vector sample dataset. +`Param` values define the exact filter parameters exposed to the model. Their native Python types and +constraints become JSON Schema without requiring another Pydantic parameter model. + +Set `OPENAI_API_KEY` before starting the sample. +""" + +HOTELS_URL = "https://raw.githubusercontent.com/Azure/azure-search-vector-samples/refs/heads/main/data/hotels.json" + + +def to_pascal(name: str) -> str: + """Convert a Python field name to the dataset's PascalCase naming.""" + return "".join(part.capitalize() for part in name.split("_")) + + +class Room(BaseModel): + """Describe one hotel room type.""" + + type: str + description: str + description_fr: str = Field(alias="Description_fr") + base_rate: float + bed_options: str + sleeps_count: int + smoking_allowed: bool + tags: list[str] + + model_config = ConfigDict(alias_generator=to_pascal, populate_by_name=True, extra="ignore") + + +class Address(BaseModel): + """Describe a hotel address.""" + + street_address: str + city: str | None + state_province: str | None + postal_code: str | None + country: str | None + + model_config = ConfigDict(alias_generator=to_pascal, populate_by_name=True, extra="ignore") + + +@vectorstoremodel(collection_name="hotels") +class Hotel(BaseModel): + """Represent one hotel from the Azure AI Search sample dataset.""" + + hotel_id: Annotated[str, VectorStoreField("key")] + hotel_name: Annotated[str | None, VectorStoreField("data")] = None + description: Annotated[str, VectorStoreField("data", is_full_text_indexed=True)] + description_vector: Annotated[ + list[float] | str | None, + VectorStoreField("vector", dimensions=1536, distance_function="cosine_similarity"), + ] = None + description_fr: Annotated[ + str, + Field(alias="Description_fr"), + VectorStoreField("data", is_full_text_indexed=True), + ] + description_fr_vector: Annotated[ + list[float] | str | None, + VectorStoreField("vector", dimensions=1536, distance_function="cosine_similarity"), + ] = None + category: Annotated[str, VectorStoreField("data")] + tags: Annotated[list[str], VectorStoreField("data", is_indexed=True)] + parking_included: Annotated[bool | None, VectorStoreField("data")] = None + last_renovation_date: Annotated[str | None, VectorStoreField("data")] = None + rating: Annotated[float, VectorStoreField("data")] + location: Annotated[dict[str, Any], VectorStoreField("data")] + address: Annotated[Address, VectorStoreField("data")] + rooms: Annotated[list[Room], VectorStoreField("data")] + + model_config = ConfigDict(alias_generator=to_pascal, populate_by_name=True, extra="ignore") + + def model_post_init(self, context: Any) -> None: + """Use the descriptions as embedding inputs when vectors are absent.""" + if self.description_vector is None: + self.description_vector = self.description + if self.description_fr_vector is None: + self.description_fr_vector = self.description_fr + + +def load_hotels() -> list[Hotel]: + """Load the existing Azure AI Search hotel dataset.""" + with urlopen(HOTELS_URL, timeout=60) as response: # nosec B310 - fixed HTTPS sample URL + records = json.loads(response.read()) + return [Hotel.model_validate(record) for record in records] + + +async def main() -> None: + """Create an in-memory hotel search tool and give it to an agent.""" + api_key = os.environ["OPENAI_API_KEY"] + collection: InMemoryCollection[str, Hotel] = InMemoryCollection( + Hotel, + embedding_generator=OpenAIEmbeddingClient( + model="text-embedding-3-small", + api_key=api_key, + ), + ) + await collection.ensure_collection_exists() + + # 1. Load the hotel records. + hotels = await asyncio.to_thread(load_hotels) + await collection.upsert(hotels) + + # 2. Param values become optional model-visible filter arguments. + # When the allowed values are known, use Literal so the tool schema exposes + # them as an enum. + category = Param( + "category", + Literal["Boutique", "Budget", "Extended-Stay", "Luxury", "Resort and Spa", "Suite"], + description="Only return hotels in this category.", + ) + min_rating = Param( + "min_rating", + float, + description="The minimum guest rating.", + minimum=0, + maximum=5, + ) + tool = create_vector_search_tool( + collection, + description="Search the hotel dataset, optionally filtering by category and minimum rating.", + filter=FilterGroup( + "and", + ( + Filter("category", "eq", category), + Filter("rating", "gte", min_rating), + ), + ), + result_mapper=lambda result: ( + f"(hotel_id: {result['record'].hotel_id}) {result['record'].hotel_name} " + f"(rating {result['record'].rating}) - {result['record'].description}. " + f"Address: {result['record'].address.city}, {result['record'].address.country}." + ), + ) + + # 3. The agent chooses whether to supply the exposed category and minimum-rating filters. + async with Agent( + client=OpenAIChatClient( + model="gpt-5.4-nano", + api_key=api_key, + ), + name="HotelAgent", + instructions=( + "Always use the search tool to answer hotel questions. " + "Use category and minimum rating filters when the request provides them. " + "Include the hotel_id in the answer." + ), + tools=[tool], + ) as agent: + result = await agent.run("Find a resort and spa with a rating of at least 4.") + print(result) + + +if __name__ == "__main__": + asyncio.run(main()) + + +""" +Sample output: +The Grand Gaming Resort (hotel_id: 20) is a Resort and Spa rated 4.2. +""" diff --git a/python/samples/AGENTS.md b/python/samples/AGENTS.md index 739329a1250..916e550e07b 100644 --- a/python/samples/AGENTS.md +++ b/python/samples/AGENTS.md @@ -10,7 +10,7 @@ python/samples/ ├── 01-get-started/ # Progressive tutorial (steps 01–07) ├── 02-agents/ # Deep-dive concept samples │ ├── tools/ # Tool patterns (function, approval, schema, etc.) -│ ├── vector_stores/ # Vector model schemas and registration +│ ├── vector_stores/ # Vector models, filters, in-memory search, and tools │ ├── middleware/ # One file per middleware concept │ ├── conversations/ # Thread, storage, suspend/resume │ ├── providers/ # One sub-folder per provider (azure_ai/, openai/, etc.)