pandion 0.1.0__py3-none-any.whl

pandion/__init__.py

@@ -0,0 +1,45 @@
+"""Pandion: indexing and retrieval pipelines you can inspect stage by stage.
+
+The engine lives in plain modules (`pandion.graph`, `pandion.executor`,
+`pandion.storage` and the rest). A name that starts with an underscore is
+private; the RAG Playground app may import every other name.
+
+`pandion.parse`, `pandion.clean`, `pandion.chunk` and `pandion.rewrite` run one
+stage when called, and hold that stage's step classes. `pandion.retrieve`,
+`pandion.rerank` and `pandion.answer` hold the step classes of the rest.
+"""
+
+from pandion.models import ModelSettings, preload_embedder, preload_reranker
+from pandion.pipeline import (
+    Answer,
+    IndexPipeline,
+    Pipeline,
+    QueryPipeline,
+    Retrieval,
+    RetrievalPipeline,
+)
+from pandion.results import StageResult
+from pandion import answer, chunk, clean, parse, rerank, retrieve, rewrite
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "Answer",
+    "IndexPipeline",
+    "ModelSettings",
+    "Pipeline",
+    "QueryPipeline",
+    "Retrieval",
+    "RetrievalPipeline",
+    "StageResult",
+    "answer",
+    "chunk",
+    "clean",
+    "parse",
+    "preload_embedder",
+    "preload_reranker",
+    "rerank",
+    "retrieve",
+    "rewrite",
+    "__version__",
+]

pandion/answer.py

@@ -0,0 +1,16 @@
+"""The answer steps, as classes for a pipeline written in code.
+
+`Search(top_k=5)` shows the pieces retrieval found, as they are. `Chat(model=
+"claude-haiku-4-5")` asks a chat model for an answer whose claims cite those
+pieces. Each is the same as its block, for example `{"use": "search",
+"top_k": 5}` (`pandion.config.step_class`). Answering runs through a
+pipeline: `p.ask(question)`.
+"""
+
+from pandion.config import step_class
+from pandion.ports import Stage
+
+Search = step_class(Stage.USE_CASE, "search", __name__)
+Chat = step_class(Stage.USE_CASE, "chat", __name__)
+
+__all__ = ["Chat", "Search"]

pandion/artifacts.py

