DuckDB Extension¶

The DuckDB integration provides SQL UDFs for responses, embeddings, prepared tasks, parsing, and token counting. The table below shows how to use the main pandas features from DuckDB.

pandas feature	DuckDB equivalent
`Series.ai.responses()`	`responses_udf()`
`Series.ai.embeddings()`	`embeddings_udf()`
`Series.ai.task()`	`task_udf()`
`Series.ai.parse()` with an explicit model	`parse_udf(..., response_format=Model)`
`Series.ai.parse()` with schema inference	`parse_udf(..., example_table_name=..., example_field_name=...)`
`Series.ai.infer_schema()` / `DataFrame.ai.infer_schema()`	`infer_schema()` on a text column or a view of JSON rows
`Series.ai.count_tokens()`	`count_tokens_udf()`
`Series.ai.extract()`	SQL `STRUCT` field access
`DataFrame.ai.responses()` / `.task()` / `.parse()`	Pass `to_json(t)` from a table alias to the UDF; use a view of those JSON rows for schema inference
`DataFrame.ai.similarity()`	SQL `list_cosine_similarity(a, b)`

The AI-based DataFrame.ai.fillna() builds examples from the whole DataFrame and predicts only missing rows. There is no dedicated DuckDB wrapper for that table-level operation; COALESCE is useful for ordinary SQL defaults but does not perform AI imputation.

The DuckDB UDF, schema, token, similarity, and type implementations live in the openaivec.duckdb_ext package. Persistent caching is a core BatchCache feature; its DuckDB backend remains in openaivec._cache:

from openaivec._cache import DuckDBCacheBackend

backend = DuckDBCacheBackend.of("results.duckdb")
# Pass backend as the cache argument when constructing BatchCache.

Cache table names are single unqualified SQL identifiers matching [A-Za-z_][A-Za-z_0-9]*; reserved words are supported. Keys can be strings, integers, booleans, floats, bytes, or tuples of those types. Keys are encoded with type tags, so 1 and "1" persist independently; other key types raise TypeError. The current typed-key table schema is incompatible with tables created by older releases. Opening an old cache table raises ValueError instead of silently misreading its entries; choose a new table name or explicitly migrate/rebuild the old data. The backend also provides get_many, put_many, and touch_many for bounded bulk SQL operations. get_many does not update LRU order; call touch_many after consuming the fetched keys. Because get_many returns a Python dictionary, requesting distinct keys that compare equal in Python (for example, 1 and True) together raises ValueError rather than silently combining them; individual lookups remain distinct.

Schema inference and parsing¶

import duckdb
from openaivec.duckdb_ext import infer_schema, parse_udf

conn = duckdb.connect()
conn.execute("CREATE TABLE reviews (review VARCHAR)")
conn.executemany(
    "INSERT INTO reviews VALUES (?)",
    [("The camera is excellent. 5 stars.",), ("The battery failed. 1 star.",)],
)

# Infer once and reuse the resulting model and prompt for multiple UDFs.
schema = infer_schema(
    conn,
    instructions="Extract product feedback and rating.",
    example_table_name="reviews",
    example_field_name="review",
)
parse_udf(
    conn,
    "parse_review",
    instructions=schema.inference_prompt,
    response_format=schema.model,
)
conn.sql("SELECT parse_review(review) AS parsed FROM reviews")

# Or infer and register in one call.
parse_udf(
    conn,
    "parse_review_auto",
    instructions="Extract product feedback and rating.",
    example_table_name="reviews",
    example_field_name="review",
)
conn.sql("SELECT parse_review_auto(review) AS parsed FROM reviews")

parse_udf infers the schema when it is registered. The sample excludes SQL NULL values and is bounded by max_examples (default 100). An explicit response_format skips inference. max_retries controls total inference attempts; max_validation_retries controls additional extraction corrections.

Token counts and complete rows¶

from openaivec.duckdb_ext import count_tokens_udf, parse_udf, responses_udf

count_tokens_udf(conn, "count_tokens")
conn.sql("SELECT review, count_tokens(review) AS num_tokens FROM reviews")

responses_udf(conn, "describe_row", instructions="Summarize the complete row.")
conn.sql("SELECT describe_row(to_json(t)) FROM reviews AS t")

# For DataFrame-style automatic parsing, infer from JSON rows and pass the
# same expression to the registered UDF.
conn.execute("CREATE VIEW review_rows AS SELECT to_json(t) AS input FROM reviews AS t")
parse_udf(
    conn,
    "parse_review_row",
    instructions="Extract product feedback and rating from each row.",
    example_table_name="review_rows",
    example_field_name="input",
)
conn.sql("SELECT parse_review_row(to_json(t)) FROM reviews AS t")

count_tokens_udf uses the same configured tiktoken.Encoding as the pandas accessor and returns SQL NULL for SQL NULL input. Structured outputs from responses_udf, task_udf, and parse_udf are native DuckDB STRUCT values; read fields with my_udf(text).field_name or expand a stored STRUCT column with parsed.*.

Similarity search and generated DDL¶

