Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

docs: jsdocs and typedoc cleanup #1526

Merged
merged 17 commits into from
Dec 17, 2024
35 changes: 35 additions & 0 deletions js/ai/src/embedder.ts
Original file line number Diff line number Diff line change
Expand Up @@ -18,32 +18,57 @@ import { Action, defineAction, z } from '@genkit-ai/core';
import { Registry } from '@genkit-ai/core/registry';
import { Document, DocumentData, DocumentDataSchema } from './document.js';

/**
* A batch (array) of embeddings.
*/
export type EmbeddingBatch = { embedding: number[] }[];

/**
* Zod schema of an embedding vector.
*/
pavelgj marked this conversation as resolved.
Show resolved Hide resolved
export const EmbeddingSchema = z.array(z.number());

/**
* Embedding vector.
*/
export type Embedding = z.infer<typeof EmbeddingSchema>;

/**
* A function used for embedder definition, encapsulates embedder implementation.
*/
export type EmbedderFn<EmbedderOptions extends z.ZodTypeAny> = (
input: Document[],
embedderOpts?: z.infer<EmbedderOptions>
) => Promise<EmbedResponse>;

/**
* Zod schema of an embed request.
*/
const EmbedRequestSchema = z.object({
input: z.array(DocumentDataSchema),
options: z.any().optional(),
});

/**
* Zod schema of an embed response.
*/
const EmbedResponseSchema = z.object({
embeddings: z.array(z.object({ embedding: EmbeddingSchema })),
// TODO: stats, etc.
});
type EmbedResponse = z.infer<typeof EmbedResponseSchema>;

/**
* Embedder action -- a subtype of {@link Action} with input/output types for embedders.
*/
export type EmbedderAction<CustomOptions extends z.ZodTypeAny = z.ZodTypeAny> =
Action<typeof EmbedRequestSchema, typeof EmbedResponseSchema> & {
__configSchema?: CustomOptions;
};

/**
* Options of an `embed` function.
*/
export interface EmbedderParams<
CustomOptions extends z.ZodTypeAny = z.ZodTypeAny,
> {
Expand Down Expand Up @@ -105,6 +130,9 @@ export function defineEmbedder<
return ewm;
}

/**
* A union type representing all the types that can refer to an embedder.
*/
export type EmbedderArgument<
CustomOptions extends z.ZodTypeAny = z.ZodTypeAny,
> = string | EmbedderAction<CustomOptions> | EmbedderReference<CustomOptions>;
Expand Down Expand Up @@ -215,6 +243,9 @@ export async function embedMany<
return response.embeddings;
}

