[typescript] Save output.data with text content instead of response data

This comes after Sarmad's schema updates in #589 We only needed to update the `hf.py` and `openai.py`, because `palm.py` aalready returns output in the form of `string | null` type. Ran yarn automated tests, but there aren't any specifically for openai. Not sure how to run the `demo.ts` file which would also be a reasonable test to ensure everything there works too
lastmile-ai · Dec 24, 2023 · c9c2e87 · c9c2e87
1 parent 9b6768f
commit c9c2e87
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 24 deletions.
diff --git a/typescript/__tests__/parsers/hf/hf.test.ts b/typescript/__tests__/parsers/hf/hf.test.ts
@@ -257,7 +257,7 @@ describe("HuggingFaceTextGeneration ModelParser", () => {
 
     const expectedOutput = {
       output_type: "execute_result",
-      data: { generated_text: "Test text generation" },
+      data: "Test text generation",
       execution_count: 0,
       metadata: {},
     };

diff --git a/typescript/lib/parsers/hf.ts b/typescript/lib/parsers/hf.ts
@@ -203,7 +203,7 @@ export class HuggingFaceTextGenerationParser extends ParameterizedModelParser<Te
       const response = await this.hfClient.textGenerationStream(
         textGenerationArgs
       );
-      output = await ConstructStreamOutput(
+      output = await constructStreamOutput(
         response,
         options as InferenceOptions
       );
@@ -240,8 +240,7 @@ export class HuggingFaceTextGenerationParser extends ParameterizedModelParser<Te
     }
 
     if (output.output_type === "execute_result") {
-      return (output.data as TextGenerationOutput | TextGenerationStreamOutput)
-        .generated_text as string;
+      return output.data as string;
     } else {
       return "";
     }
@@ -254,7 +253,7 @@ export class HuggingFaceTextGenerationParser extends ParameterizedModelParser<Te
  * @param options
  * @returns
  */
-async function ConstructStreamOutput(
+async function constructStreamOutput(
   response: AsyncGenerator<TextGenerationStreamOutput>,
   options: InferenceOptions
 ): Promise<Output> {
@@ -272,7 +271,7 @@ async function ConstructStreamOutput(
 
     output = {
       output_type: "execute_result",
-      data: delta,
+      data: accumulatedMessage,
       execution_count: index,
       metadata: metadata,
     } as ExecuteResult;
@@ -281,16 +280,12 @@ async function ConstructStreamOutput(
 }
 
 function constructOutput(response: TextGenerationOutput): Output {
-  const metadata = {};
-  const data = response;
-
   const output = {
     output_type: "execute_result",
-    data: data,
+    data: response.generated_text,
     execution_count: 0,
-    metadata: metadata,
+    metadata: _.omit(response, ["generated_text"]),
   } as ExecuteResult;
-
   return output;
 }
 

diff --git a/typescript/lib/parsers/openai.ts b/typescript/lib/parsers/openai.ts
@@ -365,6 +365,7 @@ export class OpenAIChatModelParser extends ParameterizedModelParser<Chat.ChatCom
         const input: PromptInput =
           message.role === "user" ? message.content ?? "" : { ...message };
 
+        const responseWithoutContent = _.omit(assistantResponse, ["content"]);
         const prompt: Prompt = {
           name: `${promptName}_${prompts.length + 1}`,
           input,
@@ -378,7 +379,8 @@ export class OpenAIChatModelParser extends ParameterizedModelParser<Chat.ChatCom
               ? [
                   {
                     output_type: "execute_result",
-                    data: { ...assistantResponse },
+                    data: assistantResponse.content,
+                    metadata: { ...responseWithoutContent },
                   },
                 ]
               : undefined,
@@ -536,13 +538,15 @@ export class OpenAIChatModelParser extends ParameterizedModelParser<Chat.ChatCom
       const outputs: ExecuteResult[] = [];
       const responseWithoutChoices = _.omit(response, "choices");
       for (const choice of response.choices) {
+        const messageWithoutContent = _.omit(choice.message, ["content"]);
         const output: ExecuteResult = {
           output_type: "execute_result",
-          data: { ...choice.message },
+          data: choice.message?.content,
           execution_count: choice.index,
           metadata: {
             finish_reason: choice.finish_reason,
             ...responseWithoutChoices,
+            ...messageWithoutContent,
           },
         };
 
@@ -585,12 +589,16 @@ export class OpenAIChatModelParser extends ParameterizedModelParser<Chat.ChatCom
             /*index*/ choice.index
           );
 
+          const messageWithoutContent = _.omit(message, ["content"]);
           const output: ExecuteResult = {
             output_type: "execute_result",
-            data: { ...message },
+            // TODO (rossdanlm): Handle ChatCompletionMessage.function_call
+            // too (next diff)
+            data: message?.content,
             execution_count: choice.index,
             metadata: {
               finish_reason: choice.finish_reason,
+              ...messageWithoutContent,
             },
           };
           outputs.set(choice.index, output);
@@ -625,11 +633,10 @@ export class OpenAIChatModelParser extends ParameterizedModelParser<Chat.ChatCom
     }
 
     if (output.output_type === "execute_result") {
-      const message = output.data as Chat.ChatCompletionMessageParam;
-      if (message.content != null) {
-        return message.content;
-      } else if (message.function_call) {
-        return JSON.stringify(message.function_call);
+      if (typeof output.data === "string") {
+        return output.data;
+      } else if (output.metadata?.function_call) {
+        return JSON.stringify(output.metadata?.function_call);
       } else {
         return "";
       }
@@ -671,11 +678,17 @@ export class OpenAIChatModelParser extends ParameterizedModelParser<Chat.ChatCom
     const output = aiConfig.getLatestOutput(prompt);
     if (output != null) {
       if (output.output_type === "execute_result") {
-        const outputMessage =
-          output.data as unknown as Chat.ChatCompletionMessageParam;
         // If the prompt has output saved, add it to the messages array
-        if (outputMessage.role === "assistant") {
-          messages.push(outputMessage);
+        if (output.metadata?.role === "assistant") {
+          if (typeof output.data === "string") {
+            messages.push({
+              content: output.data,
+              role: output.metadata?.role,
+              function_call: output.metadata?.function_call,
+              name: output.metadata?.name,
+            });
+          }
+          // TODO (rossdanlm): Support function_call
         }
       }
     }