@@ -0,0 +1,44 @@
+"""Artifact types and the immutable artifact value.
+
+An artifact is what a transform produces: an immutable value identified by a
+recipe hash (see `pandion.ids`). The payload itself lives on disk and is loaded
+lazily; `meta` is small and always in memory.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import StrEnum
+from typing import Any
+
+
+class ArtifactType(StrEnum):
+    """The value types that flow between pipeline stages.
+
+    A StrEnum rather than a Literal so that `inputs = {"file": ArtifactType.RAW_FILE}`
+    typechecks in plugins without an explicit annotation, and so it serializes to
+    JSON as a plain string.
+    """
+
+    RAW_FILE = "raw_file"
+    QUERY = "query"
+    PARSED_DOC = "parsed_doc"
+    CHUNK_SET = "chunk_set"
+    INDEX = "index"
+    RETRIEVAL_RESULT = "retrieval_result"
+    OUTPUT = "output"
+    QA_SET = "qa_set"
+
+
+@dataclass(frozen=True)
+class Artifact:
+    """A committed pipeline value.
+
+    `id` is a *recipe* hash, not a content hash — it is computed from the
+    transform and its inputs before the payload exists. That is sound only for
+    deterministic transforms; see `Transform.deterministic`.
+    """
+
+    id: str
+    type: ArtifactType
+    meta: dict[str, Any] = field(default_factory=dict)

pandion/chunk.py

@@ -0,0 +1,23 @@
+"""The chunkers: call this module to chunk one document, or use a class in a pipeline.
+
+`pandion.chunk(doc, use="recursive_character", chunk_size=400)` cuts one parsed
+document into pieces (`pandion.stages.chunk`). `RecursiveCharacter(chunk_size=400)`
+is the same settings for a pipeline written in code (`pandion.config.step_class`).
+"""
+
+import sys
+
+from pandion.config import step_class
+from pandion.ports import Stage
+from pandion.stages import StageModule
+from pandion.stages import chunk as run
+
+RecursiveCharacter = step_class(Stage.CHUNK, "recursive_character", __name__)
+LayoutBlocks = step_class(Stage.CHUNK, "layout_blocks", __name__)
+MarkdownHeader = step_class(Stage.CHUNK, "markdown_header", __name__)
+TokenBased = step_class(Stage.CHUNK, "token_based", __name__)
+SentenceWindow = step_class(Stage.CHUNK, "sentence_window", __name__)
+
+__all__ = ["LayoutBlocks", "MarkdownHeader", "RecursiveCharacter", "SentenceWindow", "TokenBased"]
+
+sys.modules[__name__].__class__ = StageModule

pandion/clean.py

@@ -0,0 +1,21 @@
+"""The cleaners: call this module to clean one document, or use a class in a pipeline.
+
+`pandion.clean(doc, steps=[{"use": "dedupe_blocks"}])` cleans one parsed
+document (`pandion.stages.clean`). `DedupeBlocks()` is the same step for a
+pipeline written in code (`pandion.config.step_class`).
+"""
+
+import sys
+
+from pandion.config import step_class
+from pandion.ports import Stage
+from pandion.stages import StageModule
+from pandion.stages import clean as run
+
+HeaderFooterStrip = step_class(Stage.CLEAN, "header_footer_strip", __name__)
+DedupeBlocks = step_class(Stage.CLEAN, "dedupe_blocks", __name__)
+DropMatching = step_class(Stage.CLEAN, "drop_matching", __name__)
+
+__all__ = ["DedupeBlocks", "DropMatching", "HeaderFooterStrip"]
+
+sys.modules[__name__].__class__ = StageModule

pandion/cli.py

@@ -0,0 +1,87 @@
+"""The `pandion` command.
+
+`pandion install pipeline.yaml` reads the file, works out which extras its
+steps need, and installs them with pip into the Python that runs it.
+`--print` only prints the command. Nothing is ever installed while a pipeline
+runs.
+"""
+
+from __future__ import annotations
+
+import argparse
+import subprocess
+import sys
+
+from pandion.config import ConfigError, PipelineConfig, load
+
+#: The extra each strategy or embedder needs. A name that is not here needs
+#: only the base install.
+EXTRAS = {
+    "pdfium": "pdfium",
+    "docling": "docling",
+    "lancedb": "lancedb",
+    "qwen3-embedding-0.6b": "local",
+    "bge-small-en-v1.5": "local",
+    "mmr": "local",
+    "cross_encoder": "local",
+}
+
+#: Steps that call a chat model, by stage; their extra depends on the model's provider.
+CHAT_STEPS = {"llm_rewrite": "query", "llm_rerank": "rerank", "chat": "use_case"}
+
+
+def extras_for(config: PipelineConfig) -> list[str]:
+    """The extras the pipeline's steps need, sorted."""
+    index = config.index
+    names = [index.parse[0], *(use for use, _ in index.clean), index.chunk[0]]
+    names += ["lancedb", index.store["embedder"]]
+    extras = {EXTRAS[name] for name in names if name in EXTRAS}
+    retrieve = config.retrieve
+    if retrieve is not None:
+        steps = [*([retrieve.rewrite] if retrieve.rewrite else []), retrieve.search, *retrieve.rerank]
+        if config.answer is not None:
+            steps.append(config.answer)
+        for use, settings in steps:
+            if use in CHAT_STEPS:
+                extras.add(_chat_extra(use, settings))
+            elif use in EXTRAS:
+                extras.add(EXTRAS[use])
+    return sorted(extras)
+
+
+def _chat_extra(use: str, settings: dict) -> str:
+    """Claude models need the Anthropic SDK; every other model the OpenAI SDK."""
+    from pandion.llm import CHAT_MODELS
+    from pandion.plugins import discover
+    from pandion.ports import Stage
+    from pandion.registry import registry
+
+    discover()
+    model = registry.get(Stage(CHAT_STEPS[use]), use).config_model(**settings).model
+    return "anthropic" if CHAT_MODELS[model].provider == "anthropic" else "openai"
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(prog="pandion")
+    commands = parser.add_subparsers(dest="command", required=True)
+    install = commands.add_parser("install", help="install the extras a pipeline file needs")
+    install.add_argument("config", help="the pipeline file, for example pipeline.yaml")
+    install.add_argument("--print", action="store_true", help="only print the pip command")
+    args = parser.parse_args(argv)
+
+    try:
+        config = load(args.config)
+    except (ConfigError, OSError) as exc:
+        print(f"pandion: {exc}", file=sys.stderr)
+        return 2
+    extras = extras_for(config)
+    package = f"pandion[{','.join(extras)}]" if extras else "pandion"
+    command = [sys.executable, "-m", "pip", "install", package]
+    if args.print:
+        print(f'"{sys.executable}" -m pip install "{package}"')
+        return 0
+    return subprocess.call(command)
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

