LLM Operations

datachain.llm provides named LLM operations over chain columns, parallel to datachain.func. Where func transpiles to SQL, llm calls a model: each operation runs one (expensive) model call per row and materializes the result as a typed, cached, lineage-tracked column.

Use each operation inside the matching verb (.map() for 1:1, .gen() for 1:N); the output column is typed automatically, with no output= needed. A schema=list[Model] returns many items: fan them out with .gen() (one row each) or keep them as one list[Model] column with .map().

Quickstart

import datachain as dc
from datachain import llm
from pydantic import BaseModel

class Scene(BaseModel):
    objects: list[str]
    risk: float

(
    dc.read_storage("s3://frames", type="image")
    .settings(llm="anthropic/claude-haiku-4-5")
    .map(topic=llm.classify("file", into=["accident", "normal"]))  # -> str
    .map(risk=llm.score("file", "accident risk 0..1"))             # -> float
    .map(scene=llm.complete("file", schema=Scene))                 # -> Scene
    .save("frames")
)

Inputs

A column value is sent to the model as text, an image, or a document. Its type decides the encoding automatically, so loading a column with the matching read_storage(type=...) is usually all you need:

Column type	Sent as
str, Pydantic model	text (a model → JSON)
TextFile (type="text")	text
ImageFile (type="image"), video frame	image (needs a vision model)
File (untyped)	text (read_text; errors if the bytes are binary)
bytes	text (errors if not UTF-8)
AudioFile / VideoFile	error; decode first (extract frames or a transcript)

# ImageFile is encoded as an image automatically; no type= needed:
dc.read_storage("s3://bucket/imgs/", type="image").map(cap=llm.complete("file"))

type= is an override for when the type is ambiguous (raw bytes or an untyped File) or for a PDF (there is no document file type). Its values line up with read_storage's type (text, image), plus document for PDF input:

.map(cap=llm.complete("frames", type="image"))     # raw bytes / untyped File -> image
.map(ext=llm.complete("file",  type="document"))    # File / bytes -> PDF (document-capable model)

Rules to keep in mind:

