UNPKG

ai-sdk-guardrails

Version:

Input and output guardrails middleware for Vercel AI SDK.

122 lines (118 loc) 5.39 kB
import { generateText, generateObject, streamText, streamObject, embed } from 'ai'; interface GuardrailResult { /** Whether the guardrail was triggered (blocked the request) */ tripwireTriggered: boolean; /** Human-readable message describing why the guardrail was triggered */ message?: string; /** Detailed metadata about the guardrail execution */ metadata?: Record<string, unknown>; /** Severity level of the guardrail violation */ severity?: 'low' | 'medium' | 'high' | 'critical'; /** Suggested action to resolve the issue */ suggestion?: string; /** Additional context information */ context?: { guardrailName: string; guardrailVersion?: string; executedAt: Date; executionTimeMs?: number; environment?: string; }; } type GuardrailsParams = { inputGuardrails?: InputGuardrail[]; outputGuardrails?: OutputGuardrail[]; throwOnBlocked?: boolean; enablePerformanceMonitoring?: boolean; }; type GenerateTextParams = Parameters<typeof generateText>[0]; type GenerateObjectParams = Parameters<typeof generateObject>[0]; type StreamTextParams = Parameters<typeof streamText>[0]; type StreamObjectParams = Parameters<typeof streamObject>[0]; type EmbedParams = Parameters<typeof embed>[0]; type GenerateTextResult = Awaited<ReturnType<typeof generateText>>; type GenerateObjectResult = Awaited<ReturnType<typeof generateObject>>; type StreamTextResult = ReturnType<typeof streamText>; type StreamObjectResult = ReturnType<typeof streamObject>; type EmbedResult = ReturnType<typeof embed>; type InputGuardrailContext = GenerateTextParams | GenerateObjectParams | StreamTextParams | StreamObjectParams | EmbedParams; type AIResult = GenerateTextResult | GenerateObjectResult | StreamTextResult | StreamObjectResult | EmbedResult; type OutputGuardrailContext = { input: InputGuardrailContext; result: AIResult; }; interface InputGuardrail { /** Unique identifier for the guardrail */ name: string; /** Human-readable description of what this guardrail does */ description?: string; /** Version of the guardrail for tracking changes */ version?: string; /** Tags for categorizing guardrails */ tags?: string[]; /** Whether this guardrail is enabled */ enabled?: boolean; /** Priority level for execution order */ priority?: 'low' | 'medium' | 'high' | 'critical'; /** Configuration options for the guardrail */ config?: Record<string, string | number | boolean>; /** The main execution function */ execute: (context: InputGuardrailContext) => Promise<GuardrailResult> | GuardrailResult; /** Optional setup function called once when guardrail is initialized */ setup?: () => Promise<void> | void; /** Optional cleanup function called when guardrail is destroyed */ cleanup?: () => Promise<void> | void; } interface OutputGuardrail { /** Unique identifier for the guardrail */ name: string; /** Human-readable description of what this guardrail does */ description?: string; /** Version of the guardrail for tracking changes */ version?: string; /** Tags for categorizing guardrails */ tags?: string[]; /** Whether this guardrail is enabled */ enabled?: boolean; /** Priority level for execution order */ priority?: 'low' | 'medium' | 'high' | 'critical'; /** Configuration options for the guardrail */ config?: Record<string, string | number | boolean>; /** The main execution function */ execute: (context: OutputGuardrailContext, accumulatedText?: string) => Promise<GuardrailResult> | GuardrailResult; /** Optional setup function called once when guardrail is initialized */ setup?: () => Promise<void> | void; /** Optional cleanup function called when guardrail is destroyed */ cleanup?: () => Promise<void> | void; } interface InputGuardrailsMiddlewareConfig { /** Input guardrails to execute before AI calls */ inputGuardrails: InputGuardrail[]; /** Execution options for guardrails */ executionOptions?: { parallel?: boolean; timeout?: number; continueOnFailure?: boolean; logLevel?: 'none' | 'error' | 'warn' | 'info' | 'debug'; }; /** Callback for when input is blocked */ onInputBlocked?: (results: GuardrailResult[], originalParams: InputGuardrailContext) => void; /** Whether to throw errors when guardrails are triggered */ throwOnBlocked?: boolean; } interface OutputGuardrailsMiddlewareConfig { /** Output guardrails to execute after AI calls */ outputGuardrails: OutputGuardrail[]; /** Execution options for guardrails */ executionOptions?: { parallel?: boolean; timeout?: number; continueOnFailure?: boolean; logLevel?: 'none' | 'error' | 'warn' | 'info' | 'debug'; }; /** Callback for when output is blocked */ onOutputBlocked?: (results: GuardrailResult[], originalParams: InputGuardrailContext, result: unknown) => void; /** Whether to throw errors when guardrails are triggered */ throwOnBlocked?: boolean; } export type { AIResult as A, GuardrailResult as G, InputGuardrail as I, OutputGuardrail as O, InputGuardrailContext as a, OutputGuardrailContext as b, InputGuardrailsMiddlewareConfig as c, OutputGuardrailsMiddlewareConfig as d, GuardrailsParams as e };