From e8223f34ae015398ef060444286f1fa7dd0cd622 Mon Sep 17 00:00:00 2001
From: "Rossdan Craig rossdan@lastmileai.dev" <>
Date: Wed, 27 Dec 2023 10:13:09 -0500
Subject: [PATCH] [typescript] Save output.data with OutputData type instead of
 pure string

This will make it easier for Ryan to parse because he just has to check the `.kind` field and if it's string, he knows that the string is meant to be string only. I also kept the existing form where it will still work if we ONLY pass in string, in which case we'd just assume it's supposed to be in pure string form

I also added function call support to this diff for `openai.ts`. Sorry yea I know it's a bigger diff.

##Test plan
Pass automated tests, running these commands from `aiconfig` top-level dir
```
npx ts-node typescript/demo/function-call-stream.ts
npx ts-node typescript/demo/demo.ts
npx ts-node typescript/demo/test-hf.ts
```

Also run `yarn test` from typescript dir
---
 extensions/HuggingFace/typescript/hf.ts       |  50 ++++--
 typescript/__tests__/parsers/hf/hf.test.ts    |   5 +-
 .../__tests__/parsers/palm-text/palm.test.ts  |   5 +-
 typescript/lib/parsers/hf.ts                  |  36 ++--
 typescript/lib/parsers/openai.ts              | 155 ++++++++++++++----
 typescript/lib/parsers/palm.ts                |  27 ++-
 6 files changed, 209 insertions(+), 69 deletions(-)

diff --git a/extensions/HuggingFace/typescript/hf.ts b/extensions/HuggingFace/typescript/hf.ts
index d8a638aaa..6e0e29c29 100644
--- a/extensions/HuggingFace/typescript/hf.ts
+++ b/extensions/HuggingFace/typescript/hf.ts
@@ -6,15 +6,16 @@ import {
 } from "@huggingface/inference";
 
 import {
-  Prompt,
-  Output,
-  PromptInput,
-  ParameterizedModelParser,
-  ModelMetadata,
-  ExecuteResult,
   AIConfigRuntime,
-  InferenceOptions,
   CallbackEvent,
+  ExecuteResult,
+  InferenceOptions,
+  ModelMetadata,
+  Output,
+  OutputData,
+  ParameterizedModelParser,
+  Prompt,
+  PromptInput,
 } from "aiconfig";
 import _ from "lodash";
 import * as aiconfig from "aiconfig";