similarity_search(conn, "documents", "queries", top_k=3) returns query_id, query_text, target_text, and score, ordered by query row and descending score. query_id is a 1-based position assigned during the query table scan (not a persistent key or a guaranteed order across scans); duplicate query text still receives an independent top-k ranking. top_k must be a positive integer and is passed as a SQL parameter.

pydantic_to_duckdb_ddl(Model, "schema.table") quotes each table/column identifier, including nested STRUCT fields. Table names may be unqualified or qualified as schema.table or catalog.schema.table; use DuckDB double quotes around a part to include a literal dot (for example, "my.schema".items). Spaces and reserved words are supported. Embedded double quotes in field names are escaped; SQL expressions and statement separators in identifiers are rejected. Both typing.Optional[T] and T | None map to the same DuckDB and Spark type; unions of distinct non-null types are unsupported.

openaivec.duckdb_ext ¶

DuckDB integration for openaivec.

Provides helpers that bridge openaivec's batched AI capabilities with DuckDB:

UDF registration – register responses, embeddings, task, parse and count_tokens functions directly as DuckDB scalar UDFs for SQL queries.
Schema inference – infer a Pydantic response model once from a bounded table sample before registering a typed parse UDF.
Vector similarity – similarity_search performs top-k cosine similarity queries against an embedding table using DuckDB's built-in list_cosine_similarity.
Schema → DDL – pydantic_to_duckdb_ddl converts a Pydantic model to a CREATE TABLE statement for immediate SQL analysis of structured-output results.

Quick Start¶

import duckdb
from openaivec.duckdb_ext import responses_udf, embeddings_udf

conn = duckdb.connect()
responses_udf(conn, "translate", instructions="Translate to French", reasoning={"effort": "none"})
embeddings_udf(conn, "embed")

conn.sql("SELECT translate(review) FROM products")
conn.sql("SELECT text, embed(text) FROM documents")

Functions:¶

infer_schema ¶

infer_schema(
    conn: DuckDBPyConnection,
    instructions: str,
    example_table_name: str,
    example_field_name: str,
    max_examples: int = 100,
    *,
    max_retries: int = 8,
    retry_policy: RetryPolicy | None = None,
    **api_kwargs: Any,
) -> SchemaInferenceOutput

Infer a response schema from a bounded sample of a DuckDB column.

The table is resolved through DuckDB's relation API, which accepts regular, schema-qualified, and quoted table names without interpolating the name in SQL. NULL values are skipped before the sample limit is applied.

Parameters:

Name	Type	Description	Default
`conn`	`DuckDBPyConnection`	Connection containing the example table.	required
`instructions`	`str`	Description of the information to extract.	required
`example_table_name`	`str`	Table or view containing examples.	required
`example_field_name`	`str`	Column containing example text.	required
`max_examples`	`int`	Maximum number of non-NULL examples. Must be positive.	`100`
`max_retries`	`int`	Maximum schema inference attempts. Must be positive.	`8`
`retry_policy`	`RetryPolicy \| None`	Transport retry policy.	`None`
`**api_kwargs`	`Any`	Extra parameters for schema inference API calls.	`{}`

Returns:

Name	Type	Description
`SchemaInferenceOutput`	`SchemaInferenceOutput`	Inferred prompt and Pydantic model.

Raises:

Type	Description
`ValueError`	If limits are invalid or the sample has no non-NULL values.

Example

import duckdb conn = duckdb.connect() conn.sql("CREATE TABLE reviews(text VARCHAR)")

infer_schema(conn, "Extract sentiment", "reviews", "text")¶

Source code in src/openaivec/duckdb_ext/_schema.py

def infer_schema(
    conn: duckdb.DuckDBPyConnection,
    instructions: str,
    example_table_name: str,
    example_field_name: str,
    max_examples: int = 100,
    *,
    max_retries: int = 8,
    retry_policy: RetryPolicy | None = None,
    **api_kwargs: Any,
) -> SchemaInferenceOutput:
    """Infer a response schema from a bounded sample of a DuckDB column.

    The table is resolved through DuckDB's relation API, which accepts regular,
    schema-qualified, and quoted table names without interpolating the name in
    SQL. NULL values are skipped before the sample limit is applied.

    Args:
        conn (duckdb.DuckDBPyConnection): Connection containing the example table.
        instructions (str): Description of the information to extract.
        example_table_name (str): Table or view containing examples.
        example_field_name (str): Column containing example text.
        max_examples (int): Maximum number of non-NULL examples. Must be positive.
        max_retries (int): Maximum schema inference attempts. Must be positive.
        retry_policy (RetryPolicy | None): Transport retry policy.
        **api_kwargs: Extra parameters for schema inference API calls.

    Returns:
        SchemaInferenceOutput: Inferred prompt and Pydantic model.

    Raises:
        ValueError: If limits are invalid or the sample has no non-NULL values.

    Example:
        >>> import duckdb
        >>> conn = duckdb.connect()
        >>> conn.sql("CREATE TABLE reviews(text VARCHAR)")
        >>> # infer_schema(conn, "Extract sentiment", "reviews", "text")
    """
    if max_examples < 1:
        raise ValueError("max_examples must be >= 1")
    if max_retries < 1:
        raise ValueError("max_retries must be >= 1")
    if not isinstance(example_table_name, str) or not example_table_name:
        raise ValueError("example_table_name must be a non-empty string")

    column = _quote_identifier(example_field_name)
    examples = [
        row[0]
        for row in (
            conn.table(example_table_name)
            .filter(f"{column} IS NOT NULL")
            .select(f"CAST({column} AS VARCHAR)")
            .limit(max_examples)
            .fetchall()
        )
    ]
    if not examples:
        raise ValueError("No non-NULL examples found in the selected column")

    inference_input = SchemaInferenceInput(examples=examples, instructions=instructions)
    inferer = CONTAINER.resolve(SchemaInferer)
    return inferer.infer_schema(
        inference_input,
        max_retries=max_retries,
        retry_policy=retry_policy,
        **api_kwargs,
    )

