From 78893fd709a133f3e13d85e0fdf9c5b79a385409 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Wed, 23 Sep 2026 12:19:40 +0200 Subject: [PATCH 1/3] feat(python): add DuckDB vector store connector --- python/PACKAGE_STATUS.md | 1 + python/packages/duckdb/LICENSE | 21 + python/packages/duckdb/README.md | 154 ++++ .../duckdb/agent_framework_duckdb/__init__.py | 14 + .../agent_framework_duckdb/_vector_store.py | 763 ++++++++++++++++++ .../duckdb/agent_framework_duckdb/py.typed | 0 python/packages/duckdb/pyproject.toml | 57 ++ .../packages/duckdb/samples/duckdb_vectors.py | 58 ++ .../packages/duckdb/tests/duckdb/conftest.py | 59 ++ .../duckdb/tests/duckdb/test_collection.py | 513 ++++++++++++ .../duckdb/tests/duckdb/test_settings.py | 194 +++++ python/pyproject.toml | 1 + python/uv.lock | 54 ++ 13 files changed, 1889 insertions(+) create mode 100644 python/packages/duckdb/LICENSE create mode 100644 python/packages/duckdb/README.md create mode 100644 python/packages/duckdb/agent_framework_duckdb/__init__.py create mode 100644 python/packages/duckdb/agent_framework_duckdb/_vector_store.py create mode 100644 python/packages/duckdb/agent_framework_duckdb/py.typed create mode 100644 python/packages/duckdb/pyproject.toml create mode 100644 python/packages/duckdb/samples/duckdb_vectors.py create mode 100644 python/packages/duckdb/tests/duckdb/conftest.py create mode 100644 python/packages/duckdb/tests/duckdb/test_collection.py create mode 100644 python/packages/duckdb/tests/duckdb/test_settings.py diff --git a/python/PACKAGE_STATUS.md b/python/PACKAGE_STATUS.md index 51ff0f2f7a6..8585970e46a 100644 --- a/python/PACKAGE_STATUS.md +++ b/python/PACKAGE_STATUS.md @@ -30,6 +30,7 @@ Status is grouped into these buckets: | `agent-framework-core` | `python/packages/core` | `released` | | `agent-framework-declarative` | `python/packages/declarative` | `released` | | `agent-framework-devui` | `python/packages/devui` | `beta` | +| `agent-framework-duckdb` | `python/packages/duckdb` | `alpha` | | `agent-framework-foundry` | `python/packages/foundry` | `released` | | `agent-framework-foundry-hosting` | `python/packages/foundry_hosting` | `beta` | | `agent-framework-foundry-local` | `python/packages/foundry_local` | `beta` | diff --git a/python/packages/duckdb/LICENSE b/python/packages/duckdb/LICENSE new file mode 100644 index 00000000000..96a394c43c3 --- /dev/null +++ b/python/packages/duckdb/LICENSE @@ -0,0 +1,21 @@ + MIT License + + Copyright (c) Microsoft Corporation. + + Permission is hereby granted, free of charge, to any person obtaining a copy + of this software and associated documentation files (the "Software"), to deal + in the Software without restriction, including without limitation the rights + to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + copies of the Software, and to permit persons to whom the Software is + furnished to do so, subject to the following conditions: + + The above copyright notice and this permission notice shall be included in all + copies or substantial portions of the Software. + + THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + SOFTWARE. diff --git a/python/packages/duckdb/README.md b/python/packages/duckdb/README.md new file mode 100644 index 00000000000..d9d4a5e9a7a --- /dev/null +++ b/python/packages/duckdb/README.md @@ -0,0 +1,154 @@ +# Agent Framework DuckDB vector store + +Store and search typed records in DuckDB using the +[Microsoft Agent Framework](https://learn.microsoft.com/agent-framework/) +vector-store APIs. This alpha package uses the **DuckDB Python client only**: +local files, DuckDB-supported connection URIs (including MotherDuck), and +caller-provided `duckdb.DuckDBPyConnection` objects use the same implementation. + +## Installation + +```bash +pip install agent-framework-duckdb --pre +``` + +Python 3.10+ and DuckDB 1.4.1–1.5.x are required. The package also installs +`pytz`, which DuckDB uses when returning timezone-aware timestamps. Import +`DuckDBStore` and `DuckDBCollection` directly from `agent_framework_duckdb`; +this alpha package is not part of `agent-framework-core[all]`. + +## Connections and persistence + +`DuckDBStore()` (and a directly constructed `DuckDBCollection`) defaults to +**`agent-framework.duckdb` in the current working directory**. It is a normal +persistent database file, not an in-memory database or temporary file. Give +each application a suitable writable location with +`connection_string="path/to/vectors.duckdb"`; parent directories must already +exist. The connection is opened lazily on the first async operation. +`aclose()` or an async context manager drains pending operations and releases +the file. Opening a *new* store against the same filename restores its tables +and records. + +`DUCKDB_CONNECTION_STRING` can instead provide a filename or a URI. Settings +precedence is **explicit `connection_string` > selected `.env` file > +environment > persistent default**. Select a file with `env_file_path` and +optionally `env_file_encoding`; `.env` files are never discovered implicitly. +An explicitly empty connection string is an error. The URI is held as an AF +`SecretString` and never printed or logged by the connector. + +To use [MotherDuck](https://motherduck.com/docs/getting-started/interfaces/client-apis/python/installation-authentication/), +configure the *same* DuckDB client with a `md:` URI and its usual credentials, +for example: + +```bash +export DUCKDB_CONNECTION_STRING='md:my_db' +export MOTHERDUCK_TOKEN='' +``` + +```python +from agent_framework_duckdb import DuckDBStore + + +async def list_remote_tables() -> None: + async with DuckDBStore() as store: + print(await store.list_collection_names()) +``` + +For an application-managed token, pass +`DuckDBStore(connection_string="md:my_db", +config={"motherduck_token": SecretString(token)})` (import `SecretString` +from `agent_framework`). `config` is forwarded to `duckdb.connect`; it is +not loaded from the connector's `.env` file. Keep tokens outside source control +and avoid embedding them in a URI. Other DuckDB-supported services use their +documented URI/configuration or an injected client; no service-specific +adapter is installed. DuckDB/MotherDuck may need network access and compatible +client/extension versions; authentication and service availability are managed +by DuckDB, not by this connector. + +Alternatively, use `DuckDBStore(client=duckdb.connect(...))` or +`DuckDBCollection(Record, client=...)` to **borrow** an existing connection. +The caller owns and closes that connection; `client` cannot be combined with +`connection_string`, `config`, or `.env` options. Collections obtained from a +store borrow its client, so keep the store open while using them. All connector +operations on one store/collection run serially on one worker thread, not on +the event loop. Coordinate any *external* use of an injected connection +yourself. Cancelling an awaiter does not stop a DuckDB query already running +on that thread; `aclose()` waits for it to finish. + +## Example + +```python +from dataclasses import dataclass +from typing import Annotated + +from agent_framework import Filter, VectorStoreField, vectorstoremodel +from agent_framework_duckdb import DuckDBStore + + +@vectorstoremodel(collection_name="articles") +@dataclass +class Article: + id: Annotated[str, VectorStoreField("key")] + text: Annotated[str, VectorStoreField("data")] + embedding: Annotated[list[float] | None, VectorStoreField("vector", dimensions=3)] = None + + +async def save_and_search() -> None: + async with DuckDBStore(connection_string="articles.duckdb") as store: + collection = store.get_collection(Article) + await collection.ensure_collection_exists() + await collection.upsert( + [Article("one", "DuckDB persists records", [1, 0, 0])], + generate_vectors=False, + ) + + async with DuckDBStore(connection_string="articles.duckdb") as reopened: + collection = reopened.get_collection(Article) + assert (await collection.get(["one"]))[0].text == "DuckDB persists records" + results = await collection.search( + vector=[1, 0, 0], + filter=Filter("text", "contains_text", "persists"), + ) + async for result in results: + print(result["record"].text, result["score"]) +``` + +Run a complete local example with +`uv run --package agent-framework-duckdb python packages/duckdb/samples/duckdb_vectors.py` +from the `python/` directory. It writes a default database file, reopens it, +and cleans up only its sample table. + +## Capabilities and limits + +- Tables are created by `ensure_collection_exists()`; `collection_exists()`, + `list_collection_names()`, and `ensure_collection_deleted()` provide the + corresponding lifecycle. Existing tables are **not migrated or reindexed**. + Identifiers are quoted, values are parameterized, and collection names refer + to tables in the connection's current database and schema. +- Batch upsert, retrieval, filtered/paged listing, and delete support string, + signed 64-bit integer, and UUID keys; string and UUID keys can be generated + when declared auto-generated. Records support string, integer, float, + boolean, UUID, bytes, date, timezone-aware datetime, JSON list/dict data, + and nullable dense vectors. Use `generate_vectors=False` for supplied + embeddings or pass a local `embedding_generator`. Retrieval excludes vectors + unless `include_vectors=True`. +- Exact SQL vector search supports `DEFAULT`/`cosine_distance`, + `cosine_similarity`, `euclidean_distance`, `dot_prod`, and + `negative_dot_prod`. The default score is cosine **distance**, so lower is + better and `score_threshold` is a maximum; similarity/dot-product scores + use a minimum threshold. Null vectors are omitted. Results are filtered + and paged **in DuckDB**, with the primary key as a stable tie-breaker. +- Portable filters support scalar equality/inequality, null/presence checks, + ordered scalar comparisons, `in`/`not_in`, text contains/prefix/suffix, + and AND/OR/NOT groups. They do not coerce booleans into numbers, and text + wildcards are literal. JSON equality/collection membership and nested + paths are not supported. +- No approximate vector indexes, keyword/hybrid search, full-text or explicit + data indexes, binary/sparse vectors, auto-generated integer keys, or + server-side vectorization are provided. Unsupported options raise errors. + DuckDB is an in-process database with file locking: multiple writer + **processes** cannot concurrently write the same local file. Local records + are unencrypted; choose and protect the database file appropriately. + Connector-created batch upserts are transactional. On a borrowed connection + the connector does not start/commit/roll back a caller transaction; without + one, a failed batch may have partially persisted. diff --git a/python/packages/duckdb/agent_framework_duckdb/__init__.py b/python/packages/duckdb/agent_framework_duckdb/__init__.py new file mode 100644 index 00000000000..4c1049a30d6 --- /dev/null +++ b/python/packages/duckdb/agent_framework_duckdb/__init__.py @@ -0,0 +1,14 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""DuckDB vector stores for Microsoft Agent Framework.""" + +import importlib.metadata + +from ._vector_store import DuckDBCollection, DuckDBSettings, DuckDBStore + +try: + __version__ = importlib.metadata.version(__name__) +except importlib.metadata.PackageNotFoundError: + __version__ = "0.0.0" + +__all__ = ["DuckDBCollection", "DuckDBSettings", "DuckDBStore", "__version__"] diff --git a/python/packages/duckdb/agent_framework_duckdb/_vector_store.py b/python/packages/duckdb/agent_framework_duckdb/_vector_store.py new file mode 100644 index 00000000000..c5ff2987c00 --- /dev/null +++ b/python/packages/duckdb/agent_framework_duckdb/_vector_store.py @@ -0,0 +1,763 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Persistent DuckDB collections with exact dense-vector search.""" + +from __future__ import annotations + +import asyncio +import json +import math +import threading +from collections.abc import Callable, Mapping, Sequence +from concurrent.futures import Future, ThreadPoolExecutor +from datetime import date, datetime +from typing import Any, ClassVar, Generic, cast +from uuid import UUID, uuid4 + +import duckdb +from agent_framework import ( + BaseVectorCollection, + BaseVectorSearch, + BaseVectorStore, + Filter, + FilterGroup, + SearchResults, + SecretString, + VectorStoreCollectionDefinition, + VectorStoreField, + load_settings, +) +from agent_framework._vector_filters import FilterExpression +from agent_framework._vectors import EmbeddingClient, SearchType, Vector +from agent_framework.exceptions import IntegrationException +from typing_extensions import TypedDict, TypeVar + +KeyT = TypeVar("KeyT", default=Any) +ModelT = TypeVar("ModelT", default=Any) +ResultT = TypeVar("ResultT") + +_DEFAULT_DATABASE = "agent-framework.duckdb" +_DATA_TYPES = { + "str": "VARCHAR", + "int": "BIGINT", + "float": "DOUBLE", + "bool": "BOOLEAN", + "UUID": "UUID", + "bytes": "BLOB", + "date": "DATE", + "datetime": "TIMESTAMPTZ", + "list": "JSON", + "dict": "JSON", +} +_VECTOR_TYPES = {"float": "DOUBLE", "float64": "DOUBLE", "float32": "FLOAT"} +_METRICS = { + "DEFAULT": ("array_cosine_distance", False), + "cosine_distance": ("array_cosine_distance", False), + "cosine_similarity": ("array_cosine_similarity", True), + "euclidean_distance": ("array_distance", False), + "dot_prod": ("array_inner_product", True), + "negative_dot_prod": ("-array_inner_product", False), +} +_ORDER_OPERATORS = {"gt": ">", "gte": ">=", "lt": "<", "lte": "<="} +_TABLE_SCOPE = ( + "FROM information_schema.tables " + "WHERE table_catalog = current_database() AND table_schema = current_schema() " + "AND table_type = 'BASE TABLE'" +) +_LIST_TABLES = f"SELECT table_name {_TABLE_SCOPE} ORDER BY table_name" +_TABLE_EXISTS = f"SELECT 1 {_TABLE_SCOPE} AND lower(table_name) = lower(?)" + + +class DuckDBSettings(TypedDict, total=False): + """Connection settings resolved from arguments, a selected .env file, or ``DUCKDB_`` variables.""" + + connection_string: SecretString | None + """Local database filename or DuckDB URI, resolved from ``DUCKDB_CONNECTION_STRING``.""" + + +def _identifier(name: str) -> str: + if not isinstance(name, str) or not name or "\0" in name: + raise ValueError("DuckDB identifiers must be nonempty strings without NUL.") + return '"' + name.replace('"', '""') + '"' + + +def _check_options(options: Mapping[str, Any] | None) -> None: + if options: + raise NotImplementedError(f"Unsupported DuckDB operation option(s): {', '.join(sorted(options))}.") + + +def _execute_statement( + connection: duckdb.DuckDBPyConnection, + statement: str, + parameters: Sequence[Any] = (), +) -> None: + connection.execute(statement, parameters) + + +def _column_type(field: VectorStoreField) -> str: + if field.field_type == "vector": + if type(field.dimensions) is not int or field.dimensions <= 0: + raise ValueError(f"Vector field '{field.name}' needs positive integer dimensions.") + if field.type_ not in _VECTOR_TYPES: + raise NotImplementedError(f"Unsupported DuckDB vector type '{field.type_}'.") + return f"{_VECTOR_TYPES[field.type_]}[{field.dimensions}]" + if field.type_ not in _DATA_TYPES: + raise NotImplementedError(f"Field '{field.name}' needs a supported explicit type; got '{field.type_}'.") + return _DATA_TYPES[field.type_] + + +def _validate_json(value: Any) -> None: + if value is None or type(value) in (str, int, bool): + return + if type(value) is float: + if not math.isfinite(value): + raise ValueError("JSON values must contain finite numbers.") + return + if type(value) is list: + for item in cast(list[Any], value): + _validate_json(item) + return + if type(value) is dict: + for key, item in cast(dict[Any, Any], value).items(): + if not isinstance(key, str): + raise TypeError("JSON object keys must be strings.") + _validate_json(item) + return + raise TypeError("JSON fields require JSON-compatible lists or string-keyed dictionaries.") + + +def _prepare_vector(field: VectorStoreField, value: Any) -> list[float] | None: + if value is None: + return None + if not isinstance(value, Sequence) or isinstance(value, (str, bytes, bytearray)): + raise TypeError(f"Vector field '{field.name}' requires a dense numeric sequence.") + elements = cast(Sequence[Any], value) + if len(elements) != field.dimensions: + raise ValueError(f"Vector field '{field.name}' expects {field.dimensions} dimensions; got {len(elements)}.") + if any(type(item) not in (float, int) or not math.isfinite(item) for item in elements): + raise ValueError(f"Vector field '{field.name}' requires finite numeric elements without NULLs.") + if field.distance_function in (None, "DEFAULT", "cosine_distance", "cosine_similarity") and not any(elements): + raise ValueError(f"Cosine vector field '{field.name}' requires a nonzero vector.") + return [float(item) for item in elements] + + +def _prepare_value(field: VectorStoreField, value: Any) -> Any: + if field.field_type == "vector": + return _prepare_vector(field, value) + if value is None: + if field.field_type == "key": + raise ValueError("DuckDB keys cannot be null.") + return None + kind = field.type_ + if kind == "str" and isinstance(value, str): + return value + if kind == "int" and type(value) is int: + if not -(2**63) <= value < 2**63: + raise ValueError(f"Field '{field.name}' exceeds the BIGINT range.") + return value + if kind == "float" and type(value) in (float, int): + if not math.isfinite(value): + raise ValueError(f"Field '{field.name}' requires a finite number.") + return float(value) + if kind == "bool" and type(value) is bool: + return value + if kind == "UUID" and isinstance(value, (str, UUID)): + return UUID(str(value)) + if kind == "bytes" and isinstance(value, bytes): + return value + if kind == "date" and type(value) is date: + return value + if kind == "date" and isinstance(value, str): + return date.fromisoformat(value) + if kind == "datetime": + parsed = datetime.fromisoformat(value.replace("Z", "+00:00")) if isinstance(value, str) else value + if isinstance(parsed, datetime): + if parsed.tzinfo is None or parsed.utcoffset() is None: + raise ValueError(f"Datetime field '{field.name}' requires a timezone.") + return parsed + if kind in ("list", "dict") and type(value) is (list if kind == "list" else dict): + _validate_json(value) + return json.dumps(value, allow_nan=False) + raise TypeError(f"Field '{field.name}' requires a value of type '{kind}'.") + + +def _create_client( + connection_string: str | SecretString | None, + *, + client: duckdb.DuckDBPyConnection | None, + config: Mapping[str, str | SecretString] | None, + env_file_path: str | None, + env_file_encoding: str | None, +) -> _Client: + if client is not None: + if any(value is not None for value in (connection_string, config, env_file_path, env_file_encoding)): + raise ValueError("client cannot be combined with connection_string, config, or .env options.") + if not isinstance(client, duckdb.DuckDBPyConnection): + raise TypeError("client must be a DuckDBPyConnection.") + return _Client(None, client=client, config=None) + settings = load_settings( + DuckDBSettings, + env_prefix="DUCKDB_", + connection_string=connection_string, + env_file_path=env_file_path, + env_file_encoding=env_file_encoding, + ) + address = settings.get("connection_string") + if address is not None and not isinstance(address, SecretString): + raise TypeError("connection_string must be a string or SecretString.") + if address is None: + address = SecretString(_DEFAULT_DATABASE) + if not address.get_secret_value().strip(): + raise ValueError("connection_string must not be empty.") + if config is not None: + if not isinstance(config, Mapping): + raise TypeError("config must be a mapping of DuckDB settings.") + for key, value in config.items(): + if not isinstance(key, str) or not key or not isinstance(value, (str, SecretString)): + raise TypeError("DuckDB config keys and values must be nonempty string keys and string values.") + return _Client(address, client=None, config=config) + + +class _Client: + """Run every operation on one worker, including connection setup and shutdown.""" + + def __init__( + self, + address: SecretString | None, + *, + client: duckdb.DuckDBPyConnection | None, + config: Mapping[str, str | SecretString] | None, + ) -> None: + self.owned = client is None + self._address = address + self._connection = client + self._config = {key: SecretString(value) for key, value in (config or {}).items()} + self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="af-duckdb") + self._lifecycle_lock = threading.Lock() + self._close_future: Future[None] | None = None + + def _execute(self, operation: Callable[[duckdb.DuckDBPyConnection], ResultT]) -> ResultT: + try: + if self._connection is None: + if self._address is None: + raise RuntimeError("A DuckDB connection or database path is required.") + self._connection = duckdb.connect( + database=self._address.get_secret_value(), + config={key: value.get_secret_value() for key, value in self._config.items()}, + ) + return operation(self._connection) + except duckdb.Error as exc: + raise IntegrationException("DuckDB operation failed; inspect the chained driver exception.") from exc + + async def run(self, operation: Callable[[duckdb.DuckDBPyConnection], ResultT]) -> ResultT: + """Submit one complete database operation without blocking the event loop.""" + with self._lifecycle_lock: + if self._close_future is not None: + raise RuntimeError("The DuckDB client is closed.") + future = self._executor.submit(self._execute, operation) + return await asyncio.wrap_future(future) + + def _close(self) -> None: + if self.owned and self._connection is not None: + try: + self._connection.close() + except duckdb.Error as exc: + raise IntegrationException("DuckDB close failed; inspect the chained driver exception.") from exc + + async def aclose(self) -> None: + """Wait for queued work, close owned resources, and leave injected connections open.""" + with self._lifecycle_lock: + if self._close_future is None: + self._close_future = self._executor.submit(self._close) + self._executor.shutdown(wait=False) + close_future = self._close_future + await asyncio.shield(asyncio.wrap_future(close_future)) + + +class _FilterCompiler: + """Compile supported portable filters into two-valued, parameterized SQL.""" + + def __init__(self, definition: VectorStoreCollectionDefinition) -> None: + self.definition = definition + self.parameters: list[Any] = [] + + def compile(self, expression: FilterExpression | None) -> tuple[str, list[Any]]: + """Return a SQL predicate and its bound arguments.""" + if expression is None: + return "TRUE", [] + return self._condition(expression), self.parameters + + def _equality(self, field: VectorStoreField, column: str, value: Any) -> str: + if field.type_ in ("list", "dict"): + raise NotImplementedError("DuckDB JSON equality and collection filters are not supported.") + if value is None: + return f"{column} IS NULL" + if field.type_ in ("int", "float") and type(value) in (int, float): + if not math.isfinite(value): + raise ValueError("Numeric filter operands must be finite.") + adapted = value + else: + try: + adapted = _prepare_value(field, value) + except TypeError: + return "FALSE" + self.parameters.append(adapted) + return f"{column} IS NOT DISTINCT FROM ?" + + def _condition(self, expression: FilterExpression) -> str: + if isinstance(expression, FilterGroup): + parts = [self._condition(child) for child in expression.filters] + if expression.operator == "not": + return f"(NOT ({parts[0]}))" + joiner = " AND " if expression.operator == "and" else " OR " + return f"({joiner.join(parts)})" + if not isinstance(expression, Filter): + raise TypeError("filter must be a Filter or FilterGroup.") + if "." in expression.field_name: + raise NotImplementedError("DuckDB does not support nested filter paths.") + field = self.definition.try_get_field(expression.field_name) + if field is None: + raise ValueError(f"Unknown DuckDB field '{expression.field_name}'.") + if field.field_type == "vector": + raise NotImplementedError("Filtering vector columns is not supported.") + column = f"t.{_identifier(field.storage_name or field.name)}" + op, value = expression.operator, expression.value + if op == "exists": + return "TRUE" + if op == "is_null": + return f"{column} IS NULL" + if op == "is_not_null": + return f"{column} IS NOT NULL" + if op in ("eq", "ne"): + equality = self._equality(field, column, value) + return equality if op == "eq" else f"(NOT ({equality}))" + if op in ("in", "not_in"): + parts = [self._equality(field, column, item) for item in value] + choices = f"({' OR '.join(parts)})" if parts else "FALSE" + membership = choices if op == "in" else f"(NOT {choices})" + return f"({column} IS NOT NULL AND {membership})" + if op in (*_ORDER_OPERATORS, "between"): + if field.type_ not in ("str", "int", "float", "date", "datetime"): + raise NotImplementedError(f"Ordered filtering is not supported for '{field.type_}'.") + operands: Sequence[Any] = value if op == "between" else [value] + for operand in operands: + if operand is None or type(operand) is bool: + raise TypeError("Ordered filter operands must be non-null scalars of the column's type.") + if field.type_ in ("int", "float") and type(operand) in (int, float): + if not math.isfinite(operand): + raise ValueError("Ordered filter numbers must be finite.") + self.parameters.append(operand) + else: + self.parameters.append(_prepare_value(field, operand)) + if op == "between": + return f"({column} BETWEEN ? AND ?) IS TRUE" + return f"({column} {_ORDER_OPERATORS[op]} ?) IS TRUE" + if op in ("contains_text", "starts_with", "ends_with"): + if field.type_ != "str": + raise TypeError("Text filters require a string field.") + if not isinstance(value, str): + raise TypeError("Text filters require a string operand.") + escaped = value.replace("!", "!!").replace("%", "!%").replace("_", "!_") + pattern = ("%" if op != "starts_with" else "") + escaped + ("%" if op != "ends_with" else "") + self.parameters.append(pattern) + return f"({column} LIKE ? ESCAPE '!') IS TRUE" + raise NotImplementedError(f"Unsupported DuckDB filter operator '{op}'.") + + +class DuckDBCollection(BaseVectorCollection[KeyT, ModelT], BaseVectorSearch[KeyT, ModelT], Generic[KeyT, ModelT]): + """One DuckDB table with exact dense-vector search and typed record codecs.""" + + supported_key_types: ClassVar[set[str] | None] = {"str", "int", "UUID"} + supported_vector_types: ClassVar[set[str] | None] = set(_VECTOR_TYPES) + supported_search_types: ClassVar[set[SearchType]] = {"vector"} + + def __init__( + self, + record_type: type[ModelT], + *, + connection_string: str | SecretString | None = None, + client: duckdb.DuckDBPyConnection | None = None, + config: Mapping[str, str | SecretString] | None = None, + definition: VectorStoreCollectionDefinition | None = None, + collection_name: str | None = None, + embedding_generator: EmbeddingClient | None = None, + env_file_path: str | None = None, + env_file_encoding: str | None = None, + _shared_client: _Client | None = None, + ) -> None: + """Initialize a collection; the first async operation opens an owned connection. + + Args: + record_type: Registered typed model or ``dict`` with an explicit definition. + connection_string: Persistent filename or DuckDB URI. Defaults to ``agent-framework.duckdb``. + client: Borrowed DuckDB connection; mutually exclusive with connection settings. + config: DuckDB connection settings, optionally wrapping secret values in ``SecretString``. + definition: Explicit definition for dictionary records. + collection_name: Table name, overriding the model definition. + embedding_generator: Local embedding client for generated vectors. + env_file_path: Optional selected .env file for ``DUCKDB_CONNECTION_STRING``. + env_file_encoding: Encoding of the selected .env file. + _shared_client: Internal connection owned by a parent store. + """ + super().__init__( + record_type, + definition=definition, + collection_name=collection_name, + embedding_generator=embedding_generator, + managed_client=client is None and _shared_client is None, + ) + self._table = _identifier(self.collection_name) + names = [field.storage_name or field.name for field in self.definition.fields] + if len({name.casefold() for name in names}) != len(names): + raise ValueError("DuckDB column names must be unique ignoring case.") + self._fields = tuple(self.definition.fields) + for name in names: + _identifier(name) + if _shared_client is not None and any( + value is not None for value in (connection_string, client, config, env_file_path, env_file_encoding) + ): + raise ValueError("A store collection cannot override its shared connection settings.") + self._client = _shared_client or _create_client( + connection_string, + client=client, + config=config, + env_file_path=env_file_path, + env_file_encoding=env_file_encoding, + ) + self._owns_wrapper = _shared_client is None + + def _validate_data_model(self) -> None: + super()._validate_data_model() + for field in self.definition.fields: + _column_type(field) + if field.is_full_text_indexed or field.is_indexed: + raise NotImplementedError("DuckDB full-text and explicit data indexes are not supported.") + if field.field_type == "vector": + if field.index_kind not in ("default", "flat"): + raise NotImplementedError(f"Unsupported DuckDB vector index kind '{field.index_kind}'.") + if field.distance_function not in _METRICS: + raise NotImplementedError(f"Unsupported DuckDB distance function '{field.distance_function}'.") + if any(name.startswith("duckdb.") for name in field.provider_annotations): + raise NotImplementedError("DuckDB provider annotations are not supported.") + key = self.definition.key_field + if key.is_auto_generated and key.type_ == "int": + raise NotImplementedError("DuckDB auto-generated integer keys are not supported.") + + async def __aexit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: + """Release a direct collection's connection; store collections leave ownership to the store.""" + await self.aclose() + + async def aclose(self) -> None: + """Close this collection's wrapper, leaving injected or store-owned connections open.""" + if self._owns_wrapper: + await self._client.aclose() + + async def collection_exists(self, *, operation_options: Mapping[str, Any] | None = None) -> bool: + """Check for this table in the active DuckDB database and schema.""" + _check_options(operation_options) + return await self._client.run( + lambda connection: connection.execute(_TABLE_EXISTS, [self.collection_name]).fetchone() is not None + ) + + async def ensure_collection_exists(self, *, operation_options: Mapping[str, Any] | None = None) -> None: + """Create the table if absent; existing schemas are not modified.""" + _check_options(operation_options) + columns = [ + f"{_identifier(field.storage_name or field.name)} {_column_type(field)}" + + (" PRIMARY KEY" if field.field_type == "key" else "") + for field in self._fields + ] + await self._client.run( + lambda connection: _execute_statement( + connection, f"CREATE TABLE IF NOT EXISTS {self._table} ({', '.join(columns)})" + ) + ) + + async def ensure_collection_deleted(self, *, operation_options: Mapping[str, Any] | None = None) -> None: + """Drop this table if present, without affecting other tables.""" + _check_options(operation_options) + await self._client.run(lambda connection: _execute_statement(connection, f"DROP TABLE IF EXISTS {self._table}")) + + def _deserialize_store_models_to_dicts( + self, + records: Sequence[Any], + *, + context: Mapping[str, Any] | None = None, + ) -> Sequence[dict[str, Any]]: + decoded = super()._deserialize_store_models_to_dicts(records, context=context) + for record in decoded: + for field in self._fields: + name = field.storage_name or field.name + value = record.get(name) + if field.field_type == "vector" and isinstance(value, tuple): + record[name] = list(cast(tuple[Any, ...], value)) + elif field.type_ in ("list", "dict") and value is not None: + record[name] = json.loads(value) + return decoded + + def _columns(self, include_vectors: bool) -> list[str]: + return [ + field.storage_name or field.name + for field in self._fields + if include_vectors or field.field_type != "vector" + ] + + def _order_by(self, order_by: Mapping[str, bool] | None) -> str: + parts: list[str] = [] + for name, ascending in (order_by or {}).items(): + if type(ascending) is not bool: + raise TypeError("Order directions must be booleans.") + if "." in name: + raise NotImplementedError("DuckDB does not support nested order paths.") + field = self.definition.try_get_field(name) + if field is None: + raise ValueError(f"Unknown DuckDB order field '{name}'.") + if field.field_type == "vector" or field.type_ in ("list", "dict"): + raise NotImplementedError(f"Ordering field '{name}' is not supported.") + direction = "ASC" if ascending else "DESC" + parts.append(f"t.{_identifier(field.storage_name or field.name)} {direction} NULLS LAST") + if self.definition.key_field.name not in (order_by or {}): + parts.append(f"t.{_identifier(self.definition.key_field_storage_name)} ASC") + return ", ".join(parts) + + async def _inner_upsert( + self, + records: Sequence[Any], + *, + operation_options: Mapping[str, Any] | None = None, + ) -> Sequence[KeyT]: + _check_options(operation_options) + key = self.definition.key_field + key_name = self.definition.key_field_storage_name + names = [field.storage_name or field.name for field in self._fields] + rows: list[tuple[Any, ...]] = [] + keys: list[KeyT] = [] + for record in records: + if key_name not in record: + if not key.is_auto_generated: + raise ValueError(f"Record is missing vector store field '{key.name}'.") + record = {**record, key_name: str(uuid4()) if key.type_ == "str" else uuid4()} + row = tuple(_prepare_value(field, record[name]) for field, name in zip(self._fields, names, strict=True)) + rows.append(row) + keys.append(cast(KeyT, row[names.index(key_name)])) + if not rows: + return [] + identifiers = ", ".join(map(_identifier, names)) + placeholders = ", ".join("?" for _ in names) + updates = [name for name in names if name != key_name] + conflict = ( + f"DO UPDATE SET {', '.join(f'{_identifier(name)} = EXCLUDED.{_identifier(name)}' for name in updates)}" # nosec B608 + if updates + else "DO NOTHING" + ) + statement = ( + f"INSERT INTO {self._table} ({identifiers}) VALUES ({placeholders}) " # ruff: ignore[hardcoded-sql-expression] # nosec B608 + f"ON CONFLICT ({_identifier(key_name)}) {conflict}" + ) + + def write(connection: duckdb.DuckDBPyConnection) -> None: + if self._client.owned: + connection.execute("BEGIN TRANSACTION") + try: + connection.executemany(statement, rows) + if self._client.owned: + connection.execute("COMMIT") + except Exception: + if self._client.owned: + connection.execute("ROLLBACK") + raise + + await self._client.run(write) + 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]: + _check_options(operation_options) + columns = self._columns(include_vectors) + selected = ", ".join(f"t.{_identifier(name)}" for name in columns) + statement = f"SELECT {selected} FROM {self._table} AS t WHERE " # ruff: ignore[hardcoded-sql-expression] # nosec B608 + adapted: list[Any] = [] + if keys is not None: + if order_by: + raise ValueError("order_by applies only to filtered retrieval, not key lookup.") + adapted = [_prepare_value(self.definition.key_field, key) for key in keys] + if not adapted: + return [] + placeholders = ", ".join("?" for _ in adapted) + statement += f"t.{_identifier(self.definition.key_field_storage_name)} IN ({placeholders})" + params: list[Any] = adapted + else: + predicate, params = _FilterCompiler(self.definition).compile(filter) + statement += f"{predicate} ORDER BY {self._order_by(order_by)} LIMIT ? OFFSET ?" + params.extend((top, skip)) + + def read(connection: duckdb.DuckDBPyConnection) -> list[dict[str, Any]]: + rows = connection.execute(statement, params).fetchall() + return [dict(zip(columns, row, strict=True)) for row in rows] + + records = await self._client.run(read) + if keys is not None: + key_name = self.definition.key_field_storage_name + by_key = {record[key_name]: record for record in records} + return [by_key[key] for key in adapted if key in by_key] + return records + + async def _inner_delete(self, keys: Sequence[KeyT], *, operation_options: Mapping[str, Any] | None = None) -> None: + _check_options(operation_options) + adapted = [_prepare_value(self.definition.key_field, key) for key in keys] + if adapted: + statement = ( + f"DELETE FROM {self._table} WHERE {_identifier(self.definition.key_field_storage_name)} " # ruff: ignore[hardcoded-sql-expression] # nosec B608 + f"IN ({', '.join('?' for _ in adapted)})" + ) + await self._client.run(lambda connection: _execute_statement(connection, statement, adapted)) + + 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]: + _check_options(operation_options) + if search_type != "vector" or additional_property_name is not None: + raise NotImplementedError("DuckDB supports dense vector search only, not keyword-hybrid search.") + if vector is None: + raise NotImplementedError("DuckDB cannot vectorize values; supply a vector or embedding generator.") + field = self.definition.try_get_vector_field(vector_property_name) + if field is None: + raise ValueError("Select a vector_property_name from the collection definition.") + query_vector = _prepare_vector(field, vector) + metric, descending = _METRICS[field.distance_function or "DEFAULT"] + column = f"t.{_identifier(field.storage_name or field.name)}" + distance = f"{metric}({column}, CAST(? AS {_column_type(field)}))" + if metric in ("array_cosine_distance", "array_cosine_similarity"): + distance = f"CASE WHEN array_inner_product({column}, {column}) = 0 THEN NULL ELSE {distance} END" + predicate, filter_params = _FilterCompiler(self.definition).compile(filter) + columns = self._columns(include_vectors) + selected = ", ".join(f"t.{_identifier(name)}" for name in columns) + comparison = ">=" if descending else "<=" + order = "DESC" if descending else "ASC" + statement = ( + f"SELECT {selected}, scores.value FROM {self._table} AS t " # ruff: ignore[hardcoded-sql-expression] # nosec B608 + f"CROSS JOIN LATERAL (SELECT {distance} AS value) AS scores " + f"WHERE {column} IS NOT NULL AND ({predicate}) AND isfinite(scores.value)" + ) + parameters: list[Any] = [query_vector, *filter_params] + if score_threshold is not None: + if type(score_threshold) not in (float, int) or not math.isfinite(score_threshold): + raise ValueError("score_threshold must be a finite number.") + statement += f" AND scores.value {comparison} ?" + parameters.append(score_threshold) + statement += ( + f" ORDER BY scores.value {order}, t.{_identifier(self.definition.key_field_storage_name)} ASC " + "LIMIT ? OFFSET ?" + ) + parameters.extend((top, skip)) + + def search(connection: duckdb.DuckDBPyConnection) -> list[dict[str, Any]]: + rows = connection.execute(statement, parameters).fetchall() + return [{"record": dict(zip(columns, row[:-1], strict=True)), "score": float(row[-1])} for row in rows] + + return SearchResults( + await self._client.run(search), + metadata={"distance_function": field.distance_function or "DEFAULT", "approximate": False}, + ) + + def _get_record_from_result(self, result: Any) -> Any: + return result["record"] + + def _get_score_from_result(self, result: Any) -> float | None: + return float(result["score"]) + + +class DuckDBStore(BaseVectorStore): + """Create DuckDB collections sharing one serial, worker-owned connection.""" + + def __init__( + self, + *, + connection_string: str | SecretString | None = None, + client: duckdb.DuckDBPyConnection | None = None, + config: Mapping[str, str | SecretString] | None = None, + embedding_generator: EmbeddingClient | None = None, + env_file_path: str | None = None, + env_file_encoding: str | None = None, + ) -> None: + """Initialize a persistent local store or a DuckDB-client-supported service. + + Args: + connection_string: On-disk filename or DuckDB URI; defaults to ``agent-framework.duckdb``. + client: Borrowed DuckDB connection; cannot be combined with other connection settings. + config: Additional DuckDB connection settings. Wrap credential values in ``SecretString``. + embedding_generator: Default local embedding client for collections. + env_file_path: Optional selected .env file for ``DUCKDB_CONNECTION_STRING``. + env_file_encoding: Encoding of the selected .env file. + """ + super().__init__(embedding_generator=embedding_generator, managed_client=client is None) + self._client = _create_client( + connection_string, + client=client, + config=config, + env_file_path=env_file_path, + env_file_encoding=env_file_encoding, + ) + + def get_collection( + self, + record_type: type[ModelT], + *, + definition: VectorStoreCollectionDefinition | None = None, + collection_name: str | None = None, + embedding_generator: EmbeddingClient | None = None, + ) -> DuckDBCollection[Any, ModelT]: + """Create a collection that borrows this store's connection and lifecycle.""" + return DuckDBCollection( + record_type, + definition=definition, + collection_name=collection_name, + embedding_generator=embedding_generator if embedding_generator is not None else self.embedding_generator, + _shared_client=self._client, + ) + + async def list_collection_names(self, *, operation_options: Mapping[str, Any] | None = None) -> Sequence[str]: + """List base tables in the active database and schema.""" + _check_options(operation_options) + return await self._client.run( + lambda connection: [str(row[0]) for row in connection.execute(_LIST_TABLES).fetchall()] + ) + + async def _inner_ensure_collection_deleted( + self, + collection_name: str, + *, + operation_options: Mapping[str, Any] | None = None, + ) -> None: + _check_options(operation_options) + table = _identifier(collection_name) + await self._client.run(lambda connection: _execute_statement(connection, f"DROP TABLE IF EXISTS {table}")) + + async def __aexit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: + """Close owned connections and wait for the database file to be released.""" + await self.aclose() + + async def aclose(self) -> None: + """Drain queued work and close only a connector-created connection.""" + await self._client.aclose() diff --git a/python/packages/duckdb/agent_framework_duckdb/py.typed b/python/packages/duckdb/agent_framework_duckdb/py.typed new file mode 100644 index 00000000000..e69de29bb2d diff --git a/python/packages/duckdb/pyproject.toml b/python/packages/duckdb/pyproject.toml new file mode 100644 index 00000000000..f127b550840 --- /dev/null +++ b/python/packages/duckdb/pyproject.toml @@ -0,0 +1,57 @@ +[project] +name = "agent-framework-duckdb" +description = "DuckDB vector stores for Microsoft Agent Framework." +authors = [{ name = "Microsoft", email = "af-support@microsoft.com" }] +readme = "README.md" +requires-python = ">=3.10" +version = "1.0.0a260923" +license-files = ["LICENSE"] +urls.homepage = "https://aka.ms/agent-framework" +urls.source = "https://github.com/microsoft/agent-framework/tree/main/python" +urls.issues = "https://github.com/microsoft/agent-framework/issues" +classifiers = [ + "License :: OSI Approved :: MIT License", + "Development Status :: 3 - Alpha", + "Intended Audience :: Developers", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Programming Language :: Python :: 3.14", + "Typing :: Typed", +] +dependencies = [ + "agent-framework-core>=1.19.0,<2", + "duckdb>=1.4.1,<1.6", + "pytz>=2024.1,<2027", +] + +[tool.ruff] +extend = "../../pyproject.toml" + +[tool.ruff.lint.per-file-ignores] +"samples/**" = ["D", "INP", "TD", "ERA001", "RUF", "S", "T201", "CPY"] +"tests/**" = ["D", "INP", "TD", "ERA001", "RUF", "S"] + +[tool.pyright] +extends = "../../pyproject.toml" +include = ["agent_framework_duckdb"] + +[tool.pytest.ini_options] +testpaths = ["tests"] +addopts = "-ra -q -r fEX" +asyncio_mode = "auto" +asyncio_default_fixture_loop_scope = "function" +timeout = 60 + +[tool.coverage.run] +omit = ["**/__init__.py"] + +[tool.poe] +executor.type = "uv" +include = "../../shared_tasks.toml" + +[build-system] +requires = ["flit-core>=3.11,<4.0"] +build-backend = "flit_core.buildapi" diff --git a/python/packages/duckdb/samples/duckdb_vectors.py b/python/packages/duckdb/samples/duckdb_vectors.py new file mode 100644 index 00000000000..66cb605b9c1 --- /dev/null +++ b/python/packages/duckdb/samples/duckdb_vectors.py @@ -0,0 +1,58 @@ +# Copyright (c) Microsoft. All rights reserved. + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass +from typing import Annotated +from uuid import uuid4 + +from agent_framework import Filter, VectorStoreField, vectorstoremodel + +from agent_framework_duckdb import DuckDBStore + +""" +Persist and reopen typed vector records in the default on-disk DuckDB database. + +Run from python/: + uv run --package agent-framework-duckdb python packages/duckdb/samples/duckdb_vectors.py + +The default agent-framework.duckdb file remains in the working directory. +Only the uniquely named sample table is deleted. No service or API key is needed. +""" + + +@vectorstoremodel +@dataclass +class Note: + id: Annotated[str, VectorStoreField("key")] + text: Annotated[str, VectorStoreField("data")] + embedding: Annotated[list[float] | None, VectorStoreField("vector", dimensions=3)] = None + + +async def main() -> None: + table_name = f"af_duckdb_example_{uuid4().hex}" + async with DuckDBStore() as store: + collection = store.get_collection(Note, collection_name=table_name) + await collection.ensure_collection_exists() + await collection.upsert( + [Note("one", "A persistent DuckDB record", [1, 0, 0])], + generate_vectors=False, + ) + + async with DuckDBStore() as reopened: + collection = reopened.get_collection(Note, collection_name=table_name) + try: + print((await collection.get(["one"]))[0].text) + results = await collection.search( + vector=[1, 0, 0], + filter=Filter("text", "contains_text", "DuckDB"), + ) + async for result in results: + print(result["record"].text, result["score"]) + finally: + await collection.ensure_collection_deleted() + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/python/packages/duckdb/tests/duckdb/conftest.py b/python/packages/duckdb/tests/duckdb/conftest.py new file mode 100644 index 00000000000..9357e6cd9d5 --- /dev/null +++ b/python/packages/duckdb/tests/duckdb/conftest.py @@ -0,0 +1,59 @@ +# Copyright (c) Microsoft. All rights reserved. + +from __future__ import annotations + +from collections.abc import AsyncIterator +from typing import Any + +import pytest +from agent_framework import VectorStoreCollectionDefinition, VectorStoreField + +from agent_framework_duckdb import DuckDBCollection, DuckDBStore + + +@pytest.fixture +def definition() -> VectorStoreCollectionDefinition: + return VectorStoreCollectionDefinition( + [ + VectorStoreField("key", name="id", type_="str", storage_name='record "id"'), + VectorStoreField("data", name="text", type_="str", storage_name='body "text"'), + VectorStoreField("data", name="priority", type_="int"), + VectorStoreField("data", name="active", type_="bool"), + VectorStoreField("data", name="tags", type_="list"), + VectorStoreField("vector", name="embedding", type_="float", storage_name="dense vector", dimensions=3), + ], + collection_name="documents", + ) + + +@pytest.fixture +def record_factory(): + def make( + record_id: str = "one", + *, + text: str | None = "DuckDB stores vectors", + priority: int | None = 1, + active: bool = True, + tags: list[Any] | None = None, + embedding: list[float] | None = None, + ) -> dict[str, Any]: + return { + "id": record_id, + "text": text, + "priority": priority, + "active": active, + "tags": ["example", True] if tags is None else tags, + "embedding": [1.0, 0.0, 0.0] if embedding is None else embedding, + } + + return make + + +@pytest.fixture +async def collection( + tmp_path, definition: VectorStoreCollectionDefinition +) -> AsyncIterator[DuckDBCollection[Any, Any]]: + async with DuckDBStore(connection_string=str(tmp_path / "vectors.duckdb")) as store: + selected = store.get_collection(dict, definition=definition) + await selected.ensure_collection_exists() + yield selected diff --git a/python/packages/duckdb/tests/duckdb/test_collection.py b/python/packages/duckdb/tests/duckdb/test_collection.py new file mode 100644 index 00000000000..d5d0b4933c7 --- /dev/null +++ b/python/packages/duckdb/tests/duckdb/test_collection.py @@ -0,0 +1,513 @@ +# Copyright (c) Microsoft. All rights reserved. + +from __future__ import annotations + +import asyncio +import math +import threading +import time +from dataclasses import dataclass +from datetime import date, datetime, timezone +from typing import Annotated, Any, cast +from unittest.mock import AsyncMock, MagicMock +from uuid import UUID + +import duckdb +import pytest +from agent_framework import ( + Embedding, + Filter, + FilterGroup, + GeneratedEmbeddings, + VectorStoreCollectionDefinition, + VectorStoreField, + vectorstoremodel, +) +from agent_framework.exceptions import IntegrationException + +from agent_framework_duckdb import DuckDBCollection, DuckDBStore +from agent_framework_duckdb._vector_store import _identifier + + +async def test_default_database_persists_across_closed_stores(definition, record_factory, monkeypatch, tmp_path): + monkeypatch.chdir(tmp_path) + monkeypatch.delenv("DUCKDB_CONNECTION_STRING", raising=False) + filename = tmp_path / "agent-framework.duckdb" + async with DuckDBStore() as store: + assert not filename.exists() + collection = store.get_collection(dict, definition=definition) + await collection.ensure_collection_exists() + assert await collection.upsert([record_factory()], generate_vectors=False) == ["one"] + assert filename.is_file() + async with DuckDBStore() as reopened: + collection = reopened.get_collection(dict, definition=definition) + assert (await collection.get(["one"]))[0]["text"] == "DuckDB stores vectors" + assert (await collection.get(["one"], include_vectors=True))[0]["embedding"] == [1.0, 0.0, 0.0] + + +async def test_lifecycle_and_direct_collection(definition, tmp_path): + path = str(tmp_path / "collections.duckdb") + async with DuckDBCollection(dict, connection_string=path, definition=definition) as collection: + assert not await collection.collection_exists() + await collection.ensure_collection_exists() + await collection.ensure_collection_exists() + assert await collection.collection_exists() + async with DuckDBStore(connection_string=path) as store: + assert await store.collection_exists("documents") + assert "documents" in await store.list_collection_names() + await store.ensure_collection_deleted("documents") + assert not await store.collection_exists("documents") + await store.ensure_collection_deleted("documents") + + +async def test_batch_crud_preserves_key_order_and_upserts(collection: DuckDBCollection[Any, Any], record_factory): + first = record_factory("first", priority=1) + second = record_factory("second", text="second text", priority=2, embedding=[0.0, 1.0, 0.0]) + assert await collection.upsert([first, second], generate_vectors=False) == ["first", "second"] + await collection.upsert([record_factory("first", priority=3)], generate_vectors=False) + records = await collection.get(["second", "missing", "first", "second"], include_vectors=True) + assert [item["id"] for item in records] == ["second", "first", "second"] + assert records[1]["priority"] == 3 + assert records[0]["embedding"] == [0.0, 1.0, 0.0] + assert records[1]["tags"] == ["example", True] + assert [item["id"] for item in await collection.get(order_by={"priority": False})] == ["first", "second"] + await collection.delete(["missing", "first", "second"]) + assert await collection.get(["first", "second"]) == [] + assert await collection.upsert([], generate_vectors=False) == [] + assert await collection.get([]) == [] + await collection.delete([]) + + +async def test_typed_model_roundtrips_and_searches(tmp_path): + @vectorstoremodel(collection_name="articles") + @dataclass + class Article: + id: Annotated[str, VectorStoreField("key", storage_name="article_id")] + text: Annotated[str, VectorStoreField("data")] + embedding: Annotated[list[float] | None, VectorStoreField("vector", dimensions=3)] = None + + async with DuckDBStore(connection_string=str(tmp_path / "typed.duckdb")) as store: + collection = store.get_collection(Article) + await collection.ensure_collection_exists() + assert await collection.upsert([Article("one", "typed", [1, 0, 0])], generate_vectors=False) == ["one"] + assert await collection.get(["one"]) == [Article("one", "typed", None)] + assert await collection.get(["one"], include_vectors=True) == [Article("one", "typed", [1, 0, 0])] + results = [result async for result in await collection.search(vector=[1, 0, 0])] + assert results[0]["record"] == Article("one", "typed", None) + assert results[0]["score"] == pytest.approx(0) + + +async def test_typed_uuid_date_datetime_bytes_json_and_float32_vectors(tmp_path): + @vectorstoremodel(collection_name="typed_values") + @dataclass + class TypedRecord: + id: Annotated[UUID, VectorStoreField("key")] + created: Annotated[datetime, VectorStoreField("data")] + event_date: Annotated[date, VectorStoreField("data")] + attachment: Annotated[bytes, VectorStoreField("data")] + payload: Annotated[dict[str, Any], VectorStoreField("data")] + reading: Annotated[float, VectorStoreField("data")] + embedding: Annotated[ + list[float] | None, + VectorStoreField("vector", type_="float32", dimensions=3, distance_function="euclidean_distance"), + ] = None + + original = TypedRecord( + UUID("a16b0f43-39b2-4c5a-a0fa-9a7f266685c6"), + datetime(2026, 9, 23, 9, 30, tzinfo=timezone.utc), + date(2026, 9, 23), + b"\x00\xff", + {"labels": [True, 3, "mixed"]}, + 1.25, + [0.1, 0.2, 0.3], + ) + async with DuckDBCollection(TypedRecord, connection_string=str(tmp_path / "types.duckdb")) as collection: + await collection.ensure_collection_exists() + assert await collection.upsert([original], generate_vectors=False) == [original.id] + restored = (await collection.get([original.id], include_vectors=True))[0] + assert restored.id == original.id + assert restored.created == original.created + assert restored.event_date == original.event_date + assert restored.attachment == original.attachment + assert restored.payload == original.payload + assert restored.reading == original.reading + assert restored.embedding == pytest.approx(original.embedding) + assert len(await collection.get(filter=Filter("event_date", "eq", original.event_date))) == 1 + assert len(await collection.get(filter=Filter("created", "gte", original.created))) == 1 + assert len(await collection.get(filter=Filter("id", "eq", str(original.id)))) == 1 + assert [result async for result in await collection.search(vector=[0, 0, 0])] + + +async def test_local_embeddings_are_generated_for_writes_and_value_search(definition, record_factory, tmp_path): + generator = MagicMock( + get_embeddings=AsyncMock(return_value=GeneratedEmbeddings([Embedding(vector=[1.0, 0.0, 0.0])])) + ) + async with DuckDBStore( + connection_string=str(tmp_path / "generated-vectors.duckdb"), embedding_generator=generator + ) as store: + collection = store.get_collection(dict, definition=definition) + await collection.ensure_collection_exists() + await collection.upsert([record_factory(embedding=[0.0, 0.0, 0.0])]) + assert (await collection.get(["one"], include_vectors=True))[0]["embedding"] == [1.0, 0.0, 0.0] + found = [item async for item in await collection.search("query text")] + assert [item["record"]["id"] for item in found] == ["one"] + assert generator.get_embeddings.await_count == 2 + assert generator.get_embeddings.await_args_list[0].kwargs["options"] == {"dimensions": 3} + + +async def test_search_selects_requested_vector_field(tmp_path): + definition = VectorStoreCollectionDefinition( + [ + VectorStoreField("key", name="id", type_="str"), + VectorStoreField("vector", name="primary", type_="float", dimensions=3), + VectorStoreField("vector", name="secondary", type_="float", dimensions=3), + ], + collection_name="two_vectors", + ) + async with DuckDBCollection( + dict, definition=definition, connection_string=str(tmp_path / "two-vectors.duckdb") + ) as collection: + await collection.ensure_collection_exists() + await collection.upsert( + [ + {"id": "a", "primary": [1, 0, 0], "secondary": [0, 1, 0]}, + {"id": "b", "primary": [0, 1, 0], "secondary": [1, 0, 0]}, + ], + generate_vectors=False, + ) + primary = [item async for item in await collection.search(vector=[1, 0, 0], top=1)] + secondary = [ + item + async for item in await collection.search( + vector=[1, 0, 0], vector_property_name="secondary", include_vectors=True, top=1 + ) + ] + assert primary[0]["record"]["id"] == "a" + assert secondary[0]["record"]["id"] == "b" + assert secondary[0]["record"]["primary"] == [0.0, 1.0, 0.0] + + +async def test_generated_string_and_uuid_keys(tmp_path): + for key_type in ("str", "UUID"): + definition = VectorStoreCollectionDefinition( + [ + VectorStoreField("key", name="id", type_=key_type, is_auto_generated=True), + VectorStoreField("data", name="text", type_="str"), + ], + collection_name=f"generated_{key_type}", + ) + async with DuckDBCollection( + dict, definition=definition, connection_string=str(tmp_path / f"{key_type}.duckdb") + ) as collection: + await collection.ensure_collection_exists() + keys = await collection.upsert([{"text": "first"}, {"id": None, "text": "second"}]) + assert len(set(keys)) == 2 + assert all(isinstance(key, str if key_type == "str" else UUID) for key in keys) + assert [item["text"] for item in await collection.get(keys)] == ["first", "second"] + + +async def test_scalar_filters_and_ordering(collection: DuckDBCollection[Any, Any], record_factory): + await collection.upsert( + [ + record_factory("a", text="alpha%_!", priority=3), + record_factory("b", text="alpha", priority=1, active=False, embedding=[0, 1, 0]), + record_factory("c", text=None, priority=None, embedding=None), + ], + generate_vectors=False, + ) + + async def matching(expression) -> list[str]: + return [item["id"] for item in await collection.get(filter=expression, order_by={"priority": False}, top=10)] + + assert await matching(Filter("text", "contains_text", "%_!")) == ["a"] + assert await matching(Filter("text", "starts_with", "alpha")) == ["a", "b"] + assert await matching(Filter("text", "ends_with", "alpha")) == ["b"] + assert await matching(Filter("text", "is_null")) == ["c"] + assert await matching(Filter("text", "is_not_null")) == ["a", "b"] + assert await matching(Filter("text", "exists")) == ["a", "b", "c"] + assert await matching(Filter("priority", "gte", 2)) == ["a"] + assert await matching(Filter("priority", "between", [1, 3])) == ["a", "b"] + assert await matching(Filter("priority", "in", [None, 1])) == ["b"] + assert await matching(Filter("priority", "not_in", [None, 1])) == ["a"] + assert await matching(Filter("priority", "in", [])) == [] + assert await matching(Filter("priority", "not_in", [])) == ["a", "b"] + assert await matching(Filter("active", "eq", 1)) == [] + assert await matching(Filter("active", "ne", 1)) == ["a", "b", "c"] + assert await matching( + FilterGroup( + "and", + [ + Filter("text", "starts_with", "alpha"), + FilterGroup("not", [Filter("active", "eq", False)]), + ], + ) + ) == ["a"] + assert await matching(FilterGroup("not", [Filter("priority", "gt", 2)])) == ["b", "c"] + assert [item["id"] for item in await collection.get(top=1, skip=1, order_by={"priority": False})] == ["b"] + + +@pytest.mark.parametrize( + "expression,error", + [ + (Filter("tags", "eq", ["example"]), NotImplementedError), + (Filter("tags", "contains", "example"), NotImplementedError), + (Filter("tags.nested", "eq", "example"), NotImplementedError), + (Filter("embedding", "is_null"), NotImplementedError), + (Filter("text", "duckdb.raw", "anything"), NotImplementedError), + (Filter("priority", "starts_with", "1"), TypeError), + (Filter("active", "gt", 1), NotImplementedError), + (Filter("priority", "gt", True), TypeError), + ], +) +async def test_unsupported_filters_fail_loudly(collection: DuckDBCollection[Any, Any], expression, error): + with pytest.raises(error): + await collection.get(filter=expression) + + +async def test_invalid_order_and_operation_options_fail_loudly(collection: DuckDBCollection[Any, Any]): + for order_by in ({"tags": True}, {"embedding": False}, {"text.nested": True}): + with pytest.raises(NotImplementedError): + await collection.get(order_by=order_by) + with pytest.raises(ValueError, match="Unknown"): + await collection.get(order_by={"missing": True}) + with pytest.raises(TypeError, match="booleans"): + await collection.get(order_by={"text": cast(Any, "ASC; DROP TABLE documents")}) + with pytest.raises(ValueError, match="order_by"): + await collection.get(["one"], order_by={"text": True}) + with pytest.raises(NotImplementedError, match="Unsupported DuckDB operation"): + await collection.get(operation_options={"timeout": 4}) + with pytest.raises(NotImplementedError, match="Unsupported DuckDB operation"): + await collection.ensure_collection_exists(operation_options={"index": True}) + + +@pytest.mark.parametrize( + "metric,first_score,threshold", + [ + ("DEFAULT", 0.0, 0.1), + ("cosine_distance", 0.0, 0.1), + ("cosine_similarity", 1.0, 0.5), + ("euclidean_distance", 0.0, 0.1), + ("dot_prod", 1.0, 0.5), + ("negative_dot_prod", -1.0, -0.5), + ], +) +async def test_metric_ranking_threshold_filter_and_paging(tmp_path, metric, first_score, threshold): + definition = VectorStoreCollectionDefinition( + [ + VectorStoreField("key", name="id", type_="str"), + VectorStoreField("data", name="category", type_="str"), + VectorStoreField("vector", name="embedding", type_="float", dimensions=3, distance_function=metric), + ], + collection_name="scores", + ) + async with DuckDBCollection( + dict, definition=definition, connection_string=str(tmp_path / f"{metric}.duckdb") + ) as collection: + await collection.ensure_collection_exists() + await collection.upsert( + [ + {"id": "a", "category": "keep", "embedding": [1, 0, 0]}, + {"id": "b", "category": "keep", "embedding": [0, 1, 0]}, + {"id": "c", "category": "exclude", "embedding": [1, 0, 0]}, + {"id": "d", "category": "keep", "embedding": None}, + ], + generate_vectors=False, + ) + found = [ + item + async for item in await collection.search( + vector=[1, 0, 0], filter=Filter("category", "eq", "keep"), score_threshold=threshold + ) + ] + assert [item["record"]["id"] for item in found] == ["a"] + assert found[0]["score"] == pytest.approx(first_score) + unfiltered = [item async for item in await collection.search(vector=[1, 0, 0], top=1, skip=1)] + assert unfiltered[0]["record"]["id"] == "c" + with pytest.raises(ValueError, match="finite"): + await collection.search(vector=[1, 0, 0], score_threshold=math.nan) + + +async def test_zero_and_invalid_vectors_and_unavailable_search_types( + collection: DuckDBCollection[Any, Any], record_factory +): + await collection.upsert([record_factory()], generate_vectors=False) + for vector in ([0, 0, 0], [1, math.inf, 0], [1, True, 0], b"binary"): + with pytest.raises((TypeError, ValueError)): + await collection.search(vector=vector) + with pytest.raises(ValueError, match="dimensions"): + await collection.upsert([record_factory(embedding=[1, 0])], generate_vectors=False) + assert len(await collection.get(["one"])) == 1 + with pytest.raises(NotImplementedError, match="vectorize"): + await collection.search("unembedded text") + with pytest.raises(NotImplementedError, match="keyword_hybrid"): + await collection.search("query", search_type="keyword_hybrid") + with pytest.raises(NotImplementedError, match="hybrid"): + await collection.search(vector=[1, 0, 0], additional_property_name="text") + with pytest.raises(NotImplementedError, match="Unsupported DuckDB operation"): + await collection.search(vector=[1, 0, 0], operation_options={"exact": False}) + + +async def test_existing_zero_cosine_vectors_are_skipped_without_failing_search(): + connection = duckdb.connect(":memory:") + definition = VectorStoreCollectionDefinition( + [ + VectorStoreField("key", name="id", type_="str"), + VectorStoreField("vector", name="embedding", type_="float", dimensions=3), + ], + collection_name="external_rows", + ) + try: + async with DuckDBStore(client=connection) as store: + collection = store.get_collection(dict, definition=definition) + await collection.ensure_collection_exists() + await collection.upsert([{"id": "valid", "embedding": [1, 0, 0]}], generate_vectors=False) + connection.execute("INSERT INTO external_rows VALUES ('zero', [0, 0, 0])") + found = [item async for item in await collection.search(vector=[1, 0, 0])] + assert [item["record"]["id"] for item in found] == ["valid"] + finally: + connection.close() + + +def test_model_options_rejected_before_connecting(definition): + def model(*, key_type="str", index_kind="flat", distance="DEFAULT", data_indexed=False, generated=False): + return VectorStoreCollectionDefinition( + [ + VectorStoreField("key", name="id", type_=key_type, is_auto_generated=generated), + VectorStoreField("data", name="text", type_="str", is_indexed=data_indexed), + VectorStoreField( + "vector", + name="embedding", + type_="float", + dimensions=3, + index_kind=index_kind, + distance_function=distance, + ), + ], + collection_name="rows", + ) + + for invalid_definition in ( + model(key_type="float"), + model(key_type="int", generated=True), + model(data_indexed=True), + model(index_kind="hnsw"), + model(distance="hamming"), + ): + with pytest.raises((NotImplementedError, ValueError)): + DuckDBCollection(dict, definition=invalid_definition) + with pytest.raises(ValueError, match="ignoring case"): + DuckDBCollection( + dict, + definition=VectorStoreCollectionDefinition( + [ + VectorStoreField("key", name="id", type_="str"), + VectorStoreField("data", name="text", storage_name="ID", type_="str"), + ], + collection_name="rows", + ), + ) + with pytest.raises(ValueError, match="NUL"): + DuckDBCollection(dict, definition=definition, collection_name="unsafe\0name") + + +async def test_all_identifiers_and_values_are_safe(tmp_path): + assert _identifier('foo"; DROP TABLE other; --') == '"foo""; DROP TABLE other; --"' + definition = VectorStoreCollectionDefinition( + [ + VectorStoreField("key", name="id", type_="str"), + VectorStoreField("data", name="text", type_="str", storage_name='body "; DROP TABLE safe; --'), + ], + collection_name='unsafe"; DROP TABLE safe; --', + ) + payload = "' OR 1=1; DROP TABLE safe; --%_!" + async with DuckDBStore(connection_string=str(tmp_path / "safe.duckdb")) as store: + unsafe = store.get_collection(dict, definition=definition) + safe = store.get_collection( + dict, + definition=VectorStoreCollectionDefinition( + [VectorStoreField("key", name="id", type_="str")], collection_name="safe" + ), + ) + await safe.ensure_collection_exists() + await unsafe.ensure_collection_exists() + await unsafe.upsert([{"id": "x", "text": payload}], generate_vectors=False) + assert [record["id"] for record in await unsafe.get(filter=Filter("text", "eq", payload))] == ["x"] + assert [record["id"] for record in await unsafe.get(filter=Filter("text", "contains_text", "%_!"))] == ["x"] + assert await safe.collection_exists() + assert await unsafe.collection_exists() + await unsafe.ensure_collection_deleted() + assert await safe.collection_exists() + + +async def test_owned_upsert_rolls_back_on_driver_failure(tmp_path): + path = str(tmp_path / "transactions.duckdb") + connection = duckdb.connect(path) + connection.execute("CREATE TABLE checked (id VARCHAR PRIMARY KEY, text VARCHAR CHECK (text <> 'fail'))") + connection.close() + definition = VectorStoreCollectionDefinition( + [VectorStoreField("key", name="id", type_="str"), VectorStoreField("data", name="text", type_="str")], + collection_name="checked", + ) + async with DuckDBStore(connection_string=path) as store: + collection = store.get_collection(dict, definition=definition) + with pytest.raises(IntegrationException, match="DuckDB operation failed"): + await collection.upsert([{"id": "one", "text": "ok"}, {"id": "two", "text": "fail"}]) + assert await collection.get(["one", "two"]) == [] + + +async def test_borrowed_connection_preserves_caller_transactions(): + connection = duckdb.connect(":memory:") + definition = VectorStoreCollectionDefinition( + [VectorStoreField("key", name="id", type_="str"), VectorStoreField("data", name="text", type_="str")], + collection_name="checked", + ) + try: + connection.execute("CREATE TABLE checked (id VARCHAR PRIMARY KEY, text VARCHAR CHECK (text <> 'fail'))") + async with DuckDBStore(client=connection) as store: + collection = store.get_collection(dict, definition=definition) + connection.execute("BEGIN TRANSACTION") + await collection.upsert([{"id": "one", "text": "pending"}]) + connection.execute("ROLLBACK") + assert await collection.get(["one"]) == [] + with pytest.raises(IntegrationException, match="DuckDB operation failed"): + await collection.upsert([{"id": "one", "text": "ok"}, {"id": "two", "text": "fail"}]) + assert [item["id"] for item in await collection.get(["one"])] == ["one"] + finally: + assert connection.execute("SELECT 42").fetchone() == (42,) + connection.close() + + +async def test_io_is_offloaded_serial_and_close_drains_cancelled_work(tmp_path): + store = DuckDBStore(connection_string=str(tmp_path / "worker.duckdb")) + started = threading.Event() + completed = threading.Event() + worker_thread_ids: list[int] = [] + + def slow(connection: duckdb.DuckDBPyConnection) -> None: + started.set() + time.sleep(0.1) + worker_thread_ids.append(threading.get_ident()) + connection.execute("CREATE TABLE done (id INTEGER)") + completed.set() + + task = asyncio.create_task(store._client.run(slow)) + try: + assert await asyncio.to_thread(started.wait, 1) + await asyncio.sleep(0.01) + assert not completed.is_set() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + other = await store._client.run(lambda _connection: threading.get_ident()) + assert other == worker_thread_ids[0] != threading.get_ident() + await store.aclose() + assert completed.is_set() + with pytest.raises(RuntimeError, match="closed"): + await store.list_collection_names() + reopened = duckdb.connect(str(tmp_path / "worker.duckdb")) + try: + assert reopened.execute("SELECT count(*) FROM done").fetchone() == (0,) + finally: + reopened.close() + finally: + await store.aclose() diff --git a/python/packages/duckdb/tests/duckdb/test_settings.py b/python/packages/duckdb/tests/duckdb/test_settings.py new file mode 100644 index 00000000000..e5289a336ab --- /dev/null +++ b/python/packages/duckdb/tests/duckdb/test_settings.py @@ -0,0 +1,194 @@ +# Copyright (c) Microsoft. All rights reserved. + +from __future__ import annotations + +from functools import partial +from typing import Any, cast, get_type_hints +from unittest.mock import patch + +import duckdb +import pytest +from agent_framework import SecretString, VectorStoreCollectionDefinition, VectorStoreField, load_settings + +from agent_framework_duckdb import DuckDBCollection, DuckDBSettings, DuckDBStore +from agent_framework_duckdb import _vector_store as module + + +@pytest.fixture(params=["store", "collection"]) +def constructor(request): + if request.param == "store": + return DuckDBStore + definition = VectorStoreCollectionDefinition( + [VectorStoreField("key", name="id", type_="str")], collection_name="test" + ) + return partial(DuckDBCollection, dict, definition=definition) + + +def test_public_settings_keep_uri_masked(monkeypatch): + monkeypatch.setenv("DUCKDB_CONNECTION_STRING", "md:db?motherduck_token=private-test-token") + assert get_type_hints(DuckDBSettings) == {"connection_string": SecretString | None} + settings = load_settings(DuckDBSettings, env_prefix="DUCKDB_") + uri = settings["connection_string"] + assert isinstance(uri, SecretString) + assert uri.get_secret_value().endswith("private-test-token") + assert "private-test-token" not in str(uri) + assert "private-test-token" not in repr(settings) + + +@pytest.mark.parametrize("wrapped", [False, True]) +def test_explicit_precedes_selected_file_and_environment(constructor, monkeypatch, tmp_path, wrapped): + monkeypatch.setenv("DUCKDB_CONNECTION_STRING", "environment.duckdb") + env_file = tmp_path / "selected.env" + env_file.write_text("DUCKDB_CONNECTION_STRING=file.duckdb\n", encoding="utf-8") + value = SecretString("explicit.duckdb") if wrapped else "explicit.duckdb" + instance = constructor(connection_string=value, env_file_path=str(env_file)) + assert isinstance(instance._client._address, SecretString) + assert instance._client._address.get_secret_value() == "explicit.duckdb" + assert instance._client._connection is None + + +def test_selected_file_precedes_environment_and_uses_encoding(constructor, monkeypatch, tmp_path): + monkeypatch.setenv("DUCKDB_CONNECTION_STRING", "environment.duckdb") + env_file = tmp_path / "selected.env" + env_file.write_text("DUCKDB_CONNECTION_STRING=file.duckdb\n", encoding="utf-16") + instance = constructor(env_file_path=str(env_file), env_file_encoding="utf-16") + assert instance._client._address.get_secret_value() == "file.duckdb" + + +def test_environment_precedes_persistent_default(constructor, monkeypatch): + monkeypatch.setenv("DUCKDB_CONNECTION_STRING", "from-env.duckdb") + assert constructor()._client._address.get_secret_value() == "from-env.duckdb" + monkeypatch.delenv("DUCKDB_CONNECTION_STRING") + assert constructor()._client._address.get_secret_value() == "agent-framework.duckdb" + + +def test_dotenv_not_discovered_and_no_early_file_created(constructor, monkeypatch, tmp_path): + (tmp_path / ".env").write_text("DUCKDB_CONNECTION_STRING=from-dotenv.duckdb\n", encoding="utf-8") + monkeypatch.chdir(tmp_path) + monkeypatch.delenv("DUCKDB_CONNECTION_STRING", raising=False) + instance = constructor() + assert instance._client._address.get_secret_value() == "agent-framework.duckdb" + assert not (tmp_path / "agent-framework.duckdb").exists() + + +@pytest.mark.parametrize("value", ["", " \t ", SecretString(""), SecretString(" ")]) +def test_empty_explicit_uri_does_not_fall_back(constructor, monkeypatch, value): + monkeypatch.setenv("DUCKDB_CONNECTION_STRING", "valid.duckdb") + with pytest.raises(ValueError, match="must not be empty"): + constructor(connection_string=value) + + +def test_empty_selected_uri_does_not_fall_back(constructor, monkeypatch, tmp_path): + monkeypatch.setenv("DUCKDB_CONNECTION_STRING", "valid.duckdb") + env_file = tmp_path / "selected.env" + env_file.write_text("DUCKDB_CONNECTION_STRING=''\n", encoding="utf-8") + with pytest.raises(ValueError, match="must not be empty"): + constructor(env_file_path=str(env_file)) + + +def test_empty_environment_uri_does_not_fall_back(constructor, monkeypatch): + monkeypatch.setenv("DUCKDB_CONNECTION_STRING", "") + with pytest.raises(ValueError, match="must not be empty"): + constructor() + + +def test_missing_explicit_env_file_is_not_ignored(constructor, tmp_path): + with pytest.raises(FileNotFoundError): + constructor(connection_string="vectors.duckdb", env_file_path=str(tmp_path / "missing")) + + +@pytest.mark.parametrize("value", [123, b"uri", {"password": "private-test-token"}]) +def test_invalid_uri_rejected_without_echoing_value(constructor, value): + with pytest.raises((ValueError, TypeError)) as error: + constructor(connection_string=value) + assert "private-test-token" not in str(error.value) + + +@pytest.mark.parametrize("bad", [{"threads": 2}, {"": "x"}, {"token": object()}]) +def test_invalid_config_rejected_before_connection(constructor, bad): + with pytest.raises(TypeError, match="config"): + constructor(config=cast(Any, bad)) + + +async def test_uri_and_config_forwarded_to_single_duckdb_client_without_remote_credentials(): + fake = duckdb.connect(":memory:") + synthetic_token = SecretString("synthetic-not-a-real-token") + with patch.object(module.duckdb, "connect", return_value=fake) as connect: + async with DuckDBStore( + connection_string=SecretString("md:example"), + config={"motherduck_token": synthetic_token, "threads": "2"}, + ) as store: + assert await store.list_collection_names() == [] + assert "synthetic-not-a-real-token" not in repr(store._client._config) + connect.assert_called_once_with( + database="md:example", + config={"motherduck_token": "synthetic-not-a-real-token", "threads": "2"}, + ) + + +async def test_service_uri_from_environment_reaches_duckdb_without_remote_credentials(monkeypatch): + fake = duckdb.connect(":memory:") + monkeypatch.setenv("DUCKDB_CONNECTION_STRING", "md:from-env") + monkeypatch.setenv("MOTHERDUCK_TOKEN", "synthetic-not-a-real-token") + with patch.object(module.duckdb, "connect", return_value=fake) as connect: + async with DuckDBStore() as store: + assert await store.list_collection_names() == [] + connect.assert_called_once_with(database="md:from-env", config={}) + + +async def test_duckdb_connection_config_is_applied_locally(tmp_path): + async with DuckDBStore(connection_string=str(tmp_path / "configured.duckdb"), config={"threads": "2"}) as store: + assert await store.list_collection_names() == [] + + +@pytest.mark.parametrize( + "kwargs", + [ + {"connection_string": "file.duckdb"}, + {"connection_string": ""}, + {"config": {"threads": "2"}}, + {"env_file_path": "missing.env"}, + {"env_file_encoding": "utf-16"}, + ], +) +def test_injected_client_rejects_conflicting_settings(constructor, kwargs): + connection = duckdb.connect(":memory:") + try: + with ( + patch.object(module, "load_settings", side_effect=AssertionError("settings must not be loaded")), + pytest.raises(ValueError, match="cannot be combined"), + ): + constructor(client=connection, **kwargs) + finally: + connection.close() + + +async def test_injected_connection_remains_caller_owned(constructor, monkeypatch): + monkeypatch.setenv("DUCKDB_CONNECTION_STRING", "md:should-not-be-opened") + connection = duckdb.connect(":memory:") + try: + with patch.object(module, "load_settings", side_effect=AssertionError("injected connection bypasses settings")): + instance = constructor(client=connection) + assert not instance.managed_client + await instance.aclose() + await instance.aclose() + assert connection.execute("SELECT 42").fetchone() == (42,) + with pytest.raises(RuntimeError, match="closed"): + await instance._client.run(lambda client: client.execute("SELECT 1")) + finally: + connection.close() + + +async def test_store_collections_share_one_resolved_connection(tmp_path): + store = DuckDBStore(connection_string=str(tmp_path / "shared.duckdb")) + definition = VectorStoreCollectionDefinition( + [VectorStoreField("key", name="id", type_="str")], collection_name="shared" + ) + child = store.get_collection(dict, definition=definition) + assert child._client is store._client + await child.ensure_collection_exists() + await child.aclose() + assert await child.collection_exists() + await store.aclose() + with pytest.raises(RuntimeError, match="closed"): + await child.collection_exists() diff --git a/python/pyproject.toml b/python/pyproject.toml index 35066954f0e..4c76671f702 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -87,6 +87,7 @@ agent-framework-claude = { workspace = true } agent-framework-copilotstudio = { workspace = true } agent-framework-declarative = { workspace = true } agent-framework-devui = { workspace = true } +agent-framework-duckdb = { workspace = true } agent-framework-foundry = { workspace = true } agent-framework-foundry-hosting = { workspace = true } agent-framework-foundry-local = { workspace = true } diff --git a/python/uv.lock b/python/uv.lock index 67edc89074c..5045af0d8ce 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -45,6 +45,7 @@ members = [ "agent-framework-core", "agent-framework-declarative", "agent-framework-devui", + "agent-framework-duckdb", "agent-framework-foundry", "agent-framework-foundry-hosting", "agent-framework-foundry-local", @@ -609,6 +610,23 @@ requires-dist = [ ] provides-extras = ["dev", "all"] +[[package]] +name = "agent-framework-duckdb" +version = "1.0.0a260923" +source = { editable = "packages/duckdb" } +dependencies = [ + { name = "agent-framework-core", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "duckdb", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "pytz", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, +] + +[package.metadata] +requires-dist = [ + { name = "agent-framework-core", editable = "packages/core" }, + { name = "duckdb", specifier = ">=1.4.1,<1.6" }, + { name = "pytz", specifier = ">=2024.1,<2027" }, +] + [[package]] name = "agent-framework-durabletask" version = "1.0.0b260730" @@ -2138,6 +2156,42 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/32/91/30151a39f7570f448ed84529390628a651d7f27c87d73c9b887f8189695e/docutils-0.23-py3-none-any.whl", hash = "sha256:25d013af9bf23bc1c7b2b093dff4208166c53a94786c9e447808335ef1185fea", size = 634701, upload-time = "2026-05-27T17:40:58.442Z" }, ] +[[package]] +name = "duckdb" +version = "1.5.5" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/7d/19/e57151753576373c6696a12022648546cca6038e8833fda2908ee2342d9b/duckdb-1.5.5.tar.gz", hash = "sha256:72f33ee57ca7595b23957671a2cc7f7fe2be0ecc2d68f63abedcfcaa3a5c1238", size = 18066741, upload-time = "2026-07-22T10:55:17.819Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9c/c2/b62ec24d57bb8df4e24b0b58f7f8facb32f5fdb9f1895aed9e9fcdded168/duckdb-1.5.5-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:1b543841b0ae18a9c982345cfa3987e9c065d3a4b0f067daa473d92d1e65f528", size = 32708371, upload-time = "2026-07-22T10:53:41.642Z" }, + { url = "https://files.pythonhosted.org/packages/8a/ce/769171ba45f0b73632dc3bc3108d891e81dd6c6bbfba630a34a75b4dcc0f/duckdb-1.5.5-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:1a925d06c2a4c3b64553d6cc1aced5028d376d4479bed689a7d47e9b1dccd80a", size = 17343979, upload-time = "2026-07-22T10:53:44.951Z" }, + { url = "https://files.pythonhosted.org/packages/46/59/a8e3384ee916e00d5dcf985194c1511d61978540778a1e96fa47f9fb3e0d/duckdb-1.5.5-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:0c42757cb34722144bd4dfb94b6f336339e7b2468f6813fa7fa9a319ba07bab4", size = 15493704, upload-time = "2026-07-22T10:53:47.912Z" }, + { url = "https://files.pythonhosted.org/packages/6f/1d/9840179c2607b90523a2884a129c4d4e6dbdc1178ba62a976c1043beba88/duckdb-1.5.5-cp311-cp311-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2e72f9e1a4f90a5c8483ad4d540e495bf0834ba61c360b52499a573d7ed62a3f", size = 19366574, upload-time = "2026-07-22T10:53:51.876Z" }, + { url = "https://files.pythonhosted.org/packages/b5/55/f9641a4eebcc2f4df631287d6c3b9ed2eea3b92644f93acbad825e3972b6/duckdb-1.5.5-cp311-cp311-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b9b6f86ed85d4ef5e0211eaebf75d057bd8bb520bba438a95dd0f4e42234bbfe", size = 21477952, upload-time = "2026-07-22T10:53:55.575Z" }, + { url = "https://files.pythonhosted.org/packages/3c/3a/07c3556e37a5c97b95917b029c8fdde4a25fbd76a660bacdac195cf20dcb/duckdb-1.5.5-cp311-cp311-win_amd64.whl", hash = "sha256:9f4287f97ccf0c1f3d471e7115be2b067cbf99627e2d34bffd462dd64703cddc", size = 13156986, upload-time = "2026-07-22T10:53:58.823Z" }, + { url = "https://files.pythonhosted.org/packages/4f/ff/07b48eef2078ca033847e9caa46cc7633b714c5f91ad1ce091c8ca89d792/duckdb-1.5.5-cp311-cp311-win_arm64.whl", hash = "sha256:179633a3fc6296c75d57c69c1e239fa9e5cdcb670fd1dbff88a02663f932905c", size = 14001317, upload-time = "2026-07-22T10:54:01.724Z" }, + { url = "https://files.pythonhosted.org/packages/d6/40/2e05d324400fdaa5656c9f48d6298da421cb034d85e509fa0e6e325cf04b/duckdb-1.5.5-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:d4dd65f8941a604b947e0b9b4b4f7165988e29a23ec0b69b4038520956d9933e", size = 32753858, upload-time = "2026-07-22T10:54:05.514Z" }, + { url = "https://files.pythonhosted.org/packages/79/15/5ceb58ffb5bb8a62b3fd7abb39c41467cdf94850ece02e6d88664dfc75ce/duckdb-1.5.5-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:33db46679b071f108d57139493dee2d37e1f5efcf5c5c039c2969eed11a6c8a7", size = 17368293, upload-time = "2026-07-22T10:54:09.139Z" }, + { url = "https://files.pythonhosted.org/packages/bf/5c/bf02da0b354fe83cca4f95a4fbf762181af466f7d551ab2a093f7698882a/duckdb-1.5.5-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:f0b88535a5d86fdd63dba6ea02ab68c003dfb9e4892b11256ef24c4da208baae", size = 15509131, upload-time = "2026-07-22T10:54:12.228Z" }, + { url = "https://files.pythonhosted.org/packages/ea/a9/5f1f09da421d8e930e0b063d11c1b3f90363f40ede74438cd188afdd13a2/duckdb-1.5.5-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f316eae2323d9a851883fdf2dee91c1f9efe251ab33e14a2272f82a913422ed6", size = 19391959, upload-time = "2026-07-22T10:54:15.551Z" }, + { url = "https://files.pythonhosted.org/packages/4f/98/6549769f158126fa64fd6c1ac2eb59a18282146c939867a3eb31b7c1db07/duckdb-1.5.5-cp312-cp312-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7a6d2d11859d82a936ebdcb30ce3d8a1cbb3e990bff05c12abb9b54c44fa7bd1", size = 21510909, upload-time = "2026-07-22T10:54:19.681Z" }, + { url = "https://files.pythonhosted.org/packages/af/b7/5753b41d3124838f868f9f523362812d9fc45409e9e4dd70dcbb0a25826e/duckdb-1.5.5-cp312-cp312-win_amd64.whl", hash = "sha256:ddfbdb096c11d51ee22492397d342c90a82e62c5d09961477895934d0a25372f", size = 13168544, upload-time = "2026-07-22T10:54:22.789Z" }, + { url = "https://files.pythonhosted.org/packages/5c/28/44b679c7d46245f8398feae7edac959d1b83d4eb143e25b3fce0630b78bd/duckdb-1.5.5-cp312-cp312-win_arm64.whl", hash = "sha256:2725d2b9ace3a4e75d72fc5a239f6a44b502c580edadb8fb2676db772c5f9282", size = 13988684, upload-time = "2026-07-22T10:54:26.003Z" }, + { url = "https://files.pythonhosted.org/packages/47/37/4a38116e7700720fd152c666292214fd3abdf916496991296d8d1f66efbf/duckdb-1.5.5-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:cd98829b67788609017e65c761bd42a5dd0f9129441bed8bda4d6881ccf819f0", size = 32754294, upload-time = "2026-07-22T10:54:29.822Z" }, + { url = "https://files.pythonhosted.org/packages/66/42/7d392f1ba1eee0eaf4ab4c8c7a604bfe3536cd63f979cf5c98798664f807/duckdb-1.5.5-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:feead93c56679b79592d437c62975d39cb67adedffa7592c763baf8160ac7366", size = 17368211, upload-time = "2026-07-22T10:54:33.359Z" }, + { url = "https://files.pythonhosted.org/packages/9f/a5/0a6f4fa60562faa615e55e15bd1953a2f2b17a8edd8105e5cda215e43457/duckdb-1.5.5-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:49c963d9469373d7aba8d750d9ea565ab823e94166efed953f184dd9b169b98c", size = 15509136, upload-time = "2026-07-22T10:54:36.369Z" }, + { url = "https://files.pythonhosted.org/packages/e4/cb/023c89f51978545b9fab318581bba0c457a58e7530d2d933e54ae7d8647c/duckdb-1.5.5-cp313-cp313-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a736217825461732b5442d05a220f3da2e23a0dae114efbf08c9bf171b53098a", size = 19392147, upload-time = "2026-07-22T10:54:39.551Z" }, + { url = "https://files.pythonhosted.org/packages/3e/c5/41bef391fb8b23dbc133c9f2ba016e7a7a8124513d2cc1b430f1897d87e4/duckdb-1.5.5-cp313-cp313-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:078e6a60dd8eedde5832f45422ca5c4a6b8c837aeabd8a56ca0b7d933f588053", size = 21511060, upload-time = "2026-07-22T10:54:42.788Z" }, + { url = "https://files.pythonhosted.org/packages/07/9f/c44dfc1f924ac29b3252dc1b91393c01d009dbfe9f8ed33f10b986151bd1/duckdb-1.5.5-cp313-cp313-win_amd64.whl", hash = "sha256:6826504277dba513c0c5d71d828456c94d729c9d2482f94b2e289f90a9167e28", size = 13168028, upload-time = "2026-07-22T10:54:46.127Z" }, + { url = "https://files.pythonhosted.org/packages/ca/88/591384b2cd59abddd6f5dc175e60374f9abae6064429f0c4402854c10f44/duckdb-1.5.5-cp313-cp313-win_arm64.whl", hash = "sha256:baa9c5702002fabb559ded2a39008f9f421fcbc7237d388b8213eff1e08858de", size = 13989955, upload-time = "2026-07-22T10:54:49.262Z" }, + { url = "https://files.pythonhosted.org/packages/3e/56/12c65bfa2d2605b81981b264788891bcf11ec72227889554cead5d8d13b9/duckdb-1.5.5-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:8e6413dd40facb7b8ab21bd844450cd8f549b29e138635be9cf090ef4d2049e2", size = 32761946, upload-time = "2026-07-22T10:54:53.412Z" }, + { url = "https://files.pythonhosted.org/packages/b9/46/682ce155f17e0d2822d4f13ee3db9ca4b5b7c2da61b841b2629035e1f4bc/duckdb-1.5.5-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:64078acfd16541132ac6e191eb81b2845554444a0305cc1aa581ba107e514aa8", size = 17375069, upload-time = "2026-07-22T10:54:57.269Z" }, + { url = "https://files.pythonhosted.org/packages/39/ce/a24bcbd3289c8f305a430759c5fc12242740b4af3e17f7593f3a34e333d2/duckdb-1.5.5-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:8c11775cc99a447618d5f1840126db17f2652f3eae05529df4f81f40e2df7151", size = 15519791, upload-time = "2026-07-22T10:55:00.681Z" }, + { url = "https://files.pythonhosted.org/packages/d9/76/3a01afbc615c1d418c0de58a6b68ac5ce2a8563232c0464bfbc2ce552398/duckdb-1.5.5-cp314-cp314-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:77bbc1e6ba12e1e06f9020117bdf848627ecfdf36f907550e62e008e6109dece", size = 19398251, upload-time = "2026-07-22T10:55:04.168Z" }, + { url = "https://files.pythonhosted.org/packages/a1/43/3a5e81d1728f4d234c79bfe385808ee7c04834f7c37a4b5c257459c25614/duckdb-1.5.5-cp314-cp314-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fbf0f2d48b43c6c304d00463b463c27ead6c4b01c3c1816b750f728decf71afe", size = 21513851, upload-time = "2026-07-22T10:55:07.864Z" }, + { url = "https://files.pythonhosted.org/packages/91/41/fc7c829172c60ca22485251eab285f4f1a0d87b486a024c726f21471d86e/duckdb-1.5.5-cp314-cp314-win_amd64.whl", hash = "sha256:9dc826c4b50e64f6c4e4d07a3a9cb075ef70ba3899dc43ec5493dc3d7b04b353", size = 13691858, upload-time = "2026-07-22T10:55:11.181Z" }, + { url = "https://files.pythonhosted.org/packages/e1/2c/95d9216b79e9273689d7ebce125a54503ed0c9bd7da931f0265888e99779/duckdb-1.5.5-cp314-cp314-win_arm64.whl", hash = "sha256:63e48d4b74b15aeacd688976432a7225163df8c226eddeb8536bba2d4d4ff433", size = 14470180, upload-time = "2026-07-22T10:55:14.445Z" }, +] + [[package]] name = "durabletask" version = "1.10.1" From 84228281b257a7f2bc18050ca99949e6d2bae1c0 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Thu, 24 Sep 2026 09:20:30 +0200 Subject: [PATCH 2/3] fix(python): address DuckDB review blockers --- .github/workflows/python-merge-tests.yml | 27 ++++++++++++++++ .../agent_framework_duckdb/_vector_store.py | 31 +++++++++++++++---- python/packages/duckdb/pyproject.toml | 4 +++ .../duckdb/tests/duckdb/test_collection.py | 20 +++++++++++- .../duckdb/tests/duckdb/test_settings.py | 25 +++++++++++++++ 5 files changed, 100 insertions(+), 7 deletions(-) diff --git a/.github/workflows/python-merge-tests.yml b/.github/workflows/python-merge-tests.yml index 2d75ad02844..f4850f05143 100644 --- a/.github/workflows/python-merge-tests.yml +++ b/.github/workflows/python-merge-tests.yml @@ -39,6 +39,7 @@ jobs: foundryChanged: ${{ steps.filter.outputs.foundry }} foundryHostingChanged: ${{ steps.filter.outputs.foundry_hosting }} cosmosChanged: ${{ steps.filter.outputs.cosmos }} + duckdbChanged: ${{ steps.filter.outputs.duckdb }} qdrantChanged: ${{ steps.filter.outputs.qdrant }} mongodbChanged: ${{ steps.filter.outputs.mongodb }} redisChanged: ${{ steps.filter.outputs.redis }} @@ -88,6 +89,9 @@ jobs: - 'python/packages/foundry_hosting/**' cosmos: - 'python/packages/azure-cosmos/**' + duckdb: + - 'python/packages/duckdb/**' + - '.github/workflows/python-merge-tests.yml' qdrant: - 'python/packages/qdrant/**' - '.github/workflows/python-merge-tests.yml' @@ -115,6 +119,28 @@ jobs: - name: not python tests if: steps.filter.outputs.python != 'true' run: echo "NOT python file" + + # The aggregate unit-test job is skipped for pull requests, so run DuckDB tests directly. + python-tests-duckdb: + name: Python Tests - DuckDB + needs: paths-filter + if: > + github.event_name == 'pull_request' && + needs.paths-filter.outputs.duckdbChanged == 'true' + runs-on: ubuntu-latest + defaults: + run: + working-directory: python + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + - name: Set up python and install the project + uses: ./.github/actions/python-setup + with: + python-version: ${{ env.UV_PYTHON }} + os: ${{ runner.os }} + - name: Test DuckDB package + run: uv run poe test -P duckdb + python-tests-qdrant: name: Python Integration Tests - Qdrant needs: paths-filter @@ -947,6 +973,7 @@ jobs: runs-on: ubuntu-latest needs: [ + python-tests-duckdb, python-tests-unit, python-tests-openai, python-tests-azure-openai, diff --git a/python/packages/duckdb/agent_framework_duckdb/_vector_store.py b/python/packages/duckdb/agent_framework_duckdb/_vector_store.py index c5ff2987c00..2f20ecd4b79 100644 --- a/python/packages/duckdb/agent_framework_duckdb/_vector_store.py +++ b/python/packages/duckdb/agent_framework_duckdb/_vector_store.py @@ -237,15 +237,22 @@ def __init__( self._close_future: Future[None] | None = None def _execute(self, operation: Callable[[duckdb.DuckDBPyConnection], ResultT]) -> ResultT: - try: - if self._connection is None: - if self._address is None: - raise RuntimeError("A DuckDB connection or database path is required.") - self._connection = duckdb.connect( + connection = self._connection + if connection is None: + if self._address is None: + raise RuntimeError("A DuckDB connection or database path is required.") + try: + connection = duckdb.connect( database=self._address.get_secret_value(), config={key: value.get_secret_value() for key, value in self._config.items()}, ) - return operation(self._connection) + except duckdb.Error: + connection = None + if connection is None: + raise IntegrationException("DuckDB connection failed. Verify the connection string and configuration.") + self._connection = connection + try: + return operation(connection) except duckdb.Error as exc: raise IntegrationException("DuckDB operation failed; inspect the chained driver exception.") from exc @@ -744,6 +751,18 @@ async def list_collection_names(self, *, operation_options: Mapping[str, Any] | lambda connection: [str(row[0]) for row in connection.execute(_LIST_TABLES).fetchall()] ) + async def collection_exists( + self, + collection_name: str, + *, + operation_options: Mapping[str, Any] | None = None, + ) -> bool: + """Check for a table using DuckDB's case-insensitive identifier semantics.""" + _check_options(operation_options) + return await self._client.run( + lambda connection: connection.execute(_TABLE_EXISTS, [collection_name]).fetchone() is not None + ) + async def _inner_ensure_collection_deleted( self, collection_name: str, diff --git a/python/packages/duckdb/pyproject.toml b/python/packages/duckdb/pyproject.toml index f127b550840..3439921d2d8 100644 --- a/python/packages/duckdb/pyproject.toml +++ b/python/packages/duckdb/pyproject.toml @@ -52,6 +52,10 @@ omit = ["**/__init__.py"] executor.type = "uv" include = "../../shared_tasks.toml" +[tool.poe.tasks.test] +help = "Run DuckDB connector unit tests." +cmd = 'pytest -m "not integration" --cov=agent_framework_duckdb --cov-report=term-missing:skip-covered tests' + [build-system] requires = ["flit-core>=3.11,<4.0"] build-backend = "flit_core.buildapi" diff --git a/python/packages/duckdb/tests/duckdb/test_collection.py b/python/packages/duckdb/tests/duckdb/test_collection.py index d5d0b4933c7..d92580d9530 100644 --- a/python/packages/duckdb/tests/duckdb/test_collection.py +++ b/python/packages/duckdb/tests/duckdb/test_collection.py @@ -60,6 +60,23 @@ async def test_lifecycle_and_direct_collection(definition, tmp_path): await store.ensure_collection_deleted("documents") +async def test_store_collection_checks_and_deletion_are_case_insensitive(tmp_path): + definition = VectorStoreCollectionDefinition( + [VectorStoreField("key", name="id", type_="str")], collection_name="MiXeD" + ) + async with DuckDBStore(connection_string=str(tmp_path / "mixed-case.duckdb")) as store: + collection = store.get_collection(dict, definition=definition) + await collection.ensure_collection_exists() + assert await store.list_collection_names() == ["MiXeD"] + assert await store.collection_exists("mixed") + assert await store.collection_exists("MIXED") + + await store.ensure_collection_deleted("mIxEd") + + assert not await store.collection_exists("MiXeD") + assert "MiXeD" not in await store.list_collection_names() + + async def test_batch_crud_preserves_key_order_and_upserts(collection: DuckDBCollection[Any, Any], record_factory): first = record_factory("first", priority=1) second = record_factory("second", text="second text", priority=2, embedding=[0.0, 1.0, 0.0]) @@ -450,8 +467,9 @@ async def test_owned_upsert_rolls_back_on_driver_failure(tmp_path): ) async with DuckDBStore(connection_string=path) as store: collection = store.get_collection(dict, definition=definition) - with pytest.raises(IntegrationException, match="DuckDB operation failed"): + with pytest.raises(IntegrationException, match="DuckDB operation failed") as error: await collection.upsert([{"id": "one", "text": "ok"}, {"id": "two", "text": "fail"}]) + assert isinstance(error.value.__cause__, duckdb.Error) assert await collection.get(["one", "two"]) == [] diff --git a/python/packages/duckdb/tests/duckdb/test_settings.py b/python/packages/duckdb/tests/duckdb/test_settings.py index e5289a336ab..e248d3a7809 100644 --- a/python/packages/duckdb/tests/duckdb/test_settings.py +++ b/python/packages/duckdb/tests/duckdb/test_settings.py @@ -2,6 +2,7 @@ from __future__ import annotations +import traceback from functools import partial from typing import Any, cast, get_type_hints from unittest.mock import patch @@ -9,6 +10,7 @@ import duckdb import pytest from agent_framework import SecretString, VectorStoreCollectionDefinition, VectorStoreField, load_settings +from agent_framework.exceptions import IntegrationException from agent_framework_duckdb import DuckDBCollection, DuckDBSettings, DuckDBStore from agent_framework_duckdb import _vector_store as module @@ -141,6 +143,29 @@ async def test_duckdb_connection_config_is_applied_locally(tmp_path): assert await store.list_collection_names() == [] +async def test_connection_failure_does_not_chain_unwrapped_secrets(): + address_secret = "synthetic-address-secret" + config_secret = "synthetic-config-secret" + driver_error = duckdb.IOException(f"Could not connect to {address_secret} with token {config_secret}") + store = DuckDBStore( + connection_string=SecretString(address_secret), + config={"motherduck_token": SecretString(config_secret)}, + ) + try: + with ( + patch.object(module.duckdb, "connect", side_effect=driver_error), + pytest.raises(IntegrationException, match="DuckDB connection failed") as error, + ): + await store.list_collection_names() + formatted = "".join(traceback.format_exception(type(error.value), error.value, error.value.__traceback__)) + assert error.value.__cause__ is None + assert error.value.__context__ is None + assert address_secret not in formatted + assert config_secret not in formatted + finally: + await store.aclose() + + @pytest.mark.parametrize( "kwargs", [ From fea5580aa441498872f0fa4151ead63b8ba51b6b Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Thu, 24 Sep 2026 11:58:08 +0200 Subject: [PATCH 3/3] fix(python): match DuckDB identifier semantics and run tests for core changes --- .github/workflows/python-merge-tests.yml | 3 +- .../agent_framework_duckdb/_vector_store.py | 10 ++++- .../duckdb/tests/duckdb/test_collection.py | 45 +++++++++++++++++++ 3 files changed, 55 insertions(+), 3 deletions(-) diff --git a/.github/workflows/python-merge-tests.yml b/.github/workflows/python-merge-tests.yml index f4850f05143..8b7573d72a8 100644 --- a/.github/workflows/python-merge-tests.yml +++ b/.github/workflows/python-merge-tests.yml @@ -126,7 +126,8 @@ jobs: needs: paths-filter if: > github.event_name == 'pull_request' && - needs.paths-filter.outputs.duckdbChanged == 'true' + (needs.paths-filter.outputs.duckdbChanged == 'true' || + needs.paths-filter.outputs.coreChanged == 'true') runs-on: ubuntu-latest defaults: run: diff --git a/python/packages/duckdb/agent_framework_duckdb/_vector_store.py b/python/packages/duckdb/agent_framework_duckdb/_vector_store.py index 2f20ecd4b79..fd7d5b5262e 100644 --- a/python/packages/duckdb/agent_framework_duckdb/_vector_store.py +++ b/python/packages/duckdb/agent_framework_duckdb/_vector_store.py @@ -59,13 +59,19 @@ "negative_dot_prod": ("-array_inner_product", False), } _ORDER_OPERATORS = {"gt": ">", "gte": ">=", "lt": "<", "lte": "<="} +_ASCII_FOLD = str.maketrans("ABCDEFGHIJKLMNOPQRSTUVWXYZ", "abcdefghijklmnopqrstuvwxyz") _TABLE_SCOPE = ( "FROM information_schema.tables " "WHERE table_catalog = current_database() AND table_schema = current_schema() " "AND table_type = 'BASE TABLE'" ) _LIST_TABLES = f"SELECT table_name {_TABLE_SCOPE} ORDER BY table_name" -_TABLE_EXISTS = f"SELECT 1 {_TABLE_SCOPE} AND lower(table_name) = lower(?)" +# DuckDB folds only ASCII identifier case; C collation avoids a connection's Unicode NOCASE setting. +_TABLE_EXISTS = ( + f"SELECT 1 {_TABLE_SCOPE} " + "AND (translate(table_name, 'ABCDEFGHIJKLMNOPQRSTUVWXYZ', 'abcdefghijklmnopqrstuvwxyz') COLLATE \"C\") = " + "(translate(?, 'ABCDEFGHIJKLMNOPQRSTUVWXYZ', 'abcdefghijklmnopqrstuvwxyz') COLLATE \"C\")" +) class DuckDBSettings(TypedDict, total=False): @@ -415,7 +421,7 @@ def __init__( ) self._table = _identifier(self.collection_name) names = [field.storage_name or field.name for field in self.definition.fields] - if len({name.casefold() for name in names}) != len(names): + if len({name.translate(_ASCII_FOLD) for name in names}) != len(names): raise ValueError("DuckDB column names must be unique ignoring case.") self._fields = tuple(self.definition.fields) for name in names: diff --git a/python/packages/duckdb/tests/duckdb/test_collection.py b/python/packages/duckdb/tests/duckdb/test_collection.py index d92580d9530..9d099f3b717 100644 --- a/python/packages/duckdb/tests/duckdb/test_collection.py +++ b/python/packages/duckdb/tests/duckdb/test_collection.py @@ -77,6 +77,51 @@ async def test_store_collection_checks_and_deletion_are_case_insensitive(tmp_pat assert "MiXeD" not in await store.list_collection_names() +@pytest.mark.parametrize( + ("table_name", "same_table", "different_table"), + [("ÄBC", "Äbc", "äbc"), ("äbc", "äBC", "Äbc")], +) +@pytest.mark.parametrize("config", [None, {"default_collation": "NOCASE"}]) +async def test_non_ascii_table_existence_matches_duckdb_identifiers( + tmp_path, table_name: str, same_table: str, different_table: str, config: dict[str, str] | None +): + definition = VectorStoreCollectionDefinition( + [VectorStoreField("key", name="id", type_="str")], collection_name=table_name + ) + async with DuckDBStore(connection_string=str(tmp_path / "unicode.duckdb"), config=config) as store: + original = store.get_collection(dict, definition=definition) + same = store.get_collection(dict, definition=definition, collection_name=same_table) + different = store.get_collection(dict, definition=definition, collection_name=different_table) + await original.ensure_collection_exists() + + assert await original.collection_exists() + assert await same.collection_exists() + assert await store.collection_exists(same_table) + assert not await different.collection_exists() + assert not await store.collection_exists(different_table) + await store.ensure_collection_deleted(different_table) + assert await original.collection_exists() + await store.ensure_collection_deleted(same_table) + assert not await original.collection_exists() + + +async def test_non_ascii_columns_are_distinct(tmp_path): + definition = VectorStoreCollectionDefinition( + [ + VectorStoreField("key", name="id", type_="str"), + VectorStoreField("data", name="first", type_="str", storage_name="Ä"), + VectorStoreField("data", name="second", type_="str", storage_name="ä"), + ], + collection_name="unicode_columns", + ) + async with DuckDBCollection( + dict, definition=definition, connection_string=str(tmp_path / "unicode-columns.duckdb") + ) as collection: + await collection.ensure_collection_exists() + await collection.upsert([{"id": "one", "first": "upper", "second": "lower"}]) + assert await collection.get(["one"]) == [{"id": "one", "first": "upper", "second": "lower"}] + + async def test_batch_crud_preserves_key_order_and_upserts(collection: DuckDBCollection[Any, Any], record_factory): first = record_factory("first", priority=1) second = record_factory("second", text="second text", priority=2, embedding=[0.0, 1.0, 0.0])