@genkit-ai/ai
Version:
Genkit AI framework generative AI APIs.
355 lines (335 loc) • 10.4 kB
text/typescript
/**
* 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 {
getContext,
run,
z,
type ActionContext,
type Operation,
} from '@genkit-ai/core';
import { type Registry } from '@genkit-ai/core/registry';
import { cancelOperation } from './cancel-operation.js';
import { checkOperation } from './check-operation.js';
import { type DocumentData } from './document.js';
import {
embed,
embedMany,
type EmbedderArgument,
type EmbedderParams,
type Embedding,
type EmbeddingBatch,
} from './embedder.js';
import {
generate,
generateStream,
type GenerateOptions,
type GenerateResponse,
type GenerateStreamOptions,
type GenerateStreamResponse,
} from './generate.js';
import { GenerationCommonConfigSchema, type Part } from './model-types.js';
/**
* `GenkitAI` encapsulates Genkit's AI APIs.
*/
export class GenkitAI {
readonly registry: Registry;
constructor(registry: Registry) {
this.registry = registry;
}
/**
* Embeds the given `content` using the specified `embedder`.
*/
embed<CustomOptions extends z.ZodTypeAny>(
params: EmbedderParams<CustomOptions>
): Promise<Embedding[]> {
return embed(this.registry, params);
}
/**
* A veneer for interacting with embedder models in bulk.
*/
embedMany<ConfigSchema extends z.ZodTypeAny = z.ZodTypeAny>(params: {
embedder: EmbedderArgument<ConfigSchema>;
content: string[] | DocumentData[];
metadata?: Record<string, unknown>;
options?: z.infer<ConfigSchema>;
}): Promise<EmbeddingBatch> {
return embedMany(this.registry, params);
}
/**
* Make a generate call to the default model with a simple text prompt.
*
* ```ts
* const ai = genkit({
* plugins: [googleAI()],
* model: googleAI.model('gemini-flash-latest'), // default model
* })
*
* const { text } = await ai.generate('hi');
* ```
*/
generate<O extends z.ZodTypeAny = z.ZodTypeAny>(
strPrompt: string
): Promise<GenerateResponse<z.infer<O>>>;
/**
* Make a generate call to the default model with a multipart request.
*
* ```ts
* const ai = genkit({
* plugins: [googleAI()],
* model: googleAI.model('gemini-flash-latest'), // default model
* })
*
* const { text } = await ai.generate([
* { media: {url: 'http://....'} },
* { text: 'describe this image' }
* ]);
* ```
*/
generate<O extends z.ZodTypeAny = z.ZodTypeAny>(
parts: Part[]
): Promise<GenerateResponse<z.infer<O>>>;
/**
* Generate calls a generative model based on the provided prompt and configuration. If
* `messages` is provided, the generation will include a conversation history in its
* request. If `tools` are provided, the generate method will automatically resolve
* tool calls returned from the model unless `returnToolRequests` is set to `true`.
*
* See {@link GenerateOptions} for detailed information about available options.
*
* ```ts
* const ai = genkit({
* plugins: [googleAI()],
* })
*
* const { text } = await ai.generate({
* system: 'talk like a pirate',
* prompt: [
* { media: { url: 'http://....' } },
* { text: 'describe this image' }
* ],
* messages: conversationHistory,
* tools: [ userInfoLookup ],
* model: googleAI.model('gemini-flash-latest'),
* });
* ```
*/
generate<
O extends z.ZodTypeAny = z.ZodTypeAny,
CustomOptions extends z.ZodTypeAny = typeof GenerationCommonConfigSchema,
>(
opts:
| GenerateOptions<O, CustomOptions>
| PromiseLike<GenerateOptions<O, CustomOptions>>
): Promise<GenerateResponse<z.infer<O>>>;
async generate<
O extends z.ZodTypeAny = z.ZodTypeAny,
CustomOptions extends z.ZodTypeAny = typeof GenerationCommonConfigSchema,
>(
options:
| string
| Part[]
| GenerateOptions<O, CustomOptions>
| PromiseLike<GenerateOptions<O, CustomOptions>>
): Promise<GenerateResponse<z.infer<O>>> {
let resolvedOptions: GenerateOptions<O, CustomOptions>;
if (options instanceof Promise) {
resolvedOptions = await options;
} else if (typeof options === 'string' || Array.isArray(options)) {
resolvedOptions = {
prompt: options,
};
} else {
resolvedOptions = options as GenerateOptions<O, CustomOptions>;
}
return generate(this.registry, resolvedOptions);
}
/**
* Make a streaming generate call to the default model with a simple text prompt.
*
* ```ts
* const ai = genkit({
* plugins: [googleAI()],
* model: googleAI.model('gemini-flash-latest'), // default model
* })
*
* const { response, stream } = ai.generateStream('hi');
* for await (const chunk of stream) {
* console.log(chunk.text);
* }
* console.log((await response).text);
* ```
*/
generateStream<O extends z.ZodTypeAny = z.ZodTypeAny>(
strPrompt: string
): GenerateStreamResponse<z.infer<O>>;
/**
* Make a streaming generate call to the default model with a multipart request.
*
* ```ts
* const ai = genkit({
* plugins: [googleAI()],
* model: googleAI.model('gemini-flash-latest'), // default model
* })
*
* const { response, stream } = ai.generateStream([
* { media: {url: 'http://....'} },
* { text: 'describe this image' }
* ]);
* for await (const chunk of stream) {
* console.log(chunk.text);
* }
* console.log((await response).text);
* ```
*/
generateStream<O extends z.ZodTypeAny = z.ZodTypeAny>(
parts: Part[]
): GenerateStreamResponse<z.infer<O>>;
/**
* Streaming generate calls a generative model based on the provided prompt and configuration. If
* `messages` is provided, the generation will include a conversation history in its
* request. If `tools` are provided, the generate method will automatically resolve
* tool calls returned from the model unless `returnToolRequests` is set to `true`.
*
* See {@link GenerateOptions} for detailed information about available options.
*
* ```ts
* const ai = genkit({
* plugins: [googleAI()],
* })
*
* const { response, stream } = ai.generateStream({
* system: 'talk like a pirate',
* prompt: [
* { media: { url: 'http://....' } },
* { text: 'describe this image' }
* ],
* messages: conversationHistory,
* tools: [ userInfoLookup ],
* model: googleAI.model('gemini-flash-latest'),
* });
* for await (const chunk of stream) {
* console.log(chunk.text);
* }
* console.log((await response).text);
* ```
*/
generateStream<
O extends z.ZodTypeAny = z.ZodTypeAny,
CustomOptions extends z.ZodTypeAny = typeof GenerationCommonConfigSchema,
>(
opts:
| GenerateStreamOptions<O, CustomOptions>
| PromiseLike<GenerateStreamOptions<O, CustomOptions>>
): GenerateStreamResponse<z.infer<O>>;
generateStream<
O extends z.ZodTypeAny = z.ZodTypeAny,
CustomOptions extends z.ZodTypeAny = typeof GenerationCommonConfigSchema,
>(
options:
| string
| Part[]
| GenerateStreamOptions<O, CustomOptions>
| PromiseLike<GenerateStreamOptions<O, CustomOptions>>
): GenerateStreamResponse<z.infer<O>> {
if (typeof options === 'string' || Array.isArray(options)) {
options = { prompt: options };
}
return generateStream(this.registry, options);
}
/**
* Checks the status of of a given operation. Returns a new operation which will contain the updated status.
*
* ```ts
* let operation = await ai.generateOperation({
* model: googleAI.model('veo-2.0-generate-001'),
* prompt: 'A banana riding a bicycle.',
* });
*
* while (!operation.done) {
* operation = await ai.checkOperation(operation!);
* await new Promise((resolve) => setTimeout(resolve, 5000));
* }
* ```
*
* @param operation
* @returns
*/
checkOperation<T>(operation: Operation<T>): Promise<Operation<T>> {
return checkOperation(this.registry, operation);
}
/**
* Cancels a given operation. Returns a new operation which will contain the updated status.
*
* @param operation
* @returns
*/
cancelOperation<T>(operation: Operation<T>): Promise<Operation<T>> {
return cancelOperation(this.registry, operation);
}
/**
* A flow step that executes the provided function. Each run step is recorded separately in the trace.
*
* ```ts
* ai.defineFlow('hello', async() => {
* await ai.run('step1', async () => {
* // ... step 1
* });
* await ai.run('step2', async () => {
* // ... step 2
* });
* return result;
* })
* ```
*/
run<T>(name: string, func: () => Promise<T>): Promise<T>;
/**
* A flow step that executes the provided function. Each run step is recorded separately in the trace.
*
* ```ts
* ai.defineFlow('hello', async(name) => {
* const greeting = await ai.run('step1', name, async (input) => {
* return `Hello, ${input}!`;
* });
* const result = await ai.run('step2', greeting, async (input) => {
* // ... step 2
* });
* return result;
* })
*/
run<T>(
name: string,
input: any,
func: (input?: any) => Promise<T>
): Promise<T>;
run<T>(
name: string,
funcOrInput: () => Promise<T> | any,
maybeFunc?: (input?: any) => Promise<T>
): Promise<T> {
if (maybeFunc) {
return run(name, funcOrInput, maybeFunc, this.registry);
}
return run(name, funcOrInput, this.registry);
}
/**
* Returns current action (or flow) invocation context. Can be used to access things like auth
* data set by HTTP server frameworks. If invoked outside of an action (e.g. flow or tool) will
* return `undefined`.
*/
currentContext(): ActionContext | undefined {
return getContext();
}
}