parse_udf ¶

parse_udf(
    conn: DuckDBPyConnection,
    name: str,
    *,
    instructions: str,
    response_format: type[ResponseFormat] | None = None,
    example_table_name: str | None = None,
    example_field_name: str | None = None,
    max_examples: int = 100,
    model_name: str | None = None,
    batch_size: int | None = 64,
    max_concurrency: int = 8,
    multimodal: bool = False,
    max_retries: int = 8,
    max_validation_retries: int = 3,
    retry_policy: RetryPolicy | None = None,
    **api_kwargs: Any,
) -> None

Register a DuckDB parse UDF with an explicit or inferred output schema.

DuckDB fixes UDF return types at registration. If response_format is omitted, the schema is inferred once from example_table_name and example_field_name before registering the UDF. The generated UDF then returns a typed DuckDB STRUCT for each non-NULL input.

Parameters:

Name	Type	Description	Default
`conn`	`DuckDBPyConnection`	Connection on which to register the UDF.	required
`name`	`str`	UDF name visible in SQL.	required
`instructions`	`str`	Description of the information to parse.	required
`response_format`	`type[ResponseFormat] \| None`	`str` or a Pydantic `BaseModel` subclass. If None, infer the schema from examples.	`None`
`example_table_name`	`str \| None`	Source table for schema inference.	`None`
`example_field_name`	`str \| None`	Text column for schema inference.	`None`
`max_examples`	`int`	Maximum non-NULL examples to inspect.	`100`
`model_name`	`str \| None`	Model or deployment name.	`None`
`batch_size`	`int \| None`	Inputs per API batch; None enables auto-tuning.	`64`
`max_concurrency`	`int`	Maximum concurrent API requests.	`8`
`multimodal`	`bool`	Whether file paths and URLs are multimodal inputs.	`False`
`max_retries`	`int`	Maximum schema inference attempts.	`8`
`max_validation_retries`	`int`	Additional extraction corrections.	`3`
`retry_policy`	`RetryPolicy \| None`	Transport retry policy.	`None`
`**api_kwargs`	`Any`	Extra parameters for inference and parsing API calls.	`{}`

Raises:

Type	Description
`ValueError`	If limits are invalid or inference inputs are missing.
`TypeError`	If an explicit response format is unsupported.

Example

import duckdb from pydantic import BaseModel class Parsed(BaseModel): ... label: str conn = duckdb.connect() parse_udf(conn, "parse_label", instructions="Extract label", response_format=Parsed)

Source code in src/openaivec/duckdb_ext/_schema.py

def parse_udf(
    conn: duckdb.DuckDBPyConnection,
    name: str,
    *,
    instructions: str,
    response_format: type[ResponseFormat] | None = None,
    example_table_name: str | None = None,
    example_field_name: str | None = None,
    max_examples: int = 100,
    model_name: str | None = None,
    batch_size: int | None = 64,
    max_concurrency: int = 8,
    multimodal: bool = False,
    max_retries: int = 8,
    max_validation_retries: int = 3,
    retry_policy: RetryPolicy | None = None,
    **api_kwargs: Any,
) -> None:
    """Register a DuckDB parse UDF with an explicit or inferred output schema.

    DuckDB fixes UDF return types at registration. If ``response_format`` is
    omitted, the schema is inferred once from ``example_table_name`` and
    ``example_field_name`` before registering the UDF. The generated UDF then
    returns a typed DuckDB ``STRUCT`` for each non-NULL input.

    Args:
        conn (duckdb.DuckDBPyConnection): Connection on which to register the UDF.
        name (str): UDF name visible in SQL.
        instructions (str): Description of the information to parse.
        response_format (type[ResponseFormat] | None): ``str`` or a Pydantic
            ``BaseModel`` subclass. If None, infer the schema from examples.
        example_table_name (str | None): Source table for schema inference.
        example_field_name (str | None): Text column for schema inference.
        max_examples (int): Maximum non-NULL examples to inspect.
        model_name (str | None): Model or deployment name.
        batch_size (int | None): Inputs per API batch; None enables auto-tuning.
        max_concurrency (int): Maximum concurrent API requests.
        multimodal (bool): Whether file paths and URLs are multimodal inputs.
        max_retries (int): Maximum schema inference attempts.
        max_validation_retries (int): Additional extraction corrections.
        retry_policy (RetryPolicy | None): Transport retry policy.
        **api_kwargs: Extra parameters for inference and parsing API calls.

    Raises:
        ValueError: If limits are invalid or inference inputs are missing.
        TypeError: If an explicit response format is unsupported.

    Example:
        >>> import duckdb
        >>> from pydantic import BaseModel
        >>> class Parsed(BaseModel):
        ...     label: str
        >>> conn = duckdb.connect()
        >>> parse_udf(conn, "parse_label", instructions="Extract label", response_format=Parsed)
    """
    if max_examples < 1:
        raise ValueError("max_examples must be >= 1")
    if max_retries < 1:
        raise ValueError("max_retries must be >= 1")
    if max_validation_retries < 0:
        raise ValueError("max_validation_retries must be >= 0")

    if response_format is None:
        if not example_table_name or not example_field_name:
            raise ValueError("Both example_table_name and example_field_name are required for schema inference")
        schema = infer_schema(
            conn,
            instructions,
            example_table_name,
            example_field_name,
            max_examples=max_examples,
            max_retries=max_retries,
            retry_policy=retry_policy,
            **api_kwargs,
        )
        resolved_instructions = schema.inference_prompt
        resolved_response_format = schema.model
    elif response_format is str or (isinstance(response_format, type) and issubclass(response_format, BaseModel)):
        resolved_instructions = instructions
        resolved_response_format = response_format
    else:
        raise TypeError("response_format must be str or a Pydantic BaseModel subclass")

    responses_udf(
        conn,
        name,
        instructions=resolved_instructions,
        response_format=resolved_response_format,
        model_name=model_name,
        batch_size=batch_size,
        max_concurrency=max_concurrency,
        multimodal=multimodal,
        max_validation_retries=max_validation_retries,
        retry_policy=retry_policy,
        **api_kwargs,
    )

