feat: Add conditions for conditional task execution (#195)

* refactor: Move inference_mode into GenerationSettings.

* docs: Update docs.

* feat: Add conditional task execution logic.

* docs: Update docs.

---------

Co-authored-by: Raphael Mitsch <raphael@climatiq.com>

Author: rmitsch

README.md

@@ -86,7 +86,7 @@ build modern NLP applications. It provides:
   - [`langchain`](https://github.com/langchain-ai/langchain)
   - [`outlines`](https://github.com/dottxt-ai/outlines)
   - [`transformer`](https://github.com/huggingface/transformers)
-- :arrow_forward: **Observable Pipelines:** Easy debugging and monitoring
+- :arrow_forward: **Observable Pipelines:** Easy debugging and monitoring with conditional task execution
 - :hammer_and_wrench: **Integrated Tools:**
   - Document parsing (optional via `ingestion` extra): [`docling`](https://github.com/DS4SD/docling), [`marker`](https://github.com/VikParuchuri/marker)
   - Text chunking: [`chonkie`](https://github.com/chonkie-ai/chonkie)
@@ -246,6 +246,7 @@ Encapsulates a single processing step in a pipeline.
 - Defines input arguments
 - Wraps and initializes `Bridge` instances handling task-engine-specific logic
 - Implements task-specific dataset export
+- Supports **conditional execution**: skip documents based on custom logic without materializing all docs upfront
 
 #### `GenerationSettings`
 Controls behavior of structured generation across tasks.

docs/guides/optimization.md

@@ -22,7 +22,7 @@ Optimization is valuable when:
 > Optimization involves **multiple LLM calls** during the search process. Costs depend on:
 > - Dataset size (more examples = more evaluations)
 > - DSPy optimizer configuration (`num_candidates`, `num_trials`)
-> - Model pricing (larger models cost more per call)
+> - Model pricing (larger models cost more per _call)
 >
 > Start with small datasets and conservative optimizer settings to control costs.
 
@@ -112,7 +112,7 @@ These tasks use a **generic LLM-as-judge evaluator** that compares ground truth
 - **Translation** - Evaluates translation quality
 - **Question Answering** - Evaluates answer correctness
 
-> **Note**: LLM-based evaluation adds additional costs since each evaluation requires an extra LLM call.
+> **Note**: LLM-based evaluation adds additional costs since each evaluation requires an extra LLM _call.
 
 ## Optimizer Configuration
 

docs/pipeline.md

@@ -40,4 +40,84 @@ Note: Ingestion libraries (e.g., `docling`) are optional and not installed by de
 pip install "sieves[ingestion]"
 ```
 
+## Conditional Task Execution
+
+Tasks support optional conditional execution via the `condition` parameter. This allows you to skip processing certain documents based on custom logic, without materializing all documents upfront.
+
+### Basic Usage
+
+Pass a callable `Condition[[Doc], bool]` to any task to conditionally process documents:
+
+```python
+from sieves import Pipeline, tasks, Doc
+
+docs = [
+    Doc(text="short"),
+    Doc(text="this is a much longer document that will be processed"),
+    Doc(text="med"),
+]
+
+# Define a condition function
+def is_long(doc: Doc) -> bool:
+    return len(doc.text or "") > 20
+
+# Create a task with a condition
+task = tasks.Classification(
+    labels=["science", "politics"],
+    model=model,
+    condition=is_long
+)
+
+# Run pipeline
+pipe = Pipeline([task])
+for doc in pipe(docs):
+    # doc.results[task.id] will be None for documents that failed the condition
+    print(doc.results[task.id])
+```
+
+### Key Behaviors
+
+- **Per-document evaluation**: The condition is evaluated for each document individually
+- **Lazy evaluation**: Documents are not materialized upfront; passing documents are batched together for efficient processing
+- **Result tracking**: Skipped documents have `results[task_id] = None`
+- **Order preservation**: Document order is always maintained, regardless of which documents are skipped
+- **No-op when None**: If `condition=None`, all documents are processed
+
+### Multiple Tasks with Different Conditions
+
+Different tasks in a pipeline can have different conditions:
+
+```python
+from sieves import Pipeline, tasks, Doc
+
+docs = [
+    Doc(text="short"),
+    Doc(text="this is a much longer document"),
+    Doc(text="medium text here"),
+]
+
+# Task 1: Process only documents longer than 10 characters
+task1 = tasks.Chunking(chunker, condition=lambda d: len(d.text or "") > 10)
+
+# Task 2: Process only documents longer than 20 characters
+task2 = tasks.Classification(
+    labels=["science", "politics"],
+    model=model,
+    condition=lambda d: len(d.text or "") > 20
+)
+
+# First doc: skipped by both tasks (too short)
+# Second doc: processed by both tasks (long enough)
+# Third doc: processed by task1, skipped by task2
+pipe = Pipeline([task1, task2])
+for doc in pipe(docs):
+    print(doc.results[task1.id], doc.results[task2.id])
+```
+
+### Use Cases
+
+- **Skip expensive processing** for documents that don't meet quality criteria
+- **Segment processing** by document properties (size, language, format)
+- **Optimize pipelines** by processing subsets of data through specific tasks
+
 ::: sieves.pipeline.core

docs/tasks/task.md

@@ -1,3 +1,92 @@
 # Task
 
-::: sieves.tasks.core.Task
+## Conditional Execution
+
+All tasks support optional conditional execution through the `condition` parameter. This feature allows you to skip processing certain documents based on custom criteria without materializing all documents upfront.
+
+### Overview
+
+The `condition` parameter accepts an optional callable with signature `Callable[[Doc], bool]`:
+
+```python
+def condition(doc: Doc) -> bool:
+    # Return True to process the document
+    # Return False to skip it
+    return True
+```
+
+### Implementation Details
+
+When a task is executed with a condition:
+
+1. **Per-Document Evaluation**: Each document is evaluated against the condition individually
+2. **Lazy Batching**: Only documents that pass the condition are batched together and sent to the task's `_call()` method
+3. **Order Preservation**: Documents are returned in their original order, even if some were skipped
+4. **Result Storage**: Skipped documents have `results[task_id] = None`
+
+### Examples
+
+#### Skip Documents by Size
+
+```python
+from sieves import tasks, Pipeline, Doc
+
+# Only process documents longer than 100 characters
+task = tasks.Classification(
+    labels=["positive", "negative"],
+    model=model,
+    condition=lambda doc: len(doc.text or "") > 100
+)
+
+pipe = Pipeline([task])
+docs = [Doc(text="short"), Doc(text="a very long document " * 10)]
+results = list(pipe(docs))
+
+# First doc: results[task.id] == None (skipped)
+# Second doc: results[task.id] contains classification results
+```
+
+#### Skip Documents Based on Metadata
+
+```python
+# Only process documents from specific sources
+def should_process(doc: Doc) -> bool:
+    return doc.meta.get("source") in ["source_a", "source_b"]
+
+task = tasks.NER(
+    entities=["PERSON", "LOCATION"],
+    model=model,
+    condition=should_process
+)
+```
+
+#### Multiple Conditions in Pipeline
+
+```python
+# Different conditions for different tasks
+import_task = tasks.Ingestion(export_format="markdown")
+
+# Only chunk long documents
+chunking_task = tasks.Chunking(
+    chunker,
+    condition=lambda doc: len(doc.text or "") > 500
+)
+
+# Only classify chunked documents
+classification_task = tasks.Classification(
+    labels=["science", "fiction"],
+    model=model,
+    condition=lambda doc: len(doc.text or "") > 500
+)
+
+pipe = Pipeline([import_task, chunking_task, classification_task])
+```
+
+### Technical Notes
+
+- **No Materialization**: Documents are processed using iterators; passing documents are batched together without materializing the entire document collection upfront
+- **Index-Based Tracking**: The implementation uses document indices for efficient filtering and reordering
+- **All Engines Supported**: Conditional execution works with all supported engines (DSPy, LangChain, Outlines, HuggingFace, GLiNER, etc.)
+- **Serialization**: Non-callable condition values (like `None`) serialize naturally; callable conditions are serialized as placeholders
+
+::: sieves.tasks.core.Task

sieves/engines/huggingface_.py

@@ -47,7 +47,7 @@ def build_executable(
         assert isinstance(prompt_signature, list)
 
         # Render template with few-shot examples. Note that we don't use extracted document values here, as HF zero-shot
-        # pipelines only support one hypothesis template per call - and we want to batch, so our hypothesis template
+        # pipelines only support one hypothesis template per _call - and we want to batch, so our hypothesis template
         # will be document-invariant.
         fewshot_examples_dict = HuggingFace.convert_fewshot_examples(fewshot_examples)
         # Render hypothesis template with everything but text.

sieves/serialization.py

@@ -92,7 +92,7 @@ def create(cls, cls_obj: type, attributes: dict[str, Attribute]) -> Config:
         :param attributes: Attributes to include in config.
         :return Config: Instance of dynamic config class.
         """
-        config_type = pydantic.create_model(  # type: ignore[call-overload]
+        config_type = pydantic.create_model(  # type: ignore[_call-overload]
             f"{cls_obj}Config",
             __base__=Config,
             **{attr_id: (Attribute, ...) for attr_id in attributes},
@@ -186,7 +186,7 @@ def load(cls, path: Path | str) -> Config:
         with open(path) as file:
             data = yaml.safe_load(file)
 
-        config = pydantic.create_model(  # type: ignore[call-overload]
+        config = pydantic.create_model(  # type: ignore[_call-overload]
             f"{data['cls_name']}Config",
             __base__=Config,
             cls_name=(str, ...),

sieves/tasks/core.py

@@ -3,7 +3,8 @@
 from __future__ import annotations
 
 import abc
-from collections.abc import Iterable
+import itertools
+from collections.abc import Callable, Iterable, Iterator
 from typing import TYPE_CHECKING, Any
 
 from sieves.data import Doc
@@ -17,17 +18,27 @@
 class Task(abc.ABC):
     """Abstract base class for tasks that can be executed on documents."""
 
-    def __init__(self, task_id: str | None, include_meta: bool, batch_size: int):
+    def __init__(
+        self,
+        task_id: str | None,
+        include_meta: bool,
+        batch_size: int,
+        condition: Callable[[Doc], bool] | None = None,
+    ):
         """
         Initiate new Task.
 
         :param task_id: Task ID.
         :param include_meta: Whether to include meta information generated by the task.
         :param batch_size: Batch size for processing documents. Use -1 to process all documents at once.
+        :param condition: Optional callable that determines whether to process each document.
+                          If provided, called with each Doc; if returns False, document is skipped
+                          and results[task_id] is set to None.
         """
         self._task_id = task_id if task_id else self.__class__.__name__
         self._include_meta = include_meta
         self._batch_size = batch_size
+        self._condition = condition
 
     @property
     def id(self) -> str:
@@ -39,9 +50,46 @@ def id(self) -> str:
         """
         return self._task_id
 
-    @abc.abstractmethod
     def __call__(self, docs: Iterable[Doc]) -> Iterable[Doc]:
-        """Execute task.
+        """Execute task with conditional logic.
+
+        Checks the condition for each document without materializing all docs upfront.
+        Passes all documents that pass the condition to _call() for proper batching.
+        Documents that fail the condition have results[task_id] set to None.
+
+        :param docs: Docs to process.
+        :return: Processed docs (in original order).
+        """
+        # Create three independent iterators:
+        #   1. Check which docs pass condition.
+        #   2. Yield only passing docs to _call().
+        #   3. Iterate and yield results in order.
+        docs_iters = itertools.tee(docs, 3)
+
+        # First pass: determine which docs pass the condition by index
+        passing_indices: set[int] = set()
+
+        for idx, doc in enumerate(docs_iters[0]):
+            if self._condition is None or self._condition(doc):
+                passing_indices.add(idx)
+
+        # Process all passing docs together.
+        processed = self._call(d for i, d in enumerate(docs_iters[1]) if i in passing_indices)
+        processed_iter = iter(processed) if not isinstance(processed, Iterator) else processed
+
+        # Iterate through original docs in order and yield results
+        for idx, doc in enumerate(docs_iters[2]):
+            if idx in passing_indices:
+                # Doc passed condition - use processed result.
+                yield next(processed_iter)
+            else:
+                # Doc failed condition - set None result and yield original.
+                doc.results[self.id] = None
+                yield doc
+
+    @abc.abstractmethod
+    def _call(self, docs: Iterable[Doc]) -> Iterable[Doc]:
+        """Execute task logic (to be implemented by subclasses).
 
         :param docs: Docs to process.
         :return: Processed docs.
@@ -83,6 +131,7 @@ def _state(self) -> dict[str, Any]:
             "task_id": self._task_id,
             "include_meta": self._include_meta,
             "batch_size": self._batch_size,
+            "condition": self._condition,
         }
 
     def serialize(self) -> Config:

sieves/tasks/predictive/classification/bridges.py

@@ -415,7 +415,7 @@ def _prompt_conclusion(self) -> str | None:
     @cached_property
     def prompt_signature(self) -> type[pydantic.BaseModel] | list[str]:
         if self._multi_label:
-            prompt_sig = pydantic.create_model(  # type: ignore[call-overload]
+            prompt_sig = pydantic.create_model(  # type: ignore[_call-overload]
                 "MultilabelClassification",
                 __base__=pydantic.BaseModel,
                 __doc__="Result of multi-label classification.",

sieves/tasks/predictive/classification/core.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import json
-from collections.abc import Iterable, Sequence
+from collections.abc import Callable, Iterable, Sequence
 from pathlib import Path
 from typing import Any, override
 
@@ -97,6 +97,7 @@ def __init__(
         label_descriptions: dict[str, str] | None = None,
         multi_label: bool = True,
         generation_settings: GenerationSettings = GenerationSettings(),
+        condition: Callable[[Doc], bool] | None = None,
     ) -> None:
         """Initialize new PredictiveTask.
 
@@ -112,6 +113,7 @@ def __init__(
             most likely class label. In the latter case label forcing mechanisms are utilized, which can lead to higher
             accuracy.
         :param generation_settings: Generation settings.
+        :param condition: Optional callable that determines whether to process each document.
         """
         self._labels = labels
         self._label_descriptions = label_descriptions or {}
@@ -127,6 +129,7 @@ def __init__(
             prompt_instructions=prompt_instructions,
             fewshot_examples=fewshot_examples,
             generation_settings=generation_settings,
+            condition=condition,
         )
         self._fewshot_examples: Sequence[FewshotExample]