Skip to content

Commit a02bde3

Browse files
committed
watsonx-integration
1 parent 4ebe352 commit a02bde3

25 files changed

+742
-30
lines changed

src/main/java/ee/carlrobert/codegpt/Icons.java

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ public final class Icons {
2424
public static final Icon You = IconLoader.getIcon("/icons/you.svg", Icons.class);
2525
public static final Icon YouSmall = IconLoader.getIcon("/icons/you_small.png", Icons.class);
2626
public static final Icon Ollama = IconLoader.getIcon("/icons/ollama.svg", Icons.class);
27+
public static final Icon Watsonx = IconLoader.getIcon("/icons/watsonx.svg", Icons.class);
2728
public static final Icon User = IconLoader.getIcon("/icons/user.svg", Icons.class);
2829
public static final Icon Upload = IconLoader.getIcon("/icons/upload.svg", Icons.class);
2930
public static final Icon GreenCheckmark =

src/main/java/ee/carlrobert/codegpt/completions/CompletionClientProvider.java

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import ee.carlrobert.codegpt.settings.service.llama.LlamaSettings;
1111
import ee.carlrobert.codegpt.settings.service.ollama.OllamaSettings;
1212
import ee.carlrobert.codegpt.settings.service.openai.OpenAISettings;
13+
import ee.carlrobert.codegpt.settings.service.watsonx.WatsonxSettings;
1314
import ee.carlrobert.llm.client.anthropic.ClaudeClient;
1415
import ee.carlrobert.llm.client.azure.AzureClient;
1516
import ee.carlrobert.llm.client.azure.AzureCompletionRequestParams;
@@ -18,6 +19,7 @@
1819
import ee.carlrobert.llm.client.llama.LlamaClient;
1920
import ee.carlrobert.llm.client.ollama.OllamaClient;
2021
import ee.carlrobert.llm.client.openai.OpenAIClient;
22+
import ee.carlrobert.llm.client.watsonx.WatsonxClient;
2123
import java.net.InetSocketAddress;
2224
import java.net.Proxy;
2325
import java.util.concurrent.TimeUnit;
@@ -32,6 +34,16 @@ public static CodeGPTClient getCodeGPTClient() {
3234
getDefaultClientBuilder());
3335
}
3436

37+
public static WatsonxClient getWatsonxClient() {
38+
return new WatsonxClient.Builder(getCredential(CredentialKey.WATSONX_API_KEY))
39+
.setApiVersion(WatsonxSettings.getCurrentState().getApiVersion())
40+
.setIsOnPrem(WatsonxSettings.getCurrentState().isOnPrem())
41+
.setHost(WatsonxSettings.getCurrentState().getOnPremHost())
42+
.setUsername(WatsonxSettings.getCurrentState().getUsername())
43+
.setIsZenApiKey(WatsonxSettings.getCurrentState().isZenApiKey())
44+
.build(getDefaultClientBuilder());
45+
}
46+
3547
public static OpenAIClient getOpenAIClient() {
3648
return new OpenAIClient.Builder(getCredential(CredentialKey.OPENAI_API_KEY))
3749
.setOrganization(OpenAISettings.getCurrentState().getOrganization())

src/main/java/ee/carlrobert/codegpt/completions/CompletionRequestProvider.java

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
import ee.carlrobert.codegpt.settings.service.llama.LlamaSettings;
2929
import ee.carlrobert.codegpt.settings.service.ollama.OllamaSettings;
3030
import ee.carlrobert.codegpt.settings.service.openai.OpenAISettings;
31+
import ee.carlrobert.codegpt.settings.service.watsonx.WatsonxSettings;
3132
import ee.carlrobert.codegpt.util.file.FileUtil;
3233
import ee.carlrobert.llm.client.anthropic.completion.ClaudeBase64Source;
3334
import ee.carlrobert.llm.client.anthropic.completion.ClaudeCompletionDetailedMessage;
@@ -55,6 +56,7 @@
5556
import ee.carlrobert.llm.client.openai.completion.request.OpenAIMessageImageURLContent;
5657
import ee.carlrobert.llm.client.openai.completion.request.OpenAIMessageTextContent;
5758
import ee.carlrobert.llm.client.openai.completion.request.RequestDocumentationDetails;
59+
import ee.carlrobert.llm.client.watsonx.completion.WatsonxCompletionRequest;
5860
import java.io.IOException;
5961
import java.nio.charset.StandardCharsets;
6062
import java.nio.file.Files;
@@ -307,6 +309,29 @@ private static Request buildCustomOpenAIChatCompletionRequest(
307309
}
308310
}
309311

