genkitx-aws-bedrock
Version:
Genkit AI framework plugin for AWS Bedrock APIs.
434 lines (433 loc) • 13.7 kB
TypeScript
/**
* Copyright 2026 Xavier Portilla Edo
* Copyright 2026 Google LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
import type { APIGatewayProxyEvent, APIGatewayProxyResult, APIGatewayProxyEventV2, APIGatewayProxyResultV2, Context as LambdaContext, StreamifyHandler } from "aws-lambda";
import { type Flow, type z, type ActionContext } from "genkit";
import { type ContextProvider, type RequestData } from "genkit/context";
export type { ContextProvider, RequestData, ActionContext };
/**
* Type helpers to extract input/output types from Flow
*/
type FlowInput<F extends Flow> = F extends Flow<infer I, z.ZodTypeAny, z.ZodTypeAny> ? z.infer<I> : never;
type FlowOutput<F extends Flow> = F extends Flow<z.ZodTypeAny, infer O, z.ZodTypeAny> ? z.infer<O> : never;
type FlowStream<F extends Flow> = F extends Flow<z.ZodTypeAny, z.ZodTypeAny, infer S> ? z.infer<S> : never;
/**
* CORS configuration options
*/
export interface CorsOptions {
/**
* Allowed origins for CORS requests.
* Can be a string, array of strings, or '*' for all origins.
* @default '*'
*/
origin?: string | string[];
/**
* Allowed HTTP methods.
* @default ['POST', 'OPTIONS']
*/
methods?: string[];
/**
* Allowed headers in requests.
* @default ['Content-Type', 'Authorization']
*/
allowedHeaders?: string[];
/**
* Headers exposed to the client.
*/
exposedHeaders?: string[];
/**
* Whether to allow credentials.
* @default false
*/
credentials?: boolean;
/**
* Max age for preflight cache (in seconds).
* @default 86400 (24 hours)
*/
maxAge?: number;
}
/**
* Extended action context that includes Lambda-specific information
*/
export interface LambdaActionContext extends ActionContext {
/** Lambda-specific context data */
lambda?: {
event: {
requestContext: Record<string, unknown>;
headers: Record<string, string | undefined>;
queryStringParameters: Record<string, string | undefined> | null;
pathParameters: Record<string, string | undefined> | null;
};
context: {
functionName: string;
functionVersion: string;
invokedFunctionArn: string;
memoryLimitInMB: string;
awsRequestId: string;
};
};
}
/**
* Options for configuring the Lambda handler
*/
export interface LambdaOptions<C extends ActionContext = ActionContext, T = unknown> {
/**
* CORS configuration. Set to false to disable CORS headers.
* @default { origin: '*', methods: ['POST', 'OPTIONS'] }
*/
cors?: CorsOptions | boolean;
/**
* Context provider that parses request data and returns context for the flow.
* This follows the same pattern as express, next.js, and other Genkit integrations.
*
* The context provider receives a RequestData object containing:
* - method: HTTP method ('GET', 'POST', etc.)
* - headers: Lowercase headers from the request
* - input: Parsed request body
*
* Return an ActionContext object that will be available via getContext() in the flow.
* Throw UserFacingError for authentication/authorization failures.
*
* @example
* ```typescript
* import { UserFacingError } from 'genkit';
*
* const authProvider: ContextProvider = async (req) => {
* const token = req.headers['authorization'];
* if (!token) {
* throw new UserFacingError('UNAUTHENTICATED', 'Missing auth token');
* }
* const user = await verifyToken(token);
* return { auth: { user } };
* };
*
* export const handler = onCallGenkit(
* { contextProvider: authProvider },
* myFlow
* );
* ```
*/
contextProvider?: ContextProvider<C, T>;
/**
* Custom error handler for transforming errors before response.
*/
onError?: (error: Error) => {
statusCode: number;
message: string;
} | Promise<{
statusCode: number;
message: string;
}>;
/**
* Whether to log incoming events (for debugging).
* @default false
*/
debug?: boolean;
/**
* Whether to return a streaming Lambda handler instead of a standard one.
* When true, `onCallGenkit` returns a `StreamifyHandler` that uses
* `awslambda.streamifyResponse` for real incremental streaming via
* Lambda Function URLs with `InvokeMode: RESPONSE_STREAM`.
*
* The streaming handler is compatible with `streamFlow` from `genkit/beta/client`.
* For clients sending `Accept: text/event-stream`, it writes SSE chunks
* incrementally. Otherwise it falls back to a buffered JSON response.
*
* @default false
*
* @example
* ```typescript
* export const handler = onCallGenkit(
* { streaming: true },
* myStreamingFlow
* );
* ```
*/
streaming?: boolean;
}
/**
* Response wrapper for successful flow execution (callable protocol).
* Follows the same format as express and other Genkit integrations.
*/
export interface FlowResponse<T> {
result: T;
}
/**
* Response wrapper for failed flow execution (callable protocol).
* Shape matches genkit's getCallableJSON output.
*/
export interface FlowErrorResponse {
error: {
status: string;
message: string;
details?: unknown;
};
}
/**
* Union type for flow responses
*/
export type LambdaFlowResponse<T> = FlowResponse<T> | FlowErrorResponse;
/**
* Lambda handler function type for API Gateway v1
*/
export type LambdaHandler = (event: APIGatewayProxyEvent, context: LambdaContext) => Promise<APIGatewayProxyResult>;
/**
* Lambda handler function type for API Gateway v2
*/
export type LambdaHandlerV2 = (event: APIGatewayProxyEventV2, context: LambdaContext) => Promise<APIGatewayProxyResultV2>;
/**
* Run options for flow execution
*/
export interface FlowRunOptions {
context?: Record<string, unknown>;
}
/**
* Callable function type that includes the raw handler and metadata
*/
export interface CallableLambdaFunction<F extends Flow> extends LambdaHandler {
/**
* The underlying Genkit flow
*/
flow: F;
/**
* Execute the flow directly (for testing)
*/
run: (input: FlowInput<F>, options?: FlowRunOptions) => Promise<FlowOutput<F>>;
/**
* Stream the flow directly (for testing)
*/
stream: (input: FlowInput<F>, options?: FlowRunOptions) => {
stream: AsyncIterable<FlowStream<F>>;
output: Promise<FlowOutput<F>>;
};
/**
* Flow name
*/
flowName: string;
}
/**
* Creates a Lambda handler for a Genkit flow.
*
* This function wraps a Genkit flow to create an AWS Lambda handler that:
* - Handles CORS automatically
* - Supports ContextProvider for authentication/authorization
* - Provides proper error handling
* - Returns standardized response format
*
* @example Basic usage
* ```typescript
* import { genkit, z } from 'genkit';
* import { onCallGenkit } from 'genkitx-aws-bedrock';
* import { awsBedrock, amazonNovaProV1 } from 'genkitx-aws-bedrock';
*
* const ai = genkit({
* plugins: [awsBedrock()],
* model: amazonNovaProV1(),
* });
*
* const myFlow = ai.defineFlow(
* { name: 'myFlow', inputSchema: z.string(), outputSchema: z.string() },
* async (input) => {
* const { text } = await ai.generate({ prompt: input });
* return text;
* }
* );
*
* export const handler = onCallGenkit(myFlow);
* ```
*
* @example With ContextProvider for authentication
* ```typescript
* import { UserFacingError, getContext } from 'genkit';
* import type { ContextProvider } from 'genkit/context';
*
* interface AuthContext {
* auth: { user: { id: string; name: string } };
* }
*
* const authProvider: ContextProvider<AuthContext> = async (req) => {
* const token = req.headers['authorization'];
* if (!token) {
* throw new UserFacingError('UNAUTHENTICATED', 'Missing auth token');
* }
* const user = await verifyToken(token);
* return { auth: { user } };
* };
*
* // In your flow, access context via getContext()
* const myFlow = ai.defineFlow(
* { name: 'myFlow', inputSchema: z.string(), outputSchema: z.string() },
* async (input, { context }) => {
* const { auth } = context;
* console.log('User:', auth.user.name);
* // ...
* }
* );
*
* export const handler = onCallGenkit(
* { contextProvider: authProvider },
* myFlow
* );
* ```
*
* @param flow - The Genkit flow to wrap
* @returns A Lambda handler function
*/
export declare function onCallGenkit<F extends Flow>(flow: F): CallableLambdaFunction<F>;
/**
* Creates a Lambda handler for a Genkit flow with options.
*
* @param opts - Configuration options for the Lambda handler
* @param flow - The Genkit flow to wrap
* @returns A Lambda handler function, or a StreamifyHandler when `streaming: true`
*/
export declare function onCallGenkit<C extends ActionContext, F extends Flow>(opts: LambdaOptions<C, FlowInput<F>> & {
streaming: true;
}, flow: F): StreamifyHandler<APIGatewayProxyEventV2, void>;
export declare function onCallGenkit<C extends ActionContext, F extends Flow>(opts: LambdaOptions<C, FlowInput<F>>, flow: F): CallableLambdaFunction<F>;
/**
* Context with API key authentication
*/
export interface ApiKeyContext extends ActionContext {
auth: {
apiKey: string;
};
}
/**
* Context with bearer token authentication
*/
export interface BearerTokenContext extends ActionContext {
auth: {
token: string;
};
}
/**
* Creates a context provider that requires an API key in a specific header.
*
* @example
* ```typescript
* // Require API key to match a specific value
* export const handler = onCallGenkit(
* { contextProvider: requireApiKey('X-API-Key', process.env.API_KEY!) },
* myFlow
* );
*
* // Or with a custom validation function
* export const handler = onCallGenkit(
* {
* contextProvider: requireApiKey('X-API-Key', async (key) => {
* const valid = await validateApiKey(key);
* if (!valid) {
* throw new UserFacingError('PERMISSION_DENIED', 'Invalid API key');
* }
* })
* },
* myFlow
* );
* ```
*/
export declare function requireApiKey(headerName: string, expectedValueOrValidator: string | ((apiKey: string) => void | Promise<void>)): ContextProvider<ApiKeyContext>;
/**
* Creates a context provider that requires Bearer token authentication.
*
* @example
* ```typescript
* // With custom token validation
* export const handler = onCallGenkit(
* {
* contextProvider: requireBearerToken(async (token) => {
* const user = await verifyJWT(token);
* return { auth: { user } };
* })
* },
* myFlow
* );
* ```
*/
export declare function requireBearerToken<C extends ActionContext = BearerTokenContext>(validateToken: (token: string) => C | Promise<C>): ContextProvider<C>;
/**
* Creates a context provider that requires a specific header to be present.
*
* @example
* ```typescript
* // Require header to exist
* export const handler = onCallGenkit(
* { contextProvider: requireHeader('X-Request-ID') },
* myFlow
* );
*
* // Require header to have specific value
* export const handler = onCallGenkit(
* { contextProvider: requireHeader('X-API-Version', '2.0') },
* myFlow
* );
* ```
*/
export declare function requireHeader(headerName: string, expectedValue?: string): ContextProvider<ActionContext>;
/**
* Creates a context provider that always allows requests (no authentication).
* Useful for public endpoints.
*
* @example
* ```typescript
* export const handler = onCallGenkit(
* { contextProvider: allowAll() },
* myPublicFlow
* );
* ```
*/
export declare function allowAll(): ContextProvider<ActionContext>;
/**
* Combines multiple context providers. All providers must succeed.
* The returned context is a merge of all provider contexts.
*
* @example
* ```typescript
* export const handler = onCallGenkit(
* {
* contextProvider: allOf(
* requireHeader('X-Request-ID'),
* requireApiKey('X-API-Key', process.env.API_KEY!)
* )
* },
* myFlow
* );
* ```
*/
export declare function allOf<C extends ActionContext = ActionContext>(...providers: ContextProvider<ActionContext>[]): ContextProvider<C>;
/**
* Tries context providers in order, returning the first one that succeeds.
* If all providers fail, throws the error from the last provider.
*
* @example
* ```typescript
* // Accept either API key or Bearer token
* export const handler = onCallGenkit(
* {
* contextProvider: anyOf(
* requireApiKey('X-API-Key', process.env.API_KEY!),
* requireBearerToken(async (token) => {
* const user = await verifyJWT(token);
* return { auth: { user } };
* })
* )
* },
* myFlow
* );
* ```
*/
export declare function anyOf<C extends ActionContext = ActionContext>(...providers: ContextProvider<ActionContext>[]): ContextProvider<C>;
export default onCallGenkit;