similarity_search ¶

similarity_search(
    conn: DuckDBPyConnection,
    target_table: str,
    query_table: str,
    *,
    target_column: str = "embedding",
    query_column: str = "embedding",
    target_text_column: str = "text",
    query_text_column: str = "text",
    top_k: int = 10,
) -> duckdb.DuckDBPyRelation

Perform top-k cosine similarity search between two DuckDB tables.

Uses DuckDB's built-in list_cosine_similarity for efficient vector comparison without leaving SQL.

Parameters:

Name	Type	Description	Default
`conn`	`DuckDBPyConnection`	An open DuckDB connection.	required
`target_table`	`str`	Table containing candidate embeddings.	required
`query_table`	`str`	Table containing query embeddings.	required
`target_column`	`str`	Embedding column in target_table.	`'embedding'`
`query_column`	`str`	Embedding column in query_table.	`'embedding'`
`target_text_column`	`str`	Text identifier column in target_table.	`'text'`
`query_text_column`	`str`	Text identifier column in query_table.	`'text'`
`top_k`	`int`	Number of results per query.	`10`

Returns:

Type	Description
`DuckDBPyRelation`	duckdb.DuckDBPyRelation: A DuckDB relation with `query_id` (1-based
`DuckDBPyRelation`	query row position), `query_text`, `target_text`, and `score`,
`DuckDBPyRelation`	ordered by query row and descending similarity.

Raises:

Type	Description
`ValueError`	If an identifier or `top_k` is invalid.

Example

import duckdb from openaivec.duckdb_ext import similarity_search conn = duckdb.connect()

(after populating docs and queries tables with embeddings)¶

results = similarity_search(conn, "docs", "queries", top_k=5) results.df()

Source code in src/openaivec/duckdb_ext/_similarity.py

def similarity_search(
    conn: duckdb.DuckDBPyConnection,
    target_table: str,
    query_table: str,
    *,
    target_column: str = "embedding",
    query_column: str = "embedding",
    target_text_column: str = "text",
    query_text_column: str = "text",
    top_k: int = 10,
) -> duckdb.DuckDBPyRelation:
    """Perform top-k cosine similarity search between two DuckDB tables.

    Uses DuckDB's built-in ``list_cosine_similarity`` for efficient
    vector comparison without leaving SQL.

    Args:
        conn (duckdb.DuckDBPyConnection): An open DuckDB connection.
        target_table (str): Table containing candidate embeddings.
        query_table (str): Table containing query embeddings.
        target_column (str): Embedding column in *target_table*.
        query_column (str): Embedding column in *query_table*.
        target_text_column (str): Text identifier column in *target_table*.
        query_text_column (str): Text identifier column in *query_table*.
        top_k (int): Number of results per query.

    Returns:
        duckdb.DuckDBPyRelation: A DuckDB relation with ``query_id`` (1-based
        query row position), ``query_text``, ``target_text``, and ``score``,
        ordered by query row and descending similarity.

    Raises:
        ValueError: If an identifier or ``top_k`` is invalid.

    Example:
        >>> import duckdb
        >>> from openaivec.duckdb_ext import similarity_search
        >>> conn = duckdb.connect()
        >>> # (after populating docs and queries tables with embeddings)
        >>> results = similarity_search(conn, "docs", "queries", top_k=5)
        >>> results.df()
    """
    if isinstance(top_k, bool) or not isinstance(top_k, int) or top_k < 1:
        raise ValueError("top_k must be a positive integer")
    target = _quote_table_name(target_table)
    query = _quote_table_name(query_table)
    target_vector = _quote_identifier(target_column)
    query_vector = _quote_identifier(query_column)
    target_text = _quote_identifier(target_text_column)
    query_text = _quote_identifier(query_text_column)

    sql = f"""
        WITH queries AS (
            SELECT row_number() OVER () AS query_id,
                   q.{query_text} AS query_text,
                   q.{query_vector} AS query_vector
            FROM {query} AS q
        ), candidates AS (
            SELECT q.query_id, q.query_text, t.{target_text} AS target_text,
                   list_cosine_similarity(t.{target_vector}::FLOAT[], q.query_vector::FLOAT[]) AS score
            FROM queries AS q
            CROSS JOIN {target} AS t
        )
        SELECT
            query_id, query_text, target_text, score
        FROM candidates
        QUALIFY row_number() OVER (
            PARTITION BY query_id
            ORDER BY score DESC
        ) <= ?
        ORDER BY query_id, score DESC
    """
    return conn.sql(sql, params=[top_k])

