feat(templates): add support for defaults

i-am-bee · Sep 11, 2024 · 8be1adb · 8be1adb
1 parent a10d432
commit 8be1adb
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 16 deletions.
diff --git a/examples/template.ts b/examples/template.ts
@@ -63,16 +63,13 @@ const logger = new Logger({ name: "template" });
   const modified = original.fork((oldConfig) => ({
     ...oldConfig,
     template: `${oldConfig.template} Your answers must be concise.`,
-    schema: z.object({
-      name: z.string().default("123"),
-      age: z.number(),
-      objective: z.string(),
-    }),
+    defaults: {
+      name: "Alex",
+    },
   }));
 
   const output = modified.render({
     name: undefined,
-    age: 12,
     objective: "fulfill the user needs",
   });
   logger.info(output);

diff --git a/src/template.ts b/src/template.ts
@@ -23,27 +23,30 @@ import { z, ZodType } from "zod";
 import { createSchemaValidator, toJsonSchema } from "@/internals/helpers/schema.js";
 import type { SchemaObject, ValidateFunction } from "ajv";
 import { shallowCopy } from "@/serializer/utils.js";
+import { pickBy } from "remeda";
+import { getProp } from "@/internals/helpers/object.js";
 
-export type PromptTemplateRenderFn<K extends ZodType> = (
-  this: K extends ZodType<infer A> ? A : never,
-) => any;
+export type InferValue<T> = T extends ZodType<infer A> ? A : never;
+export type PromptTemplateRenderFn<K extends ZodType> = (this: InferValue<K>) => any;
 
 export type PromptTemplateRenderInput<T extends ZodType, T2 extends z.input<T> = z.input<T>> = {
-  [K in keyof T2]: T2[K] | PromptTemplateRenderFn<T>;
+  [K in keyof T2]: T2[K] | PromptTemplateRenderFn<T> | undefined;
 };
 
 export interface PromptTemplateInput<T extends ZodType> {
   template: string;
   customTags?: [string, string];
   escape?: boolean;
   schema: SchemaObject;
+  defaults?: Partial<InferValue<T>>;
   functions?: Record<string, PromptTemplateRenderFn<T>>;
 }
 
 type PromptTemplateConstructor<T extends ZodType, N> = N extends ZodType
-  ? Omit<PromptTemplateInput<N>, "schema" | "functions"> & {
+  ? Omit<PromptTemplateInput<N>, "schema" | "functions" | "defaults"> & {
       schema: N;
       functions?: Record<string, PromptTemplateRenderFn<N | T>>;
+      defaults?: Partial<InferValue<N | T>>;
     }
   : Omit<PromptTemplateInput<T>, "schema"> & { schema: T | SchemaObject };
 
@@ -77,6 +80,7 @@ export class PromptTemplate<T extends ZodType> extends Serializable {
     super();
     this.config = {
       ...config,
+      defaults: (config.defaults ?? {}) as Partial<InferValue<T>>,
       schema: toJsonSchema(config.schema),
       escape: Boolean(config.escape),
       customTags: config.customTags ?? ["{{", "}}"],
@@ -92,8 +96,7 @@ export class PromptTemplate<T extends ZodType> extends Serializable {
   }
 
   protected validateInput(input: unknown): asserts input is T {
-    const schema = toJsonSchema(this.config.schema);
-    const validator = createSchemaValidator(schema, {
+    const validator = createSchemaValidator(this.config.schema, {
       coerceTypes: false,
     }) as ValidateFunction<T>;
 
@@ -119,11 +122,17 @@ export class PromptTemplate<T extends ZodType> extends Serializable {
     return new PromptTemplate(newConfig);
   }
 
-  render(inputs: PromptTemplateRenderInput<T>): string {
-    this.validateInput(inputs);
+  render(input: PromptTemplateRenderInput<T>): string {
+    const updatedInput: typeof input = { ...input };
+    Object.assign(
+      updatedInput,
+      pickBy(this.config.defaults, (_, k) => getProp(updatedInput, [k]) === undefined),
+    );
+
+    this.validateInput(updatedInput);
     const view: Record<string, any> = {
       ...this.config.functions,
-      ...inputs,
+      ...updatedInput,
     };
 
     const output = Mustache.render(