• A type mismatch (type="image" on non-image bytes, type="document" on a non-PDF) raises a clear error when the row is processed.
• type is the send modality (aligned with read_storage's type), distinct from an object's content_type (its exact MIME type, e.g. image/png).
• A str is sent verbatim, so a column holding a path sends the path, not the file's contents; read it as a File (read_storage(...)) to send the content.
• type="document" covers "summarize/extract from this PDF"; heavy document work (chunk by section or clause, pull embedded figures) is better done by decoding first, then llm.*.
• A None (missing) input skips the model call and yields None, so an llm.* output column is always optional. In .gen() a None input emits no rows.

embed takes text only (str, TextFile, or a model); embedding an image or document raises. Embed a text column or a caption you generated, then vector-search:

.map(vec=llm.embed("caption"))  # -> list[float]

Model selection

The model is chosen once with .settings(llm="provider/model") and inherited by every operation below it. A per-call llm= argument overrides it for that call.

Routing is handled by LiteLLM, so any provider-prefixed string works: anthropic/claude-haiku-4-5, openai/gpt-5-mini, bedrock/anthropic.claude-3-5-sonnet-v1:0, vertex_ai/gemini-2.0-flash, etc. Credentials are read from the environment / cloud IAM by default; pass them explicitly with .settings(llm_params=...). A dict of params is part of the cache key; a callable (resolved per worker, e.g. for secrets) is not, so put output-affecting params in the dict form or in per-call keyword arguments.

Token usage

Pass include_usage=True to capture token counts for a call. The function then returns a tuple (value, dc.llm.Usage), so you name both columns with the multi-output form (the value column keeps its plain type, no wrapper):

.map(
    llm.complete("file", schema=Scene, include_usage=True),
    output={"res": Scene, "tok": dc.llm.Usage},
)
# res: Scene   (res.risk works),  tok: Usage

Usage has input_tokens and output_tokens. Aggregate on it afterwards:

chain.agg(in_tok=dc.func.sum("tok.input_tokens"))

Notes:

• Embeddings report no output tokens, so output_tokens stays 0.
• In .gen() (1:N) a call's usage is attributed to one emitted row and zero on the rest, so summing it counts each call once (a naive sum is correct).
• Transient provider errors are retried by LiteLLM (with backoff); structured output relies on the provider's response_format and is not re-asked on a bad parse (that raises a clear error).

Scaling and caching

A model call is the most expensive step, and each worker processes rows sequentially, so throughput comes from more workers: .settings(parallel=N) for N processes on one machine, .settings(workers=N) to distribute.

Reliability is layered:

• File data: DataChain prefetches inputs (overlapping downloads with in-flight calls), caches them, and retries transient fetch errors; tune with .settings(prefetch=N) (default 2).
• Model calls: guarded per call with retries= and fallback=.

Materialized llm.* columns are cached and versioned, so re-running a chain reads the stored result instead of re-calling the model; the cache invalidates when any output-affecting input changes (model, prompt, schema, the input column, type, params, ...).

No fused predicate

There is intentionally no llm.if / fused filter. Materialize a column with an llm.* operation (cached and versioned), then filter on it with a plain .filter(), a cheap recall, not a model rerun. Prefer a continuous score over a hard boolean so re-thresholding stays free:

.map(spoiler_score=llm.score("file", "likelihood this is a spoiler, 0..1"))
.filter("spoiler_score > 0.7")

Functions

complete

complete(
    col: str,
    prompt: str | None = None,
    *,
    schema: Any = None,
    context: str | None = None,
    type: Media | None = None,
    llm: str | None = None,
    retries: int = 1,
    fallback: str | list[str] | None = None,
    include_usage: bool = False,
    **params: Any
) -> LLMSpec

Generate text or extract a structured object from a column.

Output is str when no schema is given and the Pydantic schema model when it is. A list[Model] schema returns many items: use .gen() to fan them out (one row each) or .map() to keep them as one list[Model] column.

Parameters:

• col (str) –

  Input column. Its type decides the encoding (text files and strings as text, images/frames as vision input); for raw bytes or an untyped File set type. See the Inputs section.

• prompt (str | None, default: None ) –

  Instruction text added before the input.

• schema (type | None, default: None ) –

  Pydantic model (or list[Model]) for structured output. When omitted, the output is plain str.

• context (str | None, default: None ) –

  Column whose value is serialized into the prompt.

• type (text | image | document | None, default: None ) –

  Override how col is encoded. Needed for raw bytes or an untyped File: e.g. type="image" for an image-bytes column, type="document" for a PDF (sent to a document-capable model).

• llm (str | None, default: None ) –

  Per-call model override, taking precedence over settings(llm=...).

• retries (int, default: 1 ) –

  Transient and schema-validation retry budget.

• fallback (str | list[str] | None, default: None ) –

  Model string(s) tried if the primary model fails.

• include_usage (bool, default: False ) –

  When True, the call returns a tuple (value, dc.llm.Usage), used with the multi-output form that names both columns: .map(llm.complete(...), output={"res": Scene, "tok": dc.llm.Usage}).

• params (Any, default: {} ) –

  Extra arguments forwarded to the underlying model call.

Returns:

• LLMSpec ( LLMSpec ) –

  A spec used inside .map() (1:1) or .gen() (1:N).

Example

.map(scene=llm.complete("file", schema=Scene))
.gen(chunk=llm.complete("file", schema=list[Chunk], prompt="split"))

Source code in datachain/llm/__init__.py

def complete(
    col: str,
    prompt: str | None = None,
    *,
    schema: Any = None,
    context: str | None = None,
    type: Media | None = None,
    llm: str | None = None,
    retries: int = 1,
    fallback: str | list[str] | None = None,
    include_usage: bool = False,
    **params: Any,
) -> LLMSpec:
    """Generate text or extract a structured object from a column.

    Output is ``str`` when no ``schema`` is given and the Pydantic ``schema`` model
    when it is. A ``list[Model]`` schema returns many items: use ``.gen()`` to fan
    them out (one row each) or ``.map()`` to keep them as one ``list[Model]`` column.

    Args:
        col (str): Input column. Its type decides the encoding (text files and
            strings as text, images/frames as vision input); for raw ``bytes`` or
            an untyped ``File`` set ``type``. See the Inputs section.
        prompt (str | None): Instruction text added before the input.
        schema (type | None): Pydantic model (or ``list[Model]``) for structured
            output. When omitted, the output is plain ``str``.
        context (str | None): Column whose value is serialized into the prompt.
        type ("text" | "image" | "document" | None): Override how ``col`` is
            encoded. Needed for raw ``bytes`` or an untyped ``File``: e.g.
            ``type="image"`` for an image-bytes column, ``type="document"`` for
            a PDF (sent to a document-capable model).
        llm (str | None): Per-call model override, taking precedence over
            ``settings(llm=...)``.
        retries (int): Transient and schema-validation retry budget.
        fallback (str | list[str] | None): Model string(s) tried if the primary
            model fails.
        include_usage (bool): When True, the call returns a tuple
            ``(value, dc.llm.Usage)``, used with the multi-output form that names
            both columns:
            ``.map(llm.complete(...), output={"res": Scene, "tok": dc.llm.Usage})``.
        params (Any): Extra arguments forwarded to the underlying model call.

    Returns:
        LLMSpec: A spec used inside ``.map()`` (1:1) or ``.gen()`` (1:N).

    Example:
        ```py
        .map(scene=llm.complete("file", schema=Scene))
        .gen(chunk=llm.complete("file", schema=list[Chunk], prompt="split"))
        ```
    """
    return LLMSpec(
        kind="complete",
        col=col,
        prompt=prompt,
        schema=schema,
        context_col=context,
        type=type,
        llm=llm,
        retries=retries,
        fallback=fallback,
        include_usage=include_usage,
        params=params,
    )

classify

classify(
    col: str,
    into: Sequence[str],
    prompt: str | None = None,
    *,
    context: str | None = None,
    type: Media | None = None,
    llm: str | None = None,
    retries: int = 1,
    fallback: str | list[str] | None = None,
    include_usage: bool = False,
    **params: Any
) -> LLMSpec

Categorize a column into exactly one of the given labels.

Parameters:

• col (str) –

  Input column passed to the model.

• into (list[str]) –

  Allowed categories; the output is constrained to one.

• prompt (str | None, default: None ) –

  Optional extra guidance added to the instruction.

• context (str | None, default: None ) –

  Column whose value is serialized into the prompt.

• type (text | image | document | None, default: None ) –

  Override how col is encoded (see complete).

• llm (str | None, default: None ) –

  Per-call model override.

• retries (int, default: 1 ) –

  Transient and schema-validation retry budget.

• fallback (str | list[str] | None, default: None ) –

  Model string(s) tried on failure.

• include_usage (bool, default: False ) –

  Also emit a dc.llm.Usage column (see complete).

• params (Any, default: {} ) –

  Extra arguments forwarded to the underlying model call.

Returns:

• LLMSpec ( LLMSpec ) –

  A spec whose output type is str.

Source code in datachain/llm/__init__.py

def classify(
    col: str,
    into: Sequence[str],
    prompt: str | None = None,
    *,
    context: str | None = None,
    type: Media | None = None,
    llm: str | None = None,
    retries: int = 1,
    fallback: str | list[str] | None = None,
    include_usage: bool = False,
    **params: Any,
) -> LLMSpec:
    """Categorize a column into exactly one of the given labels.

    Args:
        col (str): Input column passed to the model.
        into (list[str]): Allowed categories; the output is constrained to one.
        prompt (str | None): Optional extra guidance added to the instruction.
        context (str | None): Column whose value is serialized into the prompt.
        type ("text" | "image" | "document" | None): Override how ``col`` is encoded
            (see ``complete``).
        llm (str | None): Per-call model override.
        retries (int): Transient and schema-validation retry budget.
        fallback (str | list[str] | None): Model string(s) tried on failure.
        include_usage (bool): Also emit a ``dc.llm.Usage`` column (see ``complete``).
        params (Any): Extra arguments forwarded to the underlying model call.

    Returns:
        LLMSpec: A spec whose output type is ``str``.
    """
    if isinstance(into, str):
        raise ValueError(  # noqa: TRY004 - a config value error, not a type guard
            "llm.classify(into=...) expects a list of categories, not a single string"
        )
    return LLMSpec(
        kind="classify",
        col=col,
        prompt=prompt,
        into=list(into),
        context_col=context,
        type=type,
        llm=llm,
        retries=retries,
        fallback=fallback,
        include_usage=include_usage,
        params=params,
    )

score

score(
    col: str,
    prompt: str | None = None,
    *,
    context: str | None = None,
    type: Media | None = None,
    llm: str | None = None,
    retries: int = 1,
    fallback: str | list[str] | None = None,
    include_usage: bool = False,
    **params: Any
) -> LLMSpec

Numeric scoring of a column against a prompt.

Parameters:

• col (str) –

  Input column passed to the model.

• prompt (str | None, default: None ) –

  The scoring criterion (e.g. "accident risk 0..1").

• context (str | None, default: None ) –

  Column whose value is serialized into the prompt.

• type (text | image | document | None, default: None ) –

  Override how col is encoded (see complete).

• llm (str | None, default: None ) –

  Per-call model override.

• retries (int, default: 1 ) –

  Transient and schema-validation retry budget.

• fallback (str | list[str] | None, default: None ) –

  Model string(s) tried on failure.

• include_usage (bool, default: False ) –

  Also emit a dc.llm.Usage column (see complete).

• params (Any, default: {} ) –

  Extra arguments forwarded to the underlying model call.

Returns:

• LLMSpec ( LLMSpec ) –

  A spec whose output type is float.

Source code in datachain/llm/__init__.py

def score(
    col: str,
    prompt: str | None = None,
    *,
    context: str | None = None,
    type: Media | None = None,
    llm: str | None = None,
    retries: int = 1,
    fallback: str | list[str] | None = None,
    include_usage: bool = False,
    **params: Any,
) -> LLMSpec:
    """Numeric scoring of a column against a prompt.

    Args:
        col (str): Input column passed to the model.
        prompt (str | None): The scoring criterion (e.g. ``"accident risk 0..1"``).
        context (str | None): Column whose value is serialized into the prompt.
        type ("text" | "image" | "document" | None): Override how ``col`` is encoded
            (see ``complete``).
        llm (str | None): Per-call model override.
        retries (int): Transient and schema-validation retry budget.
        fallback (str | list[str] | None): Model string(s) tried on failure.
        include_usage (bool): Also emit a ``dc.llm.Usage`` column (see ``complete``).
        params (Any): Extra arguments forwarded to the underlying model call.

    Returns:
        LLMSpec: A spec whose output type is ``float``.
    """
    return LLMSpec(
        kind="score",
        col=col,
        prompt=prompt,
        context_col=context,
        type=type,
        llm=llm,
        retries=retries,
        fallback=fallback,
        include_usage=include_usage,
        params=params,
    )

embed

embed(
    col: str,
    *,
    llm: str | None = None,
    retries: int = 1,
    fallback: str | list[str] | None = None,
    include_usage: bool = False,
    **params: Any
) -> LLMSpec

Embed a column into a vector.

Parameters:

• col (str) –

  Input column to embed.

• llm (str | None, default: None ) –

  Per-call embedding-model override.

• retries (int, default: 1 ) –

  Transient retry budget.

• fallback (str | list[str] | None, default: None ) –

  Model string(s) tried on failure.

• include_usage (bool, default: False ) –

  Also emit a dc.llm.Usage column (see complete). Embeddings report no output tokens, so output_tokens stays 0.

• params (Any, default: {} ) –

  Extra arguments forwarded to the underlying model call.

Returns:

• LLMSpec ( LLMSpec ) –

  A spec whose output type is list[float].

Source code in datachain/llm/__init__.py

def embed(
    col: str,
    *,
    llm: str | None = None,
    retries: int = 1,
    fallback: str | list[str] | None = None,
    include_usage: bool = False,
    **params: Any,
) -> LLMSpec:
    """Embed a column into a vector.

    Args:
        col (str): Input column to embed.
        llm (str | None): Per-call embedding-model override.
        retries (int): Transient retry budget.
        fallback (str | list[str] | None): Model string(s) tried on failure.
        include_usage (bool): Also emit a ``dc.llm.Usage`` column (see ``complete``).
            Embeddings report no output tokens, so ``output_tokens`` stays 0.
        params (Any): Extra arguments forwarded to the underlying model call.

    Returns:
        LLMSpec: A spec whose output type is ``list[float]``.
    """
    return LLMSpec(
        kind="embed",
        col=col,
        llm=llm,
        retries=retries,
        fallback=fallback,
        include_usage=include_usage,
        params=params,
    )