From 8efd95f5766316e72206e94f9be8327e2d6e0cd9 Mon Sep 17 00:00:00 2001
From: PR Bot <pr-bot@minimaxi.com>
Date: Tue, 24 Mar 2026 15:36:08 +0800
Subject: [PATCH] Add configurable LLM provider for prompt rewriter with
 MiniMax support
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The prompt rewriter previously used a hardcoded Azure OpenAI endpoint
(ByteDance internal), making it unusable for external users. This adds a
configurable LLM provider system (tools/llm_provider.py) that supports
OpenAI, Azure OpenAI, MiniMax, and any OpenAI-compatible API via
environment variables. Provider auto-detection, temperature clamping for
MiniMax, and think-tag stripping are included.

- New: tools/llm_provider.py — configurable LLM provider factory
- Updated: tools/prompt_rewriter.py — uses new provider system
- Updated: README.md — documents LLM provider configuration
- New: 22 unit tests + 3 integration tests
---
 README.md                              |  28 ++-
 tests/conftest.py                      |  31 +++
 tests/test_llm_provider.py             | 270 +++++++++++++++++++++++++
 tests/test_llm_provider_integration.py |  74 +++++++
 tools/llm_provider.py                  | 169 ++++++++++++++++
 tools/prompt_rewriter.py               |  66 +++---
 6 files changed, 601 insertions(+), 37 deletions(-)
 create mode 100644 tests/conftest.py
 create mode 100644 tests/test_llm_provider.py
 create mode 100644 tests/test_llm_provider_integration.py
 create mode 100644 tools/llm_provider.py

