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 7f0d59ac..a06385f4 100644 --- a/pr_agent/algo/__init__.py +++ b/pr_agent/algo/__init__.py @@ -84,6 +84,13 @@ 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, + 'xai/grok-3-beta': 131072, + 'xai/grok-3-fast-beta': 131072, + 'xai/grok-3-mini-beta': 131072, + 'xai/grok-3-mini-fast-beta': 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