ai-functions
Version:
Core AI primitives for building intelligent applications
233 lines • 8.76 kB
JavaScript
/**
* BatchProvider Port
*
* Defines the explicit port (interface) that every batch provider adapter must
* satisfy, plus helpers that concentrate logic genuinely shared across adapters
* (polling, concurrent processing, in-memory job tracking, JSON-Schema conversion).
*
* Each provider file (`anthropic.ts`, `openai.ts`, `google.ts`, `bedrock.ts`,
* `cloudflare.ts`, `memory.ts`) is now a small adapter that satisfies this port
* and concentrates on provider-specific HTTP/SDK calls and request/response shapes.
*
* Port + 4 adapters = real seam. The port lives here so provider files don't
* import shared logic from each other and can't accidentally diverge.
*
* @packageDocumentation
*/
export { registerBatchAdapter, registerFlexAdapter } from '../batch-queue.js';
// ============================================================================
// Polling helper (shared by all adapters that have async batch jobs)
// ============================================================================
/** Terminal states for a batch job — polling stops here. */
const TERMINAL_STATUSES = new Set([
'completed',
'failed',
'cancelled',
'expired',
]);
/** Statuses that should trigger result fetch. */
const RESULT_STATUSES = new Set(['completed', 'cancelled', 'failed']);
/**
* Default `waitForCompletion` implementation built on top of `getStatus` +
* `getResults`. Adapters with non-standard completion semantics can still
* override `waitForCompletion`, but most don't need to.
*
* @param adapter The batch adapter to poll
* @param batchId The batch id to poll
* @param options.pollInterval Poll interval in ms (default: 5000)
* @param options.fetchResultsOn Statuses for which results should be fetched.
* Defaults to `completed`, `cancelled`, `failed`.
* @param options.throwOn Statuses that should throw rather than fetch results.
* Useful for OpenAI which throws on `cancelled`/`expired`.
*/
export async function pollUntilComplete(adapter, batchId, options = {}) {
const pollInterval = options.pollInterval ?? 5000;
const fetchResultsOn = options.fetchResultsOn ?? RESULT_STATUSES;
const throwOn = options.throwOn;
while (true) {
const status = await adapter.getStatus(batchId);
if (throwOn?.has(status.status)) {
throw new Error(`Batch ${status.status}`);
}
if (fetchResultsOn.has(status.status) || TERMINAL_STATUSES.has(status.status)) {
return adapter.getResults(batchId);
}
await sleep(pollInterval);
}
}
// ============================================================================
// Concurrent processing helper (shared by flex adapters and "local" providers)
// ============================================================================
/**
* Run `processItem` over `items` with bounded concurrency, optionally
* sleeping between waves to respect provider rate limits. Per-item failures
* are caught and emitted as `{ status: 'failed', error }` results so a single
* bad item never poisons the whole batch.
*
* Used by flex adapters (OpenAI, Google, Bedrock) and by the "local" providers
* (Google, Bedrock, Cloudflare) that fake batch processing with concurrent
* direct API calls.
*/
export async function processConcurrently(items, processItem, options = {}) {
const concurrency = options.concurrency ?? 10;
const delay = options.delayBetweenWaves ?? 0;
const results = [];
for (let i = 0; i < items.length; i += concurrency) {
const wave = items.slice(i, i + concurrency);
const waveResults = await Promise.all(wave.map(async (item) => {
try {
return await processItem(item);
}
catch (error) {
return failedResult(item, error);
}
}));
results.push(...waveResults);
options.onWaveComplete?.(results);
if (delay > 0 && i + concurrency < items.length) {
await sleep(delay);
}
}
return results;
}
/** Build a `failed` BatchResult from an unknown thrown value. */
export function failedResult(item, error) {
return {
id: item.id,
customId: item.id,
status: 'failed',
error: error instanceof Error ? error.message : 'Unknown error',
};
}
/**
* Per-provider in-memory job registry. Encapsulates the
* `Map<jobId, state>` + counter + status/result lookup pattern that
* google/bedrock/cloudflare were each duplicating.
*/
export class LocalJobStore {
idPrefix;
jobs = new Map();
counter = 0;
constructor(idPrefix) {
this.idPrefix = idPrefix;
}
create(items, options) {
const id = `${this.idPrefix}_${++this.counter}_${Date.now()}`;
const state = {
items,
options,
results: [],
status: 'pending',
createdAt: new Date(),
};
this.jobs.set(id, state);
return { id, state };
}
get(id) {
const state = this.jobs.get(id);
if (!state) {
throw new Error(`Batch not found: ${id}`);
}
return state;
}
has(id) {
return this.jobs.has(id);
}
/** Build a `BatchJob` snapshot for a tracked job. */
snapshot(id, provider) {
const state = this.get(id);
const completedItems = state.results.filter((r) => r.status === 'completed').length;
const failedItems = state.results.filter((r) => r.status === 'failed').length;
return {
id,
provider,
status: state.status,
totalItems: state.items.length,
completedItems,
failedItems,
createdAt: state.createdAt,
...(state.completedAt && { completedAt: state.completedAt }),
};
}
/**
* Wait for a tracked job to reach a terminal status by polling its in-memory
* state. Adapters that drive the state machine in a background promise can
* call this from `waitForCompletion`.
*/
async waitForCompletion(id, pollInterval = 1000) {
const state = this.get(id);
while (state.status !== 'completed' &&
state.status !== 'failed' &&
state.status !== 'cancelled') {
await sleep(pollInterval);
}
return state.results;
}
/** For tests: drop everything. */
clear() {
this.jobs.clear();
this.counter = 0;
}
}
/**
* Minimal Zod -> JSON Schema converter.
*
* This is the same simplified converter that previously lived (duplicated)
* inside `anthropic.ts` and `openai.ts`. Extracted here so both adapters call
* the same implementation. For richer conversion use `zod-to-json-schema`.
*/
export function zodToJsonSchema(zodSchema) {
const schema = zodSchema;
if (!schema._def) {
return { type: 'object' };
}
switch (schema._def.typeName) {
case 'ZodString':
return { type: 'string' };
case 'ZodNumber':
return { type: 'number' };
case 'ZodBoolean':
return { type: 'boolean' };
case 'ZodArray':
return { type: 'array', items: zodToJsonSchema(schema._def.type) };
case 'ZodObject': {
const shape = schema._def.shape?.() ?? {};
const properties = {};
for (const [key, value] of Object.entries(shape)) {
properties[key] = zodToJsonSchema(value);
}
return { type: 'object', properties, required: Object.keys(properties) };
}
default:
return { type: 'object' };
}
}
// ============================================================================
// JSON parsing (used by every adapter that returns raw text)
// ============================================================================
/**
* Try to parse `text` as JSON when it looks like JSON or a schema is expected,
* otherwise return the text unchanged. Never throws.
*/
export function tryParseJson(text, expectJson = false) {
if (!text)
return text;
const trimmed = text.trim();
const looksLikeJson = trimmed.startsWith('{') || trimmed.startsWith('[');
if (!expectJson && !looksLikeJson)
return text;
try {
return JSON.parse(text);
}
catch {
return text;
}
}
// ============================================================================
// Misc
// ============================================================================
/** Promise-based setTimeout. */
export function sleep(ms) {
return new Promise((resolve) => setTimeout(resolve, ms));
}
//# sourceMappingURL=provider.js.map