/**
* Zod schema of embedder info object.
*/
export const EmbedderInfoSchema = z.object({
/** Friendly label for this model (e.g. "Google AI - Gemini Pro") */
label: z.string().optional(),
Expand All @@ -232,6 +263,10 @@ export const EmbedderInfoSchema = z.object({
});
export type EmbedderInfo = z.infer<typeof EmbedderInfoSchema>;

/**
* A reference object that can used to resolve an embedder instance. Include additional type information
* about the specific embedder, e.g. custom config options schema.
*/
export interface EmbedderReference<
CustomOptions extends z.ZodTypeAny = z.ZodTypeAny,
> {
Expand Down
3 changes: 3 additions & 0 deletions js/ai/src/extract.ts
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,9 @@
import JSON5 from 'json5';
import { Allow, parse } from 'partial-json';

/**
* Parses partially complete JSON string.
*/
export function parsePartialJson<T = unknown>(jsonString: string): T {
return JSON5.parse<T>(JSON.stringify(parse(jsonString, Allow.ALL)));
}
Expand Down
121 changes: 120 additions & 1 deletion js/ai/src/model.ts
Original file line number Diff line number Diff line change
Expand Up @@ -44,12 +44,22 @@ const EmptyPartSchema = z.object({
custom: z.record(z.unknown()).optional(),
});

/**
* Zod schema for a text part.
*/
export const TextPartSchema = EmptyPartSchema.extend({
/** The text of the message. */
text: z.string(),
});

/**
* Text part.
*/
export type TextPart = z.infer<typeof TextPartSchema>;

/**
* Zod schema of a media part.
*/
export const MediaPartSchema = EmptyPartSchema.extend({
media: z.object({
/** The media content type. Inferred from data uri if not provided. */
Expand All @@ -58,8 +68,15 @@ export const MediaPartSchema = EmptyPartSchema.extend({
url: z.string(),
}),
});

/**
* Media part.
*/
export type MediaPart = z.infer<typeof MediaPartSchema>;

/**
* Zod schema of a tool request part.
*/
export const ToolRequestPartSchema = EmptyPartSchema.extend({
/** A request for a tool to be executed, usually provided by a model. */
toolRequest: z.object({
Expand All @@ -71,8 +88,15 @@ export const ToolRequestPartSchema = EmptyPartSchema.extend({
input: z.unknown().optional(),
}),
});

/**
* Tool part.
*/
export type ToolRequestPart = z.infer<typeof ToolRequestPartSchema>;

/**
* Zod schema of a tool response part.
*/
export const ToolResponsePartSchema = EmptyPartSchema.extend({
/** A provided response to a tool call. */
toolResponse: z.object({
Expand All @@ -84,19 +108,39 @@ export const ToolResponsePartSchema = EmptyPartSchema.extend({
output: z.unknown().optional(),
}),
});

/**
* Tool response part.
*/
export type ToolResponsePart = z.infer<typeof ToolResponsePartSchema>;

/**
* Zod schema of a data part.
*/
export const DataPartSchema = EmptyPartSchema.extend({
data: z.unknown(),
});

/**
* Data part.
*/
export type DataPart = z.infer<typeof DataPartSchema>;

/**
* Zod schema of a custom part.
*/
export const CustomPartSchema = EmptyPartSchema.extend({
custom: z.record(z.any()),
pavelgj marked this conversation as resolved.
Show resolved Hide resolved
});

/**
* Custom part.
*/
export type CustomPart = z.infer<typeof CustomPartSchema>;

/**
* Zod schema of message part.
*/
export const PartSchema = z.union([
TextPartSchema,
MediaPartSchema,
Expand All @@ -106,18 +150,38 @@ export const PartSchema = z.union([
CustomPartSchema,
]);

/**
* Message part.
*/
export type Part = z.infer<typeof PartSchema>;

/**
* Zod schema of a message role.
*/
export const RoleSchema = z.enum(['system', 'user', 'model', 'tool']);

/**
* Message role.
*/
export type Role = z.infer<typeof RoleSchema>;

/**
* Zod schema of a message.
*/
export const MessageSchema = z.object({
role: RoleSchema,
content: z.array(PartSchema),
metadata: z.record(z.unknown()).optional(),
});

/**
* Model message data.
*/
export type MessageData = z.infer<typeof MessageSchema>;

/**
* Zod schema of model info metadata.
*/
export const ModelInfoSchema = z.object({
/** Acceptable names for this model (e.g. different versions). */
versions: z.array(z.string()).optional(),
Expand Down Expand Up @@ -153,8 +217,15 @@ export const ModelInfoSchema = z.object({
.enum(['featured', 'stable', 'unstable', 'legacy', 'deprecated'])
.optional(),
});

/**
* Model info metadata.
*/
export type ModelInfo = z.infer<typeof ModelInfoSchema>;

/**
* Zod schema of a tool definition.
*/
export const ToolDefinitionSchema = z.object({
name: z.string(),
description: z.string(),
Expand All @@ -171,8 +242,15 @@ export const ToolDefinitionSchema = z.object({
.describe('additional metadata for this tool definition')
.optional(),
});

/**
* Tool definition.
*/
export type ToolDefinition = z.infer<typeof ToolDefinitionSchema>;

/**
* Zod schema of a common config object.
*/
export const GenerationCommonConfigSchema = z.object({
/** A specific version of a model family, e.g. `gemini-1.0-pro-001` for the `gemini-1.0-pro` family. */
version: z.string().optional(),
Expand All @@ -182,15 +260,26 @@ export const GenerationCommonConfigSchema = z.object({
topP: z.number().optional(),
stopSequences: z.array(z.string()).optional(),
});

/**
* Common config object.
*/
export type GenerationCommonConfig = typeof GenerationCommonConfigSchema;

/**
* Zod schema of output config.
*/
const OutputConfigSchema = z.object({
format: z.string().optional(),
schema: z.record(z.any()).optional(),
constrained: z.boolean().optional(),
instructions: z.string().optional(),
contentType: z.string().optional(),
});

/**
* Output config.
*/
export type OutputConfig = z.infer<typeof OutputConfigSchema>;

/** ModelRequestSchema represents the parameters that are passed to a model when generating content. */
Expand All @@ -207,19 +296,31 @@ export interface ModelRequest<
> extends z.infer<typeof ModelRequestSchema> {
config?: z.infer<CustomOptionsSchema>;
}

/**
* Zod schema of a generate request.
*/
export const GenerateRequestSchema = ModelRequestSchema.extend({
/** @deprecated All responses now return a single candidate. This will always be `undefined`. */
candidates: z.number().optional(),
});

/**
* Generate request data.
*/
export type GenerateRequestData = z.infer<typeof GenerateRequestSchema>;

/**
* Generate request.
*/
export interface GenerateRequest<
CustomOptionsSchema extends z.ZodTypeAny = z.ZodTypeAny,
> extends z.infer<typeof GenerateRequestSchema> {
config?: z.infer<CustomOptionsSchema>;
}

/**
* Zod schema of usage info from a generate request.
*/
export const GenerationUsageSchema = z.object({
inputTokens: z.number().optional(),
outputTokens: z.number().optional(),
Expand All @@ -234,6 +335,10 @@ export const GenerationUsageSchema = z.object({
outputAudioFiles: z.number().optional(),
custom: z.record(z.number()).optional(),
});

/**
* Usage info from a generate request.
*/
export type GenerationUsage = z.infer<typeof GenerationUsageSchema>;

/** @deprecated All responses now return a single candidate. Only the first candidate will be used if supplied. */
Expand All @@ -257,6 +362,9 @@ export const CandidateErrorSchema = z.object({
/** @deprecated All responses now return a single candidate. Only the first candidate will be used if supplied. */
export type CandidateError = z.infer<typeof CandidateErrorSchema>;

/**
* Zod schema of a model response.
*/
export const ModelResponseSchema = z.object({
message: MessageSchema.optional(),
finishReason: z.enum(['stop', 'length', 'blocked', 'other', 'unknown']),
Expand All @@ -268,15 +376,26 @@ export const ModelResponseSchema = z.object({
raw: z.unknown(),
request: GenerateRequestSchema.optional(),
});

/**
* Model response data.
*/
export type ModelResponseData = z.infer<typeof ModelResponseSchema>;

/**
* Zod schema of generaete response.
*/
export const GenerateResponseSchema = ModelResponseSchema.extend({
/** @deprecated All responses now return a single candidate. Only the first candidate will be used if supplied. Return `message`, `finishReason`, and `finishMessage` instead. */
candidates: z.array(CandidateSchema).optional(),
finishReason: z
.enum(['stop', 'length', 'blocked', 'other', 'unknown'])
.optional(),
});

/**
* Generate response data.
*/
export type GenerateResponseData = z.infer<typeof GenerateResponseSchema>;

/** ModelResponseChunkSchema represents a chunk of content to stream to the client. */
Expand Down
Loading
Loading