From 2822fe4bd5a08da77d9f91da3b787245b2fb450c Mon Sep 17 00:00:00 2001
From: pelikhan <jhalleux@microsoft.com>
Date: Wed, 4 Sep 2024 11:20:48 -0700
Subject: [PATCH] Refactor and enhance configurations in prompty.ts and update
 script models in sample package

---
 packages/core/src/prompty.ts        | 32 +++++++++++++++++++++--------
 packages/sample/src/basic.genai.mts |  1 +
 packages/sample/src/chat.genai.mts  |  1 +
 3 files changed, 25 insertions(+), 9 deletions(-)

diff --git a/packages/core/src/prompty.ts b/packages/core/src/prompty.ts
index de93b82bfa..430d4d1bb3 100644
--- a/packages/core/src/prompty.ts
+++ b/packages/core/src/prompty.ts
@@ -22,14 +22,14 @@ export interface PromptyFrontmatter {
     outputs?: JSONSchemaObject
     model?: {
         api?: "chat" | "completion"
-        configuration?:
-            | { type?: string; name?: string; organization?: string }
-            | {
-                  type?: string
-                  api_version?: string
-                  azure_deployment: string
-                  azure_endpoint: string
-              }
+        configuration?: {
+            type?: string
+            name?: string
+            organization?: string
+            api_version?: string
+            azure_deployment: string
+            azure_endpoint: string
+        }
         parameters?: {
             response_format: "json_object"
             max_tokens?: number
@@ -83,7 +83,13 @@ export function promptyParse(text: string): PromptyDocument {
 
 export function promptyToGenAIScript(doc: PromptyDocument) {
     const { frontmatter, messages } = doc
-    const { name, description, tags, sample, inputs, outputs } = frontmatter
+    const { name, description, tags, sample, inputs, outputs, model } =
+        frontmatter
+    const {
+        api = "chat",
+        configuration,
+        parameters: modelParameters,
+    } = model ?? {}
     const parameters = inputs ? structuredClone(inputs) : undefined
     if (parameters && sample && typeof sample === "object")
         for (const p in sample) {
@@ -91,7 +97,15 @@ export function promptyToGenAIScript(doc: PromptyDocument) {
             if (s !== undefined && parameters[p].type !== "object")
                 parameters[p].default = s
         }
+
+    let modelName: string = undefined
+    if (api !== "chat") throw new Error("completion api not supported")
+    if (configuration?.azure_deployment)
+        modelName = `azure:${configuration.azure_deployment}`
+    else if (configuration?.type) modelName = `openai:${configuration.type}`
+
     const meta = deleteUndefinedValues(<PromptArgs>{
+        model: modelName,
         title: name,
         description,
         tags,
diff --git a/packages/sample/src/basic.genai.mts b/packages/sample/src/basic.genai.mts
index 2244d08c0e..a3e10dc9f1 100644
--- a/packages/sample/src/basic.genai.mts
+++ b/packages/sample/src/basic.genai.mts
@@ -1,4 +1,5 @@
 script({
+    model: "azure:gpt-4o",
     title: "Basic Prompt",
     description: "A basic prompt that uses the chat API to answer questions",
     parameters: {
diff --git a/packages/sample/src/chat.genai.mts b/packages/sample/src/chat.genai.mts
index acbdb1530a..6de07437cd 100644
--- a/packages/sample/src/chat.genai.mts
+++ b/packages/sample/src/chat.genai.mts
@@ -1,4 +1,5 @@
 script({
+    model: "azure:gpt-35-turbo",
     title: "Contoso Chat Prompt",
     description: "A retail assistent for Contoso Outdoors products retailer.",
     parameters: {