fix(py/dotpromptz): register initial helpers and partials correctly (#260)

Author: yesudeep

js/src/dotprompt.ts

@@ -88,22 +88,8 @@ export class Dotprompt {
     this.schemaResolver = options?.schemaResolver;
     this.partialResolver = options?.partialResolver;
 
-    for (const key in helpers) {
-      this.defineHelper(key, helpers[key as keyof typeof helpers]);
-      this.handlebars.registerHelper(key, helpers[key as keyof typeof helpers]);
-    }
-
-    if (options?.helpers) {
-      for (const key in options.helpers) {
-        this.defineHelper(key, options.helpers[key]);
-      }
-    }
-
-    if (options?.partials) {
-      for (const key in options.partials) {
-        this.definePartial(key, options.partials[key]);
-      }
-    }
+    this.registerInitialHelpers(options?.helpers);
+    this.registerInitialPartials(options?.partials);
   }
 
   /**
@@ -174,6 +160,106 @@ export class Dotprompt {
     return renderer(data, options);
   }
 
+  /**
+   * Compiles a template into a reusable function for rendering prompts.
+   *
+   * @param source The template source or parsed prompt to compile
+   * @param additionalMetadata Additional metadata to include in the compiled template
+   * @return A promise resolving to a function for rendering the template
+   */
+  async compile<
+    Variables = Record<string, unknown>,
+    ModelConfig = Record<string, unknown>,
+  >(
+    source: string | ParsedPrompt<ModelConfig>,
+    additionalMetadata?: PromptMetadata<ModelConfig>
+  ): Promise<PromptFunction<ModelConfig>> {
+    let parsedSource: ParsedPrompt<ModelConfig>;
+    if (typeof source === 'string') {
+      parsedSource = this.parse<ModelConfig>(source);
+    } else {
+      parsedSource = source;
+    }
+
+    if (additionalMetadata) {
+      parsedSource = { ...parsedSource, ...additionalMetadata };
+    }
+
+    // Resolve all partials before compilation.
+    await this.resolvePartials(parsedSource.template);
+
+    const renderString = this.handlebars.compile<Variables>(
+      parsedSource.template,
+      {
+        knownHelpers: this.knownHelpers,
+        knownHelpersOnly: true,
+        noEscape: true,
+      }
+    );
+
+    const renderFunc = async (
+      data: DataArgument,
+      options?: PromptMetadata<ModelConfig>
+    ) => {
+      // Discard the input schema as once rendered it doesn't make sense.
+      const { input, ...mergedMetadata } =
+        await this.renderMetadata(parsedSource);
+
+      const renderedString = renderString(
+        { ...(options?.input?.default || {}), ...data.input },
+        {
+          data: {
+            metadata: {
+              prompt: mergedMetadata,
+              docs: data.docs,
+              messages: data.messages,
+            },
+            ...(data.context || {}),
+          },
+        }
+      );
+
+      return {
+        ...mergedMetadata,
+        messages: toMessages<ModelConfig>(renderedString, data),
+      };
+    };
+    (renderFunc as PromptFunction<ModelConfig>).prompt = parsedSource;
+    return renderFunc as PromptFunction<ModelConfig>;
+  }
+
+  /**
+   * Processes and resolves all metadata for a prompt template.
+   *
+   * @param source The template source or parsed prompt
+   * @param additionalMetadata Additional metadata to include
+   * @return A promise resolving to the fully processed metadata
+   */
+  async renderMetadata<ModelConfig>(
+    source: string | ParsedPrompt<ModelConfig>,
+    additionalMetadata?: PromptMetadata<ModelConfig>
+  ): Promise<PromptMetadata<ModelConfig>> {
+    let parsedSource: ParsedPrompt<ModelConfig>;
+    if (typeof source === 'string') {
+      parsedSource = this.parse<ModelConfig>(source);
+    } else {
+      parsedSource = source;
+    }
+
+    const model =
+      additionalMetadata?.model || parsedSource.model || this.defaultModel;
+
+    let modelConfig: ModelConfig | undefined;
+    if (model && this.modelConfigs[model]) {
+      modelConfig = this.modelConfigs[model] as ModelConfig;
+    }
+
+    return this.resolveMetadata<ModelConfig>(
+      modelConfig ? { config: modelConfig } : {},
+      parsedSource,
+      additionalMetadata
+    );
+  }
   /**
    * Processes schema definitions in picoschema format into standard JSON Schema.
    *
@@ -401,103 +487,37 @@ export class Dotprompt {
   }
 
   /**
-   * Compiles a template into a reusable function for rendering prompts.
-   *
-   * @param source The template source or parsed prompt to compile
-   * @param additionalMetadata Additional metadata to include in the compiled template
-   * @return A promise resolving to a function for rendering the template
+   * Registers initial helpers from built-in helpers and options.
+   * @private
    */
-  async compile<
-    Variables = Record<string, unknown>,
-    ModelConfig = Record<string, unknown>,
-  >(
-    source: string | ParsedPrompt<ModelConfig>,
-    additionalMetadata?: PromptMetadata<ModelConfig>
-  ): Promise<PromptFunction<ModelConfig>> {
-    let parsedSource: ParsedPrompt<ModelConfig>;
-    if (typeof source === 'string') {
-      parsedSource = this.parse<ModelConfig>(source);
-    } else {
-      parsedSource = source;
-    }
-
-    if (additionalMetadata) {
-      parsedSource = { ...parsedSource, ...additionalMetadata };
+  private registerInitialHelpers(
+    customHelpers?: Record<string, Handlebars.HelperDelegate>
+  ): void {
+    // Register built-in helpers
+    for (const key in helpers) {
+      this.defineHelper(key, helpers[key as keyof typeof helpers]);
+      this.handlebars.registerHelper(key, helpers[key as keyof typeof helpers]);
     }
 
-    // Resolve all partials before compilation.
-    await this.resolvePartials(parsedSource.template);
-
-    const renderString = this.handlebars.compile<Variables>(
-      parsedSource.template,
-      {
-        knownHelpers: this.knownHelpers,
-        knownHelpersOnly: true,
-        noEscape: true,
+    // Register custom helpers from options
+    if (customHelpers) {
+      for (const key in customHelpers) {
+        this.defineHelper(key, customHelpers[key]);
       }
-    );
-
-    const renderFunc = async (
-      data: DataArgument,
-      options?: PromptMetadata<ModelConfig>
-    ) => {
-      // Discard the input schema as once rendered it doesn't make sense.
-      const { input, ...mergedMetadata } =
-        await this.renderMetadata(parsedSource);
-
-      const renderedString = renderString(
-        { ...(options?.input?.default || {}), ...data.input },
-        {
-          data: {
-            metadata: {
-              prompt: mergedMetadata,
-              docs: data.docs,
-              messages: data.messages,
-            },
-            ...(data.context || {}),
-          },
-        }
-      );
-
-      return {
-        ...mergedMetadata,
-        messages: toMessages<ModelConfig>(renderedString, data),
-      };
-    };
-    (renderFunc as PromptFunction<ModelConfig>).prompt = parsedSource;
-    return renderFunc as PromptFunction<ModelConfig>;
+    }
   }
 
   /**
-   * Processes and resolves all metadata for a prompt template.
+   * Registers initial partials from the options.
    *
-   * @param source The template source or parsed prompt
-   * @param additionalMetadata Additional metadata to include
-   * @return A promise resolving to the fully processed metadata
+   * @param partials The partials to register
+   * @private
    */
-  async renderMetadata<ModelConfig>(
-    source: string | ParsedPrompt<ModelConfig>,
-    additionalMetadata?: PromptMetadata<ModelConfig>
-  ): Promise<PromptMetadata<ModelConfig>> {
-    let parsedSource: ParsedPrompt<ModelConfig>;
-    if (typeof source === 'string') {
-      parsedSource = this.parse<ModelConfig>(source);
-    } else {
-      parsedSource = source;
-    }
-
-    const model =
-      additionalMetadata?.model || parsedSource.model || this.defaultModel;
-
-    let modelConfig: ModelConfig | undefined;
-    if (model && this.modelConfigs[model]) {
-      modelConfig = this.modelConfigs[model] as ModelConfig;
+  private registerInitialPartials(partials?: Record<string, string>): void {
+    if (partials) {
+      for (const key in partials) {
+        this.definePartial(key, partials[key]);
+      }
     }
-
-    return this.resolveMetadata<ModelConfig>(
-      modelConfig ? { config: modelConfig } : {},
-      parsedSource,
-      additionalMetadata
-    );
   }
 }

python/dotpromptz/src/dotpromptz/dotprompt.py

@@ -44,8 +44,8 @@
 
 import anyio
 
-from dotpromptz.helpers import register_all_helpers
-from dotpromptz.parse import parse_document, to_messages
+from dotpromptz.helpers import BUILTIN_HELPERS
+from dotpromptz.parse import parse_document
 from dotpromptz.picoschema import picoschema_to_json_schema
 from dotpromptz.resolvers import resolve_json_schema, resolve_partial, resolve_tool
 from dotpromptz.typing import (
@@ -200,19 +200,11 @@ def __init__(
         self._partial_resolver: PartialResolver | None = partial_resolver
         self._store: PromptStore | None = None
 
-        self._register_initial_helpers()
-        self._register_initial_partials()
-
-    def _register_initial_helpers(self) -> None:
-        """Register the initial helpers."""
-        register_all_helpers(self._handlebars)
-        for name, fn in self._helpers.items():
-            self._handlebars.register_helper(name, fn)
-
-    def _register_initial_partials(self) -> None:
-        """Register the initial partials."""
-        for name, source in self._partials.items():
-            self._handlebars.register_partial(name, source)
+        self._register_initial_helpers(
+            builtin_helpers=BUILTIN_HELPERS,
+            custom_helpers=helpers,
+        )
+        self._register_initial_partials(partials)
 
     def define_helper(self, name: str, fn: HelperFn) -> Dotprompt:
         """Define a helper function for the template.
@@ -538,3 +530,35 @@ async def _wrapped_schema_resolver(self, name: str) -> JsonSchema | None:
 
         # TODO: Should we cache the resolved schema in self._schemas?
         return await resolve_json_schema(name, self._schema_resolver)
+
+    def _register_initial_helpers(
+        self,
+        builtin_helpers: dict[str, HelperFn] | None = None,
+        custom_helpers: dict[str, HelperFn] | None = None,
+    ) -> None:
+        """Register the initial helpers.
+
+        Built-in helpers are registered first, then custom helpers are
+        registered.
+
+        Args:
+            builtin_helpers: Built-in helpers to register.
+            custom_helpers: Custom helpers to register.
+        """
+        if builtin_helpers is not None:
+            for name, fn in builtin_helpers.items():
+                self.define_helper(name, fn)
+
+        if custom_helpers is not None:
+            for name, fn in custom_helpers.items():
+                self.define_helper(name, fn)
+
+    def _register_initial_partials(self, partials: dict[str, str] | None = None) -> None:
+        """Register the initial partials.
+
+        Args:
+            partials: Partials to register.
+        """
+        if partials is not None:
+            for name, source in partials.items():
+                self.define_partial(name, source)