From 665fb90a98af3fc9f4a91fe791c19b7bfa25323a Mon Sep 17 00:00:00 2001
From: Peter Dave Hello <hsu@peterdavehello.org>
Date: Tue, 8 Apr 2025 01:36:21 +0800
Subject: [PATCH] Add support of xAI and their Grok-2 model

Close #1630
---
 docs/docs/usage-guide/changing_a_model.md       | 14 ++++++++++++++
 pr_agent/algo/__init__.py                       |  3 +++
 pr_agent/algo/ai_handlers/litellm_ai_handler.py |  2 ++
 pr_agent/settings/.secrets_template.toml        |  3 +++
 4 files changed, 22 insertions(+)

diff --git a/docs/docs/usage-guide/changing_a_model.md b/docs/docs/usage-guide/changing_a_model.md
index 94861629..3a978f82 100644
--- a/docs/docs/usage-guide/changing_a_model.md
+++ b/docs/docs/usage-guide/changing_a_model.md
@@ -113,6 +113,20 @@ key = ... # your Groq api key
 ```
 (you can obtain a Groq key from [here](https://console.groq.com/keys))
 
+### xAI
+
+To use xAI's models with PR-Agent, set:
+```
+[config] # in configuration.toml
+model = "xai/grok-2-latest"
+fallback_models = ["xai/grok-2-latest"] # or any other model as fallback
+
+[xai] # in .secrets.toml
+key = "..." # your xAI API key
+```
+
+You can obtain an xAI API key from [xAI's console](https://console.x.ai/) by creating an account and navigating to the developer settings page.
+
 ### Vertex AI
 
 To use Google's Vertex AI platform and its associated models (chat-bison/codechat-bison) set:
diff --git a/pr_agent/algo/__init__.py b/pr_agent/algo/__init__.py
index 6bea7994..449d89e1 100644
--- a/pr_agent/algo/__init__.py
+++ b/pr_agent/algo/__init__.py
@@ -82,6 +82,9 @@ MAX_TOKENS = {
     'groq/llama-3.3-70b-versatile': 128000,
     'groq/mixtral-8x7b-32768': 32768,
     'groq/gemma2-9b-it': 8192,
+    'xai/grok-2': 131072,
+    'xai/grok-2-1212': 131072,
+    'xai/grok-2-latest': 131072,
     'ollama/llama3': 4096,
     'watsonx/meta-llama/llama-3-8b-instruct': 4096,
     "watsonx/meta-llama/llama-3-70b-instruct": 4096,
diff --git a/pr_agent/algo/ai_handlers/litellm_ai_handler.py b/pr_agent/algo/ai_handlers/litellm_ai_handler.py
index 78d5f479..2ca04ea3 100644
--- a/pr_agent/algo/ai_handlers/litellm_ai_handler.py
+++ b/pr_agent/algo/ai_handlers/litellm_ai_handler.py
@@ -67,6 +67,8 @@ class LiteLLMAIHandler(BaseAiHandler):
             litellm.api_key = get_settings().groq.key
         if get_settings().get("REPLICATE.KEY", None):
             litellm.replicate_key = get_settings().replicate.key
+        if get_settings().get("XAI.KEY", None):
+            litellm.api_key = get_settings().xai.key
         if get_settings().get("HUGGINGFACE.KEY", None):
             litellm.huggingface_key = get_settings().huggingface.key
         if get_settings().get("HUGGINGFACE.API_BASE", None) and 'huggingface' in get_settings().config.model:
diff --git a/pr_agent/settings/.secrets_template.toml b/pr_agent/settings/.secrets_template.toml
index 42946da5..28201a0d 100644
--- a/pr_agent/settings/.secrets_template.toml
+++ b/pr_agent/settings/.secrets_template.toml
@@ -32,6 +32,9 @@ key = "" # Optional, uncomment if you want to use Replicate. Acquire through htt
 [groq]
 key = "" # Acquire through https://console.groq.com/keys
 
+[xai]
+key = "" # Optional, uncomment if you want to use xAI. Acquire through https://console.x.ai/
+
 [huggingface]
 key = "" # Optional, uncomment if you want to use Huggingface Inference API. Acquire through https://huggingface.co/docs/api-inference/quicktour
 api_base = "" # the base url for your huggingface inference endpoint