diff --git a/README.md b/README.md
index 8d5df0e9..03f8206a 100644
--- a/README.md
+++ b/README.md
@@ -100,7 +100,7 @@ We provide Infinity models for you to play with, which are on <a href='https://h
 |  Infinity-8B   |  512  |  -  |    -    |  -   | [infinity_8b_512x512_weights](https://huggingface.co/FoundationVision/Infinity/tree/main/infinity_8b_512x512_weights) |
 |  Infinity-20B   |  1024  |  -  |    -    |  -   | [Coming Soon](TBD) |
 
-${\dagger}$ result is tested with a [prompt rewriter](tools/prompt_rewriter.py). 
+${\dagger}$ result is tested with a [prompt rewriter](tools/prompt_rewriter.py).
 
 You can load these models to generate images via the codes in [interactive_infer.ipynb](tools/interactive_infer.ipynb) and [interactive_infer_8b.ipynb](tools/interactive_infer_8b.ipynb) .
 
@@ -170,6 +170,32 @@ You can monitor the training process by checking the logs in `local_output/log.t
 
 If your experiment is interrupted, just rerun the command, and the training will **automatically resume** from the last checkpoint in `local_output/ckpt*.pth`.
 
+## 🤖 LLM Provider Configuration (Prompt Rewriter)
+
+The [prompt rewriter](tools/prompt_rewriter.py) uses an LLM to refine short prompts into detailed image descriptions. It supports multiple LLM providers via environment variables:
+
+| Provider | Env Vars | Default Model |
+|----------|----------|---------------|
+| OpenAI | `OPENAI_API_KEY` | gpt-4o |
+| Azure OpenAI | `AZURE_OPENAI_API_KEY`, `AZURE_OPENAI_ENDPOINT` | gpt-4o |
+| [MiniMax](https://www.minimax.io/) | `MINIMAX_API_KEY` | MiniMax-M2.7 |
+| Custom (OpenAI-compatible) | `LLM_API_KEY`, `LLM_BASE_URL` | gpt-4o |
+
+**Quick start:**
+```bash
+# Use OpenAI
+export OPENAI_API_KEY=sk-...
+
+# Or use MiniMax
+export MINIMAX_API_KEY=...
+
+# Or use any OpenAI-compatible API
+export LLM_BASE_URL=http://localhost:8000/v1
+export LLM_API_KEY=...
+```
+
+The provider is auto-detected from available environment variables. You can also set `LLM_PROVIDER` explicitly (`openai`, `azure`, `minimax`) and override the model with `LLM_MODEL`. See [tools/llm_provider.py](tools/llm_provider.py) for full documentation.
+
 ## 🍭 Evaluation
 We provide [eval.sh](scripts/eval.sh) for evaluation on various benchmarks with only one command. In particular, [eval.sh](scripts/eval.sh) supports evaluation on commonly used metrics such as [GenEval](https://github.com/djghosh13/geneval), [ImageReward](https://github.com/THUDM/ImageReward), [HPSv2.1](https://github.com/tgxs002/HPSv2), FID and Validation Loss. Please refer to [evaluation/README.md](evaluation/README.md) for more details.
 ```shell
diff --git a/tests/conftest.py b/tests/conftest.py
new file mode 100644
index 00000000..e1a94b84
--- /dev/null
+++ b/tests/conftest.py
@@ -0,0 +1,31 @@
+"""Configure sys.path for Infinity project tests.
+
+The cai-framework package installs a .pth file that adds cai/tools/ to
+sys.path. To avoid import conflicts, we insert the Infinity project root
+at position 0 and invalidate any cached 'tools' module.
+"""
+
+import importlib
+import os
+import sys
+
+_project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+
+# Ensure project root is first in sys.path
+if sys.path[0] != _project_root:
+    while _project_root in sys.path:
+        sys.path.remove(_project_root)
+    sys.path.insert(0, _project_root)
+
+# Evict stale 'tools' module from cai-framework so our tools/ wins
+for key in list(sys.modules):
+    if key == "tools" or key.startswith("tools."):
+        del sys.modules[key]
+
+# Force-create a proper package reference for our tools/ directory
+import types
+tools_pkg = types.ModuleType("tools")
+tools_pkg.__path__ = [os.path.join(_project_root, "tools")]
+tools_pkg.__file__ = os.path.join(_project_root, "tools", "__init__.py")
+tools_pkg.__package__ = "tools"
+sys.modules["tools"] = tools_pkg
diff --git a/tests/test_llm_provider.py b/tests/test_llm_provider.py
new file mode 100644
index 00000000..f0c9debe
--- /dev/null
+++ b/tests/test_llm_provider.py
@@ -0,0 +1,270 @@
+"""Tests for tools/llm_provider.py — configurable LLM provider."""
+
+import os
+import unittest
+from unittest.mock import patch, MagicMock
+
+import openai
+
+
+class TestDetectProvider(unittest.TestCase):
+    """Test automatic provider detection from environment variables."""
+
+    def _detect(self):
+        from tools.llm_provider import _detect_provider
+        return _detect_provider()
+
+    @patch.dict(os.environ, {"MINIMAX_API_KEY": "mm-key"}, clear=False)
+    def test_detect_minimax(self):
+        self.assertEqual(self._detect(), "minimax")
+
+    @patch.dict(os.environ, {"AZURE_OPENAI_API_KEY": "az-key"}, clear=False)
+    def test_detect_azure(self):
+        # Remove MINIMAX_API_KEY if present to isolate test
+        env = os.environ.copy()
+        env.pop("MINIMAX_API_KEY", None)
+        with patch.dict(os.environ, env, clear=True):
+            os.environ["AZURE_OPENAI_API_KEY"] = "az-key"
+            self.assertEqual(self._detect(), "azure")
+
+    @patch.dict(os.environ, {}, clear=True)
+    def test_detect_default_openai(self):
+        self.assertEqual(self._detect(), "openai")
+
+
+class TestGetApiKey(unittest.TestCase):
+    """Test API key resolution with fallback priorities."""
+
+    def _get_key(self, provider):
+        from tools.llm_provider import _get_api_key
+        return _get_api_key(provider)
+
+    @patch.dict(os.environ, {"OPENAI_API_KEY": "sk-openai"}, clear=True)
+    def test_openai_key(self):
+        self.assertEqual(self._get_key("openai"), "sk-openai")
+
+    @patch.dict(os.environ, {"MINIMAX_API_KEY": "mm-key"}, clear=True)
+    def test_minimax_key(self):
+        self.assertEqual(self._get_key("minimax"), "mm-key")
+
+    @patch.dict(os.environ, {"LLM_API_KEY": "generic-key"}, clear=True)
+    def test_generic_key_fallback(self):
+        self.assertEqual(self._get_key("openai"), "generic-key")
+
+    @patch.dict(os.environ, {"LLM_API_KEY": "generic", "OPENAI_API_KEY": "specific"}, clear=True)
+    def test_specific_overrides_generic(self):
+        self.assertEqual(self._get_key("openai"), "specific")
+
+
+class TestGetLlmClient(unittest.TestCase):
+    """Test client creation for different providers."""
+
+    @patch.dict(os.environ, {"OPENAI_API_KEY": "sk-test"}, clear=True)
+    def test_openai_client(self):
+        from tools.llm_provider import get_llm_client
+        client, model = get_llm_client(provider="openai")
+        self.assertIsInstance(client, openai.OpenAI)
+        self.assertEqual(model, "gpt-4o")
+
+    @patch.dict(os.environ, {"MINIMAX_API_KEY": "mm-test"}, clear=True)
+    def test_minimax_client(self):
+        from tools.llm_provider import get_llm_client
+        client, model = get_llm_client(provider="minimax")
+        self.assertIsInstance(client, openai.OpenAI)
+        self.assertEqual(model, "MiniMax-M2.7")
+        self.assertIn("minimax", client.base_url.host)
+
+    @patch.dict(os.environ, {"AZURE_OPENAI_API_KEY": "az-test"}, clear=True)
+    def test_azure_client(self):
+        from tools.llm_provider import get_llm_client
+        client, model = get_llm_client(provider="azure")
+        self.assertIsInstance(client, openai.AzureOpenAI)
+        self.assertEqual(model, "gpt-4o")
+
+    @patch.dict(os.environ, {"LLM_API_KEY": "custom-key"}, clear=True)
+    def test_custom_base_url(self):
+        from tools.llm_provider import get_llm_client
+        client, model = get_llm_client(
+            provider="openai",
+            base_url="http://localhost:8000/v1",
+        )
+        self.assertIsInstance(client, openai.OpenAI)
+        self.assertIn("localhost", str(client.base_url))
+
+    @patch.dict(os.environ, {"LLM_MODEL": "my-model", "OPENAI_API_KEY": "k"}, clear=True)
+    def test_env_model_override(self):
+        from tools.llm_provider import get_llm_client
+        _, model = get_llm_client(provider="openai")
+        self.assertEqual(model, "my-model")
+
+    def test_explicit_params_override(self):
+        from tools.llm_provider import get_llm_client
+        client, model = get_llm_client(
+            provider="minimax",
+            api_key="explicit-key",
+            model="custom-model",
+        )
+        self.assertEqual(model, "custom-model")
+
+
+class TestChatCompletion(unittest.TestCase):
+    """Test chat_completion wrapper."""
+
+    @patch("tools.llm_provider.get_llm_client")
+    def test_basic_completion(self, mock_get_client):
+        mock_client = MagicMock()
+        mock_choice = MagicMock()
+        mock_choice.message.content = "refined prompt"
+        mock_client.chat.completions.create.return_value = MagicMock(
+            choices=[mock_choice]
+        )
+        mock_get_client.return_value = (mock_client, "gpt-4o")
+
+        from tools.llm_provider import chat_completion
+        result = chat_completion(
+            messages=[{"role": "user", "content": "a tree"}],
+            provider="openai",
+            api_key="test-key",
+        )
+        self.assertEqual(result, "refined prompt")
+        mock_client.chat.completions.create.assert_called_once()
+
+    @patch("tools.llm_provider.get_llm_client")
+    def test_json_mode(self, mock_get_client):
+        mock_client = MagicMock()
+        mock_choice = MagicMock()
+        mock_choice.message.content = '{"key": "value"}'
+        mock_client.chat.completions.create.return_value = MagicMock(
+            choices=[mock_choice]
+        )
+        mock_get_client.return_value = (mock_client, "gpt-4o")
+
+        from tools.llm_provider import chat_completion
+        result = chat_completion(
+            messages=[{"role": "user", "content": "test"}],
+            provider="openai",
+            api_key="test-key",
+            return_json=True,
+        )
+        self.assertEqual(result, '{"key": "value"}')
+        call_kwargs = mock_client.chat.completions.create.call_args[1]
+        self.assertEqual(call_kwargs["response_format"], {"type": "json_object"})
+
+    @patch("tools.llm_provider.get_llm_client")
+    @patch.dict(os.environ, {"LLM_PROVIDER": "minimax"}, clear=True)
+    def test_minimax_think_tag_stripping(self, mock_get_client):
+        mock_client = MagicMock()
+        mock_choice = MagicMock()
+        mock_choice.message.content = "<think>internal reasoning</think>\nactual response"
+        mock_client.chat.completions.create.return_value = MagicMock(
+            choices=[mock_choice]
+        )
+        mock_get_client.return_value = (mock_client, "MiniMax-M2.7")
+
+        from tools.llm_provider import chat_completion
+        result = chat_completion(
+            messages=[{"role": "user", "content": "test"}],
+            provider="minimax",
+            api_key="test-key",
+        )
+        self.assertEqual(result, "actual response")
+
+    @patch("tools.llm_provider.get_llm_client")
+    @patch.dict(os.environ, {"LLM_PROVIDER": "minimax"}, clear=True)
+    def test_minimax_temperature_clamping(self, mock_get_client):
+        mock_client = MagicMock()
+        mock_choice = MagicMock()
+        mock_choice.message.content = "result"
+        mock_client.chat.completions.create.return_value = MagicMock(
+            choices=[mock_choice]
+        )
+        mock_get_client.return_value = (mock_client, "MiniMax-M2.7")
+
+        from tools.llm_provider import chat_completion
+        chat_completion(
+            messages=[{"role": "user", "content": "test"}],
+            provider="minimax",
+            api_key="test-key",
+            temperature=0.0,
+        )
+        call_kwargs = mock_client.chat.completions.create.call_args[1]
+        # temperature=0.0 should be clamped to 0.01 for MiniMax
+        self.assertGreaterEqual(call_kwargs["temperature"], 0.01)
+
+    @patch("tools.llm_provider.get_llm_client")
+    def test_openai_no_temperature_clamping(self, mock_get_client):
+        mock_client = MagicMock()
+        mock_choice = MagicMock()
+        mock_choice.message.content = "result"
+        mock_client.chat.completions.create.return_value = MagicMock(
+            choices=[mock_choice]
+        )
+        mock_get_client.return_value = (mock_client, "gpt-4o")
+
+        from tools.llm_provider import chat_completion
+        chat_completion(
+            messages=[{"role": "user", "content": "test"}],
+            provider="openai",
+            api_key="test-key",
+            temperature=0.0,
+        )
+        call_kwargs = mock_client.chat.completions.create.call_args[1]
+        self.assertEqual(call_kwargs["temperature"], 0.0)
+
+
+class TestPromptRewriter(unittest.TestCase):
+    """Test the PromptRewriter integration with llm_provider."""
+
+    @patch("tools.llm_provider.get_llm_client")
+    def test_rewriter_uses_llm_provider(self, mock_get_client):
+        mock_client = MagicMock()
+        mock_choice = MagicMock()
+        mock_choice.message.content = "<prompt:A detailed tree in sunlight><cfg:3>"
+        mock_client.chat.completions.create.return_value = MagicMock(
+            choices=[mock_choice]
+        )
+        mock_get_client.return_value = (mock_client, "gpt-4o")
+
+        from tools.prompt_rewriter import PromptRewriter
+        rewriter = PromptRewriter(system="", few_shot_history=[])
+        result = rewriter.rewrite("a tree")
+        self.assertIn("<prompt:", result)
+        self.assertIn("<cfg:", result)
+
+    @patch("tools.llm_provider.get_llm_client")
+    def test_get_gpt_result_backward_compat(self, mock_get_client):
+        mock_client = MagicMock()
+        mock_choice = MagicMock()
+        mock_choice.message.content = "response"
+        mock_client.chat.completions.create.return_value = MagicMock(
+            choices=[mock_choice]
+        )
+        mock_get_client.return_value = (mock_client, "gpt-4o")
+
+        from tools.prompt_rewriter import get_gpt_result
+        result, err = get_gpt_result(
+            messages=[{"role": "user", "content": "test"}],
+            retry=1,
+        )
+        self.assertEqual(result, "response")
+        self.assertIsNone(err)
+
+
+class TestDefaultModels(unittest.TestCase):
+    """Test default model selection per provider."""
+
+    def test_default_models_defined(self):
+        from tools.llm_provider import _DEFAULT_MODELS
+        self.assertIn("openai", _DEFAULT_MODELS)
+        self.assertIn("minimax", _DEFAULT_MODELS)
+        self.assertIn("azure", _DEFAULT_MODELS)
+        self.assertEqual(_DEFAULT_MODELS["minimax"], "MiniMax-M2.7")
+
+    def test_default_base_urls(self):
+        from tools.llm_provider import _DEFAULT_BASE_URLS
+        self.assertIn("minimax", _DEFAULT_BASE_URLS)
+        self.assertIn("minimax.io", _DEFAULT_BASE_URLS["minimax"])
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/test_llm_provider_integration.py b/tests/test_llm_provider_integration.py
new file mode 100644
index 00000000..72bebc6e
--- /dev/null
+++ b/tests/test_llm_provider_integration.py
@@ -0,0 +1,74 @@
+"""Integration tests for tools/llm_provider.py — require live API keys."""
+
+import os
+import unittest
+
+# Skip all tests if no API key is available
+MINIMAX_API_KEY = os.environ.get("MINIMAX_API_KEY", "")
+OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "")
+
+MESSAGES = [
+    {"role": "system", "content": "You refine image prompts. Output only the refined prompt."},
+    {"role": "user", "content": "a tree"},
+]
+
+
+@unittest.skipUnless(MINIMAX_API_KEY, "MINIMAX_API_KEY not set")
+class TestMiniMaxIntegration(unittest.TestCase):
+    """Live integration tests against MiniMax API."""
+
+    def test_chat_completion(self):
+        from tools.llm_provider import chat_completion
+        result = chat_completion(
+            messages=MESSAGES,
+            provider="minimax",
+            api_key=MINIMAX_API_KEY,
+        )
+        self.assertIsInstance(result, str)
+        self.assertGreater(len(result), 10)
+
+    def test_chat_completion_with_model(self):
+        from tools.llm_provider import chat_completion
+        result = chat_completion(
+            messages=MESSAGES,
+            provider="minimax",
+            api_key=MINIMAX_API_KEY,
+            model="MiniMax-M2.5-highspeed",
+        )
+        self.assertIsInstance(result, str)
+        self.assertGreater(len(result), 5)
+
+    def test_prompt_rewriter(self):
+        from tools.llm_provider import chat_completion
+        from tools.prompt_rewriter import SYSTEM, FEW_SHOT_HISTORY
+        messages = (
+            [{"role": "system", "content": SYSTEM}]
+            + FEW_SHOT_HISTORY
+            + [{"role": "user", "content": "a sunset over the ocean"}]
+        )
+        result = chat_completion(
+            messages=messages,
+            provider="minimax",
+            api_key=MINIMAX_API_KEY,
+        )
+        self.assertIsInstance(result, str)
+        self.assertGreater(len(result), 10)
+
+
+@unittest.skipUnless(OPENAI_API_KEY, "OPENAI_API_KEY not set")
+class TestOpenAIIntegration(unittest.TestCase):
+    """Live integration tests against OpenAI API."""
+
+    def test_chat_completion(self):
+        from tools.llm_provider import chat_completion
+        result = chat_completion(
+            messages=MESSAGES,
+            provider="openai",
+            api_key=OPENAI_API_KEY,
+        )
+        self.assertIsInstance(result, str)
+        self.assertGreater(len(result), 10)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tools/llm_provider.py b/tools/llm_provider.py
new file mode 100644
index 00000000..bf0fe86d
--- /dev/null
+++ b/tools/llm_provider.py
@@ -0,0 +1,169 @@
+"""
+Configurable LLM provider for prompt rewriting and other LLM-powered tools.
+
+Supports multiple providers via environment variables:
+  - LLM_PROVIDER: "openai" (default), "azure", "minimax"
+  - LLM_API_KEY: API key for the selected provider
+  - LLM_BASE_URL: Custom base URL (optional, auto-detected per provider)
+  - LLM_MODEL: Model name (optional, defaults per provider)
+
+Provider-specific env vars (take precedence over generic ones):
+  - OPENAI_API_KEY / OPENAI_BASE_URL
+  - AZURE_OPENAI_API_KEY / AZURE_OPENAI_ENDPOINT / AZURE_API_VERSION
+  - MINIMAX_API_KEY
+
+Examples:
+  # Standard OpenAI
+  export LLM_PROVIDER=openai
+  export OPENAI_API_KEY=sk-...
+
+  # MiniMax (OpenAI-compatible)
+  export LLM_PROVIDER=minimax
+  export MINIMAX_API_KEY=...
+
+  # Custom OpenAI-compatible endpoint
+  export LLM_PROVIDER=openai
+  export LLM_BASE_URL=http://localhost:8000/v1
+  export LLM_API_KEY=...
+"""
+
+import os
+import openai
+
+
+# Default models per provider
+_DEFAULT_MODELS = {
+    "openai": "gpt-4o",
+    "azure": "gpt-4o",
+    "minimax": "MiniMax-M2.7",
+}
+
+# Default base URLs per provider
+_DEFAULT_BASE_URLS = {
+    "openai": None,  # uses openai default
+    "minimax": "https://api.minimax.io/v1",
+}
+
+
+def _detect_provider():
+    """Auto-detect provider from available environment variables."""
+    if os.environ.get("MINIMAX_API_KEY"):
+        return "minimax"
+    if os.environ.get("AZURE_OPENAI_API_KEY") or os.environ.get("AZURE_OPENAI_ENDPOINT"):
+        return "azure"
+    return "openai"
+
+
+def _get_api_key(provider):
+    """Get API key for the given provider."""
+    # Generic key takes lowest precedence
+    key = os.environ.get("LLM_API_KEY", "")
+
+    if provider == "openai":
+        key = os.environ.get("OPENAI_API_KEY", key)
+    elif provider == "azure":
+        key = os.environ.get("AZURE_OPENAI_API_KEY", key)
+    elif provider == "minimax":
+        key = os.environ.get("MINIMAX_API_KEY", key)
+
+    # Fall back to conf.py GPT_AK if no env var is set
+    if not key:
+        try:
+            from conf import GPT_AK
+            if GPT_AK and GPT_AK != "[YOUR GPT_AK]":
+                key = GPT_AK
+        except ImportError:
+            pass
+
+    return key
+
+
+def get_llm_client(provider=None, api_key=None, base_url=None, model=None):
+    """
+    Create an OpenAI-compatible client for the specified provider.
+
+    Args:
+        provider: LLM provider name ("openai", "azure", "minimax").
+                  Auto-detected from env vars if not specified.
+        api_key: API key. Read from env vars if not specified.
+        base_url: Base URL. Uses provider default if not specified.
+        model: Model name. Uses provider default if not specified.
+
+    Returns:
+        tuple: (client, model_name)
+    """
+    if provider is None:
+        provider = os.environ.get("LLM_PROVIDER", "").lower() or _detect_provider()
+
+    if api_key is None:
+        api_key = _get_api_key(provider)
+
+    if model is None:
+        model = os.environ.get("LLM_MODEL", "") or _DEFAULT_MODELS.get(provider, "gpt-4o")
+
+    if base_url is None:
+        base_url = os.environ.get("LLM_BASE_URL", "") or _DEFAULT_BASE_URLS.get(provider)
+
+    if provider == "azure":
+        endpoint = base_url or os.environ.get(
+            "AZURE_OPENAI_ENDPOINT",
+            "https://search-va.byteintl.net/gpt/openapi/online/multimodal/crawl",
+        )
+        api_version = os.environ.get("AZURE_API_VERSION", "2023-07-01-preview")
+        client = openai.AzureOpenAI(
+            azure_endpoint=endpoint,
+            api_version=api_version,
+            api_key=api_key,
+        )
+    else:
+        # OpenAI, MiniMax, and any OpenAI-compatible provider
+        kwargs = {"api_key": api_key}
+        if base_url:
+            kwargs["base_url"] = base_url
+        client = openai.OpenAI(**kwargs)
+
+    return client, model
+
+
+def chat_completion(messages, provider=None, api_key=None, base_url=None,
+                    model=None, return_json=False, temperature=None, **kwargs):
+    """
+    Run a chat completion using the configured LLM provider.
+
+    Args:
+        messages: List of message dicts with "role" and "content" keys.
+        provider: LLM provider name. Auto-detected if not specified.
+        api_key: API key. Read from env vars if not specified.
+        base_url: Base URL. Uses provider default if not specified.
+        model: Model name. Uses provider default if not specified.
+        return_json: Whether to request JSON output format.
+        temperature: Sampling temperature (optional).
+
+    Returns:
+        str: The assistant's response content.
+    """
+    client, model_name = get_llm_client(provider, api_key, base_url, model)
+
+    create_kwargs = {
+        "model": model_name,
+        "messages": messages,
+    }
+    if return_json:
+        create_kwargs["response_format"] = {"type": "json_object"}
+
+    # MiniMax temperature must be in (0.0, 1.0]
+    resolved_provider = provider or os.environ.get("LLM_PROVIDER", "").lower() or _detect_provider()
+    if temperature is not None:
+        if resolved_provider == "minimax":
+            temperature = max(0.01, min(temperature, 1.0))
+        create_kwargs["temperature"] = temperature
+
+    completion = client.chat.completions.create(**create_kwargs)
+    content = completion.choices[0].message.content
+
+    # Strip MiniMax thinking tags if present
+    if resolved_provider == "minimax" and content and "<think>" in content:
+        import re
+        content = re.sub(r"<think>.*?</think>\s*", "", content, flags=re.DOTALL).strip()
+
+    return content
diff --git a/tools/prompt_rewriter.py b/tools/prompt_rewriter.py
index fe18a503..9402394f 100644
--- a/tools/prompt_rewriter.py
+++ b/tools/prompt_rewriter.py
@@ -15,17 +15,17 @@
 from PIL import Image
 import openai
 