312+
public WatsonxCompletionRequest buildWatsonxChatCompletionRequest(
313+
CallParameters callParameters) {
314+
var settings = WatsonxSettings.getCurrentState();
315+
String prompt = PersonaSettings.getSystemPrompt();
316+
prompt += "\n"+callParameters.getMessage().getPrompt();
317+
var builder = new WatsonxCompletionRequest.Builder(prompt);
318+
builder.setDecodingMethod(settings.isGreedyDecoding() ? "greedy" : "sample");
319+
builder.setModelId(settings.getModel());
320+
builder.setProjectId(settings.getProjectId());
321+
builder.setSpaceId(settings.getSpaceId());
322+
builder.setMaxNewTokens(settings.getMaxNewTokens());
323+
builder.setMinNewTokens(settings.getMinNewTokens());
324+
builder.setTemperature(settings.getTemperature());
325+
builder.setStopSequences(settings.getStopSequences().isEmpty() ? null : settings.getStopSequences().split(","));
326+
builder.setTopP(settings.getTopP());
327+
builder.setTopK(settings.getTopK());
328+
builder.setIncludeStopSequence(settings.getIncludeStopSequence());
329+
builder.setRandomSeed(settings.getRandomSeed());
330+
builder.setRepetitionPenalty(settings.getRepetitionPenalty());
331+
builder.setStream(true);
332+
return builder.build();
333+
}
334+
310335
public ClaudeCompletionRequest buildAnthropicChatCompletionRequest(
311336
CallParameters callParameters) {
312337
var configuration = ConfigurationSettings.getState();

src/main/java/ee/carlrobert/codegpt/completions/CompletionRequestService.java

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -126,6 +126,9 @@ public EventSource getChatCompletionAsync(
126126
settings.getModel(),
127127
eventListener);
128128
}
129+
case WATSONX -> CompletionClientProvider.getWatsonxClient().getCompletionAsync(
130+
requestProvider.buildWatsonxChatCompletionRequest(callParameters),
131+
eventListener);
129132
};
130133
}
131134

@@ -285,7 +288,7 @@ public static boolean isRequestAllowed(ServiceType serviceType) {
285288
AzureSettings.getCurrentState().isUseAzureApiKeyAuthentication()
286289
? CredentialKey.AZURE_OPENAI_API_KEY
287290
: CredentialKey.AZURE_ACTIVE_DIRECTORY_TOKEN);
288-
case CODEGPT, CUSTOM_OPENAI, ANTHROPIC, LLAMA_CPP, OLLAMA -> true;
291+
case CODEGPT, CUSTOM_OPENAI, ANTHROPIC, LLAMA_CPP, OLLAMA, WATSONX -> true;
289292
case GOOGLE -> CredentialsStore.INSTANCE.isCredentialSet(CredentialKey.GOOGLE_API_KEY);
290293
};
291294
}

src/main/java/ee/carlrobert/codegpt/conversations/ConversationService.java

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
import ee.carlrobert.codegpt.settings.service.llama.LlamaSettings;
1414
import ee.carlrobert.codegpt.settings.service.ollama.OllamaSettings;
1515
import ee.carlrobert.codegpt.settings.service.openai.OpenAISettings;
16+
import ee.carlrobert.codegpt.settings.service.watsonx.WatsonxSettings;
1617
import java.time.LocalDateTime;
1718
import java.util.ArrayList;
1819
import java.util.Comparator;
@@ -210,6 +211,9 @@ private static String getModelForSelectedService(ServiceType serviceType) {
210211
case GOOGLE -> application.getService(GoogleSettings.class)
211212
.getState()
212213
.getModel();
214+
case WATSONX -> application.getService(WatsonxSettings.class)
215+
.getState()
216+
.getModel();
213217
};
214218
}
215219
}

