Document pipeline (ai-pipeline)

License

The source embedded on this page is Kumiko Enterprise code. It is proprietary and not covered by the framework’s MIT license; running it in production needs a Kumiko Enterprise plan. See plans and pricing.

A domain-neutral document pipeline that composes all three ai-pipeline step kinds (#258) into one run: ai.extract pulls a title and key points out of an uploaded document, ai.classify routes it against a review-action catalogue using that extraction, ai.generate writes a two-sentence reviewer summary from both. The result is persisted only once every step succeeds. This is the reference sample for the AI-step vocabulary — it shows the building blocks themselves, not a specific use case (Joel’s DD-intake pipeline is one possible consumer of this vocabulary and is deliberately not built here).

What it shows

• Steps chain through ctx.steps — ai.classify’s input is built from ai.extract’s own output (title + key points), and ai.generate’s input reads both prior results. The test asserts the actual extracted text reaches the classify and generate provider requests, not just that all three ran.
• No workflow-runner needed today — process-document calls runPipeline directly against a synthetic workflow context, the same pattern ai-pipeline’s own integration test harness uses. The dispatcher/ctx handed to runPipeline is real; only the “a workflow engine invoked this” wrapper is stood in for, pending the real workflow-runner (M.4, kumiko-framework) — not a simplification specific to this sample.
• Persist only on full success, extract gates the rest — classify and generate sit in an r.step.branch that only runs when extract produced data, so a failed extract costs no further provider calls. The write handler then handles the results as an explicit union; the first non-success returns a structured 422 naming which step failed, and nothing is written. A disabled classify still lets generate run (it is skipped in-process, not by the branch). r.step.return has to come after the branch, not inside it.
• Extracted data is fenced as data — classify and generate receive the extraction inside <document_data> tags, and their prompts say that content there is data, not instructions. Extraction output derives from an untrusted document, so this keeps embedded text from steering later steps.
• The policy catalogue and the step usage share one const — the same extractSpec/classifySpec/generateSpec objects are passed both to aiExtractStep()/aiClassifyStep()/aiGenerateStep() (the step usage) and, as documentPipelineStepSpecs, to createAiPipelineFeature() (the policy catalogue) — AiStepRuntimeSpec is a structural superset of AiStepSpec, so one declaration serves both.
• “Upload” is just a payload field here — the document’s base64 PDF bytes ride directly on process-document’s write payload. Real file storage, OCR routing and the fileRef.created trigger are the document-ingest recipe’s job; composing the two is a separate concern from demoing the AI-step vocabulary.

When to reach for it

Any pipeline that needs to pull structured data out of a document, route it, and hand a human a summary before they act — invoice intake, contract triage, compliance document review. Swap the extraction schema, the action catalogue and the summary prompt for your domain; the three-step shape and the policy wiring stay the same.

Wiring

document-pipeline requires ai-pipeline, mounted with this recipe’s own step catalogue via createAiPipelineFeature(documentPipelineStepSpecs). ai-pipeline in turn requires ai-foundation, prompt-store and tenant. Mount an LLM provider plugin (ai-provider-anthropic, ai-provider-openai-compat, or a mock in tests) alongside, and set the ai-foundation:config:text-connection config-key (or a per-step providerId via ai-pipeline:write:set-policy) — none of this sample’s steps pin a provider by default.

import { createAiPipelineFeature } from "@cosmicdriftgamestudio/kumiko-ai-pipeline";
import { documentPipelineFeature, documentPipelineStepSpecs } from "./feature";

const features = [
  createAiPipelineFeature(documentPipelineStepSpecs),
  aiFoundationFeature,
  promptStoreFeature,
  yourLlmProviderPlugin,
  documentPipelineFeature,
  ...appFeatures,
];

The integration test under src/__tests__/ boots that exact stack against real Postgres + Redis and scripts the provider — including the failing-step and disabled-step paths — so every assertion above is proven, not described.

bun --env-file=../.env test --config=bunfig.integration.toml samples/recipes/ai-pipeline

---

Source code

The feature entry point — embedded straight from the source file, so the code here is exactly what runs. Multi-file samples keep their remaining files next to it on GitHub (link below):

// Domain-neutral extract -> classify -> generate document pipeline over the
// three ai-pipeline step kinds; see README.md for what it shows and why.

import { randomUUID } from "node:crypto";
import {
  createEntityExecutor,
  defineEntityDetailHandler,
  defineEntityListHandler,
  defineFeature,
  defineWorkflow,
  type PipelineCtx,
  runPipeline,
  stepsPipeline,
  withResponseData,
} from "@cosmicdrift/kumiko-framework/engine";
import { failUnprocessable } from "@cosmicdrift/kumiko-framework/errors";
import { fenceAsData } from "@cosmicdriftgamestudio/kumiko-ai-foundation";
import {
  type AiClassifyAction,
  type AiClassifyResult,
  type AiClassifyStepArgs,
  type AiExtractDocument,
  type AiExtractResult,
  type AiExtractStepArgs,
  type AiGenerateResult,
  type AiGenerateStepArgs,
  type AiStepSpec,
  aiClassifyStep,
  aiExtractStep,
  aiGenerateStep,
} from "@cosmicdriftgamestudio/kumiko-ai-pipeline";
import { z } from "zod";
import { pipelineRunEntity } from "./entity.ts";

const EXTRACT_STEP_KEY = "document-pipeline:extract";
const CLASSIFY_STEP_KEY = "document-pipeline:classify";
const GENERATE_STEP_KEY = "document-pipeline:generate";

const paramsSchema = z.object({ maxTokens: z.number().int().positive().optional() });

// The catalogue is the classify step's integration contract (same rule as
// the ai-triage recipe): ai.classify never invents an action type, only
// picks from what the caller hands it.
const REVIEW_ACTIONS: readonly AiClassifyAction[] = [
  { type: "route-to-finance", description: "Send to finance for invoice or payment processing." },
  { type: "route-to-legal", description: "Send to legal for contract review." },
  { type: "archive", description: "No action needed; file for reference." },
];

function payloadOf(ctx: PipelineCtx): ProcessDocumentPayload {
  // @cast-boundary pipeline resolver — defineStep's `run()` is declared once
  // for every step instance across the app, so PipelineCtx carries no
  // per-instance payload type. Narrowed here to this feature's own schema.
  return ctx.event.payload as ProcessDocumentPayload;
}

type ExtractedResult = Extract<AiExtractResult, { type: "extracted" }>;
type ClassifiedResult = Extract<AiClassifyResult, { type: "classified" }>;

function hasResultType(value: unknown, type: string): boolean {
  return typeof value === "object" && value !== null && "type" in value && value.type === type;
}

function isExtracted(value: unknown): value is ExtractedResult {
  return hasResultType(value, "extracted");
}

function isClassified(value: unknown): value is ClassifiedResult {
  return hasResultType(value, "classified");
}

function isGenerated(value: unknown): value is Extract<AiGenerateResult, { type: "generated" }> {
  return hasResultType(value, "generated");
}

const extractOutputSchema = z.object({
  title: z.string(),
  keyPoints: z.array(z.string()).max(5),
});

// Extraction output is document-derived, so it is fenced as data in later prompts.
const fenceDocumentData = (content: string): string => fenceAsData(content, "document_data");

function extractedFields(ctx: PipelineCtx): z.infer<typeof extractOutputSchema> | undefined {
  const extract = ctx.steps[EXTRACT_STEP_KEY];
  if (!isExtracted(extract)) return undefined;
  const parsed = extractOutputSchema.safeParse(extract.data);
  return parsed.success ? parsed.data : undefined;
}

const extractSpec: AiExtractStepArgs = {
  stepKey: EXTRACT_STEP_KEY,
  promptKey: "document-pipeline:extract:prompt",
  promptFallback: "Extract the document's title and up to five key points as a short bullet list.",
  paramsSchema,
  defaults: { enabled: true, params: {} },
  outputSchema: extractOutputSchema,
  instructions: () => "Extract the title and up to five key points from the attached document.",
  document: (ctx) => payloadOf(ctx).document,
};

const classifySpec: AiClassifyStepArgs = {
  stepKey: CLASSIFY_STEP_KEY,
  promptKey: "document-pipeline:classify:prompt",
  promptFallback: `Classify the document and propose the right review action.`,
  paramsSchema,
  defaults: { enabled: true, params: {} },
  actions: REVIEW_ACTIONS,
  input: (ctx) => {
    const fields = extractedFields(ctx);
    if (!fields) return fenceDocumentData("(no extraction available)");
    return fenceDocumentData(
      `Title: ${fields.title}\nKey points: ${fields.keyPoints.join("; ") || "(none extracted)"}`,
    );
  },
};

const generateSpec: AiGenerateStepArgs = {
  stepKey: GENERATE_STEP_KEY,
  promptKey: "document-pipeline:generate:prompt",
  promptFallback: `Write a two-sentence reviewer summary of the document.`,
  paramsSchema,
  defaults: { enabled: true, params: {} },
  input: (ctx) => {
    const title = extractedFields(ctx)?.title ?? "the document";
    const classify = ctx.steps[CLASSIFY_STEP_KEY];
    const category = isClassified(classify) ? classify.category : "uncategorized";
    return `Write a two-sentence reviewer summary of this document:\n${fenceDocumentData(`Title: ${title}\nCategory: ${category}`)}`;
  },
};

export const documentPipelineStepSpecs: readonly AiStepSpec[] = [
  extractSpec,
  classifySpec,
  generateSpec,
];

export type ProcessDocumentPayload = {
  readonly label: string;
  readonly document: AiExtractDocument;
};

// classify/generate are absent when extract did not succeed: the branch
// below skips them so a failed extract costs no further provider calls.
type DocumentPipelineStepResults = {
  readonly extract: AiExtractResult;
  readonly classify?: AiClassifyResult;
  readonly generate?: AiGenerateResult;
};

export const documentPipelineWorkflow = defineWorkflow<
  ProcessDocumentPayload,
  DocumentPipelineStepResults
>({
  name: "document-pipeline-demo",
  trigger: { kind: "event", eventType: "document-pipeline-sample:triggered" },
  steps: stepsPipeline<ProcessDocumentPayload, DocumentPipelineStepResults>(({ r }) => [
    aiExtractStep(extractSpec),
    r.step.branch({
      if: (ctx) => isExtracted(ctx.steps[EXTRACT_STEP_KEY]),
      onTrue: [aiClassifyStep(classifySpec), aiGenerateStep(generateSpec)],
    }),
    r.step.return((ctx) => ({
      isSuccess: true,
      data: {
        extract: ctx.steps[EXTRACT_STEP_KEY],
        classify: ctx.steps[CLASSIFY_STEP_KEY],
        generate: ctx.steps[GENERATE_STEP_KEY],
      },
    })),
  ]),
});

const processDocumentSchema = z.object({
  label: z.string().min(1),
  document: z.object({
    // ~6 MB PDF as base64 (~8e6 chars). Unbounded data would spike memory
    // and LLM cost on every TenantAdmin call through this sample.
    data: z.string().min(1).max(8_000_000),
    mimeType: z.literal("application/pdf"),
  }),
});

const adminOnly = { access: { roles: ["TenantAdmin", "SystemAdmin"] } } as const;

export const documentPipelineFeature = defineFeature("document-pipeline", (r) => {
  r.requires({ features: ["ai-pipeline"] });
  r.entity({ name: "pipeline-run", ...pipelineRunEntity });

  const { executor } = createEntityExecutor("pipeline-run", pipelineRunEntity);

  r.queryHandler(
    defineEntityListHandler("pipeline-run", pipelineRunEntity, {
      ...adminOnly,
      description: "List past document-pipeline runs (paged), newest first.",
    }),
  );
  r.queryHandler(
    defineEntityDetailHandler("pipeline-run", pipelineRunEntity, {
      ...adminOnly,
      description: "Read one document-pipeline run by id, including its summary and key points.",
    }),
  );

  r.writeHandler({
    name: "process-document",
    schema: processDocumentSchema,
    handler: async (event, ctx) => {
      const runResult = await runPipeline(documentPipelineWorkflow.pipelineDef, event, ctx, {
        runId: randomUUID(),
        workflowName: documentPipelineWorkflow.name,
        // Overwritten per step by runStepList before each step.run() call.
        stepIndex: 0,
      });
      if (!runResult.isSuccess) {
        // Unreachable: step failures surface as result types, not WriteFailures.
        return runResult;
      }
      const { extract, classify, generate } = runResult.data;

      if (!isExtracted(extract)) {
        return failUnprocessable("document extraction failed", {
          step: "extract",
          result: extract,
        });
      }
      if (!isClassified(classify)) {
        return failUnprocessable("document classification failed", {
          step: "classify",
          result: classify,
        });
      }
      if (!isGenerated(generate)) {
        return failUnprocessable("summary generation failed", {
          step: "generate",
          result: generate,
        });
      }

      const extracted = extractOutputSchema.safeParse(extract.data);
      if (!extracted.success) {
        return failUnprocessable("document extraction failed", {
          step: "extract",
          result: extract,
        });
      }

      const created = await executor.create(
        {
          label: event.payload.label,
          title: extracted.data.title,
          keyPoints: extracted.data.keyPoints,
          category: classify.category,
          confidence: classify.confidence,
          proposedActions: classify.proposedActions,
          summary: generate.text,
        },
        event.user,
        ctx.db,
      );
      if (!created.isSuccess) return created;

      return withResponseData(created, { id: created.data.id, extract, classify, generate });
    },
    ...adminOnly,
    description:
      "Run the extract -> classify -> generate document pipeline over an uploaded document via the LLM provider and store the result; up to three provider calls per run (a step can be tenant-disabled), nothing external is triggered.",
    agent: { risk: "mid" },
  });
});

Enterprise recipe source is included in the generated docs mirror; the commercial workspace is not publicly browsable.