-from conf import GPT_AK
+from tools.llm_provider import chat_completion, get_llm_client
 
 
 def encode_image(image_path, size=(512, 512)):
     """
     Resize an image and encode it as a Base64 string.
-    
+
     Args:
     - image_path (str): Path to the image file.
     - size (tuple): New size as a tuple, (width, height).
-    
+
     Returns:
     - str: Base64 encoded string of the resized image.
     """
@@ -42,8 +42,8 @@ def encode_image(image_path, size=(512, 512)):
 
 
 SYSTEM = """
-You are part of a team of bots that creates images. You work with an assistant bot that will draw anything you say. 
-For example, outputting the prompt and parameters like "<prompt:a beautiful morning in the woods with the sun peaking through the trees><cfg:3>" will trigger your partner bot to output an image of a forest morning, as described. 
+You are part of a team of bots that creates images. You work with an assistant bot that will draw anything you say.
+For example, outputting the prompt and parameters like "<prompt:a beautiful morning in the woods with the sun peaking through the trees><cfg:3>" will trigger your partner bot to output an image of a forest morning, as described.
 You will be prompted by users looking to create detailed, amazing images. The way to accomplish this is to refine their short prompts and make them extremely detailed and descriptive.
 - You will only ever output a single image description sentence per user request.
 - Each image description sentence should be consist of "<prompt:xxx><cfg:xxx>", where <prompt:xxx> is the image description, <cfg:xxx> is the parameter that control the image generation.
@@ -77,48 +77,42 @@ def __init__(self, system, few_shot_history):
 
     def rewrite(self, prompt):
         messages = self.system + self.few_shot_history + [{"role": "user", "content": prompt}]
-        result, _ = get_gpt_result(model_name='gpt-4o-2024-08-06', messages=messages, retry=5, ak=GPT_AK, return_json=False)
+        result, _ = get_gpt_result(messages=messages, retry=5, return_json=False)
         assert result
         return result
 
 
-def get_gpt_result(model_name='gpt-4o-2024-05-13', messages=None, retry=5, ak=None, return_json=False):
+def get_gpt_result(model_name=None, messages=None, retry=5, ak=None, return_json=False):
     """
