octabits

PlatformoctaflowExtending

The AI add-on

Instrumented models, token and cost capture, quota enforcement and usage rollups.

octaflow/ai wires model instrumentation, token/cost capture, quota, and daily usage rollups into the engine’s lifecycle hooks — the core stays AI-free.

import { z } from 'zod';
import { generateText } from 'ai';
import { createWorkflowEngine, createStepHandlerRegistry, createInMemoryWorkflowStore } from 'octaflow';
import { defineAiStep, buildAiWorkflow, createAiWorkflowHooks } from 'octaflow/ai';

const summarize = defineAiStep({
  type: 'summarize',
  workflowInputSchema: z.object({ text: z.string() }),
  outputSchema: z.object({ summary: z.string() }),
  retry: { maxAttempts: 3 }, // provider 429s
  handler: async (ctx) => {
    const { text } = await generateText({ model: ctx.context.model, prompt: `Summarize: ${ctx.workflowInput.text}` });
    return { summary: text };
  },
});

const hooks = createAiWorkflowHooks({
  modelResolver: {
    resolveModel: () => myModel,          // your LanguageModelV4
    // resolveHost: (args) => container.scope(args),   // optional → ctx.context.host
    // resolveKeySource: () => 'platform',             // optional → stamped on workflow metadata
  },
  usageRecorder: { recordStepUsage: async () => {}, incrementWorkflowUsage: async () => {} },
  // quotaPolicy: { checkQuota: async () => ({ ok: true, value: undefined }) },
  // costEstimator: createCostEstimator({ … }),        // defaults to DEFAULT_MODEL_PRICING
});

const store = createInMemoryWorkflowStore();
const registry = createStepHandlerRegistry();
const engine = createWorkflowEngine({ store, dispatcher, registry, partitionKey: 'tenant-1', hooks });

ctx.context.model is an instrumented model — token usage is captured automatically and the onAfterStep hook turns it into cost via a pluggable pricing table (createCostEstimator). ctx.context.host is whatever modelResolver.resolveHost returns (a DI scope, domain services), or undefined if you don’t supply one. → examples/13-ai-workflow.ts

Quota enforcement & usage aggregation

Two store-agnostic engines complete the add-on. They own the window math, the enforcement order, and the rollup shapes; the raw counts and aggregate reads come through narrow store seams (AiQuotaStore / AiUsageStore) that you implement with your own SQL — flow never touches a database (and the ai layer never depends on pg). Scoping is generic (partitionKey) and keySource is a free-form string, so both engines stay tenancy-agnostic.

import { createAiQuotaService, createAiUsageAggregationService, DEFAULT_AI_QUOTA } from 'octaflow/ai';

// Quota: limits come from an injected getQuota(partitionKey) callback — return
// null to exempt a scope entirely (e.g. bring-your-own-key), or null on any
// single window to leave it unlimited.
const quota = createAiQuotaService({
  store,                                   // your AiQuotaStore (count reads)
  getQuota: (partitionKey) => DEFAULT_AI_QUOTA,
});
const gate = await quota.checkQuota('tenant-1');  // Result<void, AiQuotaExceededError>
// wire it into the hooks: quotaPolicy: { checkQuota: () => quota.checkQuota('tenant-1') }

// Usage aggregation: roll completed workflows (and embedding batches) into a
// daily table, then read summaries back.
const usage = createAiUsageAggregationService({ store });  // your AiUsageStore (UPSERT + aggregate reads)
await usage.recordWorkflowCompletion({ partitionKey: 'tenant-1', date, workflowType: 'gen', keySource: 'platform', usage: totals, estimatedCostMicros });
const byDay = await usage.getUsageSummary({ partitionKey: 'tenant-1', startDate, endDate });

checkQuota enforces three windows in order — concurrent → per-day → per-month (UTC) — and the day/month counts include currently-running (not-yet-aggregated) workflows. AiUsageStore extends AiQuotaStore, so one implementation backs both engines.

The model layer

