finalized grok addition

README.md

@@ -15,7 +15,7 @@ Chest X-rays (CXRs) play an integral role in driving critical decisions in disea
 ## MedRAX
 MedRAX is built on a robust technical foundation:
 - **Core Architecture**: Built on LangChain and LangGraph frameworks
-- **Language Models**: Supports multiple LLM providers including OpenAI (GPT-4o) and Google (Gemini) models
+- **Language Models**: Supports multiple LLM providers including OpenAI (GPT-4o), Google (Gemini), and xAI (Grok) models
 - **Deployment**: Supports both local and cloud-based deployments
 - **Interface**: Production-ready interface built with Gradio
 - **Modular Design**: Tool-agnostic architecture allowing easy integration of new capabilities
@@ -54,7 +54,9 @@ Unzip the Eurorad figures to your local `MedMAX` directory.
 unzip chestagentbench/figures.zip
 ```
 
-To evaluate with GPT-4o, set your OpenAI API key in your `.env` file (see the "Environment Variable Setup" section for details) and run the quickstart script.
+To evaluate with different models, set the appropriate API key in your `.env` file (see the "Environment Variable Setup" section for details) and run the quickstart script.
+
+**Example with GPT-4o:**
 ```
 python quickstart.py \
     --model gpt-4o \
@@ -112,6 +114,9 @@ GOOGLE_API_KEY=
 OPENROUTER_API_KEY=
 OPENROUTER_BASE_URL= # Optional: Defaults to https://openrouter.ai/api/v1
 
+# xAI
+XAI_API_KEY=
+
 # -------------------------
 # Tool-specific API Keys
 # -------------------------
@@ -359,7 +364,10 @@ Supported prefix: `openrouter-`
 
 Access many open source and proprietary models via [OpenRouter](https://openrouter.ai/).
 
-**Note:** Tool compatibility may vary with open-source models. For best results with tools, we recommend using OpenAI or Google Gemini models.
+#### xAI Grok Models
+Supported prefix: `grok-`
+
+**Note:** Tool compatibility may vary with open-source models. For best results with tools, we recommend using OpenAI, Google Gemini, or xAI Grok models.
 
 #### Local LLMs
 If you are running a local LLM using frameworks like [Ollama](https://ollama.com/) or [LM Studio](https://lmstudio.ai/), you can configure the `OPENAI_BASE_URL` in your `.env` file to point to your local endpoint (e.g., `http://localhost:11434/v1`).

main.py

@@ -186,7 +186,7 @@ if __name__ == "__main__":
         model_dir="/model-weights",
         temp_dir="temp",  # Change this to the path of the temporary directory
         device="cuda",
-        model="gpt-4.1-2025-04-14",  # Change this to the model you want to use, e.g. gpt-4.1-2025-04-14, gemini-2.5-pro
+        model="grok-4",  # Change this to the model you want to use, e.g. gpt-4.1-2025-04-14, gemini-2.5-pro
         temperature=0.7,
         top_p=0.95,
         model_kwargs=model_kwargs,

medrax/models/model_factory.py

@@ -6,6 +6,7 @@ from typing import Dict, Any, Type
 from langchain_core.language_models import BaseLanguageModel
 from langchain_openai import ChatOpenAI
 from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_xai import ChatXAI
 
 
 class ModelFactory:
@@ -35,10 +36,8 @@ class ModelFactory:
             "default_base_url": "https://openrouter.ai/api/v1",
         },
         "grok": {
-        "class": ChatOpenAI,  # xAI uses OpenAI-compatible API
+        "class": ChatXAI,
         "env_key": "XAI_API_KEY",
-        "base_url_key": "XAI_BASE_URL", 
-        "default_base_url": "https://api.x.ai/v1"
         }
         # Add more providers with default configurations here
     }

pyproject.toml

@@ -20,6 +20,7 @@ dependencies = [
     "langchain-openai>=0.0.2",
     "langchain-anthropic>=0.0.2",
     "langchain-cohere>=0.0.2",
+    "langchain-xai>=0.0.1",
     "langchain-chroma>=0.0.10",
     "langgraph>=0.0.10",
     "python-dotenv>=0.19.0",