mistralai-workflows-plugins-search 0.1.0__tar.gz

mistralai_workflows_plugins_search-0.1.0/.gitignore

@@ -0,0 +1,32 @@
+# python sdk
+build
+*.egg-info
+*.egg
+
+# Env files
+.env
+.env.*
+!.env.example
+
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+.pytest_cache/
+.coverage
+htmlcov/
+
+# IDEs
+.idea/
+.vscode/
+*.swp
+*.swo
+
+# docs
+docs/node_modules/**
+docs/.docusaurus/**
+docs/pnpm-lock.yaml
+.docs.temporal.io/
+
+#test
+/.test_durations

mistralai_workflows_plugins_search-0.1.0/PKG-INFO

@@ -0,0 +1,22 @@
+Metadata-Version: 2.5
+Name: mistralai-workflows-plugins-search
+Version: 0.1.0
+Summary: Search Toolkit plugin for Mistral Workflows - durable document ingestion pipelines
+Project-URL: Homepage, https://mistral.ai
+Project-URL: Documentation, https://docs-internal-frameworks.mistral.ai/workflows
+Author-email: Mistral AI <support@mistral.ai>
+License: Proprietary
+Keywords: ai,ingestion,mistral,plugin,rag,search,workflows
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: License :: Other/Proprietary License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Requires-Python: <3.15,>=3.12
+Requires-Dist: mistral-common
+Requires-Dist: mistralai-search-toolkit<1.0.0,>=0.0.12
+Requires-Dist: mistralai-workflows<4.0.0,>=3.0.0a1
+Requires-Dist: pydantic-settings>=2.2
+Requires-Dist: pydantic<3.0.0,>=2.7.4

mistralai_workflows_plugins_search-0.1.0/README.md

@@ -0,0 +1,186 @@
+# mistralai-workflows-plugins-search
+
+Runs [Search Toolkit](../../../search/toolkit) ingestion pipelines as durable Temporal workflows,
+with one visible step per stage.
+
+A toolkit `Pipeline.run()` is a single opaque call: if it dies at document 800 of 1000, the whole
+run restarts and re-pays for the extraction already done. This plugin splits the same pipeline into
+four activities with independent retry boundaries, so a failure resumes at the stage that failed and
+a bad document is recorded rather than fatal.
+
+## Topology
+
+```
+IngestDocumentsWorkflow          route → batch → fan out in bounded waves
+  └── plan_batches               one activity, no document content
+  └── IngestBatchWorkflow        one child per batch, homogeneous by mimetype
+        ├── extract              references → documents
+        ├── split                documents → chunks
+        ├── process              chunk enrichers, summarisers, ...
+        └── embed_and_index      embed (unbounded) then index (bounded)
+```
+
+Documents are declared as `OffloadableField`, so that they *can* move between activities as blob
+references rather than as workflow payloads. Each stage output also carries an inline summary —
+counters and failures — because `OffloadableField.get_value()` raises inside a workflow, so the
+summary is the only thing the orchestrating workflow can read.
+
+`embed` and `index` share one activity on purpose: an embedded document is the largest payload the
+pipeline produces, and fusing the two stages is what keeps the vectors from crossing a boundary.
+
+> **Offloading is opt-in, and this plugin's defaults assume it is on.**
+>
+> Declaring a field offloadable does nothing by itself: the core SDK only uploads it when payload
+> offloading is enabled *and* pointed at a blob store, via
+> `TEMPORAL_PAYLOAD_OFFLOADING__ENABLED` and the matching storage configuration. Left unset,
+> documents travel inline through Temporal like any other payload.
+>
+> That matters here because Temporal caps a payload at 2 MB (`MAX_INPUT_SIZE_BYTES` in the SDK),
+> and the default `batch_size` of 30 weighs roughly 3.8 MB at the widest stage boundary. **Without
+> offloading, the default batch does not fit** — you either enable it, or lower `batch_size` to
+> around 10 and accept the smaller retry granularity. `tests/test_payload_sizes.py` measures both
+> regimes.
+
+## Setting up a worker
+
+See [`examples/worker.py`](examples/worker.py). Two registries are the entire setup surface.
+
+```python
+from mistralai.workflows.plugins.search import pipeline_registry
+
+pipeline_registry.register("pdf", pdf_pipeline)
+pipeline_registry.register_router("docs", RoutedPipeline({"ocr": pdf_pipeline}))
+```
+
+A `Pipeline` holds live clients and credentials, so it cannot cross a Temporal payload. Activities
+carry the *name* and resolve the object from worker state.
+
+Which of the two you need depends on the workflow. `IngestBatchWorkflow` carries a `pipeline_name`,
+so it only needs `register()`. `IngestDocumentsWorkflow` carries a `router_name` — but registering
+the router is not enough on its own: **every pipeline the router can return must also be registered
+by name**, because routing resolves a `Pipeline` object and the batch it produces can only carry a
+name. Miss one and `plan_batches` raises `UnregisteredPipelineError` naming what is registered,
+rather than routing the batch through a pipeline it guessed at. That is why `pdf_pipeline` appears
+twice above.
+
+The router keys on **protocol**, not on mimetype. `DEFAULT_MIME_REGISTRY` maps `application/pdf` to
+`ocr`, `text/plain` to `plain_text`, `message/rfc822` to `email`, plus `audio`, `html`, `image`,
+`numbers` and `xlsx`. A file whose protocol the router has no pipeline for is reported as
+unroutable, not raised — one unsupported file does not sink a run.
+
+Activities register themselves at import time, so `run_worker` takes only the workflows.
+
+### Registering your own document types
+
+If your extractors or stores use a `Document`, `DocumentChunk` or metadata subclass, register it
+under a stable id you choose and never change:
+
+```python
+from mistralai.workflows.plugins.search import document_type_registry
+
+document_type_registry.register("acme.contract.v1", ContractDocument)
+```
+
+The id is a contract that outlives the code: it is what a payload written today names, so moving or
+renaming the class must not change it. Everything the toolkit ships is already registered.
+
+Forget it and ingestion fails loudly rather than quietly dropping your fields:
+
+```
+UnregisteredDocumentTypeError: acme.documents:ContractDocument has no stable type id, so a
+payload naming it could not be read back after the class moves. Register a class of your own
+with `DocumentTypeRegistry.register(<id>, <class>)`; a class shipped by the Search Toolkit is
+missing an entry in `_BUILTIN_TYPE_IDS`.
+```
+
+That is deliberate. Without the id, the subclass would come back from the blob as a base
+`Document` — custom top-level fields rejected, custom metadata reduced to untyped extras — and the
+run would report success on documents it had silently downgraded.
+
+## Tuning throughput
+
+`IngestionExecutionPolicy` travels from the parent workflow down to every activity, so throughput is
+a deployment concern rather than a code change. Every value is validated `>= 1`: a zero would leave
+a stage waiting on a semaphore that never opens.
+
+| Field | Default | Governs |
+| --- | --- | --- |
+| `batch_size` | 30 | Documents per child workflow |
+| `max_concurrent_batches` | 5 | Child workflows in flight per wave |
+| `extract_concurrency` | 5 | Documents extracted at once — usually the OCR rate limit |
+| `split_concurrency` | 10 | Documents split at once; CPU-bound, cheap |
+| `process_concurrency` | 10 | Documents through the processor chain at once |
+| `index_concurrency` | 5 | Documents indexed at once; stores have no bulk API yet |
+
+Concurrency always applies *across* documents. The processor chain for a single document stays
+sequential because its stages feed each other; if a processor is not safe to call concurrently, set
+`process_concurrency = 1` rather than changing the activity.
+
+**There is no `embed_concurrency`, deliberately.** Mistral's embedder batches and rate-limits at the
+chunk level across concurrent calls, so the embed phase runs unbounded and its buffer fills across
+the whole batch. A limiter on top would only starve it. That holds for our embedder specifically —
+a caller who plugs in one that does not batch has no way to bound it today, which
+[CTX-3532](https://linear.app/mistral-ai/issue/CTX-3532) tracks.
+
+### Timeouts
+
+Activity and workflow timeouts cannot travel in the policy — they are decorator arguments,
+evaluated when the plugin is imported, and the policy only exists once a run starts. They are read
+from the environment instead, so a deployment still tunes them without a fork. **Changing one takes
+a worker restart.**
+
+| Variable | Default | Applies to |
+| --- | --- | --- |
+| `SEARCH_PLUGIN_EXTRACT_TIMEOUT_MINUTES` | 30 | `extract` |
+| `SEARCH_PLUGIN_SPLIT_TIMEOUT_MINUTES` | 10 | `split` |
+| `SEARCH_PLUGIN_PROCESS_TIMEOUT_MINUTES` | 30 | `process` |
+| `SEARCH_PLUGIN_EMBED_AND_INDEX_TIMEOUT_MINUTES` | 30 | `embed_and_index` |
+| `SEARCH_PLUGIN_PLAN_BATCHES_TIMEOUT_MINUTES` | 5 | `plan_batches` |
+| `SEARCH_PLUGIN_CORPUS_EXECUTION_TIMEOUT_HOURS` | 12 | `IngestDocumentsWorkflow` |
+
+The defaults are per-stage worst cases for a `batch_size` of 30. **Raising `batch_size` means
+raising these too** — nothing links them automatically, because one is a runtime value and the
+other is fixed at import.
+
+`IngestBatchWorkflow`'s own timeout is derived rather than configured: the sum of the four stage
+timeouts, multiplied by the SDK's `RETRY_POLICY_MAX_ATTEMPTS`. The multiplier is not padding. A
+stage timeout is per *attempt*, so without it a single retryable failure in `extract` would burn
+three quarters of the batch's budget and the workflow would die before `split`, `process` and
+`embed_and_index` ever ran — the parent's clock undoing the stage boundaries it exists to protect.
+
+## Error handling
+
+The toolkit decides what is retryable, via `is_retryable_exception`, which walks the `__cause__`
+chain for an `is_retryable` attribute and **assumes retryable when nothing has an opinion**.
+
+- An error that declares `is_retryable = False` becomes a per-document failure; the batch continues
+  and the document is listed in `failures`.
+- Anything else propagates, so Temporal retries the activity. A store outage must not be recorded
+  as thirty permanent document failures.
+
+There is deliberately **no** "every document failed the same way, so it must be environmental"
+rule. By the time a batch is complete, anything without an opinion has already propagated, so such
+a rule could only ever fire on errors that declared themselves non-retryable — re-raising exactly
+the failures that never succeed on retry, burning the activity's attempts and then failing the
+batch, which takes the corpus run down with it. A batch where every document is unusable reports
+every document instead.
+
+Because the default is "retryable", a failure your own code raises must say so itself. The plugin's
+`DocumentInputError` does exactly that for a reference that cannot be resolved: wrapping a
+`KeyError` would not help, since neither it nor the toolkit's `DocumentProcessingError` carries an
+opinion, and a missing object does not become present on retry.
+
+## Development
+
+```bash
+uv sync
+uv run invoke tests
+uv run invoke lint
+uv run invoke typecheck
+```
+
+## References
+
+- [Search Toolkit](../../../search/toolkit) — pipelines, extractors, stores, and the document
+  serialization API this plugin moves documents with
+- [Workflows SDK contributing guide](../../CONTRIBUTING.md)

mistralai_workflows_plugins_search-0.1.0/examples/worker.py

@@ -0,0 +1,48 @@
+"""Minimal worker wiring a toolkit pipeline into the search ingestion workflows."""
+
+import asyncio
+
+from mistralai.search.toolkit.ingestion.pipelines import Pipeline
+from mistralai.search.toolkit.ingestion.pipelines.routed_pipeline import RoutedPipeline
+
+from mistralai.workflows import run_worker
+from mistralai.workflows.plugins.search import (
+    IngestBatchWorkflow,
+    IngestDocumentsWorkflow,
+    pipeline_registry,
+)
+
+
+def build_pdf_pipeline() -> Pipeline:
+    raise NotImplementedError("Construct your toolkit Pipeline here: loader, extractor, splitter, stores.")
+
+
+async def main() -> None:
+    pdf_pipeline = build_pdf_pipeline()
+
+    # A Pipeline holds live clients and credentials, so it cannot cross a Temporal payload.
+    # Activities carry the name and resolve the object from this registry.
+    pipeline_registry.register("pdf", pdf_pipeline)
+
+    # The router keys on protocol, not on mimetype: DEFAULT_MIME_REGISTRY maps application/pdf to
+    # `ocr`, text/plain to `plain_text`, message/rfc822 to `email`, and so on.
+    #
+    # `pdf_pipeline` deliberately appears twice. Routing resolves a Pipeline *object*, and the batch
+    # it produces can only carry a *name*, so every pipeline this router can return also needs its
+    # own `register()` call above — otherwise plan_batches raises UnregisteredPipelineError.
+    pipeline_registry.register_router("docs", RoutedPipeline({"ocr": pdf_pipeline}))
+
+    # Register any Document, DocumentChunk or metadata subclass your extractors or stores rely on,
+    # each under a stable id you choose and never change. The toolkit's own classes are already
+    # registered. Skipping this fails the ingestion — see the README.
+    #
+    #   from mistralai.workflows.plugins.search import document_type_registry
+    #   document_type_registry.register("acme.contract.v1", ContractDocument)
+
+    # Activities register themselves with the worker at import time, so `run_worker` takes only the
+    # workflows; importing this plugin is what makes the five activities available.
+    await run_worker([IngestDocumentsWorkflow, IngestBatchWorkflow])
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

mistralai_workflows_plugins_search-0.1.0/mistralai/workflows/__init__.py

@@ -0,0 +1,6 @@
+# This file helps PyLance/VS Code and mypy resolve imports for our PEP 420 namespace
+# mistralai/workflows/plugins. Excluded from the wheel — see pyproject.toml.
+__path__ = __import__("pkgutil").extend_path(__path__, __name__)
+
+from mistralai.workflows.exports import *  # noqa: F403
+from mistralai.workflows.exports import __all__ as __all__

mistralai_workflows_plugins_search-0.1.0/mistralai/workflows/plugins/search/__init__.py

@@ -0,0 +1,56 @@
+"""Search Toolkit ingestion as durable Mistral Workflows."""
+
+from mistralai.workflows.plugins.search.activities import (
+    embed_and_index,
+    extract,
+    plan_batches,
+    process,
+    split,
+)
+from mistralai.workflows.plugins.search.models import (
+    DocumentFailure,
+    DocumentRef,
+    IngestBatchInput,
+    IngestBatchOutput,
+    IngestDocumentsInput,
+    IngestDocumentsOutput,
+    IngestionExecutionPolicy,
+)
+from mistralai.workflows.plugins.search.registry import (
+    PipelineRegistry,
+    UnknownPipelineError,
+    UnknownRouterError,
+    UnregisteredPipelineError,
+    document_type_registry,
+    pipeline_registry,
+)
+from mistralai.workflows.plugins.search.workflows import (
+    IngestBatchWorkflow,
+    IngestDocumentsWorkflow,
+)
+
+# Serialization types are deliberately not re-exported. They are the Search Toolkit's public API,
+# and a caller who needs them should import from `mistralai.search.toolkit.document.serialization`
+# so there is exactly one obvious home for them.
+__all__ = [
+    "DocumentFailure",
+    "DocumentRef",
+    "IngestBatchInput",
+    "IngestBatchOutput",
+    "IngestBatchWorkflow",
+    "IngestDocumentsInput",
+    "IngestDocumentsOutput",
+    "IngestDocumentsWorkflow",
+    "IngestionExecutionPolicy",
+    "PipelineRegistry",
+    "UnknownPipelineError",
+    "UnknownRouterError",
+    "UnregisteredPipelineError",
+    "document_type_registry",
+    "embed_and_index",
+    "extract",
+    "pipeline_registry",
+    "plan_batches",
+    "process",
+    "split",
+]

mistralai_workflows_plugins_search-0.1.0/mistralai/workflows/plugins/search/_stages.py

@@ -0,0 +1,87 @@
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass
+
+from mistralai.search.toolkit.document import Document
+from mistralai.search.toolkit.errors import is_retryable_exception
+
+from mistralai.workflows.plugins.search.models import DocumentFailure
+
+
+class DocumentInputError(Exception):
+    """The plugin established that a reference cannot yield a document.
+
+    Declares ``is_retryable = False`` so the toolkit predicate keeps it out of the retry path.
+    ``is_retryable_exception`` assumes retryable when nothing in the ``__cause__`` chain has an
+    opinion, which is the right default for an unrecognised failure and the wrong one here: a
+    missing object does not become present on retry. Wrapping the cause is not enough either — a
+    ``KeyError`` has no opinion, and neither does the toolkit's own ``DocumentProcessingError``.
+    """
+
+    is_retryable = False
+
+
+@dataclass(frozen=True, slots=True)
+class _Outcome:
+    document: Document | None = None
+    failure: DocumentFailure | None = None
+    retryable: Exception | None = None
+
+
+async def run_document_stage(
+    documents: list[Document],
+    *,
+    stage: str,
+    apply: Callable[[Document], Awaitable[Document]],
+    concurrency: int | None = None,
+) -> tuple[list[Document], list[DocumentFailure]]:
+    """Apply one pipeline stage across a batch, concurrently, isolating per-document failures.
+
+    ``concurrency`` bounds how many documents are in flight. ``None`` leaves the stage unbounded,
+    which is what embedding wants: the embedder batches at the chunk level across concurrent calls,
+    so a limiter on top of it only starves its buffer.
+
+    Errors the toolkit classifies as retryable propagate, so Temporal retries the activity — a store
+    outage must not be recorded as thirty permanent document failures. Only an error that declares
+    ``is_retryable = False``, directly or through its ``__cause__`` chain, becomes a per-document
+    failure that lets the batch continue; the toolkit assumes retryable when nothing has an opinion,
+    and over-retrying a doomed document is cheaper than abandoning a recoverable one.
+
+    There is deliberately no "every document failed the same way, so it must be environmental" net.
+    Such a net can only ever observe errors that already declared themselves non-retryable, because
+    anything without an opinion propagates before the batch is complete — so it would re-raise
+    precisely the failures that will never succeed on retry, burning the activity's attempts and
+    then failing the batch, which takes the whole corpus run down with it. A batch where every
+    document is genuinely unusable reports every document, which is both cheaper and more legible.
+    """
+    if not documents:
+        return [], []
+
+    semaphore = asyncio.Semaphore(concurrency) if concurrency is not None else None
+
+    async def run_one(document: Document) -> _Outcome:
+        try:
+            if semaphore is None:
+                return _Outcome(document=await apply(document))
+            async with semaphore:
+                return _Outcome(document=await apply(document))
+        except Exception as exc:  # noqa: BLE001 - classified below, never swallowed
+            if is_retryable_exception(exc):
+                return _Outcome(retryable=exc)
+            return _Outcome(failure=DocumentFailure(source_id=document.source_id, stage=stage, message=str(exc)))
+
+    # gather preserves input order, so `processed` stays aligned with `documents`. Every document is
+    # attempted before any decision is taken: the semaphore caps what an outage costs, and aborting
+    # on the first error would discard the failures already collected.
+    outcomes = await asyncio.gather(*(run_one(document) for document in documents))
+
+    for outcome in outcomes:
+        if outcome.retryable is not None:
+            raise outcome.retryable
+
+    processed = [o.document for o in outcomes if o.document is not None]
+    failures = [o.failure for o in outcomes if o.failure is not None]
+
+    return processed, failures