From 52afca6d342ccdace4026464893d50c386f98cf7 Mon Sep 17 00:00:00 2001
From: Jacob Lee <jacoblee93@gmail.com>
Date: Mon, 6 Nov 2023 18:31:50 -0800
Subject: [PATCH] Adds JSON mode support (#3175)

* Adds JSON mode support

* Update langchain/src/chat_models/tests/chatopenai.int.test.ts

Co-authored-by: Brace Sproul <braceasproul@gmail.com>

* Fix JSON mode

---------

Co-authored-by: Brace Sproul <braceasproul@gmail.com>
---
 .../src/models/chat/integration_openai.ts     | 26 +++++++++++++++++++
 langchain/src/chat_models/openai.ts           |  3 +++
 .../chat_models/tests/chatopenai.int.test.ts  | 14 ++++++++++
 3 files changed, 43 insertions(+)

diff --git a/examples/src/models/chat/integration_openai.ts b/examples/src/models/chat/integration_openai.ts
index 9ce6845c9690..c2ca565c396c 100644
--- a/examples/src/models/chat/integration_openai.ts
+++ b/examples/src/models/chat/integration_openai.ts
@@ -71,3 +71,29 @@ AIMessage {
   }
 }
 */
+
+// Coerce response type with JSON mode.
+// Requires "gpt-4-1106-preview" or later
+const jsonModeModel = new ChatOpenAI({
+  modelName: "gpt-4-1106-preview",
+  maxTokens: 128,
+}).bind({
+  response_format: {
+    type: "json_object",
+  },
+});
+
+// Must be invoked with a system message containing the string "JSON":
+// https://platform.openai.com/docs/guides/text-generation/json-mode
+const res = await jsonModeModel.invoke([
+  ["system", "Only return JSON"],
+  ["human", "Hi there!"],
+]);
+console.log(res);
+
+/*
+  AIMessage {
+    content: '```json\n{\n  "response": "Hello! How can I assist you today?"\n}\n```',
+    additional_kwargs: { function_call: undefined }
+  }
+*/
diff --git a/langchain/src/chat_models/openai.ts b/langchain/src/chat_models/openai.ts
index b58b01eca50a..fc85221433fa 100644
--- a/langchain/src/chat_models/openai.ts
+++ b/langchain/src/chat_models/openai.ts
@@ -155,6 +155,7 @@ export interface ChatOpenAICallOptions
     BaseFunctionCallOptions {
   tools?: StructuredTool[];
   promptIndex?: number;
+  response_format?: { type: "json_object" };
 }
 
 /**
@@ -194,6 +195,7 @@ export class ChatOpenAI<
       "functions",
       "tools",
       "promptIndex",
+      "response_format",
     ];
   }
 
@@ -373,6 +375,7 @@ export class ChatOpenAI<
           ? options?.tools.map(formatToOpenAIFunction)
           : undefined),
       function_call: options?.function_call,
+      response_format: options?.response_format,
       ...this.modelKwargs,
     };
   }
diff --git a/langchain/src/chat_models/tests/chatopenai.int.test.ts b/langchain/src/chat_models/tests/chatopenai.int.test.ts
index f1402158378b..88026a3f2166 100644
--- a/langchain/src/chat_models/tests/chatopenai.int.test.ts
+++ b/langchain/src/chat_models/tests/chatopenai.int.test.ts
@@ -775,3 +775,17 @@ test("Test ChatOpenAI token usage reporting for streaming calls", async () => {
     expect(streamingTokenUsed).toEqual(nonStreamingTokenUsed);
   }
 });
+
+test("Test ChatOpenAI JSON mode", async () => {
+  const chat = new ChatOpenAI({
+    modelName: "gpt-4-1106-preview",
+    maxTokens: 128,
+  }).bind({
+    response_format: {
+      type: "json_object",
+    },
+  });
+  const message = new HumanMessage("Hello!");
+  const res = await chat.invoke([["system", "Only return JSON"], message]);
+  console.log(JSON.stringify(res));
+});