count_tokens_udf ¶

count_tokens_udf(
    conn: DuckDBPyConnection, name: str
) -> None

Register a vectorized DuckDB UDF that counts tokens in text values.

The tokenizer is the same configured tiktoken.Encoding used by pandas.Series.ai.count_tokens. SQL NULL values remain NULL.

Parameters:

Name	Type	Description	Default
`conn`	`DuckDBPyConnection`	Connection on which to register the UDF.	required
`name`	`str`	Function name visible in SQL.	required

Example

import duckdb from openaivec.duckdb_ext import count_tokens_udf conn = duckdb.connect() count_tokens_udf(conn, "token_count") conn.sql("SELECT token_count('hello')").fetchone()[0] > 0 True

Source code in src/openaivec/duckdb_ext/_tokens.py

def count_tokens_udf(conn: duckdb.DuckDBPyConnection, name: str) -> None:
    """Register a vectorized DuckDB UDF that counts tokens in text values.

    The tokenizer is the same configured ``tiktoken.Encoding`` used by
    ``pandas.Series.ai.count_tokens``. SQL NULL values remain NULL.

    Args:
        conn (duckdb.DuckDBPyConnection): Connection on which to register the UDF.
        name (str): Function name visible in SQL.

    Example:
        >>> import duckdb
        >>> from openaivec.duckdb_ext import count_tokens_udf
        >>> conn = duckdb.connect()
        >>> count_tokens_udf(conn, "token_count")
        >>> conn.sql("SELECT token_count('hello')").fetchone()[0] > 0
        True
    """
    encoding = CONTAINER.resolve(tiktoken.Encoding)

    def _batch_udf(arrow_batch: pa.Array) -> pa.Array:
        texts = arrow_batch.to_pylist()
        non_null_texts = [text for text in texts if text is not None]
        if not non_null_texts:
            return pa.nulls(len(texts), type=pa.int64())

        counts = iter(len(tokens) for tokens in encoding.encode_batch(non_null_texts))
        return pa.array([next(counts) if text is not None else None for text in texts], type=pa.int64())

    conn.create_function(
        name,
        _batch_udf,
        [duckdb.sqltype("VARCHAR")],
        duckdb.sqltype("BIGINT"),
        type=PythonUDFType.ARROW,
        null_handling=FunctionNullHandling.SPECIAL,
    )

pydantic_to_duckdb_ddl ¶

pydantic_to_duckdb_ddl(
    model: type[BaseModel], table_name: str
) -> str

Generate a CREATE TABLE DDL statement from a Pydantic model.

Parameters:

Name	Type	Description	Default
`model`	`type[BaseModel]`	The Pydantic model class.	required
`table_name`	`str`	Name for the DuckDB table.	required

Returns:

Name	Type	Description
`str`	`str`	A `CREATE TABLE IF NOT EXISTS` statement.

Example

from pydantic import BaseModel from openaivec.duckdb_ext import pydantic_to_duckdb_ddl class Review(BaseModel): ... sentiment: str ... rating: int ... tags: list[str] print(pydantic_to_duckdb_ddl(Review, "reviews")) CREATE TABLE IF NOT EXISTS "reviews" ( "sentiment" VARCHAR, "rating" INTEGER, "tags" VARCHAR[] )

Source code in src/openaivec/duckdb_ext/_types.py

def pydantic_to_duckdb_ddl(model: type[BaseModel], table_name: str) -> str:
    """Generate a ``CREATE TABLE`` DDL statement from a Pydantic model.

    Args:
        model (type[BaseModel]): The Pydantic model class.
        table_name (str): Name for the DuckDB table.

    Returns:
        str: A ``CREATE TABLE IF NOT EXISTS`` statement.

    Example:
        >>> from pydantic import BaseModel
        >>> from openaivec.duckdb_ext import pydantic_to_duckdb_ddl
        >>> class Review(BaseModel):
        ...     sentiment: str
        ...     rating: int
        ...     tags: list[str]
        >>> print(pydantic_to_duckdb_ddl(Review, "reviews"))
        CREATE TABLE IF NOT EXISTS "reviews" (
            "sentiment" VARCHAR,
            "rating" INTEGER,
            "tags" VARCHAR[]
        )
    """
    table = _quote_table_name(table_name)
    columns: list[str] = []
    for field_name, field_info in model.model_fields.items():
        col_type = _python_type_to_duckdb(field_info.annotation) if field_info.annotation else "VARCHAR"
        columns.append(f"    {_quote_identifier(field_name)} {col_type}")
    body = ",\n".join(columns)
    return f"CREATE TABLE IF NOT EXISTS {table} (\n{body}\n)"

