feat(llm): add Grok (xAI) integration (#179) (#236)

Co-authored-by: An Kaisen <51148505+ankaisen@users.noreply.github.com>

Author: david-mamani

docs/integrations/grok.md

@@ -0,0 +1,53 @@
+# Grok (xAI) Integration
+
+MemU supports **Grok**, the AI model from xAI, as a first-class LLM provider.
+
+## Prerequisites
+
+1.  **xAI Account:** You need an active account with [xAI](https://x.ai/).
+2.  **API Key:** Obtain an API key from the [xAI Console](https://console.x.ai/).
+
+## Configuration
+
+To enable Grok, you need to set the `XAI_API_KEY` environment variable.
+
+### Environment Variable
+
+```bash
+export XAI_API_KEY="your-xai-api-key-here"
+```
+
+PowerShell:
+
+```powershell
+$env:XAI_API_KEY="your-xai-api-key-here"
+```
+
+## Usage
+
+To use Grok as your LLM provider, switch the `provider` setting to `grok`. This can be done in your configuration file or when initializing the application.
+
+### Python Example
+
+```python
+from memu.app.settings import LLMConfig
+
+# Configure MemU to use Grok
+config = LLMConfig(
+    provider="grok",
+    # The default API key env var is XAI_API_KEY
+    # The default model is grok-2-latest
+)
+
+print(f"Using provider: {config.provider}")
+print(f"Base URL: {config.base_url}")
+print(f"Chat Model: {config.chat_model}")
+```
+
+## Models Supported
+
+We currently support the following Grok models:
+
+*   **grok-2-latest** (Default)
+
+The integration automatically sets the base URL to `https://api.x.ai/v1`.

docs/providers/grok.md

@@ -0,0 +1,66 @@
+# Grok (xAI) Provider
+
+memU includes first-class support for [Grok](https://grok.x.ai/), allowing you to leverage xAI's powerful language models directly within your application.
+
+## Prerequisites
+
+To use this provider, you must have an active xAI account.
+
+1.  Navigate to the [xAI Console](https://console.x.ai/).
+2.  Sign up or log in.
+3.  Create a new **API Key** in the API Keys section.
+
+## Configuration
+
+The integration is designed to work out-of-the-box with minimal configuration.
+
+### Environment Variables
+
+Set the following environment variable in your `.env` file or system environment:
+
+```bash
+GROK_API_KEY=xai-YOUR_API_KEY_HERE
+```
+
+### Defaults
+
+When you select the `grok` provider, memU automatically configures the following defaults:
+
+*   **Base URL**: `https://api.x.ai/v1`
+*   **Model**: `grok-2-latest`
+
+## Usage Example
+
+You can enable the Grok provider by setting the `provider` field to `"grok"` in your application configuration.
+
+### Using Python Configuration
+
+```python
+from memu.app.settings import LLMConfig
+from memu.app.service import MemoryService
+
+# Configure the LLM provider to use Grok
+llm_config = LLMConfig(provider="grok")
+
+# Initialize the service
+service = MemoryService(llm_config=llm_config)
+print(f"Service initialized with model: {llm_config.chat_model}")
+# Output: Service initialized with model: grok-2-latest
+```
+
+## Troubleshooting
+
+### Connection Issues
+If you are unable to connect to the xAI API:
+1.  Verify that your `GROK_API_KEY` is set correctly and has not expired.
+2.  Ensure that the `base_url` is resolving to `https://api.x.ai/v1`. If you have manual overrides in your settings, they might be conflicting with the default.
+
+### Model Availability
+If you receive a `404` or "Model not found" error, xAI may have updated their model names. You can override the model manually in the config if needed:
+
+```python
+config = LLMConfig(
+    provider="grok",
+    chat_model="grok-beta" # Example override
+)
+```

src/memu/app/settings.py

@@ -114,6 +114,18 @@ class LLMConfig(BaseModel):
         description="Maximum batch size for embedding API calls (used by SDK client backends).",
     )
 
+    @model_validator(mode="after")
+    def set_provider_defaults(self) -> "LLMConfig":
+        if self.provider == "grok":
+            # If values match the OpenAI defaults, switch them to Grok defaults
+            if self.base_url == "https://api.openai.com/v1":
+                self.base_url = "https://api.x.ai/v1"
+            if self.api_key == "OPENAI_API_KEY":
+                self.api_key = "XAI_API_KEY"
+            if self.chat_model == "gpt-4o-mini":
+                self.chat_model = "grok-2-latest"
+        return self
+
 
 class BlobConfig(BaseModel):
     provider: str = Field(default="local")

src/memu/llm/backends/__init__.py

@@ -1,6 +1,7 @@
 from memu.llm.backends.base import LLMBackend
 from memu.llm.backends.doubao import DoubaoLLMBackend
+from memu.llm.backends.grok import GrokBackend
 from memu.llm.backends.openai import OpenAILLMBackend
 from memu.llm.backends.openrouter import OpenRouterLLMBackend
 
-__all__ = ["DoubaoLLMBackend", "LLMBackend", "OpenAILLMBackend", "OpenRouterLLMBackend"]
+__all__ = ["DoubaoLLMBackend", "GrokBackend", "LLMBackend", "OpenAILLMBackend", "OpenRouterLLMBackend"]

src/memu/llm/backends/grok.py

@@ -0,0 +1,11 @@
+from __future__ import annotations
+
+from memu.llm.backends.openai import OpenAILLMBackend
+
+
+class GrokBackend(OpenAILLMBackend):
+    """Backend for Grok (xAI) LLM API."""
+
+    name = "grok"
+    # Grok uses the same payload structure as OpenAI
+    # We inherits build_summary_payload, parse_summary_response, etc.

src/memu/llm/http_client.py

@@ -10,6 +10,7 @@
 
 from memu.llm.backends.base import LLMBackend
 from memu.llm.backends.doubao import DoubaoLLMBackend
+from memu.llm.backends.grok import GrokBackend
 from memu.llm.backends.openai import OpenAILLMBackend
 from memu.llm.backends.openrouter import OpenRouterLLMBackend
 
@@ -66,6 +67,7 @@ def parse_embedding_response(self, data: dict[str, Any]) -> list[list[float]]:
 LLM_BACKENDS: dict[str, Callable[[], LLMBackend]] = {
     OpenAILLMBackend.name: OpenAILLMBackend,
     DoubaoLLMBackend.name: DoubaoLLMBackend,
+    GrokBackend.name: GrokBackend,
     OpenRouterLLMBackend.name: OpenRouterLLMBackend,
 }
 
@@ -244,6 +246,7 @@ def _load_embedding_backend(self, provider: str) -> _EmbeddingBackend:
         backends: dict[str, type[_EmbeddingBackend]] = {
             _OpenAIEmbeddingBackend.name: _OpenAIEmbeddingBackend,
             _DoubaoEmbeddingBackend.name: _DoubaoEmbeddingBackend,
+            "grok": _OpenAIEmbeddingBackend,
             _OpenRouterEmbeddingBackend.name: _OpenRouterEmbeddingBackend,
         }
         factory = backends.get(provider)

tests/llm/test_grok_provider.py

@@ -0,0 +1,46 @@
+import unittest
+from unittest.mock import patch
+
+from memu.app.settings import LLMConfig
+from memu.llm.backends.grok import GrokBackend
+from memu.llm.openai_sdk import OpenAISDKClient
+
+
+class TestGrokProvider(unittest.IsolatedAsyncioTestCase):
+    def test_settings_defaults(self):
+        """Test that setting provider='grok' sets the correct defaults."""
+        config = LLMConfig(provider="grok")
+        self.assertEqual(config.base_url, "https://api.x.ai/v1")
+        self.assertEqual(config.api_key, "XAI_API_KEY")
+        self.assertEqual(config.chat_model, "grok-2-latest")
+
+    @patch("memu.llm.openai_sdk.AsyncOpenAI")
+    async def test_client_initialization_with_grok_config(self, mock_async_openai):
+        """Test that OpenAISDKClient initializes with Grok base URL when configured."""
+        # Setup config
+        config = LLMConfig(provider="grok")
+
+        # Instantiate client with Grok config
+        # We simulate what the application factory would do: pass the config values
+        client = OpenAISDKClient(
+            base_url=config.base_url,
+            api_key="fake-key",  # In real app, this would be os.getenv(config.api_key)
+            chat_model=config.chat_model,
+            embed_model=config.embed_model,
+        )
+
+        # Assert AsyncOpenAI was called with the correct base_url
+        mock_async_openai.assert_called_with(api_key="fake-key", base_url="https://api.x.ai/v1")
+
+        # Verify client attributes
+        self.assertEqual(client.chat_model, "grok-2-latest")
+
+    def test_grok_backend_payload_parsing(self):
+        """Test that GrokBackend parses responses correctly (inherited from OpenAI)."""
+        backend = GrokBackend()
+
+        # Simulate a typical OpenAI-compatible response
+        dummy_response = {"choices": [{"message": {"content": "Grok response content", "role": "assistant"}}]}
+
+        result = backend.parse_summary_response(dummy_response)
+        self.assertEqual(result, "Grok response content")