Merge pull request #877 from randy-tsukemen/support-groq-llama3

Add Groq Llama3 support
Codium-ai · Apr 22, 2024 · fae6cab · fae6cab
2 parents 19f239a + d457fa2
commit fae6cab
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 0 deletions.
diff --git a/docs/docs/usage-guide/additional_configurations.md b/docs/docs/usage-guide/additional_configurations.md
@@ -125,6 +125,19 @@ key = ...
 
 Also, review the [AiHandler](https://github.com/Codium-ai/pr-agent/blob/main/pr_agent/algo/ai_handler.py) file for instructions on how to set keys for other models.
 
+### Groq
+
+To use Llama3 model with Groq, for example, set:
+```
+[config] # in configuration.toml
+model = "llama3-70b-8192"
+model_turbo = "llama3-70b-8192"
+fallback_models = ["groq/llama3-70b-8192"] 
+[groq] # in .secrets.toml
+key = ... # your Groq api key
+```
+(you can obtain a Groq key from [here](https://console.groq.com/keys))
+
 ### 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
@@ -31,4 +31,6 @@
     'bedrock/anthropic.claude-v2:1': 100000,
     'bedrock/anthropic.claude-3-sonnet-20240229-v1:0': 100000,
     'bedrock/anthropic.claude-3-haiku-20240307-v1:0': 100000,
+    'groq/llama3-8b-8192': 8192,
+    'groq/llama3-70b-8192': 8192,
 }
diff --git a/pr_agent/algo/ai_handlers/litellm_ai_handler.py b/pr_agent/algo/ai_handlers/litellm_ai_handler.py
@@ -52,6 +52,8 @@ def __init__(self):
             litellm.anthropic_key = get_settings().anthropic.key
         if get_settings().get("COHERE.KEY", None):
             litellm.cohere_key = get_settings().cohere.key
+        if get_settings().get("GROQ.KEY", None):
+            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("HUGGINGFACE.KEY", None):

diff --git a/pr_agent/settings/.secrets_template.toml b/pr_agent/settings/.secrets_template.toml
@@ -29,6 +29,9 @@ key = "" # Optional, uncomment if you want to use Cohere. Acquire through https:
 [replicate]
 key = "" # Optional, uncomment if you want to use Replicate. Acquire through https://replicate.com/
 
+[groq]
+key = "" # Acquire through https://console.groq.com/keys
+
 [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