Update to make ready for xAI

actions.py

@@ -28,7 +28,12 @@ except ImportError:
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)  # TODO: remove manual setting and make centrally controlled
 
-SUPPORTED_VENDORS = {"OpenAI": ChatOpenAI, "Anthropic": ChatAnthropic, "Mistral": ChatMistralAI}
+SUPPORTED_VENDORS = {
+    "OpenAI": ChatOpenAI,
+    "Anthropic": ChatAnthropic,
+    "Mistral": ChatMistralAI,
+    "xAI": ChatOpenAI,
+}
 
 SUPPORTED_MODELS = {
     "OpenAI": [
@@ -43,6 +48,7 @@ SUPPORTED_MODELS = {
         "claude-3-haiku-20240307",
     ],
     "Mistral": ["mistral-large-latest", "open-mistral-nemo", "codestral-latest"],
+    "xAI": ["grok-beta"],
 }
 DEFAULT_TEMPERATURE = 0.1
 DEFAULT_RETRY = 3
@@ -62,6 +68,8 @@ def get_vizro_ai_plot(user_prompt, df, model, api_key, api_base, vendor_input):
         )
     if vendor_input == "Mistral":
         llm = vendor(model=model, mistral_api_key=api_key, mistral_api_url=api_base, temperature=DEFAULT_TEMPERATURE)
+    if vendor_input == "xAI":
+        llm = vendor(model=model, openai_api_key=api_key, openai_api_base=api_base, temperature=DEFAULT_TEMPERATURE)
 
     vizro_ai = VizroAI(model=llm)
     ai_outputs = vizro_ai.plot(df, user_prompt, max_debug_retry=DEFAULT_RETRY, return_elements=True)

app.py

@@ -70,6 +70,7 @@ SUPPORTED_MODELS = {
         "claude-3-haiku-20240307",
     ],
     "Mistral": ["mistral-large-latest", "open-mistral-nemo", "codestral-latest"],
+    "xAI": ["grok-beta"],
 }
 
 
@@ -180,7 +181,11 @@ plot_page = vm.Page(
                 MyDropdown(
                     options=SUPPORTED_MODELS["OpenAI"], value="gpt-4o-mini", multi=False, id="model-dropdown-id"
                 ),
-                OffCanvas(id="settings", options=["OpenAI", "Anthropic", "Mistral"], value="OpenAI"),
+                OffCanvas(
+                    id="settings",
+                    options=["OpenAI", "Anthropic", "Mistral", "xAI"],
+                    value="OpenAI",
+                ),
                 UserPromptTextArea(id="text-area-id"),
                 # Modal(id="modal"),
             ],

requirements.in

@@ -1,5 +1,5 @@
 gunicorn
-vizro-ai>=0.3.1
+vizro-ai>=0.3.2
 black
 openpyxl
 langchain_anthropic

requirements.txt

@@ -266,7 +266,7 @@ urllib3==2.2.3
     # via requests
 vizro==0.1.23
     # via vizro-ai
-vizro-ai==0.3.1
+vizro-ai==0.3.2
     # via -r requirements.in
 werkzeug==3.0.4
     # via