feat(ts-sdk): add xAI (Grok) LLM provider to OSS SDK (#6115)

docs/components/llms/models/xAI.mdx

@@ -5,11 +5,12 @@ description: "Configure xAI Grok models as an LLM provider in Mem0 with API key
 
 [xAI](https://x.ai/) is a new AI company founded by Elon Musk that develops large language models, including Grok. Grok is trained on real-time data from X (formerly Twitter) and aims to provide accurate, up-to-date responses with a touch of wit and humor.
 
-In order to use LLMs from xAI, go to their [platform](https://console.x.ai) and get the API key. Set the API key as `XAI_API_KEY` environment variable to use the model as given below in the example.
+In order to use LLMs from xAI, go to their [platform](https://console.x.ai) and get the API key. Set the API key as `XAI_API_KEY` environment variable to use the model as given below in the example. You can also optionally set `XAI_API_BASE` to use a different API endpoint (defaults to `https://api.x.ai/v1`).
 
 ## Usage
 
-```python
+<CodeGroup>
+```python Python
 import os
 from mem0 import Memory
 
@@ -37,6 +38,31 @@ messages = [
 m.add(messages, user_id="alice", metadata={"category": "movies"})
 ```
 
+```typescript TypeScript
+import { Memory } from 'mem0ai/oss';
+
+const config = {
+  llm: {
+    provider: 'xai',
+    config: {
+      apiKey: process.env.XAI_API_KEY || '',
+      model: 'grok-4.3',
+      temperature: 0.1,
+      maxTokens: 2000,
+    },
+  },
+};
+const memory = new Memory(config);
+const messages = [
+    {"role": "user", "content": "I'm planning to watch a movie tonight. Any recommendations?"},
+    {"role": "assistant", "content": "How about thriller movies? They can be quite engaging."},
+    {"role": "user", "content": "I’m not a big fan of thriller movies but I love sci-fi movies."},
+    {"role": "assistant", "content": "Got it! I'll avoid thriller recommendations and suggest sci-fi movies in the future."}
+];
+await memory.add(messages, { userId: 'alice', metadata: { category: 'movies' } });
+```
+</CodeGroup>
+
 ## Config
 
 All available parameters for the `xai` config are present in [Master List of All Params in Config](../config).

mem0-ts/src/oss/src/llms/xai.ts

@@ -0,0 +1,50 @@
+import { OpenAILLM } from "./openai";
+import { LLMConfig, Message } from "../types";
+import { LLMResponse } from "./base";
+
+/**
+ * xAI (Grok) LLM provider.
+ *
+ * xAI's Grok API is OpenAI-compatible, so this simply reuses {@link OpenAILLM}
+ * and overrides the connection defaults — mirroring `mem0/llms/xai.py` in the
+ * Python SDK. The API key resolves from `config.apiKey` or the `XAI_API_KEY`
+ * env var, and the base URL from `config.baseURL`, `XAI_API_BASE`, else
+ * `https://api.x.ai/v1`.
+ */
+export class XAILLM extends OpenAILLM {
+  constructor(config: LLMConfig) {
+    const apiKey = config.apiKey || process.env.XAI_API_KEY;
+    if (!apiKey) {
+      throw new Error("xAI API key is required");
+    }
+    super({
+      ...config,
+      apiKey,
+      baseURL:
+        config.baseURL || process.env.XAI_API_BASE || "https://api.x.ai/v1",
+      model: config.model || "grok-4.3",
+    });
+  }
+
+  async generateResponse(
+    messages: Message[],
+    responseFormat?: { type: string },
+    tools?: any[],
+  ): Promise<string | LLMResponse> {
+    try {
+      return await super.generateResponse(messages, responseFormat, tools);
+    } catch (err) {
+      const message = err instanceof Error ? err.message : String(err);
+      throw new Error(`xAI LLM failed: ${message}`);
+    }
+  }
+
+  async generateChat(messages: Message[]): Promise<LLMResponse> {
+    try {
+      return await super.generateChat(messages);
+    } catch (err) {
+      const message = err instanceof Error ? err.message : String(err);
+      throw new Error(`xAI LLM failed: ${message}`);
+    }
+  }
+}

mem0-ts/src/oss/src/utils/factory.ts

@@ -24,6 +24,7 @@ import { ValkeyDB } from "../vector_stores/valkey";
 import { OllamaLLM } from "../llms/ollama";
 import { LMStudioLLM } from "../llms/lmstudio";
 import { DeepSeekLLM } from "../llms/deepseek";
+import { XAILLM } from "../llms/xai";
 import { LiteLLM } from "../llms/litellm";
 import { MiniMaxLLM } from "../llms/minimax";
 import { VllmLLM } from "../llms/vllm";
@@ -102,6 +103,8 @@ export class LLMFactory {
         return new LangchainLLM(config);
       case "deepseek":
         return new DeepSeekLLM(config);
+      case "xai":
+        return new XAILLM(config);
       case "litellm":
         return new LiteLLM(config);
       case "minimax":

mem0-ts/src/oss/tests/factory.unit.test.ts

@@ -102,6 +102,11 @@ jest.mock("../src/llms/deepseek", () => ({
     .fn()
     .mockImplementation((config) => ({ type: "deepseek-llm", config })),
 }));
+jest.mock("../src/llms/xai", () => ({
+  XAILLM: jest
+    .fn()
+    .mockImplementation((config) => ({ type: "xai-llm", config })),
+}));
 jest.mock("../src/llms/litellm", () => ({
   LiteLLM: jest
     .fn()
@@ -258,6 +263,7 @@ describe("LLMFactory", () => {
     ["langchain"],
     ["lmstudio"],
     ["deepseek"],
+    ["xai"],
     ["litellm"],
     ["minimax"],
     ["vllm"],

mem0-ts/src/oss/tests/xai.test.ts

@@ -0,0 +1,153 @@
+/// <reference types="jest" />
+/**
+ * xAI (Grok) LLM — unit tests (mocked OpenAI).
+ */
+
+import { XAILLM } from "../src/llms/xai";
+
+const mockCreate = jest.fn();
+const mockOpenAICtor = jest.fn();
+
+jest.mock("openai", () => {
+  return jest.fn().mockImplementation((config) => {
+    mockOpenAICtor(config);
+    return { chat: { completions: { create: mockCreate } } };
+  });
+});
+
+describe("XAILLM (unit)", () => {
+  const ORIGINAL_ENV = process.env;
+
+  beforeEach(() => {
+    jest.clearAllMocks();
+    process.env = { ...ORIGINAL_ENV };
+    delete process.env.XAI_API_KEY;
+    delete process.env.XAI_API_BASE;
+    mockCreate.mockResolvedValue({
+      choices: [{ message: { content: "hi", role: "assistant" } }],
+    });
+  });
+
+  afterAll(() => {
+    process.env = ORIGINAL_ENV;
+  });
+
+  it("defaults to grok-4.3 and the xAI base URL (matching the Python provider)", async () => {
+    const llm = new XAILLM({ apiKey: "test-key" });
+    const result = await llm.generateResponse([
+      { role: "user", content: "hello" },
+    ]);
+
+    expect(mockOpenAICtor).toHaveBeenCalledWith(
+      expect.objectContaining({
+        apiKey: "test-key",
+        baseURL: "https://api.x.ai/v1",
+      }),
+    );
+    expect(mockCreate).toHaveBeenCalledWith(
+      expect.objectContaining({ model: "grok-4.3" }),
+    );
+    expect(result).toBe("hi");
+  });
+
+  it("resolves XAI_API_KEY / XAI_API_BASE from the environment", () => {
+    process.env.XAI_API_KEY = "env-key";
+    process.env.XAI_API_BASE = "https://custom.x.ai/v1";
+
+    new XAILLM({});
+
+    expect(mockOpenAICtor).toHaveBeenCalledWith(
+      expect.objectContaining({
+        apiKey: "env-key",
+        baseURL: "https://custom.x.ai/v1",
+      }),
+    );
+  });
+
+  it("prefers explicit config over defaults and the environment", async () => {
+    process.env.XAI_API_KEY = "env-key";
+
+    const llm = new XAILLM({
+      apiKey: "explicit-key",
+      baseURL: "https://proxy.example.com/v1",
+      model: "grok-420-reasoning",
+    });
+    await llm.generateResponse([{ role: "user", content: "hello" }]);
+
+    expect(mockOpenAICtor).toHaveBeenCalledWith(
+      expect.objectContaining({
+        apiKey: "explicit-key",
+        baseURL: "https://proxy.example.com/v1",
+      }),
+    );
+    expect(mockCreate).toHaveBeenCalledWith(
+      expect.objectContaining({ model: "grok-420-reasoning" }),
+    );
+  });
+
+  it("throws when no API key is provided", () => {
+    expect(() => new XAILLM({})).toThrow("xAI API key is required");
+  });
+
+  it("generateResponse() handles tool calls", async () => {
+    mockCreate.mockResolvedValueOnce({
+      choices: [
+        {
+          message: {
+            content: "",
+            role: "assistant",
+            tool_calls: [
+              {
+                function: { name: "get_weather", arguments: '{"city": "SF"}' },
+              },
+            ],
+          },
+        },
+      ],
+    });
+
+    const llm = new XAILLM({ apiKey: "test-key" });
+    const result = await llm.generateResponse(
+      [{ role: "user", content: "weather?" }],
+      undefined,
+      [{ type: "function", function: { name: "get_weather" } }],
+    );
+
+    expect(result).toEqual({
+      content: "",
+      role: "assistant",
+      toolCalls: [{ name: "get_weather", arguments: '{"city": "SF"}' }],
+    });
+  });
+
+  it("generateResponse() wraps downstream errors with an xAI-specific message", async () => {
+    mockCreate.mockRejectedValueOnce(new Error("Connection refused"));
+    const llm = new XAILLM({ apiKey: "test-key" });
+
+    await expect(
+      llm.generateResponse([{ role: "user", content: "hi" }]),
+    ).rejects.toThrow("xAI LLM failed: Connection refused");
+  });
+
+  it("generateChat() returns the LLMResponse shape", async () => {
+    mockCreate.mockResolvedValueOnce({
+      choices: [{ message: { content: "I can help.", role: "assistant" } }],
+    });
+
+    const llm = new XAILLM({ apiKey: "test-key" });
+    const result = await llm.generateChat([
+      { role: "user", content: "help me" },
+    ]);
+
+    expect(result).toEqual({ content: "I can help.", role: "assistant" });
+  });
+
+  it("generateChat() wraps downstream errors with an xAI-specific message", async () => {
+    mockCreate.mockRejectedValueOnce(new Error("Timeout"));
+    const llm = new XAILLM({ apiKey: "test-key" });
+
+    await expect(
+      llm.generateChat([{ role: "user", content: "hi" }]),
+    ).rejects.toThrow("xAI LLM failed: Timeout");
+  });
+});

mem0/llms/xai.py

@@ -34,7 +34,7 @@ class XAILLM(LLMBase):
         super().__init__(config)
 
         if not self.config.model:
-            self.config.model = "grok-2-latest"
+            self.config.model = "grok-4.3"
 
         api_key = self.config.api_key or os.getenv("XAI_API_KEY")
         base_url = self.config.xai_base_url or os.getenv("XAI_API_BASE") or "https://api.x.ai/v1"

tests/llms/test_xai.py

@@ -19,19 +19,19 @@ def mock_xai_client():
 
 def test_xai_llm_base_url():
     # case1: default
-    config = XAIConfig(model="grok-2-latest", api_key="api_key")
+    config = XAIConfig(model="grok-4.3", api_key="api_key")
     llm = XAILLM(config)
     assert str(llm.client.base_url) == "https://api.x.ai/v1/"
 
     # case2: XAI_API_BASE env var
     os.environ["XAI_API_BASE"] = "https://api.provider.com/v1"
-    config = XAIConfig(model="grok-2-latest", api_key="api_key")
+    config = XAIConfig(model="grok-4.3", api_key="api_key")
     llm = XAILLM(config)
     assert str(llm.client.base_url) == "https://api.provider.com/v1/"
 
     # case3: config.xai_base_url wins over env
     config = XAIConfig(
-        model="grok-2-latest",
+        model="grok-4.3",
         api_key="api_key",
         xai_base_url="https://api.config.com/v1",
     )
@@ -39,16 +39,22 @@ def test_xai_llm_base_url():
     assert str(llm.client.base_url) == "https://api.config.com/v1/"
 
 
+def test_xai_defaults_to_current_grok_model():
+    # No model supplied -> provider falls back to the current default (grok-4.3).
+    llm = XAILLM(XAIConfig(api_key="k"))
+    assert llm.config.model == "grok-4.3"
+
+
 def test_xai_accepts_base_llm_config():
     # Used to AttributeError on self.config.xai_base_url because the factory
     # wired XAI with plain BaseLlmConfig.
-    llm = XAILLM(BaseLlmConfig(model="grok-2-latest", api_key="k"))
+    llm = XAILLM(BaseLlmConfig(model="grok-4.3", api_key="k"))
     assert isinstance(llm.config, XAIConfig)
     assert llm.config.xai_base_url is None
 
 
 def test_generate_response_without_tools(mock_xai_client):
-    config = XAIConfig(model="grok-2-latest", temperature=0.7, max_tokens=100, top_p=1.0, api_key="api_key")
+    config = XAIConfig(model="grok-4.3", temperature=0.7, max_tokens=100, top_p=1.0, api_key="api_key")
     llm = XAILLM(config)
     messages = [
         {"role": "system", "content": "You are a helpful assistant."},
@@ -62,7 +68,7 @@ def test_generate_response_without_tools(mock_xai_client):
     response = llm.generate_response(messages)
 
     mock_xai_client.chat.completions.create.assert_called_once_with(
-        model="grok-2-latest",
+        model="grok-4.3",
         messages=messages,
         temperature=0.7,
         max_tokens=100,
@@ -72,7 +78,7 @@ def test_generate_response_without_tools(mock_xai_client):
 
 
 def test_generate_response_with_tools(mock_xai_client):
-    config = XAIConfig(model="grok-2-latest", temperature=0.7, max_tokens=100, top_p=1.0, api_key="api_key")
+    config = XAIConfig(model="grok-4.3", temperature=0.7, max_tokens=100, top_p=1.0, api_key="api_key")
     llm = XAILLM(config)
     messages = [
         {"role": "system", "content": "You are a helpful assistant."},
@@ -106,7 +112,7 @@ def test_generate_response_with_tools(mock_xai_client):
     response = llm.generate_response(messages, tools=tools)
 
     mock_xai_client.chat.completions.create.assert_called_once_with(
-        model="grok-2-latest",
+        model="grok-4.3",
         messages=messages,
         temperature=0.7,
         max_tokens=100,
@@ -121,7 +127,7 @@ def test_generate_response_with_tools(mock_xai_client):
 
 def test_empty_tools_list_not_forwarded(mock_xai_client):
     # tools=[] would otherwise get rejected by some OpenAI-compatible backends
-    config = XAIConfig(model="grok-2-latest", api_key="api_key")
+    config = XAIConfig(model="grok-4.3", api_key="api_key")
     llm = XAILLM(config)
 
     mock_response = Mock()
@@ -138,7 +144,7 @@ def test_empty_tools_list_not_forwarded(mock_xai_client):
 
 def test_tools_requested_but_model_returns_no_calls(mock_xai_client):
     # Model can decline to call any tool even when offered
-    config = XAIConfig(model="grok-2-latest", api_key="api_key")
+    config = XAIConfig(model="grok-4.3", api_key="api_key")
     llm = XAILLM(config)
     tools = [{"type": "function", "function": {"name": "f", "parameters": {"type": "object", "properties": {}}}}]
 
@@ -151,7 +157,7 @@ def test_tools_requested_but_model_returns_no_calls(mock_xai_client):
 
 
 def test_generate_response_with_response_format(mock_xai_client):
-    config = XAIConfig(model="grok-2-latest", temperature=0.7, max_tokens=100, top_p=1.0, api_key="api_key")
+    config = XAIConfig(model="grok-4.3", temperature=0.7, max_tokens=100, top_p=1.0, api_key="api_key")
     llm = XAILLM(config)
     messages = [
         {"role": "system", "content": "You are a memory extraction assistant."},
@@ -165,7 +171,7 @@ def test_generate_response_with_response_format(mock_xai_client):
     response = llm.generate_response(messages, response_format={"type": "json_object"})
 
     mock_xai_client.chat.completions.create.assert_called_once_with(
-        model="grok-2-latest",
+        model="grok-4.3",
         messages=messages,
         temperature=0.7,
         max_tokens=100,
@@ -176,7 +182,7 @@ def test_generate_response_with_response_format(mock_xai_client):
 
 
 def test_generate_response_without_response_format(mock_xai_client):
-    config = XAIConfig(model="grok-2-latest", api_key="api_key")
+    config = XAIConfig(model="grok-4.3", api_key="api_key")
     llm = XAILLM(config)
 
     mock_response = Mock()
@@ -194,7 +200,7 @@ def test_factory_creates_xai_from_dict():
         mock_openai.return_value = Mock()
         llm = LlmFactory.create(
             "xai",
-            {"model": "grok-2-latest", "api_key": "k", "xai_base_url": "https://example.com/v1"},
+            {"model": "grok-4.3", "api_key": "k", "xai_base_url": "https://example.com/v1"},
         )
     assert isinstance(llm, XAILLM)
     assert isinstance(llm.config, XAIConfig)
@@ -205,5 +211,5 @@ def test_factory_creates_xai_from_base_config():
     # Legacy callers still hand the factory a plain BaseLlmConfig
     with patch("mem0.llms.xai.OpenAI") as mock_openai:
         mock_openai.return_value = Mock()
-        llm = LlmFactory.create("xai", BaseLlmConfig(model="grok-2-latest", api_key="k"))
+        llm = LlmFactory.create("xai", BaseLlmConfig(model="grok-4.3", api_key="k"))
     assert isinstance(llm.config, XAIConfig)