-        Retrieves a chat response using the GPT-4 model.
-        Args:
-            model_name (str, optional): The name of the GPT model to use. Defaults to 'gpt-4'. [gpt-3.5-turbo, gpt-4]
-            retry (int, optional): The number of times to retry the chat API if there is an error. Defaults to 5.
-        Returns:
-            tuple: A tuple containing the chat response content (str) and the API usage (dict).
-        Raises:
-            Exception: If there is an error retrieving the chat response.
+    Retrieves a chat response using the configured LLM provider.
+
+    The provider is determined by environment variables (see tools/llm_provider.py).
+    Supports OpenAI, Azure OpenAI, MiniMax, and any OpenAI-compatible API.
+
+    Args:
+        model_name (str, optional): The model to use. If None, uses provider default.
+        messages (list): Chat messages.
+        retry (int): Number of retry attempts on failure.
+        ak (str, optional): API key override (deprecated, use env vars instead).
+        return_json (bool): Whether to request JSON response format.
+
+    Returns:
+        tuple: (response_content, None) on success, (None, -1) on failure.
     """
-    openai_ak = ak
-    client = openai.AzureOpenAI(
-        azure_endpoint="https://search-va.byteintl.net/gpt/openapi/online/multimodal/crawl",
-        api_version="2023-07-01-preview",
-        api_key=openai_ak
-    )
     for i in range(retry):
         try:
-            if return_json:
-                completion = client.chat.completions.create(
-                    model=model_name,
-                    messages=messages,
-                    response_format={ "type": "json_object" },
-                )
-            else:
-                completion = client.chat.completions.create(
-                    model=model_name,
-                    messages=messages,
-                )
-            result = json.loads(completion.model_dump_json())['choices'][0]['message']['content']
-            return result,None
+            result = chat_completion(
+                messages=messages,
+                model=model_name,
+                api_key=ak if ak else None,
+                return_json=return_json,
+            )
+            return result, None
         except Exception as e:
             traceback.print_exc()
-            if isinstance(e,KeyboardInterrupt):
+            if isinstance(e, KeyboardInterrupt):
                 exit(0)
-            sleep_time = 10 + random.randint(2,5)**(i+1)
+            sleep_time = 10 + random.randint(2, 5) ** (i + 1)
             time.sleep(sleep_time)
     return None, -1