@@ -248,7 +249,14 @@ export class HuggingFaceTextGenerationModelParserExtension extends Parameterized
     }
 
     if (output.output_type === "execute_result") {
-      if (typeof output.data === "string") {
+      if (output.data?.hasOwnProperty("value")) {
+        const outputData = output.data as OutputData;
+        if (typeof outputData.value === "string") {
+          return outputData.value;
+        }
+        // should never get here for this model parser, just being safe
+        return JSON.stringify(outputData.value);
+      } else if (typeof output.data === "string") {
         return output.data;
       }
 
@@ -278,30 +286,36 @@ async function constructStreamOutput(
   let output = {} as ExecuteResult;
 
   for await (const iteration of response) {
-    const data = iteration.token.text;
-    const metadata = iteration;
+    const newText = iteration.token.text;
 
-    accumulatedMessage += data;
-    const delta = data;
+    accumulatedMessage += newText;
     const index = 0;
-    options.callbacks!.streamCallback(delta, accumulatedMessage, 0);
+    options.callbacks!.streamCallback(newText, accumulatedMessage, index);
 
+    const outputData: OutputData = {
+      kind: "string",
+      // TODO: Investigate if we should use the accumulated message instead
+      // of newText: https://github.com/lastmile-ai/aiconfig/issues/620
+      value: newText,
+    };
     output = {
       output_type: "execute_result",
-      // TODO: Investigate if we should use the accumulated message instead
-      // of delta: https://github.com/lastmile-ai/aiconfig/issues/620
-      data: delta,
+      data: outputData,
       execution_count: index,
-      metadata: metadata,
+      metadata: iteration,
     } as ExecuteResult;
   }
   return output;
 }
 
 function constructOutput(response: TextGenerationOutput): Output {
+  const data: OutputData = {
+    kind: "string",
+    value: response.generated_text,
+  };
   const output = {
     output_type: "execute_result",
-    data: response.generated_text,
+    data,
     execution_count: 0,
     metadata: { rawResponse: response },
   } as ExecuteResult;
diff --git a/typescript/__tests__/parsers/hf/hf.test.ts b/typescript/__tests__/parsers/hf/hf.test.ts
index ed3dd1c51..51afb346f 100644
--- a/typescript/__tests__/parsers/hf/hf.test.ts
+++ b/typescript/__tests__/parsers/hf/hf.test.ts
@@ -257,7 +257,10 @@ describe("HuggingFaceTextGeneration ModelParser", () => {
 
     const expectedOutput = {
       output_type: "execute_result",
-      data: "Test text generation",
+      data: {
+        kind: "string",
+        value: "Test text generation",
+      },
       execution_count: 0,
       metadata: {
         rawResponse: {
diff --git a/typescript/__tests__/parsers/palm-text/palm.test.ts b/typescript/__tests__/parsers/palm-text/palm.test.ts
index cdbffe309..e204e33d1 100644
--- a/typescript/__tests__/parsers/palm-text/palm.test.ts
+++ b/typescript/__tests__/parsers/palm-text/palm.test.ts
@@ -122,6 +122,9 @@ describe("PaLM Text ModelParser", () => {
     const aiconfig = AIConfigRuntime.load(PALM_CONFIG_PATH);
 
     const [result] = (await aiconfig.run("prompt1")) as Output[];
-    expect((result as ExecuteResult).data).toEqual("Ranch");
+    expect((result as ExecuteResult).data).toEqual({
+      kind: "string",
+      value: "Ranch",
+    });
   });
 });
diff --git a/typescript/lib/parsers/hf.ts b/typescript/lib/parsers/hf.ts
index f4048e56d..3d50a6289 100644
--- a/typescript/lib/parsers/hf.ts
+++ b/typescript/lib/parsers/hf.ts
@@ -12,6 +12,7 @@ import {
   ExecuteResult,
   ModelMetadata,
   Output,
+  OutputData,
   Prompt,
   PromptInput,
 } from "../../types";
@@ -240,7 +241,15 @@ export class HuggingFaceTextGenerationParser extends ParameterizedModelParser<Te
     }
 
     if (output.output_type === "execute_result") {
-      if (typeof output.data === "string") {
+      if (output.data?.hasOwnProperty("value")) {
+        const outputData = output.data as OutputData;
+        if (typeof outputData.value === "string") {
+          return outputData.value;
+        }
+        // Sarmad + Ryan, let me know if you prefer I do something else
+        // Was basing this off existing code for function calls in openai.ts
+        return JSON.stringify(outputData.value);
+      } else if (typeof output.data === "string") {
         return output.data;
       }
 
@@ -270,30 +279,37 @@ async function constructStreamOutput(
   let output = {} as ExecuteResult;
 
   for await (const iteration of response) {
-    const data = iteration.token.text;
+    const newText = iteration.token.text;
     const metadata = iteration;
 
-    accumulatedMessage += data;
-    const delta = data;
+    accumulatedMessage += newText;
     const index = 0;
-    options.callbacks!.streamCallback(delta, accumulatedMessage, 0);
+    options.callbacks!.streamCallback(newText, accumulatedMessage, 0);
 
+    const data: OutputData = {
+      kind: "string",
+      // TODO: Investigate if we should use the accumulated message instead
+      // of newText: https://github.com/lastmile-ai/aiconfig/issues/620
+      value: newText,
+    };
     output = {
       output_type: "execute_result",
-      // TODO: Investigate if we should use the accumulated message instead
-      // of delta: https://github.com/lastmile-ai/aiconfig/issues/620
-      data: delta,
+      data,
       execution_count: index,
-      metadata: { metadata },
+      metadata,
     } as ExecuteResult;
   }
   return output;
 }
 
 function constructOutput(response: TextGenerationOutput): Output {
+  const data: OutputData = {
+    kind: "string",
+    value: response.generated_text,
+  };
   const output = {
     output_type: "execute_result",
-    data: response.generated_text,
+    data,
     execution_count: 0,
     metadata: { rawResponse: response },
   } as ExecuteResult;
diff --git a/typescript/lib/parsers/openai.ts b/typescript/lib/parsers/openai.ts
index f19386a77..b92c8d287 100644
--- a/typescript/lib/parsers/openai.ts
+++ b/typescript/lib/parsers/openai.ts
@@ -1,10 +1,13 @@
 import { JSONObject, JSONValue } from "../../common";
 import {
-  Prompt,
+  ExecuteResult,
+  FunctionData,
+  ModelMetadata,
   Output,
+  OutputData,
+  OutputDataWithFunctionValue,
+  Prompt,
   PromptInput,
-  ModelMetadata,
-  ExecuteResult,
 } from "../../types";
 import { AIConfigRuntime } from "../config";
 import { ParameterizedModelParser } from "../parameterizedModelParser";
@@ -171,9 +174,13 @@ export class OpenAIModelParser extends ParameterizedModelParser<CompletionCreate
       // Save response as Output(s) in the Prompt
       const outputs: ExecuteResult[] = [];
       for (const choice of response.choices) {
+        const data: OutputData = {
+          kind: "string",
+          value: choice.text, // CompletionChoice does not support function calls
+        };
         const output: ExecuteResult = {
           output_type: "execute_result",
-          data: choice.text,
+          data,
           execution_count: choice.index,
           metadata: {
             finish_reason: choice.finish_reason,
@@ -217,9 +224,13 @@ export class OpenAIModelParser extends ParameterizedModelParser<CompletionCreate
           );
 
           const chunkWithoutChoices = _.omit(chunk, "choices");
+          const data: OutputData = {
+            kind: "string",
+            value: accumulatedText,
+          };
           const output: ExecuteResult = {
             output_type: "execute_result",
-            data: accumulatedText,
+            data,
             execution_count: choice.index,
             metadata: {
               finish_reason: choice.finish_reason,
@@ -252,8 +263,13 @@ export class OpenAIModelParser extends ParameterizedModelParser<CompletionCreate
     }
 
     if (output.output_type === "execute_result") {
-      // TODO: Add in OutputData another way to support function calls
-      if (typeof output.data === "string") {
+      if (output.data?.hasOwnProperty("value")) {
+        const outputData = output.data as OutputData;
+        if (typeof outputData.value === "string") {
+          return outputData.value;
+        }
+        return JSON.stringify(outputData.value);
+      } else if (typeof output.data === "string") {
         return output.data;
       }
 
@@ -379,7 +395,8 @@ export class OpenAIChatModelParser extends ParameterizedModelParser<Chat.ChatCom
 
         const input: PromptInput =
           message.role === "user" ? message.content ?? "" : { ...message };
-
+        const outputs =
+          this.buildAssistantOutputDataIfNeeded(assistantResponse);
         const prompt: Prompt = {
           name: `${promptName}_${prompts.length + 1}`,
           input,
@@ -388,16 +405,7 @@ export class OpenAIChatModelParser extends ParameterizedModelParser<Chat.ChatCom
             parameters: params ?? {},
             remember_chat_context: true,
           },
-          outputs:
-            assistantResponse != null
-              ? [
-                  {
-                    output_type: "execute_result",
-                    data: assistantResponse.content,
-                    metadata: { rawResponse: assistantResponse },
-                  },
-                ]
-              : undefined,
+          outputs,
         };
 
         prompts.push(prompt);
@@ -552,9 +560,16 @@ export class OpenAIChatModelParser extends ParameterizedModelParser<Chat.ChatCom
       const outputs: ExecuteResult[] = [];
       const responseWithoutChoices = _.omit(response, "choices");
       for (const choice of response.choices) {
+        const outputData: OutputData | undefined = buildOutputData(
+          choice.message
+        );
+        if (outputData === undefined) {
+          continue;
+        }
+
         const output: ExecuteResult = {
           output_type: "execute_result",
-          data: choice.message?.content,
+          data: outputData,
           execution_count: choice.index,
           metadata: {
             finish_reason: choice.finish_reason,
@@ -562,7 +577,6 @@ export class OpenAIChatModelParser extends ParameterizedModelParser<Chat.ChatCom
             rawResponse: choice.message,
           },
         };
-
         outputs.push(output);
       }
 
@@ -601,12 +615,18 @@ export class OpenAIChatModelParser extends ParameterizedModelParser<Chat.ChatCom
             /*accumulatedData*/ message,
             /*index*/ choice.index
           );
+          if (message === undefined) {
+            continue;
+          }
+
+          const outputData: OutputData | undefined = buildOutputData(message);
+          if (outputData === undefined) {
+            continue;
+          }
 
           const output: ExecuteResult = {
             output_type: "execute_result",
-            // TODO (rossdanlm): Handle ChatCompletionMessage.function_call
-            // too (next diff)
-            data: message?.content,
+            data: outputData,
             execution_count: choice.index,
             metadata: {
               finish_reason: choice.finish_reason,
@@ -645,9 +665,16 @@ export class OpenAIChatModelParser extends ParameterizedModelParser<Chat.ChatCom
     }
 
     if (output.output_type === "execute_result") {
-      // TODO: Add in OutputData another way to support function calls
-      if (typeof output.data === "string") {
+      if (output.data?.hasOwnProperty("value")) {
+        const outputData = output.data as OutputData;
+        if (typeof outputData.value === "string") {
+          return outputData.value;
+        }
+        return JSON.stringify(outputData.value); // function_call
+      } else if (typeof output.data === "string") {
         return output.data;
+      } else if (output.metadata?.rawResponse?.function_call) {
+        return JSON.stringify(output.metadata?.rawResponse?.function_call);
       }
 
       // Doing this to be backwards-compatible with old output format
@@ -700,17 +727,32 @@ export class OpenAIChatModelParser extends ParameterizedModelParser<Chat.ChatCom
     const output = aiConfig.getLatestOutput(prompt);
     if (output != null) {
       if (output.output_type === "execute_result") {
-        // If the prompt has output saved, add it to the messages array
         if (output.metadata?.rawResponse?.role === "assistant") {
-          if (typeof output.data === "string") {
-            messages.push({
-              content: output.data,
-              // TODO (rossdanlm): Support function_call and don't rely on rawResponse
-              role: output.metadata?.rawResponse?.role,
-              function_call: output.metadata?.rawResponse?.function_call,
-              name: output.metadata?.rawResponse?.name,
-            });
+          output as ExecuteResult;
+          let content: string | null = null;
+          let function_call: FunctionData | undefined = undefined;
+
+          if (output.data?.hasOwnProperty("value")) {
+            const outputData = output.data as OutputData;
+            if (typeof outputData.value === "string") {
+              content = outputData.value;
+            } else if (outputData.kind === "tool_calls") {
+              outputData as OutputDataWithFunctionValue;
+              // Typescript schema does not support array of function calls yet,
+              // but Python does so doing this for forward compatibility
+              function_call =
+                outputData.value[outputData.value.length - 1].function;
+            }
+          } else if (typeof output.data === "string") {
+            content = output.data;
           }
+
+          messages.push({
+            content,
+            role: output.metadata?.rawResponse?.role,
+            ...(function_call !== undefined && { function_call }),
+            name: output.metadata?.rawResponse?.name,
+          });
         }
 
         // Doing this to be backwards-compatible with old output format
@@ -730,6 +772,26 @@ export class OpenAIChatModelParser extends ParameterizedModelParser<Chat.ChatCom
 
     return messages;
   }
+
+  private buildAssistantOutputDataIfNeeded(
+    assistantResponse: Chat.ChatCompletionMessageParam | null
+  ): Output[] | undefined {
+    let outputs: Output[] | undefined = undefined;
+    if (assistantResponse != null) {
+      const assistantOutputData: OutputData | undefined =
+        buildOutputData(assistantResponse);
+      if (assistantOutputData !== undefined) {
+        outputs = [
+          {
+            output_type: "execute_result",
+            data: assistantOutputData,
+            metadata: { rawResponse: assistantResponse },
+          },
+        ];
+      }
+    }
+    return outputs;
+  }
 }
 
 //#region Utils
@@ -851,6 +913,31 @@ function multiChoiceMessageReducer(
   return messages;
 }
 
+function buildOutputData(
+  message: Chat.ChatCompletionMessageParam | null
+): OutputData | undefined {
+  let outputData: OutputData | undefined = undefined;
+  if (message !== null) {
+    if (message.content !== null) {
+      outputData = {
+        kind: "string",
+        value: message.content,
+      };
+    } else if (message.function_call !== undefined) {
+      outputData = {
+        kind: "tool_calls",
+        value: [
+          {
+            type: "function",
+            function: message.function_call,
+          },
+        ],
+      };
+    }
+  }
+  return outputData;
+}
+
 export function countTokens(text: string): number {
   const gpt3TokenCount = gpt3Encode(text).length;
   return gpt3TokenCount;
diff --git a/typescript/lib/parsers/palm.ts b/typescript/lib/parsers/palm.ts
index 12916e4b5..714440bb8 100644
--- a/typescript/lib/parsers/palm.ts
+++ b/typescript/lib/parsers/palm.ts
@@ -1,6 +1,12 @@
 import { TextServiceClient, protos } from "@google-ai/generativelanguage";
 import { JSONObject } from "../../common";
-import { Prompt, Output, ModelMetadata, ExecuteResult } from "../../types";
+import {
+  ExecuteResult,
+  ModelMetadata,
+  Output,
+  OutputData,
+  Prompt,
+} from "../../types";
 import { AIConfigRuntime } from "../config";
 import { InferenceOptions } from "../modelParser";
 import { ParameterizedModelParser } from "../parameterizedModelParser";
@@ -180,10 +186,17 @@ export class PaLMTextParser extends ParameterizedModelParser {
     }
 
     if (output.output_type === "execute_result") {
-      return output.data as string;
-    } else {
-      return "";
+      if (output.data?.hasOwnProperty("value")) {
+        const outputData = output.data as OutputData;
+        if (typeof outputData.value === "string") {
+          return outputData.value;
+        }
+        return JSON.stringify(outputData.value);
+      } else if (typeof output.data === "string") {
+        return output.data;
+      }
     }
+    return "";
   }
 }
 
@@ -230,9 +243,13 @@ function constructOutputs(
 
   for (let i = 0; i < response.candidates.length; i++) {
     const candidate = response.candidates[i];
+    const data: OutputData = {
+      kind: "string",
+      value: candidate.output ?? "",
+    };
     const output: ExecuteResult = {
       output_type: "execute_result",
-      data: candidate.output,
+      data,
       execution_count: i,
       metadata: { rawResponse: candidate },
     };