Everything above assumes you already have a LanguageModelV4. The ai module also ships the layer that produces one in a bring-your-own-key deployment, moved from a production SaaS where every AI call went through it:

  • createRefusalAwareModel(model) — a Claude 5 model can decline a request rather than fail it: HTTP 200, empty content, stop_reason: "refusal". The AI SDK maps that to the unified finish reason content-filter, so nothing throws and a caller that does not check the finish reason sees an empty string or an unparseable object several frames from the cause. The wrapper turns it into a thrown AiRefusalError on doGenerate (with the provider’s category and explanation when sent), and defineAiStep treats it as terminal — retrying a refusal yields the same refusal. Streaming is deliberately untouched: a chat presents a decline as a turn. isAiRefusalError follows cause chains, because generateObject re-wraps.
  • Capability switches — supportsAdaptiveThinking, thinkingSettingsFor(providerId, modelId, { temperature, budgetTokens }), supportsToolSearch. Claude 4.6+ rejects the legacy budget_tokens shape and sampling params with a 400; older models require them. A BYOK deployment lets a customer name any model id, so the request shape has to branch on the resolved id, once, here.
  • The provider registry — AI_PROVIDERS, EMBEDDING_PROVIDERS, embeddingDimensionsFor: dependency-free metadata a browser picker can import. Each vendor carries its residency (eu or third-country, legal name, country), the fact a subprocessor annex and a residency badge both need. Import it from octaflow/ai/registry in a browser bundle: that subpath carries the registry alone, with no ai SDK import, and a lint keeps it that way. octaflow/ai re-exports the same names next to the model layer.
  • octaflow/ai/providers — createLanguageModelProvider({ provider, apiKey, defaultModel?, baseUrl? }) and createEmbeddingModelProvider(...): a key and a provider id become refusal-aware models, with an optional region-pinned endpoint. The vendor SDKs (@ai-sdk/anthropic, @ai-sdk/mistral, @ai-sdk/openai) are optional peers loaded on first use, so only the vendors you point at need to be installed.
  • Self-hosted and other OpenAI-compatible endpoints — createOpenAiCompatibleModelProvider({ baseUrl, residency, defaultModel, apiKey? }) and createOpenAiCompatibleEmbeddingProvider({ baseUrl, residency, model, dimensions, apiKey? }), in the same subpath, over the optional peer @ai-sdk/openai-compatible: vLLM, Ollama, llama.cpp’s server, TGI, or a hosted vendor with that API. Language models go through Chat Completions, which is what those servers speak, not the Responses API that @ai-sdk/openai defaults to. Nothing is assumed. The base URL has no default. The residency is declared by the deployment and checked by resolveEndpointResidency: { region: 'self-hosted' } adds no subprocessor, and a hosted vendor needs its legal name and country. An embedding model’s dimensions must be stated, because a guessed width fails at the first insert into a vector(N) column. A configuration that cannot work throws AiProviderConfigError, with a key, before any SDK loads. These factories sit beside AI_PROVIDERS rather than in it, so a picker over the fixed-residency vendors is unchanged. Give each self-hosted model id a ZERO_MODEL_PRICING row in your cost table, or quota prices it at the dearest vendor row.
  • A gateway with several models — createOpenAiCompatibleGatewayProvider({ baseUrl, residency, models, apiKey? }), in the same subpath: LiteLLM or any router that serves several model names behind one OpenAI-compatible endpoint and sends each somewhere else. The residency belongs to the alias, so models declares one per alias ({ 'approval-confidential': { residency: { region: 'self-hosted' } }, … }); the top-level residency is where the gateway itself runs and decides whether plain http is acceptable. languageModel(alias, { metadata? }) refuses an alias the declaration does not name with ai_model_undeclared, and metadata goes into the request body, where LiteLLM keeps it in its spend log — a workflow id there matches a gateway call to your own record. Ids only: metadata leaves as written. The declaration states what the gateway is configured to do; a fallback the gateway adds on its own breaks it silently, so review the two together.
import { createLanguageModelProvider } from 'octaflow/ai/providers';
import { thinkingSettingsFor } from 'octaflow/ai';

const models = await createLanguageModelProvider({ provider: 'mistral', apiKey: key.value });
const result = await generateText({
  model: models.languageModel(),
  ...thinkingSettingsFor(models.providerId, models.languageModel().modelId, { temperature: 0.3, budgetTokens: 4000 }),
  prompt,
});

// A model server in your own cluster: no key, no subprocessor, a declared width.
const local = await createOpenAiCompatibleModelProvider({
  baseUrl: 'http://vllm.ai.svc:8000/v1',
  residency: { region: 'self-hosted' },
  defaultModel: 'qwen3-32b',
});
const embeddings = await createOpenAiCompatibleEmbeddingProvider({
  baseUrl: 'http://tei.ai.svc/v1',
  residency: { region: 'self-hosted' },
  model: 'bge-m3',
  dimensions: 1024, // must match vector(1024)
});

Version 0.28.0, MIT-licensed.npm