feat(xai): xAI polish

libs/langchain-xai/package.json

@@ -35,7 +35,8 @@
   "author": "LangChain",
   "license": "MIT",
   "dependencies": {
-    "@langchain/openai": "~0.3.0"
+    "@langchain/openai": "~0.4.4",
+    "zod": "^3.24.2"
   },
   "peerDependencies": {
     "@langchain/core": ">=0.2.21 <0.4.0"
@@ -67,7 +68,6 @@
     "rollup": "^4.5.2",
     "ts-jest": "^29.1.0",
     "typescript": "<5.2.0",
-    "zod": "^3.22.4",
     "zod-to-json-schema": "^3.23.1"
   },
   "publishConfig": {

libs/langchain-xai/src/chat_models.ts

@@ -1,17 +1,22 @@
+import { BaseLanguageModelInput } from "@langchain/core/language_models/base";
 import {
   BaseChatModelCallOptions,
   BindToolsInput,
   LangSmithParams,
   type BaseChatModelParams,
 } from "@langchain/core/language_models/chat_models";
 import { Serialized } from "@langchain/core/load/serializable";
+import { AIMessageChunk, BaseMessage } from "@langchain/core/messages";
+import { Runnable } from "@langchain/core/runnables";
 import { getEnvironmentVariable } from "@langchain/core/utils/env";
 import {
   type OpenAICoreRequestOptions,
   type OpenAIClient,
   ChatOpenAI,
   OpenAIToolChoice,
+  ChatOpenAIStructuredOutputMethodOptions,
 } from "@langchain/openai";
+import { z } from "zod";
 
 type ChatXAIToolType = BindToolsInput | OpenAIClient.ChatCompletionTool;
 
@@ -494,4 +499,104 @@ export class ChatXAI extends ChatOpenAI<ChatXAICallOptions> {
 
     return super.completionWithRetry(newRequest, options);
   }
+
+  protected override _convertOpenAIDeltaToBaseMessageChunk(
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    delta: Record<string, any>,
+    rawResponse: OpenAIClient.ChatCompletionChunk,
+    defaultRole?:
+      | "function"
+      | "user"
+      | "system"
+      | "developer"
+      | "assistant"
+      | "tool"
+  ) {
+    const messageChunk: AIMessageChunk =
+      super._convertOpenAIDeltaToBaseMessageChunk(
+        delta,
+        rawResponse,
+        defaultRole
+      );
+    // Make concatenating chunks work without merge warning
+    if (!rawResponse.choices[0]?.finish_reason) {
+      delete messageChunk.response_metadata.usage;
+      delete messageChunk.usage_metadata;
+    } else {
+      messageChunk.usage_metadata = messageChunk.response_metadata.usage;
+    }
+    return messageChunk;
+  }
+
+  protected override _convertOpenAIChatCompletionMessageToBaseMessage(
+    message: OpenAIClient.ChatCompletionMessage,
+    rawResponse: OpenAIClient.ChatCompletion
+  ) {
+    const langChainMessage =
+      super._convertOpenAIChatCompletionMessageToBaseMessage(
+        message,
+        rawResponse
+      );
+    langChainMessage.additional_kwargs.reasoning_content =
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      (message as any).reasoning_content;
+    return langChainMessage;
+  }
+
+  override 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>;
+
+  override 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 }>;
+
+  override 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 }>;
+
+  override 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 };
+    if (ensuredConfig?.method === undefined) {
+      ensuredConfig.method = "functionCalling";
+    }
+    return super.withStructuredOutput<RunOutput>(outputSchema, ensuredConfig);
+  }
 }

libs/langchain-xai/src/tests/chat_models.int.test.ts