embeddings_udf ¶

embeddings_udf(
    conn: DuckDBPyConnection,
    name: str,
    *,
    model_name: str | None = None,
    batch_size: int = 128,
    max_concurrency: int = 8,
    limits: EmbeddingLimits | None = None,
    retry_policy: RetryPolicy | None = None,
    **api_kwargs: Any,
) -> None

Register a DuckDB Arrow-based UDF that returns embedding vectors.

The UDF processes rows in vectorized batches via AsyncBatchEmbeddings, leveraging deduplication and concurrency for throughput.

Parameters:

Name	Type	Description	Default
`conn`	`DuckDBPyConnection`	An open DuckDB connection.	required
`name`	`str`	UDF name visible in SQL.	required
`model_name`	`str \| None`	Embeddings model or deployment name.	`None`
`batch_size`	`int`	Rows per API batch. Defaults to 128.	`128`
`max_concurrency`	`int`	Maximum concurrent API requests. Defaults to 8.	`8`
`limits`	`EmbeddingLimits \| None`	Hard provider limits; None uses OpenAI defaults.	`None`
`retry_policy`	`RetryPolicy \| None`	Transport limits. `None` preserves SDK retries.	`None`
`**api_kwargs`	`Any`	Extra parameters forwarded to the OpenAI API.	`{}`

Example

import duckdb from openaivec.duckdb_ext import embeddings_udf conn = duckdb.connect() embeddings_udf(conn, "embed")

conn.sql("SELECT embed(text) FROM docs")¶

Source code in src/openaivec/duckdb_ext/_udfs.py

def embeddings_udf(
    conn: duckdb.DuckDBPyConnection,
    name: str,
    *,
    model_name: str | None = None,
    batch_size: int = 128,
    max_concurrency: int = 8,
    limits: EmbeddingLimits | None = None,
    retry_policy: RetryPolicy | None = None,
    **api_kwargs: Any,
) -> None:
    """Register a DuckDB Arrow-based UDF that returns embedding vectors.

    The UDF processes rows in vectorized batches via ``AsyncBatchEmbeddings``,
    leveraging deduplication and concurrency for throughput.

    Args:
        conn (duckdb.DuckDBPyConnection): An open DuckDB connection.
        name (str): UDF name visible in SQL.
        model_name (str | None): Embeddings model or deployment name.
        batch_size (int): Rows per API batch. Defaults to 128.
        max_concurrency (int): Maximum concurrent API requests. Defaults to 8.
        limits (EmbeddingLimits | None): Hard provider limits; None uses OpenAI defaults.
        retry_policy (RetryPolicy | None): Transport limits. ``None`` preserves SDK retries.
        **api_kwargs: Extra parameters forwarded to the OpenAI API.

    Example:
        >>> import duckdb
        >>> from openaivec.duckdb_ext import embeddings_udf
        >>> conn = duckdb.connect()
        >>> embeddings_udf(conn, "embed")
        >>> # conn.sql("SELECT embed(text) FROM docs")
    """

    _model_name = model_name or CONTAINER.resolve(EmbeddingsModelName).value
    async_client = CONTAINER.resolve(AsyncOpenAI)

    cache: AsyncBatchCache[str, np.ndarray] = AsyncBatchCache(
        batch_size=batch_size,
        max_concurrency=max_concurrency,
        max_cache_size=DEFAULT_MANAGED_CACHE_SIZE,
        show_progress=False,
    )
    batch_client = AsyncBatchEmbeddings(
        client=async_client,
        model_name=_model_name,
        cache=cache,
        api_kwargs=api_kwargs,
        limits=limits if limits is not None else EmbeddingLimits(),
        retry_policy=retry_policy,
    )

    def _batch_udf(arrow_batch: pa.Array) -> pa.Array:
        texts = arrow_batch.to_pylist()
        non_null_indices = [i for i, t in enumerate(texts) if t is not None]
        non_null_texts = [texts[i] for i in non_null_indices]

        if not non_null_texts:
            return pa.array([None] * len(texts))

        results = run_async(batch_client.create(non_null_texts))

        out: list[list[float] | None] = [None] * len(texts)
        for idx, vec in zip(non_null_indices, results):
            out[idx] = vec.tolist()

        return pa.array(out, type=pa.list_(pa.float32()))

    conn.create_function(
        name, _batch_udf, [duckdb.sqltype("VARCHAR")], duckdb.list_type("FLOAT"), type=PythonUDFType.ARROW
    )

responses_udf ¶

