feat: add DeepInfra support

This commit is contained in:
muhammad-asn
2025-03-06 15:49:07 +07:00
parent eb5f38b13b
commit 4f2551e0a6
4 changed files with 23 additions and 0 deletions

View File

@ -193,6 +193,19 @@ key = ...
(you can obtain a deepseek-chat key from [here](https://platform.deepseek.com))
### DeepInfra
To use DeepSeek model with DeepInfra, for example, set:
```
[config] # in configuration.toml
model = "deepinfra/deepseek-ai/DeepSeek-R1-Distill-Llama-70B"
fallback_models = ["deepinfra/deepseek-ai/DeepSeek-R1-Distill-Qwen-32B"]
[deepinfra] # in .secrets.toml
key = ... # your DeepInfra api key
```
(you can obtain a DeepInfra key from [here](https://deepinfra.com/dash/api_keys))
### Custom models
If the relevant model doesn't appear [here](https://github.com/Codium-ai/pr-agent/blob/main/pr_agent/algo/__init__.py), you can still use it as a custom model:

View File

@ -86,6 +86,9 @@ MAX_TOKENS = {
"watsonx/ibm/granite-13b-chat-v2": 8191,
"watsonx/ibm/granite-34b-code-instruct": 8191,
"watsonx/mistralai/mistral-large": 32768,
"deepinfra/deepseek-ai/DeepSeek-R1-Distill-Qwen-32B": 128000,
"deepinfra/deepseek-ai/DeepSeek-R1-Distill-Llama-70B": 128000,
"deepinfra/deepseek-ai/DeepSeek-R1": 128000,
}
USER_MESSAGE_ONLY_MODELS = [

View File

@ -96,6 +96,10 @@ class LiteLLMAIHandler(BaseAiHandler):
if get_settings().get("DEEPSEEK.KEY", None):
os.environ['DEEPSEEK_API_KEY'] = get_settings().get("DEEPSEEK.KEY")
# Support deepinfra models
if get_settings().get("DEEPINFRA.KEY", None):
os.environ['DEEPINFRA_API_KEY'] = get_settings().get("DEEPINFRA.KEY")
# Models that only use user meessage
self.user_message_only_models = USER_MESSAGE_ONLY_MODELS

View File

@ -94,3 +94,6 @@ pat = ""
[deepseek]
key = ""
[deepinfra]
key = ""