vrr/ProxyAI
2
0
Fork 0
mirror of https://github.com/carlrobertoh/ProxyAI.git synced 2026-10-19 01:33:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

refactor: remove llm-client dependency in favor of Koog and refresh models
Some checks failed
Build / Build (push) Has been cancelled
Details
Build / Verify Plugin (push) Has been cancelled
Details

Browse source
This commit is contained in:
Carl-Robert Linnupuu 2026-03-02 14:20:46 +00:00
parent 0eaf6f8668
commit 00775d38da
241 changed files with 7964 additions and 10007 deletions
Download patch file Download diff file Expand all files Collapse all files

10
build.gradle.kts
View file

@ -64,7 +64,6 @@ dependencies {
bundledPlugin("com.intellij.java")
bundledPlugin("org.jetbrains.kotlin")
bundledPlugin("Git4Idea")
plugin("PythonCore:251.25410.109")
testFramework(TestFrameworkType.Platform)
testFramework(TestFrameworkType.JUnit5)
@ -73,7 +72,6 @@ dependencies {
implementation(project(":proxyai-telemetry"))
implementation(project(":proxyai-treesitter"))
implementation(platform(libs.okhttp.bom))
implementation(platform(libs.slf4j.bom))
implementation(platform(libs.jackson.bom))
implementation(platform(libs.mcp.sdk.bom))
@ -81,14 +79,18 @@ dependencies {
implementation(libs.jackson.datatype.jdk8)
implementation(libs.jackson.datatype.jsr310)
implementation(libs.jackson.module.kotlin)
implementation(libs.kotlin.stdlib)
implementation(libs.flexmark.all) {
// vulnerable transitive dependency
exclude(group = "org.jsoup", module = "jsoup")
}
implementation(libs.kotlin.stdlib)
implementation(libs.kotlin.reflect)
implementation(libs.koog.agents) {
exclude(group = "org.jetbrains.kotlin", module = "kotlin-stdlib")
exclude(group = "org.jetbrains.kotlin", module = "kotlin-stdlib-common")
exclude(group = "org.jetbrains.kotlin", module = "kotlin-stdlib-jdk7")
exclude(group = "org.jetbrains.kotlin", module = "kotlin-stdlib-jdk8")
exclude(group = "org.jetbrains.kotlin", module = "kotlin-reflect")
exclude(group = "org.jetbrains.kotlinx", module = "kotlinx-coroutines-core")
exclude(group = "org.jetbrains.kotlinx", module = "kotlinx-coroutines-jdk8")
}

16
buildSrc/src/main/kotlin/proxyai.java-conventions.gradle.kts
View file

@ -1,5 +1,6 @@
import org.jetbrains.intellij.platform.gradle.TestFrameworkType
import org.jetbrains.kotlin.gradle.dsl.JvmTarget
import org.jetbrains.kotlin.gradle.dsl.KotlinJvmProjectExtension
import org.jetbrains.kotlin.gradle.tasks.KotlinCompile
val libs = versionCatalogs.named("libs")
@ -34,13 +35,7 @@ dependencies {
}
}
implementation(lib("llm.client"))
implementation(lib("kotlinx.serialization.json"))
constraints {
implementation(lib("okio")) {
because("llm-client 0.7.0 uses okio 3.2.0: https://avd.aquasec.com/nvd/cve-2023-3635")
}
}
testImplementation("junit:junit:4.13.2")
testImplementation(platform(lib("junit.bom")))
@ -53,12 +48,21 @@ dependencies {
tasks {
properties("javaVersion").let {
java {
toolchain {
languageVersion.set(JavaLanguageVersion.of(it.toInt()))
}
}
extensions.configure<KotlinJvmProjectExtension>("kotlin") {
jvmToolchain(it.toInt())
}
withType<JavaCompile> {
sourceCompatibility = it
targetCompatibility = it
}
withType<KotlinCompile> {
compilerOptions.jvmTarget.set(JvmTarget.fromTarget(it))
compilerOptions.freeCompilerArgs.add("-opt-in=kotlin.time.ExperimentalTime")
}
}
}

7
gradle.properties
View file

@ -18,10 +18,10 @@ platformVersion = 2025.1.1
# Example: platformPlugins = com.intellij.java, com.jetbrains.php:203.4449.22
platformPlugins =
javaVersion = 17
javaVersion = 21
# Gradle Releases -> https://github.com/gradle/gradle/releases
gradleVersion = 8.7
gradleVersion = 8.13
# Opt-out flag for bundling Kotlin standard library -> https://jb.gg/intellij-platform-kotlin-stdlib
kotlin.stdlib.default.dependency = false
@ -42,3 +42,6 @@ org.gradle.caching = true
# Enable Gradle Kotlin DSL Lazy Property Assignment -> https://docs.gradle.org/current/userguide/kotlin_dsl.html#kotdsl:assignment
systemProp.org.gradle.unsafe.kotlin.assignment = true
# Allow Gradle toolchains to auto-provision JDKs (needed for Java 21 on newer IntelliJ platform)
org.gradle.java.installations.auto-download = true

12
gradle/libs.versions.toml
View file

@ -5,19 +5,16 @@ changelog = "2.2.1"
checkstyle = "10.15.0"
commons-text = "1.13.1"
flexmark = "0.64.8"
gradle-intellij-plugin-version = "2.6.0"
gradle-intellij-plugin-version = "2.10.4"
gson = "2.13.1"
jackson = "2.18.3"
mcp = "0.8.0"
jsoup = "1.21.2"
jtokkit = "1.1.0"
junit = "5.13.4"
kotlin = "2.2.10"
kotlin = "2.2.20"
kotlinx-serialization = "1.6.3"
llm-client = "0.8.58"
koog = "0.6.2"
okhttp = "4.12.0"
okio = "3.15.0"
koog = "0.6.3"
slf4j = "2.0.17"
tree-sitter = "0.24.5"
grpc = "1.73.0"
@ -46,9 +43,6 @@ kotlin-reflect = { module = "org.jetbrains.kotlin:kotlin-reflect", version.ref =
koog-agents = { module = "ai.koog:koog-agents", version.ref = "koog" }
kotlin-gradle-plugin = { module = "org.jetbrains.kotlin:kotlin-gradle-plugin", version.ref = "kotlin" }
kotlinx-serialization-json = { module = "org.jetbrains.kotlinx:kotlinx-serialization-json", version.ref = "kotlinx-serialization" }
llm-client = { module = "ee.carlrobert:llm-client", version.ref = "llm-client" }
okhttp-bom = { module = "com.squareup.okhttp3:okhttp-bom", version.ref = "okhttp" }
okio = { module = "com.squareup.okio:okio", version.ref = "okio" }
slf4j-bom = { module = "org.slf4j:slf4j-bom", version.ref = "slf4j" }
tree-sitter = { module = "io.github.bonede:tree-sitter", version.ref = "tree-sitter" }
grpc-protobuf = { module = "io.grpc:grpc-protobuf", version.ref = "grpc" }

2
gradle/wrapper/gradle-wrapper.properties vendored
View file

@ -1,6 +1,6 @@
distributionBase=GRADLE_USER_HOME
distributionPath=wrapper/dists
distributionUrl=https\://services.gradle.org/distributions/gradle-8.11-bin.zip
distributionUrl=https\://services.gradle.org/distributions/gradle-8.13-bin.zip
networkTimeout=10000
validateDistributionUrl=true
zipStoreBase=GRADLE_USER_HOME

2
src/main/java/ee/carlrobert/codegpt/CodeGPTKeys.java
View file

@ -4,8 +4,8 @@ import com.intellij.openapi.util.Key;
import ee.carlrobert.codegpt.inlineedit.InlineEditSession;
import ee.carlrobert.codegpt.nextedit.NextEditDiffViewer;
import ee.carlrobert.codegpt.inlineedit.InlineEditInlayRenderer;
import ee.carlrobert.codegpt.settings.service.codegpt.CodeGPTUserDetails;
import ee.carlrobert.codegpt.toolwindow.chat.editor.ToolWindowEditorFileDetails;
import ee.carlrobert.llm.client.codegpt.CodeGPTUserDetails;
import ee.carlrobert.service.NextEditResponse;
import com.intellij.openapi.util.TextRange;
import ee.carlrobert.service.PartialCodeCompletionResponse;

23
src/main/java/ee/carlrobert/codegpt/EncodingManager.java
View file

@ -10,11 +10,7 @@ import com.knuddels.jtokkit.api.EncodingType;
import com.knuddels.jtokkit.api.IntArrayList;
import ee.carlrobert.codegpt.conversations.Conversation;
import ee.carlrobert.codegpt.conversations.message.Message;
import ee.carlrobert.llm.client.openai.completion.request.OpenAIChatCompletionDetailedMessage;
import ee.carlrobert.llm.client.openai.completion.request.OpenAIChatCompletionMessage;
import ee.carlrobert.llm.client.openai.completion.request.OpenAIChatCompletionStandardMessage;
import ee.carlrobert.llm.client.openai.completion.request.OpenAIMessageTextContent;
import ee.carlrobert.llm.client.openai.completion.response.ToolCall;
import ee.carlrobert.codegpt.completions.ChatToolCall;
import java.util.List;
import java.util.Map;
import java.util.stream.Stream;
@ -55,9 +51,9 @@ public final class EncodingManager {
int inputTokens = 0;
int outputTokens = 0;
List<ToolCall> toolCalls = message.getToolCalls();
List<ChatToolCall> toolCalls = message.getToolCalls();
if (toolCalls != null) {
for (ToolCall toolCall : toolCalls) {
for (ChatToolCall toolCall : toolCalls) {
if (toolCall.getFunction() != null && toolCall.getFunction().getArguments() != null) {
inputTokens += countTokens(toolCall.getFunction().getArguments());
}
@ -103,19 +99,6 @@ public final class EncodingManager {
}
}
public int countMessageTokens(OpenAIChatCompletionMessage message) {
if (message instanceof OpenAIChatCompletionStandardMessage standardMessage) {
return countMessageTokens(standardMessage.getRole(), standardMessage.getContent());
}
return ((OpenAIChatCompletionDetailedMessage) message).getContent().stream()
.filter(OpenAIMessageTextContent.class::isInstance)
.mapToInt(it -> countMessageTokens(
((OpenAIChatCompletionDetailedMessage) message).getRole(),
((OpenAIMessageTextContent) it).getText()))
.sum();
}
public int countMessageTokens(String role, String content) {
var tokensPerMessage = 4; // every message follows <|start|>{role/name}\n{content}<|end|>\n
return countTokens(role + content) + tokensPerMessage;

2
src/main/java/ee/carlrobert/codegpt/Icons.java
View file

@ -25,8 +25,6 @@ public final class Icons {
public static final Icon Meta = IconLoader.getIcon("/icons/meta.svg", Icons.class);
public static final Icon Mistral = IconLoader.getIcon("/icons/mistral.svg", Icons.class);
public static final Icon Moonshot = IconLoader.getIcon("/icons/moonshot.svg", Icons.class);
public static final Icon Morph = IconLoader.getIcon("/icons/morph.svg", Icons.class);
public static final Icon Relace = IconLoader.getIcon("/icons/relace.svg", Icons.class);
public static final Icon Inception = IconLoader.getIcon("/icons/inception.svg", Icons.class);
public static final Icon Send = IconLoader.getIcon("/icons/send.svg", Icons.class);
public static final Icon Sparkle = IconLoader.getIcon("/icons/sparkle.svg", Icons.class);

92
src/main/java/ee/carlrobert/codegpt/completions/ChatCompletionEventListener.java
View file

@ -1,92 +0,0 @@
package ee.carlrobert.codegpt.completions;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.intellij.openapi.diagnostic.Logger;
import com.intellij.openapi.project.Project;
import ee.carlrobert.codegpt.codecompletions.CompletionProgressNotifier;
import ee.carlrobert.codegpt.events.CodeGPTEvent;
import ee.carlrobert.codegpt.telemetry.TelemetryAction;
import ee.carlrobert.llm.client.openai.completion.ErrorDetails;
import ee.carlrobert.llm.completion.CompletionEventListener;
import okhttp3.sse.EventSource;
public class ChatCompletionEventListener implements CompletionEventListener<String> {
private static final Logger LOG = Logger.getInstance(ChatCompletionEventListener.class);
private final Project project;
private final ChatCompletionParameters callParameters;
private final CompletionResponseEventListener eventListener;
private final StringBuilder messageBuilder = new StringBuilder();
public ChatCompletionEventListener(
Project project,
ChatCompletionParameters callParameters,
CompletionResponseEventListener eventListener) {
this.project = project;
this.callParameters = callParameters;
this.eventListener = eventListener;
}
@Override
public void onOpen() {
eventListener.handleRequestOpen();
}
@Override
public void onEvent(String data) {
try {
var event = new ObjectMapper().readValue(data, CodeGPTEvent.class);
eventListener.handleCodeGPTEvent(event);
} catch (JsonProcessingException exception) {
LOG.debug("Failed to parse CodeGPTEvent from data: " + data, exception);
}
}
@Override
public void onMessage(String message, EventSource eventSource) {
messageBuilder.append(message);
callParameters.getMessage().setResponse(messageBuilder.toString());
eventListener.handleMessage(message);
}
@Override
public void onComplete(StringBuilder messageBuilder) {
handleCompleted(messageBuilder);
}
@Override
public void onCancelled(StringBuilder messageBuilder) {
handleCompleted(messageBuilder);
}
@Override
public void onError(ErrorDetails error, Throwable ex) {
try {
callParameters.getConversation().addMessage(callParameters.getMessage());
eventListener.handleError(error, ex);
} finally {
sendError(error, ex);
}
}
private void handleCompleted(StringBuilder messageBuilder) {
CompletionProgressNotifier.Companion.update(project, false);
eventListener.handleCompleted(messageBuilder.toString(), callParameters);
}
private void sendError(ErrorDetails error, Throwable ex) {
var telemetryMessage = TelemetryAction.COMPLETION_ERROR.createActionMessage();
if ("insufficient_quota".equals(error.getCode())) {
telemetryMessage
.property("type", "USER")
.property("code", "INSUFFICIENT_QUOTA");
} else {
telemetryMessage
.property("conversationId", callParameters.getConversation().getId().toString())
.error(new RuntimeException(error.toString(), ex));
}
telemetryMessage.send();
}
}

129
src/main/java/ee/carlrobert/codegpt/completions/CompletionClientProvider.java
View file

@ -1,129 +0,0 @@
package ee.carlrobert.codegpt.completions;
import static ee.carlrobert.codegpt.credentials.CredentialsStore.getCredential;
import com.intellij.openapi.application.ApplicationManager;
import com.intellij.util.net.ssl.CertificateManager;
import ee.carlrobert.codegpt.credentials.CredentialsStore.CredentialKey;
import ee.carlrobert.codegpt.settings.advanced.AdvancedSettings;
import ee.carlrobert.codegpt.settings.service.anthropic.AnthropicSettings;
import ee.carlrobert.codegpt.settings.service.llama.LlamaSettings;
import ee.carlrobert.codegpt.settings.service.ollama.OllamaSettings;
import ee.carlrobert.codegpt.settings.service.openai.OpenAISettings;
import ee.carlrobert.llm.client.anthropic.ClaudeClient;
import ee.carlrobert.llm.client.codegpt.CodeGPTClient;
import ee.carlrobert.llm.client.google.GoogleClient;
import ee.carlrobert.llm.client.inception.InceptionClient;
import ee.carlrobert.llm.client.llama.LlamaClient;
import ee.carlrobert.llm.client.mistral.MistralClient;
import ee.carlrobert.llm.client.ollama.OllamaClient;
import ee.carlrobert.llm.client.openai.OpenAIClient;
import ee.carlrobert.llm.client.openai.completion.ErrorDetails;
import java.net.InetSocketAddress;
import java.net.Proxy;
import java.util.concurrent.TimeUnit;
import java.util.function.Consumer;
import javax.net.ssl.X509TrustManager;
import com.intellij.openapi.diagnostic.Logger;
import okhttp3.Credentials;
import okhttp3.OkHttpClient;
import ee.carlrobert.codegpt.settings.configuration.ConfigurationSettings;
import okhttp3.logging.HttpLoggingInterceptor;
public class CompletionClientProvider {
public static CodeGPTClient getCodeGPTClient() {
return new CodeGPTClient(
getCredential(CredentialKey.CodeGptApiKey.INSTANCE),
getDefaultClientBuilder());
}
public static OpenAIClient getOpenAIClient() {
return new OpenAIClient.Builder(getCredential(CredentialKey.OpenaiApiKey.INSTANCE))
.setOrganization(OpenAISettings.getCurrentState().getOrganization())
.build(getDefaultClientBuilder());
}
public static ClaudeClient getClaudeClient() {
var builder = new ClaudeClient.Builder(getCredential(CredentialKey.AnthropicApiKey.INSTANCE),
AnthropicSettings.getCurrentState().getApiVersion());
if (AnthropicSettings.getCurrentState().hasCustomBaseHost()) {
builder.setHost(AnthropicSettings.getCurrentState().getBaseHost());
}
return builder.build(getDefaultClientBuilder());
}
public static LlamaClient getLlamaClient() {
var llamaSettings = LlamaSettings.getCurrentState();
return new LlamaClient.Builder()
.setPort(llamaSettings.getServerPort())
.build(getDefaultClientBuilder());
}
public static OllamaClient getOllamaClient() {
var host = ApplicationManager.getApplication()
.getService(OllamaSettings.class)
.getState()
.getHost();
var builder = new OllamaClient.Builder()
.setHost(host);
String apiKey = getCredential(CredentialKey.OllamaApikey.INSTANCE);
if (apiKey != null && !apiKey.isBlank()) {
builder.setApiKey(apiKey);
}
return builder.build(getDefaultClientBuilder());
}
public static GoogleClient getGoogleClient() {
return new GoogleClient.Builder(getCredential(CredentialKey.GoogleApiKey.INSTANCE))
.build(getDefaultClientBuilder());
}
public static MistralClient getMistralClient() {
return new MistralClient(getCredential(CredentialKey.MistralApiKey.INSTANCE),
getDefaultClientBuilder());
}
public static InceptionClient getInceptionClient() {
return new InceptionClient.Builder(getCredential(CredentialKey.InceptionApiKey.INSTANCE))
.build(getDefaultClientBuilder());
}
public static OkHttpClient.Builder getDefaultClientBuilder() {
OkHttpClient.Builder builder = new OkHttpClient.Builder();
CertificateManager certificateManager = CertificateManager.getInstance();
X509TrustManager trustManager = certificateManager.getTrustManager();
builder.sslSocketFactory(certificateManager.getSslContext().getSocketFactory(), trustManager);
var advancedSettings = AdvancedSettings.getCurrentState();
var proxyHost = advancedSettings.getProxyHost();
var proxyPort = advancedSettings.getProxyPort();
if (!proxyHost.isEmpty() && proxyPort != 0) {
builder.proxy(
new Proxy(advancedSettings.getProxyType(), new InetSocketAddress(proxyHost, proxyPort)));
if (advancedSettings.isProxyAuthSelected()) {
builder.proxyAuthenticator((route, response) ->
response.request()
.newBuilder()
.header("Proxy-Authorization", Credentials.basic(
advancedSettings.getProxyUsername(),
advancedSettings.getProxyPassword()))
.build());
}
}
if (ConfigurationSettings.getState().getDebugModeEnabled()) {
var ideLogger = Logger.getInstance(CompletionClientProvider.class);
var httpLogger = new HttpLoggingInterceptor(message -> ideLogger.info("[HTTP] " + message));
httpLogger.setLevel(HttpLoggingInterceptor.Level.BODY);
httpLogger.redactHeader("Authorization");
httpLogger.redactHeader("X-API-Key");
httpLogger.redactHeader("Api-Key");
builder.addInterceptor(httpLogger);
}
return builder
.connectTimeout(advancedSettings.getConnectTimeout(), TimeUnit.SECONDS)
.readTimeout(advancedSettings.getReadTimeout(), TimeUnit.SECONDS);
}
}

341
src/main/java/ee/carlrobert/codegpt/completions/CompletionRequestService.java
View file

@ -1,341 +0,0 @@
package ee.carlrobert.codegpt.completions;
import com.intellij.openapi.application.ApplicationManager;
import com.intellij.openapi.components.Service;
import ee.carlrobert.codegpt.completions.factory.CustomOpenAIRequest;
import ee.carlrobert.codegpt.credentials.CredentialsStore;
import ee.carlrobert.codegpt.credentials.CredentialsStore.CredentialKey;
import ee.carlrobert.codegpt.settings.service.FeatureType;
import ee.carlrobert.codegpt.settings.service.ModelSelectionService;
import ee.carlrobert.codegpt.settings.service.ServiceType;
import ee.carlrobert.llm.client.DeserializationUtil;
import ee.carlrobert.llm.client.anthropic.completion.ClaudeCompletionRequest;
import ee.carlrobert.llm.client.codegpt.request.InlineEditRequest;
import ee.carlrobert.llm.client.codegpt.request.chat.ChatCompletionRequest;
import ee.carlrobert.llm.client.google.completion.GoogleCompletionRequest;
import ee.carlrobert.llm.client.openai.completion.ErrorDetails;
import ee.carlrobert.llm.client.openai.completion.OpenAIChatCompletionEventSourceListener;
import ee.carlrobert.llm.client.openai.completion.OpenAITextCompletionEventSourceListener;
import ee.carlrobert.llm.client.openai.completion.request.OpenAIChatCompletionRequest;
import ee.carlrobert.llm.client.openai.completion.response.OpenAIChatCompletionResponse;
import ee.carlrobert.llm.client.openai.completion.response.OpenAIChatCompletionResponseChoice;
import ee.carlrobert.llm.client.openai.completion.response.OpenAIChatCompletionResponseChoiceDelta;
import ee.carlrobert.llm.completion.CompletionEventListener;
import ee.carlrobert.llm.completion.CompletionRequest;
import java.io.IOException;
import java.util.Collection;
import java.util.Objects;
import java.util.Optional;
import java.util.concurrent.ExecutionException;
import java.util.stream.Stream;
import okhttp3.Call;
import okhttp3.Callback;
import okhttp3.Request;
import okhttp3.RequestBody;
import okhttp3.Response;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSources;
import okio.Buffer;
import org.jetbrains.annotations.NotNull;
@Service
public final class CompletionRequestService {
private CompletionRequestService() {
}
public static CompletionRequestService getInstance() {
return ApplicationManager.getApplication().getService(CompletionRequestService.class);
}
public EventSource getCustomOpenAICompletionAsync(
Request customRequest,
CompletionEventListener<String> eventListener) {
var httpClient = CompletionClientProvider.getDefaultClientBuilder().build();
return EventSources.createFactory(httpClient).newEventSource(
customRequest,
new OpenAITextCompletionEventSourceListener(eventListener));
}
public EventSource getCustomOpenAIChatCompletionAsync(
Request customRequest,
CompletionEventListener<String> eventListener) {
if (!isStreamingRequest(customRequest)) {
return getCustomOpenAINonStreamingChatCompletionAsync(customRequest, eventListener);
}
var httpClient = CompletionClientProvider.getDefaultClientBuilder().build();
return EventSources.createFactory(httpClient).newEventSource(
customRequest,
new OpenAIChatCompletionEventSourceListener(eventListener));
}
private EventSource getCustomOpenAINonStreamingChatCompletionAsync(
Request customRequest,
CompletionEventListener<String> eventListener) {
var httpClient = CompletionClientProvider.getDefaultClientBuilder().build();
var call = httpClient.newCall(customRequest);
eventListener.onOpen();
call.enqueue(new Callback() {
@Override
public void onFailure(@NotNull Call call, @NotNull IOException e) {
if (call.isCanceled()) {
eventListener.onCancelled(new StringBuilder());
return;
}
eventListener.onError(new ErrorDetails(e.getMessage()), e);
}
@Override
public void onResponse(@NotNull Call call, @NotNull Response response) {
try (response) {
var completion =
DeserializationUtil.mapResponse(response, OpenAIChatCompletionResponse.class);
var content = tryExtractContent(completion).orElse("");
eventListener.onComplete(new StringBuilder(content));
} catch (Exception e) {
eventListener.onError(new ErrorDetails(e.getMessage()), e);
}
}
});
return new EventSource() {
@Override
public @NotNull Request request() {
return customRequest;
}
@Override
public void cancel() {
call.cancel();
}
};
}
private boolean isStreamingRequest(Request customRequest) {
RequestBody body = customRequest.body();
if (body == null) {
return true;
}
try {
var buffer = new Buffer();
body.writeTo(buffer);
var json = DeserializationUtil.OBJECT_MAPPER.readTree(buffer.readUtf8());
var stream = json.get("stream");
if (stream == null || stream.isNull()) {
return true;
}
if (stream.isBoolean()) {
return stream.booleanValue();
}
if (stream.isTextual()) {
return Boolean.parseBoolean(stream.asText());
}
} catch (Exception ignored) {
return true;
}
return true;
}
public String getLookupCompletion(LookupCompletionParameters params) {
var serviceType =
ModelSelectionService.getInstance().getServiceForFeature(FeatureType.LOOKUP);
var request = CompletionRequestFactory
.getFactory(serviceType)
.createLookupRequest(params);
return getChatCompletion(request, serviceType, FeatureType.LOOKUP);
}
public EventSource autoApplyAsync(
AutoApplyParameters params,
CompletionEventListener<String> eventListener) {
var selectedService =
ModelSelectionService.getInstance().getServiceForFeature(FeatureType.AUTO_APPLY);
var request = CompletionRequestFactory
.getFactory(selectedService)
.createAutoApplyRequest(params);
return getChatCompletionAsync(request, eventListener, selectedService, FeatureType.AUTO_APPLY);
}
public EventSource getCommitMessageAsync(
CommitMessageCompletionParameters params,
CompletionEventListener<String> eventListener) {
var serviceType =
ModelSelectionService.getInstance().getServiceForFeature(FeatureType.COMMIT_MESSAGE);
var request = CompletionRequestFactory
.getFactory(serviceType)
.createCommitMessageRequest(params);
return getChatCompletionAsync(request, eventListener, serviceType, FeatureType.COMMIT_MESSAGE);
}
public EventSource getInlineEditCompletionAsync(
InlineEditCompletionParameters params,
CompletionEventListener<String> eventListener) {
var serviceType =
ModelSelectionService.getInstance().getServiceForFeature(FeatureType.INLINE_EDIT);
var request = CompletionRequestFactory
.getFactory(serviceType)
.createInlineEditRequest(params);
return getChatCompletionAsync(request, eventListener, serviceType, FeatureType.INLINE_EDIT);
}
public EventSource getChatCompletionAsync(
CompletionRequest request,
CompletionEventListener<String> eventListener) {
var serviceType = ModelSelectionService.getInstance().getServiceForFeature(FeatureType.CHAT);
return getChatCompletionAsync(request, eventListener, serviceType, FeatureType.CHAT);
}
public EventSource getChatCompletionAsync(
CompletionRequest request,
CompletionEventListener<String> eventListener,
ServiceType serviceType,
FeatureType featureType) {
if (request instanceof InlineEditRequest completionRequest) {
return CompletionClientProvider.getCodeGPTClient()
.getInlineEditAsync(completionRequest, eventListener);
}
if (request instanceof OpenAIChatCompletionRequest completionRequest) {
return switch (serviceType) {
case OPENAI -> CompletionClientProvider.getOpenAIClient()
.getChatCompletionAsync(completionRequest, eventListener);
case OLLAMA -> CompletionClientProvider.getOllamaClient()
.getChatCompletionAsync(completionRequest, eventListener);
case MISTRAL -> CompletionClientProvider.getMistralClient()
.getChatCompletionAsync(completionRequest, eventListener);
case LLAMA_CPP -> CompletionClientProvider.getLlamaClient()
.getChatCompletionAsync(completionRequest, eventListener);
case INCEPTION -> CompletionClientProvider.getInceptionClient()
.getChatCompletionAsync(completionRequest, eventListener);
default -> throw new RuntimeException("Unknown service selected");
};
}
if (request instanceof ChatCompletionRequest completionRequest) {
return CompletionClientProvider.getCodeGPTClient()
.getChatCompletionAsync(completionRequest, eventListener);
}
if (request instanceof CustomOpenAIRequest completionRequest) {
return getCustomOpenAIChatCompletionAsync(completionRequest.getRequest(), eventListener);
}
if (request instanceof ClaudeCompletionRequest completionRequest) {
return CompletionClientProvider.getClaudeClient().getCompletionAsync(
completionRequest,
eventListener);
}
if (request instanceof GoogleCompletionRequest completionRequest) {
return CompletionClientProvider.getGoogleClient().getChatCompletionAsync(
completionRequest,
ModelSelectionService.getInstance().getModelForFeature(featureType, null),
eventListener);
}
throw new IllegalStateException("Unknown request type: " + request.getClass());
}
public String getChatCompletion(CompletionRequest request, ServiceType serviceType,
FeatureType featureType) {
if (request instanceof OpenAIChatCompletionRequest completionRequest) {
var response = switch (serviceType) {
case OPENAI -> CompletionClientProvider.getOpenAIClient()
.getChatCompletion(completionRequest);
case OLLAMA -> CompletionClientProvider.getOllamaClient()
.getChatCompletion(completionRequest);
case MISTRAL -> CompletionClientProvider.getMistralClient()
.getChatCompletion(completionRequest);
case LLAMA_CPP -> CompletionClientProvider.getLlamaClient()
.getChatCompletion(completionRequest);
case INCEPTION -> CompletionClientProvider.getInceptionClient()
.getChatCompletion(completionRequest);
default -> throw new RuntimeException("Unknown service selected");
};
return tryExtractContent(response).orElseThrow();
}
if (request instanceof ChatCompletionRequest completionRequest) {
var response =
CompletionClientProvider.getCodeGPTClient().getChatCompletion(completionRequest);
return tryExtractContent(response).orElseThrow();
}
if (request instanceof CustomOpenAIRequest completionRequest) {
var httpClient = CompletionClientProvider.getDefaultClientBuilder().build();
try (var response = httpClient.newCall(completionRequest.getRequest()).execute()) {
return DeserializationUtil.mapResponse(response, OpenAIChatCompletionResponse.class)
.getChoices().get(0)
.getMessage()
.getContent();
} catch (IOException e) {
throw new RuntimeException(e);
}
}
if (request instanceof ClaudeCompletionRequest completionRequest) {
return CompletionClientProvider.getClaudeClient()
.getCompletion(completionRequest)
.getContent().get(0)
.getText();
}
if (request instanceof GoogleCompletionRequest completionRequest) {
return CompletionClientProvider.getGoogleClient().getChatCompletion(
completionRequest,
ApplicationManager.getApplication()
.getService(ModelSelectionService.class)
.getModelForFeature(featureType, null))
.getCandidates().get(0)
.getContent().getParts().get(0)
.getText();
}
throw new IllegalStateException("Unknown request type: " + request.getClass());
}
public static boolean isRequestAllowed(FeatureType featureType) {
try {
return ApplicationManager.getApplication().executeOnPooledThread(() -> {
var serviceType = ModelSelectionService.getInstance().getServiceForFeature(featureType);
return isRequestAllowed(serviceType);
})
.get();
} catch (InterruptedException | ExecutionException e) {
throw new RuntimeException(e);
}
}
public static boolean isRequestAllowed(ServiceType serviceType) {
return switch (serviceType) {
case OPENAI -> CredentialsStore.INSTANCE.isCredentialSet(CredentialKey.OpenaiApiKey.INSTANCE);
case ANTHROPIC -> CredentialsStore.INSTANCE.isCredentialSet(
CredentialKey.AnthropicApiKey.INSTANCE
);
case GOOGLE -> CredentialsStore.INSTANCE.isCredentialSet(CredentialKey.GoogleApiKey.INSTANCE);
case MISTRAL ->
CredentialsStore.INSTANCE.isCredentialSet(CredentialKey.MistralApiKey.INSTANCE);
case INCEPTION ->
CredentialsStore.INSTANCE.isCredentialSet(CredentialKey.InceptionApiKey.INSTANCE);
case PROXYAI, CUSTOM_OPENAI, LLAMA_CPP, OLLAMA -> true;
};
}
/**
* Content of the first choice.
* <ul>
* <li>Search all choices which are not null</li>
* <li>Search all messages which are not null</li>
* <li>Use first content which is not null or blank (whitespace)</li>
* </ul>
*
* @return First non-blank content or {@code Optional.empty()}
*/
private Optional<String> tryExtractContent(OpenAIChatCompletionResponse response) {
return Stream.ofNullable(response.getChoices())
.flatMap(Collection::stream)
.filter(Objects::nonNull)
.map(OpenAIChatCompletionResponseChoice::getMessage)
.filter(Objects::nonNull)
.map(OpenAIChatCompletionResponseChoiceDelta::getContent)
.filter(c -> c != null && !c.isBlank())
.findFirst();
}
}

3
src/main/java/ee/carlrobert/codegpt/completions/CompletionResponseEventListener.java
View file

@ -3,7 +3,6 @@ package ee.carlrobert.codegpt.completions;
import ee.carlrobert.codegpt.conversations.Conversation;
import ee.carlrobert.codegpt.conversations.message.Message;
import ee.carlrobert.codegpt.events.CodeGPTEvent;
import ee.carlrobert.llm.client.openai.completion.ErrorDetails;
import javax.swing.JPanel;
public interface CompletionResponseEventListener {
@ -11,7 +10,7 @@ public interface CompletionResponseEventListener {
default void handleMessage(String message) {
}
default void handleError(ErrorDetails error, Throwable ex) {
default void handleError(ChatError error, Throwable ex) {
}
default void handleTokensExceeded(Conversation conversation, Message message) {

103
src/main/java/ee/carlrobert/codegpt/completions/ToolwindowChatCompletionRequestHandler.java
View file

@ -2,29 +2,22 @@ package ee.carlrobert.codegpt.completions;
import com.intellij.openapi.project.Project;
import ee.carlrobert.codegpt.codecompletions.CompletionProgressNotifier;
import ee.carlrobert.codegpt.mcp.McpToolCallEventListener;
import ee.carlrobert.codegpt.settings.service.FeatureType;
import ee.carlrobert.codegpt.settings.service.ModelSelectionService;
import ee.carlrobert.codegpt.settings.models.ModelSettings;
import ee.carlrobert.codegpt.telemetry.TelemetryAction;
import ee.carlrobert.codegpt.toolwindow.chat.ChatToolWindowTabPanel;
import ee.carlrobert.codegpt.toolwindow.ui.ResponseMessagePanel;
import ee.carlrobert.codegpt.toolwindow.ui.mcp.McpApprovalPanel;
import ee.carlrobert.codegpt.toolwindow.ui.mcp.ToolGroupPanel;
import ee.carlrobert.llm.client.openai.completion.ErrorDetails;
import ee.carlrobert.llm.completion.CompletionEventListener;
import ee.carlrobert.codegpt.toolwindow.agent.ui.SimpleAgentApprovalPanel;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.atomic.AtomicBoolean;
import okhttp3.sse.EventSource;
public class ToolwindowChatCompletionRequestHandler {
private final Project project;
private final CompletionResponseEventListener completionResponseEventListener;
private final ChatToolWindowTabPanel tabPanel;
private final List<EventSource> activeEventSources = new ArrayList<>();
private final List<CancellableRequest> activeRequests = new ArrayList<>();
private final AtomicBoolean isCancelled = new AtomicBoolean(false);
private ResponseMessagePanel responseMessagePanel;
public ToolwindowChatCompletionRequestHandler(
Project project,
@ -35,21 +28,17 @@ public class ToolwindowChatCompletionRequestHandler {
this.tabPanel = tabPanel;
}
public void setResponseMessagePanel(ResponseMessagePanel panel) {
this.responseMessagePanel = panel;
}
public void call(ChatCompletionParameters callParameters) {
isCancelled.set(false);
synchronized (activeEventSources) {
activeEventSources.clear();
synchronized (activeRequests) {
activeRequests.clear();
}
try {
EventSource eventSource = startCall(callParameters);
if (eventSource != null) {
synchronized (activeEventSources) {
activeEventSources.add(eventSource);
CancellableRequest request = startCall(callParameters);
if (request != null) {
synchronized (activeRequests) {
activeRequests.add(request);
}
}
} catch (TotalUsageExceededException e) {
@ -64,13 +53,13 @@ public class ToolwindowChatCompletionRequestHandler {
public void cancel() {
isCancelled.set(true);
synchronized (activeEventSources) {
for (EventSource source : activeEventSources) {
if (source != null) {
source.cancel();
synchronized (activeRequests) {
for (CancellableRequest request : activeRequests) {
if (request != null) {
request.cancel();
}
}
activeEventSources.clear();
activeRequests.clear();
}
}
@ -78,54 +67,36 @@ public class ToolwindowChatCompletionRequestHandler {
return isCancelled.get();
}
public void addEventSource(EventSource eventSource) {
if (eventSource != null && !isCancelled.get()) {
synchronized (activeEventSources) {
activeEventSources.add(eventSource);
}
}
}
private CompletionEventListener<String> getEventListener(
ChatCompletionParameters callParameters) {
if (callParameters.getMcpTools() != null && !callParameters.getMcpTools().isEmpty()) {
return new McpToolCallEventListener(
project,
callParameters,
completionResponseEventListener,
panel -> {
if (panel instanceof McpApprovalPanel || panel instanceof ToolGroupPanel) {
tabPanel.addToolCallApprovalPanel(panel);
} else {
tabPanel.addToolCallStatusPanel(panel);
}
return kotlin.Unit.INSTANCE;
},
this);
}
private ChatStreamEventListener getEventListener(ChatCompletionParameters callParameters) {
return new ChatCompletionEventListener(
project,
callParameters,
completionResponseEventListener);
}
private EventSource startCall(ChatCompletionParameters callParameters) {
private CancellableRequest startCall(ChatCompletionParameters callParameters) {
try {
CompletionProgressNotifier.Companion.update(project, true);
var serviceType =
ModelSelectionService.getInstance().getServiceForFeature(FeatureType.CHAT);
var eventListener = getEventListener(callParameters);
var request = CompletionRequestFactory.getFactory(serviceType)
.createChatRequest(callParameters);
try {
return CompletionRequestService.getInstance()
.getChatCompletionAsync(request, eventListener);
} catch (Exception e) {
eventListener.onError(new ErrorDetails("Failed to start request: " + e.getMessage()), e);
return null;
}
ModelSettings.getInstance().getServiceForFeature(FeatureType.CHAT);
var modelSelection =
ModelSettings.getInstance().getModelSelectionForFeature(FeatureType.CHAT);
var factory = CompletionRequestFactory.getFactory(serviceType);
var prompt = factory.createChatCompletionPrompt(callParameters);
return CompletionRequestService.getChatCompletionAsync(
serviceType,
prompt,
modelSelection,
callParameters,
getEventListener(callParameters),
panel -> {
if (panel instanceof SimpleAgentApprovalPanel) {
tabPanel.addToolCallApprovalPanel(panel);
} else {
tabPanel.addToolCallStatusPanel(panel);
}
return kotlin.Unit.INSTANCE;
});
} catch (Throwable ex) {
handleCallException(ex);
}
@ -139,11 +110,11 @@ public class ToolwindowChatCompletionRequestHandler {
"The length of the context exceeds the maximum limit that the model can handle. "
+ "Try reducing the input message or maximum completion token size.";
}
completionResponseEventListener.handleError(new ErrorDetails(errorMessage), ex);
completionResponseEventListener.handleError(new ChatError(errorMessage), ex);
}
private void sendInfo(ChatCompletionParameters callParameters) {
var service = ModelSelectionService.getInstance()
var service = ModelSettings.getInstance()
.getServiceForFeature(FeatureType.CHAT);
TelemetryAction.COMPLETION.createActionMessage()
.property("conversationId", callParameters.getConversation().getId().toString())

8
src/main/java/ee/carlrobert/codegpt/conversations/ConversationService.java
View file

@ -4,9 +4,9 @@ import com.intellij.openapi.application.ApplicationManager;
import com.intellij.openapi.components.Service;
import com.intellij.openapi.diagnostic.Logger;
import com.intellij.openapi.project.Project;
import ee.carlrobert.codegpt.completions.ChatToolCall;
import ee.carlrobert.codegpt.completions.ChatCompletionParameters;
import ee.carlrobert.codegpt.conversations.message.Message;
import ee.carlrobert.llm.client.openai.completion.response.ToolCall;
import java.time.LocalDateTime;
import java.util.Comparator;
import java.util.List;
@ -72,9 +72,9 @@ public final class ConversationService {
saveConversation(conversation);
}
public void saveAssistantMessageWithToolCalls(@NotNull Conversation conversation, @NotNull Message message, List<ToolCall> toolCalls) {
public void saveAssistantMessageWithToolCalls(@NotNull Conversation conversation, @NotNull Message message, List<ChatToolCall> toolCalls) {
if (toolCalls != null && !toolCalls.isEmpty()) {
for (ToolCall toolCall : toolCalls) {
for (ChatToolCall toolCall : toolCalls) {
message.addToolCall(toolCall);
}
}
@ -82,7 +82,7 @@ public final class ConversationService {
saveConversation(conversation);
}
public void saveAssistantMessageWithToolCalls(ChatCompletionParameters callParameters, List<ToolCall> toolCalls) {
public void saveAssistantMessageWithToolCalls(ChatCompletionParameters callParameters, List<ChatToolCall> toolCalls) {
saveAssistantMessageWithToolCalls(callParameters.getConversation(), callParameters.getMessage(), toolCalls);
}

25
src/main/java/ee/carlrobert/codegpt/conversations/message/Message.java
View file

@ -3,12 +3,13 @@ package ee.carlrobert.codegpt.conversations.message;
import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
import com.fasterxml.jackson.annotation.JsonProperty;
import ee.carlrobert.llm.client.openai.completion.response.ToolCall;
import java.util.HashMap;
import ee.carlrobert.codegpt.completions.ChatToolCall;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArrayList;
import org.jetbrains.annotations.Nullable;
@JsonIgnoreProperties(ignoreUnknown = true)
@ -22,7 +23,7 @@ public class Message {
private String imageFilePath;
private boolean webSearchIncluded;
private String personaName;
private List<ToolCall> toolCalls;
private List<ChatToolCall> toolCalls;
private Map<String, String> toolCallResults;
public Message() {
@ -101,18 +102,18 @@ public class Message {
}
@JsonProperty("tool_calls")
public @Nullable List<ToolCall> getToolCalls() {
public @Nullable List<ChatToolCall> getToolCalls() {
return toolCalls;
}
@JsonProperty("tool_calls")
public void setToolCalls(@Nullable List<ToolCall> toolCalls) {
this.toolCalls = toolCalls;
public void setToolCalls(@Nullable List<ChatToolCall> toolCalls) {
this.toolCalls = toolCalls == null ? null : new CopyOnWriteArrayList<>(toolCalls);
}
public void addToolCall(ToolCall toolCall) {
public synchronized void addToolCall(ChatToolCall toolCall) {
if (toolCalls == null) {
toolCalls = new java.util.ArrayList<>();
toolCalls = new CopyOnWriteArrayList<>();
}
toolCalls.add(toolCall);
}
@ -128,12 +129,12 @@ public class Message {
@JsonProperty("tool_call_results")
public void setToolCallResults(@Nullable Map<String, String> toolCallResults) {
this.toolCallResults = toolCallResults;
this.toolCallResults = toolCallResults == null ? null : new ConcurrentHashMap<>(toolCallResults);
}
public void addToolCallResult(String toolCallId, String executionOutput) {
public synchronized void addToolCallResult(String toolCallId, String executionOutput) {
if (toolCallResults == null) {
toolCallResults = new HashMap<>();
toolCallResults = new ConcurrentHashMap<>();
}
toolCallResults.put(toolCallId, executionOutput);
}
@ -151,6 +152,6 @@ public class Message {
@Override
public int hashCode() {
return Objects.hash(id, prompt);
return Objects.hash(id);
}
}

4
src/main/java/ee/carlrobert/codegpt/settings/service/llama/LlamaSettings.java
View file

@ -13,7 +13,7 @@ import ee.carlrobert.codegpt.completions.HuggingFaceModel;
import ee.carlrobert.codegpt.completions.llama.LlamaModel;
import ee.carlrobert.codegpt.completions.llama.PromptTemplate;
import ee.carlrobert.codegpt.settings.service.FeatureType;
import ee.carlrobert.codegpt.settings.service.ModelSelectionService;
import ee.carlrobert.codegpt.settings.models.ModelSettings;
import ee.carlrobert.codegpt.settings.service.llama.form.LlamaSettingsForm;
import java.io.File;
import java.nio.file.Path;
@ -69,7 +69,7 @@ public class LlamaSettings implements PersistentStateComponent<LlamaSettingsStat
public static boolean isRunnable() {
return (IS_OS_MAC_OSX || IS_OS_LINUX)
&& ModelSelectionService.getInstance().getServiceForFeature(FeatureType.CHAT) == LLAMA_CPP;
&& ModelSettings.getInstance().getServiceForFeature(FeatureType.CHAT) == LLAMA_CPP;
}
public static boolean isRunnable(HuggingFaceModel model) {

3
src/main/java/ee/carlrobert/codegpt/settings/service/openai/OpenAISettingsState.java
View file

@ -1,12 +1,11 @@
package ee.carlrobert.codegpt.settings.service.openai;
import ee.carlrobert.llm.client.openai.completion.OpenAIChatCompletionModel;
import java.util.Objects;
public class OpenAISettingsState {
private String organization = "";
private String model = OpenAIChatCompletionModel.GPT_4_1.getCode();
private String model = "gpt-4.1";
private boolean codeCompletionsEnabled = false;
public String getOrganization() {

13
src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ChatToolWindowPanel.java
View file

@ -26,7 +26,7 @@ import ee.carlrobert.codegpt.conversations.Conversation;
import ee.carlrobert.codegpt.conversations.ConversationService;
import ee.carlrobert.codegpt.conversations.ConversationsState;
import ee.carlrobert.codegpt.settings.service.FeatureType;
import ee.carlrobert.codegpt.settings.service.ModelSelectionService;
import ee.carlrobert.codegpt.settings.models.ModelSettings;
import ee.carlrobert.codegpt.settings.prompts.PersonaPromptDetailsState;
import ee.carlrobert.codegpt.settings.prompts.PromptsConfigurable;
import ee.carlrobert.codegpt.settings.prompts.PromptsSettings;
@ -35,7 +35,6 @@ import ee.carlrobert.codegpt.settings.service.ServiceType;
import ee.carlrobert.codegpt.settings.service.codegpt.CodeGPTUserDetailsNotifier;
import ee.carlrobert.codegpt.toolwindow.chat.ui.ToolWindowFooterNotification;
import ee.carlrobert.codegpt.toolwindow.chat.ui.textarea.AttachImageNotifier;
import ee.carlrobert.llm.client.codegpt.PricingPlan;
import java.nio.file.Path;
import javax.swing.JComponent;
import org.jetbrains.annotations.NotNull;
@ -89,8 +88,7 @@ public class ChatToolWindowPanel extends SimpleToolWindowPanel {
(ProviderChangeNotifier) provider -> {
if (provider == ServiceType.PROXYAI) {
var userDetails = CodeGPTKeys.CODEGPT_USER_DETAILS.get(project);
upgradePlanLink.setVisible(
userDetails != null && userDetails.getPricingPlan() != PricingPlan.INDIVIDUAL);
upgradePlanLink.setVisible(userDetails != null);
} else {
upgradePlanLink.setVisible(false);
}
@ -98,10 +96,9 @@ public class ChatToolWindowPanel extends SimpleToolWindowPanel {
messageBusConnection.subscribe(CodeGPTUserDetailsNotifier.getCODEGPT_USER_DETAILS_TOPIC(),
(CodeGPTUserDetailsNotifier) userDetails -> {
if (userDetails != null) {
var provider = ModelSelectionService.getInstance()
var provider = ModelSettings.getInstance()
.getServiceForFeature(FeatureType.CHAT);
upgradePlanLink.setVisible(provider == ServiceType.PROXYAI
&& userDetails.getPricingPlan() != PricingPlan.INDIVIDUAL);
upgradePlanLink.setVisible(provider == ServiceType.PROXYAI);
}
});
}
@ -224,4 +221,4 @@ public class ChatToolWindowPanel extends SimpleToolWindowPanel {
.getSelectedPersona();
}
}
}
}

9
src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ChatToolWindowTabPanel.java
View file

@ -22,8 +22,9 @@ import ee.carlrobert.codegpt.ReferencedFile;
import ee.carlrobert.codegpt.actions.ActionType;
import ee.carlrobert.codegpt.codecompletions.CompletionProgressNotifier;
import ee.carlrobert.codegpt.completions.ChatCompletionParameters;
import ee.carlrobert.codegpt.completions.CompletionRequestService;
import ee.carlrobert.codegpt.completions.ChatError;
import ee.carlrobert.codegpt.completions.CompletionRequestUtil;
import ee.carlrobert.codegpt.completions.CompletionService;
import ee.carlrobert.codegpt.completions.ConversationType;
import ee.carlrobert.codegpt.completions.ToolApprovalMode;
import ee.carlrobert.codegpt.completions.ToolwindowChatCompletionRequestHandler;
@ -61,7 +62,6 @@ import ee.carlrobert.codegpt.ui.textarea.header.tag.TagDetails;
import ee.carlrobert.codegpt.ui.textarea.header.tag.TagManager;
import ee.carlrobert.codegpt.util.EditorUtil;
import ee.carlrobert.codegpt.util.coroutines.CoroutineDispatchers;
import ee.carlrobert.llm.client.openai.completion.ErrorDetails;
import java.awt.BorderLayout;
import java.awt.GridBagConstraints;
import java.awt.GridBagLayout;
@ -526,7 +526,7 @@ public class ChatToolWindowTabPanel implements Disposable {
return;
}
if (!CompletionRequestService.isRequestAllowed(FeatureType.CHAT)) {
if (!CompletionService.isRequestAllowed(FeatureType.CHAT)) {
responseContent.displayMissingCredential();
return;
}
@ -560,7 +560,7 @@ public class ChatToolWindowTabPanel implements Disposable {
}
@Override
public void handleError(ErrorDetails error, Throwable ex) {
public void handleError(ChatError error, Throwable ex) {
try {
super.handleError(error, ex);
} finally {
@ -575,7 +575,6 @@ public class ChatToolWindowTabPanel implements Disposable {
},
this);
requestHandler.setResponseMessagePanel(responseMessagePanel);
showInputLoading(CodeGPTBundle.get("toolwindow.chat.loading"));
requestHandler.call(callParameters);
}

39
src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ToolWindowCompletionResponseEventListener.java
View file

@ -7,6 +7,7 @@ import com.intellij.openapi.project.Project;
import ee.carlrobert.codegpt.EncodingManager;
import ee.carlrobert.codegpt.codecompletions.CompletionProgressNotifier;
import ee.carlrobert.codegpt.completions.ChatCompletionParameters;
import ee.carlrobert.codegpt.completions.ChatError;
import ee.carlrobert.codegpt.completions.CompletionResponseEventListener;
import ee.carlrobert.codegpt.conversations.Conversation;
import ee.carlrobert.codegpt.conversations.ConversationService;
@ -19,7 +20,6 @@ import ee.carlrobert.codegpt.toolwindow.ui.ResponseMessagePanel;
import ee.carlrobert.codegpt.toolwindow.ui.UserMessagePanel;
import ee.carlrobert.codegpt.ui.OverlayUtil;
import ee.carlrobert.codegpt.ui.textarea.UserInputPanel;
import ee.carlrobert.llm.client.openai.completion.ErrorDetails;
import java.util.concurrent.ConcurrentLinkedQueue;
import javax.swing.Timer;
@ -29,7 +29,6 @@ abstract class ToolWindowCompletionResponseEventListener implements
private static final int UPDATE_INTERVAL_MS = 20;
private final Project project;
private final StringBuilder messageBuilder = new StringBuilder();
private final EncodingManager encodingManager;
private final ResponseMessagePanel responsePanel;
private final UserMessagePanel userMessagePanel;
@ -60,10 +59,7 @@ abstract class ToolWindowCompletionResponseEventListener implements
private ChatMessageResponseBody getResponseContainer() {
if (responseContainer == null && responsePanel != null) {
var content = responsePanel.getResponseComponent();
if (content instanceof ChatMessageResponseBody) {
responseContainer = (ChatMessageResponseBody) content;
}
responseContainer = responsePanel.getResponseComponent();
}
return responseContainer;
}
@ -78,7 +74,6 @@ abstract class ToolWindowCompletionResponseEventListener implements
streamResponseReceived = true;
try {
messageBuilder.append(partialMessage);
var ongoingTokens = encodingManager.countTokens(partialMessage);
messageBuffer.offer(partialMessage);
ApplicationManager.getApplication().invokeLater(() ->
@ -94,7 +89,7 @@ abstract class ToolWindowCompletionResponseEventListener implements
}
@Override
public void handleError(ErrorDetails error, Throwable ex) {
public void handleError(ChatError error, Throwable ex) {
ApplicationManager.getApplication().invokeLater(() -> {
try {
if ("insufficient_quota".equals(error.getCode())) {
@ -133,23 +128,23 @@ abstract class ToolWindowCompletionResponseEventListener implements
@Override
public void handleCompleted(String fullMessage, ChatCompletionParameters callParameters) {
if (fullMessage != null) {
ConversationService.getInstance().saveMessage(fullMessage, callParameters);
}
ApplicationManager.getApplication().invokeLater(() -> {
try {
responsePanel.enableAllActions(true);
if (!streamResponseReceived && !fullMessage.isEmpty()) {
var container = getResponseContainer();
if (container != null) {
container.withResponse(fullMessage);
if (fullMessage != null) {
ConversationService.getInstance().saveMessage(fullMessage, callParameters);
try {
responsePanel.enableAllActions(true);
if (!streamResponseReceived && !fullMessage.isEmpty()) {
var container = getResponseContainer();
if (container != null) {
container.withResponse(fullMessage);
}
}
totalTokensPanel.updateUserPromptTokens(userInputPanel.getText());
totalTokensPanel.updateConversationTokens(callParameters.getConversation());
} finally {
stopStreaming(responseContainer);
}
totalTokensPanel.updateUserPromptTokens(userInputPanel.getText());
totalTokensPanel.updateConversationTokens(callParameters.getConversation());
} finally {
stopStreaming(responseContainer);
}
});
}

4
src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/ChatMessageResponseBody.java
View file

@ -35,7 +35,7 @@ import ee.carlrobert.codegpt.events.EventDetails;
import ee.carlrobert.codegpt.events.WebSearchEventDetails;
import ee.carlrobert.codegpt.settings.GeneralSettingsConfigurable;
import ee.carlrobert.codegpt.settings.service.FeatureType;
import ee.carlrobert.codegpt.settings.service.ModelSelectionService;
import ee.carlrobert.codegpt.settings.models.ModelSettings;
import ee.carlrobert.codegpt.settings.service.ServiceType;
import ee.carlrobert.codegpt.telemetry.TelemetryAction;
import ee.carlrobert.codegpt.toolwindow.chat.editor.ResponseEditorPanel;
@ -115,7 +115,7 @@ public class ChatMessageResponseBody extends JPanel {
contentPanel.setOpaque(false);
add(contentPanel, BorderLayout.CENTER);
if (ModelSelectionService.getInstance().getServiceForFeature(FeatureType.CHAT)
if (ModelSettings.getInstance().getServiceForFeature(FeatureType.CHAT)
== ServiceType.PROXYAI) {
if (withProgress) {
contentPanel.add(progressPanel);

4
src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/ChatToolWindowScrollablePanel.java
View file

@ -10,7 +10,7 @@ import com.intellij.util.ui.JBUI;
import ee.carlrobert.codegpt.credentials.CredentialsStore;
import ee.carlrobert.codegpt.credentials.CredentialsStore.CredentialKey;
import ee.carlrobert.codegpt.settings.service.FeatureType;
import ee.carlrobert.codegpt.settings.service.ModelSelectionService;
import ee.carlrobert.codegpt.settings.models.ModelSettings;
import ee.carlrobert.codegpt.settings.service.ServiceType;
import ee.carlrobert.codegpt.settings.service.codegpt.CodeGPTServiceConfigurable;
import ee.carlrobert.codegpt.toolwindow.ui.ResponseMessagePanel;
@ -38,7 +38,7 @@ public class ChatToolWindowScrollablePanel extends ScrollablePanel {
clearAll();
add(landingView);
landingViewVisible = true;
if (ModelSelectionService.getInstance().getServiceForFeature(FeatureType.CHAT)
if (ModelSettings.getInstance().getServiceForFeature(FeatureType.CHAT)
== ServiceType.PROXYAI
&& !CredentialsStore.INSTANCE.isCredentialSet(CredentialKey.CodeGptApiKey.INSTANCE)) {

327
src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/textarea/ModelComboBoxAction.java
View file

@ -1,6 +1,5 @@
package ee.carlrobert.codegpt.toolwindow.chat.ui.textarea;
import static ee.carlrobert.codegpt.CodeGPTKeys.CODEGPT_USER_DETAILS;
import static ee.carlrobert.codegpt.settings.service.ServiceType.ANTHROPIC;
import static ee.carlrobert.codegpt.settings.service.ServiceType.CUSTOM_OPENAI;
import static ee.carlrobert.codegpt.settings.service.ServiceType.GOOGLE;
@ -32,7 +31,6 @@ import com.intellij.openapi.ui.popup.ListPopup;
import ee.carlrobert.codegpt.Icons;
import ee.carlrobert.codegpt.completions.llama.LlamaModel;
import ee.carlrobert.codegpt.settings.models.ModelDetailsState;
import ee.carlrobert.codegpt.settings.models.ModelRegistry;
import ee.carlrobert.codegpt.settings.models.ModelSelection;
import ee.carlrobert.codegpt.settings.models.ModelSettings;
import ee.carlrobert.codegpt.settings.models.ModelSettingsConfigurable;
@ -40,17 +38,11 @@ import ee.carlrobert.codegpt.settings.service.FeatureType;
import ee.carlrobert.codegpt.settings.service.ModelChangeNotifier;
import ee.carlrobert.codegpt.settings.service.ModelChangeNotifierAdapter;
import ee.carlrobert.codegpt.settings.service.ServiceType;
import ee.carlrobert.codegpt.settings.service.custom.CustomServiceSettingsState;
import ee.carlrobert.codegpt.settings.service.custom.CustomServicesSettings;
import ee.carlrobert.codegpt.settings.service.llama.LlamaSettings;
import ee.carlrobert.codegpt.settings.service.ollama.OllamaSettings;
import ee.carlrobert.codegpt.toolwindow.ui.CodeGPTModelsListPopupAction;
import ee.carlrobert.codegpt.toolwindow.ui.ModelListPopup;
import ee.carlrobert.llm.client.codegpt.PricingPlan;
import ee.carlrobert.llm.client.google.models.GoogleModel;
import ee.carlrobert.llm.client.openai.completion.OpenAIChatCompletionModel;
import java.awt.Color;
import java.util.Arrays;
import java.util.List;
import java.util.function.Consumer;
import javax.swing.Icon;
@ -63,37 +55,16 @@ public class ModelComboBoxAction extends ComboBoxAction {
private static final Logger LOG = Logger.getInstance(ModelComboBoxAction.class);
private final Consumer<ServiceType> onModelChange;
private final Project project;
private final List<ServiceType> availableProviders;
private final boolean showConfigureModels;
private final FeatureType featureType;
public ModelComboBoxAction(
Project project,
Consumer<ServiceType> onModelChange,
ServiceType selectedService) {
this(project, onModelChange, selectedService, Arrays.asList(ServiceType.values()), true,
FeatureType.CHAT);
}
public ModelComboBoxAction(
Project project,
Consumer<ServiceType> onModelChange,
ServiceType selectedProvider,
List<ServiceType> availableProviders,
boolean showConfigureModels) {
this(project, onModelChange, selectedProvider, availableProviders, showConfigureModels,
FeatureType.CHAT);
}
public ModelComboBoxAction(
Project project,
Consumer<ServiceType> onModelChange,
ServiceType selectedProvider,
List<ServiceType> availableProviders,
boolean showConfigureModels,
FeatureType featureType) {
this.project = project;
this.onModelChange = onModelChange;
this.availableProviders = availableProviders;
this.showConfigureModels = showConfigureModels;
@ -149,18 +120,9 @@ public class ModelComboBoxAction extends ComboBoxAction {
return popup;
}
private AnAction[] getCodeGPTModelActions(Project project, Presentation presentation) {
var registry = ModelRegistry.getInstance();
if (featureType == FeatureType.AGENT) {
return registry.getAllModelsForFeature(featureType).stream()
.filter(model -> model.getProvider() == PROXYAI)
.map(model -> createCodeGPTModelAction(model, presentation))
.toArray(AnAction[]::new);
}
var userDetails = CODEGPT_USER_DETAILS.get(project);
return registry.getProxyAIChatModelsForPricingPlan(
userDetails == null ? null : userDetails.getPricingPlan()).stream()
private AnAction[] getProxyAIModelActions(Presentation presentation) {
return getAvailableModelsForFeature().stream()
.filter(model -> model.getProvider() == PROXYAI)
.map(model -> createCodeGPTModelAction(model, presentation))
.toArray(AnAction[]::new);
}
@ -175,85 +137,93 @@ public class ModelComboBoxAction extends ComboBoxAction {
if (availableProviders.contains(PROXYAI)) {
var proxyAIGroup = DefaultActionGroup.createPopupGroup(() -> "ProxyAI");
proxyAIGroup.getTemplatePresentation().setIcon(Icons.DefaultSmall);
proxyAIGroup.addAll(getCodeGPTModelActions(project, presentation));
proxyAIGroup.addAll(getProxyAIModelActions(presentation));
actionGroup.add(proxyAIGroup);
}
if (availableProviders.contains(ANTHROPIC)) {
var anthropicGroup = DefaultActionGroup.createPopupGroup(() -> "Anthropic");
anthropicGroup.getTemplatePresentation().setIcon(Icons.Anthropic);
ModelRegistry.getInstance().getAgentModels(ANTHROPIC).forEach(item ->
anthropicGroup.add(createModelAction(
ANTHROPIC,
item.getName(),
Icons.Anthropic,
presentation,
() -> ApplicationManager.getApplication().getService(ModelSettings.class)
.setModel(featureType, item.getModel().getId(), ANTHROPIC))));
getAvailableModelsForFeature().stream()
.filter(model -> model.getProvider() == ANTHROPIC)
.forEach(item -> {
anthropicGroup.add(createModelAction(
ANTHROPIC,
item.getDisplayName(),
Icons.Anthropic,
presentation,
() -> ApplicationManager.getApplication().getService(ModelSettings.class)
.setModel(featureType, item.getModel(), ANTHROPIC)));
});
actionGroup.add(anthropicGroup);
}
if (availableProviders.contains(OPENAI)) {
var openaiGroup = DefaultActionGroup.createPopupGroup(() -> "OpenAI");
openaiGroup.getTemplatePresentation().setIcon(Icons.OpenAI);
if (featureType == FeatureType.AGENT) {
addOpenAIGroupForAgent(openaiGroup, presentation);
} else {
addOpenAIGroupForChat(openaiGroup, presentation);
}
getAvailableModelsForFeature().stream()
.filter(model -> model.getProvider() == OPENAI)
.forEach(item -> {
openaiGroup.add(createModelAction(
OPENAI,
item.getDisplayName(),
Icons.OpenAI,
presentation,
() -> ApplicationManager.getApplication().getService(ModelSettings.class)
.setModel(featureType, item.getModel(), OPENAI)));
});
actionGroup.add(openaiGroup);
}
if (availableProviders.contains(CUSTOM_OPENAI)) {
List<CustomServiceSettingsState> services = ApplicationManager.getApplication()
.getService(CustomServicesSettings.class)
.getState()
.getServices();
var customGroup = DefaultActionGroup.createPopupGroup(() -> "Custom OpenAI");
customGroup.getTemplatePresentation().setIcon(Icons.OpenAI);
services.forEach(model ->
customGroup.add(createCustomOpenAIModelAction(model, presentation))
);
getAvailableModelsForFeature().stream()
.filter(model -> model.getProvider() == CUSTOM_OPENAI)
.forEach(model -> customGroup.add(createModelAction(
CUSTOM_OPENAI,
model.getDisplayName(),
Icons.OpenAI,
presentation,
() -> ApplicationManager.getApplication().getService(ModelSettings.class)
.setModel(featureType, model.getModel(), CUSTOM_OPENAI))));
actionGroup.add(customGroup);
}
if (availableProviders.contains(GOOGLE)) {
var googleGroup = DefaultActionGroup.createPopupGroup(() -> "Google");
googleGroup.getTemplatePresentation().setIcon(Icons.Google);
ModelRegistry.getInstance().getAgentModels(GOOGLE).forEach(item ->
googleGroup.add(createModelAction(
getAvailableModelsForFeature().stream()
.filter(model -> model.getProvider() == GOOGLE)
.forEach(item -> googleGroup.add(createModelAction(
GOOGLE,
item.getName(),
item.getDisplayName(),
Icons.Google,
presentation,
() -> {
var application = ApplicationManager.getApplication();
application.getService(ModelSettings.class)
.setModel(featureType, item.getModel().getId(), GOOGLE);
})));
() -> ApplicationManager.getApplication().getService(ModelSettings.class)
.setModel(featureType, item.getModel(), GOOGLE))));
actionGroup.add(googleGroup);
}
if (availableProviders.contains(MISTRAL)) {
var mistralGroup = DefaultActionGroup.createPopupGroup(() -> "Mistral");
mistralGroup.getTemplatePresentation().setIcon(Icons.Mistral);
ModelRegistry.getInstance().getAgentModels(MISTRAL).forEach(item ->
mistralGroup.add(createModelAction(
getAvailableModelsForFeature().stream()
.filter(model -> model.getProvider() == MISTRAL)
.forEach(item -> mistralGroup.add(createModelAction(
MISTRAL,
item.getName(),
item.getDisplayName(),
Icons.Mistral,
presentation,
() -> ApplicationManager.getApplication().getService(ModelSettings.class)
.setModel(featureType, item.getModel().getId(), MISTRAL))));
.setModel(featureType, item.getModel(), MISTRAL))));
actionGroup.add(mistralGroup);
}
if (availableProviders.contains(INCEPTION)) {
var inceptionGroup = DefaultActionGroup.createPopupGroup(() -> "Inception");
inceptionGroup.getTemplatePresentation().setIcon(Icons.Inception);
inceptionGroup.add(createInceptionModelAction(ModelRegistry.MERCURY, presentation));
inceptionGroup.add(createInceptionModelAction(presentation));
actionGroup.add(inceptionGroup);
}
@ -312,7 +282,7 @@ public class ModelComboBoxAction extends ComboBoxAction {
switch (selectedService) {
case PROXYAI:
var proxyAIModel = ModelRegistry.getInstance().getAllModelsForFeature(featureType).stream()
var proxyAIModel = getAvailableModelsForFeature().stream()
.filter(it -> it.getProvider() == PROXYAI)
.filter(it -> modelCode != null && it.getModel().equals(modelCode))
.findFirst();
@ -323,36 +293,36 @@ public class ModelComboBoxAction extends ComboBoxAction {
break;
case OPENAI:
templatePresentation.setIcon(Icons.OpenAI);
var openAIModelName = ModelRegistry.getInstance().getModelDisplayName(OPENAI, modelCode);
var openAIModelName = getModelSettings().getModelDisplayName(OPENAI, modelCode);
templatePresentation.setText(openAIModelName);
break;
case CUSTOM_OPENAI:
ModelRegistry.getInstance().getCustomOpenAIModels().stream()
.filter(it -> {
var id = it.getId();
return id != null && id.equals(modelCode);
})
getAvailableModelsForFeature().stream()
.filter(it -> it.getProvider() == CUSTOM_OPENAI)
.filter(it -> modelCode != null && modelCode.equals(it.getModel()))
.findFirst()
.ifPresent(selection -> {
templatePresentation.setIcon(Icons.OpenAI);
templatePresentation.setText(selection.getDisplayName());
});
.ifPresentOrElse(selection -> {
templatePresentation.setIcon(Icons.OpenAI);
templatePresentation.setText(selection.getDisplayName());
},
() -> {
templatePresentation.setIcon(Icons.OpenAI);
templatePresentation.setText(
getModelSettings().getModelDisplayName(CUSTOM_OPENAI, modelCode));
});
break;
case ANTHROPIC:
templatePresentation.setIcon(Icons.Anthropic);
var anthropicModelName = ModelRegistry.getInstance()
.getModelDisplayName(ANTHROPIC, modelCode);
var anthropicModelName = getModelSettings().getModelDisplayName(ANTHROPIC, modelCode);
templatePresentation.setText(anthropicModelName);
break;
case LLAMA_CPP:
templatePresentation.setText(getLlamaCppPresentationText());
templatePresentation.setText(getModelSettings().getModelDisplayName(LLAMA_CPP, modelCode));
templatePresentation.setIcon(Icons.Llama);
break;
case OLLAMA:
templatePresentation.setIcon(Icons.Ollama);
templatePresentation.setText(application.getService(OllamaSettings.class)
.getState()
.getModel());
templatePresentation.setText(getModelSettings().getModelDisplayName(OLLAMA, modelCode));
break;
case GOOGLE:
templatePresentation.setText(getGooglePresentationText());
@ -364,8 +334,7 @@ public class ModelComboBoxAction extends ComboBoxAction {
break;
case INCEPTION:
templatePresentation.setIcon(Icons.Inception);
var inceptionModelName = ModelRegistry.getInstance()
.getModelDisplayName(INCEPTION, modelCode);
var inceptionModelName = getModelSettings().getModelDisplayName(INCEPTION, modelCode);
templatePresentation.setText(inceptionModelName);
break;
default:
@ -378,12 +347,16 @@ public class ModelComboBoxAction extends ComboBoxAction {
.getService(ModelSettings.class)
.getState()
.getModelSelection(featureType);
return ModelRegistry.getInstance().getModelDisplayName(GOOGLE, getGoogleModelCode(chatModel));
return getModelSettings().getModelDisplayName(GOOGLE, getGoogleModelCode(chatModel));
}
private String getGoogleModelCode(@Nullable ModelDetailsState chatModel) {
if (chatModel == null || chatModel.getModel() == null || chatModel.getModel().isBlank()) {
return ModelRegistry.GEMINI_PRO_2_5;
return getAvailableModelsForFeature().stream()
.filter(model -> model.getProvider() == GOOGLE)
.map(ModelSelection::getModel)
.findFirst()
.orElse("");
}
return chatModel.getModel();
@ -434,36 +407,20 @@ public class ModelComboBoxAction extends ComboBoxAction {
private AnAction createCodeGPTModelAction(ModelSelection model,
Presentation comboBoxPresentation) {
var selected = model.getDisplayName().equals(comboBoxPresentation.getText());
var locked = shouldLockModel(model);
return new CodeGPTModelsListPopupAction(model.getDisplayName(), model.getModel(),
var selected = isModelSelected(PROXYAI, model.getModel());
return new CodeGPTModelsListPopupAction(
model.getDisplayName(),
model.getModel(),
model.getIcon() != null ? model.getIcon() : Icons.DefaultSmall,
model.getPricingPlan() != null ? model.getPricingPlan() : PricingPlan.ANONYMOUS,
locked, selected, () -> {
var application = ApplicationManager.getApplication();
application
false,
selected,
() -> {
ApplicationManager.getApplication()
.getService(ModelSettings.class)
.setModel(featureType, model.getModel(), PROXYAI);
handleModelChange(PROXYAI);
});
}
private boolean shouldLockModel(ModelSelection model) {
if (featureType == FeatureType.AGENT) {
return false;
}
var userDetails = CODEGPT_USER_DETAILS.get(project);
if (userDetails != null) {
if (userDetails.getPricingPlan() == PricingPlan.INDIVIDUAL) {
return false;
}
if (userDetails.getPricingPlan() == PricingPlan.FREE) {
return model.getPricingPlan() == PricingPlan.INDIVIDUAL;
}
}
return model.getPricingPlan() != PricingPlan.ANONYMOUS;
});
}
private AnAction createOllamaModelAction(String model, Presentation comboBoxPresentation) {
@ -480,92 +437,6 @@ public class ModelComboBoxAction extends ComboBoxAction {
});
}
private AnAction createOpenAIModelAction(
OpenAIChatCompletionModel model,
Presentation comboBoxPresentation) {
var modelName = ModelRegistry.getInstance().getModelDisplayName(OPENAI, model.getCode());
return createModelAction(
OPENAI,
modelName,
Icons.OpenAI,
comboBoxPresentation,
() -> ApplicationManager.getApplication().getService(ModelSettings.class)
.setModel(featureType, model.getCode(), OPENAI));
}
private void addOpenAIGroupForAgent(DefaultActionGroup openaiGroup, Presentation presentation) {
ModelRegistry.getInstance().getAgentModels(OPENAI).forEach(item ->
openaiGroup.add(createModelAction(
OPENAI,
item.getName(),
Icons.OpenAI,
presentation,
() -> ApplicationManager.getApplication().getService(ModelSettings.class)
.setModel(featureType, item.getModel().getId(), OPENAI))));
}
private void addOpenAIGroupForChat(DefaultActionGroup openaiGroup, Presentation presentation) {
List.of(
OpenAIChatCompletionModel.GPT_5,
OpenAIChatCompletionModel.GPT_5_MINI,
OpenAIChatCompletionModel.O_4_MINI,
OpenAIChatCompletionModel.O_3,
OpenAIChatCompletionModel.O_3_MINI,
OpenAIChatCompletionModel.O_1_PREVIEW,
OpenAIChatCompletionModel.O_1_MINI,
OpenAIChatCompletionModel.GPT_4_1,
OpenAIChatCompletionModel.GPT_4_1_MINI,
OpenAIChatCompletionModel.GPT_4_1_NANO,
OpenAIChatCompletionModel.GPT_4_O,
OpenAIChatCompletionModel.GPT_4_O_MINI,
OpenAIChatCompletionModel.GPT_4_0125_128k)
.forEach(model -> openaiGroup.add(createOpenAIModelAction(model, presentation)));
}
private AnAction createCustomOpenAIModelAction(
CustomServiceSettingsState state,
Presentation comboBoxPresentation) {
var model = state.getChatCompletionSettings().getBody().get("model");
var displayName =
state.getName()
+ ((model instanceof String && !((String) model).isEmpty()) ? " (" + model + ")" : "");
return createModelAction(
CUSTOM_OPENAI,
displayName,
Icons.OpenAI,
comboBoxPresentation,
() -> ApplicationManager.getApplication().getService(ModelSettings.class)
.setModel(featureType, state.getId(), CUSTOM_OPENAI));
}
private AnAction createGoogleModelAction(GoogleModel model, Presentation comboBoxPresentation) {
var modelName = ModelRegistry.getInstance().getModelDisplayName(GOOGLE, model.getCode());
return createModelAction(
GOOGLE,
modelName,
Icons.Google,
comboBoxPresentation,
() -> {
var application = ApplicationManager.getApplication();
application.getService(ModelSettings.class)
.setModel(featureType, model.getCode(), GOOGLE);
});
}
private AnAction createAnthropicModelAction(
String modelCode,
Presentation comboBoxPresentation) {
var modelName = ModelRegistry.getInstance().getModelDisplayName(ANTHROPIC, modelCode);
return createModelAction(
ANTHROPIC,
modelName,
Icons.Anthropic,
comboBoxPresentation,
() -> ApplicationManager.getApplication().getService(ModelSettings.class)
.setModel(featureType, modelCode, ANTHROPIC));
}
private AnAction createLlamaModelAction(Presentation comboBoxPresentation) {
return createModelAction(
LLAMA_CPP,
@ -577,19 +448,13 @@ public class ModelComboBoxAction extends ComboBoxAction {
LlamaSettings.getCurrentState().getHuggingFaceModel().getCode(), LLAMA_CPP));
}
private AnAction createMistralModelAction(String modelCode, Presentation comboBoxPresentation) {
var modelName = ModelRegistry.getInstance().getModelDisplayName(MISTRAL, modelCode);
return createModelAction(
MISTRAL,
modelName,
Icons.Mistral,
comboBoxPresentation,
() -> ApplicationManager.getApplication().getService(ModelSettings.class)
.setModel(featureType, modelCode, MISTRAL));
}
private AnAction createInceptionModelAction(String modelCode, Presentation comboBoxPresentation) {
var modelName = ModelRegistry.getInstance().getModelDisplayName(INCEPTION, modelCode);
private AnAction createInceptionModelAction(Presentation comboBoxPresentation) {
var modelCode = getAvailableModelsForFeature().stream()
.filter(model -> model.getProvider() == INCEPTION)
.map(ModelSelection::getModel)
.findFirst()
.orElse("mercury");
var modelName = getModelSettings().getModelDisplayName(INCEPTION, modelCode);
return createModelAction(
INCEPTION,
modelName,
@ -603,6 +468,22 @@ public class ModelComboBoxAction extends ComboBoxAction {
var chatModel = ApplicationManager.getApplication().getService(ModelSettings.class).getState()
.getModelSelection(featureType);
var modelCode = chatModel != null ? chatModel.getModel() : null;
return ModelRegistry.getInstance().getModelDisplayName(MISTRAL, modelCode);
return getModelSettings().getModelDisplayName(MISTRAL, modelCode);
}
private ModelSettings getModelSettings() {
return ApplicationManager.getApplication().getService(ModelSettings.class);
}
private List<ModelSelection> getAvailableModelsForFeature() {
return getModelSettings().getAvailableModels(featureType);
}
private boolean isModelSelected(ServiceType serviceType, @Nullable String modelCode) {
var current = getModelSettings().getState().getModelSelection(featureType);
if (current == null || current.getProvider() != serviceType || modelCode == null) {
return false;
}
return modelCode.equals(current.getModel());
}
}

8
src/main/java/ee/carlrobert/codegpt/ui/ModelIconLabel.java
View file

@ -3,8 +3,8 @@ package ee.carlrobert.codegpt.ui;
import com.intellij.ui.components.JBLabel;
import com.intellij.util.ui.JBFont;
import ee.carlrobert.codegpt.Icons;
import ee.carlrobert.llm.client.openai.completion.OpenAIChatCompletionModel;
import java.util.NoSuchElementException;
import ee.carlrobert.codegpt.settings.models.ModelSettings;
import ee.carlrobert.codegpt.settings.service.ServiceType;
import javax.swing.SwingConstants;
public class ModelIconLabel extends JBLabel {
@ -29,8 +29,8 @@ public class ModelIconLabel extends JBLabel {
private String formatModelName(String modelCode) {
try {
return OpenAIChatCompletionModel.findByCode(modelCode).getDescription();
} catch (NoSuchElementException e) {
return ModelSettings.getInstance().getModelDisplayName(ServiceType.OPENAI, modelCode);
} catch (Exception e) {
return modelCode;
}
}

27
src/main/kotlin/ee/carlrobert/codegpt/actions/BaseCommitWorkflowAction.kt
View file

@ -19,16 +19,15 @@ import com.intellij.openapi.vcs.changes.Change
import com.intellij.vcs.commit.CommitWorkflowUi
import ee.carlrobert.codegpt.EncodingManager
import ee.carlrobert.codegpt.codecompletions.CompletionProgressNotifier
import ee.carlrobert.codegpt.completions.CompletionRequestService
import ee.carlrobert.codegpt.completions.CompletionError
import ee.carlrobert.codegpt.completions.CompletionStreamEventListener
import ee.carlrobert.codegpt.completions.CompletionService
import ee.carlrobert.codegpt.settings.service.FeatureType
import ee.carlrobert.codegpt.util.ThinkingOutputParser
import ee.carlrobert.codegpt.ui.OverlayUtil
import ee.carlrobert.codegpt.util.CommitWorkflowChanges
import ee.carlrobert.codegpt.util.GitUtil.getProjectRepository
import ee.carlrobert.llm.client.openai.completion.ErrorDetails
import ee.carlrobert.llm.completion.CompletionEventListener
import ee.carlrobert.codegpt.util.ThinkingOutputParser
import git4idea.repo.GitRepository
import okhttp3.sse.EventSource
import java.io.StringWriter
import java.nio.file.Path
@ -48,7 +47,7 @@ abstract class BaseCommitWorkflowAction : DumbAwareAction() {
override fun update(event: AnActionEvent) {
val commitWorkflowUi = event.getData(VcsDataKeys.COMMIT_WORKFLOW_UI) ?: return
val requestAllowed = CompletionRequestService.isRequestAllowed(FeatureType.COMMIT_MESSAGE)
val requestAllowed = CompletionService.isRequestAllowed(FeatureType.COMMIT_MESSAGE)
runInEdt {
event.presentation.isEnabled =
requestAllowed && CommitWorkflowChanges(commitWorkflowUi).isFilesSelected
@ -125,12 +124,12 @@ abstract class BaseCommitWorkflowAction : DumbAwareAction() {
class CommitMessageEventListener(
private val project: Project,
private val commitWorkflowUi: CommitWorkflowUi
) : CompletionEventListener<String> {
) : CompletionStreamEventListener {
private val messageBuilder = StringBuilder()
private val thinkingOutputParser = ThinkingOutputParser()
override fun onMessage(message: String, eventSource: EventSource) {
override fun onMessage(message: String) {
val processedChunk = thinkingOutputParser.processChunk(message)
if (processedChunk.isNotEmpty() && thinkingOutputParser.isFinished) {
messageBuilder.append(message)
@ -138,14 +137,14 @@ class CommitMessageEventListener(
}
}
override fun onComplete(result: StringBuilder) {
override fun onComplete(messageBuilder: StringBuilder) {
if (messageBuilder.isEmpty()) {
updateCommitMessage(result.toString())
updateCommitMessage(messageBuilder.toString())
}
stopLoading()
}
override fun onError(error: ErrorDetails, ex: Throwable) {
override fun onError(error: CompletionError, ex: Throwable) {
Notifications.Bus.notify(
Notification(
"proxyai.notification.group",
@ -157,6 +156,10 @@ class CommitMessageEventListener(
stopLoading()
}
override fun onCancelled(messageBuilder: StringBuilder) {
stopLoading()
}
private fun stopLoading() {
CompletionProgressNotifier.update(project, false)
}
@ -168,4 +171,4 @@ class CommitMessageEventListener(
}
}
}
}
}

6
src/main/kotlin/ee/carlrobert/codegpt/actions/CodeCompletionFeatureToggleActions.kt
View file

@ -6,7 +6,7 @@ import com.intellij.openapi.components.service
import com.intellij.openapi.project.DumbAwareAction
import ee.carlrobert.codegpt.codecompletions.CodeCompletionService
import ee.carlrobert.codegpt.settings.service.FeatureType
import ee.carlrobert.codegpt.settings.service.ModelSelectionService
import ee.carlrobert.codegpt.settings.models.ModelSettings
import ee.carlrobert.codegpt.settings.service.ServiceType.*
import ee.carlrobert.codegpt.settings.service.codegpt.CodeGPTServiceSettings
import ee.carlrobert.codegpt.settings.service.custom.CustomServicesSettings
@ -22,7 +22,7 @@ abstract class CodeCompletionFeatureToggleActions(
override fun actionPerformed(e: AnActionEvent) {
val serviceType =
service<ModelSelectionService>().getServiceForFeature(FeatureType.CODE_COMPLETION)
service<ModelSettings>().getServiceForFeature(FeatureType.CODE_COMPLETION)
when (serviceType) {
PROXYAI -> {
service<CodeGPTServiceSettings>().state.codeCompletionSettings.codeCompletionsEnabled =
@ -63,7 +63,7 @@ abstract class CodeCompletionFeatureToggleActions(
override fun update(e: AnActionEvent) {
val selectedService =
service<ModelSelectionService>().getServiceForFeature(FeatureType.CODE_COMPLETION)
service<ModelSettings>().getServiceForFeature(FeatureType.CODE_COMPLETION)
val codeCompletionEnabled =
service<CodeCompletionService>().isCodeCompletionsEnabled(selectedService) == true
e.presentation.isVisible = codeCompletionEnabled != enableFeatureAction

7
src/main/kotlin/ee/carlrobert/codegpt/actions/GenerateCommitMessageAction.kt
View file

@ -1,10 +1,11 @@
package ee.carlrobert.codegpt.actions
import com.intellij.openapi.components.service
import com.intellij.openapi.project.Project
import com.intellij.vcs.commit.CommitWorkflowUi
import ee.carlrobert.codegpt.codecompletions.CompletionProgressNotifier
import ee.carlrobert.codegpt.completions.CommitMessageCompletionParameters
import ee.carlrobert.codegpt.completions.CompletionRequestService
import ee.carlrobert.codegpt.completions.CompletionService
import ee.carlrobert.codegpt.settings.prompts.CommitMessageTemplate
class GenerateCommitMessageAction : BaseCommitWorkflowAction() {
@ -19,10 +20,10 @@ class GenerateCommitMessageAction : BaseCommitWorkflowAction() {
gitDiff: String
) {
CompletionProgressNotifier.update(project, true)
CompletionRequestService.getInstance().getCommitMessageAsync(
service<CompletionService>().getCommitMessage(
CommitMessageCompletionParameters(
gitDiff,
project.getService(CommitMessageTemplate::class.java).getSystemPrompt()
project.service<CommitMessageTemplate>().getSystemPrompt()
),
CommitMessageEventListener(project, commitWorkflowUi)
)

5
src/main/kotlin/ee/carlrobert/codegpt/actions/GenerateCommitMessageWithAdditionalInputAction.kt
View file

@ -1,10 +1,11 @@
package ee.carlrobert.codegpt.actions
import com.intellij.openapi.components.service
import com.intellij.openapi.project.Project
import com.intellij.openapi.ui.Messages
import com.intellij.vcs.commit.CommitWorkflowUi
import ee.carlrobert.codegpt.completions.CommitMessageCompletionParameters
import ee.carlrobert.codegpt.completions.CompletionRequestService
import ee.carlrobert.codegpt.completions.CompletionService
import ee.carlrobert.codegpt.settings.prompts.CommitMessageTemplate
class GenerateCommitMessageWithAdditionalInputAction : BaseCommitWorkflowAction() {
@ -29,7 +30,7 @@ class GenerateCommitMessageWithAdditionalInputAction : BaseCommitWorkflowAction(
if (userInput != null) {
val systemPrompt =
project.getService(CommitMessageTemplate::class.java).getSystemPrompt()
CompletionRequestService.getInstance().getCommitMessageAsync(
service<CompletionService>().getCommitMessage(
CommitMessageCompletionParameters(
gitDiff,
"$systemPrompt\n\nUser input: $userInput"

4
src/main/kotlin/ee/carlrobert/codegpt/actions/LlamaServerToggleActions.kt
View file

@ -12,7 +12,7 @@ import ee.carlrobert.codegpt.completions.llama.LlamaModel.findByHuggingFaceModel
import ee.carlrobert.codegpt.completions.llama.LlamaServerAgent
import ee.carlrobert.codegpt.completions.llama.LlamaServerStartupParams
import ee.carlrobert.codegpt.settings.service.FeatureType
import ee.carlrobert.codegpt.settings.service.ModelSelectionService
import ee.carlrobert.codegpt.settings.models.ModelSettings
import ee.carlrobert.codegpt.settings.service.ServiceType.LLAMA_CPP
import ee.carlrobert.codegpt.settings.service.llama.LlamaSettings
import ee.carlrobert.codegpt.completions.llama.logging.NoOpLoggingStrategy
@ -50,7 +50,7 @@ abstract class LlamaServerToggleActions(
var notification: Notification? = null
override fun actionPerformed(e: AnActionEvent) {
val modelSelectionService = ModelSelectionService.getInstance()
val modelSelectionService = ModelSettings.getInstance()
val isLlamaUsed =
(modelSelectionService.getServiceForFeature(FeatureType.CHAT) == LLAMA_CPP ||
modelSelectionService.getServiceForFeature(FeatureType.CODE_COMPLETION) == LLAMA_CPP)

6
src/main/kotlin/ee/carlrobert/codegpt/agent/AgentEvents.kt
View file

@ -1,9 +1,7 @@
package ee.carlrobert.codegpt.agent
import ai.koog.prompt.executor.clients.LLMClientException
import ee.carlrobert.codegpt.agent.history.CheckpointRef
import ee.carlrobert.codegpt.agent.tools.AskUserQuestionTool
import ee.carlrobert.codegpt.conversations.message.TokenUsage
import ee.carlrobert.codegpt.settings.service.ServiceType
import ee.carlrobert.codegpt.toolwindow.agent.AgentCreditsEvent
import ee.carlrobert.codegpt.toolwindow.agent.ui.approval.ToolApprovalRequest
@ -23,12 +21,12 @@ interface AgentEvents {
) {
}
fun onRetry(attempt: Int, maxAttempts: Int, reason: String? = null) {}
fun onRetry(attempt: Int, maxAttempts: Int) {}
fun onRunCheckpointUpdated(runMessageId: UUID, ref: CheckpointRef?) {}
fun onQueuedMessagesResolved()
fun onTokenUsageAvailable(tokenUsage: Long) {}
fun onCreditsAvailable(event: AgentCreditsEvent) {}
fun onClientException(provider: ServiceType, ex: LLMClientException) {}
fun onAgentException(provider: ServiceType, throwable: Throwable) {}
fun onHistoryCompressionStateChanged(isCompressing: Boolean) {}
suspend fun approveToolCall(request: ToolApprovalRequest): Boolean = false

156
src/main/kotlin/ee/carlrobert/codegpt/agent/AgentFactory.kt
View file

@ -21,28 +21,21 @@ import ai.koog.agents.features.tokenizer.feature.MessageTokenizer
import ai.koog.agents.features.tokenizer.feature.tokenizer
import ai.koog.prompt.dsl.prompt
import ai.koog.prompt.executor.clients.LLMClient
import ai.koog.prompt.executor.clients.anthropic.AnthropicLLMClient
import ai.koog.prompt.executor.clients.google.GoogleLLMClient
import ai.koog.prompt.executor.clients.mistralai.MistralAILLMClient
import ai.koog.prompt.executor.clients.openai.OpenAILLMClient
import ai.koog.prompt.executor.model.PromptExecutor
import ai.koog.prompt.executor.ollama.client.OllamaClient
import ai.koog.prompt.llm.LLMProvider
import ai.koog.prompt.message.Message
import ai.koog.prompt.tokenizer.Tokenizer
import com.intellij.openapi.components.service
import com.intellij.openapi.project.Project
import ee.carlrobert.codegpt.EncodingManager
import ee.carlrobert.codegpt.agent.clients.*
import ee.carlrobert.codegpt.agent.clients.RetryingPromptExecutor
import ee.carlrobert.codegpt.agent.credits.extractCreditsSnapshot
import ee.carlrobert.codegpt.agent.tools.*
import ee.carlrobert.codegpt.credentials.CredentialsStore.CredentialKey
import ee.carlrobert.codegpt.credentials.CredentialsStore.getCredential
import ee.carlrobert.codegpt.settings.hooks.HookManager
import ee.carlrobert.codegpt.settings.models.LLMClientFactory
import ee.carlrobert.codegpt.settings.models.ModelSettings
import ee.carlrobert.codegpt.settings.service.FeatureType
import ee.carlrobert.codegpt.settings.service.ModelSelectionService
import ee.carlrobert.codegpt.settings.service.ServiceType
import ee.carlrobert.codegpt.settings.service.custom.CustomServicesSettings
import ee.carlrobert.codegpt.settings.service.ollama.OllamaSettings
import ee.carlrobert.codegpt.settings.skills.SkillDiscoveryService
import ee.carlrobert.codegpt.settings.skills.SkillPromptFormatter
import ee.carlrobert.codegpt.toolwindow.agent.AgentCreditsEvent
@ -52,6 +45,8 @@ import kotlin.time.Duration.Companion.seconds
object AgentFactory {
private const val MAX_AGENT_ITERATIONS = 250
fun createAgent(
agentType: AgentType,
provider: ServiceType,
@ -152,8 +147,8 @@ object AgentFactory {
""".trimIndent()
)
},
model = service<ModelSelectionService>().getAgentModel(),
maxAgentIterations = 100
model = service<ModelSettings>().getAgentModel(),
maxAgentIterations = MAX_AGENT_ITERATIONS
),
toolRegistry = registry,
installFeatures = {
@ -169,72 +164,13 @@ object AgentFactory {
)
}
fun createExecutor(provider: ServiceType, events: AgentEvents? = null): PromptExecutor {
val httpClient = HttpClientProvider.createHttpClient()
return when (provider) {
ServiceType.OPENAI -> {
val apiKey = getCredential(CredentialKey.OpenaiApiKey) ?: ""
createRetryingExecutor(OpenAILLMClient(apiKey, baseClient = httpClient), events)
}
ServiceType.ANTHROPIC -> {
val apiKey = getCredential(CredentialKey.AnthropicApiKey) ?: ""
createRetryingExecutor(AnthropicLLMClient(apiKey, baseClient = httpClient), events)
}
ServiceType.GOOGLE -> {
val apiKey = getCredential(CredentialKey.GoogleApiKey) ?: ""
createRetryingExecutor(GoogleLLMClient(apiKey, baseClient = httpClient), events)
}
ServiceType.OLLAMA -> {
createRetryingExecutor(
OllamaClient(
baseUrl = service<OllamaSettings>().state.host ?: "http://localhost:11434",
baseClient = httpClient
),
events
)
}
ServiceType.MISTRAL -> {
val apiKey = getCredential(CredentialKey.MistralApiKey) ?: ""
createRetryingExecutor(MistralAILLMClient(apiKey, baseClient = httpClient), events)
}
ServiceType.CUSTOM_OPENAI -> {
val state = service<CustomServicesSettings>()
.customServiceStateForFeatureType(FeatureType.AGENT)
val serviceId = state.id
?: throw IllegalArgumentException("No custom service configured")
val apiKey = getCredential(CredentialKey.CustomServiceApiKeyById(serviceId))
?: throw IllegalArgumentException("No API key found for custom service: $serviceId")
createRetryingExecutor(
CustomOpenAILLMClient.fromSettingsState(
apiKey,
state.chatCompletionSettings,
httpClient
),
events
)
}
ServiceType.PROXYAI -> {
val apiKey = getCredential(CredentialKey.CodeGptApiKey) ?: ""
createRetryingExecutor(ProxyAILLMClient(apiKey, baseClient = httpClient), events)
}
ServiceType.INCEPTION -> {
val apiKey = getCredential(CredentialKey.InceptionApiKey) ?: ""
createRetryingExecutor(
InceptionAILLMClient(apiKey, baseClient = httpClient),
events
)
}
else -> throw UnsupportedOperationException("Provider not supported: $provider")
}
fun createExecutor(
provider: ServiceType,
events: AgentEvents? = null,
featureType: FeatureType = FeatureType.AGENT
): PromptExecutor {
val llmClient = LLMClientFactory.createClient(provider, featureType)
return createRetryingExecutor(llmClient, events)
}
private fun createRetryingExecutor(client: LLMClient, events: AgentEvents?): PromptExecutor {
@ -304,8 +240,8 @@ object AgentFactory {
""".trimIndent()
)
},
model = service<ModelSelectionService>().getAgentModel(),
maxAgentIterations = 100
model = service<ModelSettings>().getAgentModel(),
maxAgentIterations = MAX_AGENT_ITERATIONS
),
toolRegistry = createToolRegistry(
project,
@ -377,8 +313,8 @@ object AgentFactory {
""".trimIndent()
)
},
model = service<ModelSelectionService>().getAgentModel(),
maxAgentIterations = 100
model = service<ModelSettings>().getAgentModel(),
maxAgentIterations = MAX_AGENT_ITERATIONS
),
toolRegistry = createToolRegistry(
project,
@ -498,16 +434,36 @@ object AgentFactory {
): List<Message.Response> {
val preparedPrompt = config.missingToolsConversionStrategy.convertPrompt(prompt, tools)
val responses = executor.execute(preparedPrompt, model, tools)
val appendableResponses = appendableResponses(responses)
val appendableResponses = appendableResponses(responses, model.provider)
appendableResponses.forEach(appendResponse)
tokenCounter?.addAndGet(countResponseTokens(appendableResponses))
return appendableResponses
}
private fun appendableResponses(responses: List<Message.Response>): List<Message.Response> {
return responses
private fun appendableResponses(
responses: List<Message.Response>,
provider: LLMProvider
): List<Message.Response> {
val sortedResponses = responses
.sortedBy { if (it is Message.Assistant) 0 else 1 }
.filter { it !is Message.Reasoning }
.filterNot { response ->
when {
response is Message.Assistant &&
response.hasOnlyTextContent() &&
response.content.isBlank() -> true
response is Message.Reasoning &&
provider != LLMProvider.Google -> true
else -> false
}
}
return if (provider == LLMProvider.Google && sortedResponses.any { it is Message.Tool.Call }) {
sortedResponses.filterNot { it is Message.Assistant }
} else {
sortedResponses
}
}
private fun countResponseTokens(responses: List<Message.Response>): Long {
@ -543,43 +499,52 @@ object AgentFactory {
val contextService = project.service<AgentMcpContextService>()
val mcpContext = contextService.get(sessionId)
return ToolRegistry.Companion {
if (SubagentTool.READ in selected) tool(ReadTool(project, hookManager, sessionId))
if (SubagentTool.READ in selected) tool(ReadTool(project, sessionId, hookManager))
if (SubagentTool.EDIT in selected) {
tool(
ConfirmingEditTool(EditTool(project, hookManager, sessionId)) { name, details ->
ConfirmingEditTool(EditTool(project, sessionId, hookManager)) { name, details ->
approveToolCall?.invoke(name, details) ?: false
}
)
}
if (SubagentTool.WRITE in selected) {
tool(
ConfirmingWriteTool(WriteTool(project, hookManager)) { name, details ->
ConfirmingWriteTool(
WriteTool(
project,
sessionId,
hookManager
)
) { name, details ->
approveToolCall?.invoke(name, details) ?: false
}
)
}
if (SubagentTool.TODO_WRITE in selected) tool(
TodoWriteTool(
project,
sessionId,
hookManager
project = project,
sessionId = sessionId,
hookManager = hookManager
)
)
if (SubagentTool.INTELLIJ_SEARCH in selected) tool(
IntelliJSearchTool(
project = project,
hookManager = hookManager
sessionId = sessionId,
hookManager = hookManager,
)
)
if (SubagentTool.WEB_SEARCH in selected) tool(
WebSearchTool(
workingDirectory = project.basePath ?: System.getProperty("user.dir"),
sessionId = sessionId,
hookManager = hookManager,
)
)
if (SubagentTool.WEB_FETCH in selected) tool(
WebFetchTool(
workingDirectory = project.basePath ?: System.getProperty("user.dir"),
sessionId = sessionId,
hookManager = hookManager,
)
)
@ -593,18 +558,21 @@ object AgentFactory {
if (SubagentTool.KILL_SHELL in selected) tool(
KillShellTool(
workingDirectory = project.basePath ?: System.getProperty("user.dir"),
sessionId = sessionId,
hookManager = hookManager,
)
)
if (SubagentTool.RESOLVE_LIBRARY_ID in selected) tool(
ResolveLibraryIdTool(
workingDirectory = project.basePath ?: System.getProperty("user.dir"),
sessionId = sessionId,
hookManager = hookManager,
)
)
if (SubagentTool.GET_LIBRARY_DOCS in selected) tool(
GetLibraryDocsTool(
workingDirectory = project.basePath ?: System.getProperty("user.dir"),
sessionId = sessionId,
hookManager = hookManager
)
)

43
src/main/kotlin/ee/carlrobert/codegpt/agent/AgentService.kt
View file

@ -4,22 +4,21 @@ import ai.koog.agents.core.agent.AIAgent
import ai.koog.agents.core.agent.AIAgentService
import ai.koog.agents.snapshot.feature.AgentCheckpointData
import ai.koog.agents.snapshot.providers.file.JVMFilePersistenceStorageProvider
import ai.koog.prompt.executor.clients.LLMClientException
import com.intellij.openapi.components.Service
import com.intellij.openapi.components.service
import com.intellij.openapi.diagnostic.thisLogger
import com.intellij.openapi.project.Project
import ee.carlrobert.codegpt.agent.history.AgentCheckpointHistoryService
import ee.carlrobert.codegpt.agent.history.CheckpointRef
import ee.carlrobert.codegpt.settings.models.ModelSettings
import ee.carlrobert.codegpt.settings.service.FeatureType
import ee.carlrobert.codegpt.settings.service.ModelSelectionService
import ee.carlrobert.codegpt.settings.service.ServiceType
import ee.carlrobert.codegpt.toolwindow.agent.AgentToolWindowContentManager
import ee.carlrobert.codegpt.ui.textarea.header.tag.McpTagDetails
import kotlinx.coroutines.*
import kotlinx.coroutines.flow.MutableSharedFlow
import kotlinx.coroutines.flow.asSharedFlow
import kotlinx.datetime.Clock
import kotlin.time.Clock
import kotlinx.serialization.json.JsonNull
import java.util.*
import java.util.concurrent.ConcurrentHashMap
@ -111,10 +110,15 @@ class AgentService(private val project: Project) {
return
}
val provider = service<ModelSelectionService>().getServiceForFeature(FeatureType.AGENT)
val provider = service<ModelSettings>().getServiceForFeature(FeatureType.AGENT)
val contentManager = project.service<AgentToolWindowContentManager>()
val runtimeAgentId = contentManager.getSession(sessionId)?.runtimeAgentId
val runtime = ensureSessionRuntime(sessionId, provider, events)
val runtime = runCatching {
ensureSessionRuntime(sessionId, provider, events)
}.onFailure { ex ->
logger.warn("Failed to initialize runtime for session=$sessionId", ex)
events.onAgentException(provider, ex)
}.getOrNull() ?: return
val agent = runCatching {
runBlocking {
@ -122,6 +126,7 @@ class AgentService(private val project: Project) {
}
}.onFailure { ex ->
logger.warn("Failed to create managed agent for session=$sessionId", ex)
events.onAgentException(provider, ex)
}.getOrNull() ?: return
contentManager.setRuntimeAgentId(sessionId, agent.id)
@ -129,12 +134,11 @@ class AgentService(private val project: Project) {
sessionJobs[sessionId] = CoroutineScope(Dispatchers.IO).launch {
try {
agent.run(message)
} catch (ex: LLMClientException) {
events.onClientException(provider, ex)
} catch (_: CancellationException) {
return@launch
} catch (ex: Exception) {
} catch (ex: Throwable) {
logger.error(ex)
events.onAgentException(provider, ex)
} finally {
val ref = refreshSessionResumeCheckpoint(sessionId, agent.id)
events.onRunCheckpointUpdated(message.id, ref)
@ -179,7 +183,7 @@ class AgentService(private val project: Project) {
}
fun clearPendingMessages(sessionId: String) {
pendingMessages.remove(sessionId)
pendingMessages[sessionId]?.clear()
}
fun getAgentForSession(sessionId: String): AIAgent<MessageWithContext, String>? {
@ -239,17 +243,8 @@ class AgentService(private val project: Project) {
sessionId: String,
agentId: String
): CheckpointRef? {
val ref = runCatching {
historyService.loadLatestResumeCheckpoint(agentId)
?.let { CheckpointRef(agentId, it.checkpointId) }
}.onFailure { ex ->
logger.warn(
"Agent checkpoints: failed to refresh session checkpoint session=$sessionId " +
"agentId=$agentId error=${ex.message}",
ex
)
}.getOrNull()
val ref = historyService.loadLatestResumeCheckpoint(agentId)
?.let { CheckpointRef(agentId, it.checkpointId) }
if (ref != null) {
project.service<AgentToolWindowContentManager>().setResumeCheckpointRef(sessionId, ref)
}
@ -261,11 +256,10 @@ class AgentService(private val project: Project) {
provider: ServiceType,
events: AgentEvents
): SessionRuntime {
val modelId = service<ModelSelectionService>().getAgentModel().id
val modelId = service<ModelSettings>().getAgentModel().id
val selectedServerIds = project.service<AgentMcpContextService>()
.get(sessionId)
?.selectedServerIds
?: emptySet()
?.selectedServerIds ?: emptySet()
val existing = sessionRuntimes[sessionId]
if (existing != null &&
existing.provider == provider &&
@ -302,9 +296,8 @@ class AgentService(private val project: Project) {
}
private fun logCheckpointLoadFailure(sessionId: String, agentId: String, ex: Throwable) {
val sessionInfo = sessionId.let { " session=$it" }
logger.error(
"Agent checkpoints: failed to load for$sessionInfo agentId=$agentId error=${ex.message}",
"Agent checkpoints: failed to load for sessionId=$sessionId agentId=$agentId error=${ex.message}",
ex
)
}

4
src/main/kotlin/ee/carlrobert/codegpt/agent/AgentSystemPrompts.kt
View file

@ -1,6 +1,6 @@
package ee.carlrobert.codegpt.agent
import ee.carlrobert.codegpt.settings.models.ModelRegistry
import ee.carlrobert.codegpt.settings.models.ModelCatalog
import ee.carlrobert.codegpt.settings.models.ModelSelection
import ee.carlrobert.codegpt.settings.service.ServiceType
import ee.carlrobert.codegpt.settings.skills.SkillDescriptor
@ -42,7 +42,7 @@ internal object AgentSystemPrompts {
projectPath: String?
): String {
val modelId = modelSelection?.model
return if (modelId == ModelRegistry.PROXYAI_AUTO || modelId?.startsWith("claude") == true) {
return if (modelId == ModelCatalog.PROXYAI_AUTO || modelId?.startsWith("claude") == true) {
createAnthropicSystemPrompt(projectPath)
} else {
createOpenAiSystemPrompt(projectPath)

4
src/main/kotlin/ee/carlrobert/codegpt/agent/PromptEnhancer.kt
View file

@ -4,7 +4,7 @@ import ai.koog.prompt.dsl.prompt
import com.intellij.openapi.components.service
import com.intellij.openapi.project.Project
import ee.carlrobert.codegpt.settings.service.FeatureType
import ee.carlrobert.codegpt.settings.service.ModelSelectionService
import ee.carlrobert.codegpt.settings.models.ModelSettings
import ee.carlrobert.codegpt.ui.textarea.TagProcessorFactory
import ee.carlrobert.codegpt.ui.textarea.header.tag.TagDetails
import ee.carlrobert.codegpt.util.GitUtil
@ -37,7 +37,7 @@ class PromptEnhancer(private val project: Project) {
val tagsContext = buildTagsContext(tags)
val historyContext = buildHistoryContext(sessionId)
val modelService = service<ModelSelectionService>()
val modelService = service<ModelSettings>()
val provider = modelService.getServiceForFeature(FeatureType.AGENT)
val model = modelService.getAgentModel()
val executor = AgentFactory.createExecutor(provider)

104
src/main/kotlin/ee/carlrobert/codegpt/agent/ProxyAIAgent.kt
View file

@ -18,41 +18,43 @@ import ai.koog.prompt.dsl.prompt
import ai.koog.prompt.llm.LLModel
import ai.koog.prompt.message.Message
import ai.koog.prompt.message.RequestMetaInfo
import ai.koog.prompt.streaming.StreamFrame
import ai.koog.prompt.tokenizer.Tokenizer
import com.intellij.openapi.components.service
import com.intellij.openapi.project.Project
import ee.carlrobert.codegpt.EncodingManager
import ee.carlrobert.codegpt.agent.clients.shouldStream
import ee.carlrobert.codegpt.agent.strategy.CODE_AGENT_COMPRESSION
import ee.carlrobert.codegpt.agent.strategy.HistoryCompressionConfig
import ee.carlrobert.codegpt.agent.strategy.SingleRunStrategyProvider
import ee.carlrobert.codegpt.agent.strategy.buildHistoryTooBigPredicate
import ee.carlrobert.codegpt.agent.tools.*
import ee.carlrobert.codegpt.mcp.McpSessionManager
import ee.carlrobert.codegpt.settings.configuration.ConfigurationSettings
import ee.carlrobert.codegpt.settings.hooks.HookEventType
import ee.carlrobert.codegpt.settings.hooks.HookManager
import ee.carlrobert.codegpt.settings.models.ModelSettings
import ee.carlrobert.codegpt.settings.service.FeatureType
import ee.carlrobert.codegpt.settings.service.ModelSelectionService
import ee.carlrobert.codegpt.settings.service.ServiceType
import ee.carlrobert.codegpt.settings.service.custom.CustomServicesSettings
import ee.carlrobert.codegpt.settings.skills.SkillDiscoveryService
import ee.carlrobert.codegpt.toolwindow.agent.ui.approval.BashPayload
import ee.carlrobert.codegpt.toolwindow.agent.ui.approval.ToolApprovalRequest
import ee.carlrobert.codegpt.toolwindow.agent.ui.approval.ToolApprovalType
import ee.carlrobert.codegpt.util.ReasoningFrameTextAdapter
import io.github.oshai.kotlinlogging.KotlinLogging
import kotlinx.datetime.Clock
import kotlinx.serialization.json.JsonObject
import kotlinx.serialization.json.JsonPrimitive
import java.nio.file.Files
import java.nio.file.Paths
import java.util.*
import java.util.concurrent.ConcurrentHashMap
import kotlin.time.Clock
data class ToolError(val message: String)
object ProxyAIAgent {
private const val INSTRUCTION_FILE_NAME = "PROXYAI.md"
private const val MAX_AGENT_ITERATIONS = 250
private val logger = KotlinLogging.logger { }
@ -85,15 +87,15 @@ object ProxyAIAgent {
pendingMessages: ConcurrentHashMap<String, ArrayDeque<MessageWithContext>>,
): GraphAIAgentService<MessageWithContext, String> {
val modelSelection =
service<ModelSelectionService>().getModelSelectionForFeature(FeatureType.AGENT)
service<ModelSettings>().getModelSelectionForFeature(FeatureType.AGENT)
val skills = project.service<SkillDiscoveryService>().listSkills()
val stream = provider != ServiceType.CUSTOM_OPENAI
val stream = shouldStreamAgentToolLoop(project, provider)
val projectInstructions = loadProjectInstructions(project.basePath)
val executor = AgentFactory.createExecutor(provider, events)
val pendingMessageQueue = pendingMessages.getOrPut(sessionId) { ArrayDeque() }
val hookManager = HookManager(project)
val toolRegistry = createToolRegistry(project, events, sessionId, provider, hookManager)
val agentModel = service<ModelSelectionService>().getAgentModel()
val agentModel = service<ModelSettings>().getAgentModel()
return AIAgentService<MessageWithContext, String>(
promptExecutor = executor,
strategy = SingleRunStrategyProvider().build(
@ -135,7 +137,7 @@ object ProxyAIAgent {
}
},
model = agentModel,
maxAgentIterations = 100
maxAgentIterations = MAX_AGENT_ITERATIONS
),
toolRegistry = toolRegistry,
) {
@ -160,13 +162,15 @@ object ProxyAIAgent {
handleEvents {
val toolCallToUiId: MutableMap<String, String> = HashMap()
val anonymousToolIds: ArrayDeque<String> = ArrayDeque()
val frameAdapter = ReasoningFrameTextAdapter()
onLLMStreamingFrameReceived { ctx ->
if (!stream) return@onLLMStreamingFrameReceived
val frame = ctx.streamFrame
if (frame is StreamFrame.Append) {
events.onTextReceived(frame.text)
frameAdapter.consume(ctx.streamFrame).forEach { chunk ->
if (chunk.isNotEmpty()) {
events.onTextReceived(chunk)
}
}
}
@ -178,7 +182,9 @@ object ProxyAIAgent {
events.onTextReceived(it.content)
}
(msg as? Message.Reasoning)?.let {
events.onTextReceived(it.content)
if (it.content.isNotBlank()) {
events.onTextReceived("<think>${it.content}</think>")
}
}
}
}
@ -255,12 +261,30 @@ object ProxyAIAgent {
),
sessionId = sessionId
)
events.onAgentCompleted(it.agentId)
events.onAgentException(provider, it.throwable)
}
}
}
}
private fun shouldStreamAgentToolLoop(
project: Project,
provider: ServiceType,
): Boolean {
return when (provider) {
ServiceType.CUSTOM_OPENAI -> {
val selectedServiceId =
project.service<ModelSettings>().getStoredModelForFeature(FeatureType.AGENT)
project.service<CustomServicesSettings>().state.services
.firstOrNull { it.id == selectedServiceId }?.chatCompletionSettings?.shouldStream()
?: false
}
ServiceType.GOOGLE -> false
else -> true
}
}
private fun createToolRegistry(
project: Project,
events: AgentEvents,
@ -270,67 +294,75 @@ object ProxyAIAgent {
): ToolRegistry {
val workingDirectory = project.basePath ?: System.getProperty("user.dir")
val contextService = project.service<AgentMcpContextService>()
val mcpSessionManager = service<McpSessionManager>()
val standardApproval = approvalHandler(events)
val writeApproval = writeApprovalHandler(events)
val genericApproval = genericApprovalHandler(events)
return ToolRegistry {
tool(ReadTool(project, hookManager, sessionId))
tool(ConfirmingEditTool(EditTool(project, hookManager, sessionId), standardApproval))
tool(ConfirmingWriteTool(WriteTool(project, hookManager), writeApproval))
tool(ReadTool(project, sessionId, hookManager))
tool(ConfirmingEditTool(EditTool(project, sessionId, hookManager), standardApproval))
tool(ConfirmingWriteTool(WriteTool(project, sessionId, hookManager), writeApproval))
tool(TodoWriteTool(project, sessionId, hookManager))
tool(
AskUserQuestionTool(
workingDirectory = workingDirectory,
sessionId = sessionId,
hookManager = hookManager,
events = events
)
)
createMcpTools(
project = project,
sessionId = sessionId,
contextService = contextService,
sessionManager = mcpSessionManager,
hookManager = hookManager,
approve = genericApproval
).forEach { mcpTool -> tool(mcpTool) }
tool(ExitTool)
tool(IntelliJSearchTool(project = project, hookManager = hookManager))
tool(
IntelliJSearchTool(
project = project,
sessionId = sessionId,
hookManager = hookManager,
)
)
tool(
WebSearchTool(
workingDirectory = workingDirectory,
sessionId = sessionId,
hookManager = hookManager,
)
)
tool(
WebFetchTool(
workingDirectory = workingDirectory,
sessionId = sessionId,
hookManager = hookManager,
)
)
tool(
BashOutputTool(
workingDirectory = workingDirectory,
sessionId = sessionId,
hookManager = hookManager,
sessionId = sessionId
)
)
tool(
KillShellTool(
workingDirectory = workingDirectory,
sessionId = sessionId,
hookManager = hookManager,
)
)
tool(
ResolveLibraryIdTool(
workingDirectory = workingDirectory,
sessionId = sessionId,
hookManager = hookManager,
)
)
tool(
GetLibraryDocsTool(
workingDirectory = workingDirectory,
hookManager = hookManager
sessionId = sessionId,
hookManager = hookManager,
)
)
tool(
@ -412,35 +444,15 @@ object ProxyAIAgent {
}
private fun createMcpTools(
project: Project,
sessionId: String,
contextService: AgentMcpContextService,
sessionManager: McpSessionManager,
hookManager: HookManager,
approve: suspend (String, String) -> Boolean
): List<Tool<*, *>> {
val context = contextService.get(sessionId) ?: return emptyList()
if (!context.hasSelection() || context.conversationId == null) {
return emptyList()
}
val dynamicTools = McpDynamicToolRegistry.createTools(context, approve)
if (dynamicTools.isNotEmpty()) {
return dynamicTools
}
return listOf(
ConfirmingMcpTool(
McpTool(
project = project,
sessionId = sessionId,
context = context,
sessionManager = sessionManager,
hookManager = hookManager
),
approve = approve
)
)
return McpDynamicToolRegistry.createTools(context, approve)
}
private suspend fun safeApprove(
@ -452,8 +464,8 @@ object ProxyAIAgent {
}
private fun computeAvailableInput(model: LLModel): Long {
val contextLength = model.contextLength
val outputLength = model.maxOutputTokens ?: 0
val contextLength = model.contextLength ?: 128_000L
val outputLength = model.maxOutputTokens ?: 0L
return (contextLength - outputLength).coerceAtLeast(1L)
}
}

111
src/main/kotlin/ee/carlrobert/codegpt/agent/clients/CodeCompletionSupport.kt Normal file
View file

@ -0,0 +1,111 @@
package ee.carlrobert.codegpt.agent.clients
import com.fasterxml.jackson.databind.JsonNode
import ee.carlrobert.codegpt.codecompletions.InfillRequest
import ee.carlrobert.codegpt.completions.CompletionError
import ee.carlrobert.codegpt.util.JsonMapper
import io.ktor.client.HttpClient
import io.ktor.client.request.header
import io.ktor.client.request.preparePost
import io.ktor.client.request.setBody
import io.ktor.client.statement.bodyAsText
import io.ktor.http.ContentType
import io.ktor.http.HttpHeaders
import io.ktor.http.contentType
import io.ktor.http.isSuccess
internal const val MAX_COMPLETION_TOKENS = 128
internal const val CUSTOM_SERVICE_API_KEY_PLACEHOLDER = "${'$'}CUSTOM_SERVICE_API_KEY"
internal interface CodeCompletionCapable : AutoCloseable {
suspend fun getCodeCompletion(infillRequest: InfillRequest): String
}
internal class CodeCompletionClientException(
val statusCode: Int,
val body: String
) : RuntimeException("HTTP $statusCode: $body")
internal fun Throwable.toCompletionError(): CompletionError {
return when (this) {
is CodeCompletionClientException -> CompletionError(
message = extractErrorMessage(body).ifBlank {
"Request failed with status $statusCode"
}
)
else -> CompletionError(message ?: "Request failed")
}
}
internal suspend fun postCompletionJson(
client: HttpClient,
url: String,
headers: Map<String, String>,
body: String
): String {
return client.preparePost(url) {
contentType(ContentType.Application.Json)
header(HttpHeaders.Accept, ContentType.Application.Json.toString())
headers.forEach { (key, value) -> header(key, value) }
setBody(body)
}.execute { response ->
val payload = response.bodyAsText()
if (!response.status.isSuccess()) {
throw CodeCompletionClientException(response.status.value, payload)
}
payload
}
}
internal fun parseOpenAITextCompletion(payload: String): String {
val root = JsonMapper.mapper.readTree(payload)
val choice = root.path("choices").firstOrNull() ?: return ""
return choice.path("text").asText("")
}
internal fun parseOpenAIChatCompletion(payload: String): String {
val root = JsonMapper.mapper.readTree(payload)
val choice = root.path("choices").firstOrNull() ?: return ""
return choice.path("message").path("content").asText()
.ifBlank { choice.path("delta").path("content").asText() }
}
internal fun parseOllamaCompletion(payload: String): String {
val root = JsonMapper.mapper.readTree(payload)
return root.path("response").asText()
.ifBlank { root.path("message").path("content").asText() }
}
internal fun parseLlamaCompletion(payload: String): String {
return JsonMapper.mapper.readTree(payload).path("content").asText("")
}
internal fun buildOpenAIStyleCompletionBody(model: String, infillRequest: InfillRequest): String {
return JsonMapper.mapper.writeValueAsString(
mapOf(
"model" to model,
"prompt" to infillRequest.prefix,
"suffix" to infillRequest.suffix,
"stream" to false,
"max_tokens" to MAX_COMPLETION_TOKENS,
"temperature" to 0.0,
"frequency_penalty" to 0.0,
"presence_penalty" to 0.0,
"stop" to infillRequest.stopTokens.ifEmpty { null }
).filterValues { it != null }
)
}
private fun JsonNode.firstOrNull(): JsonNode? {
return if (isArray && size() > 0) get(0) else null
}
private fun extractErrorMessage(body: String): String {
return runCatching {
val root = JsonMapper.mapper.readTree(body)
root.path("error").path("message").asText()
.ifBlank { root.path("message").asText() }
.ifBlank { body }
}.getOrDefault(body)
}

38
src/main/kotlin/ee/carlrobert/codegpt/agent/clients/CodeGPTMistralAILLMClient.kt Normal file
View file

@ -0,0 +1,38 @@
package ee.carlrobert.codegpt.agent.clients
import ai.koog.prompt.executor.clients.mistralai.MistralAIClientSettings
import ai.koog.prompt.executor.clients.mistralai.MistralAILLMClient
import ee.carlrobert.codegpt.codecompletions.InfillRequest
import io.ktor.client.HttpClient
import io.ktor.http.HttpHeaders
class CodeGPTMistralAILLMClient(
private val apiKey: String,
private val baseUrl: String,
private val model: String,
private val baseClient: HttpClient = HttpClientProvider.createHttpClient()
) : MistralAILLMClient(
apiKey = apiKey,
settings = MistralAIClientSettings(baseUrl = baseUrl),
baseClient = baseClient
),
CodeCompletionCapable {
override suspend fun getCodeCompletion(infillRequest: InfillRequest): String {
val headers = buildMap {
if (apiKey.isNotBlank()) {
put(HttpHeaders.Authorization, "Bearer $apiKey")
}
put("X-LLM-Application-Tag", "codegpt")
}
return parseOpenAIChatCompletion(
postCompletionJson(
client = baseClient,
url = "$baseUrl/v1/fim/completions",
headers = headers,
body = buildOpenAIStyleCompletionBody(model, infillRequest)
)
)
}
}

86
src/main/kotlin/ee/carlrobert/codegpt/agent/clients/CodeGPTOllamaClient.kt Normal file
View file

@ -0,0 +1,86 @@
package ee.carlrobert.codegpt.agent.clients
import ai.koog.prompt.executor.clients.LLMEmbeddingProvider
import ai.koog.prompt.executor.clients.LLMClient
import ai.koog.prompt.executor.ollama.client.OllamaClient
import com.intellij.openapi.components.service
import ee.carlrobert.codegpt.codecompletions.InfillRequest
import ee.carlrobert.codegpt.settings.models.ModelSettings
import ee.carlrobert.codegpt.settings.service.FeatureType
import ee.carlrobert.codegpt.settings.service.ollama.OllamaSettingsState
import ee.carlrobert.codegpt.util.JsonMapper
import io.ktor.client.HttpClient
import io.ktor.http.HttpHeaders
class CodeGPTOllamaClient private constructor(
private val delegate: OllamaClient,
private val settings: OllamaSettingsState,
private val baseUrl: String,
private val apiKey: String?,
private val baseClient: HttpClient
) : LLMClient by delegate,
LLMEmbeddingProvider by delegate,
CodeCompletionCapable {
constructor(
settings: OllamaSettingsState,
baseUrl: String,
apiKey: String?,
baseClient: HttpClient = HttpClientProvider.createHttpClient()
) : this(
delegate = OllamaClient(baseUrl = baseUrl, baseClient = baseClient),
settings = settings,
baseUrl = baseUrl,
apiKey = apiKey,
baseClient = baseClient
)
override suspend fun getCodeCompletion(infillRequest: InfillRequest): String {
val model = service<ModelSettings>().getModelSelectionForFeature(FeatureType.CODE_COMPLETION).modelId
val stopTokens = buildList {
if (infillRequest.stopTokens.isNotEmpty()) {
addAll(infillRequest.stopTokens)
}
}.toMutableList()
val prompt = if (settings.fimOverride) {
settings.fimTemplate.stopTokens?.let { stopTokens.addAll(it) }
settings.fimTemplate.buildPrompt(infillRequest)
} else {
infillRequest.prefix
}
val body = JsonMapper.mapper.writeValueAsString(
mapOf(
"model" to model,
"prompt" to prompt,
"suffix" to if (settings.fimOverride) null else infillRequest.suffix,
"stream" to false,
"raw" to true,
"options" to mapOf(
"stop" to stopTokens.ifEmpty { null },
"num_predict" to MAX_COMPLETION_TOKENS,
"temperature" to 0.4
).filterValues { it != null }
).filterValues { it != null }
)
val headers = buildMap {
apiKey?.takeIf { it.isNotBlank() }?.let {
put(HttpHeaders.Authorization, "Bearer $it")
}
}
return parseOllamaCompletion(
postCompletionJson(
client = baseClient,
url = "$baseUrl/api/generate",
headers = headers,
body = body
)
)
}
override fun close() {
delegate.close()
}
}

42
src/main/kotlin/ee/carlrobert/codegpt/agent/clients/CodeGPTOpenAILLMClient.kt Normal file
View file

@ -0,0 +1,42 @@
package ee.carlrobert.codegpt.agent.clients
import ai.koog.prompt.executor.clients.openai.OpenAIClientSettings
import ai.koog.prompt.executor.clients.openai.OpenAILLMClient
import ee.carlrobert.codegpt.codecompletions.InfillRequest
import io.ktor.client.HttpClient
import io.ktor.http.HttpHeaders
class CodeGPTOpenAILLMClient(
private val apiKey: String,
private val organization: String?,
private val baseUrl: String,
private val model: String,
private val baseClient: HttpClient = HttpClientProvider.createHttpClient()
) : OpenAILLMClient(
apiKey = apiKey,
settings = OpenAIClientSettings(baseUrl = baseUrl),
baseClient = baseClient
),
CodeCompletionCapable {
override suspend fun getCodeCompletion(infillRequest: InfillRequest): String {
val headers = buildMap {
if (apiKey.isNotBlank()) {
put(HttpHeaders.Authorization, "Bearer $apiKey")
}
put("X-LLM-Application-Tag", "codegpt")
organization?.takeIf { it.isNotBlank() }?.let {
put("OpenAI-Organization", it)
}
}
return parseOpenAITextCompletion(
postCompletionJson(
client = baseClient,
url = "$baseUrl/v1/completions",
headers = headers,
body = buildOpenAIStyleCompletionBody(model, infillRequest)
)
)
}
}

373
src/main/kotlin/ee/carlrobert/codegpt/agent/clients/CustomOpenAILLMClient.kt
View file

@ -17,22 +17,30 @@ import ai.koog.prompt.message.ResponseMetaInfo
import ai.koog.prompt.params.LLMParams
import ai.koog.prompt.streaming.StreamFrame
import ai.koog.prompt.streaming.buildStreamFrameFlow
import com.fasterxml.jackson.core.type.TypeReference
import com.fasterxml.jackson.databind.ObjectMapper
import ee.carlrobert.codegpt.codecompletions.InfillPromptTemplate
import ee.carlrobert.codegpt.codecompletions.InfillRequest
import ee.carlrobert.codegpt.settings.Placeholder
import ee.carlrobert.codegpt.settings.service.custom.CustomServiceCodeCompletionSettingsState
import ee.carlrobert.codegpt.settings.service.custom.CustomServiceChatCompletionSettingsState
import ee.carlrobert.codegpt.settings.service.custom.CustomServicePlaceholders
import ee.carlrobert.codegpt.util.JsonMapper
import io.github.oshai.kotlinlogging.KotlinLogging
import io.ktor.client.*
import io.ktor.client.plugins.*
import io.ktor.client.request.*
import io.ktor.http.HttpHeaders
import org.apache.commons.text.StringEscapeUtils
import kotlinx.coroutines.flow.Flow
import kotlinx.datetime.Clock
import kotlinx.serialization.ExperimentalSerializationApi
import kotlinx.serialization.KSerializer
import kotlinx.serialization.builtins.ListSerializer
import kotlinx.serialization.descriptors.elementNames
import kotlinx.serialization.json.JsonElement
import kotlinx.serialization.json.JsonTransformingSerializer
import kotlinx.serialization.json.buildJsonObject
import kotlinx.serialization.json.jsonObject
import kotlinx.serialization.json.*
import java.net.URI
import kotlin.io.encoding.ExperimentalEncodingApi
import kotlin.time.Clock
/**
* Configuration settings for connecting to the CustomOpenAI API.
@ -55,10 +63,11 @@ class CustomOpenAIClientSettings(
* @param clock Clock instance used for tracking response metadata timestamps.
*/
class CustomOpenAILLMClient(
apiKey: String,
private val apiKey: String,
private val settings: CustomOpenAIClientSettings,
private val state: CustomServiceChatCompletionSettingsState,
baseClient: HttpClient = HttpClient(),
private val chatState: CustomServiceChatCompletionSettingsState? = null,
private val codeCompletionState: CustomServiceCodeCompletionSettingsState? = null,
private val baseClient: HttpClient = HttpClient(),
clock: Clock = Clock.System
) : AbstractOpenAILLMClient<CustomOpenAIChatCompletionResponse, CustomOpenAIChatCompletionStreamResponse>(
apiKey,
@ -67,7 +76,8 @@ class CustomOpenAILLMClient(
clock,
staticLogger,
OpenAICompatibleToolDescriptorSchemaGenerator()
) {
),
CodeCompletionCapable {
data object CustomOpenAI : LLMProvider("custom-openai", "Custom OpenAI")
companion object {
@ -83,12 +93,8 @@ class CustomOpenAILLMClient(
baseClient: HttpClient = HttpClient(),
timeoutConfig: ConnectionTimeoutConfig = ConnectionTimeoutConfig()
): CustomOpenAILLMClient {
val stateUrl = state.url ?: throw IllegalStateException("Url not set")
val uri = URI.create(stateUrl)
val authority = uri.authority ?: uri.host
val settings = CustomOpenAIClientSettings(
baseUrl = "${uri.scheme}://${authority}",
chatCompletionsPath = uri.path,
val settings = createClientSettings(
url = state.url ?: throw IllegalStateException("Url not set"),
timeoutConfig = timeoutConfig
)
val clientWithCustomHeaders = baseClient.config {
@ -103,17 +109,74 @@ class CustomOpenAILLMClient(
header(
normalizedKey,
value.replace($$"$CUSTOM_SERVICE_API_KEY", apiKey)
value.replace(CUSTOM_SERVICE_API_KEY_PLACEHOLDER, apiKey)
)
}
}
}
return CustomOpenAILLMClient(apiKey, settings, state, clientWithCustomHeaders)
return CustomOpenAILLMClient(
apiKey = apiKey,
settings = settings,
chatState = state,
baseClient = clientWithCustomHeaders
)
}
fun fromCodeCompletionSettingsState(
apiKey: String,
state: CustomServiceCodeCompletionSettingsState,
baseClient: HttpClient = HttpClient(),
timeoutConfig: ConnectionTimeoutConfig = ConnectionTimeoutConfig()
): CustomOpenAILLMClient {
return CustomOpenAILLMClient(
apiKey = apiKey,
settings = createClientSettings(
url = state.url ?: throw IllegalStateException("Url not set"),
timeoutConfig = timeoutConfig
),
codeCompletionState = state,
baseClient = baseClient
)
}
private fun createClientSettings(
url: String,
timeoutConfig: ConnectionTimeoutConfig
): CustomOpenAIClientSettings {
val uri = URI.create(url)
val authority = uri.authority ?: uri.host
return CustomOpenAIClientSettings(
baseUrl = "${uri.scheme}://${authority}",
chatCompletionsPath = uri.path,
timeoutConfig = timeoutConfig
)
}
}
override fun llmProvider(): LLMProvider = CustomOpenAI
override suspend fun getCodeCompletion(infillRequest: InfillRequest): String {
val state = requireNotNull(codeCompletionState) {
"Custom OpenAI code completion requested on a chat-only client"
}
val url = requireNotNull(state.url)
val payload = postCompletionJson(
client = baseClient,
url = url,
headers = replaceCredentialPlaceholders(state.headers, apiKey),
body = buildCodeCompletionRequestBody(state, infillRequest)
)
return if (
state.infillTemplate == InfillPromptTemplate.CHAT_COMPLETION ||
state.parseResponseAsChatCompletions
) {
parseOpenAIChatCompletion(payload)
} else {
parseOpenAITextCompletion(payload)
}
}
override fun serializeProviderChatRequest(
messages: List<OpenAIMessage>,
model: LLModel,
@ -122,12 +185,23 @@ class CustomOpenAILLMClient(
params: LLMParams,
stream: Boolean
): String {
val customParams = params.toCustomOpenAIParams(state)
val state = requireNotNull(chatState) {
"Custom OpenAI chat request requested on a code-completion-only client"
}
val customParams: CustomOpenAIParams = params.toCustomOpenAIParams(state)
val streamRequest = state.shouldStream()
val additionalProperties = buildCustomOpenAIAdditionalProperties(
body = state.body,
messages = messages,
streamRequest = streamRequest,
credential = apiKey,
json = json
)
val responseFormat = createResponseFormat(customParams.schema, model)
val request = CustomOpenAIChatCompletionRequest(
messages = messages,
model = model.id,
stream = stream,
stream = streamRequest,
temperature = customParams.temperature,
tools = tools,
toolChoice = customParams.toolChoice?.toOpenAIToolChoice(),
@ -148,12 +222,153 @@ class CustomOpenAILLMClient(
models = customParams.models,
route = customParams.route,
user = customParams.user,
additionalProperties = customParams.additionalProperties,
additionalProperties = additionalProperties,
)
return json.encodeToString(CustomOpenAIChatCompletionRequestSerializer, request)
}
private fun buildCodeCompletionRequestBody(
state: CustomServiceCodeCompletionSettingsState,
infillRequest: InfillRequest
): String {
return if (state.infillTemplate == InfillPromptTemplate.CHAT_COMPLETION) {
val messages = listOf(
mapOf(
"role" to "system",
"content" to (
"You are a code completion assistant. Complete the code between the given prefix and suffix. " +
"Return only the missing code that should be inserted, without any formatting, explanations, or markdown."
)
),
mapOf(
"role" to "user",
"content" to (
"<PREFIX>\n${infillRequest.prefix}\n</PREFIX>\n\n" +
"<SUFFIX>\n${infillRequest.suffix}\n</SUFFIX>\n\nComplete:"
)
)
)
val transformedBody = state.body.entries.mapNotNull { (key, value) ->
when (key.lowercase()) {
"messages" -> key to messages
"prompt", "suffix" -> null
"stream" -> key to false
else -> key to transformCodeCompletionValue(
value = value,
template = InfillPromptTemplate.CHAT_COMPLETION,
infillRequest = infillRequest
)
}
}.toMap().toMutableMap()
if (!transformedBody.containsKey("messages")) {
transformedBody["messages"] = messages
}
JsonMapper.mapper.writeValueAsString(transformedBody)
} else {
JsonMapper.mapper.writeValueAsString(
state.body.entries.associate { (key, value) ->
when (key.lowercase()) {
"stop" -> key to transformCodeCompletionStopValue(
value = value,
template = state.infillTemplate,
infillRequest = infillRequest
)
"stream" -> key to false
else -> key to transformCodeCompletionValue(
value = value,
template = state.infillTemplate,
infillRequest = infillRequest
)
}
}
)
}
}
private fun replaceCredentialPlaceholders(
headers: Map<String, String>,
credential: String?
): Map<String, String> {
return headers.mapValues { (_, value) ->
if (credential != null && value.contains(CUSTOM_SERVICE_API_KEY_PLACEHOLDER)) {
value.replace(CUSTOM_SERVICE_API_KEY_PLACEHOLDER, credential)
} else {
value
}
}
}
private fun transformCodeCompletionStopValue(
value: Any,
template: InfillPromptTemplate,
infillRequest: InfillRequest
): Any {
if (value !is String) {
return transformCodeCompletionValue(value, template, infillRequest)
}
if (value.isEmpty()) {
return value
}
return if (value.startsWith("[") && value.endsWith("]")) {
ObjectMapper().readValue(value, object : TypeReference<List<String>>() {})
} else {
value.split(",").map { StringEscapeUtils.unescapeJava(it.trim()) }
}
}
private fun transformCodeCompletionValue(
value: Any,
template: InfillPromptTemplate,
infillRequest: InfillRequest
): Any {
if (value !is String) {
return when (value) {
is Map<*, *> -> value.entries
.filter { it.key != null }
.associate { (nestedKey, nestedValue) ->
nestedKey.toString() to transformCodeCompletionValue(
value = nestedValue ?: "",
template = template,
infillRequest = infillRequest
)
}
is Iterable<*> -> value.map {
transformCodeCompletionValue(it ?: "", template, infillRequest)
}
is Array<*> -> value.map {
transformCodeCompletionValue(it ?: "", template, infillRequest)
}
else -> value
}
}
val replaced = if (value.contains(CUSTOM_SERVICE_API_KEY_PLACEHOLDER)) {
value.replace(CUSTOM_SERVICE_API_KEY_PLACEHOLDER, apiKey)
} else {
value
}
return when (replaced) {
Placeholder.FIM_PROMPT.code -> template.buildPrompt(infillRequest)
Placeholder.PREFIX.code -> infillRequest.prefix
Placeholder.SUFFIX.code -> infillRequest.suffix
else -> replaced.takeIf {
it.contains(Placeholder.PREFIX.code) || it.contains(Placeholder.SUFFIX.code)
}?.replace(Placeholder.PREFIX.code, infillRequest.prefix)
?.replace(Placeholder.SUFFIX.code, infillRequest.suffix)
?: replaced
}
}
override fun processProviderChatResponse(response: CustomOpenAIChatCompletionResponse): List<LLMChoice> {
require(response.choices.isNotEmpty()) { "Empty choices in response" }
return response.choices.map {
@ -175,14 +390,14 @@ class CustomOpenAILLMClient(
response.collect { chunk ->
chunk.choices.firstOrNull()?.let { choice ->
choice.delta.content?.let { emitAppend(it) }
choice.delta.content?.let { emitTextDelta(it) }
choice.delta.toolCalls?.forEach { openAIToolCall ->
val index = openAIToolCall.index ?: 0
val id = openAIToolCall.id
val functionName = openAIToolCall.function.name
val functionArgs = openAIToolCall.function.arguments
upsertToolCall(index, id, functionName, functionArgs)
val id = openAIToolCall.id.orEmpty()
val functionName = openAIToolCall.function.name.orEmpty()
val functionArgs = openAIToolCall.function.arguments.orEmpty()
emitToolCallDelta(id, functionName, functionArgs, index)
}
choice.finishReason?.let { finishReason = it }
@ -195,8 +410,7 @@ class CustomOpenAILLMClient(
}
override suspend fun moderate(prompt: Prompt, model: LLModel): ModerationResult {
logger.warn { "Moderation is not supported by CustomOpenAI API" }
throw UnsupportedOperationException("Moderation is not supported by CustomOpenAI API.")
throw UnsupportedOperationException("Moderation not supported.")
}
@OptIn(ExperimentalEncodingApi::class)
@ -269,6 +483,111 @@ class CustomOpenAILLMClient(
}
}
internal fun buildCustomOpenAIAdditionalProperties(
body: Map<String, Any>,
messages: List<OpenAIMessage>,
streamRequest: Boolean,
credential: String,
json: Json
): Map<String, JsonElement>? = body
.filterKeys { it !in CUSTOM_OPENAI_RESERVED_BODY_KEYS }
.mapValues { (key, value) ->
transformCustomOpenAIBodyValue(
key = key,
value = value,
streamRequest = streamRequest,
messages = messages,
credential = credential,
json = json
)
}
.takeIf { it.isNotEmpty() }
internal fun transformCustomOpenAIBodyValue(
key: String?,
value: Any?,
streamRequest: Boolean,
messages: List<OpenAIMessage>,
credential: String,
json: Json
): JsonElement {
return when (value) {
null -> JsonNull
is JsonElement -> value
is String -> when {
!streamRequest && key == "stream" -> JsonPrimitive(false)
CustomServicePlaceholders.isMessages(value) -> json.parseToJsonElement(
json.encodeToString(ListSerializer(OpenAIMessage.serializer()), messages)
)
CustomServicePlaceholders.isPrompt(value) -> JsonPrimitive(
renderCustomOpenAIPrompt(
messages,
json
)
)
value.contains($$"$CUSTOM_SERVICE_API_KEY") -> {
JsonPrimitive(value.replace($$"$CUSTOM_SERVICE_API_KEY", credential))
}
else -> JsonPrimitive(value)
}
is Number -> JsonPrimitive(value)
is Boolean -> JsonPrimitive(value)
is Map<*, *> -> JsonObject(
value.entries
.filter { (nestedKey, _) -> nestedKey != null }
.associate { (nestedKey, nestedValue) ->
nestedKey.toString() to transformCustomOpenAIBodyValue(
key = nestedKey.toString(),
value = nestedValue,
streamRequest = streamRequest,
messages = messages,
credential = credential,
json = json
)
}
)
is Iterable<*> -> JsonArray(
value.map { item ->
transformCustomOpenAIBodyValue(
key = null,
value = item,
streamRequest = streamRequest,
messages = messages,
credential = credential,
json = json
)
}
)
is Array<*> -> JsonArray(
value.map { item ->
transformCustomOpenAIBodyValue(
key = null,
value = item,
streamRequest = streamRequest,
messages = messages,
credential = credential,
json = json
)
}
)
else -> JsonPrimitive(value.toString())
}
}
internal fun renderCustomOpenAIPrompt(messages: List<OpenAIMessage>, json: Json): String {
return messages.joinToString(separator = "\n\n") { message ->
message.content?.text()?.takeIf { it.isNotBlank() }
?: json.encodeToString(OpenAIMessage.serializer(), message)
}
}
internal object CustomOpenAIChatCompletionRequestSerializer :
CustomOpenAIAdditionalPropertiesFlatteningSerializer(CustomOpenAIChatCompletionRequest.serializer())

23
src/main/kotlin/ee/carlrobert/codegpt/agent/clients/CustomOpenAIParams.kt
View file

@ -11,7 +11,7 @@ internal fun LLMParams.toCustomOpenAIParams(state: CustomServiceChatCompletionSe
val mergedAdditionalProperties = buildMap {
body
.filterKeys { it !in RESERVED_BODY_KEYS }
.filterKeys { it !in CUSTOM_OPENAI_RESERVED_BODY_KEYS }
.forEach { (key, value) -> put(key, value.toJsonElement()) }
}.takeIf { it.isNotEmpty() }
@ -33,7 +33,11 @@ internal fun LLMParams.toCustomOpenAIParams(state: CustomServiceChatCompletionSe
)
}
private val RESERVED_BODY_KEYS = setOf(
internal fun CustomServiceChatCompletionSettingsState.shouldStream(): Boolean {
return body.findValue("stream").asBoolean() == true
}
internal val CUSTOM_OPENAI_RESERVED_BODY_KEYS = setOf(
"frequencyPenalty", "frequency_penalty",
"maxTokens", "max_tokens",
"messages",
@ -41,6 +45,7 @@ private val RESERVED_BODY_KEYS = setOf(
"presencePenalty", "presence_penalty",
"repetitionPenalty", "repetition_penalty",
"stop",
"stream",
"temperature",
"toolChoice", "tool_choice",
"tools",
@ -79,6 +84,17 @@ private fun Any?.asInt(): Int? = when (this) {
else -> null
}
private fun Any?.asBoolean(): Boolean? = when (this) {
is Boolean -> this
is String -> when (trim().lowercase()) {
"true" -> true
"false" -> false
else -> null
}
else -> null
}
private fun Any?.asJsonElementMap(): Map<String, JsonElement>? {
if (this !is Map<*, *>) return null
@ -165,8 +181,6 @@ private fun Any?.toJsonElement(): JsonElement = when (this) {
* @property schema JSON Schema to constrain model output (validated when supported).
* @property toolChoice Controls if/which tool must be called (`none`/`auto`/`required`/specific).
* @property user stable end-user identifier
* @property includeThoughts Request inclusion of model “thoughts”/reasoning traces (model-dependent).
* @property thinkingBudget Soft cap on tokens spent on internal reasoning (reasoning models).
* @property additionalProperties Additional properties that can be used to store custom parameters.
* @property frequencyPenalty Number in [-2.0, 2.0]—penalizes frequent tokens to reduce repetition.
* @property presencePenalty Number in [-2.0, 2.0]—encourages introduction of new tokens/topics.
@ -184,7 +198,6 @@ private fun Any?.toJsonElement(): JsonElement = when (this) {
* Use empty list [] for no transformations.
* @property models List of allowed models for this request.
* @property route Request routing identifier.
* @property provider Model provider preferences.
*/
public open class CustomOpenAIParams(
temperature: Double? = null,

20
src/main/kotlin/ee/carlrobert/codegpt/agent/clients/HttpClientProvider.kt
View file

@ -4,9 +4,11 @@ import com.intellij.openapi.diagnostic.Logger
import com.intellij.util.net.ssl.CertificateManager
import ee.carlrobert.codegpt.settings.advanced.AdvancedSettings
import ee.carlrobert.codegpt.settings.advanced.AdvancedSettingsState
import ee.carlrobert.codegpt.settings.configuration.ConfigurationSettings
import io.ktor.client.*
import io.ktor.client.engine.apache5.*
import io.ktor.client.plugins.*
import io.ktor.client.plugins.logging.*
import org.apache.hc.client5.http.auth.AuthScope
import org.apache.hc.client5.http.auth.UsernamePasswordCredentials
import org.apache.hc.client5.http.impl.async.HttpAsyncClientBuilder
@ -28,6 +30,7 @@ object HttpClientProvider {
*/
fun createHttpClient(): HttpClient {
val advancedSettings = AdvancedSettings.getCurrentState()
val debugModeEnabled = ConfigurationSettings.getState().debugModeEnabled
return HttpClient(Apache5) {
engine {
configureProxy(advancedSettings)
@ -38,6 +41,23 @@ object HttpClientProvider {
requestTimeoutMillis = advancedSettings.readTimeout.toLong() * 1000
socketTimeoutMillis = advancedSettings.readTimeout.toLong() * 1000
}
if (debugModeEnabled) {
install(Logging) {
logger = object : io.ktor.client.plugins.logging.Logger {
override fun log(message: String) {
HttpClientProvider.logger.info("[KtorHTTP] $message")
}
}
level = LogLevel.ALL
sanitizeHeader { header ->
header.equals("Authorization", ignoreCase = true)
|| header.equals("Proxy-Authorization", ignoreCase = true)
|| header.equals("X-Api-Key", ignoreCase = true)
|| header.equals("Api-Key", ignoreCase = true)
}
}
}
}
}

109
src/main/kotlin/ee/carlrobert/codegpt/agent/clients/InceptionAILLMClient.kt
View file

@ -1,32 +1,111 @@
package ee.carlrobert.codegpt.agent.clients
import ai.koog.prompt.dsl.ModerationResult
import ai.koog.prompt.dsl.Prompt
import ai.koog.prompt.executor.clients.openai.OpenAIClientSettings
import ai.koog.prompt.executor.clients.openai.OpenAILLMClient
import ai.koog.prompt.llm.LLMProvider
import ai.koog.prompt.llm.LLModel
import com.intellij.openapi.components.service
import ee.carlrobert.codegpt.codecompletions.InfillRequest
import ee.carlrobert.codegpt.completions.inception.InceptionApplyRequest
import ee.carlrobert.codegpt.completions.inception.InceptionChatCompletionResponse
import ee.carlrobert.codegpt.completions.inception.InceptionNextEditRequest
import ee.carlrobert.codegpt.settings.models.ModelSettings
import ee.carlrobert.codegpt.settings.service.FeatureType
import ee.carlrobert.codegpt.util.JsonMapper
import io.ktor.client.*
import io.ktor.client.request.*
import io.ktor.client.statement.*
import io.ktor.http.*
import kotlinx.coroutines.runBlocking
class InceptionAILLMClient(
apiKey: String,
baseClient: HttpClient = HttpClient()
private val apiKey: String,
private val settings: OpenAIClientSettings = OpenAIClientSettings(
baseUrl = System.getProperty(INCEPTION_BASE_URL_PROPERTY) ?: DEFAULT_BASE_URL,
chatCompletionsPath = DEFAULT_CHAT_COMPLETIONS_PATH
),
private val baseClient: HttpClient = HttpClient()
) : OpenAILLMClient(
apiKey = apiKey,
settings = OpenAIClientSettings(
baseUrl = "https://api.inceptionlabs.ai",
chatCompletionsPath = "v1/chat/completions"
),
settings = settings,
baseClient = baseClient
) {
),
CodeCompletionCapable {
data object Inception : LLMProvider("inception", "Inception")
fun getApplyEditCompletion(request: InceptionApplyRequest): InceptionChatCompletionResponse {
return runBlocking {
val requestBody = JsonMapper.mapper.writeValueAsString(request)
val response = baseClient.post("${settings.baseUrl}$APPLY_EDIT_COMPLETION_PATH") {
contentType(ContentType.Application.Json)
if (apiKey.isNotBlank()) {
header(HttpHeaders.Authorization, "Bearer $apiKey")
}
setBody(requestBody)
}
val body = response.bodyAsText()
if (!response.status.isSuccess()) {
throw RuntimeException(
"Inception apply request failed with status ${response.status.value}: $body"
)
}
JsonMapper.mapper.readValue(body, InceptionChatCompletionResponse::class.java)
}
}
fun getNextEditCompletion(request: InceptionNextEditRequest): InceptionChatCompletionResponse {
return runBlocking {
val requestBody = JsonMapper.mapper.writeValueAsString(request)
val response = baseClient.post("${settings.baseUrl}$NEXT_EDIT_COMPLETION_PATH") {
contentType(ContentType.Application.Json)
if (apiKey.isNotBlank()) {
header(HttpHeaders.Authorization, "Bearer $apiKey")
}
setBody(requestBody)
}
val body = response.bodyAsText()
if (!response.status.isSuccess()) {
throw RuntimeException(
"Inception next-edit request failed with status ${response.status.value}: $body"
)
}
JsonMapper.mapper.readValue(body, InceptionChatCompletionResponse::class.java)
}
}
override fun llmProvider(): LLMProvider = Inception
override suspend fun moderate(
prompt: Prompt,
model: LLModel
): ModerationResult {
throw UnsupportedOperationException("Moderation is not supported by Inception API.")
override suspend fun getCodeCompletion(infillRequest: InfillRequest): String {
val body = JsonMapper.mapper.writeValueAsString(
mapOf(
"model" to service<ModelSettings>()
.getModelSelectionForFeature(FeatureType.CODE_COMPLETION)
.modelId,
"prompt" to infillRequest.prefix,
"suffix" to infillRequest.suffix,
"stream" to false
)
)
val headers = buildMap {
if (apiKey.isNotBlank()) {
put(HttpHeaders.Authorization, "Bearer $apiKey")
}
}
return parseOpenAITextCompletion(
postCompletionJson(
client = baseClient,
url = "${settings.baseUrl}/v1/fim/completions",
headers = headers,
body = body
)
)
}
companion object {
private const val DEFAULT_BASE_URL = "https://api.inceptionlabs.ai"
private const val DEFAULT_CHAT_COMPLETIONS_PATH = "v1/chat/completions"
private const val APPLY_EDIT_COMPLETION_PATH = "/v1/apply/completions"
private const val NEXT_EDIT_COMPLETION_PATH = "/v1/edit/completions"
private const val INCEPTION_BASE_URL_PROPERTY = "inception.baseUrl"
}
}

65
src/main/kotlin/ee/carlrobert/codegpt/agent/clients/LlamaCppLLMClient.kt Normal file
View file

@ -0,0 +1,65 @@
package ee.carlrobert.codegpt.agent.clients
import ai.koog.prompt.executor.clients.openai.OpenAIClientSettings
import ai.koog.prompt.executor.clients.openai.OpenAILLMClient
import ai.koog.prompt.llm.LLMProvider
import ee.carlrobert.codegpt.codecompletions.InfillRequest
import ee.carlrobert.codegpt.completions.llama.LlamaModel
import ee.carlrobert.codegpt.settings.service.llama.LlamaSettingsState
import ee.carlrobert.codegpt.util.JsonMapper
import io.ktor.client.*
class LlamaCppLLMClient(
private val baseUrl: String,
private val state: LlamaSettingsState,
private val baseClient: HttpClient = HttpClient()
) : OpenAILLMClient(
apiKey = "",
settings = OpenAIClientSettings(
baseUrl = baseUrl,
chatCompletionsPath = "v1/chat/completions"
),
baseClient = baseClient
) ,
CodeCompletionCapable {
data object LlamaCpp : LLMProvider("llama_cpp", "llama.cpp")
override fun llmProvider(): LLMProvider = LlamaCpp
override suspend fun getCodeCompletion(infillRequest: InfillRequest): String {
val promptTemplate = if (state.isUseCustomModel) {
state.localModelInfillPromptTemplate
} else {
LlamaModel.findByHuggingFaceModel(state.huggingFaceModel).infillPromptTemplate
}
val stopTokens = buildList {
promptTemplate.stopTokens?.let { addAll(it) }
if (infillRequest.stopTokens.isNotEmpty()) {
addAll(infillRequest.stopTokens)
}
}.ifEmpty { null }
val body = JsonMapper.mapper.writeValueAsString(
mapOf(
"prompt" to promptTemplate.buildPrompt(infillRequest),
"stream" to false,
"n_predict" to MAX_COMPLETION_TOKENS,
"temperature" to 0.0,
"top_k" to state.topK,
"top_p" to state.topP,
"min_p" to state.minP,
"repeat_penalty" to state.repeatPenalty,
"stop" to stopTokens
).filterValues { it != null }
)
return parseLlamaCompletion(
postCompletionJson(
client = baseClient,
url = "$baseUrl/completion",
headers = emptyMap(),
body = body
)
)
}
}

78
src/main/kotlin/ee/carlrobert/codegpt/agent/clients/ProxyAILLMClient.kt
View file

@ -18,17 +18,27 @@ import ai.koog.prompt.message.Message
import ai.koog.prompt.message.ResponseMetaInfo
import ai.koog.prompt.params.LLMParams
import ai.koog.prompt.streaming.StreamFrame
import ai.koog.prompt.streaming.StreamFrameFlowBuilder
import ai.koog.prompt.streaming.buildStreamFrameFlow
import ee.carlrobert.codegpt.agent.clients.ProxyAIClientSettings.Companion.DEFAULT_AUTO_APPLY_PATH
import ee.carlrobert.codegpt.agent.clients.ProxyAIClientSettings.Companion.DEFAULT_USER_DETAILS_PATH
import ee.carlrobert.codegpt.completions.autoapply.AutoApplyRequest
import ee.carlrobert.codegpt.completions.autoapply.AutoApplyResponse
import ee.carlrobert.codegpt.settings.service.codegpt.CodeGPTApiException
import ee.carlrobert.codegpt.settings.service.codegpt.CodeGPTUserDetails
import ee.carlrobert.codegpt.util.JsonMapper
import io.github.oshai.kotlinlogging.KotlinLogging
import io.ktor.client.*
import io.ktor.client.request.*
import io.ktor.client.statement.*
import io.ktor.http.*
import kotlinx.coroutines.flow.Flow
import kotlinx.datetime.Clock
import kotlinx.coroutines.runBlocking
import kotlinx.serialization.json.JsonObject
import kotlinx.serialization.json.JsonObjectBuilder
import kotlinx.serialization.json.JsonPrimitive
import kotlinx.serialization.json.buildJsonObject
import kotlin.io.encoding.ExperimentalEncodingApi
import kotlin.time.Clock
/**
* Configuration settings for connecting to the ProxyAI API.
@ -41,6 +51,8 @@ public class ProxyAIClientSettings(
public companion object {
public const val DEFAULT_BASE_URL: String = "https://codegpt-api.carlrobert.ee"
public const val DEFAULT_CHAT_COMPLETIONS_PATH: String = "/v2/chat/completions"
public const val DEFAULT_AUTO_APPLY_PATH: String = "/v1/code/apply"
public const val DEFAULT_USER_DETAILS_PATH: String = "/v1/users/details"
}
}
@ -48,9 +60,9 @@ public class ProxyAIClientSettings(
* Implementation of [LLMClient] for ProxyAI API.
*/
public class ProxyAILLMClient(
apiKey: String,
private val apiKey: String,
private val settings: ProxyAIClientSettings = ProxyAIClientSettings(),
baseClient: HttpClient = HttpClient(),
private val baseClient: HttpClient = HttpClientProvider.createHttpClient(),
clock: Clock = Clock.System
) : AbstractOpenAILLMClient<ProxyAIChatCompletionResponse, ProxyAIChatCompletionStreamResponse>(
apiKey,
@ -72,6 +84,50 @@ public class ProxyAILLMClient(
override fun llmProvider(): LLMProvider = ProxyAI
fun getApplyEditCompletion(request: AutoApplyRequest): AutoApplyResponse {
return runBlocking {
val response = baseClient.post("${settings.baseUrl}$DEFAULT_AUTO_APPLY_PATH") {
contentType(ContentType.Application.Json)
header(HttpHeaders.Authorization, "Bearer $apiKey")
setBody(JsonMapper.mapper.writeValueAsString(request))
}
val body = response.bodyAsText()
if (!response.status.isSuccess()) {
throw toApiException(response.status.value, body)
}
JsonMapper.mapper.readValue(body, AutoApplyResponse::class.java)
}
}
fun getUserDetails(requestApiKey: String = apiKey): CodeGPTUserDetails {
return runBlocking {
val response = baseClient.get("${settings.baseUrl}$DEFAULT_USER_DETAILS_PATH") {
header(HttpHeaders.Authorization, "Bearer $requestApiKey")
}
val body = response.bodyAsText()
if (!response.status.isSuccess()) {
throw toApiException(response.status.value, body)
}
JsonMapper.mapper.readValue(body, CodeGPTUserDetails::class.java)
}
}
private fun toApiException(status: Int, body: String): CodeGPTApiException {
return runCatching {
JsonMapper.mapper.readValue(body, CodeGPTApiException::class.java).apply {
this.status = status
if (detail.isNullOrBlank()) {
detail = body
}
}
}.getOrElse {
CodeGPTApiException(
status = status,
detail = body.ifBlank { "Request failed with status $status" }
)
}
}
override fun serializeProviderChatRequest(
messages: List<OpenAIMessage>,
model: LLModel,
@ -111,7 +167,6 @@ public class ProxyAILLMClient(
throw LLMClientException(clientName, errorMsg)
}
// Check for errors in choices
response.choices.forEach { choice ->
choice.error?.let { error ->
val errorMsg = "ProxyAI API error (code: ${error.code}): ${error.message}"
@ -139,14 +194,14 @@ public class ProxyAILLMClient(
response.collect { chunk ->
chunk.choices.firstOrNull()?.let { choice ->
choice.delta?.content?.let { emitAppend(it) }
choice.delta?.content?.let { emitTextDelta(it) }
choice.delta?.toolCalls?.forEach { openAIToolCall ->
val index = openAIToolCall.index ?: 0
val id = openAIToolCall.id
val functionName = openAIToolCall.function.name
val functionArgs = openAIToolCall.function.arguments
upsertToolCall(index, id, functionName, functionArgs)
val id = openAIToolCall.id.orEmpty()
val functionName = openAIToolCall.function.name.orEmpty()
val functionArgs = openAIToolCall.function.arguments.orEmpty()
emitToolCallDelta(id, functionName, functionArgs, index)
}
choice.finishReason?.let { finishReason = it }
@ -159,8 +214,7 @@ public class ProxyAILLMClient(
}
override suspend fun moderate(prompt: Prompt, model: LLModel): ModerationResult {
logger.warn { "Moderation is not supported by ProxyAI API" }
throw UnsupportedOperationException("Moderation is not supported by ProxyAI API.")
throw UnsupportedOperationException("Moderation not supported.")
}
@OptIn(ExperimentalEncodingApi::class)

144
src/main/kotlin/ee/carlrobert/codegpt/agent/clients/RetryingPromptExecutor.kt
View file

@ -1,23 +1,32 @@
package ee.carlrobert.codegpt.agent.clients
import ai.koog.agents.core.tools.ToolDescriptor
import ai.koog.http.client.KoogHttpClientException
import ai.koog.prompt.dsl.ModerationResult
import ai.koog.prompt.dsl.Prompt
import ai.koog.prompt.executor.clients.LLMClient
import ai.koog.prompt.executor.clients.anthropic.AnthropicParams
import ai.koog.prompt.executor.clients.anthropic.models.AnthropicThinking
import ai.koog.prompt.executor.clients.openai.OpenAIResponsesParams
import ai.koog.prompt.executor.clients.openai.base.models.ReasoningEffort
import ai.koog.prompt.executor.clients.openai.models.ReasoningConfig
import ai.koog.prompt.executor.clients.openai.models.ReasoningSummary
import ai.koog.prompt.executor.llms.SingleLLMPromptExecutor
import ai.koog.prompt.executor.model.PromptExecutor
import ai.koog.prompt.llm.LLMProvider
import ai.koog.prompt.llm.LLModel
import ai.koog.prompt.message.Message
import ai.koog.prompt.params.LLMParams
import ai.koog.prompt.streaming.StreamFrame
import ee.carlrobert.codegpt.agent.AgentEvents
import io.github.oshai.kotlinlogging.KotlinLogging
import kotlinx.coroutines.CancellationException
import kotlinx.coroutines.TimeoutCancellationException
import kotlinx.coroutines.delay
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.catch
import kotlinx.coroutines.flow.flow
import kotlinx.coroutines.flow.onEach
import kotlinx.serialization.SerializationException
import kotlin.math.pow
import kotlin.random.Random
import kotlin.time.Duration
@ -33,7 +42,10 @@ class RetryingPromptExecutor(
) : PromptExecutor {
companion object {
private const val ANTHROPIC_MIN_THINKING_BUDGET = 1_024
private const val ANTHROPIC_DEFAULT_THINKING_BUDGET = 2_048
private val logger = KotlinLogging.logger { }
private val RETRYABLE_HTTP_STATUS_CODES = setOf(408, 409, 425, 429, 500, 502, 503, 504)
fun fromClient(
client: LLMClient,
@ -46,15 +58,29 @@ class RetryingPromptExecutor(
events = events
)
}
internal fun isRetryableFailure(error: Throwable): Boolean {
if (error is CancellationException && error !is TimeoutCancellationException) {
return false
}
val causes = generateSequence(error) { it.cause }.take(10).toList()
val statusCode = causes
.filterIsInstance<KoogHttpClientException>()
.firstNotNullOfOrNull { it.statusCode }
return statusCode in RETRYABLE_HTTP_STATUS_CODES
}
}
private class RetryStreamingRequestException(cause: Throwable) :
Exception("Retrying streaming request", cause)
override fun executeStreaming(
prompt: Prompt,
model: LLModel,
tools: List<ToolDescriptor>
): Flow<StreamFrame> {
var attempt = 1
var refinedPrompt = prompt
val refinedPrompt = prompt.withReasoningParams(model)
var hasReceivedData: Boolean
fun createStream(attemptNum: Int): Flow<StreamFrame> {
@ -63,22 +89,23 @@ class RetryingPromptExecutor(
return delegate.executeStreaming(refinedPrompt, model, tools)
.onEach { hasReceivedData = true }
.catch { error ->
val shouldRetry = !hasReceivedData && attemptNum < retryPolicy.maxAttempts
val retryable = isRetryableFailure(error)
val shouldRetry =
retryable && !hasReceivedData && attemptNum < retryPolicy.maxAttempts
logger.warn {
"Stream error: ${error.javaClass.simpleName}, " +
"hasReceivedData=$hasReceivedData, " +
"attempt=$attemptNum/${retryPolicy.maxAttempts}, " +
"retryable=$retryable, " +
"willRetry=$shouldRetry"
}
if (shouldRetry) {
attempt++
logger.warn { "Retrying streaming request (attempt $attempt/${retryPolicy.maxAttempts})" }
events?.onRetry(
attempt,
retryPolicy.maxAttempts,
error.javaClass.simpleName
)
val nextAttempt = attemptNum + 1
logger.warn {
"Retrying streaming request (attempt $nextAttempt/${retryPolicy.maxAttempts})"
}
events?.onRetry(nextAttempt, retryPolicy.maxAttempts)
val delayMs = if (attemptNum == 1) {
retryPolicy.initialDelay.inWholeMilliseconds
@ -92,7 +119,7 @@ class RetryingPromptExecutor(
delayMs + (if (jitterMs > 0) Random.nextInt(jitterMs.toInt()) else 0)
delay(jitteredDelay)
throw CancellationException("Retrying streaming request", error)
throw RetryStreamingRequestException(error)
} else {
logger.error { "Streaming failed: $error" }
throw error
@ -101,22 +128,20 @@ class RetryingPromptExecutor(
}
val retryLoop = flow {
var currentAttempt = attempt
var currentAttempt = 1
while (currentAttempt <= retryPolicy.maxAttempts) {
try {
createStream(currentAttempt).collect { emit(it) }
break
} catch (ce: CancellationException) {
if (ce.message == "Retrying streaming request" && currentAttempt < retryPolicy.maxAttempts) {
} catch (_: RetryStreamingRequestException) {
if (currentAttempt < retryPolicy.maxAttempts) {
currentAttempt++
continue
} else {
throw ce
throw IllegalStateException("Retry attempt overflow in streaming executor")
}
} catch (_: SerializationException) {
refinedPrompt = prompt.copy(messages = prompt.messages.dropLast(1))
currentAttempt++
continue
} catch (ce: CancellationException) {
throw ce
}
}
}
@ -127,7 +152,7 @@ class RetryingPromptExecutor(
override suspend fun moderate(prompt: Prompt, model: LLModel): ModerationResult =
delegate.moderate(prompt, model)
override suspend fun models(): List<String> = delegate.models()
override suspend fun models(): List<LLModel> = delegate.models()
data class RetryPolicy(
val maxAttempts: Int,
@ -142,25 +167,20 @@ class RetryingPromptExecutor(
model: LLModel,
tools: List<ToolDescriptor>
): List<Message.Response> {
val promptWithReasoning = prompt.withReasoningParams(model)
var attempt = 1
var delay = retryPolicy.initialDelay
var lastError: Throwable? = null
while (attempt <= retryPolicy.maxAttempts) {
try {
if (attempt > 1) {
events?.onRetry(
attempt,
retryPolicy.maxAttempts,
lastError?.javaClass?.simpleName
)
}
return delegate.execute(prompt, model, tools)
return delegate.execute(promptWithReasoning, model, tools)
} catch (t: Throwable) {
lastError = t
if (attempt >= retryPolicy.maxAttempts) throw t
val retryable = isRetryableFailure(t)
if (!retryable || attempt >= retryPolicy.maxAttempts) throw t
events?.onRetry(attempt + 1, retryPolicy.maxAttempts, t.javaClass.simpleName)
events?.onRetry(attempt + 1, retryPolicy.maxAttempts)
val jitter = (delay.inWholeMilliseconds * retryPolicy.jitterFactor).toLong()
val jitteredMs =
@ -178,4 +198,66 @@ class RetryingPromptExecutor(
override fun close() {
(delegate as? AutoCloseable)?.close()
}
private fun Prompt.withReasoningParams(model: LLModel): Prompt {
val params = when (model.provider) {
LLMProvider.OpenAI -> params.withOpenAIReasoning()
LLMProvider.Anthropic -> params.withAnthropicReasoning()
else -> params
}
return withParams(params)
}
private fun LLMParams.withOpenAIReasoning(): LLMParams {
val base = when (this) {
is OpenAIResponsesParams -> this
else -> OpenAIResponsesParams(
temperature = temperature,
maxTokens = maxTokens,
numberOfChoices = numberOfChoices,
speculation = speculation,
schema = schema,
toolChoice = toolChoice,
user = user,
additionalProperties = additionalProperties
)
}
val reasoning = base.reasoning ?: ReasoningConfig(
effort = ReasoningEffort.MEDIUM,
summary = ReasoningSummary.AUTO
)
return base.copy(reasoning = reasoning)
}
private fun LLMParams.withAnthropicReasoning(): LLMParams {
val base = when (this) {
is AnthropicParams -> this
else -> AnthropicParams(
temperature = temperature,
maxTokens = maxTokens,
numberOfChoices = numberOfChoices,
speculation = speculation,
schema = schema,
toolChoice = toolChoice,
user = user,
additionalProperties = additionalProperties
)
}
if (base.thinking != null) return base
val thinkingBudget = resolveAnthropicThinkingBudget(base.maxTokens) ?: return base
return base.copy(thinking = AnthropicThinking.Enabled(budgetTokens = thinkingBudget))
}
private fun resolveAnthropicThinkingBudget(maxTokens: Int?): Int? {
val limit = maxTokens ?: ANTHROPIC_DEFAULT_THINKING_BUDGET
if (limit <= ANTHROPIC_MIN_THINKING_BUDGET) {
return null
}
return (limit / 2)
.coerceAtLeast(ANTHROPIC_MIN_THINKING_BUDGET)
.coerceAtMost(ANTHROPIC_DEFAULT_THINKING_BUDGET)
}
}

13
src/main/kotlin/ee/carlrobert/codegpt/agent/history/AgentCheckpointConversationMapper.kt
View file

@ -1,11 +1,11 @@
package ee.carlrobert.codegpt.agent.history
import ai.koog.agents.snapshot.feature.AgentCheckpointData
import ee.carlrobert.codegpt.completions.ChatToolCall
import ee.carlrobert.codegpt.completions.ChatToolFunction
import ee.carlrobert.codegpt.conversations.Conversation
import ee.carlrobert.codegpt.conversations.message.Message
import ee.carlrobert.codegpt.util.StringUtil.stripThinkingBlocks
import ee.carlrobert.llm.client.openai.completion.response.ToolCall
import ee.carlrobert.llm.client.openai.completion.response.ToolFunctionResponse
object AgentCheckpointConversationMapper {
@ -16,12 +16,13 @@ object AgentCheckpointConversationMapper {
val conversation = Conversation()
val turns = AgentCheckpointTurnSequencer.toVisibleTurns(
history = checkpoint.messageHistory,
projectInstructions = projectInstructions
projectInstructions = projectInstructions,
preserveSyntheticContinuation = true
)
turns.forEach { turn ->
val response = StringBuilder()
val toolCalls = mutableListOf<ToolCall>()
val toolCalls = mutableListOf<ChatToolCall>()
val toolResults = LinkedHashMap<String, String>()
var syntheticToolIdIndex = 0
@ -39,11 +40,11 @@ object AgentCheckpointConversationMapper {
val callId = event.id?.takeIf { it.isNotBlank() }
?: "tool-call-${++syntheticToolIdIndex}"
toolCalls.add(
ToolCall(
ChatToolCall(
null,
callId,
"function",
ToolFunctionResponse(event.tool, event.content.trim())
ChatToolFunction(event.tool, event.content.trim())
)
)
}

8
src/main/kotlin/ee/carlrobert/codegpt/agent/history/AgentCheckpointTurnSequencer.kt
View file

@ -75,7 +75,13 @@ object AgentCheckpointTurnSequencer {
when (message) {
is PromptMessage.User -> {
if (isSyntheticTimelineUserMessage(message)) {
if (preserveSyntheticContinuation && currentPrompt != null) {
if (
preserveSyntheticContinuation &&
currentPrompt != null &&
currentEvents.any {
it is TurnEvent.ToolCall || it is TurnEvent.ToolResult
}
) {
// Keep collecting events for the currently active visible turn.
// Synthetic todo prompts are injected internally mid-run.
return@forEachIndexed

2
src/main/kotlin/ee/carlrobert/codegpt/agent/history/AgentHistoryThreadSummary.kt
View file

@ -1,6 +1,6 @@
package ee.carlrobert.codegpt.agent.history
import kotlinx.datetime.Instant
import kotlin.time.Instant
enum class ThreadStatus {
COMPLETED,

85
src/main/kotlin/ee/carlrobert/codegpt/agent/strategy/SingleRunStrategyProvider.kt
View file

@ -13,6 +13,7 @@ import ai.koog.agents.core.environment.result
import ai.koog.agents.features.tokenizer.feature.tokenizer
import ai.koog.prompt.dsl.Prompt
import ai.koog.prompt.executor.model.PromptExecutor
import ai.koog.prompt.llm.LLMProvider
import ai.koog.prompt.message.Message
import ai.koog.prompt.streaming.StreamFrame
import ai.koog.prompt.streaming.toMessageResponses
@ -203,7 +204,7 @@ private suspend fun AIAgentLLMWriteSession.requestResponses(
val preparedPrompt = config.missingToolsConversionStrategy.convertPrompt(prompt, tools)
executor.execute(preparedPrompt, model, tools)
}
val appendableResponses = appendableResponses(responses)
val appendableResponses = appendableResponses(responses, model.provider)
appendableResponses.forEach(appendResponse)
return if (stream) responses else appendableResponses
}
@ -222,10 +223,56 @@ private suspend fun AIAgentLLMWriteSession.requestAndPublish(
return response
}
private fun appendableResponses(responses: List<Message.Response>): List<Message.Response> {
return responses
internal fun appendableResponses(
responses: List<Message.Response>,
provider: LLMProvider
): List<Message.Response> {
val sortedResponses = responses
.sortedBy { if (it is Message.Assistant) 0 else 1 }
.filter { it !is Message.Reasoning }
.filterNot { response ->
response is Message.Assistant &&
response.hasOnlyTextContent() &&
response.content.isBlank()
}
if (provider == LLMProvider.Google) {
return if (sortedResponses.any { it is Message.Tool.Call }) {
sortedResponses.filterNot { it is Message.Assistant }
} else {
sortedResponses
}
}
val nonReasoningResponses = sortedResponses.filterNot { it is Message.Reasoning }
if (nonReasoningResponses.isNotEmpty()) {
return nonReasoningResponses
}
// Keep reasoning when it is the only non-empty model output.
val reasoningFallback = sortedResponses
.filterIsInstance<Message.Reasoning>()
.filter { it.content.isNotBlank() }
return if (reasoningFallback.isNotEmpty()) reasoningFallback else nonReasoningResponses
}
internal fun extractAssistantOutput(responses: List<Message.Response>): String {
val assistantText = responses
.filterIsInstance<Message.Assistant>()
.map { it.content.trim() }
.filter { it.isNotEmpty() }
.joinToString("\n")
.trim()
if (assistantText.isNotEmpty()) {
return assistantText
}
val reasoningText = responses
.filterIsInstance<Message.Reasoning>()
.map { it.content.trim() }
.filter { it.isNotEmpty() }
.joinToString("\n")
.trim()
return if (reasoningText.isNotEmpty()) "<think>$reasoningText</think>" else ""
}
private fun publishUsageAndCredits(
@ -236,11 +283,22 @@ private fun publishUsageAndCredits(
provider: ServiceType,
responses: List<Message.Response>,
) {
events.onTokenUsageAvailable(promptTokenCount(prompt, tokenizer).toLong())
events.onTokenUsageAvailable(promptTokenCount(prompt, tokenizer, responses).toLong())
publishCreditsIfAvailable(sessionId, provider, responses, events)
}
private fun promptTokenCount(prompt: Prompt, tokenizer: PromptTokenizer): Int {
private fun promptTokenCount(
prompt: Prompt,
tokenizer: PromptTokenizer,
responses: List<Message.Response>
): Int {
val reportedInputCount = responses.firstNotNullOfOrNull { response ->
response.metaInfo.inputTokensCount?.takeIf { it > 0 }
?: response.metaInfo.totalTokensCount?.takeIf { it > 0 }
}
if (reportedInputCount != null) {
return reportedInputCount
}
return if (prompt.latestTokenUsage == 0) {
tokenizer.tokenCountFor(prompt)
} else {
@ -305,14 +363,17 @@ private fun publishCreditsIfAvailable(
@EdgeTransformationDslMarker
private infix fun <IncomingOutput, OutgoingInput> AIAgentEdgeBuilderIntermediate<IncomingOutput, List<Message.Response>, OutgoingInput>.onSingleAssistantResponse(
block: suspend (Message.Response) -> Boolean
block: suspend (String) -> Boolean
): AIAgentEdgeBuilderIntermediate<IncomingOutput, String, OutgoingInput> {
return onIsInstance(List::class)
.transformed { response ->
response.filter { item -> item is Message.Response && item !is Message.Reasoning }
.filterIsInstance<Message.Response>()
val modelResponses = response.filterIsInstance<Message.Response>()
when {
modelResponses.any { it is Message.Tool.Call } -> null
else -> extractAssistantOutput(modelResponses)
}
}
.onCondition { it.size == 1 && it[0] is Message.Assistant }
.onCondition { messages -> block(messages[0]) }
.transformed { it[0].content }
.onCondition { it != null }
.transformed { it!! }
.onCondition { message -> block(message) }
}

2
src/main/kotlin/ee/carlrobert/codegpt/agent/tools/AskUserQuestionTool.kt
View file

@ -7,6 +7,7 @@ import kotlinx.serialization.Serializable
class AskUserQuestionTool(
workingDirectory: String,
sessionId: String,
hookManager: HookManager,
private val events: AgentEvents,
) : BaseTool<AskUserQuestionTool.Args, AskUserQuestionTool.Result>(
@ -29,6 +30,7 @@ class AskUserQuestionTool(
argsClass = Args::class,
resultClass = Result::class,
hookManager = hookManager,
sessionId = sessionId
) {
@Serializable

4
src/main/kotlin/ee/carlrobert/codegpt/agent/tools/BaseTool.kt
View file

@ -30,7 +30,7 @@ abstract class BaseTool<Args : Any, Result : Any>(
description: String,
protected val workingDirectory: String,
private val hookManager: HookManager,
private val sessionId: String? = null,
private val sessionId: String,
private val argsClass: KClass<Args>,
private val resultClass: KClass<Result>
) : Tool<Args, Result>(
@ -58,7 +58,7 @@ abstract class BaseTool<Args : Any, Result : Any>(
* @return The result of executing the tool
*/
final override suspend fun execute(args: Args): Result {
val toolId = sessionId?.let { ToolRunContext.getToolId(it) }
val toolId = sessionId.let { ToolRunContext.getToolId(it) }
var effectiveArgs = args
val beforePayload = mapOf(
"tool_name" to name,

3
src/main/kotlin/ee/carlrobert/codegpt/agent/tools/BashOutputTool.kt
View file

@ -14,7 +14,7 @@ import kotlinx.serialization.Serializable
class BashOutputTool(
workingDirectory: String,
hookManager: HookManager,
private val sessionId: String = "global",
private val sessionId: String,
) :
BaseTool<BashOutputTool.Args, BashOutputTool.Result>(
workingDirectory = workingDirectory,
@ -33,6 +33,7 @@ class BashOutputTool(
argsClass = Args::class,
resultClass = Result::class,
hookManager = hookManager,
sessionId = sessionId,
) {
@Serializable

3
src/main/kotlin/ee/carlrobert/codegpt/agent/tools/BashTool.kt
View file

@ -174,7 +174,8 @@ class BashTool(
""".trimIndent(),
argsClass = Args::class,
resultClass = Result::class,
hookManager = hookManager
hookManager = hookManager,
sessionId = sessionId,
) {
@Serializable

30
src/main/kotlin/ee/carlrobert/codegpt/agent/tools/ConfirmingTools.kt
View file

@ -83,33 +83,3 @@ class ConfirmingEditTool(
return delegate.execute(args)
}
}
class ConfirmingMcpTool(
private val delegate: McpTool,
private val approve: suspend (name: String, details: String) -> Boolean
) : Tool<McpTool.Args, McpTool.Result>(
argsSerializer = McpTool.Args.serializer(),
resultSerializer = McpTool.Result.serializer(),
name = delegate.name,
description = delegate.descriptor.description
) {
override suspend fun execute(args: McpTool.Args): McpTool.Result {
val details = buildMcpApprovalDetails(
serverId = args.serverId,
serverName = args.serverName,
toolName = args.toolName,
arguments = args.arguments
)
val ok = approve("MCP", details)
if (!ok) {
return McpTool.Result.error(
toolName = args.toolName,
output = "User rejected MCP tool execution",
serverId = args.serverId,
serverName = args.serverName
)
}
return delegate.execute(args)
}
}

5
src/main/kotlin/ee/carlrobert/codegpt/agent/tools/EditTool.kt
View file

@ -22,8 +22,8 @@ import java.io.File
class EditTool(
private val project: Project,
private val sessionId: String,
private val hookManager: HookManager,
private val sessionId: String? = null
) : BaseTool<EditTool.Args, EditTool.Result>(
workingDirectory = project.basePath ?: System.getProperty("user.dir"),
argsSerializer = Args.serializer(),
@ -52,7 +52,8 @@ class EditTool(
""".trimIndent(),
argsClass = Args::class,
resultClass = Result::class,
hookManager = hookManager
hookManager = hookManager,
sessionId = sessionId,
) {
private fun getLineAndColumn(document: Document, offset: Int): Pair<Int, Int> {

4
src/main/kotlin/ee/carlrobert/codegpt/agent/tools/GetLibraryDocsTool.kt
View file

@ -19,6 +19,7 @@ import java.nio.charset.StandardCharsets
*/
class GetLibraryDocsTool(
workingDirectory: String,
sessionId: String,
hookManager: HookManager
) : BaseTool<GetLibraryDocsTool.Args, GetLibraryDocsTool.Result>(
workingDirectory = workingDirectory,
@ -30,7 +31,8 @@ class GetLibraryDocsTool(
""".trimIndent(),
argsClass = Args::class,
resultClass = Result::class,
hookManager = hookManager
hookManager = hookManager,
sessionId = sessionId,
) {
@Serializable

4
src/main/kotlin/ee/carlrobert/codegpt/agent/tools/IntelliJSearchTool.kt
View file

@ -33,6 +33,7 @@ import java.nio.file.Paths
*/
class IntelliJSearchTool(
hookManager: HookManager,
sessionId: String,
private val project: Project,
) : BaseTool<IntelliJSearchTool.Args, IntelliJSearchTool.Result>(
workingDirectory = project.basePath ?: System.getProperty("user.dir"),
@ -53,7 +54,8 @@ class IntelliJSearchTool(
""".trimIndent(),
argsClass = Args::class,
resultClass = Result::class,
hookManager = hookManager
hookManager = hookManager,
sessionId = sessionId,
) {
companion object {

4
src/main/kotlin/ee/carlrobert/codegpt/agent/tools/KillShellTool.kt
View file

@ -10,6 +10,7 @@ import kotlinx.serialization.Serializable
class KillShellTool(
workingDirectory: String,
sessionId: String,
hookManager: HookManager,
) : BaseTool<KillShellTool.Args, KillShellTool.Result>(
workingDirectory = workingDirectory,
@ -25,7 +26,8 @@ class KillShellTool(
""".trimIndent(),
argsClass = Args::class,
resultClass = Result::class,
hookManager = hookManager
hookManager = hookManager,
sessionId = sessionId
) {
@Serializable

153
src/main/kotlin/ee/carlrobert/codegpt/agent/tools/McpDynamicToolRegistry.kt
View file

@ -2,15 +2,14 @@ package ee.carlrobert.codegpt.agent.tools
import ai.koog.agents.core.tools.Tool
import ai.koog.agents.core.tools.ToolDescriptor
import ai.koog.agents.core.tools.ToolParameterDescriptor
import ai.koog.agents.core.tools.ToolParameterType
import com.intellij.openapi.components.service
import com.intellij.openapi.diagnostic.thisLogger
import com.intellij.openapi.project.Project
import ee.carlrobert.codegpt.agent.AgentMcpContext
import ee.carlrobert.codegpt.mcp.ConnectionStatus
import ee.carlrobert.codegpt.mcp.McpToolAliasResolver
import ee.carlrobert.codegpt.mcp.McpSessionAttachment
import ee.carlrobert.codegpt.mcp.McpSessionManager
import ee.carlrobert.codegpt.mcp.ToolSchemaParser
import io.modelcontextprotocol.client.McpSyncClient
import io.modelcontextprotocol.spec.McpSchema
import kotlinx.serialization.json.JsonObject
@ -33,59 +32,23 @@ object McpDynamicToolRegistry {
if (attachments.isEmpty()) return emptyList()
val namesByTool = attachments
.flatMap { attachment -> attachment.availableTools.map { it.name.lowercase() } }
.groupingBy { it }
.eachCount()
val usedNames = mutableSetOf<String>()
val tools = mutableListOf<Tool<*, *>>()
attachments.forEach { attachment ->
attachment.availableTools.forEach { tool ->
val exposedName = uniqueToolName(
toolName = tool.name,
attachment = attachment,
hasCollisions = (namesByTool[tool.name.lowercase()] ?: 0) > 1,
usedNames = usedNames
)
tools.add(
SessionBoundMcpTool(
conversationId = conversationId,
serverId = attachment.serverId,
serverName = attachment.serverName,
sourceTool = tool,
exposedName = exposedName,
approve = approve
)
)
}
return McpToolAliasResolver.resolve(
items = attachments.flatMap { attachment ->
attachment.availableTools.map { tool -> attachment to tool }
},
toolName = { (_, tool) -> tool.name },
scopeName = { (attachment, _) -> attachment.serverName.ifBlank { attachment.serverId } }
).map { (entry, exposedName) ->
val (attachment, tool) = entry
SessionBoundMcpTool(
conversationId = conversationId,
serverId = attachment.serverId,
serverName = attachment.serverName,
sourceTool = tool,
exposedName = exposedName,
approve = approve
)
}
return tools
}
private fun uniqueToolName(
toolName: String,
attachment: McpSessionAttachment,
hasCollisions: Boolean,
usedNames: MutableSet<String>
): String {
val base = if (hasCollisions) {
"${toolName}_${normalizeName(attachment.serverName.ifBlank { attachment.serverId })}"
} else {
toolName
}
if (usedNames.add(base)) return base
var index = 2
while (true) {
val candidate = "${base}_$index"
if (usedNames.add(candidate)) return candidate
index++
}
}
private fun normalizeName(value: String): String {
val normalized = value.lowercase().replace(Regex("[^a-z0-9]+"), "_").trim('_')
return normalized.ifBlank { "server" }
}
private fun ensureConnectedAttachment(
@ -220,8 +183,8 @@ private fun buildDescriptor(
append(serverName)
append(')')
}
val allParameters = parseParameterDescriptors(sourceTool.schema)
val requiredNames = parseRequiredNames(sourceTool.schema).toSet()
val allParameters = ToolSchemaParser.parseParameterDescriptors(sourceTool.schema)
val requiredNames = ToolSchemaParser.parseRequiredNames(sourceTool.schema).toSet()
val required = allParameters.filter { it.name in requiredNames }
val optional = allParameters.filterNot { it.name in requiredNames }
return ToolDescriptor(
@ -231,79 +194,3 @@ private fun buildDescriptor(
optionalParameters = optional
)
}
private fun parseRequiredNames(schema: Map<String, Any?>): List<String> {
return (schema["required"] as? List<*>)
?.mapNotNull { it?.toString() }
?: emptyList()
}
private fun parseParameterDescriptors(schema: Map<String, Any?>): List<ToolParameterDescriptor> {
val properties = schema["properties"] as? Map<*, *> ?: emptyMap<Any, Any?>()
return properties.mapNotNull { (nameRaw, propertyRaw) ->
val name = nameRaw?.toString() ?: return@mapNotNull null
val property = propertyRaw as? Map<*, *> ?: emptyMap<Any, Any?>()
ToolParameterDescriptor(
name = name,
description = property["description"]?.toString().orEmpty(),
type = parseParameterType(property)
)
}
}
private fun parseParameterType(definition: Map<*, *>): ToolParameterType {
val enumValues = (definition["enum"] as? List<*>)?.mapNotNull { it?.toString() }.orEmpty()
if (enumValues.isNotEmpty()) {
return ToolParameterType.Enum(enumValues.toTypedArray())
}
val anyOf = definition["anyOf"] as? List<*>
if (!anyOf.isNullOrEmpty()) {
val mapped = anyOf.mapNotNull { it as? Map<*, *> }
if (mapped.size == 2 && mapped.any { it["type"]?.toString()?.lowercase() == "null" }) {
return parseParameterType(mapped.first {
it["type"]?.toString()?.lowercase() != "null"
})
}
val types = mapped.map { element ->
ToolParameterDescriptor(
name = "",
description = element["description"]?.toString().orEmpty(),
type = parseParameterType(element)
)
}.toTypedArray()
return ToolParameterType.AnyOf(types)
}
return when (definition["type"]?.toString()?.lowercase()) {
"string" -> ToolParameterType.String
"integer" -> ToolParameterType.Integer
"number" -> ToolParameterType.Float
"boolean" -> ToolParameterType.Boolean
"null" -> ToolParameterType.Null
"array" -> {
val itemDefinition = definition["items"] as? Map<*, *> ?: emptyMap<Any, Any?>()
ToolParameterType.List(parseParameterType(itemDefinition))
}
"object" -> {
val properties = parseParameterDescriptors(toStringKeyedMap(definition["properties"]))
val required = (definition["required"] as? List<*>)
?.mapNotNull { it?.toString() }
?: emptyList()
ToolParameterType.Object(
properties = properties,
requiredProperties = required
)
}
else -> ToolParameterType.String
}
}
private fun toStringKeyedMap(value: Any?): Map<String, Any?> {
val map = value as? Map<*, *> ?: return emptyMap()
return map.entries.mapNotNull { (key, nestedValue) ->
key?.toString()?.let { it to nestedValue }
}.toMap()
}

294
src/main/kotlin/ee/carlrobert/codegpt/agent/tools/McpTool.kt
View file

@ -1,59 +1,24 @@
package ee.carlrobert.codegpt.agent.tools
import ai.koog.agents.core.tools.annotations.LLMDescription
import com.intellij.openapi.components.service
import com.intellij.openapi.project.Project
import ee.carlrobert.codegpt.agent.AgentMcpContext
import ee.carlrobert.codegpt.agent.AgentMcpContextService
import ee.carlrobert.codegpt.mcp.ConnectionStatus
import ee.carlrobert.codegpt.mcp.McpSessionAttachment
import ee.carlrobert.codegpt.mcp.McpSessionManager
import ee.carlrobert.codegpt.settings.hooks.HookManager
import io.modelcontextprotocol.client.McpSyncClient
import io.modelcontextprotocol.spec.McpSchema
import kotlinx.serialization.SerialName
import kotlinx.serialization.Serializable
import kotlinx.serialization.json.JsonElement
import java.util.*
class McpTool(
private val project: Project,
private val sessionId: String,
context: AgentMcpContext,
private val sessionManager: McpSessionManager,
hookManager: HookManager
) : BaseTool<McpTool.Args, McpTool.Result>(
argsSerializer = Args.serializer(),
resultSerializer = Result.serializer(),
name = "MCP",
description = buildToolDescription(context, sessionManager),
workingDirectory = project.basePath ?: System.getProperty("user.dir"),
hookManager = hookManager,
sessionId = sessionId,
argsClass = Args::class,
resultClass = Result::class
) {
object McpTool {
@Serializable
data class Args(
@property:LLMDescription(
"Tool name to execute on the MCP server"
)
@property:LLMDescription("Tool name to execute on the MCP server")
@SerialName("tool_name")
val toolName: String,
@property:LLMDescription(
"Optional MCP server id when multiple servers expose the same tool"
)
@property:LLMDescription("Optional MCP server id when multiple servers expose the same tool")
@SerialName("server_id")
val serverId: String? = null,
@property:LLMDescription(
"Optional MCP server name alternative to server_id"
)
@property:LLMDescription("Optional MCP server name alternative to server_id")
@SerialName("server_name")
val serverName: String? = null,
@property:LLMDescription(
"JSON object with tool arguments"
)
@property:LLMDescription("JSON object with tool arguments")
val arguments: Map<String, JsonElement> = emptyMap()
)
@ -85,253 +50,4 @@ class McpTool(
}
}
}
override suspend fun doExecute(args: Args): Result {
val contextService = project.service<AgentMcpContextService>()
val context = contextService.get(sessionId)
?: return args.error("No MCP context found for current agent session")
val conversationId = context.conversationId
?: return args.error("No MCP conversation is bound to the current session")
if (!context.hasSelection()) {
return args.error("No MCP servers selected. Add MCP server tags in the input panel.")
}
val selectedAttachments = loadAttachments(conversationId)
.filter { it.serverId in context.selectedServerIds }
val target = resolveTargetAttachment(args, context, selectedAttachments)
?: return args.error("Unable to resolve target MCP server for tool '${args.toolName}'")
val attached = runCatching {
ensureAttached(conversationId, target.serverId)
}.getOrElse { error ->
return args.error(
output = "Failed to attach MCP server '${target.serverName}': " +
(error.message ?: "unknown error"),
serverId = target.serverId,
serverName = target.serverName
)
}
if (!attached.isConnected()) {
val error = attached.lastError ?: "MCP server is not connected"
return args.error(
output = "Failed to connect MCP server '${attached.serverName}': $error",
serverId = attached.serverId,
serverName = attached.serverName
)
}
val resolvedToolName = resolveToolName(attached, args.toolName)
?: return args.error(
output = "Tool '${args.toolName}' is not available on MCP server '${attached.serverName}'",
serverId = attached.serverId,
serverName = attached.serverName
)
val client = resolveClient(conversationId, attached.serverId)
?: return args.error(
output = "MCP server '${attached.serverName}' is not connected",
serverId = attached.serverId,
serverName = attached.serverName
)
return executeTool(client, attached, args, resolvedToolName)
}
override fun createDeniedResult(originalArgs: Args, deniedReason: String): Result {
return Result.error(
toolName = originalArgs.toolName,
output = deniedReason,
serverId = originalArgs.serverId,
serverName = originalArgs.serverName
)
}
override fun encodeResultToString(result: Result): String {
val status = if (result.success) "success" else "error"
val server = result.serverName ?: result.serverId ?: "unknown"
return buildString {
appendLine("Tool: ${result.toolName}")
appendLine("Server: $server")
appendLine("Status: $status")
appendLine("Output:")
append(result.output)
}
}
private fun loadAttachments(conversationId: UUID): List<McpSessionAttachment> {
return runCatching { sessionManager.getSessionAttachments(conversationId) }
.getOrDefault(emptyList())
}
private fun resolveTargetAttachment(
args: Args,
context: AgentMcpContext,
selectedAttachments: List<McpSessionAttachment>
): McpSessionAttachment? {
if (args.serverId.hasText()) {
return resolveByServerId(args.serverId, context.selectedServerIds, selectedAttachments)
}
if (args.serverName.hasText()) {
return resolveByServerName(args.serverName, selectedAttachments)
}
return resolveByToolOrFallback(
toolName = args.toolName,
selectedServerIds = context.selectedServerIds,
selectedAttachments = selectedAttachments
)
}
private fun resolveByServerId(
serverId: String?,
selectedServerIds: Set<String>,
selectedAttachments: List<McpSessionAttachment>
): McpSessionAttachment? {
val requestedId = serverId.normalizedOrNull() ?: return null
if (requestedId !in selectedServerIds) {
return null
}
return selectedAttachments.firstOrNull { it.serverId == requestedId }
?: disconnectedAttachment(requestedId)
}
private fun resolveByServerName(
serverName: String?,
selectedAttachments: List<McpSessionAttachment>
): McpSessionAttachment? {
val requestedName = serverName.normalizedOrNull() ?: return null
return selectedAttachments.firstOrNull {
it.serverName.equals(requestedName, ignoreCase = true)
}
}
private fun resolveByToolOrFallback(
toolName: String,
selectedServerIds: Set<String>,
selectedAttachments: List<McpSessionAttachment>
): McpSessionAttachment? {
if (selectedAttachments.isEmpty()) {
return if (selectedServerIds.size == 1) disconnectedAttachment(selectedServerIds.first())
else null
}
val byTool = selectedAttachments.filter { attachment ->
attachment.availableTools.any { it.name.equals(toolName, ignoreCase = true) }
}
return when {
byTool.size == 1 -> byTool.first()
selectedAttachments.size == 1 -> selectedAttachments.first()
else -> null
}
}
private fun disconnectedAttachment(serverId: String): McpSessionAttachment {
return McpSessionAttachment(
serverId = serverId,
serverName = serverId,
connectionStatus = ConnectionStatus.DISCONNECTED
)
}
private fun ensureAttached(conversationId: UUID, serverId: String): McpSessionAttachment {
val current = loadAttachments(conversationId).firstOrNull { it.serverId == serverId }
if (current != null && current.isConnected()) {
return current
}
return sessionManager.attachServerToSession(conversationId, serverId).get()
}
private fun resolveToolName(attachment: McpSessionAttachment, requestedTool: String): String? {
return attachment.availableTools
.firstOrNull { it.name.equals(requestedTool, ignoreCase = true) }
?.name
}
private fun resolveClient(conversationId: UUID, serverId: String): McpSyncClient? {
val clientKey = "$conversationId:$serverId"
return runCatching { sessionManager.ensureClientConnected(clientKey).get() }.getOrNull()
}
private fun executeTool(
client: McpSyncClient,
attachment: McpSessionAttachment,
args: Args,
resolvedToolName: String
): Result {
return runCatching {
val callArgs = args.arguments.toMcpArguments()
val request = McpSchema.CallToolRequest(resolvedToolName, callArgs)
val toolResult = client.callTool(request)
val content = toolResult.formatMcpContent()
if (toolResult.isError == true) {
Result.error(
toolName = resolvedToolName,
output = content,
serverId = attachment.serverId,
serverName = attachment.serverName
)
} else {
Result(
serverId = attachment.serverId,
serverName = attachment.serverName,
toolName = resolvedToolName,
success = true,
output = content
)
}
}.getOrElse { error ->
Result.error(
toolName = resolvedToolName,
output = error.message ?: "MCP tool execution failed",
serverId = attachment.serverId,
serverName = attachment.serverName
)
}
}
private fun Args.error(
output: String,
serverId: String? = null,
serverName: String? = null
): Result = Result.error(toolName, output, serverId, serverName)
}
private fun String?.hasText(): Boolean = this?.trim()?.isNotEmpty() == true
private fun String?.normalizedOrNull(): String? = this?.trim()?.takeIf { it.isNotEmpty() }
private fun buildToolDescription(
context: AgentMcpContext?,
sessionManager: McpSessionManager
): String {
val base = """
Execute a tool on an attached MCP server.
- Use tool_name and arguments for the target MCP tool call.
- Use server_id when multiple selected servers expose the same tool.
""".trimIndent()
val conversationId = context?.conversationId ?: return base
val selectedServers = context.selectedServerIds
if (selectedServers.isEmpty()) return base
val attachments = runCatching { sessionManager.getSessionAttachments(conversationId) }
.getOrDefault(emptyList())
.filter { it.serverId in selectedServers }
if (attachments.isEmpty()) return base
val availableLines = attachments.map { attachment ->
val tools = attachment.availableTools.joinToString(", ") { it.name }
if (tools.isBlank()) {
"- ${attachment.serverName} (${attachment.serverId}): no tools discovered yet"
} else {
"- ${attachment.serverName} (${attachment.serverId}): $tools"
}
}
return buildString {
appendLine(base)
appendLine()
appendLine("Selected MCP servers and tools:")
availableLines.forEach { appendLine(it) }
}.trim()
}

16
src/main/kotlin/ee/carlrobert/codegpt/agent/tools/ReadTool.kt
View file

@ -6,8 +6,8 @@ import com.intellij.openapi.components.service
import com.intellij.openapi.diagnostic.thisLogger
import com.intellij.openapi.fileEditor.FileDocumentManager
import com.intellij.openapi.fileTypes.FileTypeManager
import com.intellij.openapi.vfs.LocalFileSystem
import com.intellij.openapi.project.Project
import com.intellij.openapi.vfs.LocalFileSystem
import ee.carlrobert.codegpt.settings.ProxyAISettingsService
import ee.carlrobert.codegpt.settings.ToolPermissionPolicy
import ee.carlrobert.codegpt.settings.hooks.HookEventType
@ -17,18 +17,18 @@ import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.withContext
import kotlinx.serialization.SerialName
import kotlinx.serialization.Serializable
import java.nio.charset.StandardCharsets
import java.nio.file.Files
import java.nio.file.Path
import java.nio.file.Paths
import java.nio.charset.StandardCharsets
/**
* Reads a file from the local filesystem using IntelliJ's Document and VirtualFile APIs.
*/
class ReadTool(
private val project: Project,
private val sessionId: String,
private val hookManager: HookManager,
private val sessionId: String? = null
) : BaseTool<ReadTool.Args, ReadTool.Result>(
workingDirectory = project.basePath ?: System.getProperty("user.dir"),
argsSerializer = Args.serializer(),
@ -53,7 +53,8 @@ class ReadTool(
""".trimIndent(),
argsClass = Args::class,
resultClass = Result::class,
hookManager = hookManager
hookManager = hookManager,
sessionId = sessionId,
) {
companion object {
@ -143,8 +144,11 @@ class ReadTool(
val virtualFile = LocalFileSystem.getInstance().findFileByPath(fileSystemPath)
val fileType = when {
virtualFile != null -> FileTypeManager.getInstance().getFileTypeByFile(virtualFile)
else -> FileTypeManager.getInstance().getFileTypeByFileName(path.fileName.toString())
virtualFile != null -> FileTypeManager.getInstance()
.getFileTypeByFile(virtualFile)
else -> FileTypeManager.getInstance()
.getFileTypeByFileName(path.fileName.toString())
}
when {

4
src/main/kotlin/ee/carlrobert/codegpt/agent/tools/ResolveLibraryIdTool.kt
View file

@ -17,6 +17,7 @@ import java.nio.charset.StandardCharsets
*/
class ResolveLibraryIdTool(
workingDirectory: String,
sessionId: String,
hookManager: HookManager,
) : BaseTool<ResolveLibraryIdTool.Args, ResolveLibraryIdTool.Result>(
workingDirectory = workingDirectory,
@ -28,7 +29,8 @@ class ResolveLibraryIdTool(
""".trimIndent(),
argsClass = Args::class,
resultClass = Result::class,
hookManager = hookManager
hookManager = hookManager,
sessionId = sessionId,
) {
@Serializable

1
src/main/kotlin/ee/carlrobert/codegpt/agent/tools/TaskTool.kt
View file

@ -40,6 +40,7 @@ class TaskTool(
argsClass = Args::class,
resultClass = Result::class,
hookManager = hookManager,
sessionId = sessionId,
) {
@Serializable

3
src/main/kotlin/ee/carlrobert/codegpt/agent/tools/TodoWriteTool.kt
View file

@ -44,7 +44,8 @@ NOTE that you should not use this tool if there is only one trivial task to do.
""".trimIndent(),
argsClass = Args::class,
resultClass = String::class,
hookManager = hookManager
hookManager = hookManager,
sessionId = sessionId,
) {
private val agentTabTitlePublisher =

4
src/main/kotlin/ee/carlrobert/codegpt/agent/tools/WebFetchTool.kt
View file

@ -14,6 +14,7 @@ import java.net.URI
class WebFetchTool(
workingDirectory: String,
private val userAgent: String = "Mozilla/5.0 (compatible; ProxyAI/1.0; +https://tryproxy.io)",
sessionId: String,
hookManager: HookManager,
) : BaseTool<WebFetchTool.Args, WebFetchTool.Result>(
workingDirectory = workingDirectory,
@ -29,7 +30,8 @@ Use this tool when you already have a specific URL and need the page content as
""".trimIndent(),
argsClass = Args::class,
resultClass = Result::class,
hookManager = hookManager
hookManager = hookManager,
sessionId = sessionId,
) {
@Serializable

4
src/main/kotlin/ee/carlrobert/codegpt/agent/tools/WebSearchTool.kt
View file

@ -16,6 +16,7 @@ import java.time.format.DateTimeFormatter
class WebSearchTool(
workingDirectory: String,
private val userAgent: String = "Mozilla/5.0 (compatible; ProxyAI/1.0; +https://tryproxy.io)",
sessionId: String,
hookManager: HookManager,
) : BaseTool<WebSearchTool.Args, WebSearchTool.Result>(
workingDirectory = workingDirectory,
@ -58,7 +59,8 @@ IMPORTANT - Use the correct year in search queries:
""".trimIndent(),
argsClass = Args::class,
resultClass = Result::class,
hookManager = hookManager
hookManager = hookManager,
sessionId = sessionId,
) {
@Serializable

6
src/main/kotlin/ee/carlrobert/codegpt/agent/tools/WriteTool.kt
View file

@ -9,8 +9,8 @@ import com.intellij.openapi.fileEditor.FileDocumentManager
import com.intellij.openapi.project.Project
import com.intellij.openapi.util.text.StringUtil
import com.intellij.openapi.vfs.LocalFileSystem
import com.intellij.openapi.vfs.VirtualFileManager
import com.intellij.openapi.vfs.VfsUtil
import com.intellij.openapi.vfs.VirtualFileManager
import com.intellij.psi.PsiDocumentManager
import ee.carlrobert.codegpt.settings.ProxyAISettingsService
import ee.carlrobert.codegpt.settings.hooks.HookManager
@ -31,6 +31,7 @@ import java.nio.file.Paths
*/
class WriteTool(
private val project: Project,
sessionId: String,
hookManager: HookManager,
) : BaseTool<WriteTool.Args, WriteTool.Result>(
workingDirectory = project.basePath ?: System.getProperty("user.dir"),
@ -57,7 +58,8 @@ class WriteTool(
""".trimIndent(),
argsClass = Args::class,
resultClass = Result::class,
hookManager = hookManager
hookManager = hookManager,
sessionId = sessionId,
) {
@Serializable

21
src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionEventListener.kt
View file

@ -10,22 +10,21 @@ import com.intellij.openapi.diagnostic.thisLogger
import com.intellij.openapi.editor.Editor
import com.intellij.openapi.util.TextRange
import ee.carlrobert.codegpt.CodeGPTKeys
import ee.carlrobert.codegpt.completions.CompletionError
import ee.carlrobert.codegpt.completions.CompletionStreamEventListener
import ee.carlrobert.codegpt.nextedit.NextEditCoordinator
import ee.carlrobert.codegpt.settings.service.FeatureType
import ee.carlrobert.codegpt.settings.service.ModelSelectionService
import ee.carlrobert.codegpt.settings.models.ModelSettings
import ee.carlrobert.codegpt.settings.service.ServiceType
import ee.carlrobert.codegpt.settings.service.codegpt.CodeGPTServiceSettings
import ee.carlrobert.codegpt.ui.OverlayUtil.showNotification
import ee.carlrobert.llm.client.openai.completion.ErrorDetails
import ee.carlrobert.llm.completion.CompletionEventListener
import kotlinx.coroutines.channels.ProducerScope
import okhttp3.sse.EventSource
import java.util.concurrent.atomic.AtomicBoolean
class CodeCompletionEventListener(
private val editor: Editor,
private val channel: ProducerScope<InlineCompletionElement>,
) : CompletionEventListener<String> {
) : CompletionStreamEventListener {
companion object {
private val logger = thisLogger()
@ -43,7 +42,7 @@ class CodeCompletionEventListener(
setLoading(true)
}
override fun onMessage(message: String, eventSource: EventSource) {
override fun onMessage(message: String) {
if (cancelled.get()) {
return
}
@ -51,16 +50,16 @@ class CodeCompletionEventListener(
messageBuilder.append(message)
}
override fun onComplete(result: StringBuilder) {
override fun onComplete(messageBuilder: StringBuilder) {
try {
CodeGPTKeys.REMAINING_CODE_COMPLETION.set(editor, null)
CodeGPTKeys.REMAINING_NEXT_EDITS.set(editor, null)
if (cancelled.get() || result.isEmpty()) {
if (cancelled.get() || messageBuilder.isEmpty()) {
return
}
var finalResult = CodeCompletionFormatter(editor).format(result.toString())
val finalResult = CodeCompletionFormatter(editor).format(messageBuilder.toString())
cache?.setCache(prefix, suffix, finalResult)
runInEdt { channel.trySend(InlineCompletionGrayTextElement(finalResult)) }
} finally {
@ -73,9 +72,9 @@ class CodeCompletionEventListener(
handleCompleted()
}
override fun onError(error: ErrorDetails, ex: Throwable) {
override fun onError(error: CompletionError, ex: Throwable) {
val isCodeGPTService =
service<ModelSelectionService>().getServiceForFeature(FeatureType.CODE_COMPLETION) == ServiceType.PROXYAI
service<ModelSettings>().getServiceForFeature(FeatureType.CODE_COMPLETION) == ServiceType.PROXYAI
if (isCodeGPTService && "RATE_LIMIT_EXCEEDED" == error.code) {
service<CodeGPTServiceSettings>().state
.codeCompletionSettings

31
src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionPromptFactory.kt Normal file
View file

@ -0,0 +1,31 @@
package ee.carlrobert.codegpt.codecompletions
import ai.koog.prompt.dsl.Prompt
import ai.koog.prompt.dsl.prompt
import ai.koog.prompt.params.LLMParams
internal object CodeCompletionPromptFactory {
private const val MAX_TOKENS = 128
private const val SYSTEM_PROMPT = "You are a code completion assistant. " +
"Complete the code between the given prefix and suffix. " +
"Return only the missing code that should be inserted, without any formatting, explanations, or markdown."
fun createPrompt(infillRequest: InfillRequest): Prompt {
val params = LLMParams(
maxTokens = MAX_TOKENS,
temperature = 0.0
)
return prompt("code-completion-fim", params = params) {
system(SYSTEM_PROMPT)
user(buildFIMUserPrompt(infillRequest))
}
}
private fun buildFIMUserPrompt(infillRequest: InfillRequest): String {
return "<PREFIX>\n${infillRequest.prefix}\n</PREFIX>\n\n" +
"<SUFFIX>\n${infillRequest.suffix}\n</SUFFIX>\n\n" +
"Complete:"
}
}

281
src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionRequestFactory.kt
View file

@ -1,281 +0,0 @@
package ee.carlrobert.codegpt.codecompletions
import com.fasterxml.jackson.core.JsonProcessingException
import com.fasterxml.jackson.core.type.TypeReference
import com.fasterxml.jackson.databind.ObjectMapper
import com.intellij.openapi.components.service
import ee.carlrobert.codegpt.completions.llama.LlamaModel
import ee.carlrobert.codegpt.credentials.CredentialsStore.CredentialKey
import ee.carlrobert.codegpt.credentials.CredentialsStore.getCredential
import ee.carlrobert.codegpt.settings.Placeholder.*
import ee.carlrobert.codegpt.settings.service.FeatureType
import ee.carlrobert.codegpt.settings.service.ModelSelectionService
import ee.carlrobert.codegpt.settings.service.custom.CustomServicesSettings
import ee.carlrobert.codegpt.settings.service.llama.LlamaSettings
import ee.carlrobert.codegpt.settings.service.llama.LlamaSettingsState
import ee.carlrobert.codegpt.settings.service.ollama.OllamaSettings
import ee.carlrobert.llm.client.inception.request.InceptionFIMRequest
import ee.carlrobert.llm.client.llama.completion.LlamaCompletionRequest
import ee.carlrobert.llm.client.ollama.completion.request.OllamaCompletionRequest
import ee.carlrobert.llm.client.ollama.completion.request.OllamaParameters
import ee.carlrobert.llm.client.openai.completion.request.OpenAIChatCompletionRequest
import ee.carlrobert.llm.client.openai.completion.request.OpenAIChatCompletionStandardMessage
import ee.carlrobert.llm.client.openai.completion.request.OpenAITextCompletionRequest
import okhttp3.MediaType.Companion.toMediaType
import okhttp3.Request
import okhttp3.RequestBody.Companion.toRequestBody
import org.apache.commons.text.StringEscapeUtils
import java.nio.charset.StandardCharsets
object CodeCompletionRequestFactory {
private const val MAX_TOKENS = 128
@JvmStatic
fun buildOpenAIRequest(details: InfillRequest): OpenAITextCompletionRequest {
return OpenAITextCompletionRequest.Builder(details.prefix)
.setModel(
ModelSelectionService.getInstance().getModelForFeature(FeatureType.CODE_COMPLETION)
)
.setSuffix(details.suffix)
.setStream(true)
.setMaxTokens(MAX_TOKENS)
.setTemperature(0.0)
.setPresencePenalty(0.0)
.setStop(details.stopTokens.ifEmpty { null })
.build()
}
@JvmStatic
fun buildChatBasedFIMRequest(details: InfillRequest): OpenAIChatCompletionRequest {
val systemMessage = OpenAIChatCompletionStandardMessage(
"system",
"You are a code completion assistant. Complete the code between the given prefix and suffix. " +
"Return only the missing code that should be inserted, without any formatting, explanations, or markdown."
)
val userMessage = OpenAIChatCompletionStandardMessage(
"user",
"<PREFIX>\n${details.prefix}\n</PREFIX>\n\n<SUFFIX>\n${details.suffix}\n</SUFFIX>\n\nComplete:"
)
return OpenAIChatCompletionRequest.Builder(listOf(systemMessage, userMessage))
.setModel(
ModelSelectionService.getInstance().getModelForFeature(FeatureType.CODE_COMPLETION)
)
.setStream(true)
.setMaxTokens(MAX_TOKENS)
.setTemperature(0.0)
.build()
}
@JvmStatic
fun buildChatBasedFIMHttpRequest(
details: InfillRequest,
url: String,
headers: Map<String, String>,
body: Map<String, Any>,
credential: String?
): Request {
val requestBuilder = Request.Builder().url(url)
for (entry in headers.entries) {
var value = entry.value
if (credential != null && value.contains("\$CUSTOM_SERVICE_API_KEY")) {
value = value.replace("\$CUSTOM_SERVICE_API_KEY", credential)
}
requestBuilder.addHeader(entry.key, value)
}
// Create chat completion messages using the improved prompt template
val systemMessage = mapOf<String, String>(
"role" to "system",
"content" to "You are a code completion assistant. Complete the code between the given prefix and suffix. " +
"Return only the missing code that should be inserted, without any formatting, explanations, or markdown."
)
val userMessage = mapOf<String, String>(
"role" to "user",
"content" to "<PREFIX>\n${details.prefix}\n</PREFIX>\n\n<SUFFIX>\n${details.suffix}\n</SUFFIX>\n\nComplete:"
)
// Transform the custom body configuration, excluding completion-specific parameters
val transformedBody = body.entries.mapNotNull { (key, value) ->
when (key.lowercase()) {
"messages" -> key to listOf(systemMessage, userMessage)
// Exclude completion-specific parameters that don't apply to chat completions
"prompt", "suffix" -> null
else -> key to transformValue(value, InfillPromptTemplate.CHAT_COMPLETION, details)
}
}.toMap().toMutableMap()
// Ensure we have messages for chat completion
if (!transformedBody.containsKey("messages")) {
transformedBody["messages"] = listOf(systemMessage, userMessage)
}
try {
val jsonBody = ObjectMapper()
.writerWithDefaultPrettyPrinter()
.writeValueAsString(transformedBody)
.toByteArray(StandardCharsets.UTF_8)
.toRequestBody("application/json".toMediaType())
return requestBuilder.post(jsonBody).build()
} catch (e: JsonProcessingException) {
throw RuntimeException(e)
}
}
@JvmStatic
fun buildCustomRequest(details: InfillRequest): Request {
val activeService = service<CustomServicesSettings>()
.customServiceStateForFeatureType(FeatureType.CODE_COMPLETION)
val settings = activeService.codeCompletionSettings
val credential = getCredential(CredentialKey.CustomServiceApiKeyById(requireNotNull(activeService.id)))
return buildCustomRequest(
details,
settings.url!!,
settings.headers,
settings.body,
settings.infillTemplate,
credential
)
}
@JvmStatic
fun buildCustomRequest(
details: InfillRequest,
url: String,
headers: Map<String, String>,
body: Map<String, Any>,
infillTemplate: InfillPromptTemplate,
credential: String?
): Request {
// For chat-based FIM, we should not use this method
// The routing logic in CodeCompletionService will handle it
if (infillTemplate == InfillPromptTemplate.CHAT_COMPLETION) {
throw IllegalArgumentException("Chat-based FIM should use buildChatBasedFIMRequest instead")
}
val requestBuilder = Request.Builder().url(url)
for (entry in headers.entries) {
var value = entry.value
if (credential != null && value.contains("\$CUSTOM_SERVICE_API_KEY")) {
value = value.replace("\$CUSTOM_SERVICE_API_KEY", credential)
}
requestBuilder.addHeader(entry.key, value)
}
val transformedBody = body.entries.associate { (key, value) ->
when (key.lowercase()) {
"stop" -> {
if (value is String) {
if (value.isEmpty())
null
if (value.startsWith("[") && value.endsWith("]"))
key to ObjectMapper().readValue(value, object : TypeReference<List<String>>() {})
else
key to value.split(",").map { StringEscapeUtils.unescapeJava(it.trim()) }
} else {
key to value
}
}
else -> key to transformValue(value, infillTemplate, details)
}
}
try {
val requestBody = ObjectMapper()
.writerWithDefaultPrettyPrinter()
.writeValueAsString(transformedBody)
.toByteArray(StandardCharsets.UTF_8)
.toRequestBody("application/json".toMediaType())
return requestBuilder.post(requestBody).build()
} catch (e: JsonProcessingException) {
throw RuntimeException(e)
}
}
@JvmStatic
fun buildLlamaRequest(details: InfillRequest): LlamaCompletionRequest {
val settings = LlamaSettings.getCurrentState()
val promptTemplate = getLlamaInfillPromptTemplate(settings)
val prompt = promptTemplate.buildPrompt(details)
val stopTokens = buildList {
if (promptTemplate.stopTokens != null) addAll(promptTemplate.stopTokens)
if (details.stopTokens.isNotEmpty()) addAll(details.stopTokens)
}.ifEmpty { null }
return LlamaCompletionRequest.Builder(prompt)
.setN_predict(MAX_TOKENS)
.setStream(true)
.setTemperature(0.0)
.setStop(stopTokens)
.build()
}
fun buildOllamaRequest(details: InfillRequest): OllamaCompletionRequest {
val settings = service<OllamaSettings>().state
val model = service<ModelSelectionService>().getModelForFeature(FeatureType.CODE_COMPLETION)
val stopTokens = buildList {
if (details.stopTokens.isNotEmpty()) addAll(details.stopTokens)
}.toMutableList()
val prompt = if (settings.fimOverride) {
settings.fimTemplate.stopTokens?.let { stopTokens.addAll(it) }
settings.fimTemplate.buildPrompt(details)
} else {
details.prefix
}
return OllamaCompletionRequest.Builder(
model,
prompt
)
.setSuffix(if (settings.fimOverride) null else details.suffix)
.setStream(true)
.setOptions(
OllamaParameters.Builder()
.stop(stopTokens.ifEmpty { null })
.numPredict(MAX_TOKENS)
.temperature(0.4)
.build()
)
.setRaw(true)
.build()
}
fun buildInceptionRequest(details: InfillRequest): InceptionFIMRequest {
val model = service<ModelSelectionService>().getModelForFeature(FeatureType.CODE_COMPLETION)
return InceptionFIMRequest.Builder()
.setPrompt(details.prefix)
.setSuffix(details.suffix)
.setModel(model)
.setStream(true)
.build()
}
private fun getLlamaInfillPromptTemplate(settings: LlamaSettingsState): InfillPromptTemplate {
if (settings.isUseCustomModel) {
return settings.localModelInfillPromptTemplate
}
return LlamaModel.findByHuggingFaceModel(settings.huggingFaceModel).infillPromptTemplate
}
private fun transformValue(
value: Any,
template: InfillPromptTemplate,
details: InfillRequest
): Any {
if (value !is String) return value
return when (value) {
FIM_PROMPT.code -> template.buildPrompt(details)
PREFIX.code -> details.prefix
SUFFIX.code -> details.suffix
else -> {
return value.takeIf { it.contains(PREFIX.code) || it.contains(SUFFIX.code) }
?.replace(PREFIX.code, details.prefix)
?.replace(SUFFIX.code, details.suffix) ?: value
}
}
}
}

175
src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionService.kt
View file

@ -2,41 +2,42 @@ package ee.carlrobert.codegpt.codecompletions
import com.intellij.openapi.components.Service
import com.intellij.openapi.components.service
import ee.carlrobert.codegpt.codecompletions.CodeCompletionRequestFactory.buildChatBasedFIMHttpRequest
import ee.carlrobert.codegpt.codecompletions.CodeCompletionRequestFactory.buildCustomRequest
import ee.carlrobert.codegpt.codecompletions.CodeCompletionRequestFactory.buildInceptionRequest
import ee.carlrobert.codegpt.codecompletions.CodeCompletionRequestFactory.buildLlamaRequest
import ee.carlrobert.codegpt.codecompletions.CodeCompletionRequestFactory.buildOllamaRequest
import ee.carlrobert.codegpt.codecompletions.CodeCompletionRequestFactory.buildOpenAIRequest
import ee.carlrobert.codegpt.completions.CompletionClientProvider
import ee.carlrobert.codegpt.credentials.CredentialsStore.CredentialKey
import ee.carlrobert.codegpt.credentials.CredentialsStore.getCredential
import ee.carlrobert.codegpt.agent.clients.CodeCompletionCapable
import ee.carlrobert.codegpt.agent.clients.CustomOpenAILLMClient
import ee.carlrobert.codegpt.agent.clients.toCompletionError
import ee.carlrobert.codegpt.codecompletions.edit.GrpcClientService
import ee.carlrobert.codegpt.completions.CancellableRequest
import ee.carlrobert.codegpt.completions.CompletionStreamEventListener
import ee.carlrobert.codegpt.settings.models.ModelSettings
import ee.carlrobert.codegpt.settings.models.LLMClientFactory
import ee.carlrobert.codegpt.settings.service.FeatureType
import ee.carlrobert.codegpt.settings.service.ModelSelectionService
import ee.carlrobert.codegpt.settings.service.ServiceType
import ee.carlrobert.codegpt.settings.service.ServiceType.*
import ee.carlrobert.codegpt.settings.service.codegpt.CodeGPTServiceSettings
import ee.carlrobert.codegpt.settings.service.custom.CustomServiceCodeCompletionSettingsState
import ee.carlrobert.codegpt.settings.service.custom.CustomServicesSettings
import ee.carlrobert.codegpt.settings.service.inception.InceptionSettings
import ee.carlrobert.codegpt.settings.service.llama.LlamaSettings
import ee.carlrobert.codegpt.settings.service.mistral.MistralSettings
import ee.carlrobert.codegpt.settings.service.ollama.OllamaSettings
import ee.carlrobert.codegpt.settings.service.openai.OpenAISettings
import ee.carlrobert.llm.client.openai.completion.OpenAIChatCompletionEventSourceListener
import ee.carlrobert.llm.client.openai.completion.OpenAITextCompletionEventSourceListener
import ee.carlrobert.llm.completion.CompletionEventListener
import okhttp3.sse.EventSource
import okhttp3.sse.EventSources.createFactory
import kotlinx.coroutines.CancellationException
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.SupervisorJob
import kotlinx.coroutines.launch
@Service
class CodeCompletionService {
fun getSelectedModelCode(): String? {
return ModelSelectionService.getInstance().getModelForFeature(FeatureType.CODE_COMPLETION)
private val scope = CoroutineScope(SupervisorJob() + Dispatchers.IO)
fun getSelectedModelCode(): String {
return ModelSettings.getInstance().getModelForFeature(FeatureType.CODE_COMPLETION)
}
fun isCodeCompletionsEnabled(): Boolean = isCodeCompletionsEnabled(
ModelSelectionService.getInstance().getServiceForFeature(FeatureType.CODE_COMPLETION)
ModelSettings.getInstance().getServiceForFeature(FeatureType.CODE_COMPLETION)
)
fun isCodeCompletionsEnabled(selectedService: ServiceType): Boolean =
@ -57,66 +58,86 @@ class CodeCompletionService {
fun getCodeCompletionAsync(
infillRequest: InfillRequest,
eventListener: CompletionEventListener<String>
): EventSource {
return when (val selectedService =
ModelSelectionService.getInstance().getServiceForFeature(FeatureType.CODE_COMPLETION)) {
OPENAI -> {
CompletionClientProvider.getOpenAIClient()
.getCompletionAsync(buildOpenAIRequest(infillRequest), eventListener)
}
eventListener: CompletionStreamEventListener
): CancellableRequest {
val selectedService = ModelSettings.getInstance().getServiceForFeature(FeatureType.CODE_COMPLETION)
return when (selectedService) {
OPENAI, CUSTOM_OPENAI, MISTRAL, OLLAMA, LLAMA_CPP, INCEPTION ->
executeCodeCompletion(selectedService, infillRequest, eventListener)
CUSTOM_OPENAI -> {
val activeService =
service<CustomServicesSettings>().customServiceStateForFeatureType(FeatureType.CODE_COMPLETION)
val customSettings = activeService.codeCompletionSettings
val isChatBasedFIM =
customSettings.infillTemplate == InfillPromptTemplate.CHAT_COMPLETION
if (isChatBasedFIM) {
val credential = getCredential(
CredentialKey.CustomServiceApiKeyById(
requireNotNull(activeService.id)
)
)
createFactory(
CompletionClientProvider.getDefaultClientBuilder().build()
).newEventSource(
buildChatBasedFIMHttpRequest(
infillRequest,
customSettings.url!!,
customSettings.headers,
customSettings.body,
credential
),
OpenAIChatCompletionEventSourceListener(eventListener)
)
} else {
createFactory(
CompletionClientProvider.getDefaultClientBuilder().build()
).newEventSource(
buildCustomRequest(infillRequest),
if (customSettings.parseResponseAsChatCompletions) {
OpenAIChatCompletionEventSourceListener(eventListener)
} else {
OpenAITextCompletionEventSourceListener(eventListener)
}
)
}
}
MISTRAL -> CompletionClientProvider.getMistralClient()
.getCodeCompletionAsync(buildOpenAIRequest(infillRequest), eventListener)
OLLAMA -> CompletionClientProvider.getOllamaClient()
.getCompletionAsync(buildOllamaRequest(infillRequest), eventListener)
LLAMA_CPP -> CompletionClientProvider.getLlamaClient()
.getCodeCompletionAsync(buildLlamaRequest(infillRequest), eventListener)
INCEPTION -> CompletionClientProvider.getInceptionClient()
.getFimCompletionAsync(buildInceptionRequest(infillRequest), eventListener)
else -> throw IllegalArgumentException("Code completion not supported for ${selectedService.name}")
PROXYAI -> executeProxyAICodeCompletion(infillRequest, eventListener)
ANTHROPIC, GOOGLE -> throw IllegalArgumentException("Code completion not supported for $selectedService")
}
}
fun testCustomOpenAIConnectionAsync(
settings: CustomServiceCodeCompletionSettingsState,
apiKey: String?,
eventListener: CompletionStreamEventListener
): CancellableRequest {
return executeWithClient(
clientProvider = {
CustomOpenAILLMClient.fromCodeCompletionSettingsState(
apiKey.orEmpty(),
settings
)
},
infillRequest = InfillRequest.Builder("Hello", "!", 0).build(),
eventListener = eventListener
)
}
private fun executeCodeCompletion(
serviceType: ServiceType,
infillRequest: InfillRequest,
eventListener: CompletionStreamEventListener
): CancellableRequest {
return executeWithClient(
clientProvider = {
val client = LLMClientFactory.createClient(serviceType, FeatureType.CODE_COMPLETION)
client as? CodeCompletionCapable
?: error("Code completion not supported by ${client::class.simpleName}")
},
infillRequest = infillRequest,
eventListener = eventListener
)
}
private fun executeProxyAICodeCompletion(
infillRequest: InfillRequest,
eventListener: CompletionStreamEventListener
): CancellableRequest {
val project = infillRequest.editor?.project
?: throw IllegalArgumentException("ProxyAI code completion requires an active editor project")
return project.service<GrpcClientService>()
.getCodeCompletionAsync(infillRequest, eventListener)
}
private fun executeWithClient(
clientProvider: () -> CodeCompletionCapable,
infillRequest: InfillRequest,
eventListener: CompletionStreamEventListener
): CancellableRequest {
val job = scope.launch {
val messageBuilder = StringBuilder()
val client = clientProvider()
try {
val completion = client.getCodeCompletion(infillRequest)
eventListener.onOpen()
if (completion.isNotEmpty()) {
messageBuilder.append(completion)
eventListener.onMessage(completion)
}
eventListener.onComplete(messageBuilder)
} catch (_: CancellationException) {
eventListener.onCancelled(messageBuilder)
} catch (ex: Throwable) {
eventListener.onError(ex.toCompletionError(), ex)
} finally {
runCatching { client.close() }
}
}
return CancellableRequest { job.cancel() }
}
}

21
src/main/kotlin/ee/carlrobert/codegpt/codecompletions/DebouncedCodeCompletionProvider.kt
View file

@ -6,9 +6,9 @@ import com.intellij.codeInsight.inline.completion.suggestion.InlineCompletionSin
import com.intellij.codeInsight.inline.completion.suggestion.InlineCompletionSuggestion
import com.intellij.openapi.components.service
import ee.carlrobert.codegpt.CodeGPTKeys.REMAINING_CODE_COMPLETION
import ee.carlrobert.codegpt.codecompletions.edit.GrpcClientService
import ee.carlrobert.codegpt.completions.CancellableRequest
import ee.carlrobert.codegpt.settings.service.FeatureType
import ee.carlrobert.codegpt.settings.service.ModelSelectionService
import ee.carlrobert.codegpt.settings.models.ModelSettings
import ee.carlrobert.codegpt.settings.service.ServiceType
import ee.carlrobert.codegpt.settings.service.codegpt.CodeGPTServiceSettings
import ee.carlrobert.codegpt.settings.service.custom.CustomServicesSettings
@ -20,7 +20,6 @@ import ee.carlrobert.codegpt.settings.service.openai.OpenAISettings
import kotlinx.coroutines.channels.awaitClose
import kotlinx.coroutines.flow.channelFlow
import kotlinx.coroutines.flow.emptyFlow
import okhttp3.sse.EventSource
import java.util.concurrent.atomic.AtomicReference
import kotlin.time.Duration
import kotlin.time.DurationUnit
@ -28,7 +27,7 @@ import kotlin.time.toDuration
class DebouncedCodeCompletionProvider : DebouncedInlineCompletionProvider() {
private val currentCallRef = AtomicReference<EventSource?>(null)
private val currentCallRef = AtomicReference<CancellableRequest?>(null)
override val id: InlineCompletionProviderID
get() = InlineCompletionProviderID("CodeGPTInlineCompletionProvider")
@ -70,19 +69,11 @@ class DebouncedCodeCompletionProvider : DebouncedInlineCompletionProvider() {
CompletionProgressNotifier.update(project, true)
var eventListener = CodeCompletionEventListener(request.editor, this)
val eventListener = CodeCompletionEventListener(request.editor, this)
val infillRequest = InfillRequestUtil.buildInfillRequest(request)
if (service<ModelSelectionService>().getServiceForFeature(FeatureType.CODE_COMPLETION) == ServiceType.PROXYAI) {
val grpcClient = project.service<GrpcClientService>()
grpcClient.cancelCodeCompletion()
grpcClient.getCodeCompletionAsync(infillRequest, eventListener, this)
return@channelFlow
}
val call = service<CodeCompletionService>().getCodeCompletionAsync(
infillRequest,
CodeCompletionEventListener(request.editor, this)
eventListener
)
currentCallRef.set(call)
} finally {
@ -110,7 +101,7 @@ class DebouncedCodeCompletionProvider : DebouncedInlineCompletionProvider() {
override fun isEnabled(event: InlineCompletionEvent): Boolean {
val selectedService =
service<ModelSelectionService>().getServiceForFeature(FeatureType.CODE_COMPLETION)
service<ModelSettings>().getServiceForFeature(FeatureType.CODE_COMPLETION)
val codeCompletionsEnabled = when (selectedService) {
ServiceType.PROXYAI -> service<CodeGPTServiceSettings>().state.codeCompletionSettings.codeCompletionsEnabled
ServiceType.OPENAI -> OpenAISettings.getCurrentState().isCodeCompletionsEnabled

8
src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequest.kt
View file

@ -62,7 +62,7 @@ class InfillRequest private constructor(
this.caretOffset = caretOffset
this.stopTokens = getStopTokens()
this.editor = editor
this.fileDetails = FileDetails(editor.document.text, editor.virtualFile.path)
this.fileDetails = FileDetails(editor.document.text, editor.virtualFile?.path ?: "")
}
fun gitDiff(gitDiff: String) =
@ -100,8 +100,8 @@ class InfillRequest private constructor(
prefix
}
return InfillRequest(
modifiedPrefix,
suffix,
modifiedPrefix.truncateText(MAX_PROMPT_TOKENS, false),
suffix.truncateText(MAX_PROMPT_TOKENS),
caretOffset,
editor,
fileDetails,
@ -132,4 +132,4 @@ class ContextElement(val psiElement: PsiElement) {
fun String.truncateText(maxTokens: Int, fromStart: Boolean = true): String {
return service<EncodingManager>().truncateText(this, maxTokens, fromStart)
}
}

31
src/main/kotlin/ee/carlrobert/codegpt/codecompletions/edit/CodeCompletionStreamObserver.kt
View file

@ -1,25 +1,20 @@
package ee.carlrobert.codegpt.codecompletions.edit
import com.intellij.codeInsight.inline.completion.elements.InlineCompletionElement
import com.intellij.notification.NotificationType
import com.intellij.openapi.diagnostic.thisLogger
import com.intellij.openapi.editor.Editor
import ee.carlrobert.codegpt.CodeGPTKeys
import ee.carlrobert.codegpt.codecompletions.CodeCompletionEventListener
import ee.carlrobert.codegpt.completions.CompletionError
import ee.carlrobert.codegpt.completions.CompletionStreamEventListener
import ee.carlrobert.codegpt.ui.OverlayUtil
import ee.carlrobert.llm.client.openai.completion.ErrorDetails
import ee.carlrobert.service.PartialCodeCompletionResponse
import io.grpc.Status
import io.grpc.StatusRuntimeException
import io.grpc.stub.StreamObserver
import kotlinx.coroutines.channels.ProducerScope
import okhttp3.Request
import okhttp3.sse.EventSource
class CodeCompletionStreamObserver(
private val editor: Editor,
private val channel: ProducerScope<InlineCompletionElement>,
private val eventListener: CodeCompletionEventListener,
private val eventListener: CompletionStreamEventListener,
) : StreamObserver<PartialCodeCompletionResponse> {
companion object {
@ -27,19 +22,10 @@ class CodeCompletionStreamObserver(
}
private val messageBuilder = StringBuilder()
private val emptyEventSource = object : EventSource {
override fun cancel() {
}
override fun request(): Request {
return Request.Builder().build()
}
}
override fun onNext(value: PartialCodeCompletionResponse) {
CodeGPTKeys.LAST_COMPLETION_RESPONSE_ID.set(editor, value.id)
messageBuilder.append(value.partialCompletion)
eventListener.onMessage(value.partialCompletion, emptyEventSource)
eventListener.onMessage(value.partialCompletion)
}
override fun onError(t: Throwable?) {
@ -51,8 +37,7 @@ class CodeCompletionStreamObserver(
}
if (code == Status.Code.UNAVAILABLE) {
eventListener.onError(ErrorDetails("Connection unavailable"), t)
channel.close(t)
eventListener.onError(CompletionError("Connection unavailable"), t)
return
}
}
@ -62,8 +47,10 @@ class CodeCompletionStreamObserver(
t?.message ?: "Something went wrong",
NotificationType.ERROR
)
eventListener.onError(ErrorDetails(t?.message ?: "Code completion error"), t ?: Throwable())
channel.close(t)
eventListener.onError(
CompletionError(t?.message ?: "Code completion error"),
t ?: Throwable()
)
}
override fun onCompleted() {

69
src/main/kotlin/ee/carlrobert/codegpt/codecompletions/edit/GrpcClientService.kt
View file

@ -1,6 +1,5 @@
package ee.carlrobert.codegpt.codecompletions.edit
import com.intellij.codeInsight.inline.completion.elements.InlineCompletionElement
import com.intellij.openapi.Disposable
import com.intellij.openapi.components.Service
import com.intellij.openapi.diagnostic.thisLogger
@ -8,12 +7,13 @@ import com.intellij.openapi.editor.Editor
import com.intellij.openapi.project.Project
import com.intellij.util.net.ssl.CertificateManager
import ee.carlrobert.codegpt.CodeGPTPlugin
import ee.carlrobert.codegpt.codecompletions.CodeCompletionEventListener
import ee.carlrobert.codegpt.codecompletions.InfillRequest
import ee.carlrobert.codegpt.completions.CancellableRequest
import ee.carlrobert.codegpt.completions.CompletionStreamEventListener
import ee.carlrobert.codegpt.credentials.CredentialsStore
import ee.carlrobert.codegpt.credentials.CredentialsStore.CredentialKey.CodeGptApiKey
import ee.carlrobert.codegpt.settings.service.FeatureType
import ee.carlrobert.codegpt.settings.service.ModelSelectionService
import ee.carlrobert.codegpt.settings.models.ModelSettings
import ee.carlrobert.codegpt.util.GitUtil
import ee.carlrobert.codegpt.util.RecentlyViewedFilesUtil
import ee.carlrobert.codegpt.util.file.FileUtil
@ -23,7 +23,6 @@ import io.grpc.ManagedChannel
import io.grpc.netty.shaded.io.grpc.netty.GrpcSslContexts
import io.grpc.netty.shaded.io.grpc.netty.NettyChannelBuilder
import io.grpc.netty.shaded.io.netty.channel.ChannelOption
import kotlinx.coroutines.channels.ProducerScope
import java.util.concurrent.TimeUnit
@Service(Service.Level.PROJECT)
@ -42,21 +41,24 @@ class GrpcClientService(private val project: Project) : Disposable {
private const val HOST = "grpc.tryproxy.io"
private const val PORT = 9090
private const val SHUTDOWN_TIMEOUT_SECONDS = 5L
private const val HOST_PROPERTY = "proxyai.grpc.host"
private const val PORT_PROPERTY = "proxyai.grpc.port"
private const val PLAINTEXT_PROPERTY = "proxyai.grpc.plaintext"
private val logger = thisLogger()
}
fun getCodeCompletionAsync(
request: InfillRequest,
eventListener: CodeCompletionEventListener,
channel: ProducerScope<InlineCompletionElement>
) {
eventListener: CompletionStreamEventListener
): CancellableRequest {
ensureCodeCompletionConnection()
val editor = request.editor ?: return
val editor = request.editor ?: return CancellableRequest {}
val grpcRequest = createCodeCompletionGrpcRequest(request)
eventListener.onOpen()
codeCompletionObserver =
CodeCompletionStreamObserver(editor, channel, eventListener)
CodeCompletionStreamObserver(editor, eventListener)
codeCompletionContext?.cancel(null)
val ctx = Context.current().withCancellation()
codeCompletionContext = ctx
@ -68,6 +70,7 @@ class GrpcClientService(private val project: Project) : Disposable {
} finally {
ctx.detach(prev)
}
return CancellableRequest { cancelCodeCompletion() }
}
@Synchronized
@ -191,7 +194,7 @@ class GrpcClientService(private val project: Project) : Disposable {
val gitDiff = request.gitDiff ?: ""
return GrpcCodeCompletionRequest.newBuilder()
.setModel(
ModelSelectionService.getInstance().getModelForFeature(FeatureType.CODE_COMPLETION)
ModelSettings.getInstance().getModelForFeature(FeatureType.CODE_COMPLETION)
)
.setFilePath(fileDetails.filePath)
.setFileContent(fileDetails.fileContent)
@ -210,7 +213,7 @@ class GrpcClientService(private val project: Project) : Disposable {
RecentlyViewedFilesUtil.orderedFiles(project, editor.virtualFile, 3)
.map { it.path to FileUtil.readContent(it) }
return NextEditRequest.newBuilder()
.setFileName(editor.virtualFile.name)
.setFileName(editor.virtualFile?.name ?: "")
.setFileContent(fileContent)
.setGitDiff(GitUtil.getCurrentChanges(project) ?: "")
.setCursorPosition(caretOffset)
@ -219,20 +222,36 @@ class GrpcClientService(private val project: Project) : Disposable {
.build()
}
private fun createChannel(): ManagedChannel = NettyChannelBuilder.forAddress(HOST, PORT)
.useTransportSecurity()
.sslContext(
GrpcSslContexts.forClient()
.trustManager(CertificateManager.getInstance().trustManager)
.build()
)
.withOption(ChannelOption.CONNECT_TIMEOUT_MILLIS, 5_000)
.keepAliveTime(2, TimeUnit.MINUTES)
.keepAliveTimeout(20, TimeUnit.SECONDS)
.keepAliveWithoutCalls(true)
.idleTimeout(30, TimeUnit.MINUTES)
.maxInboundMessageSize(32 * 1024 * 1024)
.build()
private fun createChannel(): ManagedChannel {
val builder = NettyChannelBuilder.forAddress(resolveHost(), resolvePort())
.withOption(ChannelOption.CONNECT_TIMEOUT_MILLIS, 5_000)
.keepAliveTime(2, TimeUnit.MINUTES)
.keepAliveTimeout(20, TimeUnit.SECONDS)
.keepAliveWithoutCalls(true)
.idleTimeout(30, TimeUnit.MINUTES)
.maxInboundMessageSize(32 * 1024 * 1024)
if (java.lang.Boolean.getBoolean(PLAINTEXT_PROPERTY)) {
builder.usePlaintext()
} else {
builder.useTransportSecurity()
builder.sslContext(
GrpcSslContexts.forClient()
.trustManager(CertificateManager.getInstance().trustManager)
.build()
)
}
return builder.build()
}
private fun resolveHost(): String {
return System.getProperty(HOST_PROPERTY) ?: HOST
}
private fun resolvePort(): Int {
return Integer.getInteger(PORT_PROPERTY, PORT)
}
private fun ensureActiveChannel() {
if (channel == null || channel?.isShutdown == true || channel?.isTerminated == true) {

1
src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/CompletionContextService.kt
View file

@ -15,7 +15,6 @@ class CompletionContextService {
companion object {
private val CONTEXT_FINDERS = mapOf(
"JAVA" to JavaContextFinder::class.java,
// "Python" to PythonContextFinder::class.java
)
}

2
src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/LanguageContextFinder.kt
View file

@ -8,4 +8,4 @@ interface LanguageContextFinder {
* Determines relevant enclosing [PsiElement] and [PsiElement]s relevant to the context and returns their source code [PsiElement].
*/
fun findContext(psiElement: PsiElement): InfillContext
}
}

143
src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/PythonContextFinder.kt
View file

@ -1,143 +0,0 @@
package ee.carlrobert.codegpt.codecompletions.psi
import com.intellij.openapi.project.Project
import com.intellij.psi.PsiElement
import com.intellij.psi.PsiRecursiveElementWalkingVisitor
import com.intellij.psi.util.PsiTreeUtil
import com.intellij.psi.util.findParentOfType
import com.jetbrains.python.psi.*
import com.jetbrains.python.psi.impl.PyBuiltinCache
import com.jetbrains.python.psi.resolve.PyResolveContext
import com.jetbrains.python.psi.resolve.PyResolveUtil
import com.jetbrains.python.psi.types.TypeEvalContext
import ee.carlrobert.codegpt.codecompletions.ContextElement
import ee.carlrobert.codegpt.codecompletions.InfillContext
class PythonContextFinder : LanguageContextFinder {
/**
* Finds enclosing [PyFunction] or [PyClass] of [psiElement] and
* determines source code elements of all used [PyReferenceExpression]s for the context.
*/
override fun findContext(psiElement: PsiElement): InfillContext {
val enclosingElement = findEnclosingElement(psiElement)
val referenceExpressions = findRelevantElements(enclosingElement, enclosingElement)
val declarations =
referenceExpressions.map { findDeclarations(it, psiElement.containingFile.project) }
.flatten().distinct()
.filter {
// Filter out elements whose source code is inside the enclosingElement
// e.g. for something like this: [i for i in range(10)] findRelevantElements()
// would return a "PyReferenceExpression: i" which is irrelevant
!it.containingFile.equals(enclosingElement.containingFile) || !enclosingElement.textRange.contains(
it.textRange
)
}
val sourceElements = declarations.mapNotNull { findSourceElement(it) }
return InfillContext(
ContextElement(enclosingElement),
sourceElements.map { ContextElement(it) }.toSet()
)
}
fun findEnclosingElement(psiElement: PsiElement): PsiElement = findEnclosingContext(psiElement)
?: PsiTreeUtil.prevCodeLeaf(psiElement)?.let { findEnclosingContext(it) } ?: psiElement
private fun findEnclosingContext(psiElement: PsiElement) =
PsiTreeUtil.findFirstContext(psiElement, true) { it is PyFunction || it is PyClass }
private fun findRelevantElements(
psiElement: Collection<PsiElement>,
rootElement: PsiElement
): Set<PyReferenceExpression> =
psiElement.map { findRelevantElements(it, rootElement) }.flatten().distinctBy { it.name }
.toSet()
/**
* Finds [PyReferenceExpression]s inside of [psiElement].
* If [psiElement] is a [PyFunction] inside of a [PyClass] it also adds all [PyReferenceExpression] of any class/instance fields.
*/
fun findRelevantElements(
psiElement: PsiElement,
rootElement: PsiElement
): Set<PyReferenceExpression> {
val resultSet = mutableSetOf<PyReferenceExpression>()
psiElement.accept(object : PsiRecursiveElementWalkingVisitor() {
override fun visitElement(element: PsiElement) {
when (element) {
is PyReferenceExpression -> resultSet.add(element)
is PyFunction -> {
// If the cursor was not inside a PyFunction but inside a PyClass, do not look into the functions
if (rootElement is PyClass) {
return
}
val enclosingContext = findEnclosingContext(element)
if (enclosingContext is PyClass) {
// add class and instance fields of enclosing class
// TODO: class fields declarations have to be present in the infillPrompt
// (same file as enclosingElement) as well
resultSet.addAll(
findRelevantElements(enclosingContext.classAttributes.mapNotNull {
findTargetExpressionAssignment(
it
)
}, rootElement)
)
resultSet.addAll(
findRelevantElements(
enclosingContext.instanceAttributes.mapNotNull {
findTargetExpressionAssignment(it)
},
rootElement
)
)
}
super.visitElement(element)
}
else -> super.visitElement(element)
}
}
})
return resultSet.distinctBy { it.name }.toSet()
}
private fun findTargetExpressionAssignment(targetExpression: PyTargetExpression): PyExpression? {
return targetExpression.findParentOfType<PyAssignmentStatement>()?.assignedValue
}
private fun findDeclarations(
pyReference: PyReferenceExpression,
project: Project
): Set<PsiElement> {
// see https://github.com/JetBrains/intellij-community/blob/ae5290861a1f41b93c48c475239a52faa94b97b0/python/python-psi-impl/src/com/jetbrains/python/codeInsight/PyTargetElementEvaluator.java#L51-L54
return PyResolveUtil.resolveDeclaration(
pyReference.reference,
PyResolveContext.defaultContext(
TypeEvalContext.codeAnalysis(
project,
pyReference.containingFile
)
)
)?.let {
if (PyBuiltinCache.getInstance(pyReference).isBuiltin(it) || it.filePath()
.contains("/stdlib/")
) {
null
} else {
setOf(it)
}
} ?: emptySet()
}
private fun findSourceElement(psiElement: PsiElement): PsiElement? {
val navigationElement = psiElement.navigationElement
val file = navigationElement.containingFile.virtualFile
return if (file.isInLocalFileSystem) {
navigationElement
} else {
null
}
}
}

229
src/main/kotlin/ee/carlrobert/codegpt/completions/AgentCompletionRunner.kt Normal file
View file

@ -0,0 +1,229 @@
package ee.carlrobert.codegpt.completions
import ai.koog.agents.core.agent.AIAgentService
import ai.koog.agents.core.agent.GraphAIAgentService
import ai.koog.agents.core.agent.config.AIAgentConfig
import ai.koog.agents.core.tools.ToolRegistry
import ai.koog.agents.features.eventHandler.feature.handleEvents
import ai.koog.agents.features.tokenizer.feature.MessageTokenizer
import ai.koog.prompt.dsl.prompt
import ai.koog.prompt.tokenizer.Tokenizer
import com.intellij.openapi.components.service
import com.intellij.openapi.project.Project
import ee.carlrobert.codegpt.EncodingManager
import ee.carlrobert.codegpt.agent.AgentEvents
import ee.carlrobert.codegpt.agent.MessageWithContext
import ee.carlrobert.codegpt.agent.clients.shouldStream
import ee.carlrobert.codegpt.agent.strategy.CODE_AGENT_COMPRESSION
import ee.carlrobert.codegpt.agent.strategy.HistoryCompressionConfig
import ee.carlrobert.codegpt.agent.strategy.SingleRunStrategyProvider
import ee.carlrobert.codegpt.mcp.McpTool
import ee.carlrobert.codegpt.mcp.McpToolAliasResolver
import ee.carlrobert.codegpt.mcp.McpToolCallHandler
import ee.carlrobert.codegpt.settings.models.ModelSettings
import ee.carlrobert.codegpt.settings.service.ServiceType
import ee.carlrobert.codegpt.settings.service.custom.CustomServicesSettings
import ee.carlrobert.codegpt.util.ReasoningFrameTextAdapter
import kotlinx.coroutines.*
import java.util.*
import java.util.concurrent.atomic.AtomicBoolean
import java.util.concurrent.atomic.AtomicReference
internal object AgentCompletionRunner : CompletionRunner {
private const val MAX_CHAT_AGENT_ITERATIONS = 32
override fun run(request: CompletionRunnerRequest): CancellableRequest {
val chatRequest = request as? CompletionRunnerRequest.Chat
?: throw IllegalArgumentException("AgentCompletionRunner can only run Chat requests")
return executeChat(chatRequest)
}
private fun executeChat(request: CompletionRunnerRequest.Chat): CancellableRequest {
val availableTools = request.callParameters.mcpTools
?.takeIf {
it.isNotEmpty() && request.callParameters.toolApprovalMode != ToolApprovalMode.BLOCK_ALL
}
?: emptyList()
val scope = CoroutineScope(SupervisorJob() + Dispatchers.IO)
val cancelled = AtomicBoolean(false)
val jobRef = AtomicReference<Job?>()
val messageBuilder = StringBuilder()
val project = request.callParameters.project!!
val stream = shouldStreamAgentToolLoop(project, request)
val toolCallHandler = project.let { McpToolCallHandler.getInstance(it) }
val toolRegistry = createChatToolRegistry(
callParameters = request.callParameters,
mcpTools = availableTools,
toolCallHandler = toolCallHandler,
onToolCallUIUpdate = request.onToolCallUIUpdate
)
val events = object : AgentEvents {
override fun onTextReceived(text: String) {
if (text.isNotEmpty()) {
messageBuilder.append(text)
request.eventListener.onMessage(text)
}
}
override fun onQueuedMessagesResolved() = Unit
}
val cancellableRequest = CancellableRequest {
cancelled.set(true)
val conversationId = request.callParameters.conversation.id
toolCallHandler.cancelPendingApprovals(conversationId)
toolCallHandler.cancelExecutions(conversationId)
jobRef.get()?.cancel(CancellationException("Cancelled by user"))
}
request.eventListener.onOpen()
val job = scope.launch {
var service: GraphAIAgentService<MessageWithContext, String>? = null
try {
service = AIAgentService<MessageWithContext, String>(
promptExecutor = request.executor,
strategy = SingleRunStrategyProvider().build(
project = project,
executor = request.executor,
pendingMessageQueue = ArrayDeque<MessageWithContext>(),
historyCompressionConfig = HistoryCompressionConfig(
isLimitExceeded = { _, _ -> false },
compressionStrategy = CODE_AGENT_COMPRESSION
),
events = events,
sessionId = request.callParameters.sessionId?.toString()
?: request.callParameters.conversation.id.toString(),
provider = request.serviceType,
stream = stream
),
agentConfig = AIAgentConfig(
prompt = prompt("chat-completion-agent") {
request.prompt.messages.forEach { message(it) }
},
model = request.model,
maxAgentIterations = MAX_CHAT_AGENT_ITERATIONS
),
toolRegistry = toolRegistry
) {
val frameAdapter = ReasoningFrameTextAdapter()
install(MessageTokenizer) {
tokenizer = object : Tokenizer {
override fun countTokens(text: String): Int {
return EncodingManager.getInstance().countTokens(text)
}
}
enableCaching = false
}
handleEvents {
onLLMStreamingFrameReceived { ctx ->
if (!stream) return@onLLMStreamingFrameReceived
frameAdapter.consume(ctx.streamFrame).forEach { chunk ->
if (chunk.isNotEmpty()) {
events.onTextReceived(chunk)
}
}
}
onNodeExecutionCompleted { ctx ->
if (stream) return@onNodeExecutionCompleted
(ctx.output as? List<*>)?.forEach { msg ->
(msg as? ai.koog.prompt.message.Message.Assistant)?.let {
events.onTextReceived(it.content)
}
(msg as? ai.koog.prompt.message.Message.Reasoning)?.let {
if (it.content.isNotBlank()) {
events.onTextReceived("<think>${it.content}</think>")
}
}
}
}
}
}
val agent = service.createAgent()
agent.run(MessageWithContext(""))
if (cancelled.get()) {
request.eventListener.onCancelled(StringBuilder(messageBuilder))
return@launch
}
request.eventListener.onComplete(StringBuilder(messageBuilder))
} catch (_: CancellationException) {
request.eventListener.onCancelled(StringBuilder(messageBuilder))
} catch (exception: Throwable) {
request.eventListener.onError(
ChatError(exception.message ?: "Failed to complete request"),
exception
)
} finally {
runCatching {
service?.closeAll()
request.executor.close()
toolCallHandler.clearConversationState(request.callParameters.conversation.id)
scope.cancel()
}
}
}
jobRef.set(job)
return cancellableRequest
}
internal fun shouldStreamAgentToolLoop(
project: Project,
request: CompletionRunnerRequest.Chat,
): Boolean {
val provider = request.serviceType
return when (provider) {
ServiceType.CUSTOM_OPENAI -> {
val selectedServiceId = project.service<ModelSettings>()
.getStoredModelForFeature(request.callParameters.featureType)
project.service<CustomServicesSettings>().state.services
.firstOrNull { it.id == selectedServiceId }
?.chatCompletionSettings
?.shouldStream()
?: false
}
ServiceType.GOOGLE -> false
else -> true
}
}
private fun createChatToolRegistry(
callParameters: ChatCompletionParameters,
mcpTools: List<McpTool>,
toolCallHandler: McpToolCallHandler?,
onToolCallUIUpdate: (javax.swing.JPanel) -> Unit
): ToolRegistry {
if (mcpTools.isEmpty() || toolCallHandler == null) {
return ToolRegistry {}
}
val aliases = McpToolAliasResolver.resolve(
items = mcpTools,
toolName = { it.name },
scopeName = { it.serverId }
)
return ToolRegistry {
aliases.forEach { (tool, alias) ->
tool(
McpToolAdapter(
tool = tool,
exposedName = alias,
callParameters = callParameters,
toolCallHandler = toolCallHandler,
onToolCallUIUpdate = onToolCallUIUpdate
)
)
}
}
}
}

95
src/main/kotlin/ee/carlrobert/codegpt/completions/AutoApplyService.kt Normal file
View file

@ -0,0 +1,95 @@
package ee.carlrobert.codegpt.completions
import com.intellij.openapi.components.Service
import com.intellij.openapi.components.service
import ee.carlrobert.codegpt.agent.clients.HttpClientProvider
import ee.carlrobert.codegpt.agent.clients.InceptionAILLMClient
import ee.carlrobert.codegpt.agent.clients.ProxyAILLMClient
import ee.carlrobert.codegpt.credentials.CredentialsStore.CredentialKey
import ee.carlrobert.codegpt.credentials.CredentialsStore.getCredential
import ee.carlrobert.codegpt.settings.models.ModelSettings
import ee.carlrobert.codegpt.settings.service.FeatureType
import ee.carlrobert.codegpt.settings.service.ServiceType.INCEPTION
import ee.carlrobert.codegpt.settings.service.ServiceType.PROXYAI
import ee.carlrobert.codegpt.completions.autoapply.AutoApplyRequest
import ee.carlrobert.codegpt.completions.inception.InceptionApplyRequest
import ee.carlrobert.codegpt.completions.inception.InceptionChatMessage
import java.util.regex.Pattern
@Service
class AutoApplyService {
companion object {
private val UPDATED_CODE_PATTERN =
Pattern.compile("<\\|updated_code\\|>(.*?)<\\|/updated_code\\|>", Pattern.DOTALL)
}
fun applyCode(params: AutoApplyParameters, originalCode: String): String {
val modelSelection =
service<ModelSettings>().getModelSelectionForFeature(FeatureType.AUTO_APPLY)
return when (modelSelection.provider) {
PROXYAI -> applyCodeWithProxyAI(modelSelection.model, originalCode, params.source)
INCEPTION -> applyCodeWithInception(modelSelection.model, originalCode, params.source)
else -> applyCodeWithDefault(params)
}
}
private fun applyCodeWithProxyAI(
model: String,
originalCode: String,
updateSnippet: String
): String {
val request = AutoApplyRequest(model, originalCode, updateSnippet)
return ProxyAILLMClient(getCredential(CredentialKey.CodeGptApiKey) ?: "")
.use { client -> client.getApplyEditCompletion(request).mergedCode }
}
private fun applyCodeWithInception(
model: String,
originalCode: String,
updateSnippet: String
): String {
val prompt = buildInceptionPrompt(originalCode, updateSnippet)
val request = InceptionApplyRequest.Builder()
.setModel(model)
.setMessages(listOf(InceptionChatMessage("user", prompt)))
.build()
return InceptionAILLMClient(
apiKey = getCredential(CredentialKey.InceptionApiKey) ?: "",
baseClient = HttpClientProvider.createHttpClient()
).use { client ->
client.getApplyEditCompletion(request)
.choices
?.firstOrNull()
?.message
?.content
.orEmpty()
}.let { extractUpdatedCode(it) }
}
private fun applyCodeWithDefault(params: AutoApplyParameters): String {
return extractUpdatedCode(service<CompletionService>().autoApply(params))
}
private fun buildInceptionPrompt(originalCode: String, updateSnippet: String): String {
return """
<|original_code|>
$originalCode
<|/original_code|>
<|update_snippet|>
$updateSnippet
<|/update_snippet|>
""".trimIndent()
}
private fun extractUpdatedCode(content: String): String {
val matcher = UPDATED_CODE_PATTERN.matcher(content)
return if (matcher.find()) {
matcher.group(1).trim()
} else {
content
}
}
}

23
src/main/kotlin/ee/carlrobert/codegpt/completions/BaseChatStreamEventListener.kt Normal file
View file

@ -0,0 +1,23 @@
package ee.carlrobert.codegpt.completions
abstract class BaseChatStreamEventListener(
protected val callParameters: ChatCompletionParameters,
protected val wrappedListener: CompletionResponseEventListener,
private val syncMessageResponse: Boolean = true
) : ChatStreamEventListener {
protected val messageBuilder = StringBuilder()
override fun onOpen() {
wrappedListener.handleRequestOpen()
}
override fun onMessage(message: String) {
messageBuilder.append(message)
if (syncMessageResponse) {
callParameters.message.response = messageBuilder.toString()
}
wrappedListener.handleMessage(message)
}
}

67
src/main/kotlin/ee/carlrobert/codegpt/completions/ChatCompletionEventListener.kt Normal file
View file

@ -0,0 +1,67 @@
package ee.carlrobert.codegpt.completions
import com.fasterxml.jackson.core.JsonProcessingException
import com.fasterxml.jackson.databind.ObjectMapper
import com.intellij.openapi.diagnostic.Logger
import com.intellij.openapi.project.Project
import ee.carlrobert.codegpt.codecompletions.CompletionProgressNotifier
import ee.carlrobert.codegpt.events.CodeGPTEvent
import ee.carlrobert.codegpt.telemetry.TelemetryAction
class ChatCompletionEventListener(
private val project: Project,
callParameters: ChatCompletionParameters,
wrappedListener: CompletionResponseEventListener
) : BaseChatStreamEventListener(callParameters, wrappedListener, syncMessageResponse = true) {
companion object {
private val LOG = Logger.getInstance(ChatCompletionEventListener::class.java)
}
private val objectMapper = ObjectMapper()
override fun onEvent(data: String) {
try {
val event = objectMapper.readValue(data, CodeGPTEvent::class.java)
wrappedListener.handleCodeGPTEvent(event)
} catch (exception: JsonProcessingException) {
LOG.debug("Failed to parse CodeGPTEvent from data: $data", exception)
}
}
override fun onComplete(messageBuilder: StringBuilder) {
handleCompleted(messageBuilder)
}
override fun onCancelled(messageBuilder: StringBuilder) {
handleCompleted(messageBuilder)
}
override fun onError(error: ChatError, ex: Throwable) {
try {
callParameters.conversation.addMessage(callParameters.message)
wrappedListener.handleError(error, ex)
} finally {
sendError(error, ex)
}
}
private fun handleCompleted(messageBuilder: StringBuilder) {
CompletionProgressNotifier.update(project, false)
wrappedListener.handleCompleted(messageBuilder.toString(), callParameters)
}
private fun sendError(error: ChatError, ex: Throwable) {
val telemetryMessage = TelemetryAction.COMPLETION_ERROR.createActionMessage()
if (error.code == "insufficient_quota") {
telemetryMessage
.property("type", "USER")
.property("code", "INSUFFICIENT_QUOTA")
} else {
telemetryMessage
.property("conversationId", callParameters.conversation.id.toString())
.error(RuntimeException(error.toString(), ex))
}
telemetryMessage.send()
}
}

85
src/main/kotlin/ee/carlrobert/codegpt/completions/ChatStreamingCompletionRunner.kt Normal file
View file

@ -0,0 +1,85 @@
package ee.carlrobert.codegpt.completions
import ai.koog.prompt.streaming.StreamFrame
import ee.carlrobert.codegpt.util.ReasoningFrameTextAdapter
import kotlinx.coroutines.CancellationException
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.Job
import kotlinx.coroutines.SupervisorJob
import kotlinx.coroutines.cancel
import kotlinx.coroutines.launch
import java.util.concurrent.atomic.AtomicBoolean
import java.util.concurrent.atomic.AtomicReference
internal object ChatStreamingCompletionRunner : CompletionRunner {
override fun run(request: CompletionRunnerRequest): CancellableRequest {
val chatRequest = request as? CompletionRunnerRequest.Chat
?: throw IllegalArgumentException("ChatStreamingCompletionRunner can only run Chat requests")
return executeAsync(chatRequest)
}
private fun executeAsync(request: CompletionRunnerRequest.Chat): CancellableRequest {
val scope = CoroutineScope(SupervisorJob() + Dispatchers.IO)
val cancelled = AtomicBoolean(false)
val jobRef = AtomicReference<Job?>()
val cancellableRequest = CancellableRequest {
cancelled.set(true)
jobRef.get()?.cancel(CancellationException("Cancelled by user"))
}
request.eventListener.onOpen()
val job = scope.launch {
val messageBuilder = StringBuilder()
try {
streamOrFallback(request, messageBuilder)
if (cancelled.get()) {
request.eventListener.onCancelled(StringBuilder(messageBuilder))
return@launch
}
request.eventListener.onComplete(StringBuilder(messageBuilder))
} catch (_: CancellationException) {
request.eventListener.onCancelled(StringBuilder(messageBuilder))
} catch (exception: Throwable) {
request.eventListener.onError(
ChatError(exception.message ?: "Failed to complete request"),
exception
)
} finally {
runCatching { request.executor.close() }
scope.cancel()
}
}
jobRef.set(job)
return cancellableRequest
}
private suspend fun streamOrFallback(
request: CompletionRunnerRequest.Chat,
messageBuilder: StringBuilder
) {
val frameAdapter = ReasoningFrameTextAdapter()
runCatching {
request.executor.executeStreaming(request.prompt, request.model, emptyList())
.collect { frame ->
frameAdapter.consume(frame).forEach { chunk ->
if (chunk.isNotEmpty()) {
messageBuilder.append(chunk)
request.eventListener.onMessage(chunk)
}
}
}
}.getOrElse {
val responses = request.executor.execute(request.prompt, request.model, emptyList())
val text = CompletionTextExtractor.extract(responses)
if (text.isNotBlank()) {
messageBuilder.append(text)
request.eventListener.onMessage(text)
}
}
}
}

36
src/main/kotlin/ee/carlrobert/codegpt/completions/ChatTypes.kt Normal file
View file

@ -0,0 +1,36 @@
package ee.carlrobert.codegpt.completions
data class ChatError @JvmOverloads constructor(
val message: String,
val code: String? = null
)
data class ChatToolFunction(
val name: String? = null,
val arguments: String? = null
)
data class ChatToolCall(
val index: Int? = null,
val id: String = "",
val type: String = "function",
val function: ChatToolFunction = ChatToolFunction()
)
fun interface CancellableRequest {
fun cancel()
}
interface ChatStreamEventListener {
fun onOpen() {}
fun onEvent(data: String) {}
fun onMessage(message: String)
fun onComplete(messageBuilder: StringBuilder)
fun onCancelled(messageBuilder: StringBuilder)
fun onError(error: ChatError, ex: Throwable)
}

6
src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionError.kt Normal file
View file

@ -0,0 +1,6 @@
package ee.carlrobert.codegpt.completions
data class CompletionError @JvmOverloads constructor(
val message: String,
val code: String? = null
)

7
src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionParameters.kt
View file

@ -11,7 +11,6 @@ import ee.carlrobert.codegpt.settings.configuration.ChatMode
import ee.carlrobert.codegpt.settings.prompts.PersonaDetails
import ee.carlrobert.codegpt.settings.service.FeatureType
import ee.carlrobert.codegpt.util.file.FileUtil
import ee.carlrobert.llm.client.openai.completion.response.ToolCall
import java.nio.file.Files
import java.nio.file.Path
import java.util.*
@ -32,7 +31,7 @@ class ChatCompletionParameters private constructor(
var mcpAttachedServerIds: List<String>?,
var mcpTools: List<McpTool>?,
var toolApprovalMode: ToolApprovalMode,
var toolResults: List<Pair<ToolCall, String>>? = null,
var toolResults: List<Pair<ChatToolCall, String>>? = null,
var project: Project?,
var chatMode: ChatMode = ChatMode.ASK,
var featureType: FeatureType = FeatureType.CHAT,
@ -72,7 +71,7 @@ class ChatCompletionParameters private constructor(
private var mcpAttachedServerIds: List<String>? = null
private var mcpTools: List<McpTool>? = null
private var toolApprovalMode: ToolApprovalMode = ToolApprovalMode.AUTO_APPROVE
private var toolResults: List<Pair<ToolCall, String>>? = null
private var toolResults: List<Pair<ChatToolCall, String>>? = null
private var project: Project? = null
private var chatMode: ChatMode = ChatMode.ASK
private var featureType: FeatureType = FeatureType.CHAT
@ -111,7 +110,7 @@ class ChatCompletionParameters private constructor(
fun mcpTools(tools: List<McpTool>?) = apply { this.mcpTools = tools }
fun toolApprovalMode(mode: ToolApprovalMode) = apply { this.toolApprovalMode = mode }
fun toolResults(results: List<Pair<ToolCall, String>>?) =
fun toolResults(results: List<Pair<ChatToolCall, String>>?) =
apply { this.toolResults = results }
fun project(project: Project?) = apply { this.project = project }

3
src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionRequest.kt Normal file
View file

@ -0,0 +1,3 @@
package ee.carlrobert.codegpt.completions
interface CompletionRequest

287
src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionRequestFactory.kt
View file

@ -1,5 +1,10 @@
package ee.carlrobert.codegpt.completions
import ai.koog.prompt.dsl.Prompt
import ai.koog.prompt.dsl.PromptBuilder
import ai.koog.prompt.dsl.prompt
import ai.koog.prompt.message.AttachmentContent
import ai.koog.prompt.message.ContentPart
import com.intellij.openapi.components.service
import com.intellij.openapi.project.Project
import com.intellij.openapi.vfs.LocalFileSystem
@ -8,28 +13,28 @@ import ee.carlrobert.codegpt.completions.CompletionRequestFactory.Companion.RECE
import ee.carlrobert.codegpt.completions.factory.*
import ee.carlrobert.codegpt.conversations.message.Message
import ee.carlrobert.codegpt.nextedit.NextEditPromptUtil
import ee.carlrobert.codegpt.psistructure.ClassStructureSerializer
import ee.carlrobert.codegpt.settings.configuration.ChatMode
import ee.carlrobert.codegpt.settings.prompts.CoreActionsState
import ee.carlrobert.codegpt.settings.prompts.FilteredPromptsService
import ee.carlrobert.codegpt.settings.prompts.PersonaDetails
import ee.carlrobert.codegpt.settings.prompts.PromptsSettings
import ee.carlrobert.codegpt.settings.prompts.addProjectPath
import ee.carlrobert.codegpt.settings.service.FeatureType
import ee.carlrobert.codegpt.settings.service.ModelSelectionService
import ee.carlrobert.codegpt.settings.service.ServiceType
import ee.carlrobert.codegpt.ui.textarea.ConversationTagProcessor
import ee.carlrobert.codegpt.util.EditWindowFormatter.FormatResult
import ee.carlrobert.codegpt.util.EditorUtil
import ee.carlrobert.codegpt.util.GitUtil
import ee.carlrobert.codegpt.util.file.FileUtil
import ee.carlrobert.llm.completion.CompletionRequest
import ee.carlrobert.codegpt.mcp.McpToolPromptFormatter
interface CompletionRequestFactory {
fun createChatRequest(params: ChatCompletionParameters): CompletionRequest
fun createInlineEditRequest(params: InlineEditCompletionParameters): CompletionRequest
fun createInlineEditQuestionRequest(parameters: ChatCompletionParameters): CompletionRequest
fun createAutoApplyRequest(params: AutoApplyParameters): CompletionRequest
fun createCommitMessageRequest(params: CommitMessageCompletionParameters): CompletionRequest
fun createLookupRequest(params: LookupCompletionParameters): CompletionRequest
fun createChatCompletionPrompt(callParameters: ChatCompletionParameters): Prompt
fun createInlineEditPrompt(params: InlineEditCompletionParameters): Prompt
fun createInlineEditQuestionPrompt(parameters: ChatCompletionParameters): Prompt
fun createAutoApplyPrompt(params: AutoApplyParameters): Prompt
fun createCommitMessagePrompt(params: CommitMessageCompletionParameters): Prompt
fun createLookupPrompt(params: LookupCompletionParameters): Prompt
fun createNextEditRequest(
params: NextEditParameters,
formatResult: FormatResult
@ -44,7 +49,7 @@ interface CompletionRequestFactory {
@JvmStatic
fun getFactory(serviceType: ServiceType): CompletionRequestFactory {
return when (serviceType) {
ServiceType.PROXYAI -> CodeGPTRequestFactory(ClassStructureSerializer)
ServiceType.PROXYAI -> CodeGPTRequestFactory()
ServiceType.OPENAI -> OpenAIRequestFactory()
ServiceType.CUSTOM_OPENAI -> CustomOpenAIRequestFactory()
ServiceType.ANTHROPIC -> ClaudeRequestFactory()
@ -55,24 +60,191 @@ interface CompletionRequestFactory {
ServiceType.INCEPTION -> InceptionRequestFactory()
}
}
@JvmStatic
fun getFactoryForFeature(featureType: FeatureType): CompletionRequestFactory {
val serviceType = ModelSelectionService.getInstance().getServiceForFeature(featureType)
return getFactory(serviceType)
}
}
}
abstract class BaseRequestFactory : CompletionRequestFactory {
companion object {
private const val LOOKUP_MAX_TOKENS = 512
private const val AUTO_APPLY_MAX_TOKENS = 8192
private const val DEFAULT_MAX_TOKENS = 4096
override fun createChatCompletionPrompt(callParameters: ChatCompletionParameters): Prompt {
val systemPrompt = buildKoogSystemPrompt(callParameters)
val currentMessage = callParameters.message
return prompt("chat-completion") {
if (systemPrompt.isNotBlank()) {
system(systemPrompt)
}
val currentMessageId = currentMessage.id
callParameters.conversation.messages.forEach { msg ->
val isCurrent = msg.id == currentMessageId
if (isCurrent && callParameters.retry) {
return@forEach
}
if (isCurrent) {
return@forEach
}
appendUserPrompt(msg.prompt)
appendAssistantAndToolState(msg)
}
when (callParameters.requestType) {
RequestType.NORMAL_REQUEST -> {
appendCurrentUserMessage(callParameters)
}
RequestType.TOOL_CALL_REQUEST -> {
appendCurrentUserMessage(callParameters)
appendAssistantAndToolState(currentMessage)
}
RequestType.TOOL_CALL_CONTINUATION -> {
appendCurrentUserMessage(callParameters)
appendAssistantAndToolState(currentMessage, includeToolResults = false)
currentMessage.toolCallResults?.forEach { (callId, result) ->
val toolName = currentMessage.toolCalls
?.firstOrNull { it.id == callId }
?.function
?.name
?: "unknown-function"
tool { result(callId, toolName, result) }
}
if (!callParameters.retry && !currentMessage.response.isNullOrBlank()) {
assistant(currentMessage.response)
}
}
}
}
}
override fun createInlineEditQuestionRequest(parameters: ChatCompletionParameters): CompletionRequest {
private fun PromptBuilder.appendCurrentUserMessage(
callParameters: ChatCompletionParameters
) {
val promptWithContext = getPromptWithFilesContext(callParameters)
val imageDetails = callParameters.imageDetails
if (imageDetails != null) {
val format = imageDetails.mediaType
.substringAfter('/')
.substringBefore(';')
.ifBlank { "png" }
val imagePart = ContentPart.Image(
content = AttachmentContent.Binary.Bytes(imageDetails.data),
format = format,
mimeType = imageDetails.mediaType
)
val parts = buildList {
add(imagePart)
if (promptWithContext.isNotBlank()) {
add(ContentPart.Text(promptWithContext))
}
}
user(parts)
return
}
appendUserPrompt(promptWithContext)
}
private fun PromptBuilder.appendUserPrompt(prompt: String?) {
if (!prompt.isNullOrBlank()) {
user(prompt)
}
}
private fun PromptBuilder.appendAssistantAndToolState(
message: Message,
includeToolResults: Boolean = true
) {
if (message.hasToolCalls()) {
if (!message.response.isNullOrBlank()) {
assistant(message.response)
}
message.toolCalls?.forEach { toolCall ->
tool {
call(
toolCall.id,
toolCall.function.name ?: "unknown-function",
toolCall.function.arguments
?.takeIf { it.isNotBlank() }
?: "{}"
)
}
}
if (includeToolResults) {
message.toolCallResults?.forEach { (callId, result) ->
val toolName = message.toolCalls
?.firstOrNull { it.id == callId }
?.function
?.name
?: "unknown-function"
tool { result(callId, toolName, result) }
}
}
} else if (!message.response.isNullOrBlank()) {
assistant(message.response)
}
}
private fun buildKoogSystemPrompt(callParameters: ChatCompletionParameters): String {
val promptsSettings = service<PromptsSettings>().state
val filteredPrompts = service<FilteredPromptsService>()
val systemParts = mutableListOf<String>()
when (callParameters.conversationType) {
ConversationType.DEFAULT -> {
val selectedPersona = promptsSettings.personas.selectedPersona
if (!selectedPersona.disabled) {
val baseInstructions = callParameters.personaDetails?.instructions?.addProjectPath()
?: filteredPrompts
.getFilteredPersonaPrompt(callParameters.chatMode)
.addProjectPath()
val clickableInstructions = filteredPrompts.applyClickableLinks(baseInstructions)
if (clickableInstructions.isNotBlank()) {
systemParts.add(clickableInstructions)
}
}
}
ConversationType.REVIEW_CHANGES -> {
promptsSettings.coreActions.reviewChanges.instructions
?.takeIf { it.isNotBlank() }
?.let(systemParts::add)
}
ConversationType.FIX_COMPILE_ERRORS -> {
promptsSettings.coreActions.fixCompileErrors.instructions
?.takeIf { it.isNotBlank() }
?.let(systemParts::add)
}
else -> Unit
}
val history = callParameters.history
?.takeIf { it.isNotEmpty() }
?.joinToString("\n\n") { ConversationTagProcessor.formatConversation(it) }
.orEmpty()
if (history.isNotBlank()) {
systemParts.add("Conversation history:\n$history")
}
if (!callParameters.mcpTools.isNullOrEmpty() &&
callParameters.toolApprovalMode != ToolApprovalMode.BLOCK_ALL
) {
val toolsPrompt = McpToolPromptFormatter().formatToolsForSystemPrompt(callParameters.mcpTools!!)
if (toolsPrompt.isNotBlank()) {
systemParts.add(toolsPrompt)
}
}
return systemParts.joinToString("\n\n").trim()
}
override fun createInlineEditQuestionPrompt(parameters: ChatCompletionParameters): Prompt {
return createChatCompletionPrompt(buildInlineEditQuestionParams(parameters))
}
private fun buildInlineEditQuestionParams(parameters: ChatCompletionParameters): ChatCompletionParameters {
val systemPrompt = """
You are an Inline Edit assistant for a single open file.
Respond in two parts:
@ -108,7 +280,7 @@ abstract class BaseRequestFactory : CompletionRequestFactory {
.featureType(FeatureType.INLINE_EDIT)
.build()
return createChatRequest(newParams)
return newParams
}
protected fun prepareInlineEditSystemPrompt(params: InlineEditCompletionParameters): String {
@ -142,16 +314,15 @@ abstract class BaseRequestFactory : CompletionRequestFactory {
systemPrompt = systemPrompt.replace("{{CURRENT_FILE_CONTEXT}}", currentFileBlock)
val externalContext = buildString {
val currentPath = filePath
val unique = mutableSetOf<String>()
val hasRefs = params.referencedFiles
?.filter { it.filePath != currentPath }
?.filter { it.filePath != filePath }
?.any { !it.fileContent.isNullOrBlank() } == true
if (hasRefs) {
append("\n\n### Referenced Files")
params.referencedFiles
.filter { it.filePath != currentPath }
.filter { it.filePath != filePath }
.forEach {
if (!it.fileContent.isNullOrBlank() && unique.add(it.filePath)) {
append("\n\n```${it.fileExtension}:${it.filePath}\n")
@ -198,39 +369,38 @@ abstract class BaseRequestFactory : CompletionRequestFactory {
}
}
override fun createInlineEditRequest(params: InlineEditCompletionParameters): CompletionRequest {
override fun createInlineEditPrompt(params: InlineEditCompletionParameters): Prompt {
val systemPrompt = prepareInlineEditSystemPrompt(params)
return createBasicCompletionRequest(
systemPrompt,
"systemPrompt.userPrompt",
AUTO_APPLY_MAX_TOKENS,
true,
FeatureType.INLINE_EDIT
)
return prompt("inline-edit-completion") {
if (systemPrompt.isNotBlank()) {
system(systemPrompt)
}
params.conversation?.messages?.forEach { message ->
appendUserPrompt(message.prompt)
if (!message.response.isNullOrBlank()) {
assistant(message.response)
}
}
user("Implement.")
}
}
override fun createCommitMessageRequest(params: CommitMessageCompletionParameters): CompletionRequest {
return createBasicCompletionRequest(
params.systemPrompt,
params.gitDiff,
512,
true,
FeatureType.COMMIT_MESSAGE
)
override fun createCommitMessagePrompt(params: CommitMessageCompletionParameters): Prompt {
return createBasicPrompt(params.systemPrompt, params.gitDiff, "commit-message-completion")
}
override fun createLookupRequest(params: LookupCompletionParameters): CompletionRequest {
return createBasicCompletionRequest(
override fun createLookupPrompt(params: LookupCompletionParameters): Prompt {
return createBasicPrompt(
service<PromptsSettings>().state.coreActions.generateNameLookups.instructions
?: CoreActionsState.DEFAULT_GENERATE_NAME_LOOKUPS_PROMPT,
params.prompt,
LOOKUP_MAX_TOKENS,
false,
FeatureType.LOOKUP
"lookup-completion"
)
}
override fun createAutoApplyRequest(params: AutoApplyParameters): CompletionRequest {
override fun createAutoApplyPrompt(params: AutoApplyParameters): Prompt {
val destination = params.destination
val language = FileUtil.getFileExtension(destination.path)
@ -249,22 +419,27 @@ abstract class BaseRequestFactory : CompletionRequestFactory {
.replace("{{changes_to_merge}}", formattedSource)
.replace("{{destination_file}}", formattedDestination)
return createBasicCompletionRequest(
return createBasicPrompt(
systemPrompt,
"Merge the following changes to the destination file.",
AUTO_APPLY_MAX_TOKENS,
true,
FeatureType.AUTO_APPLY
"auto-apply-completion"
)
}
abstract fun createBasicCompletionRequest(
protected fun createBasicPrompt(
systemPrompt: String,
userPrompt: String,
maxTokens: Int = DEFAULT_MAX_TOKENS,
stream: Boolean = false,
featureType: FeatureType
): CompletionRequest
name: String = "completion"
): Prompt {
return prompt(name) {
if (systemPrompt.isNotBlank()) {
system(systemPrompt)
}
if (userPrompt.isNotBlank()) {
user(userPrompt)
}
}
}
protected fun getPromptWithFilesContext(callParameters: ChatCompletionParameters): String {
return callParameters.referencedFiles?.let {

117
src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionRequestService.kt Normal file
View file

@ -0,0 +1,117 @@
package ee.carlrobert.codegpt.completions
import ai.koog.agents.core.tools.ToolDescriptor
import ai.koog.prompt.dsl.Prompt
import ai.koog.prompt.dsl.prompt
import ai.koog.prompt.executor.model.PromptExecutor
import ai.koog.prompt.llm.LLModel
import ee.carlrobert.codegpt.agent.AgentFactory
import ee.carlrobert.codegpt.agent.clients.CustomOpenAILLMClient
import ee.carlrobert.codegpt.agent.clients.HttpClientProvider
import ee.carlrobert.codegpt.agent.clients.RetryingPromptExecutor
import ee.carlrobert.codegpt.settings.models.ModelSelection
import ee.carlrobert.codegpt.settings.service.FeatureType
import ee.carlrobert.codegpt.settings.service.ServiceType
import ee.carlrobert.codegpt.settings.service.custom.CustomServiceChatCompletionSettingsState
import kotlinx.coroutines.runBlocking
import javax.swing.JPanel
import kotlin.time.Duration.Companion.seconds
object CompletionRequestService {
fun getCompletion(
serviceType: ServiceType,
featureType: FeatureType,
prompt: Prompt,
modelSelection: ModelSelection,
tools: List<ToolDescriptor> = emptyList()
): String {
val executor = AgentFactory.createExecutor(serviceType, featureType = featureType)
return try {
runBlocking {
val responses = executor.execute(prompt, modelSelection.llmModel, tools)
CompletionTextExtractor.extract(responses)
}
} finally {
runCatching { executor.close() }
}
}
fun getCompletionAsync(
serviceType: ServiceType,
featureType: FeatureType,
prompt: Prompt,
modelSelection: ModelSelection,
eventListener: CompletionStreamEventListener,
tools: List<ToolDescriptor> = emptyList()
): CancellableRequest {
val request = CompletionRunnerRequest.Streaming(
executor = AgentFactory.createExecutor(serviceType, featureType = featureType),
model = modelSelection.llmModel,
prompt = prompt,
eventListener = eventListener,
tools = tools,
mode = StreamingMode.STREAMING,
cancellationResultBuilder = { StringBuilder(it) }
)
return CompletionRunnerFactory.create(request).run(request)
}
@JvmStatic
fun getChatCompletionAsync(
serviceType: ServiceType,
prompt: Prompt,
modelSelection: ModelSelection,
callParameters: ChatCompletionParameters,
eventListener: ChatStreamEventListener,
onToolCallUIUpdate: ((JPanel) -> Unit)? = null
): CancellableRequest {
val request = CompletionRunnerRequest.Chat(
serviceType = serviceType,
executor = AgentFactory.createExecutor(serviceType, featureType = FeatureType.CHAT),
model = modelSelection.llmModel,
prompt = prompt,
callParameters = callParameters,
eventListener = eventListener,
onToolCallUIUpdate = onToolCallUIUpdate ?: {}
)
return CompletionRunnerFactory.create(request).run(request)
}
fun testCustomServiceConnectionAsync(
settings: CustomServiceChatCompletionSettingsState,
apiKey: String?,
modelId: String?,
eventListener: CompletionStreamEventListener
): CancellableRequest {
val client = CustomOpenAILLMClient.fromSettingsState(
apiKey.orEmpty(),
settings,
HttpClientProvider.createHttpClient()
)
val retryPolicy = RetryingPromptExecutor.RetryPolicy(
maxAttempts = 2,
initialDelay = 1.seconds,
maxDelay = 4.seconds,
backoffMultiplier = 2.0,
jitterFactor = 0.1
)
val request = CompletionRunnerRequest.Streaming(
executor = RetryingPromptExecutor.fromClient(client, retryPolicy, null),
model = LLModel(
id = modelId?.takeIf { it.isNotBlank() } ?: "gpt-4.1-mini",
provider = CustomOpenAILLMClient.CustomOpenAI,
capabilities = emptyList(),
contextLength = 128_000,
maxOutputTokens = 4_096
),
prompt = prompt("custom-service-test-connection") {
user("Test connection")
},
eventListener = eventListener,
mode = StreamingMode.SINGLE_RESPONSE,
cancellationResultBuilder = { StringBuilder() }
)
return CompletionRunnerFactory.create(request).run(request)
}
}

39
src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionRunner.kt Normal file
View file

@ -0,0 +1,39 @@
package ee.carlrobert.codegpt.completions
import ai.koog.agents.core.tools.ToolDescriptor
import ai.koog.prompt.dsl.Prompt
import ai.koog.prompt.executor.model.PromptExecutor
import ai.koog.prompt.llm.LLModel
import ee.carlrobert.codegpt.settings.service.ServiceType
import javax.swing.JPanel
internal interface CompletionRunner {
fun run(request: CompletionRunnerRequest): CancellableRequest
}
internal enum class StreamingMode {
STREAMING,
SINGLE_RESPONSE
}
internal sealed interface CompletionRunnerRequest {
data class Streaming(
val executor: PromptExecutor,
val model: LLModel,
val prompt: Prompt,
val eventListener: CompletionStreamEventListener,
val tools: List<ToolDescriptor> = emptyList(),
val mode: StreamingMode = StreamingMode.STREAMING,
val cancellationResultBuilder: (StringBuilder) -> StringBuilder = { StringBuilder(it) }
) : CompletionRunnerRequest
data class Chat(
val serviceType: ServiceType,
val executor: PromptExecutor,
val model: LLModel,
val prompt: Prompt,
val callParameters: ChatCompletionParameters,
val eventListener: ChatStreamEventListener,
val onToolCallUIUpdate: (JPanel) -> Unit = {}
) : CompletionRunnerRequest
}

19
src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionRunnerFactory.kt Normal file
View file

@ -0,0 +1,19 @@
package ee.carlrobert.codegpt.completions
internal object CompletionRunnerFactory {
fun create(request: CompletionRunnerRequest): CompletionRunner {
return when (request) {
is CompletionRunnerRequest.Chat ->
if (shouldUseAgentRunner(request)) AgentCompletionRunner else ChatStreamingCompletionRunner
is CompletionRunnerRequest.Streaming -> StreamingCompletionRunner
}
}
private fun shouldUseAgentRunner(request: CompletionRunnerRequest.Chat): Boolean {
if (request.callParameters.requestType != RequestType.NORMAL_REQUEST) {
return true
}
return !request.callParameters.mcpTools.isNullOrEmpty() &&
request.callParameters.toolApprovalMode != ToolApprovalMode.BLOCK_ALL
}
}

119
src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionService.kt Normal file
View file

@ -0,0 +1,119 @@
package ee.carlrobert.codegpt.completions
import ai.koog.prompt.dsl.Prompt
import com.intellij.openapi.application.ApplicationManager
import com.intellij.openapi.components.Service
import com.intellij.openapi.components.service
import ee.carlrobert.codegpt.completions.CompletionRequestFactory.Companion.getFactory
import ee.carlrobert.codegpt.completions.CompletionRequestService.getCompletion
import ee.carlrobert.codegpt.credentials.CredentialsStore.CredentialKey
import ee.carlrobert.codegpt.credentials.CredentialsStore.isCredentialSet
import ee.carlrobert.codegpt.settings.models.ModelSettings
import ee.carlrobert.codegpt.settings.service.FeatureType
import ee.carlrobert.codegpt.settings.service.ServiceType
import java.util.concurrent.Callable
import java.util.concurrent.ExecutionException
@Service
class CompletionService {
fun getLookupCompletion(params: LookupCompletionParameters): String {
return getChatCompletion(FeatureType.LOOKUP) { serviceType ->
getFactory(serviceType).createLookupPrompt(params)
}
}
fun autoApply(params: AutoApplyParameters): String {
return getChatCompletion(FeatureType.AUTO_APPLY) { serviceType ->
getFactory(serviceType).createAutoApplyPrompt(params)
}
}
fun autoApply(
params: AutoApplyParameters,
eventListener: CompletionStreamEventListener
): CancellableRequest {
return getChatCompletion(FeatureType.AUTO_APPLY, eventListener) { serviceType ->
getFactory(serviceType).createAutoApplyPrompt(params)
}
}
fun getCommitMessage(
params: CommitMessageCompletionParameters,
eventListener: CompletionStreamEventListener
): CancellableRequest {
return getChatCompletion(FeatureType.COMMIT_MESSAGE, eventListener) { serviceType ->
getFactory(serviceType).createCommitMessagePrompt(params)
}
}
fun getInlineEditCompletion(
params: InlineEditCompletionParameters,
eventListener: CompletionStreamEventListener
): CancellableRequest {
return getChatCompletion(FeatureType.INLINE_EDIT, eventListener) { serviceType ->
getFactory(serviceType).createInlineEditPrompt(params)
}
}
private fun getChatCompletion(
featureType: FeatureType,
prompt: (ServiceType) -> Prompt
): String {
val modelSelection = service<ModelSettings>().getModelSelectionForFeature(featureType)
return getCompletion(
modelSelection.provider,
featureType,
prompt(modelSelection.provider),
modelSelection
)
}
private fun getChatCompletion(
featureType: FeatureType,
eventListener: CompletionStreamEventListener,
prompt: (ServiceType) -> Prompt
): CancellableRequest {
val modelSelection = service<ModelSettings>().getModelSelectionForFeature(featureType)
return CompletionRequestService.getCompletionAsync(
modelSelection.provider,
featureType,
prompt(modelSelection.provider),
modelSelection,
eventListener
)
}
companion object {
@JvmStatic
fun isRequestAllowed(featureType: FeatureType): Boolean {
try {
return ApplicationManager.getApplication()
.executeOnPooledThread(Callable {
val serviceType = service<ModelSettings>().getServiceForFeature(featureType)
isRequestAllowed(serviceType)
})
.get()
} catch (e: InterruptedException) {
throw RuntimeException(e)
} catch (e: ExecutionException) {
throw RuntimeException(e)
}
}
@JvmStatic
fun isRequestAllowed(serviceType: ServiceType): Boolean {
return when (serviceType) {
ServiceType.OPENAI -> isCredentialSet(CredentialKey.OpenaiApiKey)
ServiceType.ANTHROPIC -> isCredentialSet(
CredentialKey.AnthropicApiKey
)
ServiceType.GOOGLE -> isCredentialSet(CredentialKey.GoogleApiKey)
ServiceType.MISTRAL -> isCredentialSet(CredentialKey.MistralApiKey)
ServiceType.INCEPTION -> isCredentialSet(CredentialKey.InceptionApiKey)
ServiceType.PROXYAI, ServiceType.CUSTOM_OPENAI, ServiceType.LLAMA_CPP, ServiceType.OLLAMA -> true
}
}
}
}

13
src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionStreamEventListener.kt Normal file
View file

@ -0,0 +1,13 @@
package ee.carlrobert.codegpt.completions
interface CompletionStreamEventListener {
fun onOpen() {}
fun onMessage(message: String)
fun onComplete(messageBuilder: StringBuilder)
fun onCancelled(messageBuilder: StringBuilder)
fun onError(error: CompletionError, ex: Throwable)
}

33
src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionTextExtractor.kt Normal file
View file

@ -0,0 +1,33 @@
package ee.carlrobert.codegpt.completions
import ai.koog.prompt.message.Message
internal object CompletionTextExtractor {
fun extract(responses: List<Message.Response>): String {
val reasoningText = responses
.filterIsInstance<Message.Reasoning>()
.map { it.content.trim() }
.filter { it.isNotEmpty() }
.joinToString("\n")
.trim()
val assistantText = responses
.filterIsInstance<Message.Assistant>()
.joinToString("\n") { it.content.trim() }
.trim()
if (reasoningText.isNotEmpty() && assistantText.isNotEmpty()) {
return "<think>$reasoningText</think>\n$assistantText"
}
if (assistantText.isNotEmpty()) {
return assistantText
}
if (reasoningText.isNotEmpty()) {
return "<think>$reasoningText</think>"
}
return responses.joinToString("\n") { it.content.trim() }.trim()
}
}

62
src/main/kotlin/ee/carlrobert/codegpt/completions/McpToolAdapter.kt Normal file
View file

@ -0,0 +1,62 @@
package ee.carlrobert.codegpt.completions
import ai.koog.agents.core.tools.Tool
import ee.carlrobert.codegpt.mcp.McpTool
import ee.carlrobert.codegpt.mcp.McpToolCallHandler
import kotlinx.coroutines.suspendCancellableCoroutine
import kotlinx.serialization.builtins.serializer
import kotlinx.serialization.json.Json
import kotlinx.serialization.json.JsonObject
import java.util.UUID
import java.util.concurrent.CompletableFuture
import javax.swing.JPanel
import kotlin.coroutines.resume
import kotlin.coroutines.resumeWithException
internal class McpToolAdapter(
private val tool: McpTool,
private val exposedName: String,
private val callParameters: ChatCompletionParameters,
private val toolCallHandler: McpToolCallHandler,
private val onToolCallUIUpdate: (JPanel) -> Unit
) : Tool<JsonObject, String>(
argsSerializer = JsonObject.serializer(),
resultSerializer = String.serializer(),
name = exposedName,
description = tool.description.ifBlank { "MCP tool" }
) {
override suspend fun execute(args: JsonObject): String {
val callId = UUID.randomUUID().toString()
val toolCall = ChatToolCall(
id = callId,
function = ChatToolFunction(
name = tool.name,
arguments = Json.encodeToString(JsonObject.serializer(), args)
)
)
callParameters.message.addToolCall(toolCall)
val result = awaitResult(
toolCallHandler.executeToolCall(
toolCall = toolCall,
serverId = tool.serverId,
conversationId = callParameters.conversation.id,
approvalMode = callParameters.toolApprovalMode,
onUIUpdate = onToolCallUIUpdate
)
)
callParameters.message.addToolCallResult(callId, result)
return result
}
private suspend fun awaitResult(future: CompletableFuture<String>): String =
suspendCancellableCoroutine { continuation ->
future.whenComplete { value, throwable ->
if (throwable != null) {
continuation.resumeWithException(throwable)
} else {
continuation.resume(value)
}
}
continuation.invokeOnCancellation { future.cancel(true) }
}
}

85
src/main/kotlin/ee/carlrobert/codegpt/completions/StreamingCompletionRunner.kt Normal file
View file

@ -0,0 +1,85 @@
package ee.carlrobert.codegpt.completions
import ai.koog.prompt.streaming.StreamFrame
import ee.carlrobert.codegpt.util.ReasoningFrameTextAdapter
import kotlinx.coroutines.CancellationException
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.Job
import kotlinx.coroutines.SupervisorJob
import kotlinx.coroutines.cancel
import kotlinx.coroutines.launch
import java.util.concurrent.atomic.AtomicBoolean
import java.util.concurrent.atomic.AtomicReference
internal object StreamingCompletionRunner : CompletionRunner {
override fun run(request: CompletionRunnerRequest): CancellableRequest {
val streamRequest = request as? CompletionRunnerRequest.Streaming
?: throw IllegalArgumentException("StreamingCompletionRunner can only run Streaming requests")
return executeAsync(streamRequest)
}
private fun executeAsync(request: CompletionRunnerRequest.Streaming): CancellableRequest {
val scope = CoroutineScope(SupervisorJob() + Dispatchers.IO)
val cancelled = AtomicBoolean(false)
val jobRef = AtomicReference<Job?>()
val cancellableRequest = CancellableRequest {
cancelled.set(true)
jobRef.get()?.cancel(CancellationException("Cancelled by user"))
}
request.eventListener.onOpen()
val job = scope.launch {
val messageBuilder = StringBuilder()
try {
when (request.mode) {
StreamingMode.STREAMING -> {
val frameAdapter = ReasoningFrameTextAdapter()
request.executor.executeStreaming(
request.prompt,
request.model,
request.tools
).collect { frame ->
frameAdapter.consume(frame).forEach { chunk ->
if (chunk.isNotEmpty()) {
messageBuilder.append(chunk)
request.eventListener.onMessage(chunk)
}
}
}
}
StreamingMode.SINGLE_RESPONSE -> {
val responses = request.executor.execute(request.prompt, request.model, request.tools)
val text = CompletionTextExtractor.extract(responses)
if (text.isNotBlank()) {
messageBuilder.append(text)
request.eventListener.onMessage(text)
}
}
}
if (cancelled.get()) {
request.eventListener.onCancelled(StringBuilder(messageBuilder))
return@launch
}
request.eventListener.onComplete(StringBuilder(messageBuilder))
} catch (_: CancellationException) {
request.eventListener.onCancelled(request.cancellationResultBuilder(messageBuilder))
} catch (exception: Throwable) {
request.eventListener.onError(
CompletionError(exception.message ?: "Failed to complete request"),
exception
)
} finally {
runCatching { request.executor.close() }
scope.cancel()
}
}
jobRef.set(job)
return cancellableRequest
}
}

14
src/main/kotlin/ee/carlrobert/codegpt/completions/autoapply/AutoApplyModels.kt Normal file
View file

@ -0,0 +1,14 @@
package ee.carlrobert.codegpt.completions.autoapply
import com.fasterxml.jackson.annotation.JsonIgnoreProperties
data class AutoApplyRequest(
val model: String,
val code: String,
val update: String,
)
@JsonIgnoreProperties(ignoreUnknown = true)
data class AutoApplyResponse(
val mergedCode: String = "",
)

166
src/main/kotlin/ee/carlrobert/codegpt/completions/factory/ClaudeRequestFactory.kt
View file

@ -1,169 +1,5 @@
package ee.carlrobert.codegpt.completions.factory
import com.intellij.openapi.components.service
import ee.carlrobert.codegpt.completions.BaseRequestFactory
import ee.carlrobert.codegpt.completions.ChatCompletionParameters
import ee.carlrobert.codegpt.completions.ToolApprovalMode
import ee.carlrobert.codegpt.mcp.McpToolPromptFormatter
import ee.carlrobert.codegpt.settings.configuration.ConfigurationSettings
import ee.carlrobert.codegpt.settings.prompts.FilteredPromptsService
import ee.carlrobert.codegpt.settings.prompts.PromptsSettings
import ee.carlrobert.codegpt.settings.service.FeatureType
import ee.carlrobert.codegpt.settings.service.ModelSelectionService
import ee.carlrobert.llm.client.anthropic.completion.*
import ee.carlrobert.llm.completion.CompletionRequest
class ClaudeRequestFactory : BaseRequestFactory() {
private val mcpToolPromptFormatter = McpToolPromptFormatter()
override fun createChatRequest(params: ChatCompletionParameters): ClaudeCompletionRequest {
return ClaudeCompletionRequest().apply {
model = ModelSelectionService.getInstance().getModelForFeature(FeatureType.CHAT)
maxTokens = service<ConfigurationSettings>().state.maxTokens
isStream = true
var systemPrompt = ""
val selectedPersona = service<PromptsSettings>().state.personas.selectedPersona
if (!selectedPersona.disabled) {
val base =
service<FilteredPromptsService>().getFilteredPersonaPrompt(params.chatMode)
systemPrompt = service<FilteredPromptsService>().applyClickableLinks(base)
}
if (!params.mcpTools.isNullOrEmpty() && params.toolApprovalMode != ToolApprovalMode.BLOCK_ALL) {
val toolsPrompt =
mcpToolPromptFormatter.formatToolsForSystemPrompt(params.mcpTools!!)
if (toolsPrompt.isNotEmpty()) {
systemPrompt = if (systemPrompt.isEmpty()) {
toolsPrompt
} else {
"$systemPrompt\n\n$toolsPrompt"
}
}
tools = params.mcpTools!!.map { mcpTool ->
ClaudeTool(
mcpTool.name,
mcpTool.description,
convertMcpSchemaToClaudeInputSchema(mcpTool.schema)
)
}
toolChoice = when (params.toolApprovalMode) {
ToolApprovalMode.AUTO_APPROVE -> ClaudeToolChoice.auto()
ToolApprovalMode.REQUIRE_APPROVAL -> ClaudeToolChoice.auto()
else -> null
}
}
if (systemPrompt.isNotEmpty()) {
system = systemPrompt
}
messages = buildClaudeMessages(params)
if (params.toolResults.isNullOrEmpty()) {
when {
params.imageDetails != null -> {
messages.add(
ClaudeCompletionDetailedMessage(
"user",
listOf(
ClaudeMessageImageContent(
ClaudeBase64Source(
params.imageDetails!!.mediaType,
params.imageDetails!!.data
)
),
ClaudeMessageTextContent(params.message.prompt)
)
)
)
}
else -> {
val promptWithContext = getPromptWithFilesContext(params)
if (promptWithContext.isNotBlank()) {
messages.add(
ClaudeCompletionStandardMessage(
"user", promptWithContext
)
)
}
}
}
}
}
}
override fun createBasicCompletionRequest(
systemPrompt: String,
userPrompt: String,
maxTokens: Int,
stream: Boolean,
featureType: FeatureType
): CompletionRequest {
return ClaudeCompletionRequest().apply {
system = systemPrompt
isStream = stream
model = ModelSelectionService.getInstance().getModelForFeature(featureType)
messages =
listOf<ClaudeCompletionMessage>(ClaudeCompletionStandardMessage("user", userPrompt))
this.maxTokens = maxTokens
}
}
private fun buildClaudeMessages(params: ChatCompletionParameters): MutableList<ClaudeCompletionMessage> {
val messages = mutableListOf<ClaudeCompletionMessage>()
for (prevMessage in params.conversation.messages) {
if (prevMessage.id == params.message.id && params.toolResults.isNullOrEmpty()) {
if (params.retry) {
break
} else {
continue
}
}
if (prevMessage.prompt.isNotEmpty()) {
messages.add(ClaudeCompletionStandardMessage("user", prevMessage.prompt))
}
val response = prevMessage.response
if (!response.isNullOrEmpty()) {
if (!params.toolResults.isNullOrEmpty()) {
val isCurrentMessageWithToolResults =
prevMessage.id == params.message.id && !params.toolResults.isNullOrEmpty()
if (params.toolResults.isNullOrEmpty() && !isCurrentMessageWithToolResults) {
continue
}
try {
if (isCurrentMessageWithToolResults) {
return mutableListOf()
}
} catch (e: Exception) {
messages.add(ClaudeCompletionStandardMessage("assistant", response))
}
} else {
messages.add(ClaudeCompletionStandardMessage("assistant", response))
}
}
}
return messages
}
private fun convertMcpSchemaToClaudeInputSchema(mcpSchema: Map<String, Any>): Map<String, Any> {
if (mcpSchema.containsKey("type")) {
return mcpSchema
}
return mapOf(
"type" to "object",
"properties" to mcpSchema,
"required" to mcpSchema.keys.toList()
)
}
}
class ClaudeRequestFactory : BaseRequestFactory()

323
src/main/kotlin/ee/carlrobert/codegpt/completions/factory/CodeGPTRequestFactory.kt
View file

@ -1,326 +1,5 @@
package ee.carlrobert.codegpt.completions.factory
import com.intellij.openapi.application.ApplicationInfo
import com.intellij.openapi.components.service
import com.intellij.openapi.vfs.LocalFileSystem
import com.intellij.openapi.vfs.VirtualFile
import ee.carlrobert.codegpt.CodeGPTPlugin
import ee.carlrobert.codegpt.ReferencedFile
import ee.carlrobert.codegpt.completions.BaseRequestFactory
import ee.carlrobert.codegpt.completions.ChatCompletionParameters
import ee.carlrobert.codegpt.completions.InlineEditCompletionParameters
import ee.carlrobert.codegpt.completions.ToolApprovalMode
import ee.carlrobert.codegpt.completions.factory.OpenAIRequestFactory.Companion.buildOpenAIMessages
import ee.carlrobert.codegpt.mcp.McpToolConverter
import ee.carlrobert.codegpt.psistructure.ClassStructureSerializer
import ee.carlrobert.codegpt.psistructure.models.ClassStructure
import ee.carlrobert.codegpt.settings.configuration.ConfigurationSettings
import ee.carlrobert.codegpt.settings.prompts.CoreActionsState
import ee.carlrobert.codegpt.settings.prompts.PromptsSettings
import ee.carlrobert.codegpt.settings.service.FeatureType
import ee.carlrobert.codegpt.settings.service.ModelSelectionService
import ee.carlrobert.codegpt.ui.textarea.ConversationTagProcessor
import ee.carlrobert.codegpt.util.file.FileUtil
import ee.carlrobert.llm.client.codegpt.request.InlineEditRequest
import ee.carlrobert.llm.client.codegpt.request.chat.*
import ee.carlrobert.llm.client.openai.completion.request.*
import ee.carlrobert.llm.completion.CompletionRequest
class CodeGPTRequestFactory(private val classStructureSerializer: ClassStructureSerializer) :
BaseRequestFactory() {
override fun createChatRequest(params: ChatCompletionParameters): ChatCompletionRequest {
val model = ModelSelectionService.getInstance().getModelForFeature(FeatureType.CHAT)
val configuration = service<ConfigurationSettings>().state
val messages = buildCodeGPTMessages(model, params, emptyList(), null)
val requestBuilder: ChatCompletionRequest.Builder =
ChatCompletionRequest.Builder(messages)
.setModel(model)
.setSessionId(params.sessionId)
.setStream(true)
.setTools(params.mcpTools?.map { McpToolConverter.convertToOpenAITool(it) }
?: emptyList())
.setMetadata(
Metadata(
CodeGPTPlugin.getVersion(),
service<ApplicationInfo>().build.asString()
)
)
if ("o4-mini" == model) {
requestBuilder
.setMaxTokens(null)
.setTemperature(null)
} else {
requestBuilder
.setMaxTokens(configuration.maxTokens)
.setTemperature(configuration.temperature.toDouble())
}
if (params.message.isWebSearchIncluded) {
requestBuilder.setWebSearchIncluded(true)
}
val contextFiles = params.referencedFiles
?.mapNotNull { file ->
LocalFileSystem.getInstance().findFileByPath(file.filePath)?.let {
if (it.isDirectory) {
val children = mutableListOf<ContextFile>()
processFolder(it, children)
children
} else {
listOf(ContextFile(file.fileName(), file.filePath(), file.fileContent()))
}
}
}
?.flatten()
.orEmpty()
val psiContext = params.psiStructure?.map { classStructure ->
ContextFile(
classStructure.virtualFile.name,
classStructure.virtualFile.path,
classStructureSerializer.serialize(classStructure)
)
}.orEmpty()
val contextFilesWithPsi = contextFiles + psiContext
if (!params.mcpTools.isNullOrEmpty() && params.toolApprovalMode != ToolApprovalMode.BLOCK_ALL) {
val codegptTools = params.mcpTools!!.map { mcpTool ->
Tool().apply {
function = ToolFunction().apply {
name = mcpTool.name
description = mcpTool.description
parameters = convertMcpSchemaToCodeGPTParameters(mcpTool.schema)
}
}
}
requestBuilder.setTools(codegptTools)
val toolChoice = when (params.toolApprovalMode) {
ToolApprovalMode.AUTO_APPROVE -> "auto"
ToolApprovalMode.REQUIRE_APPROVAL -> "require_approval"
else -> null
}
if (toolChoice != null) {
requestBuilder.setToolChoice(toolChoice)
}
}
val conversationsHistory = params.history?.joinToString("\n\n") {
ConversationTagProcessor.formatConversation(it)
}
requestBuilder.setContext(
AdditionalRequestContext(
contextFilesWithPsi,
conversationsHistory
)
)
return requestBuilder.build()
}
private fun processFolder(folder: VirtualFile, contextFiles: MutableList<ContextFile>) {
folder.children.forEach { child ->
when {
child.isDirectory -> processFolder(child, contextFiles)
else -> contextFiles.add(
ContextFile(
child.name,
child.path,
FileUtil.readContent(child)
)
)
}
}
}
override fun createBasicCompletionRequest(
systemPrompt: String,
userPrompt: String,
maxTokens: Int,
stream: Boolean,
featureType: FeatureType
): ChatCompletionRequest {
val model = ModelSelectionService.getInstance().getModelForFeature(featureType)
if (model == "o4-mini") {
return buildBasicO1Request(model, userPrompt, systemPrompt, maxTokens, stream = stream)
}
return ChatCompletionRequest.Builder(
listOf(
OpenAIChatCompletionStandardMessage("system", systemPrompt),
OpenAIChatCompletionStandardMessage("user", userPrompt)
)
)
.setModel(model)
.setStream(stream)
.build()
}
override fun createInlineEditRequest(params: InlineEditCompletionParameters): CompletionRequest {
val model = ModelSelectionService.getInstance().getModelForFeature(FeatureType.INLINE_EDIT)
val systemTemplate = service<PromptsSettings>().state.coreActions.inlineEdit.instructions
?: CoreActionsState.DEFAULT_INLINE_EDIT_PROMPT
val contextFiles = mutableListOf<InlineEditRequest.ContextFile>()
params.referencedFiles
?.filter { !it.fileContent.isNullOrBlank() }
?.forEach { ref ->
contextFiles.add(
InlineEditRequest.ContextFile(ref.fileName(), ref.filePath(), ref.fileContent())
)
}
val conversationHistory = params.conversationHistory?.map { conversation ->
conversation.messages.flatMap { msg ->
listOfNotNull(
msg.prompt?.takeIf { it.isNotBlank() }?.let {
InlineEditRequest.ConversationMessage("user", it)
},
msg.response?.takeIf { it.isNotBlank() }?.let {
InlineEditRequest.ConversationMessage("assistant", it)
}
)
}
} ?: emptyList()
val fileContent = params.filePath?.let { path ->
LocalFileSystem.getInstance().findFileByPath(path)?.let { vf ->
FileUtil.readContent(vf)
}
}
val currentConversationMessages = params.conversation?.messages?.flatMap { msg ->
listOfNotNull(
msg.prompt?.takeIf { it.isNotBlank() }?.let {
InlineEditRequest.ConversationMessage("user", it)
},
msg.response?.takeIf { it.isNotBlank() }?.let {
InlineEditRequest.ConversationMessage("assistant", it)
}
)
}
return InlineEditRequest.Builder(model, systemTemplate)
.setMetadata(
Metadata(CodeGPTPlugin.getVersion(), service<ApplicationInfo>().build.asString())
)
.setConversation(currentConversationMessages)
.setSelectedText(params.selectedText)
.setCursorOffset(params.cursorOffset ?: 0)
.setFilePath(params.filePath)
.setFileContent(fileContent)
.setProjectBasePath(params.projectBasePath)
.setContextFiles(contextFiles)
.setGitDiff(params.gitDiff)
.setConversationHistory(conversationHistory.toMutableList())
.setDiagnosticsInfo(params.diagnosticsInfo)
.build()
}
private fun buildBasicO1Request(
model: String,
prompt: String,
systemPrompt: String = "",
maxCompletionTokens: Int = 4096,
stream: Boolean = false
): ChatCompletionRequest {
val messages = if (systemPrompt.isEmpty()) {
listOf(OpenAIChatCompletionStandardMessage("user", prompt))
} else {
listOf(
OpenAIChatCompletionStandardMessage("user", systemPrompt),
OpenAIChatCompletionStandardMessage("user", prompt)
)
}
return ChatCompletionRequest.Builder(messages)
.setModel(model)
.setMaxTokens(maxCompletionTokens)
.setStream(stream)
.setTemperature(null)
.build()
}
private fun convertMcpSchemaToCodeGPTParameters(mcpSchema: Map<String, Any>): ToolFunctionParameters {
return ToolFunctionParameters().apply {
if (mcpSchema.containsKey("type")) {
type = mcpSchema["type"] as? String ?: "object"
properties = mcpSchema["properties"] as? Map<String, Any> ?: mcpSchema
required = mcpSchema["required"] as? List<String> ?: emptyList()
} else {
type = "object"
properties = mcpSchema
required = mcpSchema.keys.toList()
}
}
}
private fun buildCodeGPTMessages(
model: String?,
params: ChatCompletionParameters,
referencedFiles: List<ReferencedFile>?,
psiStructure: Set<ClassStructure>?
): List<OpenAIChatCompletionMessage> {
val messages =
buildOpenAIMessages(model, params, referencedFiles, emptyList(), psiStructure)
/*if (params.toolResults.isNullOrEmpty()) {
return messages
}
val convertedMessages = mutableListOf<OpenAIChatCompletionMessage>()
var pendingToolResults = mutableListOf<Pair<String, String>>()
for (message in messages) {
when (message) {
is OpenAIChatCompletionToolMessage -> {
val callId = message.callId ?: ""
val content = message.content ?: ""
pendingToolResults.add(Pair(callId, content))
}
is OpenAIChatCompletionAssistantMessage -> {
if (!message.toolCalls.isNullOrEmpty() && pendingToolResults.isNotEmpty()) {
val toolCallsInfo = message.toolCalls.joinToString("\n") { toolCall ->
val result = pendingToolResults.find { it.first == toolCall.id }?.second
?: "No result"
"Tool ${toolCall.function.name} (${toolCall.id}): $result"
}
val newContent = if (message.content.isNullOrEmpty()) {
"Tool execution results:\n$toolCallsInfo"
} else {
"${message.content}\n\nTool execution results:\n$toolCallsInfo"
}
convertedMessages.add(
OpenAIChatCompletionStandardMessage("assistant", newContent)
)
pendingToolResults.clear()
} else {
convertedMessages.add(message)
}
}
else -> {
convertedMessages.add(message)
}
}
}
if (pendingToolResults.isNotEmpty()) {
val toolResultsContent = pendingToolResults.joinToString("\n") { (toolCallId, result) ->
"Tool result ($toolCallId): $result"
}
convertedMessages.add(
OpenAIChatCompletionStandardMessage(
"assistant",
"Tool execution results:\n$toolResultsContent"
)
)
}*/
return messages
}
}
class CodeGPTRequestFactory : BaseRequestFactory()

Some files were not shown because too many files have changed in this diff Show more