@@ -68,7 +68,7 @@ test("streaming", async () => {
 test("invoke with bound tools", async () => {
   const chat = new ChatXAI({
     maxRetries: 0,
-    model: "grok-beta",
+    model: "grok-2-1212",
   });
   const message = new HumanMessage("What is the current weather in Hawaii?");
   const res = await chat
@@ -144,7 +144,7 @@ test("stream with bound tools, yielding a single chunk", async () => {
 
 test("Few shotting with tool calls", async () => {
   const chat = new ChatXAI({
-    model: "grok-beta",
+    model: "grok-2-1212",
     temperature: 0,
   }).bind({
     tools: [
@@ -194,9 +194,9 @@ test("Few shotting with tool calls", async () => {
   expect(res.content).toContain("24");
 });
 
-test("Groq can stream tool calls", async () => {
+test("xAI can stream tool calls", async () => {
   const model = new ChatXAI({
-    model: "grok-beta",
+    model: "grok-2-1212",
     temperature: 0,
   });
 

libs/langchain-xai/src/tests/chat_models_structured_output.int.test.ts

@@ -7,7 +7,7 @@ import { ChatXAI } from "../chat_models.js";
 test("withStructuredOutput zod schema function calling", async () => {
   const model = new ChatXAI({
     temperature: 0,
-    model: "grok-beta",
+    model: "grok-2-1212",
   });
 
   const calculatorSchema = z.object({
@@ -37,7 +37,7 @@ test("withStructuredOutput zod schema function calling", async () => {
 test("withStructuredOutput zod schema JSON mode", async () => {
   const model = new ChatXAI({
     temperature: 0,
-    model: "grok-beta",
+    model: "grok-2-1212",
   });
 
   const calculatorSchema = z.object({
@@ -76,7 +76,7 @@ Respond with a JSON object containing three keys:
 test("withStructuredOutput JSON schema function calling", async () => {
   const model = new ChatXAI({
     temperature: 0,
-    model: "grok-beta",
+    model: "grok-2-1212",
   });
 
   const calculatorSchema = z.object({
@@ -106,7 +106,7 @@ test("withStructuredOutput JSON schema function calling", async () => {
 test("withStructuredOutput OpenAI function definition function calling", async () => {
   const model = new ChatXAI({
     temperature: 0,
-    model: "grok-beta",
+    model: "grok-2-1212",
   });
 
   const calculatorSchema = z.object({
@@ -120,14 +120,12 @@ test("withStructuredOutput OpenAI function definition function calling", async (
   });
 
   const prompt = ChatPromptTemplate.fromMessages([
-    "system",
-    `You are VERY bad at math and must always use a calculator.`,
-    "human",
-    "Please help me!! What is 2 + 2?",
+    ["system", `You are VERY bad at math and must always use a calculator.`],
+    ["human", "Please help me!! What is 2 + 2?"],
   ]);
   const chain = prompt.pipe(modelWithStructuredOutput);
   const result = await chain.invoke({});
-  // console.log(result);
+
   expect("operation" in result).toBe(true);
   expect("number1" in result).toBe(true);
   expect("number2" in result).toBe(true);
@@ -136,7 +134,7 @@ test("withStructuredOutput OpenAI function definition function calling", async (
 test("withStructuredOutput JSON schema JSON mode", async () => {
   const model = new ChatXAI({
     temperature: 0,
-    model: "grok-beta",
+    model: "grok-2-1212",
   });
 
   const calculatorSchema = z.object({
@@ -175,7 +173,7 @@ Respond with a JSON object containing three keys:
 test("withStructuredOutput JSON schema", async () => {
   const model = new ChatXAI({
     temperature: 0,
-    model: "grok-beta",
+    model: "grok-2-1212",
   });
 
   const jsonSchema = {
@@ -216,7 +214,7 @@ Respond with a JSON object containing three keys:
 test("withStructuredOutput includeRaw true", async () => {
   const model = new ChatXAI({
     temperature: 0,
-    model: "grok-beta",
+    model: "grok-2-1212",
   });
 
   const calculatorSchema = z.object({

yarn.lock

@@ -13335,7 +13335,7 @@ __metadata:
     rollup: ^4.5.2
     ts-jest: ^29.1.0
     typescript: <5.2.0
-    zod: ^3.22.4
+    zod: ^3.24.2
     zod-to-json-schema: ^3.23.1
   peerDependencies:
     "@langchain/core": ">=0.2.21 <0.4.0"