pandion/config.py

@@ -0,0 +1,381 @@
+"""Read a pipeline.yaml and check it before anything runs.
+
+```yaml
+path: ./my-index            # where sources, artifacts and the embedding cache go
+models:                     # optional: pandion.models.ModelSettings
+  download: auto
+index:
+  parse: {use: docling, do_ocr: false}
+  clean:                    # optional: a list of cleaners, run in order
+    - use: dedupe_blocks
+  chunk: {use: recursive_character, chunk_size: 400}
+  embed: {use: qwen3-embedding-0.6b, truncate_dim: null}
+  store: {use: lancedb, metric: cosine, build_fts: true}
+retrieve:                   # optional: what p.retrieve runs
+  rewrite: none             # optional: none, or {use: llm_rewrite, style: keywords}
+  search: {use: hybrid_rrf, top_k: 20}
+  rerank:                   # optional: none, one step, or a list run in order
+    - use: cross_encoder
+      top_k: 5
+answer:                     # optional: what p.ask runs after retrieval
+  use: chat                 # search (the pieces as they are) or chat
+  model: claude-haiku-4-5
+```
+
+Every key under a step other than `use` is a setting of that step, with the
+same name as in its config model, so a file and the same pipeline built in the
+RAG Playground hash to the same artifact ids. `embed` and `store` are one step
+when it runs: the lancedb index embeds while it writes, so `embed.use` becomes
+its `embedder` setting. A relative `path` or `models.cache_dir` is read from
+the folder the file is in.
+
+The question is not part of the file: it is passed to `p.retrieve` or `p.ask`,
+which put it in the query step's `text`. A model's key is never in the file
+either: it is passed to `p.retrieve` or `p.ask` as `credentials`.
+
+The same pipeline can be written in code (`build`): each stage as a dict that
+is the file's block for that stage, as a step class such as
+`pandion.parse.Docling(do_ocr=False)` (`step_class`), or as a plain name, which
+is short for `{"use": name}`. Every way goes through `from_dict`, so it is
+checked the same way and gives the same `PipelineConfig`. `to_dict` turns a
+`PipelineConfig` back into the file's shape.
+
+Anything the file gets wrong is a `ConfigError` that names the step and what
+it accepts.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, ClassVar
+
+import yaml
+from pydantic import BaseModel, ValidationError
+
+from pandion.embeddings import EMBEDDER_NAMES
+from pandion.models import ModelSettings
+from pandion.plugins import discover
+from pandion.ports import Stage
+from pandion.registry import UnknownTransformError, registry
+from pandion.stages import _strategy
+
+TOP_KEYS = ("path", "models", "index", "retrieve", "answer")
+RETRIEVE_KEYS = ("rewrite", "search", "rerank")
+INDEX_KEYS = ("parse", "clean", "chunk", "embed", "store")
+EMBED_KEYS = ("truncate_dim",)
+STORE_KEYS = ("metric", "build_fts")
+#: The answer steps. `eval` is a use case too, but it scores, it does not answer.
+ANSWER_STEPS = ("search", "chat")
+
+#: Where each stage's step sits in the file, for messages.
+WHERE = {
+    Stage.PARSE: "index.parse",
+    Stage.CLEAN: "index.clean",
+    Stage.CHUNK: "index.chunk",
+    Stage.QUERY: "retrieve.rewrite",
+    Stage.RETRIEVE: "retrieve.search",
+    Stage.RERANK: "retrieve.rerank",
+    Stage.USE_CASE: "answer",
+}
+
+
+class ConfigError(ValueError):
+    """A pipeline file that cannot run, with what is wrong and where."""
+
+
+@dataclass
+class IndexConfig:
+    """The index steps, each as (strategy name, its settings)."""
+
+    parse: tuple[str, dict[str, Any]]
+    chunk: tuple[str, dict[str, Any]]
+    #: The lancedb index's settings, with `embedder` taken from `embed.use`.
+    store: dict[str, Any]
+    clean: list[tuple[str, dict[str, Any]]] = field(default_factory=list)
+
+
+@dataclass
+class RetrieveConfig:
+    """The retrieval steps. No rewrite means the question is searched as asked."""
+
+    search: tuple[str, dict[str, Any]]
+    rewrite: tuple[str, dict[str, Any]] | None = None
+    rerank: list[tuple[str, dict[str, Any]]] = field(default_factory=list)
+
+
+@dataclass
+class PipelineConfig:
+    path: Path
+    models: ModelSettings
+    index: IndexConfig
+    retrieve: RetrieveConfig | None = None
+    #: The answer step, as (strategy name, its settings).
+    answer: tuple[str, dict[str, Any]] | None = None
+
+
+def load(file: str | Path) -> PipelineConfig:
+    """Read and check the pipeline file at `file`."""
+    file = Path(file)
+    data = yaml.safe_load(file.read_text(encoding="utf-8"))
+    return from_dict(data, base=file.parent)
+
+
+def from_dict(data: Any, *, base: Path) -> PipelineConfig:
+    """Check a loaded pipeline file. Relative paths are read from `base`."""
+    _keys(data, TOP_KEYS, "The pipeline file")
+    if not data.get("path"):
+        raise ConfigError(
+            "The pipeline file has no path: the folder where Pandion keeps the "
+            "sources, the artifacts and the embedding cache."
+        )
+    models = dict(data.get("models") or {})
+    if models.get("cache_dir"):
+        models["cache_dir"] = _resolve(models["cache_dir"], base)
+    try:
+        settings = ModelSettings(**models)
+    except ValidationError as exc:
+        problems = "; ".join(
+            f"models.{'.'.join(map(str, e['loc']))}: {e['msg']}" for e in exc.errors()
+        )
+        raise ConfigError(problems) from None
+    if "index" not in data:
+        raise ConfigError("The pipeline file has no index section.")
+    if "answer" in data and "retrieve" not in data:
+        raise ConfigError(
+            "The answer needs a retrieve section: the answer step reads what "
+            "retrieval found."
+        )
+    return PipelineConfig(
+        path=_resolve(data["path"], base),
+        models=settings,
+        index=_index(data["index"]),
+        retrieve=_retrieve(data["retrieve"]) if "retrieve" in data else None,
+        answer=_answer(data["answer"]) if "answer" in data else None,
+    )
+
+
+def build(
+    *,
+    path: str | Path,
+    models: ModelSettings | dict[str, Any] | None = None,
+    index: dict[str, Any],
+    retrieve: dict[str, Any] | None = None,
+    answer: Any = None,
+) -> PipelineConfig:
+    """Check a pipeline written in code; relative paths are read from the
+    current folder.
+
+    Each stage is a dict (the file's block for that stage), a step class
+    (`step_class`), a plain name for `{"use": name}`, or a list of these
+    where the file takes a list. A stage that is None is left out.
+    """
+    if isinstance(models, BaseModel):
+        models = models.model_dump(exclude_unset=True)
+    data: dict[str, Any] = {"path": path, "models": models or {}}
+    data["index"] = {k: _block(v) for k, v in index.items() if v is not None}
+    if isinstance(data["index"].get("clean"), dict):
+        data["index"]["clean"] = [data["index"]["clean"]]
+    if retrieve is not None and any(v is not None for v in retrieve.values()):
+        data["retrieve"] = {k: _block(v) for k, v in retrieve.items() if v is not None}
+    if answer is not None:
+        data["answer"] = _block(answer)
+    return from_dict(data, base=Path("."))
+
+
+def _block(value: Any) -> Any:
+    """One stage written in code, as the file's block for it."""
+    if isinstance(value, str):
+        return value if value == "none" else {"use": value}
+    if isinstance(value, BaseModel) and isinstance(getattr(type(value), "use", None), str):
+        return {"use": type(value).use, **value.model_dump(mode="json", exclude_unset=True)}
+    if isinstance(value, (list, tuple)):
+        return [_block(v) for v in value]
+    return value
+
+
+def to_dict(config: PipelineConfig) -> dict[str, Any]:
+    """`config` in the file's shape: `from_dict` of it gives `config` again."""
+    index = config.index
+    store = dict(index.store)
+    embed = {"use": store.pop("embedder")}
+    embed.update({key: store.pop(key) for key in EMBED_KEYS if key in store})
+    data: dict[str, Any] = {"path": config.path.as_posix()}
+    models = config.models.model_dump(mode="json", exclude_defaults=True)
+    if models:
+        data["models"] = models
+    data["index"] = {"parse": _out(index.parse)}
+    if index.clean:
+        data["index"]["clean"] = [_out(step) for step in index.clean]
+    data["index"].update(chunk=_out(index.chunk), embed=embed, store={"use": "lancedb", **store})
+    if config.retrieve is not None:
+        retrieve = config.retrieve
+        data["retrieve"] = {}
+        if retrieve.rewrite is not None:
+            data["retrieve"]["rewrite"] = _out(retrieve.rewrite)
+        data["retrieve"]["search"] = _out(retrieve.search)
+        if retrieve.rerank:
+            data["retrieve"]["rerank"] = [_out(step) for step in retrieve.rerank]
+    if config.answer is not None:
+        data["answer"] = _out(config.answer)
+    return data
+
+
+def _out(step: tuple[str, dict[str, Any]]) -> dict[str, Any]:
+    use, settings = step
+    return {"use": use, **settings}
+
+
+def step_class(stage: Stage, use: str, module: str) -> type[BaseModel]:
+    """A typed class for the `use` strategy of `stage`: its config model, which
+    also knows its strategy's name, so a pipeline can take it in place of a dict.
+
+    `Docling(do_ocr=False)` has the same fields and validation as the docling
+    config model, keeps only the settings it was given (so it makes the same
+    node as `{"use": "docling", "do_ocr": False}`), and refuses an unknown
+    setting at once, naming the stage and listing the settings it has.
+    """
+    discover()
+    strategy = registry.get(stage, use)
+    model = strategy.config_model
+    where = WHERE[stage]
+
+    def __init__(self, **settings: Any) -> None:
+        try:
+            _strategy(stage, use, settings, where=where)
+        except ValueError as exc:
+            raise ConfigError(str(exc)) from None
+        model.__init__(self, **settings)
+
+    name = "".join(part.capitalize() for part in use.split("_"))
+    return type(
+        name,
+        (model,),
+        {
+            "__init__": __init__,
+            "__module__": module,
+            "__qualname__": name,
+            "__doc__": f"The {use} {stage} step. {strategy.summary}",
+            "__annotations__": {"use": ClassVar[str]},
+            "use": use,
+        },
+    )
+
+
+def _answer(answer: Any) -> tuple[str, dict[str, Any]]:
+    use, settings = _use(answer, "answer")
+    if use not in ANSWER_STEPS:
+        raise ConfigError(
+            f"answer uses '{use}', which is not an answer step. "
+            f"The answer steps are: {', '.join(ANSWER_STEPS)}."
+        )
+    discover()
+    return _step(Stage.USE_CASE, answer, "answer")
+
+
+def _index(index: Any) -> IndexConfig:
+    _keys(index, INDEX_KEYS, "index")
+    for stage in ("parse", "chunk", "embed", "store"):
+        if stage not in index:
+            raise ConfigError(
+                f"The index has no {stage} step. An index needs parse, chunk, "
+                "embed and store; clean may be left out."
+            )
+    discover()
+    steps = index.get("clean") or []
+    if not isinstance(steps, list):
+        raise ConfigError("index.clean must be a list of cleaners, each with use.")
+    return IndexConfig(
+        parse=_step(Stage.PARSE, index["parse"], "index.parse"),
+        clean=[
+            _step(Stage.CLEAN, step, f"index.clean step {n}")
+            for n, step in enumerate(steps, start=1)
+        ],
+        chunk=_step(Stage.CHUNK, index["chunk"], "index.chunk"),
+        store=_store(index["embed"], index["store"]),
+    )
+
+
+def _retrieve(retrieve: Any) -> RetrieveConfig:
+    _keys(retrieve, RETRIEVE_KEYS, "retrieve")
+    if not retrieve.get("search"):
+        raise ConfigError(
+            "retrieve has no search step. Retrieval needs search; rewrite and "
+            "rerank may be left out."
+        )
+    discover()
+    rewrite = retrieve.get("rewrite")
+    if _none(rewrite):
+        rewrite = None
+    else:
+        if isinstance(rewrite, dict) and "text" in rewrite:
+            raise ConfigError(
+                "retrieve.rewrite sets text, but the question is not part of the "
+                "file: it is passed to p.retrieve."
+            )
+        rewrite = _step(Stage.QUERY, rewrite, "retrieve.rewrite")
+    rerank = retrieve.get("rerank")
+    steps = [] if _none(rerank) else rerank if isinstance(rerank, list) else [rerank]
+    return RetrieveConfig(
+        rewrite=rewrite,
+        search=_step(Stage.RETRIEVE, retrieve["search"], "retrieve.search"),
+        rerank=[
+            _step(Stage.RERANK, step, "retrieve.rerank" if len(steps) == 1 else f"retrieve.rerank step {n}")
+            for n, step in enumerate(steps, start=1)
+        ],
+    )
+
+
+def _none(step: Any) -> bool:
+    """Left out, or written as `none`."""
+    return step is None or step == "none"
+
+
+def _step(stage: Stage, step: Any, where: str) -> tuple[str, dict[str, Any]]:
+    use, settings = _use(step, where)
+    try:
+        _strategy(stage, use, settings, where=where)
+    except (UnknownTransformError, ValueError) as exc:
+        raise ConfigError(str(exc)) from None
+    return use, settings
+
+
+def _store(embed: Any, store: Any) -> dict[str, Any]:
+    embedder, embed_settings = _use(embed, "index.embed")
+    kind, store_settings = _use(store, "index.store")
+    if embedder not in EMBEDDER_NAMES:
+        raise ConfigError(
+            f"index.embed uses '{embedder}', which is not an embedder. "
+            f"The embedders are: {', '.join(EMBEDDER_NAMES)}."
+        )
+    if kind != "lancedb":
+        raise ConfigError(f"index.store uses '{kind}'. The only store is lancedb.")
+    _keys(embed_settings, EMBED_KEYS, "index.embed")
+    _keys(store_settings, STORE_KEYS, "index.store")
+    settings = {"embedder": embedder, **embed_settings, **store_settings}
+    try:
+        _strategy(Stage.INDEX, "lancedb", settings, where="index.store")
+    except ValueError as exc:
+        raise ConfigError(str(exc)) from None
+    return settings
+
+
+def _use(step: Any, where: str) -> tuple[str, dict[str, Any]]:
+    if not isinstance(step, dict) or not step.get("use"):
+        raise ConfigError(f"{where} must say which strategy to use, as use: <name>.")
+    settings = dict(step)
+    return settings.pop("use"), settings
+
+
+def _keys(data: Any, valid: tuple[str, ...], where: str) -> None:
+    if not isinstance(data, dict):
+        raise ConfigError(f"{where} must be a mapping of settings.")
+    unknown = [key for key in data if key not in valid]
+    if unknown:
+        names = ", ".join(f"'{key}'" for key in unknown)
+        raise ConfigError(f"{where} has no setting {names}. Its settings are: {', '.join(valid)}.")
+
+
+def _resolve(value: Any, base: Path) -> Path:
+    path = Path(str(value)).expanduser()
+    return path if path.is_absolute() else base / path