google · pavelgj · Jan 24, 2025 · Jan 24, 2025 · Jan 24, 2025 · Jan 24, 2025
diff --git a/js/src/environment.ts b/js/src/environment.ts
@@ -133,16 +133,26 @@ export class DotpromptEnvironment {
   private async renderPicoschema<ModelConfig>(
     meta: PromptMetadata<ModelConfig>
   ): Promise<PromptMetadata<ModelConfig>> {
-    if (!meta.output?.schema) return meta;
-    return {
-      ...meta,
-      output: {
+    if (!meta.output?.schema && !meta.input?.schema) return meta;
+
+    const newMeta = { ...meta };
+    if (meta.input?.schema) {
+      newMeta.input = {
+        ...meta.input,
+        schema: await picoschema(meta.input.schema, {
+          schemaResolver: this.wrappedSchemaResolver.bind(this),
+        }),
+      };
+    }
+    if (meta.output?.schema) {
+      newMeta.output = {
         ...meta.output,
         schema: await picoschema(meta.output.schema, {
           schemaResolver: this.wrappedSchemaResolver.bind(this),
         }),
-      },
-    };
+      };
+    }
+    return newMeta;
   }
 
   private async wrappedSchemaResolver(name: string): Promise<JSONSchema | null> {
@@ -151,7 +161,7 @@ export class DotpromptEnvironment {
     return null;
   }
 
-  private async renderMetadata<ModelConfig = Record<string, any>>(
+  private async resolveMetadata<ModelConfig = Record<string, any>>(
     base: PromptMetadata<ModelConfig>,
     ...merges: (PromptMetadata<ModelConfig> | undefined)[]
   ): Promise<PromptMetadata<ModelConfig>> {
@@ -162,7 +172,6 @@ export class DotpromptEnvironment {
       out = { ...out, ...merges[i] };
       out.config = { ...config, ...(merges[i]?.config || {}) };
     }
-    delete out.input;
     delete (out as any).template;
     out = removeUndefinedFields(out);
     out = await this.resolveTools(out);
@@ -258,14 +267,9 @@ export class DotpromptEnvironment {
       knownHelpersOnly: true,
     });
 
-    const outFunc = async (data: DataArgument, options?: PromptMetadata<ModelConfig>) => {
-      const selectedModel = options?.model || source.model || this.defaultModel;
-      const modelConfig = this.modelConfigs[selectedModel!] as ModelConfig;
-      const mergedMetadata = await this.renderMetadata<ModelConfig>(
-        modelConfig ? { co
8000
nfig: modelConfig } : {},
-        source,
-        options
-      );
+    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(source);
 
       const renderedString = renderString(
         { ...(options?.input?.default || {}), ...data.input },
@@ -276,13 +280,29 @@ export class DotpromptEnvironment {
           },
         }
       );
+
       return {
         ...mergedMetadata,
         messages: toMessages<ModelConfig>(renderedString, data),
       };
     };
-    (outFunc as PromptFunction<ModelConfig>).prompt = source;
-    return outFunc as PromptFunction<ModelConfig>;
+    (renderFunc as PromptFunction<ModelConfig>).prompt = source;
+    return renderFunc as PromptFunction<ModelConfig>;
+  }
+
+  async renderMetadata<ModelConfig>(
+    source: string | ParsedPrompt<ModelConfig>,
+    additionalMetadata?: PromptMetadata<ModelConfig>
+  ): Promise<PromptMetadata<ModelConfig>> {
+    if (typeof source === "string") source = this.parse<ModelConfig>(source);
+
+    const selectedModel = additionalMetadata?.model || source.model || this.defaultModel;
+    const modelConfig = this.modelConfigs[selectedModel!] as ModelConfig;
+    return this.resolveMetadata<ModelConfig>(
+      modelConfig ? { config: modelConfig } : {},
+      source,
+      additionalMetadata
+    );
   }
 
   get<Variables = any, ModelConfig = Record<string, any>>(

diff --git a/js/src/parse.ts b/js/src/parse.ts
@@ -17,7 +17,6 @@
 import { parse } from "yaml";
 import {
   DataArgument,
-  Document,
   MediaPart,
   Message,
   ParsedPrompt,
@@ -41,10 +40,10 @@ const RESERVED_METADATA_KEYWORDS: (keyof PromptMetadata)[] = [
   "ext",
 ];
 
-const BASE_METADATA: PromptMetadata<ModelConfig> = {
+const BASE_METADATA: PromptMetadata<any> = {
   ext: {},
   metadata: {},
-  config: {} as ModelConfig,
+  config: {},
 };
 
 export function parseDocument<ModelConfig = Record<string, any>>(

diff --git a/js/src/picoschema.ts b/js/src/picoschema.ts
@@ -55,8 +55,8 @@ export class PicoschemaParser {
     if (typeof schema === "string") {
       const [type, description] = extractDescription(schema);
       if (JSON_SCHEMA_SCALAR_TYPES.includes(type)) {
-        const out: JSONSchema = { type };
-        if (description) out.description = description;
+        let out: JSONSchema = { type };
+        if (description) out = { ...out, description };
         return out;
       }
       const resolvedSchema = await this.mustResolveSchema(type);
@@ -80,7 +80,7 @@ export class PicoschemaParser {
       const [type, description] = extractDescription(obj);
       if (!JSON_SCHEMA_SCALAR_TYPES.includes(type)) {
         let resolvedSchema = await this.mustResolveSchema(type);
-        if (description) resolvedSchema.description = description;
+        if (description) resolvedSchema = { ...resolvedSchema, description };
         return resolvedSchema;
       }
 
@@ -139,7 +139,7 @@ export class PicoschemaParser {
         schema.properties[propertyName] = prop;
       } else if (type === "enum") {
         const prop = { enum: obj[key] };
-        if (isOptional) prop.enum.push(null);
+        if (isOptional && !prop.enum.includes(null)) prop.enum.push(null);
         schema.properties[propertyName] = prop;
       } else {
         throw new Error(

diff --git a/js/test/spec.test.ts b/js/test/spec.test.ts
@@ -64,8 +64,8 @@ files
 
               const result = await env.render(s.template, { ...s.data, ...tc.data }, tc.options);
               const { raw, ...prunedResult } = result;
-              const { raw: expectRaw, ...expected } = tc.expect;
-              expect(prunedResult).toEqual({
+              const { raw: expectRaw, input: discardInputForRender, ...expected } = tc.expect;
+              expect(prunedResult, "render should produce the expected result").toEqual({
                 ...expected,
                 ext: expected.ext || {},
                 config: expected.config || {},
@@ -75,6 +75,19 @@ files
               if (tc.expect.raw) {
                 expect(raw).toEqual(expectRaw);
               }
+
+              const metadataResult = await env.renderMetadata(s.template, tc.options);
+              const { raw: metadataResultRaw, ...prunedMetadataResult } = metadataResult;
+              const { messages, raw: metadataExpectRaw, ...expectedMetadata } = tc.expect;
+              expect(
+                prunedMetadataResult,
+                "renderMetadata should produce the expected result"
+              ).toEqual({
+                ...expectedMetadata,
+                ext: expectedMetadata.ext || {},
+                config: expectedMetadata.config || {},
+                metadata: expectedMetadata.metadata || {},
+              });
             });
           });
         });

diff --git a/spec/picoschema.yaml b/spec/picoschema.yaml
@@ -24,6 +24,22 @@
         messages: []
         output:
           schema: { type: string }
+- name: input_and_output
+  template: |
+    ---
+    input:
+      schema: string
+    output:
+      schema: string
+    ---
+  tests:
+    - desc: returns as expected
+      expect:
+        messages: []
+        output:
+          schema: { type: string }
+        input:
+          schema: { type: string }
 - name: simple_scalar_description
   template: |
     ---
@@ -322,9 +338,9 @@
             additionalProperties: false
             required: [foo]
             properties:
-              foo: 
+              foo:
                 type: number
                 description: a foo
               foo2:
-                type: [number, 'null']
-                description: this one is optional
+                type: [number, "null"]
+                description: this one is optional
diff --git a/spec/variables.yaml b/spec/variables.yaml
@@ -29,6 +29,9 @@
       options:
         input: { default: { name: "User" } }
       expect:
+        input:
+          default:
+            name: User
         messages:
           - role: user
             content: [{ text: "Hello, User!\n" }]
@@ -38,6 +41,9 @@
       options:
         input: { default: { name: "User" } }
       expect:
+        input:
+          default:
+            name: User
         messages:
           - role: user
             content: [{ text: "Hello, Pavel!\n" }]