responses_udf(
    conn: DuckDBPyConnection,
    name: str,
    *,
    instructions: str,
    response_format: type = str,
    model_name: str | None = None,
    batch_size: int | None = 64,
    max_concurrency: int = 8,
    multimodal: bool = False,
    max_validation_retries: int = 3,
    retry_policy: RetryPolicy | None = None,
    **api_kwargs: Any,
) -> None

Register a DuckDB Arrow-based UDF that calls the OpenAI Responses API.

The UDF processes rows in vectorized batches via AsyncBatchResponses, leveraging deduplication and concurrency for throughput.

When response_format is a Pydantic BaseModel, the UDF returns a DuckDB STRUCT whose fields match the model, allowing direct field access in SQL (e.g. SELECT udf(text).sentiment FROM ...). When response_format is str, the UDF returns VARCHAR.

SQL NULL inputs are not sent to the API. Missing parsed responses remain SQL NULL, including structured outputs; all-NULL Arrow batches retain the declared return type.

Parameters:

Name	Type	Description	Default
`conn`	`DuckDBPyConnection`	An open DuckDB connection.	required
`name`	`str`	UDF name visible in SQL.	required
`instructions`	`str`	System prompt for the model.	required
`response_format`	`type`	`str` for plain text or a Pydantic `BaseModel` for structured output as a DuckDB STRUCT. Defaults to `str`.	`str`
`model_name`	`str \| None`	Model or deployment name. Defaults to the container-registered `ResponsesModelName`.	`None`
`batch_size`	`int \| None`	Rows per API batch; None enables auto-tuning. Defaults to 64.	`64`
`max_concurrency`	`int`	Maximum concurrent API requests. Defaults to 8.	`8`
`max_validation_retries`	`int`	Additional schema/ID corrections per batch. Defaults to 3; 0 disables correction. Must be nonnegative.	`3`
`retry_policy`	`RetryPolicy \| None`	Transport limits. `None` preserves SDK retries.	`None`
`**api_kwargs`	`Any`	Extra parameters forwarded to the OpenAI API.	`{}`

Example

import duckdb from pydantic import BaseModel from openaivec.duckdb_ext import responses_udf class Sentiment(BaseModel): ... label: str ... score: float conn = duckdb.connect() responses_udf(conn, "sentiment", instructions="Analyze sentiment", response_format=Sentiment)

conn.sql("SELECT sentiment(text).label, sentiment(text).score FROM docs")¶

Source code in src/openaivec/duckdb_ext/_udfs.py

def responses_udf(
    conn: duckdb.DuckDBPyConnection,
    name: str,
    *,
    instructions: str,
    response_format: type = str,
    model_name: str | None = None,
    batch_size: int | None = 64,
    max_concurrency: int = 8,
    multimodal: bool = False,
    max_validation_retries: int = 3,
    retry_policy: RetryPolicy | None = None,
    **api_kwargs: Any,
) -> None:
    """Register a DuckDB Arrow-based UDF that calls the OpenAI Responses API.

    The UDF processes rows in vectorized batches via ``AsyncBatchResponses``,
    leveraging deduplication and concurrency for throughput.

    When ``response_format`` is a Pydantic ``BaseModel``, the UDF returns a
    DuckDB ``STRUCT`` whose fields match the model, allowing direct field
    access in SQL (e.g. ``SELECT udf(text).sentiment FROM ...``).
    When ``response_format`` is ``str``, the UDF returns ``VARCHAR``.

    SQL NULL inputs are not sent to the API. Missing parsed responses remain
    SQL NULL, including structured outputs; all-NULL Arrow batches retain the
    declared return type.

    Args:
        conn (duckdb.DuckDBPyConnection): An open DuckDB connection.
        name (str): UDF name visible in SQL.
        instructions (str): System prompt for the model.
        response_format (type): ``str`` for plain text or a Pydantic ``BaseModel``
            for structured output as a DuckDB STRUCT. Defaults to ``str``.
        model_name (str | None): Model or deployment name. Defaults to the
            container-registered ``ResponsesModelName``.
        batch_size (int | None): Rows per API batch; None enables auto-tuning.
            Defaults to 64.
        max_concurrency (int): Maximum concurrent API requests. Defaults to 8.
        max_validation_retries (int): Additional schema/ID corrections per batch.
            Defaults to 3; 0 disables correction. Must be nonnegative.
        retry_policy (RetryPolicy | None): Transport limits. ``None`` preserves SDK retries.
        **api_kwargs: Extra parameters forwarded to the OpenAI API.

    Example:
        >>> import duckdb
        >>> from pydantic import BaseModel
        >>> from openaivec.duckdb_ext import responses_udf
        >>> class Sentiment(BaseModel):
        ...     label: str
        ...     score: float
        >>> conn = duckdb.connect()
        >>> responses_udf(conn, "sentiment", instructions="Analyze sentiment", response_format=Sentiment)
        >>> # conn.sql("SELECT sentiment(text).label, sentiment(text).score FROM docs")
    """

    _model_name = model_name or CONTAINER.resolve(ResponsesModelName).value
    async_client = CONTAINER.resolve(AsyncOpenAI)

    cache: AsyncBatchCache = AsyncBatchCache(
        batch_size=batch_size,
        max_concurrency=max_concurrency,
        max_cache_size=DEFAULT_MANAGED_CACHE_SIZE,
        show_progress=False,
    )
    batch_client = AsyncBatchResponses(
        client=async_client,
        model_name=_model_name,
        system_message=instructions,
        response_format=response_format,
        cache=cache,
        max_validation_retries=max_validation_retries,
        retry_policy=retry_policy,
        api_kwargs=api_kwargs,
        multimodal=multimodal,
    )

    is_structured = isinstance(response_format, type) and issubclass(response_format, BaseModel)
    return_type = _pydantic_to_struct_type(response_format) if is_structured else duckdb.sqltype("VARCHAR")
    arrow_type = (
        conn.sql("SELECT NULL")
        .select(duckdb.ConstantExpression(None).cast(return_type))
        .limit(0)
        .to_arrow_table()
        .schema.field(0)
        .type
    )

    def _batch_udf(arrow_batch: pa.Array) -> pa.Array:
        texts = arrow_batch.to_pylist()
        non_null_indices = [i for i, t in enumerate(texts) if t is not None]
        non_null_texts = [texts[i] for i in non_null_indices]

        if not non_null_texts:
            return pa.nulls(len(texts), type=arrow_type)

        results = run_async(batch_client.parse(non_null_texts))

        out: list[Any] = [None] * len(texts)
        for idx, result in zip(non_null_indices, results):
            if is_structured and isinstance(result, BaseModel):
                out[idx] = _serialize_for_duckdb(result.model_dump())
            elif result is not None:
                out[idx] = str(result)

        if all(value is None for value in out):
            return pa.nulls(len(out), type=arrow_type)
        return pa.array(out)

    conn.create_function(
        name,
        _batch_udf,
        [duckdb.sqltype("VARCHAR")],
        return_type,
        type=PythonUDFType.ARROW,
        null_handling=FunctionNullHandling.SPECIAL,
    )

