fix(openai): Revert Azure default withStructuredOutput changes (#7596)

langchain-ai · Jan 25, 2025 · 0ecaad1 · 0ecaad1
1 parent 0eae6d9
commit 0ecaad1
Show file tree

Hide file tree

Showing 3 changed files with 96 additions and 1 deletion.
diff --git a/libs/langchain-openai/src/azure/chat_models.ts b/libs/langchain-openai/src/azure/chat_models.ts
@@ -4,7 +4,14 @@ import {
   type BaseChatModelParams,
 } from "@langchain/core/language_models/chat_models";
 import { getEnvironmentVariable } from "@langchain/core/utils/env";
-import { ChatOpenAI } from "../chat_models.js";
+import { BaseLanguageModelInput } from "@langchain/core/language_models/base";
+import { BaseMessage } from "@langchain/core/messages";
+import { Runnable } from "@langchain/core/runnables";
+import { z } from "zod";
+import {
+  ChatOpenAI,
+  ChatOpenAIStructuredOutputMethodOptions,
+} from "../chat_models.js";
 import { OpenAIEndpointConfig, getEndpoint } from "../utils/azure.js";
 import {
   AzureOpenAIInput,
@@ -635,4 +642,64 @@ export class AzureChatOpenAI extends ChatOpenAI {
 
     return json;
   }
+
+  withStructuredOutput<
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    RunOutput extends Record<string, any> = Record<string, any>
+  >(
+    outputSchema:
+      | z.ZodType<RunOutput>
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      | Record<string, any>,
+    config?: ChatOpenAIStructuredOutputMethodOptions<false>
+  ): Runnable<BaseLanguageModelInput, RunOutput>;
+
+  withStructuredOutput<
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    RunOutput extends Record<string, any> = Record<string, any>
+  >(
+    outputSchema:
+      | z.ZodType<RunOutput>
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      | Record<string, any>,
+    config?: ChatOpenAIStructuredOutputMethodOptions<true>
+  ): Runnable<BaseLanguageModelInput, { raw: BaseMessage; parsed: RunOutput }>;
+
+  withStructuredOutput<
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    RunOutput extends Record<string, any> = Record<string, any>
+  >(
+    outputSchema:
+      | z.ZodType<RunOutput>
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      | Record<string, any>,
+    config?: ChatOpenAIStructuredOutputMethodOptions<boolean>
+  ):
+    | Runnable<BaseLanguageModelInput, RunOutput>
+    | Runnable<BaseLanguageModelInput, { raw: BaseMessage; parsed: RunOutput }>;
+
+  withStructuredOutput<
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    RunOutput extends Record<string, any> = Record<string, any>
+  >(
+    outputSchema:
+      | z.ZodType<RunOutput>
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      | Record<string, any>,
+    config?: ChatOpenAIStructuredOutputMethodOptions<boolean>
+  ):
+    | Runnable<BaseLanguageModelInput, RunOutput>
+    | Runnable<
+        BaseLanguageModelInput,
+        { raw: BaseMessage; parsed: RunOutput }
+      > {
+    const ensuredConfig = { ...config };
+    // Not all Azure gpt-4o deployments models support jsonSchema yet
+    if (this.model.startsWith("gpt-4o")) {
+      if (ensuredConfig?.method === undefined) {
+        ensuredConfig.method = "functionCalling";
+      }
+    }
+    return super.withStructuredOutput<RunOutput>(outputSchema, ensuredConfig);
+  }
 }
diff --git a/libs/langchain-openai/src/chat_models.ts b/libs/langchain-openai/src/chat_models.ts
@@ -1884,6 +1884,19 @@ export class ChatOpenAI<
     config?: ChatOpenAIStructuredOutputMethodOptions<true>
   ): Runnable<BaseLanguageModelInput, { raw: BaseMessage; parsed: RunOutput }>;
 
+  withStructuredOutput<
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    RunOutput extends Record<string, any> = Record<string, any>
+  >(
+    outputSchema:
+      | z.ZodType<RunOutput>
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      | Record<string, any>,
+    config?: ChatOpenAIStructuredOutputMethodOptions<boolean>
+  ):
+    | Runnable<BaseLanguageModelInput, RunOutput>
+    | Runnable<BaseLanguageModelInput, { raw: BaseMessage; parsed: RunOutput }>;
+
   withStructuredOutput<
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
     RunOutput extends Record<string, any> = Record<string, any>

diff --git a/libs/langchain-openai/src/tests/azure/chat_models.int.test.ts b/libs/langchain-openai/src/tests/azure/chat_models.int.test.ts
@@ -1,6 +1,7 @@
 /* eslint-disable no-process-env */
 
 import { test, jest, expect } from "@jest/globals";
+import { z } from "zod";
 import {
   BaseMessage,
   ChatMessage,
@@ -950,3 +951,17 @@ testFn("Test Azure ChatOpenAI with bearer token provider", async () => {
   const res = await chat.invoke([["system", "Say hi"], message]);
   // console.log(res);
 });
+
+test("Test Azure ChatOpenAI withStructuredOutput", async () => {
+  const chat = new AzureChatOpenAI({
+    modelName: "gpt-4o-mini",
+  });
+  const message = new HumanMessage("Good!");
+  const model = await chat.withStructuredOutput(
+    z.object({
+      sentiment: z.string(),
+    })
+  );
+  const res = await model.invoke([message]);
+  expect(res.sentiment).toBeDefined();
+});