src/main/java/ee/carlrobert/codegpt/settings/GeneralSettings.java

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import ee.carlrobert.codegpt.settings.service.llama.LlamaSettings;
1919
import ee.carlrobert.codegpt.settings.service.ollama.OllamaSettings;
2020
import ee.carlrobert.codegpt.settings.service.openai.OpenAISettings;
21+
import ee.carlrobert.codegpt.settings.service.watsonx.WatsonxSettings;
2122
import ee.carlrobert.codegpt.util.ApplicationUtil;
2223
import org.jetbrains.annotations.NotNull;
2324

@@ -92,6 +93,10 @@ public void sync(Conversation conversation) {
9293
ApplicationManager.getApplication().getService(OllamaSettings.class).getState()
9394
.setModel(conversation.getModel());
9495
break;
96+
case WATSONX:
97+
ApplicationManager.getApplication().getService(WatsonxSettings.class).getState()
98+
.setModel(conversation.getModel());
99+
break;
95100
default:
96101
break;
97102
}
@@ -144,6 +149,11 @@ public String getModel() {
144149
.getService(GoogleSettings.class)
145150
.getState()
146151
.getModel();
152+
case WATSONX:
153+
return ApplicationManager.getApplication()
154+
.getService(WatsonxSettings.class)
155+
.getState()
156+
.getModel();
147157
default:
148158
return "Unknown";
149159
}

src/main/java/ee/carlrobert/codegpt/settings/service/ServiceType.java

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,9 @@ public enum ServiceType {
1212
AZURE("AZURE", "service.azure.title", "azure.chat.completion"),
1313
GOOGLE("GOOGLE", "service.google.title", "google.chat.completion"),
1414
LLAMA_CPP("LLAMA_CPP", "service.llama.title", "llama.chat.completion"),
15-
OLLAMA("OLLAMA", "service.ollama.title", "ollama.chat.completion");
15+
OLLAMA("OLLAMA", "service.ollama.title", "ollama.chat.completion"),
16+
WATSONX("WATSONX", "service.watsonx.title", "watsonx.chat.completion");
17+
1618

1719
private final String code;
1820
private final String label;
Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
package ee.carlrobert.codegpt.settings.service.watsonx;
2+
3+
import com.intellij.openapi.application.ApplicationManager;
4+
import com.intellij.openapi.components.PersistentStateComponent;
5+
import com.intellij.openapi.components.State;
6+
import com.intellij.openapi.components.Storage;
7+
import ee.carlrobert.codegpt.completions.llama.LlamaModel;
8+
import org.jetbrains.annotations.NotNull;
9+
10+
@State(name = "CodeGPT_WatsonxSettings", storages = @Storage("CodeGPT_WatsonxSettings.xml"))
11+
public class WatsonxSettings implements PersistentStateComponent<WatsonxSettingsState> {
12+
13+
private WatsonxSettingsState state = new WatsonxSettingsState();
14+
15+
@Override
16+
@NotNull
17+
public WatsonxSettingsState getState() {
18+
return state;
19+
}
20+
21+
@Override
22+
public void loadState(@NotNull WatsonxSettingsState state) {
23+
this.state = state;
24+
}
25+
26+
public static WatsonxSettingsState getCurrentState() {
27+
return getInstance().getState();
28+
}
29+
public static boolean isCodeCompletionsPossible() {
30+
return getInstance().getState().isCodeCompletionsEnabled();
31+
}
32+
33+
public static ee.carlrobert.codegpt.settings.service.watsonx.WatsonxSettings getInstance() {
34+
return ApplicationManager.getApplication().getService(WatsonxSettings.class);
35+
}
36+
}

0 commit comments

Comments
 (0)