task_udf ¶

task_udf(
    conn: DuckDBPyConnection,
    name: str,
    *,
    task: PreparedTask[ResponseFormat],
    model_name: str | None = None,
    batch_size: int | None = 64,
    max_concurrency: int = 8,
    multimodal: bool = False,
    max_validation_retries: int = 3,
    retry_policy: RetryPolicy | None = None,
    **api_kwargs: Any,
) -> None

Register a DuckDB UDF backed by a PreparedTask.

Parameters:

Name	Type	Description	Default
`conn`	`DuckDBPyConnection`	An open DuckDB connection.	required
`name`	`str`	UDF name visible in SQL.	required
`task`	`PreparedTask`	Pre-configured task with instructions and response format.	required
`model_name`	`str \| None`	Model or deployment name.	`None`
`batch_size`	`int \| None`	Rows per API batch; None enables auto-tuning. Defaults to 64.	`64`
`max_concurrency`	`int`	Maximum concurrent API requests. Defaults to 8.	`8`
`max_validation_retries`	`int`	Additional schema/ID corrections per batch. Defaults to 3; 0 disables correction. Must be nonnegative.	`3`
`multimodal`	`bool`	When `True`, file paths and URLs are sent as multimodal content. Defaults to `False`.	`False`
`retry_policy`	`RetryPolicy \| None`	Transport limits. `None` preserves SDK retries.	`None`
`**api_kwargs`	`Any`	Extra parameters forwarded to the OpenAI API.	`{}`

Source code in src/openaivec/duckdb_ext/_udfs.py

def task_udf(
    conn: duckdb.DuckDBPyConnection,
    name: str,
    *,
    task: PreparedTask[ResponseFormat],
    model_name: str | None = None,
    batch_size: int | None = 64,
    max_concurrency: int = 8,
    multimodal: bool = False,
    max_validation_retries: int = 3,
    retry_policy: RetryPolicy | None = None,
    **api_kwargs: Any,
) -> None:
    """Register a DuckDB UDF backed by a ``PreparedTask``.

    Args:
        conn (duckdb.DuckDBPyConnection): An open DuckDB connection.
        name (str): UDF name visible in SQL.
        task (PreparedTask): Pre-configured task with instructions and response format.
        model_name (str | None): Model or deployment name.
        batch_size (int | None): Rows per API batch; None enables auto-tuning.
            Defaults to 64.
        max_concurrency (int): Maximum concurrent API requests. Defaults to 8.
        max_validation_retries (int): Additional schema/ID corrections per batch.
            Defaults to 3; 0 disables correction. Must be nonnegative.
        multimodal (bool): When ``True``, file paths and URLs are sent as
            multimodal content. Defaults to ``False``.
        retry_policy (RetryPolicy | None): Transport limits. ``None`` preserves SDK retries.
        **api_kwargs: Extra parameters forwarded to the OpenAI API.
    """
    responses_udf(
        conn,
        name,
        instructions=task.instructions,
        response_format=task.response_format,
        model_name=model_name,
        batch_size=batch_size,
        max_concurrency=max_concurrency,
        multimodal=multimodal,
        max_validation_retries=max_validation_retries,
        retry_policy=retry_policy,
        **api_kwargs,
    )