Datachain

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 typeSent as
str, Pydantic modeltext (a model → JSON)
TextFile (type="text")text
ImageFile (type="image"), video frameimage (needs a vision model)
File (untyped)text (read_text; errors if the bytes are binary)
bytestext (errors if not UTF-8)
AudioFile / VideoFileerror; 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
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
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
76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
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
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
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
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
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,
    )