diff --git a/docs/components/llms/models/xAI.mdx b/docs/components/llms/models/xAI.mdx index b75b4530c..812a4e650 100644 --- a/docs/components/llms/models/xAI.mdx +++ b/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 + +```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' } }); +``` + + ## Config All available parameters for the `xai` config are present in [Master List of All Params in Config](../config). \ No newline at end of file diff --git a/mem0-ts/src/oss/src/llms/xai.ts b/mem0-ts/src/oss/src/llms/xai.ts new file mode 100644 index 000000000..a365a6a3a --- /dev/null +++ b/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 { + 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 { + try { + return await super.generateChat(messages); + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + throw new Error(`xAI LLM failed: ${message}`); + } + } +} diff --git a/mem0-ts/src/oss/src/utils/factory.ts b/mem0-ts/src/oss/src/utils/factory.ts index 1a9526f74..d4e5f4f61 100644 --- a/mem0-ts/src/oss/src/utils/factory.ts +++ b/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": diff --git a/mem0-ts/src/oss/tests/factory.unit.test.ts b/mem0-ts/src/oss/tests/factory.unit.test.ts index 0b4cdeeb1..490ca43f9 100644 --- a/mem0-ts/src/oss/tests/factory.unit.test.ts +++ b/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"], diff --git a/mem0-ts/src/oss/tests/xai.test.ts b/mem0-ts/src/oss/tests/xai.test.ts new file mode 100644 index 000000000..eb8ac09f6 --- /dev/null +++ b/mem0-ts/src/oss/tests/xai.test.ts @@ -0,0 +1,153 @@ +/// +/** + * 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"); + }); +}); diff --git a/mem0/llms/xai.py b/mem0/llms/xai.py index 8c38e1d07..57a7a3efc 100644 --- a/mem0/llms/xai.py +++ b/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" diff --git a/tests/llms/test_xai.py b/tests/llms/test_xai.py index 5143064ec..8c5e6c12f 100644 --- a/tests/llms/test_xai.py +++ b/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)