diff --git a/build.gradle.kts b/build.gradle.kts index 1dd14068..4dd4f88d 100644 --- a/build.gradle.kts +++ b/build.gradle.kts @@ -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") } diff --git a/buildSrc/src/main/kotlin/proxyai.java-conventions.gradle.kts b/buildSrc/src/main/kotlin/proxyai.java-conventions.gradle.kts index de892d7a..41e35a39 100644 --- a/buildSrc/src/main/kotlin/proxyai.java-conventions.gradle.kts +++ b/buildSrc/src/main/kotlin/proxyai.java-conventions.gradle.kts @@ -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("kotlin") { + jvmToolchain(it.toInt()) + } withType { sourceCompatibility = it targetCompatibility = it } withType { compilerOptions.jvmTarget.set(JvmTarget.fromTarget(it)) + compilerOptions.freeCompilerArgs.add("-opt-in=kotlin.time.ExperimentalTime") } } } diff --git a/gradle.properties b/gradle.properties index a3de177c..e5edd43e 100644 --- a/gradle.properties +++ b/gradle.properties @@ -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 diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml index 2ef86c62..22e7dd46 100644 --- a/gradle/libs.versions.toml +++ b/gradle/libs.versions.toml @@ -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" } diff --git a/gradle/wrapper/gradle-wrapper.properties b/gradle/wrapper/gradle-wrapper.properties index 94113f20..37f853b1 100644 --- a/gradle/wrapper/gradle-wrapper.properties +++ b/gradle/wrapper/gradle-wrapper.properties @@ -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 diff --git a/src/main/java/ee/carlrobert/codegpt/CodeGPTKeys.java b/src/main/java/ee/carlrobert/codegpt/CodeGPTKeys.java index 2837d5d1..d952e7b3 100644 --- a/src/main/java/ee/carlrobert/codegpt/CodeGPTKeys.java +++ b/src/main/java/ee/carlrobert/codegpt/CodeGPTKeys.java @@ -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; diff --git a/src/main/java/ee/carlrobert/codegpt/EncodingManager.java b/src/main/java/ee/carlrobert/codegpt/EncodingManager.java index a0d276c4..078f2d93 100644 --- a/src/main/java/ee/carlrobert/codegpt/EncodingManager.java +++ b/src/main/java/ee/carlrobert/codegpt/EncodingManager.java @@ -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 toolCalls = message.getToolCalls(); + List 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; diff --git a/src/main/java/ee/carlrobert/codegpt/Icons.java b/src/main/java/ee/carlrobert/codegpt/Icons.java index 15d72060..e3569f4e 100644 --- a/src/main/java/ee/carlrobert/codegpt/Icons.java +++ b/src/main/java/ee/carlrobert/codegpt/Icons.java @@ -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); diff --git a/src/main/java/ee/carlrobert/codegpt/completions/ChatCompletionEventListener.java b/src/main/java/ee/carlrobert/codegpt/completions/ChatCompletionEventListener.java deleted file mode 100644 index 51ae08a3..00000000 --- a/src/main/java/ee/carlrobert/codegpt/completions/ChatCompletionEventListener.java +++ /dev/null @@ -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 { - - 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(); - } -} diff --git a/src/main/java/ee/carlrobert/codegpt/completions/CompletionClientProvider.java b/src/main/java/ee/carlrobert/codegpt/completions/CompletionClientProvider.java deleted file mode 100644 index 1915550c..00000000 --- a/src/main/java/ee/carlrobert/codegpt/completions/CompletionClientProvider.java +++ /dev/null @@ -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); - } -} diff --git a/src/main/java/ee/carlrobert/codegpt/completions/CompletionRequestService.java b/src/main/java/ee/carlrobert/codegpt/completions/CompletionRequestService.java deleted file mode 100644 index 3cf2b6f8..00000000 --- a/src/main/java/ee/carlrobert/codegpt/completions/CompletionRequestService.java +++ /dev/null @@ -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 eventListener) { - var httpClient = CompletionClientProvider.getDefaultClientBuilder().build(); - return EventSources.createFactory(httpClient).newEventSource( - customRequest, - new OpenAITextCompletionEventSourceListener(eventListener)); - } - - public EventSource getCustomOpenAIChatCompletionAsync( - Request customRequest, - CompletionEventListener 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 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 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 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 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 eventListener) { - var serviceType = ModelSelectionService.getInstance().getServiceForFeature(FeatureType.CHAT); - return getChatCompletionAsync(request, eventListener, serviceType, FeatureType.CHAT); - } - - public EventSource getChatCompletionAsync( - CompletionRequest request, - CompletionEventListener 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. - *
    - *
  • Search all choices which are not null
    • - *
  • Search all messages which are not null
    • - *
  • Use first content which is not null or blank (whitespace)
    • - *
- * - * @return First non-blank content or {@code Optional.empty()} - */ - private Optional 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(); - } -} diff --git a/src/main/java/ee/carlrobert/codegpt/completions/CompletionResponseEventListener.java b/src/main/java/ee/carlrobert/codegpt/completions/CompletionResponseEventListener.java index 99c153ed..5c842d24 100644 --- a/src/main/java/ee/carlrobert/codegpt/completions/CompletionResponseEventListener.java +++ b/src/main/java/ee/carlrobert/codegpt/completions/CompletionResponseEventListener.java @@ -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) { diff --git a/src/main/java/ee/carlrobert/codegpt/completions/ToolwindowChatCompletionRequestHandler.java b/src/main/java/ee/carlrobert/codegpt/completions/ToolwindowChatCompletionRequestHandler.java index 593e11fe..a01eb569 100644 --- a/src/main/java/ee/carlrobert/codegpt/completions/ToolwindowChatCompletionRequestHandler.java +++ b/src/main/java/ee/carlrobert/codegpt/completions/ToolwindowChatCompletionRequestHandler.java @@ -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 activeEventSources = new ArrayList<>(); + private final List 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 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()) diff --git a/src/main/java/ee/carlrobert/codegpt/conversations/ConversationService.java b/src/main/java/ee/carlrobert/codegpt/conversations/ConversationService.java index 052632d0..ce1b6be1 100644 --- a/src/main/java/ee/carlrobert/codegpt/conversations/ConversationService.java +++ b/src/main/java/ee/carlrobert/codegpt/conversations/ConversationService.java @@ -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 toolCalls) { + public void saveAssistantMessageWithToolCalls(@NotNull Conversation conversation, @NotNull Message message, List 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 toolCalls) { + public void saveAssistantMessageWithToolCalls(ChatCompletionParameters callParameters, List toolCalls) { saveAssistantMessageWithToolCalls(callParameters.getConversation(), callParameters.getMessage(), toolCalls); } diff --git a/src/main/java/ee/carlrobert/codegpt/conversations/message/Message.java b/src/main/java/ee/carlrobert/codegpt/conversations/message/Message.java index f066dc2e..f3b27499 100644 --- a/src/main/java/ee/carlrobert/codegpt/conversations/message/Message.java +++ b/src/main/java/ee/carlrobert/codegpt/conversations/message/Message.java @@ -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 toolCalls; + private List toolCalls; private Map toolCallResults; public Message() { @@ -101,18 +102,18 @@ public class Message { } @JsonProperty("tool_calls") - public @Nullable List getToolCalls() { + public @Nullable List getToolCalls() { return toolCalls; } @JsonProperty("tool_calls") - public void setToolCalls(@Nullable List toolCalls) { - this.toolCalls = toolCalls; + public void setToolCalls(@Nullable List 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 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); } } diff --git a/src/main/java/ee/carlrobert/codegpt/settings/service/llama/LlamaSettings.java b/src/main/java/ee/carlrobert/codegpt/settings/service/llama/LlamaSettings.java index c0984bf7..e685fabc 100644 --- a/src/main/java/ee/carlrobert/codegpt/settings/service/llama/LlamaSettings.java +++ b/src/main/java/ee/carlrobert/codegpt/settings/service/llama/LlamaSettings.java @@ -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 { 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(); } } -} \ No newline at end of file +} diff --git a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ChatToolWindowTabPanel.java b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ChatToolWindowTabPanel.java index 5e9b1893..d21ea6e2 100644 --- a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ChatToolWindowTabPanel.java +++ b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ChatToolWindowTabPanel.java @@ -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); } diff --git a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ToolWindowCompletionResponseEventListener.java b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ToolWindowCompletionResponseEventListener.java index c488b492..f2043a5f 100644 --- a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ToolWindowCompletionResponseEventListener.java +++ b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ToolWindowCompletionResponseEventListener.java @@ -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); } }); } diff --git a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/ChatMessageResponseBody.java b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/ChatMessageResponseBody.java index 603e73df..af37f246 100644 --- a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/ChatMessageResponseBody.java +++ b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/ChatMessageResponseBody.java @@ -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); diff --git a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/ChatToolWindowScrollablePanel.java b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/ChatToolWindowScrollablePanel.java index 9d4865d6..7359cb64 100644 --- a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/ChatToolWindowScrollablePanel.java +++ b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/ChatToolWindowScrollablePanel.java @@ -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)) { diff --git a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/textarea/ModelComboBoxAction.java b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/textarea/ModelComboBoxAction.java index b8a46bc1..8c72ce7a 100644 --- a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/textarea/ModelComboBoxAction.java +++ b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/textarea/ModelComboBoxAction.java @@ -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 onModelChange; - private final Project project; private final List availableProviders; private final boolean showConfigureModels; private final FeatureType featureType; public ModelComboBoxAction( - Project project, - Consumer onModelChange, - ServiceType selectedService) { - this(project, onModelChange, selectedService, Arrays.asList(ServiceType.values()), true, - FeatureType.CHAT); - } - - public ModelComboBoxAction( - Project project, - Consumer onModelChange, - ServiceType selectedProvider, - List availableProviders, - boolean showConfigureModels) { - this(project, onModelChange, selectedProvider, availableProviders, showConfigureModels, - FeatureType.CHAT); - } - - public ModelComboBoxAction( - Project project, Consumer onModelChange, ServiceType selectedProvider, List 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 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 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()); } } diff --git a/src/main/java/ee/carlrobert/codegpt/ui/ModelIconLabel.java b/src/main/java/ee/carlrobert/codegpt/ui/ModelIconLabel.java index 49da1db1..2ca26311 100644 --- a/src/main/java/ee/carlrobert/codegpt/ui/ModelIconLabel.java +++ b/src/main/java/ee/carlrobert/codegpt/ui/ModelIconLabel.java @@ -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; } } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/actions/BaseCommitWorkflowAction.kt b/src/main/kotlin/ee/carlrobert/codegpt/actions/BaseCommitWorkflowAction.kt index 7c53e45b..beb6eb76 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/actions/BaseCommitWorkflowAction.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/actions/BaseCommitWorkflowAction.kt @@ -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 { +) : 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( } } } -} \ No newline at end of file +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/actions/CodeCompletionFeatureToggleActions.kt b/src/main/kotlin/ee/carlrobert/codegpt/actions/CodeCompletionFeatureToggleActions.kt index 4273eccb..e416ea6d 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/actions/CodeCompletionFeatureToggleActions.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/actions/CodeCompletionFeatureToggleActions.kt @@ -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().getServiceForFeature(FeatureType.CODE_COMPLETION) + service().getServiceForFeature(FeatureType.CODE_COMPLETION) when (serviceType) { PROXYAI -> { service().state.codeCompletionSettings.codeCompletionsEnabled = @@ -63,7 +63,7 @@ abstract class CodeCompletionFeatureToggleActions( override fun update(e: AnActionEvent) { val selectedService = - service().getServiceForFeature(FeatureType.CODE_COMPLETION) + service().getServiceForFeature(FeatureType.CODE_COMPLETION) val codeCompletionEnabled = service().isCodeCompletionsEnabled(selectedService) == true e.presentation.isVisible = codeCompletionEnabled != enableFeatureAction diff --git a/src/main/kotlin/ee/carlrobert/codegpt/actions/GenerateCommitMessageAction.kt b/src/main/kotlin/ee/carlrobert/codegpt/actions/GenerateCommitMessageAction.kt index f4246725..7aa59b94 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/actions/GenerateCommitMessageAction.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/actions/GenerateCommitMessageAction.kt @@ -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().getCommitMessage( CommitMessageCompletionParameters( gitDiff, - project.getService(CommitMessageTemplate::class.java).getSystemPrompt() + project.service().getSystemPrompt() ), CommitMessageEventListener(project, commitWorkflowUi) ) diff --git a/src/main/kotlin/ee/carlrobert/codegpt/actions/GenerateCommitMessageWithAdditionalInputAction.kt b/src/main/kotlin/ee/carlrobert/codegpt/actions/GenerateCommitMessageWithAdditionalInputAction.kt index 6117459f..bec456ec 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/actions/GenerateCommitMessageWithAdditionalInputAction.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/actions/GenerateCommitMessageWithAdditionalInputAction.kt @@ -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().getCommitMessage( CommitMessageCompletionParameters( gitDiff, "$systemPrompt\n\nUser input: $userInput" diff --git a/src/main/kotlin/ee/carlrobert/codegpt/actions/LlamaServerToggleActions.kt b/src/main/kotlin/ee/carlrobert/codegpt/actions/LlamaServerToggleActions.kt index 76202803..fe14833e 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/actions/LlamaServerToggleActions.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/actions/LlamaServerToggleActions.kt @@ -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) diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/AgentEvents.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/AgentEvents.kt index 4957d40b..edebca6d 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/AgentEvents.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/AgentEvents.kt @@ -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 diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/AgentFactory.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/AgentFactory.kt index 47673def..1f60313b 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/AgentFactory.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/AgentFactory.kt @@ -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().getAgentModel(), - maxAgentIterations = 100 + model = service().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().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() - .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().getAgentModel(), - maxAgentIterations = 100 + model = service().getAgentModel(), + maxAgentIterations = MAX_AGENT_ITERATIONS ), toolRegistry = createToolRegistry( project, @@ -377,8 +313,8 @@ object AgentFactory { """.trimIndent() ) }, - model = service().getAgentModel(), - maxAgentIterations = 100 + model = service().getAgentModel(), + maxAgentIterations = MAX_AGENT_ITERATIONS ), toolRegistry = createToolRegistry( project, @@ -498,16 +434,36 @@ object AgentFactory { ): List { 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): List { - return responses + private fun appendableResponses( + responses: List, + provider: LLMProvider + ): List { + 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): Long { @@ -543,43 +499,52 @@ object AgentFactory { val contextService = project.service() 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 ) ) diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/AgentService.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/AgentService.kt index e3efc535..e185f2a5 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/AgentService.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/AgentService.kt @@ -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().getServiceForFeature(FeatureType.AGENT) + val provider = service().getServiceForFeature(FeatureType.AGENT) val contentManager = project.service() 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? { @@ -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().setResumeCheckpointRef(sessionId, ref) } @@ -261,11 +256,10 @@ class AgentService(private val project: Project) { provider: ServiceType, events: AgentEvents ): SessionRuntime { - val modelId = service().getAgentModel().id + val modelId = service().getAgentModel().id val selectedServerIds = project.service() .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 ) } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/AgentSystemPrompts.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/AgentSystemPrompts.kt index 26b7a3b6..ee271b7e 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/AgentSystemPrompts.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/AgentSystemPrompts.kt @@ -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) diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/PromptEnhancer.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/PromptEnhancer.kt index aa605fc5..3a89d477 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/PromptEnhancer.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/PromptEnhancer.kt @@ -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() + val modelService = service() val provider = modelService.getServiceForFeature(FeatureType.AGENT) val model = modelService.getAgentModel() val executor = AgentFactory.createExecutor(provider) diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/ProxyAIAgent.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/ProxyAIAgent.kt index faa14baf..fd915050 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/ProxyAIAgent.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/ProxyAIAgent.kt @@ -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>, ): GraphAIAgentService { val modelSelection = - service().getModelSelectionForFeature(FeatureType.AGENT) + service().getModelSelectionForFeature(FeatureType.AGENT) val skills = project.service().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().getAgentModel() + val agentModel = service().getAgentModel() return AIAgentService( 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 = HashMap() val anonymousToolIds: ArrayDeque = 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("${it.content}") + } } } } @@ -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().getStoredModelForFeature(FeatureType.AGENT) + project.service().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() - val mcpSessionManager = service() 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> { 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) } } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/CodeCompletionSupport.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/CodeCompletionSupport.kt new file mode 100644 index 00000000..ed5de5f3 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/CodeCompletionSupport.kt @@ -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, + 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) +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/CodeGPTMistralAILLMClient.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/CodeGPTMistralAILLMClient.kt new file mode 100644 index 00000000..e795ea1e --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/CodeGPTMistralAILLMClient.kt @@ -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) + ) + ) + } +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/CodeGPTOllamaClient.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/CodeGPTOllamaClient.kt new file mode 100644 index 00000000..96c11023 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/CodeGPTOllamaClient.kt @@ -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().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() + } +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/CodeGPTOpenAILLMClient.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/CodeGPTOpenAILLMClient.kt new file mode 100644 index 00000000..ff806b74 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/CodeGPTOpenAILLMClient.kt @@ -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) + ) + ) + } +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/CustomOpenAILLMClient.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/CustomOpenAILLMClient.kt index 923e86ab..a980787d 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/CustomOpenAILLMClient.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/CustomOpenAILLMClient.kt @@ -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( 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, 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 ( + "\n${infillRequest.prefix}\n\n\n" + + "\n${infillRequest.suffix}\n\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, + credential: String? + ): Map { + 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>() {}) + } 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 { 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, + messages: List, + streamRequest: Boolean, + credential: String, + json: Json +): Map? = 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, + 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, 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()) diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/CustomOpenAIParams.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/CustomOpenAIParams.kt index 4627ae5a..76f38866 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/CustomOpenAIParams.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/CustomOpenAIParams.kt @@ -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? { 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, diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/HttpClientProvider.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/HttpClientProvider.kt index b9ad25d4..5298c712 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/HttpClientProvider.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/HttpClientProvider.kt @@ -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) + } + } + } } } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/InceptionAILLMClient.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/InceptionAILLMClient.kt index 5b369afd..f5afb64d 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/InceptionAILLMClient.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/InceptionAILLMClient.kt @@ -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() + .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" } } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/LlamaCppLLMClient.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/LlamaCppLLMClient.kt new file mode 100644 index 00000000..14a3cc79 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/LlamaCppLLMClient.kt @@ -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 + ) + ) + } +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/ProxyAILLMClient.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/ProxyAILLMClient.kt index 44bae851..4b27250e 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/ProxyAILLMClient.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/ProxyAILLMClient.kt @@ -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( 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, 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) diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/RetryingPromptExecutor.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/RetryingPromptExecutor.kt index e90ac518..9b8ae955 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/RetryingPromptExecutor.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/RetryingPromptExecutor.kt @@ -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() + .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 ): Flow { - var attempt = 1 - var refinedPrompt = prompt + val refinedPrompt = prompt.withReasoningParams(model) var hasReceivedData: Boolean fun createStream(attemptNum: Int): Flow { @@ -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 = delegate.models() + override suspend fun models(): List = delegate.models() data class RetryPolicy( val maxAttempts: Int, @@ -142,25 +167,20 @@ class RetryingPromptExecutor( model: LLModel, tools: List ): List { + 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) + } } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/history/AgentCheckpointConversationMapper.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/history/AgentCheckpointConversationMapper.kt index 0c74637b..7bd363a4 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/history/AgentCheckpointConversationMapper.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/history/AgentCheckpointConversationMapper.kt @@ -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() + val toolCalls = mutableListOf() val toolResults = LinkedHashMap() 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()) ) ) } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/history/AgentCheckpointTurnSequencer.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/history/AgentCheckpointTurnSequencer.kt index 01d152d1..0443d92f 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/history/AgentCheckpointTurnSequencer.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/history/AgentCheckpointTurnSequencer.kt @@ -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 diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/history/AgentHistoryThreadSummary.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/history/AgentHistoryThreadSummary.kt index 8950bb2e..9e0a8a51 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/history/AgentHistoryThreadSummary.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/history/AgentHistoryThreadSummary.kt @@ -1,6 +1,6 @@ package ee.carlrobert.codegpt.agent.history -import kotlinx.datetime.Instant +import kotlin.time.Instant enum class ThreadStatus { COMPLETED, diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/strategy/SingleRunStrategyProvider.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/strategy/SingleRunStrategyProvider.kt index 420165dc..d52bfe83 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/strategy/SingleRunStrategyProvider.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/strategy/SingleRunStrategyProvider.kt @@ -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): List { - return responses +internal fun appendableResponses( + responses: List, + provider: LLMProvider +): List { + 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() + .filter { it.content.isNotBlank() } + return if (reasoningFallback.isNotEmpty()) reasoningFallback else nonReasoningResponses +} + +internal fun extractAssistantOutput(responses: List): String { + val assistantText = responses + .filterIsInstance() + .map { it.content.trim() } + .filter { it.isNotEmpty() } + .joinToString("\n") + .trim() + if (assistantText.isNotEmpty()) { + return assistantText + } + + val reasoningText = responses + .filterIsInstance() + .map { it.content.trim() } + .filter { it.isNotEmpty() } + .joinToString("\n") + .trim() + return if (reasoningText.isNotEmpty()) "$reasoningText" else "" } private fun publishUsageAndCredits( @@ -236,11 +283,22 @@ private fun publishUsageAndCredits( provider: ServiceType, responses: List, ) { - 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 +): 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 AIAgentEdgeBuilderIntermediate, OutgoingInput>.onSingleAssistantResponse( - block: suspend (Message.Response) -> Boolean + block: suspend (String) -> Boolean ): AIAgentEdgeBuilderIntermediate { return onIsInstance(List::class) .transformed { response -> - response.filter { item -> item is Message.Response && item !is Message.Reasoning } - .filterIsInstance() + val modelResponses = response.filterIsInstance() + 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) } } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/AskUserQuestionTool.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/AskUserQuestionTool.kt index 30c3b5da..ad411b39 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/AskUserQuestionTool.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/AskUserQuestionTool.kt @@ -7,6 +7,7 @@ import kotlinx.serialization.Serializable class AskUserQuestionTool( workingDirectory: String, + sessionId: String, hookManager: HookManager, private val events: AgentEvents, ) : BaseTool( @@ -29,6 +30,7 @@ class AskUserQuestionTool( argsClass = Args::class, resultClass = Result::class, hookManager = hookManager, + sessionId = sessionId ) { @Serializable diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/BaseTool.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/BaseTool.kt index 4d0c7e6f..12d3cf0a 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/BaseTool.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/BaseTool.kt @@ -30,7 +30,7 @@ abstract class BaseTool( description: String, protected val workingDirectory: String, private val hookManager: HookManager, - private val sessionId: String? = null, + private val sessionId: String, private val argsClass: KClass, private val resultClass: KClass ) : Tool( @@ -58,7 +58,7 @@ abstract class BaseTool( * @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, diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/BashOutputTool.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/BashOutputTool.kt index 661e2ecf..0bed29fc 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/BashOutputTool.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/BashOutputTool.kt @@ -14,7 +14,7 @@ import kotlinx.serialization.Serializable class BashOutputTool( workingDirectory: String, hookManager: HookManager, - private val sessionId: String = "global", + private val sessionId: String, ) : BaseTool( workingDirectory = workingDirectory, @@ -33,6 +33,7 @@ class BashOutputTool( argsClass = Args::class, resultClass = Result::class, hookManager = hookManager, + sessionId = sessionId, ) { @Serializable diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/BashTool.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/BashTool.kt index 55812708..9cf426d6 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/BashTool.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/BashTool.kt @@ -174,7 +174,8 @@ class BashTool( """.trimIndent(), argsClass = Args::class, resultClass = Result::class, - hookManager = hookManager + hookManager = hookManager, + sessionId = sessionId, ) { @Serializable diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/ConfirmingTools.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/ConfirmingTools.kt index 925f59c9..4bdee86b 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/ConfirmingTools.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/ConfirmingTools.kt @@ -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( - 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) - } -} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/EditTool.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/EditTool.kt index 1801c17c..3d7abc75 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/EditTool.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/EditTool.kt @@ -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( 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 { diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/GetLibraryDocsTool.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/GetLibraryDocsTool.kt index 361c58bc..c0f3eb4f 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/GetLibraryDocsTool.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/GetLibraryDocsTool.kt @@ -19,6 +19,7 @@ import java.nio.charset.StandardCharsets */ class GetLibraryDocsTool( workingDirectory: String, + sessionId: String, hookManager: HookManager ) : BaseTool( workingDirectory = workingDirectory, @@ -30,7 +31,8 @@ class GetLibraryDocsTool( """.trimIndent(), argsClass = Args::class, resultClass = Result::class, - hookManager = hookManager + hookManager = hookManager, + sessionId = sessionId, ) { @Serializable diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/IntelliJSearchTool.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/IntelliJSearchTool.kt index 02474601..2887dc33 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/IntelliJSearchTool.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/IntelliJSearchTool.kt @@ -33,6 +33,7 @@ import java.nio.file.Paths */ class IntelliJSearchTool( hookManager: HookManager, + sessionId: String, private val project: Project, ) : BaseTool( 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 { diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/KillShellTool.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/KillShellTool.kt index c8a5a8a2..4f77e97d 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/KillShellTool.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/KillShellTool.kt @@ -10,6 +10,7 @@ import kotlinx.serialization.Serializable class KillShellTool( workingDirectory: String, + sessionId: String, hookManager: HookManager, ) : BaseTool( workingDirectory = workingDirectory, @@ -25,7 +26,8 @@ class KillShellTool( """.trimIndent(), argsClass = Args::class, resultClass = Result::class, - hookManager = hookManager + hookManager = hookManager, + sessionId = sessionId ) { @Serializable diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/McpDynamicToolRegistry.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/McpDynamicToolRegistry.kt index 32939f07..50b635c9 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/McpDynamicToolRegistry.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/McpDynamicToolRegistry.kt @@ -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() - val tools = mutableListOf>() - 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 { - 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): List { - return (schema["required"] as? List<*>) - ?.mapNotNull { it?.toString() } - ?: emptyList() -} - -private fun parseParameterDescriptors(schema: Map): List { - val properties = schema["properties"] as? Map<*, *> ?: emptyMap() - return properties.mapNotNull { (nameRaw, propertyRaw) -> - val name = nameRaw?.toString() ?: return@mapNotNull null - val property = propertyRaw as? Map<*, *> ?: emptyMap() - 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() - 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 { - val map = value as? Map<*, *> ?: return emptyMap() - return map.entries.mapNotNull { (key, nestedValue) -> - key?.toString()?.let { it to nestedValue } - }.toMap() -} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/McpTool.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/McpTool.kt index ad3edde5..72a9ff4f 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/McpTool.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/McpTool.kt @@ -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( - 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 = emptyMap() ) @@ -85,253 +50,4 @@ class McpTool( } } } - - override suspend fun doExecute(args: Args): Result { - val contextService = project.service() - 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 { - return runCatching { sessionManager.getSessionAttachments(conversationId) } - .getOrDefault(emptyList()) - } - - private fun resolveTargetAttachment( - args: Args, - context: AgentMcpContext, - selectedAttachments: List - ): 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, - selectedAttachments: List - ): 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? { - val requestedName = serverName.normalizedOrNull() ?: return null - return selectedAttachments.firstOrNull { - it.serverName.equals(requestedName, ignoreCase = true) - } - } - - private fun resolveByToolOrFallback( - toolName: String, - selectedServerIds: Set, - selectedAttachments: List - ): 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() } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/ReadTool.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/ReadTool.kt index b8e20c84..a1b8fb0b 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/ReadTool.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/ReadTool.kt @@ -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( 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 { diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/ResolveLibraryIdTool.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/ResolveLibraryIdTool.kt index dfd0db4c..bdb144e2 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/ResolveLibraryIdTool.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/ResolveLibraryIdTool.kt @@ -17,6 +17,7 @@ import java.nio.charset.StandardCharsets */ class ResolveLibraryIdTool( workingDirectory: String, + sessionId: String, hookManager: HookManager, ) : BaseTool( workingDirectory = workingDirectory, @@ -28,7 +29,8 @@ class ResolveLibraryIdTool( """.trimIndent(), argsClass = Args::class, resultClass = Result::class, - hookManager = hookManager + hookManager = hookManager, + sessionId = sessionId, ) { @Serializable diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/TaskTool.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/TaskTool.kt index 3ba79b21..218c5b46 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/TaskTool.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/TaskTool.kt @@ -40,6 +40,7 @@ class TaskTool( argsClass = Args::class, resultClass = Result::class, hookManager = hookManager, + sessionId = sessionId, ) { @Serializable diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/TodoWriteTool.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/TodoWriteTool.kt index b5483df1..c69b02f5 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/TodoWriteTool.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/TodoWriteTool.kt @@ -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 = diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/WebFetchTool.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/WebFetchTool.kt index 4d77abc1..11a3474e 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/WebFetchTool.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/WebFetchTool.kt @@ -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( 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 diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/WebSearchTool.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/WebSearchTool.kt index 7d1e88bb..1805d393 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/WebSearchTool.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/WebSearchTool.kt @@ -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( 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 diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/WriteTool.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/WriteTool.kt index e3778d20..a56645ac 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/WriteTool.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/WriteTool.kt @@ -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( 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 diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionEventListener.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionEventListener.kt index 428a62fc..c22044d2 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionEventListener.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionEventListener.kt @@ -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, -) : CompletionEventListener { +) : 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().getServiceForFeature(FeatureType.CODE_COMPLETION) == ServiceType.PROXYAI + service().getServiceForFeature(FeatureType.CODE_COMPLETION) == ServiceType.PROXYAI if (isCodeGPTService && "RATE_LIMIT_EXCEEDED" == error.code) { service().state .codeCompletionSettings diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionPromptFactory.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionPromptFactory.kt new file mode 100644 index 00000000..20469e1f --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionPromptFactory.kt @@ -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 "\n${infillRequest.prefix}\n\n\n" + + "\n${infillRequest.suffix}\n\n\n" + + "Complete:" + } +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionRequestFactory.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionRequestFactory.kt deleted file mode 100644 index 411ee6a4..00000000 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionRequestFactory.kt +++ /dev/null @@ -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", - "\n${details.prefix}\n\n\n\n${details.suffix}\n\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, - body: Map, - 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( - "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( - "role" to "user", - "content" to "\n${details.prefix}\n\n\n\n${details.suffix}\n\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() - .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, - body: Map, - 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>() {}) - 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().state - val model = service().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().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 - } - } - } -} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionService.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionService.kt index 580ac06d..0a1d4690 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionService.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionService.kt @@ -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 - ): 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().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() + .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() } + } } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/DebouncedCodeCompletionProvider.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/DebouncedCodeCompletionProvider.kt index fe0eb214..0536861e 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/DebouncedCodeCompletionProvider.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/DebouncedCodeCompletionProvider.kt @@ -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(null) + private val currentCallRef = AtomicReference(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().getServiceForFeature(FeatureType.CODE_COMPLETION) == ServiceType.PROXYAI) { - val grpcClient = project.service() - grpcClient.cancelCodeCompletion() - grpcClient.getCodeCompletionAsync(infillRequest, eventListener, this) - return@channelFlow - } - val call = service().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().getServiceForFeature(FeatureType.CODE_COMPLETION) + service().getServiceForFeature(FeatureType.CODE_COMPLETION) val codeCompletionsEnabled = when (selectedService) { ServiceType.PROXYAI -> service().state.codeCompletionSettings.codeCompletionsEnabled ServiceType.OPENAI -> OpenAISettings.getCurrentState().isCodeCompletionsEnabled diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequest.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequest.kt index c50e2436..1cd4d417 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequest.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequest.kt @@ -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().truncateText(this, maxTokens, fromStart) -} \ No newline at end of file +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/edit/CodeCompletionStreamObserver.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/edit/CodeCompletionStreamObserver.kt index a3e8443c..65a3c0bd 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/edit/CodeCompletionStreamObserver.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/edit/CodeCompletionStreamObserver.kt @@ -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, - private val eventListener: CodeCompletionEventListener, + private val eventListener: CompletionStreamEventListener, ) : StreamObserver { 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() { diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/edit/GrpcClientService.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/edit/GrpcClientService.kt index 393b8ecd..a7c80f47 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/edit/GrpcClientService.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/edit/GrpcClientService.kt @@ -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 - ) { + 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) { diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/CompletionContextService.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/CompletionContextService.kt index 1841fc11..0b173365 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/CompletionContextService.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/CompletionContextService.kt @@ -15,7 +15,6 @@ class CompletionContextService { companion object { private val CONTEXT_FINDERS = mapOf( "JAVA" to JavaContextFinder::class.java, - // "Python" to PythonContextFinder::class.java ) } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/LanguageContextFinder.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/LanguageContextFinder.kt index 1983f45d..24d6b4b1 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/LanguageContextFinder.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/LanguageContextFinder.kt @@ -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 -} \ No newline at end of file +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/PythonContextFinder.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/PythonContextFinder.kt deleted file mode 100644 index ab6e5966..00000000 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/PythonContextFinder.kt +++ /dev/null @@ -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, - rootElement: PsiElement - ): Set = - 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 { - val resultSet = mutableSetOf() - 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()?.assignedValue - } - - private fun findDeclarations( - pyReference: PyReferenceExpression, - project: Project - ): Set { - // 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 - } - - } -} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/completions/AgentCompletionRunner.kt b/src/main/kotlin/ee/carlrobert/codegpt/completions/AgentCompletionRunner.kt new file mode 100644 index 00000000..8d05ccb4 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/completions/AgentCompletionRunner.kt @@ -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() + 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? = null + try { + service = AIAgentService( + promptExecutor = request.executor, + strategy = SingleRunStrategyProvider().build( + project = project, + executor = request.executor, + pendingMessageQueue = ArrayDeque(), + 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("${it.content}") + } + } + } + } + } + } + + 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() + .getStoredModelForFeature(request.callParameters.featureType) + project.service().state.services + .firstOrNull { it.id == selectedServiceId } + ?.chatCompletionSettings + ?.shouldStream() + ?: false + } + + ServiceType.GOOGLE -> false + else -> true + } + } + + private fun createChatToolRegistry( + callParameters: ChatCompletionParameters, + mcpTools: List, + 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 + ) + ) + } + } + } +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/completions/AutoApplyService.kt b/src/main/kotlin/ee/carlrobert/codegpt/completions/AutoApplyService.kt new file mode 100644 index 00000000..a9f1210c --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/completions/AutoApplyService.kt @@ -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().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().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 + } + } +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/completions/BaseChatStreamEventListener.kt b/src/main/kotlin/ee/carlrobert/codegpt/completions/BaseChatStreamEventListener.kt new file mode 100644 index 00000000..71572006 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/completions/BaseChatStreamEventListener.kt @@ -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) + } + +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/completions/ChatCompletionEventListener.kt b/src/main/kotlin/ee/carlrobert/codegpt/completions/ChatCompletionEventListener.kt new file mode 100644 index 00000000..f87251b0 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/completions/ChatCompletionEventListener.kt @@ -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() + } +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/completions/ChatStreamingCompletionRunner.kt b/src/main/kotlin/ee/carlrobert/codegpt/completions/ChatStreamingCompletionRunner.kt new file mode 100644 index 00000000..121b7f52 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/completions/ChatStreamingCompletionRunner.kt @@ -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() + 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) + } + } + } +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/completions/ChatTypes.kt b/src/main/kotlin/ee/carlrobert/codegpt/completions/ChatTypes.kt new file mode 100644 index 00000000..86edac1b --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/completions/ChatTypes.kt @@ -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) +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionError.kt b/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionError.kt new file mode 100644 index 00000000..02bd3643 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionError.kt @@ -0,0 +1,6 @@ +package ee.carlrobert.codegpt.completions + +data class CompletionError @JvmOverloads constructor( + val message: String, + val code: String? = null +) diff --git a/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionParameters.kt b/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionParameters.kt index 405c49a1..4438d2d3 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionParameters.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionParameters.kt @@ -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?, var mcpTools: List?, var toolApprovalMode: ToolApprovalMode, - var toolResults: List>? = null, + var toolResults: List>? = 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? = null private var mcpTools: List? = null private var toolApprovalMode: ToolApprovalMode = ToolApprovalMode.AUTO_APPROVE - private var toolResults: List>? = null + private var toolResults: List>? = 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?) = apply { this.mcpTools = tools } fun toolApprovalMode(mode: ToolApprovalMode) = apply { this.toolApprovalMode = mode } - fun toolResults(results: List>?) = + fun toolResults(results: List>?) = apply { this.toolResults = results } fun project(project: Project?) = apply { this.project = project } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionRequest.kt b/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionRequest.kt new file mode 100644 index 00000000..ea3246ee --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionRequest.kt @@ -0,0 +1,3 @@ +package ee.carlrobert.codegpt.completions + +interface CompletionRequest diff --git a/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionRequestFactory.kt b/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionRequestFactory.kt index bef5887c..cd0dbe29 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionRequestFactory.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionRequestFactory.kt @@ -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().state + val filteredPrompts = service() + val systemParts = mutableListOf() + + 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() 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().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 { diff --git a/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionRequestService.kt b/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionRequestService.kt new file mode 100644 index 00000000..c8149cd3 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionRequestService.kt @@ -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 = 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 = 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) + } +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionRunner.kt b/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionRunner.kt new file mode 100644 index 00000000..2aa69103 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionRunner.kt @@ -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 = 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 +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionRunnerFactory.kt b/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionRunnerFactory.kt new file mode 100644 index 00000000..58978f0b --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionRunnerFactory.kt @@ -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 + } +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionService.kt b/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionService.kt new file mode 100644 index 00000000..7894df77 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionService.kt @@ -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().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().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().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 + } + } + } +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionStreamEventListener.kt b/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionStreamEventListener.kt new file mode 100644 index 00000000..d77cada6 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionStreamEventListener.kt @@ -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) +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionTextExtractor.kt b/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionTextExtractor.kt new file mode 100644 index 00000000..601e4751 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionTextExtractor.kt @@ -0,0 +1,33 @@ +package ee.carlrobert.codegpt.completions + +import ai.koog.prompt.message.Message + +internal object CompletionTextExtractor { + fun extract(responses: List): String { + val reasoningText = responses + .filterIsInstance() + .map { it.content.trim() } + .filter { it.isNotEmpty() } + .joinToString("\n") + .trim() + + val assistantText = responses + .filterIsInstance() + .joinToString("\n") { it.content.trim() } + .trim() + + if (reasoningText.isNotEmpty() && assistantText.isNotEmpty()) { + return "$reasoningText\n$assistantText" + } + + if (assistantText.isNotEmpty()) { + return assistantText + } + + if (reasoningText.isNotEmpty()) { + return "$reasoningText" + } + + return responses.joinToString("\n") { it.content.trim() }.trim() + } +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/completions/McpToolAdapter.kt b/src/main/kotlin/ee/carlrobert/codegpt/completions/McpToolAdapter.kt new file mode 100644 index 00000000..9fd1231c --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/completions/McpToolAdapter.kt @@ -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( + 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 = + suspendCancellableCoroutine { continuation -> + future.whenComplete { value, throwable -> + if (throwable != null) { + continuation.resumeWithException(throwable) + } else { + continuation.resume(value) + } + } + continuation.invokeOnCancellation { future.cancel(true) } + } +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/completions/StreamingCompletionRunner.kt b/src/main/kotlin/ee/carlrobert/codegpt/completions/StreamingCompletionRunner.kt new file mode 100644 index 00000000..4ef0af6f --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/completions/StreamingCompletionRunner.kt @@ -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() + 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 + } +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/completions/autoapply/AutoApplyModels.kt b/src/main/kotlin/ee/carlrobert/codegpt/completions/autoapply/AutoApplyModels.kt new file mode 100644 index 00000000..ef119a68 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/completions/autoapply/AutoApplyModels.kt @@ -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 = "", +) diff --git a/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/ClaudeRequestFactory.kt b/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/ClaudeRequestFactory.kt index d43b7001..dbc42e20 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/ClaudeRequestFactory.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/ClaudeRequestFactory.kt @@ -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().state.maxTokens - isStream = true - - var systemPrompt = "" - val selectedPersona = service().state.personas.selectedPersona - if (!selectedPersona.disabled) { - val base = - service().getFilteredPersonaPrompt(params.chatMode) - systemPrompt = service().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(ClaudeCompletionStandardMessage("user", userPrompt)) - this.maxTokens = maxTokens - } - } - - private fun buildClaudeMessages(params: ChatCompletionParameters): MutableList { - val messages = mutableListOf() - - 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): Map { - if (mcpSchema.containsKey("type")) { - return mcpSchema - } - - return mapOf( - "type" to "object", - "properties" to mcpSchema, - "required" to mcpSchema.keys.toList() - ) - } -} +class ClaudeRequestFactory : BaseRequestFactory() diff --git a/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/CodeGPTRequestFactory.kt b/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/CodeGPTRequestFactory.kt index c6e5cc1e..cb653e6d 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/CodeGPTRequestFactory.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/CodeGPTRequestFactory.kt @@ -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().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().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() - 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) { - 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().state.coreActions.inlineEdit.instructions - ?: CoreActionsState.DEFAULT_INLINE_EDIT_PROMPT - - val contextFiles = mutableListOf() - 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().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): ToolFunctionParameters { - return ToolFunctionParameters().apply { - if (mcpSchema.containsKey("type")) { - type = mcpSchema["type"] as? String ?: "object" - properties = mcpSchema["properties"] as? Map ?: mcpSchema - required = mcpSchema["required"] as? List ?: emptyList() - } else { - type = "object" - properties = mcpSchema - required = mcpSchema.keys.toList() - } - } - } - - private fun buildCodeGPTMessages( - model: String?, - params: ChatCompletionParameters, - referencedFiles: List?, - psiStructure: Set? - ): List { - val messages = - buildOpenAIMessages(model, params, referencedFiles, emptyList(), psiStructure) - - /*if (params.toolResults.isNullOrEmpty()) { - return messages - } - - val convertedMessages = mutableListOf() - var pendingToolResults = mutableListOf>() - - 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() diff --git a/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/CustomOpenAIRequestFactory.kt b/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/CustomOpenAIRequestFactory.kt index f98bb618..1dde4a2f 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/CustomOpenAIRequestFactory.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/CustomOpenAIRequestFactory.kt @@ -1,210 +1,5 @@ package ee.carlrobert.codegpt.completions.factory -import com.fasterxml.jackson.databind.ObjectMapper -import com.intellij.openapi.components.service 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.credentials.CredentialsStore.CredentialKey -import ee.carlrobert.codegpt.credentials.CredentialsStore.getCredential -import ee.carlrobert.codegpt.mcp.McpToolConverter -import ee.carlrobert.codegpt.settings.service.FeatureType -import ee.carlrobert.codegpt.settings.service.custom.CustomServiceChatCompletionSettingsState -import ee.carlrobert.codegpt.settings.service.custom.CustomServicePlaceholders -import ee.carlrobert.codegpt.settings.service.custom.CustomServicesSettings -import ee.carlrobert.llm.client.openai.completion.request.OpenAIChatCompletionMessage -import ee.carlrobert.llm.client.openai.completion.request.OpenAIChatCompletionStandardMessage -import ee.carlrobert.llm.completion.CompletionRequest -import okhttp3.Request -import okhttp3.RequestBody.Companion.toRequestBody -import java.nio.charset.StandardCharsets -class CustomOpenAIRequest(val request: Request) : CompletionRequest - -class CustomOpenAIRequestFactory : BaseRequestFactory() { - - override fun createChatRequest(params: ChatCompletionParameters): CustomOpenAIRequest { - val serviceState = - service().customServiceStateForFeatureType(params.featureType) - val messages = OpenAIRequestFactory.buildOpenAIMessages( - null, - params, - params.referencedFiles, - params.history, - params.psiStructure - ) - val request = buildCustomOpenAIChatCompletionRequest( - serviceState.chatCompletionSettings, - messages, - true, - getCredential(CredentialKey.CustomServiceApiKeyById(requireNotNull(serviceState.id))), - params - ) - return CustomOpenAIRequest(request) - } - - override fun createBasicCompletionRequest( - systemPrompt: String, - userPrompt: String, - maxTokens: Int, - stream: Boolean, - featureType: FeatureType - ): CompletionRequest { - val service = - service().customServiceStateForFeatureType(featureType) - - val request = buildCustomOpenAIChatCompletionRequest( - service.chatCompletionSettings, - listOf( - OpenAIChatCompletionStandardMessage("system", systemPrompt), - OpenAIChatCompletionStandardMessage("user", userPrompt) - ), - stream, - getCredential(CredentialKey.CustomServiceApiKeyById(requireNotNull(service.id))) - ) - return CustomOpenAIRequest(request) - } - - override fun createInlineEditRequest(params: InlineEditCompletionParameters): CompletionRequest { - val service = - service().customServiceStateForFeatureType(FeatureType.INLINE_EDIT) - val systemPrompt = prepareInlineEditSystemPrompt(params) - val messages = - OpenAIRequestFactory.buildInlineEditMessages(systemPrompt, params.conversation) - val request = buildCustomOpenAIChatCompletionRequest( - service.chatCompletionSettings, - messages, - true, - getCredential(CredentialKey.CustomServiceApiKeyById(requireNotNull(service.id))) - ) - return CustomOpenAIRequest(request) - } - - companion object { - fun buildCustomOpenAICompletionRequest( - context: String, - url: String, - headers: MutableMap, - body: MutableMap, - credential: String? - ): Request { - val usedSettings = CustomServiceChatCompletionSettingsState() - usedSettings.body = body - usedSettings.headers = headers - usedSettings.url = url - return buildCustomOpenAIChatCompletionRequest( - usedSettings, - listOf(OpenAIChatCompletionStandardMessage("user", context)), - true, - credential, - null - ) - } - - fun buildCustomOpenAIChatCompletionRequest( - settings: CustomServiceChatCompletionSettingsState, - messages: List, - streamRequest: Boolean, - credential: String?, - params: ChatCompletionParameters? = null - ): Request { - val requestBuilder = Request.Builder().url(requireNotNull(settings.url).trim()) - - settings.headers.forEach { (key, value) -> - val headerValue = when { - credential != null && value.contains($$"$CUSTOM_SERVICE_API_KEY") -> - value.replace($$"$CUSTOM_SERVICE_API_KEY", credential) - - else -> value - } - requestBuilder.addHeader(key, headerValue) - } - - val body = settings.body - .mapValues { (key, value) -> - transformBodyValue( - key = key, - value = value, - streamRequest = streamRequest, - messages = messages, - credential = credential - ) - } - .toMutableMap() - .apply { - if (params != null && !params.mcpTools.isNullOrEmpty() && params.toolApprovalMode != ToolApprovalMode.BLOCK_ALL) { - val openAITools = params.mcpTools!!.map { mcpTool -> - McpToolConverter.convertToOpenAITool(mcpTool) - } - - put("tools", openAITools) - } - } - - return try { - val requestBodyString = ObjectMapper().writerWithDefaultPrettyPrinter() - .writeValueAsString(body) - val requestBody = requestBodyString - .toByteArray(StandardCharsets.UTF_8) - .toRequestBody() - - requestBuilder.post(requestBody).build() - } catch (e: Exception) { - throw RuntimeException("Failed to build CustomOpenAI request", e) - } - } - - private fun transformBodyValue( - key: String?, - value: Any?, - streamRequest: Boolean, - messages: List, - credential: String? - ): Any? { - return when (value) { - null -> null - is String -> when { - !streamRequest && key == "stream" -> false - CustomServicePlaceholders.isMessages(value) -> messages - CustomServicePlaceholders.isPrompt(value) -> renderPrompt(messages) - credential != null && value.contains($$"$CUSTOM_SERVICE_API_KEY") -> - value.replace($$"$CUSTOM_SERVICE_API_KEY", credential) - - else -> value - } - - is Map<*, *> -> value.entries.associate { (nestedKey, nestedValue) -> - nestedKey.toString() to transformBodyValue( - key = nestedKey?.toString(), - value = nestedValue, - streamRequest = streamRequest, - messages = messages, - credential = credential - ) - } - - is List<*> -> value.map { item -> - transformBodyValue( - key = null, - value = item, - streamRequest = streamRequest, - messages = messages, - credential = credential - ) - } - - else -> value - } - } - - private fun renderPrompt(messages: List): String { - return messages.joinToString(separator = "\n\n") { message -> - when (message) { - is OpenAIChatCompletionStandardMessage -> message.content - else -> ObjectMapper().writeValueAsString(message) - } - } - } - } -} +class CustomOpenAIRequestFactory : BaseRequestFactory() diff --git a/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/GoogleRequestFactory.kt b/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/GoogleRequestFactory.kt index cae415dd..fa9f18ef 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/GoogleRequestFactory.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/GoogleRequestFactory.kt @@ -1,209 +1,5 @@ package ee.carlrobert.codegpt.completions.factory -import com.intellij.openapi.components.service -import ee.carlrobert.codegpt.EncodingManager import ee.carlrobert.codegpt.completions.BaseRequestFactory -import ee.carlrobert.codegpt.completions.ChatCompletionParameters -import ee.carlrobert.codegpt.completions.ConversationType -import ee.carlrobert.codegpt.completions.TotalUsageExceededException -import ee.carlrobert.codegpt.conversations.ConversationsState -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.codegpt.util.file.FileUtil -import ee.carlrobert.llm.client.google.completion.GoogleCompletionContent -import ee.carlrobert.llm.client.google.completion.GoogleCompletionRequest -import ee.carlrobert.llm.client.google.completion.GoogleContentPart -import ee.carlrobert.llm.client.google.completion.GoogleGenerationConfig -import ee.carlrobert.llm.client.google.models.GoogleModel -import java.io.IOException -import java.nio.file.Files -import java.nio.file.Path -class GoogleRequestFactory : BaseRequestFactory() { - - override fun createChatRequest(params: ChatCompletionParameters): GoogleCompletionRequest { - val configuration = service().state - val selectedModel = ModelSelectionService.getInstance().getModelForFeature(FeatureType.CHAT) - val messages = buildGoogleMessages(selectedModel, params) - return GoogleCompletionRequest.Builder(messages) - .generationConfig( - GoogleGenerationConfig.Builder() - .maxOutputTokens(configuration.maxTokens) - .temperature(configuration.temperature.toDouble()).build() - ) - .systemInstruction(buildSystemInstruction(params)) - .build() - } - - override fun createBasicCompletionRequest( - systemPrompt: String, - userPrompt: String, - maxTokens: Int, - stream: Boolean, - featureType: FeatureType - ): GoogleCompletionRequest { - val configuration = service().state - return GoogleCompletionRequest.Builder( - listOf( - GoogleCompletionContent("user", listOf(systemPrompt)), - GoogleCompletionContent("model", listOf("Understood.")), - GoogleCompletionContent("user", listOf(userPrompt)) - ) - ) - .generationConfig( - GoogleGenerationConfig.Builder() - .maxOutputTokens(maxTokens) - .temperature(configuration.temperature.toDouble()).build() - ) - .build() - } - - private fun buildGoogleMessages( - model: String?, - params: ChatCompletionParameters - ): List { - val messages = buildGoogleMessages(params) - - if (model == null) { - return messages - } - - val encodingManager = service() - val totalUsage = messages.parallelStream() - .mapToInt { message -> - encodingManager.countMessageTokens( - message.role, - message.parts.joinToString(",") { it.text ?: "" } - ) - } - .sum() + service().state.maxTokens - - return GoogleModel.findByCode(model)?.let { googleModel -> - if (totalUsage <= googleModel.maxTokens) { - messages - } else { - tryReducingGoogleMessagesOrThrow( - messages, - params.conversation.isDiscardTokenLimit, - totalUsage, - googleModel.maxTokens - ) - } - } ?: messages - } - - private fun buildGoogleMessages(params: ChatCompletionParameters): List { - val message = params.message - val messages = mutableListOf() - - for (prevMessage in params.conversation.messages) { - if (params.retry && prevMessage.id == message.id) { - break - } - - prevMessage.imageFilePath?.takeIf { it.isNotEmpty() }?.let { imagePath -> - try { - val imageData = Files.readAllBytes(Path.of(imagePath)) - val imageMediaType = - FileUtil.getImageMediaType(Path.of(imagePath).fileName.toString()) - messages.add( - GoogleCompletionContent( - listOf( - GoogleContentPart( - null, - GoogleContentPart.Blob(imageMediaType, imageData) - ), - GoogleContentPart(prevMessage.prompt) - ), "user" - ) - ) - } catch (e: IOException) { - throw RuntimeException(e) - } - } ?: messages.add(GoogleCompletionContent("user", listOf(prevMessage.prompt))) - - messages.add(GoogleCompletionContent("model", listOf(prevMessage.response))) - } - - val imageDetails = params.imageDetails - if (imageDetails != null) { - messages.add( - GoogleCompletionContent( - listOf( - GoogleContentPart( - null, - GoogleContentPart.Blob( - imageDetails.mediaType, - imageDetails.data - ) - ), - GoogleContentPart(message.prompt) - ), "user" - ) - ) - } else { - messages.add( - GoogleCompletionContent( - "user", - listOf(getPromptWithFilesContext(params)) - ) - ) - } - - return messages - } - - private fun tryReducingGoogleMessagesOrThrow( - messages: List, - discardTokenLimit: Boolean, - totalUsage: Int, - modelMaxTokens: Int - ): List { - if (!service().state!!.discardAllTokenLimits) { - if (!discardTokenLimit) { - throw TotalUsageExceededException() - } - } - - val encodingManager = EncodingManager.getInstance() - var currentUsage = totalUsage - - // skip the system prompt - val updatedMessages = messages.mapIndexed { index, message -> - if (index == 0 || currentUsage <= modelMaxTokens) { - message - } else { - currentUsage -= encodingManager.countMessageTokens( - message.role, - message.parts.joinToString(",") { it.text } - ) - null - } - } - - return updatedMessages.filterNotNull() - } - - private fun buildSystemInstruction(params: ChatCompletionParameters): String? { - return when (params.conversationType) { - ConversationType.DEFAULT -> { - val selectedPersona = service().state.personas.selectedPersona - if (!selectedPersona.disabled) { - val base = service().getFilteredPersonaPrompt(params.chatMode) - service().applyClickableLinks(base) - } else null - } - - ConversationType.FIX_COMPILE_ERRORS -> service() - .applyClickableLinks(service().state.coreActions.fixCompileErrors.instructions.orEmpty()) - - ConversationType.REVIEW_CHANGES -> service() - .applyClickableLinks(service().state.coreActions.reviewChanges.instructions.orEmpty()) - - else -> null - } - } -} +class GoogleRequestFactory : BaseRequestFactory() diff --git a/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/InceptionRequestFactory.kt b/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/InceptionRequestFactory.kt index 759c85b8..7b472240 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/InceptionRequestFactory.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/InceptionRequestFactory.kt @@ -1,114 +1,21 @@ package ee.carlrobert.codegpt.completions.factory -import com.intellij.openapi.components.service import ee.carlrobert.codegpt.completions.* -import ee.carlrobert.codegpt.settings.prompts.CoreActionsState -import ee.carlrobert.codegpt.settings.prompts.PromptsSettings +import ee.carlrobert.codegpt.completions.inception.InceptionChatMessage +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.settings.service.ModelSelectionService import ee.carlrobert.codegpt.util.EditWindowFormatter.FormatResult -import ee.carlrobert.codegpt.util.EditorUtil -import ee.carlrobert.llm.client.inception.request.InceptionApplyRequest -import ee.carlrobert.llm.client.inception.request.InceptionNextEditRequest -import ee.carlrobert.llm.client.openai.completion.request.OpenAIChatCompletionRequest -import ee.carlrobert.llm.client.openai.completion.request.OpenAIChatCompletionStandardMessage -import ee.carlrobert.llm.completion.CompletionRequest class InceptionRequestFactory : BaseRequestFactory() { - override fun createChatRequest(params: ChatCompletionParameters): OpenAIChatCompletionRequest { - val model = ModelSelectionService.getInstance().getModelForFeature(FeatureType.CHAT) - val messages = OpenAIRequestFactory.buildOpenAIMessages( - model = model, - callParameters = params - ) - return OpenAIChatCompletionRequest.Builder(messages) - .setModel(model) - .setStream(true) - .build() - } - - override fun createInlineEditRequest(params: InlineEditCompletionParameters): OpenAIChatCompletionRequest { - val model = ModelSelectionService.getInstance().getModelForFeature(FeatureType.INLINE_EDIT) - val systemPrompt = prepareInlineEditSystemPrompt(params) - val messages = - OpenAIRequestFactory.buildInlineEditMessages(systemPrompt, params.conversation) - return OpenAIChatCompletionRequest.Builder(messages) - .setModel(model) - .setStream(true) - .build() - } - - override fun createAutoApplyRequest(params: AutoApplyParameters): InceptionApplyRequest { - val model = ModelSelectionService.getInstance().getModelForFeature(FeatureType.AUTO_APPLY) - val prompt = - "<|original_code|>\n" + EditorUtil.getFileContent(params.destination) + "\n<|/original_code|>\n\n<|update_snippet|>\n" + params.source + "\n<|/update_snippet|>" - - return InceptionApplyRequest.Builder() - .setModel(model) - .setMessages(listOf(OpenAIChatCompletionStandardMessage("user", prompt))) - .build() - } - override fun createNextEditRequest(params: NextEditParameters, formatResult: FormatResult): InceptionNextEditRequest { - val model = ModelSelectionService.getInstance().getModelForFeature(FeatureType.NEXT_EDIT) + val model = ModelSettings.getInstance().getModelForFeature(FeatureType.NEXT_EDIT) val content = composeNextEditMessage(params, formatResult) - val message = OpenAIChatCompletionStandardMessage("user", content) + val message = InceptionChatMessage("user", content) return InceptionNextEditRequest.Builder() .setModel(model) .setMessages(listOf(message)) .build() } - - override fun createBasicCompletionRequest( - systemPrompt: String, - userPrompt: String, - maxTokens: Int, - stream: Boolean, - featureType: FeatureType - ): CompletionRequest { - val model = ModelSelectionService.getInstance().getModelForFeature(featureType) - return OpenAIChatCompletionRequest.Builder( - listOf( - OpenAIChatCompletionStandardMessage("system", systemPrompt), - OpenAIChatCompletionStandardMessage("user", userPrompt) - ) - ) - .setModel(model) - .setStream(stream) - .build() - } - - override fun createCommitMessageRequest(params: CommitMessageCompletionParameters): OpenAIChatCompletionRequest { - val model = - ModelSelectionService.getInstance().getModelForFeature(FeatureType.COMMIT_MESSAGE) - val (gitDiff, systemPrompt) = params - return OpenAIChatCompletionRequest.Builder( - listOf( - OpenAIChatCompletionStandardMessage("system", systemPrompt), - OpenAIChatCompletionStandardMessage("user", gitDiff) - ) - ) - .setModel(model) - .setStream(true) - .build() - } - - override fun createLookupRequest(params: LookupCompletionParameters): OpenAIChatCompletionRequest { - val model = ModelSelectionService.getInstance().getModelForFeature(FeatureType.LOOKUP) - val (prompt) = params - val systemPrompt = - service().state.coreActions.generateNameLookups.instructions - ?: CoreActionsState.DEFAULT_GENERATE_NAME_LOOKUPS_PROMPT - - return OpenAIChatCompletionRequest.Builder( - listOf( - OpenAIChatCompletionStandardMessage("system", systemPrompt), - OpenAIChatCompletionStandardMessage("user", prompt) - ) - ) - .setModel(model) - .setStream(false) - .build() - } } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/LlamaRequestFactory.kt b/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/LlamaRequestFactory.kt index 92a332d2..178038ed 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/LlamaRequestFactory.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/LlamaRequestFactory.kt @@ -1,64 +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.InlineEditCompletionParameters -import ee.carlrobert.codegpt.settings.configuration.ConfigurationSettings -import ee.carlrobert.codegpt.settings.service.FeatureType -import ee.carlrobert.codegpt.settings.service.ModelSelectionService -import ee.carlrobert.llm.client.openai.completion.request.OpenAIChatCompletionRequest -import ee.carlrobert.llm.client.openai.completion.request.OpenAIChatCompletionStandardMessage -class LlamaRequestFactory : BaseRequestFactory() { - - override fun createChatRequest(params: ChatCompletionParameters): OpenAIChatCompletionRequest { - val model = ModelSelectionService.getInstance().getModelForFeature(FeatureType.CHAT) - val configuration = service().state - return OpenAIChatCompletionRequest.Builder( - OpenAIRequestFactory.buildOpenAIMessages(model, params) - ) - .setModel(model) - .setStream(true) - .setMaxTokens(null) - .setMaxCompletionTokens(configuration.maxTokens) - .setTemperature(configuration.temperature.toDouble()) - .build() - } - - override fun createBasicCompletionRequest( - systemPrompt: String, - userPrompt: String, - maxTokens: Int, - stream: Boolean, - featureType: FeatureType - ): OpenAIChatCompletionRequest { - val model = ModelSelectionService.getInstance().getModelForFeature(featureType) - val configuration = service().state - return OpenAIChatCompletionRequest.Builder( - listOf( - OpenAIChatCompletionStandardMessage("system", systemPrompt), - OpenAIChatCompletionStandardMessage("user", userPrompt) - ) - ) - .setModel(model) - .setStream(stream) - .setMaxTokens(null) - .setMaxCompletionTokens(maxTokens) - .setTemperature(configuration.temperature.toDouble()) - .build() - } - - override fun createInlineEditRequest(params: InlineEditCompletionParameters): OpenAIChatCompletionRequest { - val model = ModelSelectionService.getInstance().getModelForFeature(FeatureType.INLINE_EDIT) - val systemPrompt = prepareInlineEditSystemPrompt(params) - val messages = - OpenAIRequestFactory.buildInlineEditMessages(systemPrompt, params.conversation) - val configuration = service().state - return OpenAIChatCompletionRequest.Builder(messages) - .setModel(model) - .setStream(true) - .setTemperature(configuration.temperature.toDouble()) - .build() - } -} +class LlamaRequestFactory : BaseRequestFactory() diff --git a/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/MistralRequestFactory.kt b/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/MistralRequestFactory.kt index 16097588..7cbebca4 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/MistralRequestFactory.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/MistralRequestFactory.kt @@ -1,50 +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.factory.OpenAIRequestFactory.Companion.buildOpenAIMessages -import ee.carlrobert.codegpt.settings.configuration.ConfigurationSettings -import ee.carlrobert.codegpt.settings.service.FeatureType -import ee.carlrobert.codegpt.settings.service.ModelSelectionService -import ee.carlrobert.llm.client.openai.completion.request.OpenAIChatCompletionRequest -class MistralRequestFactory : BaseRequestFactory() { - - override fun createChatRequest(params: ChatCompletionParameters): OpenAIChatCompletionRequest { - val model = ModelSelectionService.getInstance().getModelForFeature(FeatureType.CHAT) - val configuration = service().state - - return OpenAIChatCompletionRequest.Builder( - buildOpenAIMessages( - model = model, - callParameters = params, - referencedFiles = params.referencedFiles, - conversationsHistory = params.history, - psiStructure = params.psiStructure, - ) - ) - .setModel(model) - .setMaxTokens(configuration.maxTokens) - .setMaxCompletionTokens(null) - .setStream(true) - .setTemperature(configuration.temperature.toDouble()) - .build() - } - - override fun createBasicCompletionRequest( - systemPrompt: String, - userPrompt: String, - maxTokens: Int, - stream: Boolean, - featureType: FeatureType - ): OpenAIChatCompletionRequest { - val model = ModelSelectionService.getInstance().getModelForFeature(featureType) - return OpenAIRequestFactory.createBasicCompletionRequest( - systemPrompt, - userPrompt, - model = model, - isStream = stream - ) - } -} \ No newline at end of file +class MistralRequestFactory : BaseRequestFactory() diff --git a/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/OllamaRequestFactory.kt b/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/OllamaRequestFactory.kt index eb4a6539..13fb5f6b 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/OllamaRequestFactory.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/OllamaRequestFactory.kt @@ -1,51 +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.factory.OpenAIRequestFactory.Companion.buildOpenAIMessages -import ee.carlrobert.codegpt.settings.configuration.ConfigurationSettings -import ee.carlrobert.codegpt.settings.models.ModelSettings -import ee.carlrobert.codegpt.settings.service.FeatureType -import ee.carlrobert.codegpt.settings.service.ModelSelectionService -import ee.carlrobert.llm.client.openai.completion.request.OpenAIChatCompletionRequest -import ee.carlrobert.llm.completion.CompletionRequest -class OllamaRequestFactory : BaseRequestFactory() { - - override fun createChatRequest(params: ChatCompletionParameters): CompletionRequest { - val model = ModelSelectionService.getInstance().getModelForFeature(FeatureType.CHAT) - val configuration = service().state - - return OpenAIChatCompletionRequest.Builder( - buildOpenAIMessages( - model = model, - callParameters = params, - referencedFiles = params.referencedFiles, - conversationsHistory = params.history, - psiStructure = params.psiStructure, - ) - ) - .setModel(model) - .setMaxCompletionTokens(configuration.maxTokens) - .setStream(true) - .setTemperature(configuration.temperature.toDouble()) - .build() - } - - override fun createBasicCompletionRequest( - systemPrompt: String, - userPrompt: String, - maxTokens: Int, - stream: Boolean, - featureType: FeatureType - ): CompletionRequest { - val model = ModelSelectionService.getInstance().getModelForFeature(featureType) - return OpenAIRequestFactory.createBasicCompletionRequest( - systemPrompt, - userPrompt, - model = model, - isStream = stream - ) - } -} +class OllamaRequestFactory : BaseRequestFactory() diff --git a/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/OpenAIRequestFactory.kt b/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/OpenAIRequestFactory.kt index 41423b88..aa8b7c15 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/OpenAIRequestFactory.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/OpenAIRequestFactory.kt @@ -1,611 +1,5 @@ package ee.carlrobert.codegpt.completions.factory -import com.intellij.openapi.components.service -import com.intellij.openapi.vfs.readText -import ee.carlrobert.codegpt.EncodingManager -import ee.carlrobert.codegpt.ReferencedFile -import ee.carlrobert.codegpt.completions.* -import ee.carlrobert.codegpt.conversations.Conversation -import ee.carlrobert.codegpt.conversations.ConversationsState -import ee.carlrobert.codegpt.mcp.McpToolConverter -import ee.carlrobert.codegpt.mcp.McpToolPromptFormatter -import ee.carlrobert.codegpt.psistructure.models.ClassStructure -import ee.carlrobert.codegpt.settings.configuration.ConfigurationSettings -import ee.carlrobert.codegpt.settings.configuration.ConfigurationSettings.Companion.getState -import ee.carlrobert.codegpt.settings.prompts.CoreActionsState -import ee.carlrobert.codegpt.settings.prompts.FilteredPromptsService -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.ui.textarea.ConversationTagProcessor -import ee.carlrobert.codegpt.util.file.FileUtil.getImageMediaType -import ee.carlrobert.llm.client.openai.completion.OpenAIChatCompletionModel.* -import ee.carlrobert.llm.client.openai.completion.request.* -import ee.carlrobert.llm.completion.CompletionRequest +import ee.carlrobert.codegpt.completions.BaseRequestFactory -class OpenAIRequestFactory : BaseRequestFactory() { - - override fun createChatRequest(params: ChatCompletionParameters): OpenAIChatCompletionRequest { - val model = ModelSelectionService.getInstance().getModelForFeature(params.featureType) - val configuration = service().state - val requestBuilder: OpenAIChatCompletionRequest.Builder = - OpenAIChatCompletionRequest.Builder(buildOpenAIChatMessages(model, params)) - .setModel(model) - .setStream(true) - .setMaxTokens(null) - .setMaxCompletionTokens(configuration.maxTokens) - if (isReasoningModel(model)) { - requestBuilder - .setTemperature(null) - .setPresencePenalty(null) - .setFrequencyPenalty(null) - } else { - requestBuilder.setTemperature(configuration.temperature.toDouble()) - } - if (!params.mcpTools.isNullOrEmpty() && params.toolApprovalMode != ToolApprovalMode.BLOCK_ALL) { - val openAITools = params.mcpTools!!.map { McpToolConverter.convertToOpenAITool(it) } - requestBuilder.setTools(openAITools) - - when (params.toolApprovalMode) { - ToolApprovalMode.AUTO_APPROVE -> { - requestBuilder.setToolChoice("auto") - } - - ToolApprovalMode.REQUIRE_APPROVAL -> { - requestBuilder.setToolChoice("auto") - } - - else -> {} - } - } - return requestBuilder.build() - } - - override fun createInlineEditRequest(params: InlineEditCompletionParameters): OpenAIChatCompletionRequest { - val model = ModelSelectionService.getInstance().getModelForFeature(FeatureType.INLINE_EDIT) - val systemPrompt = prepareInlineEditSystemPrompt(params) - val messages = buildInlineEditMessages(systemPrompt, params.conversation) - - val configuration = service().state - return if (isReasoningModel(model)) { - val collapsed = messages.joinToString("\n\n") { msg -> - when (msg) { - is OpenAIChatCompletionStandardMessage -> msg.content - else -> "" - } - } - buildBasicO1Request(model, collapsed, systemPrompt = "", stream = true) - } else { - OpenAIChatCompletionRequest.Builder(messages) - .setModel(model) - .setStream(true) - .setTemperature(configuration.temperature.toDouble()) - .build() - } - } - - override fun createAutoApplyRequest(params: AutoApplyParameters): CompletionRequest { - val model = ModelSelectionService.getInstance().getModelForFeature(FeatureType.AUTO_APPLY) - val systemPrompt = - service().getFilteredAutoApplyPrompt(params.chatMode) - .replace("{{changes_to_merge}}", CompletionRequestUtil.formatCode(params.source)) - .replace( - "{{destination_file}}", - CompletionRequestUtil.formatCode( - params.destination.readText(), - params.destination.path - ) - ) - val prompt = "Merge the following changes to the destination file." - - if (isReasoningModel(model)) { - return buildBasicO1Request(model, prompt, systemPrompt, stream = true) - } - return createBasicCompletionRequest(systemPrompt, prompt, model, true) - } - - override fun createBasicCompletionRequest( - systemPrompt: String, - userPrompt: String, - maxTokens: Int, - stream: Boolean, - featureType: FeatureType - ): CompletionRequest { - val model = ModelSelectionService.getInstance().getModelForFeature(featureType) - return if (isReasoningModel(model)) { - buildBasicO1Request( - model, - userPrompt, - systemPrompt, - maxCompletionTokens = maxTokens, - stream = stream - ) - } else { - OpenAIChatCompletionRequest.Builder( - listOf( - OpenAIChatCompletionStandardMessage("system", systemPrompt), - OpenAIChatCompletionStandardMessage("user", userPrompt) - ) - ) - .setModel(model) - .setStream(stream) - .build() - } - } - - override fun createCommitMessageRequest(params: CommitMessageCompletionParameters): OpenAIChatCompletionRequest { - val model = - ModelSelectionService.getInstance().getModelForFeature(FeatureType.COMMIT_MESSAGE) - val (gitDiff, systemPrompt) = params - if (isReasoningModel(model)) { - return buildBasicO1Request(model, gitDiff, systemPrompt, stream = true) - } - return createBasicCompletionRequest(systemPrompt, gitDiff, model, true) - } - - override fun createLookupRequest(params: LookupCompletionParameters): OpenAIChatCompletionRequest { - val model = ModelSelectionService.getInstance().getModelForFeature(FeatureType.LOOKUP) - val (prompt) = params - if (isReasoningModel(model)) { - return buildBasicO1Request( - model, - prompt, - service().state.coreActions.generateNameLookups.instructions - ?: CoreActionsState.DEFAULT_GENERATE_NAME_LOOKUPS_PROMPT - ) - } - return createBasicCompletionRequest( - service().state.coreActions.generateNameLookups.instructions - ?: CoreActionsState.DEFAULT_GENERATE_NAME_LOOKUPS_PROMPT, prompt, model - ) - } - - companion object { - fun buildInlineEditMessages( - systemPrompt: String, - conversation: Conversation? - ): MutableList { - val messages = mutableListOf() - messages.add(OpenAIChatCompletionStandardMessage("system", systemPrompt)) - conversation?.messages?.forEach { m -> - val p = m.prompt?.trim().orEmpty() - if (p.isNotEmpty()) messages.add(OpenAIChatCompletionStandardMessage("user", p)) - val r = m.response?.trim().orEmpty() - if (r.isNotEmpty()) messages.add( - OpenAIChatCompletionStandardMessage( - "assistant", - r - ) - ) - } - messages.add(OpenAIChatCompletionStandardMessage("user", "Implement.")) - return messages - } - - fun isReasoningModel(model: String?) = - listOf( - O_4_MINI.code, - O_3.code, - O_3_MINI.code, - O_1_MINI.code, - O_1_PREVIEW.code, - GPT_5_MINI.code, - GPT_5.code, - ).contains(model) - - fun buildBasicO1Request( - model: String, - prompt: String, - systemPrompt: String = "", - maxCompletionTokens: Int = 4096, - stream: Boolean = false, - ): OpenAIChatCompletionRequest { - val messages = if (systemPrompt.isEmpty()) { - listOf(OpenAIChatCompletionStandardMessage("user", prompt)) - } else { - listOf( - OpenAIChatCompletionStandardMessage("user", systemPrompt), - OpenAIChatCompletionStandardMessage("user", prompt) - ) - } - return OpenAIChatCompletionRequest.Builder(messages) - .setModel(model) - .setMaxCompletionTokens(maxCompletionTokens) - .setMaxTokens(null) - .setStream(stream) - .setTemperature(null) - .setFrequencyPenalty(null) - .setPresencePenalty(null) - .build() - } - - fun buildOpenAIMessages( - model: String?, - callParameters: ChatCompletionParameters, - referencedFiles: List? = null, - conversationsHistory: List? = null, - psiStructure: Set? = null - ): List { - val messages = buildOpenAIChatMessages( - model = model, - callParameters = callParameters, - referencedFiles = referencedFiles ?: callParameters.referencedFiles, - conversationsHistory = conversationsHistory ?: callParameters.history, - psiStructure = psiStructure, - ) - - if (model == null) { - return messages - } - - val encodingManager = EncodingManager.getInstance() - val totalUsage = messages.parallelStream() - .mapToInt { message: OpenAIChatCompletionMessage? -> - when (message) { - is OpenAIChatCompletionToolMessage -> { - countToolMessageTokens(message, encodingManager) - } - - is OpenAIChatCompletionAssistantMessage -> { - countAssistantMessageTokens(message, encodingManager) - } - - else -> encodingManager.countMessageTokens(message) - } - } - .sum() + getState().maxTokens - val modelMaxTokens: Int - try { - modelMaxTokens = findByCode(model).maxTokens - - if (totalUsage <= modelMaxTokens) { - return messages - } - } catch (ex: NoSuchElementException) { - return messages - } - return tryReducingMessagesOrThrow( - messages, - callParameters.conversation.isDiscardTokenLimit, - totalUsage, - modelMaxTokens - ) - } - - private fun buildOpenAIChatMessages( - model: String?, - callParameters: ChatCompletionParameters, - referencedFiles: List? = null, - conversationsHistory: List? = null, - psiStructure: Set? = null - ): MutableList { - val messages = mutableListOf() - val role = if (isReasoningModel(model)) "user" else "system" - - val effectiveHistory = conversationsHistory ?: callParameters.history - val effectiveRefs = referencedFiles ?: callParameters.referencedFiles - val effectivePsi = psiStructure ?: callParameters.psiStructure - - addSystemMessage(messages, callParameters, effectiveHistory, role, effectiveRefs) - - when (callParameters.requestType) { - RequestType.TOOL_CALL_REQUEST -> { - handleToolCallRequest(messages, callParameters, effectiveRefs, effectivePsi) - } - - RequestType.TOOL_CALL_CONTINUATION -> { - handleToolCallContinuation(messages, callParameters) - } - - RequestType.NORMAL_REQUEST -> { - handleNormalRequest(messages, callParameters, effectiveRefs, effectivePsi) - } - } - - return messages - } - - private fun addSystemMessage( - messages: MutableList, - callParameters: ChatCompletionParameters, - conversationsHistory: List?, - role: String, - referencedFiles: List? - ) { - val selectedPersona = service().state.personas.selectedPersona - if (callParameters.conversationType == ConversationType.DEFAULT && !selectedPersona.disabled) { - val sessionPersonaDetails = callParameters.personaDetails - val baseInstructions = sessionPersonaDetails?.instructions?.addProjectPath() - ?: service() - .getFilteredPersonaPrompt(callParameters.chatMode) - .addProjectPath() - val instructions = service() - .applyClickableLinks(baseInstructions) - val history = if (conversationsHistory.isNullOrEmpty()) { - "" - } else { - conversationsHistory.joinToString("\n\n") { - ConversationTagProcessor.formatConversation(it) - } - } - - if (instructions.isNotEmpty()) { - val filesBlock = buildString { - val unique = mutableSetOf() - val hasRefs = referencedFiles - ?.any { !it.fileContent().isNullOrBlank() && unique.add(it.filePath()) } == true - if (hasRefs) { - append("\n\n") - referencedFiles!!.forEach { rf -> - if (!rf.fileContent().isNullOrBlank()) { - append("\n") - append(CompletionRequestUtil.formatCode(rf.fileContent(), rf.filePath())) - } - } - append("\n") - } - } - - val base = if (history.isBlank()) instructions else instructions.trimEnd() + "\n" + history - val content = base + filesBlock - messages.add(OpenAIChatCompletionStandardMessage(role, content)) - } - } - if (callParameters.conversationType == ConversationType.REVIEW_CHANGES) { - val base = service().state.coreActions.reviewChanges.instructions - messages.add(OpenAIChatCompletionStandardMessage(role, base)) - } - if (callParameters.conversationType == ConversationType.FIX_COMPILE_ERRORS) { - val base = - service().state.coreActions.fixCompileErrors.instructions - messages.add(OpenAIChatCompletionStandardMessage(role, base)) - } - - if (!callParameters.mcpTools.isNullOrEmpty() && callParameters.toolApprovalMode != ToolApprovalMode.BLOCK_ALL) { - val toolsPrompt = - McpToolPromptFormatter().formatToolsForSystemPrompt(callParameters.mcpTools!!) - if (toolsPrompt.isNotEmpty()) { - val lastMessage = messages.lastOrNull() - if (lastMessage is OpenAIChatCompletionStandardMessage && lastMessage.role == "system") { - messages[messages.size - 1] = OpenAIChatCompletionStandardMessage( - "system", - lastMessage.content + "\n" + toolsPrompt - ) - } else { - messages.add(OpenAIChatCompletionStandardMessage("system", toolsPrompt)) - } - } - } - } - - private fun addConversationHistory( - messages: MutableList, - callParameters: ChatCompletionParameters - ) { - callParameters.conversation.messages.forEach { msg -> - val isCurrent = msg.id == callParameters.message.id - - if (msg.prompt.isNotEmpty() && !isCurrent) { - messages.add(OpenAIChatCompletionStandardMessage("user", msg.prompt)) - } - - if (msg.hasToolCalls()) { - if (!isCurrent || !callParameters.retry) { - messages.add( - OpenAIChatCompletionAssistantMessage( - msg.response ?: "", - msg.toolCalls - ) - ) - msg.toolCallResults?.forEach { result -> - val functionName = - msg.toolCalls?.find { it.id == result.key }?.function?.name - ?: "unknown-function" - messages.add( - OpenAIChatCompletionToolMessage( - result.key, - functionName, - result.value - ) - ) - } - } - } else { - if (!isCurrent || !callParameters.retry) { - messages.add( - OpenAIChatCompletionStandardMessage( - "assistant", - msg.response ?: "" - ) - ) - } - } - } - } - - private fun handleNormalRequest( - messages: MutableList, - callParameters: ChatCompletionParameters, - referencedFiles: List?, - psiStructure: Set? - ) { - addConversationHistory(messages, callParameters) - - val message = callParameters.message - val imageDetails = callParameters.imageDetails - if (imageDetails != null) { - messages.add( - OpenAIChatCompletionDetailedMessage( - "user", - listOf( - OpenAIMessageImageURLContent( - OpenAIImageUrl( - imageDetails.mediaType, - imageDetails.data - ) - ), - OpenAIMessageTextContent(message.prompt) - ) - ) - ) - } else { - if (message.prompt.isNotEmpty()) { - messages.add(OpenAIChatCompletionStandardMessage("user", message.prompt)) - } - } - } - - private fun handleToolCallRequest( - messages: MutableList, - callParameters: ChatCompletionParameters, - referencedFiles: List?, - psiStructure: Set? - ) { - addConversationHistory(messages, callParameters) - - val message = callParameters.message - if (message.prompt.isNotEmpty()) { - messages.add(OpenAIChatCompletionStandardMessage("user", message.prompt)) - } - - if (message.hasToolCalls()) { - messages.add( - OpenAIChatCompletionAssistantMessage( - message.response ?: "", - message.toolCalls - ) - ) - } - } - - private fun handleToolCallContinuation( - messages: MutableList, - callParameters: ChatCompletionParameters - ) { - addConversationHistory(messages, callParameters) - - val message = callParameters.message - if (message.prompt.isNotEmpty()) { - messages.add(OpenAIChatCompletionStandardMessage("user", message.prompt)) - } - - if (!message.toolCalls.isNullOrEmpty()) { - messages.add( - OpenAIChatCompletionAssistantMessage( - message.response, - message.toolCalls - ) - ) - } - - message.toolCallResults?.forEach { (callId, result) -> - val functionName = - message.toolCalls?.find { it.id == callId }?.function?.name - ?: "unknown-function" - messages.add(OpenAIChatCompletionToolMessage(callId, functionName, result)) - } - - if (!callParameters.retry && message.response != null && message.response.isNotEmpty()) { - messages.add(OpenAIChatCompletionStandardMessage("assistant", message.response)) - } - } - - private fun tryReducingMessagesOrThrow( - messages: MutableList, - discardTokenLimit: Boolean, - totalInputUsage: Int, - modelMaxTokens: Int - ): List { - val result: MutableList = messages.toMutableList() - var totalUsage = totalInputUsage - if (!ConversationsState.getInstance().discardAllTokenLimits) { - if (!discardTokenLimit) { - throw TotalUsageExceededException() - } - } - val encodingManager = EncodingManager.getInstance() - // skip the system prompt - for (i in 1 until result.size - 1) { - if (totalUsage <= modelMaxTokens) { - break - } - - val message = result[i] - when (message) { - is OpenAIChatCompletionStandardMessage -> { - totalUsage -= encodingManager.countMessageTokens(message) - result[i] = null - } - - is OpenAIChatCompletionToolMessage -> { - totalUsage -= encodingManager.countMessageTokens(message) - } - - is OpenAIChatCompletionDetailedMessage -> { - totalUsage -= encodingManager.countMessageTokens(message) - result[i] = null - } - } - } - - return result.filterNotNull() - } - - fun createBasicCompletionRequest( - systemPrompt: String, - userPrompt: String, - model: String? = null, - isStream: Boolean = false - ): OpenAIChatCompletionRequest { - return OpenAIChatCompletionRequest.Builder( - listOf( - OpenAIChatCompletionStandardMessage("system", systemPrompt), - OpenAIChatCompletionStandardMessage("user", userPrompt) - ) - ) - .setModel(model) - .setStream(isStream) - .build() - } - - private fun countToolMessageTokens( - message: OpenAIChatCompletionToolMessage, - encodingManager: EncodingManager - ): Int { - val content = message.content ?: "" - val toolCallId = message.callId ?: "" - - val contentTokens = encodingManager.countTokens(content) - val toolCallIdTokens = encodingManager.countTokens(toolCallId) - - return contentTokens + toolCallIdTokens + 10 - } - - private fun countAssistantMessageTokens( - message: OpenAIChatCompletionAssistantMessage, - encodingManager: EncodingManager - ): Int { - val content = message.content ?: "" - var totalTokens = encodingManager.countMessageTokens("assistant", content) - - message.toolCalls?.let { toolCalls -> - toolCalls.forEach { toolCall -> - val toolCallId = toolCall.id ?: "" - val toolType = toolCall.type ?: "" - val functionName = toolCall.function?.name ?: "" - val functionArgs = toolCall.function?.arguments ?: "" - - totalTokens += encodingManager.countTokens(toolCallId) - totalTokens += encodingManager.countTokens(toolType) - totalTokens += encodingManager.countTokens(functionName) - totalTokens += encodingManager.countTokens(functionArgs) - - totalTokens += 8 - } - } - - return totalTokens - } - } -} +class OpenAIRequestFactory : BaseRequestFactory() diff --git a/src/main/kotlin/ee/carlrobert/codegpt/completions/inception/InceptionModels.kt b/src/main/kotlin/ee/carlrobert/codegpt/completions/inception/InceptionModels.kt new file mode 100644 index 00000000..c59678ae --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/completions/inception/InceptionModels.kt @@ -0,0 +1,70 @@ +package ee.carlrobert.codegpt.completions.inception + +import com.fasterxml.jackson.annotation.JsonIgnoreProperties +import com.fasterxml.jackson.annotation.JsonProperty +import ee.carlrobert.codegpt.completions.CompletionRequest + +data class InceptionChatMessage( + val role: String, + val content: String, +) + +data class InceptionApplyRequest( + val model: String, + val messages: List, + val stream: Boolean = false, + val diffusing: Boolean = false, +) : CompletionRequest { + class Builder { + private var model: String = "" + private var messages: List = emptyList() + private var stream: Boolean = false + private var diffusing: Boolean = false + + fun setModel(model: String) = apply { this.model = model } + fun setMessages(messages: List) = apply { this.messages = messages } + fun setStream(stream: Boolean) = apply { this.stream = stream } + fun setDiffusing(diffusing: Boolean) = apply { this.diffusing = diffusing } + + fun build(): InceptionApplyRequest { + return InceptionApplyRequest(model, messages, stream, diffusing) + } + } +} + +data class InceptionNextEditRequest( + val model: String, + val messages: List, +) : CompletionRequest { + class Builder { + private var model: String = "" + private var messages: List = emptyList() + + fun setModel(model: String) = apply { this.model = model } + fun setMessages(messages: List) = apply { this.messages = messages } + + fun build(): InceptionNextEditRequest { + return InceptionNextEditRequest(model, messages) + } + } +} + +@JsonIgnoreProperties(ignoreUnknown = true) +data class InceptionChatCompletionResponse( + val id: String? = null, + val choices: List? = null, +) + +@JsonIgnoreProperties(ignoreUnknown = true) +data class InceptionChatCompletionChoice( + val message: InceptionChatCompletionMessage? = null, + val delta: InceptionChatCompletionMessage? = null, + @JsonProperty("finish_reason") + val finishReason: String? = null, +) + +@JsonIgnoreProperties(ignoreUnknown = true) +data class InceptionChatCompletionMessage( + val role: String? = null, + val content: String? = null, +) diff --git a/src/main/kotlin/ee/carlrobert/codegpt/inlineedit/InlineAskResponseListener.kt b/src/main/kotlin/ee/carlrobert/codegpt/inlineedit/InlineAskResponseListener.kt index 3a51a5ff..0fe5b60d 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/inlineedit/InlineAskResponseListener.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/inlineedit/InlineAskResponseListener.kt @@ -5,9 +5,9 @@ import com.intellij.openapi.diagnostic.Logger import com.intellij.openapi.project.Project 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.ui.OverlayUtil -import ee.carlrobert.llm.client.openai.completion.ErrorDetails class InlineAskResponseListener( private val project: Project, @@ -41,11 +41,11 @@ class InlineAskResponseListener( } } - override fun handleError(error: ErrorDetails?, ex: Throwable?) { + override fun handleError(error: ChatError, ex: Throwable) { runInEdt { inlay.onCompletionFinished() CompletionProgressNotifier.update(project, false) - val message = error?.message ?: ex?.message ?: "Something went wrong" + val message = error.message.ifBlank { ex.message ?: "Something went wrong" } OverlayUtil.showNotification(message) } } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/inlineedit/InlineEditSearchReplaceListener.kt b/src/main/kotlin/ee/carlrobert/codegpt/inlineedit/InlineEditSearchReplaceListener.kt index c7bd136e..ef3b3b82 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/inlineedit/InlineEditSearchReplaceListener.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/inlineedit/InlineEditSearchReplaceListener.kt @@ -24,6 +24,8 @@ import com.intellij.util.ui.JBUI import ee.carlrobert.codegpt.CodeGPTBundle import ee.carlrobert.codegpt.CodeGPTKeys import ee.carlrobert.codegpt.codecompletions.CompletionProgressNotifier +import ee.carlrobert.codegpt.completions.CompletionError +import ee.carlrobert.codegpt.completions.CompletionStreamEventListener import ee.carlrobert.codegpt.conversations.Conversation import ee.carlrobert.codegpt.conversations.ConversationService import ee.carlrobert.codegpt.conversations.message.Message @@ -35,9 +37,6 @@ import ee.carlrobert.codegpt.toolwindow.chat.parser.SseMessageParser import ee.carlrobert.codegpt.ui.OverlayUtil import ee.carlrobert.codegpt.ui.components.InlineEditChips import ee.carlrobert.codegpt.util.EditorDiffUtil -import ee.carlrobert.llm.client.openai.completion.ErrorDetails -import ee.carlrobert.llm.completion.CompletionEventListener -import okhttp3.sse.EventSource import java.awt.Color import java.awt.Font import java.util.* @@ -56,7 +55,7 @@ class InlineEditSearchReplaceListener( private val selectionTextRange: TextRange, private val requestId: Long, private val conversation: Conversation -) : CompletionEventListener { +) : CompletionStreamEventListener { private val project: Project = editor.project!! private val logger = Logger.getInstance(InlineEditSearchReplaceListener::class.java) @@ -70,10 +69,10 @@ class InlineEditSearchReplaceListener( private val searchHighlighters = mutableListOf() private var currentSearchPattern: String? = null - private val highlightDebounceAlarm = Alarm() + private val highlightDebounceAlarm = Alarm(Alarm.ThreadToUse.POOLED_THREAD, project) private var hintComponent: JComponent? = null private var lastHintMessage: String? = null - private val waitingAlarm = Alarm() + private val waitingAlarm = Alarm(Alarm.ThreadToUse.POOLED_THREAD, project) private val searchHighlightColor = JBColor( Color(255, 235, 59, 80), @@ -255,7 +254,7 @@ class InlineEditSearchReplaceListener( } } - override fun onMessage(message: String, eventSource: EventSource) { + override fun onMessage(message: String) { if (!isCurrentRequest()) return if (isStopping) return hasReceivedMessage = true @@ -315,7 +314,7 @@ class InlineEditSearchReplaceListener( } } - override fun onComplete(completionMessageBuilder: StringBuilder) { + override fun onComplete(messageBuilder: StringBuilder) { if (!isCurrentRequest()) return runInEdt { @@ -379,6 +378,10 @@ class InlineEditSearchReplaceListener( } } + override fun onCancelled(messageBuilder: StringBuilder) { + stopLoading() + } + private fun createDiffViewerLink(): ActionLink { return ActionLink("Open in Diff Viewer") { try { @@ -472,7 +475,7 @@ class InlineEditSearchReplaceListener( project.service().displayConversation(newConversation) } - override fun onError(error: ErrorDetails, ex: Throwable) { + override fun onError(error: CompletionError, ex: Throwable) { if (!isCurrentRequest()) return runInEdt { diff --git a/src/main/kotlin/ee/carlrobert/codegpt/inlineedit/InlineEditSubmissionHandler.kt b/src/main/kotlin/ee/carlrobert/codegpt/inlineedit/InlineEditSubmissionHandler.kt index 4cbe1756..9d5ecdf1 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/inlineedit/InlineEditSubmissionHandler.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/inlineedit/InlineEditSubmissionHandler.kt @@ -11,19 +11,18 @@ import com.intellij.openapi.editor.ex.EditorEx import com.intellij.openapi.fileEditor.FileDocumentManager import com.intellij.openapi.util.TextRange import com.intellij.openapi.util.text.StringUtil -import ee.carlrobert.codegpt.CodeGPTKeys import ee.carlrobert.codegpt.CodeGPTBundle +import ee.carlrobert.codegpt.CodeGPTKeys import ee.carlrobert.codegpt.ReferencedFile import ee.carlrobert.codegpt.codecompletions.CompletionProgressNotifier import ee.carlrobert.codegpt.completions.* import ee.carlrobert.codegpt.conversations.Conversation import ee.carlrobert.codegpt.conversations.message.Message import ee.carlrobert.codegpt.settings.configuration.ChatMode +import ee.carlrobert.codegpt.settings.models.ModelSettings import ee.carlrobert.codegpt.settings.service.FeatureType import ee.carlrobert.codegpt.settings.service.ServiceType -import ee.carlrobert.codegpt.settings.service.ModelSelectionService import ee.carlrobert.codegpt.ui.OverlayUtil -import okhttp3.sse.EventSource import java.util.concurrent.atomic.AtomicReference class InlineEditSubmissionHandler( @@ -34,7 +33,7 @@ class InlineEditSubmissionHandler( private val previousSourceRef = AtomicReference(null) private val previousPromptRef = AtomicReference(null) - private val currentEventSourceRef = AtomicReference(null) + private val currentRequestRef = AtomicReference(null) private val logger = Logger.getInstance(InlineEditSubmissionHandler::class.java) fun getSessionConversation(): Conversation = sessionConversation @@ -53,7 +52,7 @@ class InlineEditSubmissionHandler( val editorEx = editor as? EditorEx ?: return val inlay = editorEx.getUserData(InlineEditInlay.INLAY_KEY) ?: return val selectedService = - ModelSelectionService.getInstance().getServiceForFeature(FeatureType.INLINE_EDIT) + ModelSettings.getInstance().getServiceForFeature(FeatureType.INLINE_EDIT) if (!forceApply && inlay.isQuickQuestionEnabled()) { handleAskSubmission( @@ -138,21 +137,27 @@ class InlineEditSubmissionHandler( params: ChatCompletionParameters ) { try { - currentEventSourceRef.getAndSet(null)?.cancel() + currentRequestRef.getAndSet(null)?.cancel() val project = editor.project ?: return - val request = CompletionRequestFactory - .getFactory(selectedService) - .createInlineEditQuestionRequest(params) + val prompt = CompletionRequestFactory.getFactory(selectedService) + .createInlineEditQuestionPrompt(params) + val modelSelection = ModelSettings.getInstance() + .getModelSelectionForFeature(FeatureType.INLINE_EDIT) val listener = ChatCompletionEventListener( project, params, InlineAskResponseListener(project, inlay) ) - val eventSource = service() - .getChatCompletionAsync(request, listener, selectedService, null) - currentEventSourceRef.set(eventSource) + val request = CompletionRequestService.getChatCompletionAsync( + serviceType = selectedService, + prompt = prompt, + modelSelection = modelSelection, + callParameters = params, + eventListener = listener + ) + currentRequestRef.set(request) } catch (ex: Exception) { logger.warn("InlineAsk: request dispatch failed", ex) runInEdt { @@ -223,13 +228,13 @@ class InlineEditSubmissionHandler( listener: InlineEditSearchReplaceListener ) { try { - currentEventSourceRef.getAndSet(null)?.cancel() + currentRequestRef.getAndSet(null)?.cancel() - val eventSource = service().getInlineEditCompletionAsync( + val request = service().getInlineEditCompletion( parameters, listener ) - currentEventSourceRef.set(eventSource) + currentRequestRef.set(request) } catch (ex: Exception) { logger.warn("InlineEdit: request dispatch failed", ex) runInEdt { @@ -266,7 +271,7 @@ class InlineEditSubmissionHandler( observableProperties.loading.set(false) observableProperties.submitted.set(false) editor.project?.let { project -> - CompletionProgressNotifier.Companion.update(project, false) + CompletionProgressNotifier.update(project, false) } editorEx.getUserData(InlineEditInlay.INLAY_KEY)?.onCompletionFinished() } @@ -295,7 +300,7 @@ class InlineEditSubmissionHandler( observableProperties.loading.set(false) observableProperties.submitted.set(false) editor.project?.let { project -> - CompletionProgressNotifier.Companion.update(project, false) + CompletionProgressNotifier.update(project, false) } editorEx?.getUserData(InlineEditInlay.INLAY_KEY)?.onCompletionFinished() } @@ -306,7 +311,7 @@ class InlineEditSubmissionHandler( val newRequestId = System.nanoTime() editorEx?.putUserData(InlineEditSearchReplaceListener.REQUEST_ID_KEY, newRequestId) - currentEventSourceRef.getAndSet(null)?.cancel() + currentRequestRef.getAndSet(null)?.cancel() } private fun revertAllChanges(prevSource: String) { diff --git a/src/main/kotlin/ee/carlrobert/codegpt/inlineedit/engine/InlineEditApplyStrategyFactory.kt b/src/main/kotlin/ee/carlrobert/codegpt/inlineedit/engine/InlineEditApplyStrategyFactory.kt index 4e712d65..34627259 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/inlineedit/engine/InlineEditApplyStrategyFactory.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/inlineedit/engine/InlineEditApplyStrategyFactory.kt @@ -1,12 +1,12 @@ package ee.carlrobert.codegpt.inlineedit.engine 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 object InlineEditApplyStrategyFactory { fun get(): ApplyStrategy { - val serviceType = ModelSelectionService.getInstance().getServiceForFeature(FeatureType.INLINE_EDIT) + val serviceType = ModelSettings.getInstance().getServiceForFeature(FeatureType.INLINE_EDIT) return if (serviceType == ServiceType.PROXYAI) ProxyAIApplyStrategy() else SearchReplaceApplyStrategy() } } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/inlineedit/engine/ProxyAIApplyStrategy.kt b/src/main/kotlin/ee/carlrobert/codegpt/inlineedit/engine/ProxyAIApplyStrategy.kt index 4cb55f97..8c28f045 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/inlineedit/engine/ProxyAIApplyStrategy.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/inlineedit/engine/ProxyAIApplyStrategy.kt @@ -1,14 +1,13 @@ package ee.carlrobert.codegpt.inlineedit.engine import com.intellij.openapi.application.runInEdt +import com.intellij.openapi.components.service +import com.intellij.openapi.fileEditor.FileDocumentManager import com.intellij.openapi.util.TextRange -import ee.carlrobert.codegpt.completions.CompletionClientProvider import ee.carlrobert.codegpt.CodeGPTBundle +import ee.carlrobert.codegpt.completions.AutoApplyParameters +import ee.carlrobert.codegpt.completions.AutoApplyService import ee.carlrobert.codegpt.inlineedit.InlineEditSession -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.codegpt.request.AutoApplyRequest import kotlinx.coroutines.CoroutineScope import kotlinx.coroutines.Dispatchers import kotlinx.coroutines.SupervisorJob @@ -22,10 +21,7 @@ class ProxyAIApplyStrategy : ApplyStrategy { override fun apply(ctx: ApplyContext) { val blocks = MarkdownUtil.extractCodeBlocks(ctx.lastAssistantResponse) if (blocks.isEmpty()) return - - val modelSelection = - ModelSelectionService.getInstance().getModelSelectionForFeature(FeatureType.AUTO_APPLY) - if (modelSelection?.provider != ServiceType.PROXYAI) return + val destination = FileDocumentManager.getInstance().getFile(ctx.editor.document) ?: return val updateSnippet = blocks.joinToString("\n// ... existing code ...\n\n") { it.trimEnd() } val original = ctx.editor.document.text @@ -36,9 +32,13 @@ class ProxyAIApplyStrategy : ApplyStrategy { } val merged = try { - CompletionClientProvider.getCodeGPTClient() - .applyChanges(AutoApplyRequest(modelSelection.model, original, updateSnippet)) - .mergedCode + service().applyCode( + AutoApplyParameters( + source = updateSnippet, + destination = destination + ), + original + ) } catch (_: Exception) { null } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/inlineedit/engine/SearchReplaceApplyStrategy.kt b/src/main/kotlin/ee/carlrobert/codegpt/inlineedit/engine/SearchReplaceApplyStrategy.kt index 7aadd66e..feba2b85 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/inlineedit/engine/SearchReplaceApplyStrategy.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/inlineedit/engine/SearchReplaceApplyStrategy.kt @@ -1,12 +1,11 @@ package ee.carlrobert.codegpt.inlineedit.engine import com.intellij.openapi.application.runInEdt -import com.intellij.openapi.util.TextRange -import ee.carlrobert.codegpt.inlineedit.InlineEditSearchReplaceListener -import ee.carlrobert.codegpt.completions.CompletionRequestService -import ee.carlrobert.codegpt.completions.InlineEditCompletionParameters import com.intellij.openapi.components.service -import ee.carlrobert.codegpt.conversations.Conversation +import com.intellij.openapi.util.TextRange +import ee.carlrobert.codegpt.completions.CompletionService +import ee.carlrobert.codegpt.completions.InlineEditCompletionParameters +import ee.carlrobert.codegpt.inlineedit.InlineEditSearchReplaceListener class SearchReplaceApplyStrategy : ApplyStrategy { override fun apply(ctx: ApplyContext) { @@ -45,7 +44,7 @@ class SearchReplaceApplyStrategy : ApplyStrategy { editor.putUserData(InlineEditSearchReplaceListener.LISTENER_KEY, listener) try { - service().getInlineEditCompletionAsync(parameters, listener) + service().getInlineEditCompletion(parameters, listener) } catch (_: Exception) { runInEdt { inlay.setThinkingVisible(false) diff --git a/src/main/kotlin/ee/carlrobert/codegpt/mcp/McpContinuationRunner.kt b/src/main/kotlin/ee/carlrobert/codegpt/mcp/McpContinuationRunner.kt deleted file mode 100644 index 107a1157..00000000 --- a/src/main/kotlin/ee/carlrobert/codegpt/mcp/McpContinuationRunner.kt +++ /dev/null @@ -1,115 +0,0 @@ -package ee.carlrobert.codegpt.mcp - -import com.intellij.openapi.application.runInEdt -import com.intellij.openapi.diagnostic.thisLogger -import 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.codegpt.settings.service.FeatureType -import ee.carlrobert.codegpt.settings.service.ModelSelectionService -import ee.carlrobert.llm.client.openai.completion.ErrorDetails -import ee.carlrobert.llm.client.openai.completion.response.ToolCall -import kotlinx.coroutines.CoroutineScope -import kotlinx.coroutines.Dispatchers -import kotlinx.coroutines.SupervisorJob -import kotlinx.coroutines.launch -import javax.swing.JPanel - -class McpContinuationRunner( - private val callParameters: ChatCompletionParameters, - private val wrappedListener: CompletionResponseEventListener, - private val requestHandler: ToolwindowChatCompletionRequestHandler, - private val setProcessing: (Boolean) -> Unit, - private val resetProgress: () -> Unit, - private val onToolCallUIUpdate: ((JPanel) -> Unit)? = null -) { - private val logger = thisLogger() - private val scope = CoroutineScope(SupervisorJob() + Dispatchers.Default) - - fun run(toolCallResults: List>) { - val toolsToUse = callParameters.message.toolCallResults?.size?.let { - if (it >= 10) { - logger.info("Tool call limit of 10 reached. Continuing without tools for natural conclusion.") - emptyList() - } else { - callParameters.mcpTools - } - } - - val continuationParams = - ChatCompletionParameters.builder(callParameters.conversation, callParameters.message) - .sessionId(callParameters.sessionId) - .conversationType(callParameters.conversationType) - .mcpTools(toolsToUse) - .mcpAttachedServerIds(callParameters.mcpAttachedServerIds) - .toolApprovalMode(callParameters.toolApprovalMode) - .toolResults(toolCallResults) - .project(callParameters.project) - .requestType(RequestType.TOOL_CALL_CONTINUATION) - .build() - - scope.launch(Dispatchers.IO) { - try { - val service = CompletionRequestService.getInstance() - val serviceType = - ModelSelectionService.getInstance().getServiceForFeature(FeatureType.CHAT) - val factory = CompletionRequestFactory.getFactory(serviceType) - val request = factory.createChatRequest(continuationParams) - val continuationListener = McpToolCallEventListener( - project = requireNotNull(callParameters.project) { "Project is required for MCP continuation" }, - callParameters = continuationParams, - wrappedListener = object : CompletionResponseEventListener { - override fun handleRequestOpen() { - wrappedListener.handleRequestOpen() - } - - override fun handleMessage(message: String) { - wrappedListener.handleMessage(message) - } - - override fun handleError(error: ErrorDetails, ex: Throwable) { - setProcessing(false) - resetProgress() - wrappedListener.handleError(error, ex) - } - - override fun handleCompleted( - fullMessage: String, - callParameters: ChatCompletionParameters - ) { - setProcessing(false) - resetProgress() - wrappedListener.handleCompleted(fullMessage, callParameters) - } - - override fun handleTokensExceeded( - conversation: Conversation, - message: Message - ) { - wrappedListener.handleTokensExceeded(conversation, message) - } - - override fun handleCodeGPTEvent(event: CodeGPTEvent) { - wrappedListener.handleCodeGPTEvent(event) - } - }, - onToolCallUIUpdate = { panel -> onToolCallUIUpdate?.invoke(panel) ?: Unit }, - requestHandler = requestHandler - ) - - service.getChatCompletionAsync(request, continuationListener) - ?.let { requestHandler.addEventSource(it) } - } catch (e: Exception) { - logger.error("Failed to continue conversation after tool execution", e) - runInEdt { - setProcessing(false) - wrappedListener.handleError( - ErrorDetails("Failed to continue conversation after tool execution: ${e.message}"), - e - ) - } - } - } - } -} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/mcp/McpToolAliasResolver.kt b/src/main/kotlin/ee/carlrobert/codegpt/mcp/McpToolAliasResolver.kt new file mode 100644 index 00000000..901750f2 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/mcp/McpToolAliasResolver.kt @@ -0,0 +1,42 @@ +package ee.carlrobert.codegpt.mcp + +import java.util.Locale + +internal object McpToolAliasResolver { + + fun resolve( + items: List, + toolName: (T) -> String, + scopeName: (T) -> String + ): List> { + val nameCounts = items + .groupingBy { toolName(it).lowercase(Locale.ROOT) } + .eachCount() + val usedAliases = mutableSetOf() + + return items.map { item -> + val base = if ((nameCounts[toolName(item).lowercase(Locale.ROOT)] ?: 0) > 1) { + "${toolName(item)}_${normalize(scopeName(item))}" + } else { + toolName(item) + } + + var alias = base + var index = 2 + while (!usedAliases.add(alias)) { + alias = "${base}_$index" + index++ + } + + item to alias + } + } + + private fun normalize(value: String): String { + return value + .lowercase(Locale.ROOT) + .replace(Regex("[^a-z0-9]+"), "_") + .trim('_') + .ifBlank { "server" } + } +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/mcp/McpToolArgumentsParser.kt b/src/main/kotlin/ee/carlrobert/codegpt/mcp/McpToolArgumentsParser.kt new file mode 100644 index 00000000..0aa515e5 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/mcp/McpToolArgumentsParser.kt @@ -0,0 +1,42 @@ +package ee.carlrobert.codegpt.mcp + +import com.fasterxml.jackson.databind.ObjectMapper + +object McpToolArgumentsParser { + + private val objectMapper = ObjectMapper() + + data class ParseResult( + val arguments: Map = emptyMap(), + val error: String? = null + ) + + fun parse(argumentsJson: String?): ParseResult { + if (argumentsJson.isNullOrBlank()) { + return ParseResult() + } + + return try { + val raw = (objectMapper.readValue(argumentsJson, Map::class.java) as? Map<*, *>) + ?.entries + ?.mapNotNull { (key, value) -> + val strKey = key?.toString() ?: return@mapNotNull null + strKey to value + } + ?.toMap() + ?: return ParseResult() + ParseResult( + arguments = raw.mapValues { (_, value) -> + when (value) { + is Double -> if (value % 1 == 0.0) value.toInt() else value + else -> value + } + } + ) + } catch (exception: Exception) { + ParseResult( + error = exception.message ?: "Invalid tool arguments JSON" + ) + } + } +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/mcp/McpToolCallEventListener.kt b/src/main/kotlin/ee/carlrobert/codegpt/mcp/McpToolCallEventListener.kt deleted file mode 100644 index 31d8f3c1..00000000 --- a/src/main/kotlin/ee/carlrobert/codegpt/mcp/McpToolCallEventListener.kt +++ /dev/null @@ -1,560 +0,0 @@ -package ee.carlrobert.codegpt.mcp - -import com.fasterxml.jackson.databind.ObjectMapper -import com.intellij.icons.AllIcons -import com.intellij.openapi.application.invokeLater -import com.intellij.openapi.components.service -import com.intellij.openapi.diagnostic.thisLogger -import com.intellij.openapi.project.Project -import com.intellij.ui.JBColor -import com.intellij.util.ui.JBUI -import com.intellij.util.ui.UIUtil -import ee.carlrobert.codegpt.codecompletions.CompletionProgressNotifier -import ee.carlrobert.codegpt.completions.ChatCompletionParameters -import ee.carlrobert.codegpt.completions.CompletionResponseEventListener -import ee.carlrobert.codegpt.completions.ToolApprovalMode -import ee.carlrobert.codegpt.completions.ToolwindowChatCompletionRequestHandler -import ee.carlrobert.codegpt.conversations.ConversationService -import ee.carlrobert.codegpt.toolwindow.ui.mcp.McpStatusPanel -import ee.carlrobert.codegpt.toolwindow.ui.mcp.ToolGroupPanel -import ee.carlrobert.llm.client.openai.completion.ErrorDetails -import ee.carlrobert.llm.client.openai.completion.response.ToolCall -import ee.carlrobert.llm.client.openai.completion.response.ToolFunctionResponse -import ee.carlrobert.llm.completion.CompletionEventListener -import okhttp3.sse.EventSource -import java.util.concurrent.CompletableFuture -import javax.swing.BorderFactory -import javax.swing.BoxLayout -import javax.swing.JLabel -import javax.swing.JPanel - -class McpToolCallEventListener( - private val project: Project, - private val callParameters: ChatCompletionParameters, - private val wrappedListener: CompletionResponseEventListener, - private val onToolCallUIUpdate: (JPanel) -> Unit, - private val requestHandler: ToolwindowChatCompletionRequestHandler -) : CompletionEventListener { - - companion object { - private val logger = thisLogger() - } - - private val objectMapper = ObjectMapper() - private val toolCallHandler = McpToolCallHandler.getInstance(project) - private val sessionManager = service() - private val messageBuilder = StringBuilder() - private val aggregator = ToolCallAggregator() - private var isProcessingToolCalls = false - private val activeStatusPanels = mutableMapOf() - private val mcpToolsMap = mutableMapOf() - - init { - callParameters.mcpTools?.forEach { tool -> - mcpToolsMap[tool.name] = tool - } - } - - override fun onOpen() { - wrappedListener.handleRequestOpen() - } - - override fun onToolCall(toolCall: ToolCall) { - isProcessingToolCalls = true - aggregator.add(toolCall) - } - - override fun onMessage(message: String, eventSource: EventSource) { - messageBuilder.append(message) - wrappedListener.handleMessage(message) - } - - override fun onComplete(messageBuilder: StringBuilder) { - if (isProcessingToolCalls && aggregator.hasCalls()) { - val finalToolCalls = aggregator.build() - aggregator.clear() - - ConversationService.getInstance().saveAssistantMessageWithToolCalls( - callParameters, - finalToolCalls - ) - - clearMessageBuilder() - executeToolCalls(finalToolCalls) - } else { - CompletionProgressNotifier.Companion.update(project, false) - wrappedListener.handleCompleted(messageBuilder.toString(), callParameters) - } - } - - override fun onCancelled(messageBuilder: StringBuilder) { - toolCallHandler.cancelAllPendingApprovals() - toolCallHandler.cancelAllExecutions() - - val statusPanelCleanup = { - activeStatusPanels.values.forEach { statusPanel -> - try { - statusPanel.complete(false, "Cancelled") - } catch (e: Exception) { - logger.warn("Failed to complete status panel during cancellation: ${e.message}") - } - } - activeStatusPanels.clear() - callParameters.message.response = messageBuilder.toString() - } - - cleanupAndComplete(statusPanelCleanup) - } - - override fun onError(error: ErrorDetails, ex: Throwable) { - toolCallHandler.cancelAllPendingApprovals() - - val errorCleanup = { - activeStatusPanels.values.forEach { statusPanel -> - try { - statusPanel.complete(false, "Error occurred") - } catch (e: Exception) { - logger.warn("Failed to complete status panel during error: ${e.message}") - } - } - activeStatusPanels.clear() - - if (messageBuilder.isNotEmpty()) { - callParameters.message.response = messageBuilder.toString() - } - } - - cleanupAndHandleError(error, ex, errorCleanup) - } - - private fun executeToolCalls(toolCalls: List) { - if (toolCalls.isEmpty()) { - return - } - - isProcessingToolCalls = true - - val mcpToolCalls = toolCalls - .filterNot { it.function.name.isNullOrEmpty() } - .map { toolCall -> McpToolConverter.toMcpToolCall(toolCall) } - - if (mcpToolCalls.isEmpty()) { - return - } - - executeToolCallsWithApproval(mcpToolCalls, toolCalls) - } - - private data class ToolCallAggregate( - var id: String?, - var index: Int, - var name: String, - var type: String, - val arguments: StringBuilder - ) - - private class ToolCallAggregator { - - private val byIndex = mutableMapOf() - - fun add(toolCall: ToolCall) { - val idx = toolCall.index ?: return - val existing = byIndex[idx] - if (existing == null) { - byIndex[idx] = ToolCallAggregate( - id = toolCall.id, - index = idx, - name = toolCall.function.name ?: "", - type = toolCall.type ?: "function", - arguments = StringBuilder(toolCall.function.arguments ?: "") - ) - } else { - if (!toolCall.id.isNullOrBlank() && existing.id.isNullOrBlank()) existing.id = - toolCall.id - if (!toolCall.function.name.isNullOrBlank()) existing.name = toolCall.function.name - if (!toolCall.type.isNullOrBlank()) existing.type = toolCall.type - if (!toolCall.function.arguments.isNullOrEmpty()) existing.arguments.append(toolCall.function.arguments) - } - } - - fun hasCalls(): Boolean = byIndex.isNotEmpty() - fun build(): List = byIndex.values - .distinctBy { it.index } - .sortedBy { it.index } - .map { agg -> - val argsString = agg.arguments.toString() - ToolCall( - agg.index, - agg.id ?: "tool-call-${agg.index}", - agg.type, - ToolFunctionResponse(agg.name, argsString) - ) - } - - fun clear() = byIndex.clear() - } - - private fun executeToolCallsWithApproval( - mcpToolCalls: List, - originalCalls: List - ) { - val toolCallResults = mutableListOf>>() - val toolCallsWithServer = mcpToolCalls.mapNotNull { mcpToolCall -> - val serverName = findServerForTool(mcpToolCall.name) - if (serverName != null) { - val toolCall = ToolCall( - null, - mcpToolCall.id, - "function", - ToolFunctionResponse( - mcpToolCall.name, - objectMapper.writeValueAsString(mcpToolCall.arguments) - ) - ) - Triple(mcpToolCall, toolCall, serverName) - } else { - logger.warn("No server found for tool: ${mcpToolCall.name}") - null - } - } - - if (toolCallsWithServer.isEmpty()) { - logger.error("No valid tool calls to execute") - isProcessingToolCalls = false - - onComplete(messageBuilder) - return - } - - when (callParameters.toolApprovalMode) { - ToolApprovalMode.AUTO_APPROVE -> { - toolCallsWithServer.forEach { (mcpToolCall, openAIToolCall, serverName) -> - showExecutionStatus(mcpToolCall, serverName) - - val future = toolCallHandler.executeToolCall( - toolCall = mcpToolCall, - serverId = serverName, - conversationId = callParameters.conversation.id, - approvalMode = ToolApprovalMode.AUTO_APPROVE, - onUIUpdate = { panel -> invokeLater { onToolCallUIUpdate(panel) } } - ) - toolCallResults.add(openAIToolCall to future) - } - } - - ToolApprovalMode.REQUIRE_APPROVAL -> { - if (toolCallsWithServer.size > 1) { - val expectedCount = toolCallsWithServer.size - fun tryScheduleContinuation() { - if (toolCallResults.size < expectedCount) return - awaitAllAndContinue(toolCallResults, originalCalls) - } - - val groupPanel = ToolGroupPanel( - toolCalls = toolCallsWithServer.map { (_, toolCall, serverName) -> - val actualServerName = getActualServerName(serverName) - toolCall to actualServerName - }, - onApprove = { toolCall, panel, _ -> - val match = toolCallsWithServer.find { it.second.id == toolCall.id } - match?.let { (mcpToolCall, _, serverName) -> - panel.showExecutionStatus() - - val future = toolCallHandler.executeToolCall( - toolCall = mcpToolCall, - serverId = serverName, - conversationId = callParameters.conversation.id, - approvalMode = ToolApprovalMode.AUTO_APPROVE, - onUIUpdate = { panel -> invokeLater { onToolCallUIUpdate(panel) } } - ) - toolCallResults.add(toolCall to future) - future.whenComplete { result, throwable -> - invokeLater { - panel.showCompletionResult( - throwable?.let { "Error: ${it.message}" } ?: result, - throwable == null - ) - } - } - - tryScheduleContinuation() - } - }, - onReject = { toolCall, panel -> - cleanupAndComplete() - toolCallHandler.cancelAllPendingApprovals() - toolCallHandler.cancelAllExecutions() - }, - onBatchApprove = { panels, _ -> - toolCallsWithServer.forEachIndexed { index, (mcpToolCall, openAIToolCall, serverName) -> - panels[index].showExecutionStatus() - - val future = toolCallHandler.executeToolCall( - toolCall = mcpToolCall, - serverId = serverName, - conversationId = callParameters.conversation.id, - approvalMode = ToolApprovalMode.AUTO_APPROVE, - onUIUpdate = { panel -> invokeLater { onToolCallUIUpdate(panel) } } - ) - toolCallResults.add(openAIToolCall to future) - future.whenComplete { result, throwable -> - invokeLater { - panels[index].showCompletionResult( - throwable?.let { "Error: ${it.message}" } ?: result, - throwable == null - ) - } - } - } - tryScheduleContinuation() - }, - onBatchReject = { panels -> - cleanupAndComplete() - toolCallHandler.cancelAllPendingApprovals() - toolCallHandler.cancelAllExecutions() - } - ) - - invokeLater { - onToolCallUIUpdate(groupPanel) - } - return - } else { - toolCallsWithServer.forEach { (mcpToolCall, openAIToolCall, serverName) -> - val future = toolCallHandler.executeToolCall( - toolCall = mcpToolCall, - serverId = serverName, - conversationId = callParameters.conversation.id, - approvalMode = callParameters.toolApprovalMode, - onUIUpdate = { panel -> invokeLater { onToolCallUIUpdate(panel) } } - ) - toolCallResults.add(openAIToolCall to future) - future.whenComplete { result, throwable -> - invokeLater { - val results = listOf( - openAIToolCall to (throwable?.let { "Error: ${it.message}" } - ?: result) - ) - continueWithToolResults(results, originalCalls) - } - } - } - return - } - } - - ToolApprovalMode.BLOCK_ALL -> { - cleanupAndComplete() - return - } - } - - awaitAllAndContinue(toolCallResults, originalCalls) - } - - private fun awaitAllAndContinue( - toolCallResults: List>>, - originalCalls: List - ) { - val allFutures = toolCallResults.map { it.second }.toTypedArray() - CompletableFuture.allOf(*allFutures).thenRun { - invokeLater { - val results = toolCallResults.map { (toolCall, future) -> - try { - val result = future.join() - toolCall to result - } catch (e: Exception) { - logger.error("Tool '${toolCall.function.name}' failed", e) - toolCall to "Error: ${e.message}" - } - } - continueWithToolResults(results, originalCalls) - } - }.exceptionally { throwable -> - logger.error("Tool execution failed", throwable) - invokeLater { - isProcessingToolCalls = false - aggregator.clear() - clearMessageBuilder() - onToolCallUIUpdate.invoke(createErrorPanel("Tool execution failed: ${throwable.message}")) - wrappedListener.handleCompleted(messageBuilder.toString(), callParameters) - } - null - } - } - - private fun findServerForTool(toolName: String): String? { - val mcpTools = callParameters.mcpTools ?: return null - val tool = mcpTools.find { it.name == toolName } ?: return null - - val conversationId = callParameters.conversation.id - val attachments = sessionManager.getSessionAttachments(conversationId) - val hasAttachment = attachments.any { it.serverId == tool.serverId } - - if (!hasAttachment) { - logger.warn("Tool found but server not attached. Server ID: ${tool.serverId}") - try { - val attachFuture = - sessionManager.attachServerToSession(conversationId, tool.serverId) - attachFuture.get() - } catch (e: Exception) { - logger.error("Failed to attach server: ${e.message}") - } - } - - return tool.serverId - } - - private fun continueWithToolResults( - toolResults: List>, - originalToolCalls: List - ) { - toolResults.forEach { (toolCall, result) -> - val isError = result.startsWith("Error:") || result.startsWith("Tool execution failed:") - - activeStatusPanels[toolCall.id]?.let { statusPanel -> - if (isError) { - val msg = result.trim() - val formatted = if (msg.length > 100) msg.take(100) + "..." else msg - statusPanel.complete(false, formatted) - } else { - statusPanel.complete(true, "Completed successfully") - } - } ?: run { - logger.warn("No status panel found for tool ${toolCall.id} (${toolCall.function.name})") - } - } - - toolResults.forEach { (toolCall, _) -> - activeStatusPanels.remove(toolCall.id) - } - - val toolCallResultPairs = mutableListOf>() - val toolCallsForMatching: List? = originalToolCalls - - if (toolCallsForMatching.isNullOrEmpty()) { - logger.error("No tool calls found! Cannot continue with tool results without original tool calls") - return - } - - toolResults.forEach { (toolCall, result) -> - val originalToolCall = toolCallsForMatching.find { it.id == toolCall.id } - if (originalToolCall != null) { - toolCallResultPairs.add(originalToolCall to result) - } else { - logger.warn("Could not find original tool call for tool ID: ${toolCall.id}, name: ${toolCall.function.name}") - } - } - - val allToolsRejected = toolResults.all { (_, result) -> - result.contains("rejected by user") || result.contains("blocked by policy") - } - - if (allToolsRejected) { - cleanupAndComplete() - return - } - - try { - val toolExecutionResults = mutableMapOf() - toolResults.forEach { (originalToolCall, result) -> - toolExecutionResults[originalToolCall.id] = result - } - ConversationService.getInstance().saveToolExecutionResults( - callParameters.conversation, - callParameters.message, - toolExecutionResults - ) - logger.info("Saved tool execution outputs for ${toolResults.size} tool calls") - } catch (e: Exception) { - logger.error("Failed to save tool outputs in message", e) - } - - val runner = McpContinuationRunner( - callParameters, - wrappedListener, - requestHandler, - { value -> isProcessingToolCalls = value }, - { - isProcessingToolCalls = false - }, - onToolCallUIUpdate - ) - runner.run(toolCallResultPairs) - } - - private fun createErrorPanel(errorMessage: String): JPanel { - val panel = JPanel() - panel.layout = BoxLayout(panel, BoxLayout.X_AXIS) - panel.border = BorderFactory.createCompoundBorder( - BorderFactory.createLineBorder(JBColor.RED, 1), - JBUI.Borders.empty(8) - ) - panel.background = UIUtil.getPanelBackground() - - val iconLabel = JLabel(AllIcons.General.Error) - iconLabel.border = JBUI.Borders.emptyRight(8) - panel.add(iconLabel) - - val messageLabel = JLabel(errorMessage) - messageLabel.foreground = JBColor.RED - panel.add(messageLabel) - - return panel - } - - private fun getActualServerName(serverId: String): String { - return try { - val attachments = sessionManager.getSessionAttachments(callParameters.conversation.id) - val attachment = attachments.find { it.serverId == serverId } - attachment?.serverName ?: serverId - } catch (_: Exception) { - serverId - } - } - - private fun clearMessageBuilder() { - if (messageBuilder.isNotEmpty()) { - messageBuilder.clear() - } - } - - private fun cleanupAndComplete() { - isProcessingToolCalls = false - aggregator.clear() - clearMessageBuilder() - CompletionProgressNotifier.update(project, false) - wrappedListener.handleCompleted(messageBuilder.toString(), callParameters) - } - - private fun cleanupAndComplete(additionalCleanup: () -> Unit) { - isProcessingToolCalls = false - aggregator.clear() - clearMessageBuilder() - additionalCleanup() - CompletionProgressNotifier.update(project, false) - wrappedListener.handleCompleted(messageBuilder.toString(), callParameters) - } - - private fun cleanupAndHandleError(error: ErrorDetails, ex: Throwable, additionalCleanup: () -> Unit = {}) { - isProcessingToolCalls = false - aggregator.clear() - clearMessageBuilder() - additionalCleanup() - wrappedListener.handleError(error, ex) - } - - private fun showExecutionStatus(toolCall: McpToolCall, serverId: String) { - invokeLater { - val serverName = getActualServerName(serverId) - val statusPanel = McpStatusPanel( - toolName = toolCall.name, - serverName = serverName, - initialStatus = "Executing...", - args = toolCall.arguments, - ) - activeStatusPanels[toolCall.id] = statusPanel - onToolCallUIUpdate(statusPanel) - } - } -} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/mcp/McpToolCallHandler.kt b/src/main/kotlin/ee/carlrobert/codegpt/mcp/McpToolCallHandler.kt index 39396a81..a5003116 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/mcp/McpToolCallHandler.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/mcp/McpToolCallHandler.kt @@ -1,16 +1,18 @@ package ee.carlrobert.codegpt.mcp -import com.fasterxml.jackson.databind.ObjectMapper import com.intellij.openapi.application.runInEdt 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.tools.McpTool as AgentMcpTool +import ee.carlrobert.codegpt.completions.ChatToolCall import ee.carlrobert.codegpt.completions.ToolApprovalMode -import ee.carlrobert.codegpt.toolwindow.ui.mcp.McpApprovalPanel -import ee.carlrobert.codegpt.toolwindow.ui.mcp.McpStatusPanel -import ee.carlrobert.llm.client.openai.completion.response.ToolCall -import ee.carlrobert.llm.client.openai.completion.response.ToolFunctionResponse +import ee.carlrobert.codegpt.toolwindow.agent.ui.approval.DefaultApprovalPanelFactory +import ee.carlrobert.codegpt.toolwindow.agent.ui.approval.ToolApprovalRequest +import ee.carlrobert.codegpt.toolwindow.agent.ui.approval.ToolApprovalType +import ee.carlrobert.codegpt.toolwindow.agent.ui.ToolCallCard +import ee.carlrobert.codegpt.toolwindow.agent.ui.descriptor.ToolKind import io.modelcontextprotocol.client.McpSyncClient import io.modelcontextprotocol.spec.McpSchema import io.modelcontextprotocol.spec.McpSchema.ImageContent @@ -27,27 +29,56 @@ import javax.swing.JPanel * Manages the lifecycle of tool calls from request to result display. */ @Service(Service.Level.PROJECT) -class McpToolCallHandler(private val project: Project) { +class McpToolCallHandler( + private val project: Project +) { private val sessionManager = service() - private val pendingApprovals = ConcurrentHashMap>() - private val statusPanels = ConcurrentHashMap() + private val statusPanels = ConcurrentHashMap() + private val statusPanelConversationIds = ConcurrentHashMap() private val activeExecutions = ConcurrentHashMap>() + private val activeExecutionConversationIds = ConcurrentHashMap() + private val autoApprovedConversations = ConcurrentHashMap.newKeySet() + private val approvalQueues = ConcurrentHashMap>() + private val activeApprovals = ConcurrentHashMap() + private val queueLock = Any() + + private data class PendingApprovalRequest( + val toolCall: ChatToolCall, + val parsedArguments: Map, + val serverId: String, + val conversationId: UUID, + val onUIUpdate: (JPanel) -> Unit, + val resultFuture: CompletableFuture + ) fun executeToolCall( - toolCall: McpToolCall, + toolCall: ChatToolCall, serverId: String, conversationId: UUID, approvalMode: ToolApprovalMode, onUIUpdate: (JPanel) -> Unit ): CompletableFuture { + val parseResult = McpToolArgumentsParser.parse(toolCall.function.arguments) + if (parseResult.error != null) { + return CompletableFuture.completedFuture( + "Tool execution blocked: ${parseResult.error}" + ) + } + val parsedArguments = parseResult.arguments return when (approvalMode) { ToolApprovalMode.AUTO_APPROVE -> { - executeToolInternal(toolCall, serverId, conversationId, onUIUpdate) + executeToolInternal(toolCall, parsedArguments, serverId, conversationId, onUIUpdate) } ToolApprovalMode.REQUIRE_APPROVAL -> { - requestUserApproval(toolCall, serverId, conversationId, onUIUpdate) + requestUserApproval( + toolCall, + parsedArguments, + serverId, + conversationId, + onUIUpdate + ) } ToolApprovalMode.BLOCK_ALL -> { @@ -57,57 +88,126 @@ class McpToolCallHandler(private val project: Project) { } private fun requestUserApproval( - mcpToolCall: McpToolCall, + toolCall: ChatToolCall, + parsedArguments: Map, serverId: String, conversationId: UUID, onUIUpdate: (JPanel) -> Unit ): CompletableFuture { - val toolCall = ToolCall( - null, - mcpToolCall.id, - "function", - ToolFunctionResponse( - mcpToolCall.name, - ObjectMapper().writeValueAsString(mcpToolCall.arguments) - ) + if (autoApprovedConversations.contains(conversationId)) { + return executeToolInternal(toolCall, parsedArguments, serverId, conversationId, onUIUpdate) + } + + val resultFuture = CompletableFuture() + val request = PendingApprovalRequest( + toolCall = toolCall, + parsedArguments = parsedArguments, + serverId = serverId, + conversationId = conversationId, + onUIUpdate = onUIUpdate, + resultFuture = resultFuture ) - val approvalFuture = CompletableFuture() - pendingApprovals[toolCall.id] = approvalFuture + synchronized(queueLock) { + approvalQueues.computeIfAbsent(conversationId) { ArrayDeque() } + .addLast(request) + } + + maybeShowNextApproval(conversationId) + return resultFuture + } + + private fun maybeShowNextApproval(conversationId: UUID) { + val next = synchronized(queueLock) { + if (activeApprovals.containsKey(conversationId)) { + return + } + val queue = approvalQueues[conversationId] ?: return + val request = queue.pollFirst() ?: run { + approvalQueues.remove(conversationId) + return + } + activeApprovals[conversationId] = request + request + } runInEdt { val actualServerName = try { - sessionManager.getServerInfo(conversationId, serverId) ?: serverId + sessionManager.getServerInfo(next.conversationId, next.serverId) ?: next.serverId } catch (_: Exception) { - serverId + next.serverId } - val approvalPanel = McpApprovalPanel( - toolCall = toolCall, - serverName = actualServerName, + if (autoApprovedConversations.contains(next.conversationId)) { + resolveApproval(next, approved = true, autoApproveSession = true) + return@runInEdt + } + + val approvalPanel = DefaultApprovalPanelFactory.create( + project = project, + request = ToolApprovalRequest( + type = ToolApprovalType.GENERIC, + title = "Run MCP tool '${next.toolCall.function.name ?: "unknown-function"}'", + details = buildApprovalDetails(actualServerName, next.parsedArguments) + ), onApprove = { autoApprove -> - approvalFuture.complete(true) + resolveApproval(next, approved = true, autoApproveSession = autoApprove) }, onReject = { - approvalFuture.complete(false) - }, + resolveApproval(next, approved = false) + } ) - onUIUpdate(approvalPanel) - } - - return approvalFuture.thenCompose { approved -> - pendingApprovals.remove(toolCall.id) - - if (approved) { - executeToolInternal(mcpToolCall, serverId, conversationId, onUIUpdate) - } else { - CompletableFuture.completedFuture("Tool execution rejected by user") - } + next.onUIUpdate(approvalPanel) } } + private fun resolveApproval( + request: PendingApprovalRequest, + approved: Boolean, + autoApproveSession: Boolean = false + ) { + val removed = synchronized(queueLock) { + if (activeApprovals[request.conversationId] != request) { + false + } else { + activeApprovals.remove(request.conversationId) + true + } + } + if (!removed) { + return + } + + if (autoApproveSession) { + autoApprovedConversations.add(request.conversationId) + } + + if (!approved) { + request.resultFuture.complete("Tool execution rejected by user") + maybeShowNextApproval(request.conversationId) + return + } + + executeToolInternal( + toolCall = request.toolCall, + parsedArguments = request.parsedArguments, + serverId = request.serverId, + conversationId = request.conversationId, + onUIUpdate = request.onUIUpdate + ).whenComplete { result, throwable -> + if (throwable != null) { + request.resultFuture.complete("Error executing tool: ${throwable.message}") + } else { + request.resultFuture.complete(result) + } + } + + maybeShowNextApproval(request.conversationId) + } + private fun executeToolInternal( - toolCall: McpToolCall, + toolCall: ChatToolCall, + parsedArguments: Map, serverId: String, conversationId: UUID, onUIUpdate: (JPanel) -> Unit @@ -119,13 +219,14 @@ class McpToolCallHandler(private val project: Project) { } runInEdt { - val statusPanel = McpStatusPanel( - toolName = toolCall.name, - serverName = actualServerName, - initialStatus = "Executing...", - args = toolCall.arguments, + val statusPanel = ToolCallCard( + project = project, + toolName = toolCall.function.name ?: "MCP", + args = Unit, + overrideKind = ToolKind.MCP ) statusPanels[toolCall.id] = statusPanel + statusPanelConversationIds[toolCall.id] = conversationId onUIUpdate(statusPanel) } @@ -133,47 +234,55 @@ class McpToolCallHandler(private val project: Project) { try { val clientKey = "$conversationId:$serverId" val mcpClient = sessionManager.ensureClientConnected(clientKey).get() - if (mcpClient == null) { - return@supplyAsync "Error: MCP server '$actualServerName' not connected" - } + ?: return@supplyAsync "Error: MCP server '$actualServerName' not connected" - executeWithClient(mcpClient, toolCall) + executeWithClient( + client = mcpClient, + toolCall = toolCall, + parsedArguments = parsedArguments, + serverId = serverId, + serverName = actualServerName + ) } catch (e: Exception) { logger.error("Exception in executeToolInternal: ${e.message}", e) "Error executing tool: ${e.message}" } } activeExecutions[toolCall.id] = execution - execution.whenComplete { _, _ -> activeExecutions.remove(toolCall.id) } + activeExecutionConversationIds[toolCall.id] = conversationId + execution.whenComplete { _, _ -> + activeExecutions.remove(toolCall.id) + activeExecutionConversationIds.remove(toolCall.id) + } return execution } private fun executeWithClient( client: McpSyncClient, - toolCall: McpToolCall, + toolCall: ChatToolCall, + parsedArguments: Map, + serverId: String, + serverName: String, ): String { return try { - val validatedArguments = when { - toolCall.arguments.isEmpty() -> { - emptyMap() - } - - else -> { - toolCall.arguments - } - } - val callToolRequest = McpSchema.CallToolRequest(toolCall.name, validatedArguments) + val toolName = toolCall.function.name + ?: return "Tool execution error: missing tool name" + @Suppress("UNCHECKED_CAST") + val callToolRequest = McpSchema.CallToolRequest( + toolName, + parsedArguments as Map + ) val toolResult = try { val future = CompletableFuture.supplyAsync { client.callTool(callToolRequest) } future.get(30, TimeUnit.SECONDS) } catch (e: TimeoutException) { - val message = "Tool call '${toolCall.name}' timed out after 30 seconds" + val message = "Tool call '$toolName' timed out after 30 seconds" logger.error(message, e) return message } catch (e: Exception) { - val message = "Tool call '${toolCall.name}' failed: ${e.message}" + val message = "Tool call '$toolName' failed: ${e.message}" logger.error(message, e) return message } @@ -194,9 +303,19 @@ class McpToolCallHandler(private val project: Project) { statusPanels[toolCall.id]?.let { panel -> runInEdt { - panel.complete(false, null, errorMessage) + panel.complete( + false, + AgentMcpTool.Result( + serverId = serverId, + serverName = serverName, + toolName = toolName, + success = false, + output = errorMessage + ) + ) } statusPanels.remove(toolCall.id) + statusPanelConversationIds.remove(toolCall.id) } return "Tool execution failed: $errorMessage" @@ -227,41 +346,102 @@ class McpToolCallHandler(private val project: Project) { toolResult.isError -> "Error: Tool execution failed" else -> "Tool executed successfully (no content returned)" } - val resultContent = rawResultContent - statusPanels[toolCall.id]?.let { panel -> runInEdt { - val displayMessage = formatSuccessMessage(resultContent) - panel.complete(true, displayMessage, resultContent) + panel.complete( + true, + AgentMcpTool.Result( + serverId = serverId, + serverName = serverName, + toolName = toolName, + success = true, + output = rawResultContent + ) + ) } statusPanels.remove(toolCall.id) + statusPanelConversationIds.remove(toolCall.id) } - resultContent + rawResultContent } catch (e: Exception) { val errorMessage = "Tool execution error: ${e.message ?: "Unknown error"}" - logger.error("Tool execution failed for '${toolCall.name}': $errorMessage", e) + logger.error("Tool execution failed for '${toolCall.function.name}': $errorMessage", e) val formattedError = formatForStatusPanel(errorMessage) statusPanels[toolCall.id]?.let { panel -> runInEdt { - panel.complete(false, null, formattedError) + panel.complete( + false, + AgentMcpTool.Result( + serverId = serverId, + serverName = serverName, + toolName = toolCall.function.name ?: "MCP", + success = false, + output = formattedError + ) + ) } statusPanels.remove(toolCall.id) + statusPanelConversationIds.remove(toolCall.id) } errorMessage } } - fun cancelAllPendingApprovals() { - pendingApprovals.values.forEach { it.complete(false) } - pendingApprovals.clear() + fun cancelPendingApprovals(conversationId: UUID) { + val toCancel = mutableListOf() + synchronized(queueLock) { + activeApprovals.remove(conversationId)?.let(toCancel::add) + + approvalQueues.remove(conversationId)?.let { queue -> + while (queue.isNotEmpty()) { + toCancel.add(queue.removeFirst()) + } + } + } + + toCancel.forEach { request -> + request.resultFuture.complete("Tool execution rejected by user") + } } - fun cancelAllExecutions() { - activeExecutions.values.forEach { it.cancel(true) } - activeExecutions.clear() + fun cancelExecutions(conversationId: UUID) { + val executionIds = activeExecutionConversationIds + .filterValues { it == conversationId } + .keys + .toList() + executionIds.forEach { callId -> + activeExecutions.remove(callId)?.cancel(true) + activeExecutionConversationIds.remove(callId) + } + + val statusIds = statusPanelConversationIds + .filterValues { it == conversationId } + .keys + .toList() + statusIds.forEach { callId -> + val panel = statusPanels.remove(callId) ?: return@forEach + statusPanelConversationIds.remove(callId) + runInEdt { + panel.complete( + false, + AgentMcpTool.Result.error( + toolName = "MCP", + output = "Cancelled" + ) + ) + } + } + } + + fun clearConversationState(conversationId: UUID) { + autoApprovedConversations.remove(conversationId) + synchronized(queueLock) { + activeApprovals.remove(conversationId) + approvalQueues.remove(conversationId) + } } companion object { @@ -277,48 +457,23 @@ class McpToolCallHandler(private val project: Project) { return if (raw.length > 100) raw.take(100) + "..." else raw } - private fun formatSuccessMessage(result: String): String { - val cleaned = result.trim() - - return when { - cleaned.contains("search results") || cleaned.contains("results found") -> { - val countMatch = Regex("(\\d+)\\s*(?:search\\s*)?results?").find(cleaned.lowercase()) - if (countMatch != null) { - "Found ${countMatch.groupValues[1]} results" - } else { - "Search completed successfully" + private fun buildApprovalDetails(serverName: String, args: Map): String { + val arguments = if (args.isEmpty()) { + "Arguments: none" + } else { + buildString { + append("Arguments:\n") + args.entries.forEachIndexed { index, (key, value) -> + val suffix = if (index < args.size - 1) "\n" else "" + append("- ") + append(key) + append(": ") + append(value) + append(suffix) } } - - cleaned.contains("SELECT") || cleaned.contains("INSERT") || cleaned.contains("UPDATE") -> { - val rowMatch = Regex("(\\d+)\\s*rows?").find(cleaned) - if (rowMatch != null) { - "Query returned ${rowMatch.groupValues[1]} rows" - } else { - "Query completed successfully" - } - } - - cleaned.contains("file") && (cleaned.contains("created") || cleaned.contains("updated") || cleaned.contains("deleted")) -> { - when { - cleaned.contains("created") -> "File created successfully" - cleaned.contains("updated") -> "File updated successfully" - cleaned.contains("deleted") -> "File deleted successfully" - else -> "File operation completed" - } - } - - cleaned.contains("directory") && (cleaned.contains("created") || cleaned.contains("deleted")) -> { - when { - cleaned.contains("created") -> "Directory created successfully" - cleaned.contains("deleted") -> "Directory deleted successfully" - else -> "Directory operation completed" - } - } - - cleaned.isBlank() -> "Completed successfully" - cleaned.length > 100 -> "Completed successfully" - else -> "Completed successfully" } + + return "Server: $serverName\n$arguments" } } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/mcp/McpToolConverter.kt b/src/main/kotlin/ee/carlrobert/codegpt/mcp/McpToolConverter.kt index faf78943..126f6e2a 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/mcp/McpToolConverter.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/mcp/McpToolConverter.kt @@ -1,32 +1,23 @@ package ee.carlrobert.codegpt.mcp -import com.fasterxml.jackson.databind.ObjectMapper -import ee.carlrobert.llm.client.openai.completion.request.Tool -import ee.carlrobert.llm.client.openai.completion.request.ToolFunction -import ee.carlrobert.llm.client.openai.completion.request.ToolFunctionParameters -import ee.carlrobert.llm.client.openai.completion.response.ToolCall -import java.util.UUID - object McpToolConverter { - private val objectMapper = ObjectMapper() - - fun convertToOpenAITool(mcpTool: McpTool): Tool { - val tool = Tool() + fun convertToOpenAITool(mcpTool: McpTool): OpenAITool { + val tool = OpenAITool() tool.type = "function" tool.function = convertToOpenAIFunction(mcpTool) return tool } - private fun convertToOpenAIFunction(mcpTool: McpTool): ToolFunction { - val function = ToolFunction() + private fun convertToOpenAIFunction(mcpTool: McpTool): OpenAIToolFunction { + val function = OpenAIToolFunction() function.name = mcpTool.name function.description = mcpTool.description function.parameters = if (mcpTool.schema.isNotEmpty()) { convertSchemaToParameters(mcpTool.schema) } else { - val emptyParams = ToolFunctionParameters() + val emptyParams = OpenAIToolFunctionParameters() emptyParams.type = "object" emptyParams.properties = mutableMapOf() emptyParams @@ -35,8 +26,8 @@ object McpToolConverter { return function } - private fun convertSchemaToParameters(schema: Map): ToolFunctionParameters { - val parameters = ToolFunctionParameters() + private fun convertSchemaToParameters(schema: Map): OpenAIToolFunctionParameters { + val parameters = OpenAIToolFunctionParameters() parameters.type = "object" val properties = mutableMapOf() @@ -44,27 +35,13 @@ object McpToolConverter { when { schema.containsKey("type") && schema["type"] == "object" -> { - val schemaProperties = schema["properties"] as? Map - if (schemaProperties != null) { - properties.putAll(schemaProperties) - } - - val schemaRequired = schema["required"] as? List - if (schemaRequired != null) { - required.addAll(schemaRequired) - } + properties.putAll(toStringAnyMap(schema["properties"])) + required.addAll(toStringList(schema["required"])) } schema.containsKey("properties") -> { - val schemaProperties = schema["properties"] as? Map - if (schemaProperties != null) { - properties.putAll(schemaProperties) - } - - val schemaRequired = schema["required"] as? List - if (schemaRequired != null) { - required.addAll(schemaRequired) - } + properties.putAll(toStringAnyMap(schema["properties"])) + required.addAll(toStringList(schema["required"])) } else -> { @@ -77,44 +54,16 @@ object McpToolConverter { return parameters } - fun toMcpToolCall(toolCall: ToolCall): McpToolCall { - return McpToolCall( - id = toolCall.id ?: UUID.randomUUID().toString(), - name = toolCall.function.name, - arguments = parseArguments(toolCall.function.arguments) - ) + private fun toStringAnyMap(value: Any?): Map { + val map = value as? Map<*, *> ?: return emptyMap() + return map.entries.mapNotNull { (key, nestedValue) -> + val strKey = key?.toString() ?: return@mapNotNull null + nestedValue?.let { strKey to it } + }.toMap() } - private fun parseArguments(argumentsJson: String?): Map { - if (argumentsJson.isNullOrEmpty()) { - return emptyMap() - } - - return try { - val result = objectMapper.readValue(argumentsJson, Map::class.java) as Map - val processedResult = result.mapValues { (key, value) -> - when (value) { - is Double -> { - if (value % 1 == 0.0) { - value.toInt() - } else { - value - } - } - - else -> value - } - } - - processedResult - } catch (e: Exception) { - emptyMap() - } + private fun toStringList(value: Any?): List { + val list = value as? List<*> ?: return emptyList() + return list.mapNotNull { it?.toString() } } } - -data class McpToolCall( - val id: String, - val name: String, - val arguments: Map -) diff --git a/src/main/kotlin/ee/carlrobert/codegpt/mcp/OpenAIToolModels.kt b/src/main/kotlin/ee/carlrobert/codegpt/mcp/OpenAIToolModels.kt new file mode 100644 index 00000000..e432ba64 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/mcp/OpenAIToolModels.kt @@ -0,0 +1,18 @@ +package ee.carlrobert.codegpt.mcp + +data class OpenAITool( + var type: String? = null, + var function: OpenAIToolFunction? = null, +) + +data class OpenAIToolFunction( + var name: String? = null, + var description: String? = null, + var parameters: OpenAIToolFunctionParameters? = null, +) + +data class OpenAIToolFunctionParameters( + var type: String? = null, + var properties: Map = emptyMap(), + var required: List = emptyList(), +) diff --git a/src/main/kotlin/ee/carlrobert/codegpt/mcp/ToolSchemaParser.kt b/src/main/kotlin/ee/carlrobert/codegpt/mcp/ToolSchemaParser.kt new file mode 100644 index 00000000..a098b7e1 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/mcp/ToolSchemaParser.kt @@ -0,0 +1,86 @@ +package ee.carlrobert.codegpt.mcp + +import ai.koog.agents.core.tools.ToolParameterDescriptor +import ai.koog.agents.core.tools.ToolParameterType + +object ToolSchemaParser { + + fun parseRequiredNames(schema: Map): List { + return (schema["required"] as? List<*>) + ?.mapNotNull { it?.toString() } + ?: emptyList() + } + + fun parseParameterDescriptors(schema: Map): List { + val properties = schema["properties"] as? Map<*, *> ?: emptyMap() + return properties.mapNotNull { (nameRaw, propertyRaw) -> + val name = nameRaw?.toString() ?: return@mapNotNull null + val property = propertyRaw as? Map<*, *> ?: return@mapNotNull null + toToolParameterDescriptor(name, 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 -> + toToolParameterDescriptor("", 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() + 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 toToolParameterDescriptor( + name: String, + definition: Map<*, *> + ): ToolParameterDescriptor { + return ToolParameterDescriptor( + name = name, + description = definition["description"]?.toString().orEmpty(), + type = parseParameterType(definition) + ) + } + + private fun toStringKeyedMap(value: Any?): Map { + val map = value as? Map<*, *> ?: return emptyMap() + return map.entries.mapNotNull { (key, nestedValue) -> + key?.toString()?.let { it to nestedValue } + }.toMap() + } +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/nextedit/InceptionNextEditProvider.kt b/src/main/kotlin/ee/carlrobert/codegpt/nextedit/InceptionNextEditProvider.kt index f6257113..053026e9 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/nextedit/InceptionNextEditProvider.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/nextedit/InceptionNextEditProvider.kt @@ -7,11 +7,11 @@ import com.intellij.openapi.editor.Editor import com.intellij.openapi.fileEditor.FileDocumentManager import ee.carlrobert.codegpt.CodeGPTKeys import ee.carlrobert.codegpt.codecompletions.CompletionProgressNotifier -import ee.carlrobert.codegpt.completions.CompletionClientProvider import ee.carlrobert.codegpt.completions.NextEditParameters import ee.carlrobert.codegpt.completions.factory.InceptionRequestFactory +import ee.carlrobert.codegpt.nextedit.executors.InceptionNextEditExecutor 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.inception.InceptionSettings import ee.carlrobert.codegpt.util.EditWindowFormatter @@ -29,7 +29,7 @@ class InceptionNextEditProvider : NextEditProvider { caretOffset: Int, addToQueue: Boolean ) { - if (service().getServiceForFeature(FeatureType.NEXT_EDIT) != ServiceType.INCEPTION + if (service().getServiceForFeature(FeatureType.NEXT_EDIT) != ServiceType.INCEPTION || !service().state.nextEditsEnabled ) { return @@ -77,7 +77,7 @@ class InceptionNextEditProvider : NextEditProvider { EditWindowFormatter.formatWithIndices(fileContent, params.fileName, caretOffset) val request = InceptionRequestFactory().createNextEditRequest(params, formatResult) - val response = CompletionClientProvider.getInceptionClient().getNextEditCompletion(request) + val response = InceptionNextEditExecutor.execute(request) val text = response.choices?.firstOrNull()?.message?.content ?: return null val prefix = fileContent.substring(0, formatResult.editStartIndex) val editedContent = NextEditRequestProcessor.extractCodeFromBackticks(text) @@ -90,4 +90,4 @@ class InceptionNextEditProvider : NextEditProvider { } } -const val DEFAULT_CONTEXT_TOKENS: Int = 2048 \ No newline at end of file +const val DEFAULT_CONTEXT_TOKENS: Int = 2048 diff --git a/src/main/kotlin/ee/carlrobert/codegpt/nextedit/NextEditCoordinator.kt b/src/main/kotlin/ee/carlrobert/codegpt/nextedit/NextEditCoordinator.kt index 2bdb12ac..4851b946 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/nextedit/NextEditCoordinator.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/nextedit/NextEditCoordinator.kt @@ -5,7 +5,7 @@ import com.intellij.openapi.components.service import com.intellij.openapi.editor.Editor import ee.carlrobert.codegpt.codecompletions.CompletionProgressNotifier 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 object NextEditCoordinator { @@ -22,7 +22,7 @@ object NextEditCoordinator { addToQueue: Boolean = false, ) { val serviceType = - service().getServiceForFeature(FeatureType.NEXT_EDIT) + service().getServiceForFeature(FeatureType.NEXT_EDIT) val provider = providers[serviceType] ?: return editor.project?.let { CompletionProgressNotifier.Companion.update(it, true) } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/nextedit/NextEditDiffViewer.kt b/src/main/kotlin/ee/carlrobert/codegpt/nextedit/NextEditDiffViewer.kt index 8e34ec74..a226f884 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/nextedit/NextEditDiffViewer.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/nextedit/NextEditDiffViewer.kt @@ -317,7 +317,7 @@ class NextEditDiffViewer( fun createSimpleDiffRequest(editor: Editor, nextRevision: String): SimpleDiffRequest { val project = editor.project - val virtualFile = editor.virtualFile + val virtualFile = requireNotNull(editor.virtualFile) { "Editor virtual file is required for diff request." } val tempDiffFile = LightVirtualFile(virtualFile.name, nextRevision) val diffContentFactory = DiffContentFactory.getInstance() return SimpleDiffRequest( diff --git a/src/main/kotlin/ee/carlrobert/codegpt/nextedit/ProxyAINextEditProvider.kt b/src/main/kotlin/ee/carlrobert/codegpt/nextedit/ProxyAINextEditProvider.kt index 16254304..9c62b740 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/nextedit/ProxyAINextEditProvider.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/nextedit/ProxyAINextEditProvider.kt @@ -5,13 +5,13 @@ import com.intellij.openapi.editor.Editor import ee.carlrobert.codegpt.codecompletions.edit.GrpcClientService import ee.carlrobert.codegpt.nextedit.NextEditProvider 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 class ProxyAINextEditProvider : NextEditProvider { override fun request(editor: Editor, fileContent: String, caretOffset: Int, addToQueue: Boolean) { - if (service().getServiceForFeature(FeatureType.NEXT_EDIT) != ServiceType.PROXYAI + if (service().getServiceForFeature(FeatureType.NEXT_EDIT) != ServiceType.PROXYAI || !service().state.nextEditsEnabled ) { return diff --git a/src/main/kotlin/ee/carlrobert/codegpt/nextedit/TriggerNextEditAction.kt b/src/main/kotlin/ee/carlrobert/codegpt/nextedit/TriggerNextEditAction.kt index 60e6849b..a4667b30 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/nextedit/TriggerNextEditAction.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/nextedit/TriggerNextEditAction.kt @@ -10,7 +10,7 @@ import com.intellij.openapi.editor.actionSystem.EditorAction import com.intellij.openapi.editor.actionSystem.EditorWriteActionHandler import com.intellij.util.application 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 kotlin.coroutines.cancellation.CancellationException @@ -24,7 +24,7 @@ class TriggerNextEditAction : EditorAction(Handler()), HintManagerImpl.ActionToI private class Handler : EditorWriteActionHandler() { override fun doExecute(editor: Editor, caret: Caret?, dataContext: DataContext?) { - val nextEditModelProvider = ModelSelectionService.getInstance() + val nextEditModelProvider = ModelSettings.getInstance() .getServiceForFeature(FeatureType.NEXT_EDIT) if (!listOf( ServiceType.PROXYAI, diff --git a/src/main/kotlin/ee/carlrobert/codegpt/nextedit/executors/InceptionNextEditExecutor.kt b/src/main/kotlin/ee/carlrobert/codegpt/nextedit/executors/InceptionNextEditExecutor.kt new file mode 100644 index 00000000..cdeda39e --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/nextedit/executors/InceptionNextEditExecutor.kt @@ -0,0 +1,20 @@ +package ee.carlrobert.codegpt.nextedit.executors + +import ee.carlrobert.codegpt.agent.clients.HttpClientProvider +import ee.carlrobert.codegpt.agent.clients.InceptionAILLMClient +import ee.carlrobert.codegpt.completions.inception.InceptionChatCompletionResponse +import ee.carlrobert.codegpt.completions.inception.InceptionNextEditRequest +import ee.carlrobert.codegpt.credentials.CredentialsStore.CredentialKey +import ee.carlrobert.codegpt.credentials.CredentialsStore.getCredential + +internal object InceptionNextEditExecutor { + fun execute(request: InceptionNextEditRequest): InceptionChatCompletionResponse { + val apiKey = getCredential(CredentialKey.InceptionApiKey) ?: "" + return InceptionAILLMClient( + apiKey = apiKey, + baseClient = HttpClientProvider.createHttpClient() + ).use { client -> + client.getNextEditCompletion(request) + } + } +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/refactorings/RenameCompletionLookupListener.kt b/src/main/kotlin/ee/carlrobert/codegpt/refactorings/RenameCompletionLookupListener.kt index 59722ba0..7e0157da 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/refactorings/RenameCompletionLookupListener.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/refactorings/RenameCompletionLookupListener.kt @@ -13,12 +13,12 @@ import com.intellij.openapi.components.service import com.intellij.openapi.diagnostic.thisLogger import ee.carlrobert.codegpt.CodeGPTKeys.IS_PROMPT_TEXT_FIELD_DOCUMENT import ee.carlrobert.codegpt.Icons -import ee.carlrobert.codegpt.completions.CompletionRequestService +import ee.carlrobert.codegpt.completions.CompletionService import ee.carlrobert.codegpt.completions.LookupCompletionParameters import ee.carlrobert.codegpt.settings.configuration.ConfigurationSettings import ee.carlrobert.codegpt.settings.service.FeatureType +import ee.carlrobert.codegpt.settings.service.codegpt.CodeGPTApiException import ee.carlrobert.codegpt.ui.OverlayUtil -import ee.carlrobert.llm.client.codegpt.response.CodeGPTException import kotlinx.coroutines.CoroutineScope import kotlinx.coroutines.Dispatchers import kotlinx.coroutines.SupervisorJob @@ -42,7 +42,7 @@ class RenameCompletionLookupListener : LookupManagerListener { val featureEnabled = service().state.methodNameGenerationEnabled if (isPromptTextFieldDocument || !featureEnabled - || !CompletionRequestService.isRequestAllowed(FeatureType.CHAT) + || !CompletionService.isRequestAllowed(FeatureType.CHAT) ) { return } @@ -53,13 +53,16 @@ class RenameCompletionLookupListener : LookupManagerListener { val selection = runReadAction { context.text } scope.launch { try { - val lookupCompletion = service() + val lookupCompletion = service() .getLookupCompletion(LookupCompletionParameters(selection)) if (lookupCompletion.isNotEmpty()) { addCompletionLookupValues(newLookup, lookupCompletion) } - } catch (ex: CodeGPTException) { - OverlayUtil.showNotification(ex.detail, NotificationType.ERROR) + } catch (ex: CodeGPTApiException) { + OverlayUtil.showNotification( + ex.detail ?: ex.message ?: "Request failed", + NotificationType.ERROR + ) } catch (ex: Exception) { logger.error( "Something went wrong while requesting completion lookup values.", diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/agents/SubagentGenerator.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/agents/SubagentGenerator.kt index f7ca8ffe..0f89bc88 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/settings/agents/SubagentGenerator.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/agents/SubagentGenerator.kt @@ -5,7 +5,7 @@ import ai.koog.prompt.message.Message import com.intellij.openapi.components.service import ee.carlrobert.codegpt.agent.AgentFactory 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 kotlinx.serialization.Serializable import kotlinx.serialization.decodeFromString @@ -29,7 +29,7 @@ object SubagentGenerator { suspend fun generate(query: String): GeneratedSubagent { return try { - val modelService = service() + val modelService = service() val provider: ServiceType = modelService.getServiceForFeature(FeatureType.AGENT) val model = modelService.getAgentModel() val executor = AgentFactory.createExecutor(provider) diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/hooks/HookGenerator.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/hooks/HookGenerator.kt index 8070c552..4f39bdf6 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/settings/hooks/HookGenerator.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/hooks/HookGenerator.kt @@ -9,7 +9,7 @@ import com.intellij.openapi.diagnostic.thisLogger import com.intellij.openapi.project.Project import ee.carlrobert.codegpt.agent.AgentFactory 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.util.file.FileUtil import kotlinx.serialization.SerialName @@ -73,7 +73,7 @@ object HookGenerator { suspend fun generate(query: String, project: Project? = null): GeneratedHookResult { return try { - val modelService = service() + val modelService = service() val provider: ServiceType = modelService.getServiceForFeature(FeatureType.AGENT) val model = modelService.getAgentModel() val executor = AgentFactory.createExecutor(provider) diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/hooks/HookManager.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/hooks/HookManager.kt index 46987027..852ff605 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/settings/hooks/HookManager.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/hooks/HookManager.kt @@ -18,7 +18,7 @@ class HookManager( event: HookEventType, payload: Map, toolName: String? = null, - sessionId: String? = null + sessionId: String, ): List { val settings = project.service().getSettings() val configuration = settings.hooks ?: return emptyList() @@ -42,7 +42,10 @@ class HookManager( for (hook in matchingHooks) { try { logger.debug("Executing hook '${hook.command}' for event: ${event.eventName}") - val result = executionService.executeHook(hook, event, payload, projectRoot) + val adjustedPayload = payload.toMutableMap().apply { + put("session_id", sessionId) + } + val result = executionService.executeHook(hook, event, adjustedPayload, projectRoot) results.add(result) incrementLoopCount(hook, event, sessionId) @@ -79,8 +82,8 @@ class HookManager( suspend fun checkHooksForDenial( event: HookEventType, payload: Map, - toolName: String? = null, - sessionId: String? = null + toolName: String, + sessionId: String ): String? { val results = executeHooksForEvent(event, payload, toolName, sessionId) results.forEach { result -> @@ -139,18 +142,18 @@ class HookManager( private fun loopAllows( hook: HookConfig, event: HookEventType, - sessionId: String? + sessionId: String ): Boolean { val limit = hook.loopLimit ?: return true if (event != HookEventType.STOP && event != HookEventType.SUBAGENT_STOP) return true - val key = "${sessionId ?: "global"}:${event.eventName}:${hookIdentity(hook)}" + val key = "${sessionId}:${event.eventName}:${hookIdentity(hook)}" val counter = loopCounts.computeIfAbsent(key) { AtomicInteger(0) } return counter.get() < limit } - private fun incrementLoopCount(hook: HookConfig, event: HookEventType, sessionId: String?) { + private fun incrementLoopCount(hook: HookConfig, event: HookEventType, sessionId: String) { if (event != HookEventType.STOP && event != HookEventType.SUBAGENT_STOP) return - val key = "${sessionId ?: "global"}:${event.eventName}:${hookIdentity(hook)}" + val key = "${sessionId}:${event.eventName}:${hookIdentity(hook)}" loopCounts.computeIfAbsent(key) { AtomicInteger(0) }.incrementAndGet() } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/migration/LegacySettingsMigration.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/migration/LegacySettingsMigration.kt index d61d5b97..b6746f34 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/settings/migration/LegacySettingsMigration.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/migration/LegacySettingsMigration.kt @@ -3,18 +3,19 @@ package ee.carlrobert.codegpt.settings.migration import com.intellij.openapi.components.service import com.intellij.openapi.diagnostic.thisLogger import ee.carlrobert.codegpt.settings.GeneralSettings -import ee.carlrobert.codegpt.settings.models.ModelRegistry +import ee.carlrobert.codegpt.settings.models.Devstral2 +import ee.carlrobert.codegpt.settings.models.ModelCatalog +import ee.carlrobert.codegpt.settings.models.ModelSettings import ee.carlrobert.codegpt.settings.models.ModelSettingsState import ee.carlrobert.codegpt.settings.service.FeatureType import ee.carlrobert.codegpt.settings.service.ServiceType import ee.carlrobert.codegpt.settings.service.anthropic.AnthropicSettings import ee.carlrobert.codegpt.settings.service.codegpt.CodeGPTServiceSettings import ee.carlrobert.codegpt.settings.service.custom.CustomServicesSettings -import ee.carlrobert.codegpt.settings.service.google.GoogleSettings 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.google.models.GoogleModel +import ai.koog.prompt.executor.clients.google.GoogleModels object LegacySettingsMigration { @@ -38,9 +39,18 @@ object LegacySettingsMigration { } private fun createMigratedState(selectedService: ServiceType): ModelSettingsState { + val modelSettings = service() return ModelSettingsState().apply { val chatModel = getLegacyChatModelForService(selectedService) - val agentModel = getLegacyAgentModelForService(selectedService, chatModel) + val agentModel = if (modelSettings.isFeatureSupportedByProvider( + FeatureType.AGENT, + selectedService + ) + ) { + getLegacyAgentModelForService(selectedService, chatModel) + } else { + null + } setModelSelection(FeatureType.AGENT, agentModel, selectedService) setModelSelection(FeatureType.CHAT, chatModel, selectedService) @@ -52,11 +62,15 @@ object LegacySettingsMigration { val codeModel = getLegacyCodeModelForService(selectedService) setModelSelection(FeatureType.CODE_COMPLETION, codeModel, selectedService) - if (selectedService == ServiceType.PROXYAI) { - setModelSelection(FeatureType.NEXT_EDIT, ModelRegistry.MERCURY_CODER, ServiceType.PROXYAI) + val nextEditModel = if (selectedService == ServiceType.PROXYAI) { + ModelCatalog.MERCURY_CODER } else { - setModelSelection(FeatureType.NEXT_EDIT, null, selectedService) + resolveAvailableModel( + serviceType = selectedService, + featureType = FeatureType.NEXT_EDIT + ) } + setModelSelection(FeatureType.NEXT_EDIT, nextEditModel, selectedService) } } @@ -65,26 +79,46 @@ object LegacySettingsMigration { when (serviceType) { ServiceType.PROXYAI -> { val settings = service() - settings.state.chatCompletionSettings.model ?: ModelRegistry.GEMINI_FLASH_2_5 + resolveAvailableModel( + serviceType = serviceType, + featureType = FeatureType.CHAT, + preferredModel = settings.state.chatCompletionSettings.model, + fallbackModel = ModelCatalog.GEMINI_FLASH_2_5 + ) } ServiceType.OPENAI -> { - OpenAISettings.getCurrentState().model ?: ModelRegistry.GPT_5 + resolveAvailableModel( + serviceType = serviceType, + featureType = FeatureType.CHAT, + preferredModel = OpenAISettings.getCurrentState().model + ) } ServiceType.ANTHROPIC -> { - AnthropicSettings.getCurrentState().model - ?: ModelRegistry.CLAUDE_SONNET_4_20250514 + resolveAvailableModel( + serviceType = serviceType, + featureType = FeatureType.CHAT, + preferredModel = AnthropicSettings.getCurrentState().model + ) } ServiceType.GOOGLE -> { - val settings = service() - settings.state.model ?: GoogleModel.GEMINI_2_5_PRO.code + resolveAvailableModel( + serviceType = serviceType, + featureType = FeatureType.CHAT, + fallbackModel = GoogleModels.Gemini2_5Pro.id + ) } ServiceType.OLLAMA -> { val settings = service() - settings.state.model ?: ModelRegistry.LLAMA_3_2 + settings.state.model + ?: resolveAvailableModel( + serviceType = serviceType, + featureType = FeatureType.CHAT, + fallbackModel = ModelCatalog.LLAMA_3_2 + ) } ServiceType.LLAMA_CPP -> { @@ -97,20 +131,32 @@ object LegacySettingsMigration { } ServiceType.CUSTOM_OPENAI -> { - service().state.services - .map { it.name } + val preferredServiceId = service().state.services .lastOrNull() - ?.takeIf { it.isNotBlank() } ?: "Default" + ?.id + resolveAvailableModel( + serviceType = serviceType, + featureType = FeatureType.CHAT, + preferredModel = preferredServiceId + ) } ServiceType.MISTRAL -> { - ModelRegistry.DEVSTRAL_MEDIUM_2507 + resolveAvailableModel( + serviceType = serviceType, + featureType = FeatureType.CHAT, + fallbackModel = Devstral2.id + ) } ServiceType.INCEPTION -> { - ModelRegistry.MERCURY_CODER + resolveAvailableModel( + serviceType = serviceType, + featureType = FeatureType.CHAT, + fallbackModel = ModelCatalog.MERCURY_CODER + ) } - } + } ?: error("Could not resolve legacy chat model for $serviceType") } catch (e: Exception) { logger.warn("Failed to get legacy chat model for $serviceType", e) throw e @@ -121,27 +167,39 @@ object LegacySettingsMigration { serviceType: ServiceType, fallbackModel: String ): String { - val registry = ModelRegistry.getInstance() - if (!registry.isFeatureSupportedByProvider(FeatureType.AGENT, serviceType)) { + val modelSettings = service() + if (!modelSettings.isFeatureSupportedByProvider(FeatureType.AGENT, serviceType)) { return fallbackModel } if (serviceType == ServiceType.PROXYAI) { - return ModelRegistry.PROXYAI_AUTO + return ModelCatalog.PROXYAI_AUTO } - return registry.getAgentModels(serviceType).firstOrNull()?.model?.id ?: fallbackModel + return modelSettings.getAvailableModels(FeatureType.AGENT) + .firstOrNull { it.provider == serviceType } + ?.model + ?: fallbackModel } private fun getLegacyCodeModelForService(serviceType: ServiceType): String? { return try { when (serviceType) { ServiceType.PROXYAI -> { - service().state.codeCompletionSettings.model + resolveAvailableModel( + serviceType = serviceType, + featureType = FeatureType.CODE_COMPLETION, + preferredModel = service().state.codeCompletionSettings.model, + fallbackModel = ModelCatalog.MERCURY_CODER + ) } ServiceType.OPENAI -> { - ModelRegistry.GPT_3_5_TURBO_INSTRUCT + resolveAvailableModel( + serviceType = serviceType, + featureType = FeatureType.CODE_COMPLETION, + fallbackModel = ModelCatalog.GPT_3_5_TURBO_INSTRUCT + ) } ServiceType.ANTHROPIC -> { @@ -166,17 +224,30 @@ object LegacySettingsMigration { } ServiceType.CUSTOM_OPENAI -> { - service().state.services - .map { it.name } - .lastOrNull() ?: "" + val preferredServiceId = service().state.services + .lastOrNull() + ?.id + resolveAvailableModel( + serviceType = serviceType, + featureType = FeatureType.CODE_COMPLETION, + preferredModel = preferredServiceId + ) } ServiceType.MISTRAL -> { - ModelRegistry.CODESTRAL_LATEST + resolveAvailableModel( + serviceType = serviceType, + featureType = FeatureType.CODE_COMPLETION, + fallbackModel = ModelCatalog.CODESTRAL_LATEST + ) } ServiceType.INCEPTION -> { - ModelRegistry.MERCURY_CODER + resolveAvailableModel( + serviceType = serviceType, + featureType = FeatureType.CODE_COMPLETION, + fallbackModel = ModelCatalog.MERCURY_CODER + ) } } } catch (e: Exception) { @@ -184,4 +255,27 @@ object LegacySettingsMigration { null } } + + private fun resolveAvailableModel( + serviceType: ServiceType, + featureType: FeatureType, + preferredModel: String? = null, + fallbackModel: String? = null + ): String? { + val modelSettings = service() + if (!modelSettings.isFeatureSupportedByProvider(featureType, serviceType)) { + return null + } + + if (!preferredModel.isNullOrBlank() && + modelSettings.findModel(serviceType, preferredModel) != null + ) { + return preferredModel + } + + return modelSettings.getAvailableModels(featureType) + .firstOrNull { it.provider == serviceType } + ?.model + ?: fallbackModel + } } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/models/LLMClientFactory.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/models/LLMClientFactory.kt new file mode 100644 index 00000000..2234ba83 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/models/LLMClientFactory.kt @@ -0,0 +1,175 @@ +package ee.carlrobert.codegpt.settings.models + +import ai.koog.prompt.executor.clients.LLMClient +import ai.koog.prompt.executor.clients.anthropic.AnthropicClientSettings +import ai.koog.prompt.executor.clients.anthropic.AnthropicLLMClient +import ai.koog.prompt.executor.clients.anthropic.AnthropicModels.Haiku_4_5 +import ai.koog.prompt.executor.clients.anthropic.AnthropicModels.Opus_4_5 +import ai.koog.prompt.executor.clients.anthropic.AnthropicModels.Sonnet_4_5 +import ai.koog.prompt.executor.clients.google.GoogleClientSettings +import ai.koog.prompt.executor.clients.google.GoogleLLMClient +import com.intellij.openapi.components.service +import ee.carlrobert.codegpt.agent.clients.* +import ee.carlrobert.codegpt.credentials.CredentialsStore.CredentialKey +import ee.carlrobert.codegpt.credentials.CredentialsStore.getCredential +import ee.carlrobert.codegpt.settings.service.FeatureType +import ee.carlrobert.codegpt.settings.service.ServiceType +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.settings.service.openai.OpenAISettings + +object LLMClientFactory { + + private val supportedServices = setOf( + ServiceType.PROXYAI, + ServiceType.OPENAI, + ServiceType.ANTHROPIC, + ServiceType.GOOGLE, + ServiceType.MISTRAL, + ServiceType.LLAMA_CPP, + ServiceType.OLLAMA, + ServiceType.CUSTOM_OPENAI, + ServiceType.INCEPTION, + ) + + fun supportedServiceTypes(): Set = supportedServices + + fun createClient(serviceType: ServiceType, featureType: FeatureType): LLMClient { + return when (serviceType) { + ServiceType.PROXYAI -> { + val apiKey = getCredential(CredentialKey.CodeGptApiKey) ?: "" + val baseUrl = System.getProperty("proxyai.baseUrl") + ?: ProxyAIClientSettings.DEFAULT_BASE_URL + ProxyAILLMClient( + apiKey = apiKey, + settings = ProxyAIClientSettings(baseUrl = baseUrl), + baseClient = HttpClientProvider.createHttpClient() + ) + } + + ServiceType.OPENAI -> { + val apiKey = getCredential(CredentialKey.OpenaiApiKey) ?: "" + val baseUrl = System.getProperty("openai.baseUrl") ?: "https://api.openai.com" + CodeGPTOpenAILLMClient( + apiKey = apiKey, + organization = OpenAISettings.getCurrentState().organization, + baseUrl = baseUrl, + model = resolveOpenAICodeCompletionModel(), + baseClient = HttpClientProvider.createHttpClient() + ) + } + + ServiceType.ANTHROPIC -> { + val apiKey = getCredential(CredentialKey.AnthropicApiKey) ?: "" + val baseUrl = System.getProperty("anthropic.baseUrl") + ?: "https://api.anthropic.com" + val modelVersionsMap = mapOf( + Opus_4_6 to "claude-opus-4-6", + Opus_4_5 to "claude-opus-4-5-20251101", + Sonnet_4_6 to "claude-sonnet-4-6", + Sonnet_4_5 to "claude-sonnet-4-5-20250929", + Haiku_4_5 to "claude-haiku-4-5-20251001", + ) + + AnthropicLLMClient( + apiKey = apiKey, + baseClient = HttpClientProvider.createHttpClient(), + settings = AnthropicClientSettings( + modelVersionsMap = modelVersionsMap, + baseUrl = baseUrl + ) + ) + } + + ServiceType.GOOGLE -> { + val apiKey = getCredential(CredentialKey.GoogleApiKey) ?: "" + val baseUrl = + System.getProperty("google.baseUrl") + ?: "https://generativelanguage.googleapis.com" + GoogleLLMClient( + apiKey = apiKey, + settings = GoogleClientSettings(baseUrl = baseUrl), + baseClient = HttpClientProvider.createHttpClient() + ) + } + + ServiceType.MISTRAL -> { + val apiKey = getCredential(CredentialKey.MistralApiKey) ?: "" + val baseUrl = System.getProperty("mistral.baseUrl") + ?: "https://api.mistral.ai" + CodeGPTMistralAILLMClient( + apiKey = apiKey, + baseUrl = baseUrl, + model = currentCodeCompletionModel(), + baseClient = HttpClientProvider.createHttpClient() + ) + } + + ServiceType.INCEPTION -> { + val apiKey = getCredential(CredentialKey.InceptionApiKey) ?: "" + InceptionAILLMClient(apiKey, baseClient = HttpClientProvider.createHttpClient()) + } + + ServiceType.LLAMA_CPP -> { + val settings = service().state + LlamaCppLLMClient( + baseUrl = System.getProperty("llama.baseUrl") + ?: settings.serverPort?.let { "http://localhost:$it" } + ?: "http://localhost:8080", + state = settings, + baseClient = HttpClientProvider.createHttpClient() + ) + } + + ServiceType.OLLAMA -> { + val settings = service().state + val baseUrl = service().state.host + ?: System.getProperty("ollama.baseUrl") + ?: "http://localhost:11434" + CodeGPTOllamaClient( + settings = settings, + baseUrl = baseUrl, + apiKey = getCredential(CredentialKey.OllamaApikey), + baseClient = HttpClientProvider.createHttpClient() + ) + } + + ServiceType.CUSTOM_OPENAI -> { + val state = + service().customServiceStateForFeatureType(featureType) + val serviceId = state.id ?: error("No custom service configured") + val apiKey = getCredential(CredentialKey.CustomServiceApiKeyById(serviceId)) + ?: error("No API key found for custom service: $serviceId") + + when (featureType) { + FeatureType.CODE_COMPLETION -> CustomOpenAILLMClient.fromCodeCompletionSettingsState( + apiKey, + state.codeCompletionSettings, + HttpClientProvider.createHttpClient() + ) + + else -> CustomOpenAILLMClient.fromSettingsState( + apiKey, + state.chatCompletionSettings, + HttpClientProvider.createHttpClient() + ) + } + } + } + } + + private fun currentCodeCompletionModel(): String { + return service().getModelSelectionForFeature(FeatureType.CODE_COMPLETION).modelId + } + + private fun resolveOpenAICodeCompletionModel(): String { + val current = + service().getModelSelectionForFeature(FeatureType.CODE_COMPLETION) + return if (current.modelId == ModelCatalog.GPT_3_5_TURBO_INSTRUCT) { + current.modelId + } else { + ModelCatalog.GPT_3_5_TURBO_INSTRUCT + } + } +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/models/ModelCatalog.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/models/ModelCatalog.kt new file mode 100644 index 00000000..afde7a58 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/models/ModelCatalog.kt @@ -0,0 +1,110 @@ +package ee.carlrobert.codegpt.settings.models + +import ai.koog.prompt.executor.clients.anthropic.AnthropicModels.Haiku_4_5 +import ai.koog.prompt.llm.LLModel +import com.intellij.openapi.components.Service +import ee.carlrobert.codegpt.settings.service.FeatureType +import ee.carlrobert.codegpt.settings.service.ServiceType +import javax.swing.Icon + +data class ModelSelection( + val provider: ServiceType, + val llmModel: LLModel, + val displayName: String, + val icon: Icon? = null, + val serviceId: String? = null, +) { + val modelId: String + get() = llmModel.id + + val selectionId: String + get() = serviceId ?: modelId + + val model: String + get() = selectionId + + val id: String? + get() = serviceId + + val fullDisplayName: String = if (provider == ServiceType.LLAMA_CPP) { + displayName + } else { + "$provider • $displayName" + } +} + +@Service +class ModelCatalog { + + private val providerRegistry = ModelProviderRegistry() + + private data class DefaultModelRef(val provider: ServiceType, val selectionId: String) + + private val defaultModelRefs = mapOf( + FeatureType.AGENT to DefaultModelRef(ServiceType.PROXYAI, PROXYAI_AUTO), + FeatureType.CHAT to DefaultModelRef(ServiceType.PROXYAI, PROXYAI_AUTO), + FeatureType.AUTO_APPLY to DefaultModelRef(ServiceType.PROXYAI, MERCURY_CODER), + FeatureType.COMMIT_MESSAGE to DefaultModelRef(ServiceType.PROXYAI, Haiku_4_5.id), + FeatureType.INLINE_EDIT to DefaultModelRef(ServiceType.PROXYAI, Haiku_4_5.id), + FeatureType.LOOKUP to DefaultModelRef(ServiceType.PROXYAI, Haiku_4_5.id), + FeatureType.CODE_COMPLETION to DefaultModelRef(ServiceType.PROXYAI, MERCURY_CODER), + FeatureType.NEXT_EDIT to DefaultModelRef(ServiceType.PROXYAI, MERCURY_CODER), + ) + + fun getAllModelsForFeature(featureType: FeatureType): List { + return providerRegistry.modelsForFeature(featureType) + } + + fun getDefaultModelForFeature(featureType: FeatureType): ModelSelection { + val ref = defaultModelRefs.getValue(featureType) + return providerRegistry.findModel(ref.provider, ref.selectionId) + ?: error("Missing default model mapping for feature=$featureType provider=${ref.provider} id=${ref.selectionId}") + } + + fun getProvidersForFeature(featureType: FeatureType): List { + return providerRegistry.providersForFeature(featureType) + } + + fun isFeatureSupportedByProvider(featureType: FeatureType, provider: ServiceType): Boolean { + return providerRegistry.isFeatureSupported(featureType, provider) + } + + fun findModel(provider: ServiceType, modelCode: String): ModelSelection? { + return providerRegistry.findModel(provider, modelCode) + } + + fun getModelDisplayName(provider: ServiceType, modelCode: String?): String { + val code = modelCode?.takeIf { it.isNotBlank() } + return if (code != null) { + findModel(provider, code)?.displayName ?: code + } else { + getAllModels().firstOrNull { it.provider == provider }?.displayName ?: provider.name + } + } + + fun getAllModels(): List { + return providerRegistry.allModels() + } + + companion object { + // ProxyAI Models + const val PROXYAI_AUTO = "auto" + const val GEMINI_FLASH_2_5 = "gemini-flash-2.5" + const val CLAUDE_4_SONNET = "claude-4-sonnet" + const val CLAUDE_4_SONNET_THINKING = "claude-4-sonnet-thinking" + const val CLAUDE_4_5_SONNET = "claude-sonnet-4-5" + const val CLAUDE_4_5_SONNET_THINKING = "claude-sonnet-4-5-thinking" + + // OpenAI Models + const val GPT_3_5_TURBO_INSTRUCT = "gpt-3.5-turbo-instruct" + const val GPT_5_MINI = "gpt-5-mini" + + // Mistral Models + const val CODESTRAL_LATEST = "codestral-latest" + + // Ollama default models + const val LLAMA_3_2 = "llama3.2" + + const val MERCURY_CODER = "mercury-coder" + } +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/models/ModelIcons.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/models/ModelIcons.kt deleted file mode 100644 index 9906fbb6..00000000 --- a/src/main/kotlin/ee/carlrobert/codegpt/settings/models/ModelIcons.kt +++ /dev/null @@ -1,46 +0,0 @@ -package ee.carlrobert.codegpt.settings.models - -import ee.carlrobert.codegpt.Icons -import ee.carlrobert.codegpt.settings.service.ServiceType -import javax.swing.Icon - -object ModelIcons { - fun getIconForModel(model: ModelSelection): Icon? { - return when (model.provider) { - ServiceType.PROXYAI -> getProxyAIModelIcon(model.model) - else -> getIconForProvider(model.provider) - } - } - - fun getProxyAIModelIcon(modelId: String): Icon? { - return when (modelId) { - ModelRegistry.GPT_5_CODEX, ModelRegistry.GPT_5, ModelRegistry.GPT_5_MINI, - ModelRegistry.O4_MINI, ModelRegistry.GPT_4_1, ModelRegistry.GPT_4_1_MINI -> Icons.OpenAI - - ModelRegistry.CLAUDE_4_5_SONNET_THINKING, ModelRegistry.CLAUDE_4_5_SONNET -> Icons.Anthropic - ModelRegistry.GEMINI_PRO_2_5, ModelRegistry.GEMINI_FLASH_2_5 -> Icons.Google - ModelRegistry.DEEPSEEK_R1, ModelRegistry.DEEPSEEK_V3 -> Icons.DeepSeek - ModelRegistry.QWEN3_CODER -> Icons.Qwen - "qwen-2.5-32b-chat", ModelRegistry.QWEN_2_5_32B_CODE -> Icons.Qwen - "llama-3.1-405b" -> Icons.Meta - ModelRegistry.MERCURY_CODER -> Icons.Inception - ModelRegistry.MORPH -> Icons.Morph - ModelRegistry.RELACE -> Icons.Relace - else -> Icons.DefaultSmall - } - } - - fun getIconForProvider(provider: ServiceType): Icon? { - return when (provider) { - ServiceType.PROXYAI -> Icons.DefaultSmall - ServiceType.OPENAI -> Icons.OpenAI - ServiceType.ANTHROPIC -> Icons.Anthropic - ServiceType.GOOGLE -> Icons.Google - ServiceType.MISTRAL -> Icons.Mistral - ServiceType.OLLAMA -> Icons.Ollama - ServiceType.CUSTOM_OPENAI -> Icons.OpenAI - ServiceType.LLAMA_CPP -> Icons.Llama - ServiceType.INCEPTION -> Icons.Inception - } - } -} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/models/ModelProvider.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/models/ModelProvider.kt new file mode 100644 index 00000000..7ad95c70 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/models/ModelProvider.kt @@ -0,0 +1,743 @@ +package ee.carlrobert.codegpt.settings.models + +import ai.koog.prompt.executor.clients.anthropic.AnthropicModels.Haiku_4_5 +import ai.koog.prompt.executor.clients.anthropic.AnthropicModels.Opus_4_5 +import ai.koog.prompt.executor.clients.anthropic.AnthropicModels.Sonnet_4_5 +import ai.koog.prompt.executor.clients.google.GoogleModels +import ai.koog.prompt.executor.clients.google.GoogleModels.Gemini3_Pro_Preview +import ai.koog.prompt.executor.clients.mistralai.MistralAIModels +import ai.koog.prompt.executor.clients.openai.OpenAIModels +import ai.koog.prompt.executor.clients.openai.OpenAIModels.Chat.GPT5Mini +import ai.koog.prompt.executor.clients.openai.OpenAIModels.Chat.GPT5_2 +import ai.koog.prompt.llm.LLMCapability +import ai.koog.prompt.llm.LLMProvider +import ai.koog.prompt.llm.LLModel +import com.intellij.openapi.components.service +import ee.carlrobert.codegpt.Icons +import ee.carlrobert.codegpt.agent.clients.CustomOpenAILLMClient +import ee.carlrobert.codegpt.agent.clients.InceptionAILLMClient +import ee.carlrobert.codegpt.agent.clients.ProxyAILLMClient +import ee.carlrobert.codegpt.completions.llama.LlamaModel +import ee.carlrobert.codegpt.settings.models.ModelCatalog.Companion.MERCURY_CODER +import ee.carlrobert.codegpt.settings.service.FeatureType +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.ollama.OllamaSettings +import javax.swing.Icon + +interface ModelProvider { + val serviceType: ServiceType + val supportedFeatures: Set + fun modelsForFeature(featureType: FeatureType): List +} + +class ModelProviderRegistry( + private val providers: List = defaultModelProviders() +) { + + fun allProviders(): List = providers + + fun provider(serviceType: ServiceType): ModelProvider? = + providers.firstOrNull { it.serviceType == serviceType } + + fun requireProvider(serviceType: ServiceType): ModelProvider { + return provider(serviceType) + ?: throw IllegalArgumentException("No ModelProvider for $serviceType") + } + + fun providersForFeature(featureType: FeatureType): List { + return providers.filter { featureType in it.supportedFeatures }.map { it.serviceType } + } + + fun isFeatureSupported(featureType: FeatureType, serviceType: ServiceType): Boolean { + return featureType in requireProvider(serviceType).supportedFeatures + } + + fun modelsForFeature(featureType: FeatureType): List { + return providers + .asSequence() + .filter { featureType in it.supportedFeatures } + .flatMap { it.modelsForFeature(featureType).asSequence() } + .distinctBy { "${it.provider}:${it.selectionId}" } + .toList() + } + + fun allModels(): List { + return providers + .asSequence() + .flatMap { provider -> + FeatureType.entries.asSequence() + .filter { it in provider.supportedFeatures } + .flatMap { provider.modelsForFeature(it).asSequence() } + } + .distinctBy { "${it.provider}:${it.selectionId}" } + .toList() + } + + fun findModel(serviceType: ServiceType, modelCode: String): ModelSelection? { + return allModels() + .asSequence() + .filter { it.provider == serviceType } + .firstOrNull { it.selectionId == modelCode || it.modelId == modelCode } + } +} + +private val CHAT_FEATURES = setOf( + FeatureType.CHAT, + FeatureType.COMMIT_MESSAGE, + FeatureType.INLINE_EDIT, + FeatureType.LOOKUP, +) + +private fun defaultModelProviders(): List { + return listOf( + ProxyAIModelProvider(), + OpenAIModelProvider(), + AnthropicModelProvider(), + GoogleModelProvider(), + MistralModelProvider(), + LlamaCppModelProvider(), + OllamaModelProvider(), + CustomOpenAIModelProvider(), + InceptionModelProvider(), + ) +} + +private val OPENAI_CAPABILITIES = listOf( + LLMCapability.Temperature, + LLMCapability.Schema.JSON.Basic, + LLMCapability.Schema.JSON.Standard, + LLMCapability.Speculation, + LLMCapability.Tools, + LLMCapability.ToolChoice, + LLMCapability.Vision.Image, + LLMCapability.Document, + LLMCapability.Completion, + LLMCapability.MultipleChoices, + LLMCapability.OpenAIEndpoint.Completions, +) + +private fun virtualModel( + provider: LLMProvider, + modelId: String, + capabilities: List, + contextLength: Long = 200_000, + maxOutputTokens: Long? = 32_768, +): LLModel { + return LLModel( + id = modelId, + provider = provider, + capabilities = capabilities, + contextLength = contextLength, + maxOutputTokens = maxOutputTokens, + ) +} + +private fun formatCustomModelDisplayName(serviceName: String, modelName: String): String { + val safeServiceName = serviceName.trim().ifBlank { ServiceType.CUSTOM_OPENAI.label } + val safeModelName = modelName.trim() + + return if (safeModelName.length > 20) { + "$safeServiceName (...${safeModelName.takeLast(20)})" + } else { + "$safeServiceName ($safeModelName)" + } +} + +private fun defaultModelDisplayName(modelId: String): String { + return when (modelId) { + // ProxyAI + "auto" -> "Auto" + + // OpenAI + GPT5_4.id -> "GPT-5.4" + GPT5_3_Codex.id -> "GPT-5.3 Codex" + OpenAIModels.Chat.GPT5_2Pro.id -> "GPT-5.2 Pro" + GPT5_2_Codex.id -> "GPT-5.2 Codex" + GPT5_2.id -> "GPT-5.2" + OpenAIModels.Chat.GPT5_1Codex.id -> "GPT-5.1 Codex" + OpenAIModels.Chat.GPT5_1.id -> "GPT-5.1" + OpenAIModels.Chat.GPT5Codex.id -> "GPT-5 Codex" + GPT5Mini.id -> "GPT-5 Mini" + OpenAIModels.Chat.GPT4_1.id -> "GPT-4.1" + OpenAIModels.Chat.GPT4_1Mini.id -> "GPT-4.1 Mini" + ModelCatalog.GPT_3_5_TURBO_INSTRUCT -> "GPT-3.5 Turbo Instruct" + + // Anthropic + Opus_4_6.id -> "Claude Opus 4.6" + Opus_4_5.id -> "Claude Opus 4.5" + Sonnet_4_6.id -> "Claude Sonnet 4.6" + Sonnet_4_5.id -> "Claude Sonnet 4.5" + Haiku_4_5.id -> "Claude Haiku 4.5" + + // Google + Gemini3_1_Pro_Preview.id -> "Gemini 3.1 Pro Preview" + Gemini3_Pro_Preview.id -> "Gemini 3 Pro Preview" + Gemini3_Flash_Preview.id -> "Gemini 3 Flash Preview" + GoogleModels.Gemini2_5Pro.id -> "Gemini 2.5 Pro" + GoogleModels.Gemini2_5Flash.id -> "Gemini 2.5 Flash" + GoogleModels.Gemini2_5FlashLite.id -> "Gemini 2.5 Flash Lite" + + // Mistral + Devstral2.id -> "Devstral 2" + MistralAIModels.Chat.DevstralMedium.id -> "Devstral Medium" + MistralAIModels.Chat.Codestral.id -> "Codestral" + + // Inception + Mercury.id -> "Mercury" + MERCURY_CODER -> "Mercury Coder" + + else -> modelId + } +} + +private val Gemini3_Flash_Preview: LLModel = LLModel( + provider = LLMProvider.Google, + id = "gemini-3-flash-preview", + capabilities = listOf( + LLMCapability.Temperature, + LLMCapability.Completion, + LLMCapability.MultipleChoices, + LLMCapability.Vision.Image, + LLMCapability.Vision.Video, + LLMCapability.Audio, + LLMCapability.Tools, + LLMCapability.ToolChoice, + LLMCapability.Schema.JSON.Basic, + LLMCapability.Schema.JSON.Standard, + ), + contextLength = 1_048_576, + maxOutputTokens = 65_536, +) + +private val Gemini3_1_Pro_Preview: LLModel = LLModel( + provider = LLMProvider.Google, + id = "gemini-3.1-pro-preview", + capabilities = listOf( + LLMCapability.Temperature, + LLMCapability.Completion, + LLMCapability.MultipleChoices, + LLMCapability.Vision.Image, + LLMCapability.Vision.Video, + LLMCapability.Audio, + LLMCapability.Tools, + LLMCapability.ToolChoice, + LLMCapability.Schema.JSON.Basic, + LLMCapability.Schema.JSON.Standard, + ), + contextLength = 1_048_576, + maxOutputTokens = 65_536, +) + +val Sonnet_4_6: LLModel = LLModel( + provider = LLMProvider.Anthropic, + id = "claude-sonnet-4-6", + capabilities = listOf( + LLMCapability.Temperature, + LLMCapability.Tools, + LLMCapability.ToolChoice, + LLMCapability.Vision.Image, + LLMCapability.Document, + LLMCapability.Completion + ), + contextLength = 200_000, + maxOutputTokens = 64_000, +) + +val Opus_4_6: LLModel = LLModel( + provider = LLMProvider.Anthropic, + id = "claude-opus-4-6", + capabilities = listOf( + LLMCapability.Temperature, + LLMCapability.Tools, + LLMCapability.ToolChoice, + LLMCapability.Vision.Image, + LLMCapability.Document, + LLMCapability.Completion + ), + contextLength = 200_000, + maxOutputTokens = 64_000, +) + + +private val GPT5_4: LLModel = LLModel( + provider = LLMProvider.OpenAI, + id = "gpt-5.4", + capabilities = listOf( + LLMCapability.Completion, + LLMCapability.Schema.JSON.Basic, + LLMCapability.Schema.JSON.Standard, + LLMCapability.Speculation, + LLMCapability.Tools, + LLMCapability.ToolChoice, + LLMCapability.Vision.Image, + LLMCapability.Document, + LLMCapability.MultipleChoices, + LLMCapability.OpenAIEndpoint.Responses, + ), + contextLength = 1_050_000, + maxOutputTokens = 128_000, +) + +private val GPT5_3_Codex: LLModel = LLModel( + provider = LLMProvider.OpenAI, + id = "gpt-5.3-codex", + capabilities = listOf( + LLMCapability.Completion, + LLMCapability.Schema.JSON.Basic, + LLMCapability.Schema.JSON.Standard, + LLMCapability.Speculation, + LLMCapability.Tools, + LLMCapability.ToolChoice, + LLMCapability.Vision.Image, + LLMCapability.Document, + LLMCapability.MultipleChoices, + LLMCapability.OpenAIEndpoint.Responses, + ), + contextLength = 400_000, + maxOutputTokens = 128_000, +) + +private val GPT5_2_Codex: LLModel = LLModel( + provider = LLMProvider.OpenAI, + id = "gpt-5.2-codex", + capabilities = listOf( + LLMCapability.Completion, + LLMCapability.Schema.JSON.Basic, + LLMCapability.Schema.JSON.Standard, + LLMCapability.Speculation, + LLMCapability.Tools, + LLMCapability.ToolChoice, + LLMCapability.Vision.Image, + LLMCapability.Document, + LLMCapability.MultipleChoices, + LLMCapability.OpenAIEndpoint.Responses, + ), + contextLength = 400_000, + maxOutputTokens = 128_000, +) + +val Devstral2: LLModel = LLModel( + provider = LLMProvider.MistralAI, + id = "devstral-2512", + capabilities = listOf( + LLMCapability.Temperature, + LLMCapability.Completion, + LLMCapability.Tools, + LLMCapability.ToolChoice, + LLMCapability.Schema.JSON.Basic, + LLMCapability.Schema.JSON.Standard, + ), + contextLength = 256_000, + maxOutputTokens = 65_536, +) + +val Mercury: LLModel = LLModel( + provider = InceptionAILLMClient.Inception, + id = "mercury", + capabilities = listOf( + LLMCapability.Temperature, + LLMCapability.Schema.JSON.Basic, + LLMCapability.Schema.JSON.Standard, + LLMCapability.Speculation, + LLMCapability.Tools, + LLMCapability.ToolChoice, + LLMCapability.Vision.Image, + LLMCapability.Document, + LLMCapability.Completion, + LLMCapability.MultipleChoices, + LLMCapability.OpenAIEndpoint.Completions, + ), + contextLength = 200_000, + maxOutputTokens = 32_768, +) + +private class ProxyAIModelProvider : ModelProvider { + override val serviceType: ServiceType = ServiceType.PROXYAI + override val supportedFeatures: Set = setOf( + FeatureType.AGENT, + FeatureType.CHAT, + FeatureType.CODE_COMPLETION, + FeatureType.AUTO_APPLY, + FeatureType.COMMIT_MESSAGE, + FeatureType.INLINE_EDIT, + FeatureType.NEXT_EDIT, + FeatureType.LOOKUP, + ) + + override fun modelsForFeature(featureType: FeatureType): List { + return when (featureType) { + in (CHAT_FEATURES + FeatureType.AGENT) -> listOf( + defaultModel(ModelCatalog.PROXYAI_AUTO, Icons.DefaultSmall), + defaultModel(GPT5_4.id, Icons.OpenAI), + defaultModel(GPT5_3_Codex.id, Icons.OpenAI), + defaultModel(GPT5Mini.id, Icons.OpenAI), + defaultModel(Opus_4_6.id, Icons.Anthropic), + defaultModel(Sonnet_4_6.id, Icons.Anthropic), + defaultModel(Haiku_4_5.id, Icons.Anthropic), + defaultModel(Gemini3_1_Pro_Preview.id, Icons.Google), + defaultModel(Gemini3_Flash_Preview.id, Icons.Google), + ) + + FeatureType.AUTO_APPLY -> listOf( + defaultModel(MERCURY_CODER, Icons.Inception), + ) + + FeatureType.CODE_COMPLETION -> listOf( + defaultModel(MERCURY_CODER, Icons.Inception) + ) + + FeatureType.NEXT_EDIT -> listOf( + defaultModel(MERCURY_CODER) + ) + + else -> emptyList() + } + } + + private fun defaultModel( + modelId: String, + icon: Icon? = null, + ): ModelSelection { + return ModelSelection( + provider = serviceType, + llmModel = virtualModel(ProxyAILLMClient.ProxyAI, modelId, OPENAI_CAPABILITIES), + displayName = defaultModelDisplayName(modelId), + icon = icon, + ) + } +} + +private class OpenAIModelProvider : ModelProvider { + override val serviceType: ServiceType = ServiceType.OPENAI + override val supportedFeatures: Set = setOf( + FeatureType.AGENT, + FeatureType.CHAT, + FeatureType.CODE_COMPLETION, + FeatureType.AUTO_APPLY, + FeatureType.COMMIT_MESSAGE, + FeatureType.INLINE_EDIT, + FeatureType.LOOKUP, + ) + + private val models = listOf( + GPT5_4, + GPT5_3_Codex, + GPT5_2_Codex, + OpenAIModels.Chat.GPT5_2Pro, + GPT5_2, + OpenAIModels.Chat.GPT5_1Codex, + OpenAIModels.Chat.GPT5_1, + OpenAIModels.Chat.GPT5Codex, + GPT5Mini, + OpenAIModels.Chat.GPT4_1, + OpenAIModels.Chat.GPT4_1Mini, + ) + + override fun modelsForFeature(featureType: FeatureType): List { + return when (featureType) { + in CHAT_FEATURES, FeatureType.AGENT, FeatureType.AUTO_APPLY -> + models.map { defaultModel(it) } + + FeatureType.CODE_COMPLETION -> listOf( + defaultModel( + virtualModel( + LLMProvider.OpenAI, + ModelCatalog.GPT_3_5_TURBO_INSTRUCT, + OPENAI_CAPABILITIES + ) + ) + ) + + else -> emptyList() + } + } + + private fun defaultModel(llmModel: LLModel): ModelSelection { + return ModelSelection( + provider = serviceType, + llmModel = llmModel, + displayName = defaultModelDisplayName(llmModel.id), + ) + } +} + +private class AnthropicModelProvider : ModelProvider { + override val serviceType: ServiceType = ServiceType.ANTHROPIC + override val supportedFeatures: Set = setOf( + FeatureType.AGENT, + FeatureType.CHAT, + FeatureType.AUTO_APPLY, + FeatureType.COMMIT_MESSAGE, + FeatureType.INLINE_EDIT, + FeatureType.LOOKUP, + ) + + private val models = listOf( + Opus_4_6, + Opus_4_5, + Sonnet_4_6, + Sonnet_4_5, + Haiku_4_5, + ) + + override fun modelsForFeature(featureType: FeatureType): List { + return when (featureType) { + FeatureType.AGENT, in CHAT_FEATURES, FeatureType.AUTO_APPLY -> models.map { model -> + ModelSelection( + provider = serviceType, + llmModel = model, + displayName = defaultModelDisplayName(model.id), + ) + } + + else -> emptyList() + } + } +} + +private class GoogleModelProvider : ModelProvider { + override val serviceType: ServiceType = ServiceType.GOOGLE + override val supportedFeatures: Set = setOf( + FeatureType.AGENT, + FeatureType.CHAT, + FeatureType.AUTO_APPLY, + FeatureType.COMMIT_MESSAGE, + FeatureType.INLINE_EDIT, + FeatureType.LOOKUP, + ) + + private val chatModels = listOf( + Gemini3_1_Pro_Preview, + Gemini3_Pro_Preview, + Gemini3_Flash_Preview, + ) + + override fun modelsForFeature(featureType: FeatureType): List { + return when (featureType) { + FeatureType.AGENT -> listOf( + defaultModel(Gemini3_1_Pro_Preview), + defaultModel(Gemini3_Pro_Preview), + defaultModel(Gemini3_Flash_Preview), + ) + + in CHAT_FEATURES, FeatureType.AUTO_APPLY -> chatModels.map { model -> + ModelSelection( + provider = serviceType, + llmModel = model, + displayName = defaultModelDisplayName(model.id), + ) + } + + else -> emptyList() + } + } + + private fun defaultModel(llmModel: LLModel): ModelSelection { + return ModelSelection( + provider = serviceType, + llmModel = llmModel, + displayName = defaultModelDisplayName(llmModel.id), + ) + } +} + +private class MistralModelProvider : ModelProvider { + override val serviceType: ServiceType = ServiceType.MISTRAL + override val supportedFeatures: Set = setOf( + FeatureType.AGENT, + FeatureType.CHAT, + FeatureType.CODE_COMPLETION, + FeatureType.AUTO_APPLY, + FeatureType.COMMIT_MESSAGE, + FeatureType.INLINE_EDIT, + FeatureType.LOOKUP, + ) + + private val chatModels = listOf( + Devstral2, + MistralAIModels.Chat.DevstralMedium, + MistralAIModels.Chat.Codestral, + ) + + override fun modelsForFeature(featureType: FeatureType): List { + return when (featureType) { + FeatureType.AGENT -> listOf( + defaultModel(Devstral2), + defaultModel(MistralAIModels.Chat.DevstralMedium), + ) + + in CHAT_FEATURES, FeatureType.AUTO_APPLY -> chatModels.map { defaultModel(it) } + + FeatureType.CODE_COMPLETION -> listOf(defaultModel(MistralAIModels.Chat.Codestral)) + + else -> emptyList() + } + } + + private fun defaultModel(llmModel: LLModel): ModelSelection { + return ModelSelection( + provider = serviceType, + llmModel = llmModel, + displayName = defaultModelDisplayName(llmModel.id), + ) + } +} + +private class InceptionModelProvider : ModelProvider { + override val serviceType: ServiceType = ServiceType.INCEPTION + override val supportedFeatures: Set = setOf( + FeatureType.AGENT, + FeatureType.CHAT, + FeatureType.CODE_COMPLETION, + FeatureType.AUTO_APPLY, + FeatureType.COMMIT_MESSAGE, + FeatureType.INLINE_EDIT, + FeatureType.LOOKUP, + FeatureType.NEXT_EDIT, + ) + + override fun modelsForFeature(featureType: FeatureType): List { + return when (featureType) { + FeatureType.AGENT -> listOf(inceptionMercury()) + in CHAT_FEATURES, FeatureType.AUTO_APPLY -> listOf( + inceptionMercury(), + inceptionMercuryCoder(), + ) + + FeatureType.CODE_COMPLETION -> listOf(inceptionMercuryCoder()) + FeatureType.NEXT_EDIT -> listOf(inceptionMercuryCoder()) + else -> emptyList() + } + } + + private fun inceptionMercury(): ModelSelection { + return ModelSelection( + provider = serviceType, + llmModel = Mercury, + displayName = defaultModelDisplayName(Mercury.id), + ) + } + + private fun inceptionMercuryCoder(): ModelSelection { + return ModelSelection( + provider = serviceType, + llmModel = virtualModel( + InceptionAILLMClient.Inception, + MERCURY_CODER, + OPENAI_CAPABILITIES + ), + displayName = "Mercury Coder", + ) + } +} + +private class LlamaCppModelProvider : ModelProvider { + override val serviceType: ServiceType = ServiceType.LLAMA_CPP + override val supportedFeatures: Set = setOf( + FeatureType.CHAT, + FeatureType.CODE_COMPLETION, + FeatureType.AUTO_APPLY, + FeatureType.COMMIT_MESSAGE, + FeatureType.INLINE_EDIT, + FeatureType.LOOKUP, + ) + + override fun modelsForFeature(featureType: FeatureType): List { + return LlamaModel.entries.flatMap { llamaModel -> + llamaModel.huggingFaceModels.map { hfModel -> + val displayName = + "${llamaModel.label} (${hfModel.parameterSize}B) / Q${hfModel.quantization}" + ModelSelection( + provider = serviceType, + llmModel = virtualModel( + LLMProvider.OpenAI, + hfModel.name, + OPENAI_CAPABILITIES + ), + displayName = displayName, + ) + } + } + } +} + +private class OllamaModelProvider : ModelProvider { + override val serviceType: ServiceType = ServiceType.OLLAMA + override val supportedFeatures: Set = setOf( + FeatureType.AGENT, + FeatureType.CHAT, + FeatureType.CODE_COMPLETION, + FeatureType.AUTO_APPLY, + FeatureType.COMMIT_MESSAGE, + FeatureType.INLINE_EDIT, + FeatureType.LOOKUP, + ) + + private fun availableModels(): List { + return runCatching { service().state.availableModels }.getOrDefault( + emptyList() + ) + } + + override fun modelsForFeature(featureType: FeatureType): List { + return availableModels().map { model -> + ModelSelection( + provider = serviceType, + llmModel = virtualModel(LLMProvider.Ollama, model, OPENAI_CAPABILITIES), + displayName = model, + ) + } + } +} + +private class CustomOpenAIModelProvider : ModelProvider { + override val serviceType: ServiceType = ServiceType.CUSTOM_OPENAI + override val supportedFeatures: Set = setOf( + FeatureType.AGENT, + FeatureType.CHAT, + FeatureType.CODE_COMPLETION, + FeatureType.AUTO_APPLY, + FeatureType.COMMIT_MESSAGE, + FeatureType.INLINE_EDIT, + FeatureType.LOOKUP, + ) + + private fun buildModels( + modelExtractor: (CustomServiceSettingsState) -> String? + ): List { + return runCatching { + service().state.services.mapNotNull { svc -> + val serviceId = svc.id ?: return@mapNotNull null + val serviceName = svc.name ?: "" + val modelName = + modelExtractor(svc)?.takeIf { it.isNotBlank() } ?: return@mapNotNull null + val displayName = formatCustomModelDisplayName(serviceName, modelName) + + ModelSelection( + provider = serviceType, + llmModel = virtualModel( + CustomOpenAILLMClient.CustomOpenAI, + modelName, + OPENAI_CAPABILITIES + ), + displayName = displayName, + serviceId = serviceId, + ) + } + }.getOrDefault(emptyList()) + } + + override fun modelsForFeature(featureType: FeatureType): List { + return when (featureType) { + FeatureType.CODE_COMPLETION -> buildModels { svc -> + svc.codeCompletionSettings.body["model"] as? String + } + + else -> buildModels { svc -> + svc.chatCompletionSettings.body["model"] as? String + } + } + } +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/models/ModelRegistry.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/models/ModelRegistry.kt deleted file mode 100644 index 31a91016..00000000 --- a/src/main/kotlin/ee/carlrobert/codegpt/settings/models/ModelRegistry.kt +++ /dev/null @@ -1,1071 +0,0 @@ -package ee.carlrobert.codegpt.settings.models - -import ai.koog.prompt.executor.clients.anthropic.AnthropicModels.Haiku_4_5 -import ai.koog.prompt.executor.clients.anthropic.AnthropicModels.Opus_4_5 -import ai.koog.prompt.executor.clients.anthropic.AnthropicModels.Sonnet_4_5 -import ai.koog.prompt.executor.clients.google.GoogleModels.Gemini2_5Flash -import ai.koog.prompt.executor.clients.google.GoogleModels.Gemini2_5Pro -import ai.koog.prompt.executor.clients.google.GoogleModels.Gemini3_Pro_Preview -import ai.koog.prompt.executor.clients.mistralai.MistralAIModels.Chat.DevstralMedium -import ai.koog.prompt.executor.clients.openai.OpenAIModels.Chat.GPT4_1 -import ai.koog.prompt.executor.clients.openai.OpenAIModels.Chat.GPT4_1Mini -import ai.koog.prompt.executor.clients.openai.OpenAIModels.Chat.GPT5Mini -import ai.koog.prompt.executor.clients.openai.OpenAIModels.Chat.GPT5_1 -import ai.koog.prompt.executor.clients.openai.OpenAIModels.Chat.GPT5_1Codex -import ai.koog.prompt.executor.clients.openai.OpenAIModels.Chat.GPT5_2 -import ai.koog.prompt.llm.LLMCapability -import ai.koog.prompt.llm.LLMProvider -import ai.koog.prompt.llm.LLModel -import com.intellij.openapi.application.ApplicationManager -import com.intellij.openapi.components.Service -import com.intellij.openapi.components.service -import com.intellij.openapi.diagnostic.thisLogger -import ee.carlrobert.codegpt.Icons -import ee.carlrobert.codegpt.agent.clients.CustomOpenAILLMClient -import ee.carlrobert.codegpt.agent.clients.InceptionAILLMClient -import ee.carlrobert.codegpt.agent.clients.ProxyAILLMClient -import ee.carlrobert.codegpt.completions.llama.LlamaModel -import ee.carlrobert.codegpt.settings.models.ModelRegistry.Companion.MERCURY -import ee.carlrobert.codegpt.settings.service.FeatureType -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.llm.client.codegpt.PricingPlan -import ee.carlrobert.llm.client.google.models.GoogleModel -import ee.carlrobert.llm.client.openai.completion.OpenAIChatCompletionModel -import javax.swing.Icon - -data class ModelSelection( - val provider: ServiceType, - val model: String, - val displayName: String, - val icon: Icon? = null, - val pricingPlan: PricingPlan? = null, - val llmModel: LLModel? = null, - val id: String? = "", -) { - val fullDisplayName: String = if (provider == ServiceType.LLAMA_CPP) { - displayName - } else { - "$provider • $displayName" - } -} - -data class ModelCapability( - val provider: ServiceType, - val supportedFeatures: Set, - val requiresPricingPlan: PricingPlan? = null -) - -@Service -class ModelRegistry { - - private val logger = thisLogger() - - private val providerCapabilities = mapOf( - ServiceType.PROXYAI to ModelCapability( - ServiceType.PROXYAI, - setOf( - FeatureType.AGENT, - FeatureType.CHAT, - FeatureType.CODE_COMPLETION, - FeatureType.AUTO_APPLY, - FeatureType.COMMIT_MESSAGE, - FeatureType.INLINE_EDIT, - FeatureType.NEXT_EDIT, - FeatureType.LOOKUP - ) - ), - ServiceType.OPENAI to ModelCapability( - ServiceType.OPENAI, - setOf( - FeatureType.AGENT, - FeatureType.CHAT, - FeatureType.CODE_COMPLETION, - FeatureType.AUTO_APPLY, - FeatureType.COMMIT_MESSAGE, - FeatureType.INLINE_EDIT, - FeatureType.LOOKUP - ) - ), - ServiceType.ANTHROPIC to ModelCapability( - ServiceType.ANTHROPIC, - setOf( - FeatureType.AGENT, - FeatureType.CHAT, - FeatureType.AUTO_APPLY, - FeatureType.COMMIT_MESSAGE, - FeatureType.INLINE_EDIT, - FeatureType.LOOKUP - ) - ), - ServiceType.GOOGLE to ModelCapability( - ServiceType.GOOGLE, - setOf( - FeatureType.AGENT, - FeatureType.CHAT, - FeatureType.AUTO_APPLY, - FeatureType.COMMIT_MESSAGE, - FeatureType.INLINE_EDIT, - FeatureType.LOOKUP - ) - ), - ServiceType.MISTRAL to ModelCapability( - ServiceType.MISTRAL, - setOf( - FeatureType.AGENT, - FeatureType.CHAT, - FeatureType.CODE_COMPLETION, - FeatureType.AUTO_APPLY, - FeatureType.COMMIT_MESSAGE, - FeatureType.INLINE_EDIT, - FeatureType.LOOKUP - ) - ), - ServiceType.OLLAMA to ModelCapability( - ServiceType.OLLAMA, - setOf( - FeatureType.AGENT, - FeatureType.CHAT, - FeatureType.CODE_COMPLETION, - FeatureType.AUTO_APPLY, - FeatureType.COMMIT_MESSAGE, - FeatureType.INLINE_EDIT, - FeatureType.LOOKUP - ) - ), - ServiceType.LLAMA_CPP to ModelCapability( - ServiceType.LLAMA_CPP, - setOf( - FeatureType.CHAT, FeatureType.CODE_COMPLETION, FeatureType.AUTO_APPLY, - FeatureType.COMMIT_MESSAGE, FeatureType.INLINE_EDIT, FeatureType.LOOKUP - ) - ), - ServiceType.CUSTOM_OPENAI to ModelCapability( - ServiceType.CUSTOM_OPENAI, - setOf( - FeatureType.AGENT, - FeatureType.CHAT, - FeatureType.CODE_COMPLETION, - FeatureType.AUTO_APPLY, - FeatureType.COMMIT_MESSAGE, - FeatureType.INLINE_EDIT, - FeatureType.LOOKUP - ) - ), - ServiceType.INCEPTION to ModelCapability( - ServiceType.INCEPTION, - setOf( - FeatureType.AGENT, - FeatureType.CHAT, - FeatureType.CODE_COMPLETION, - FeatureType.AUTO_APPLY, - FeatureType.COMMIT_MESSAGE, - FeatureType.INLINE_EDIT, - FeatureType.LOOKUP, - FeatureType.NEXT_EDIT - ) - ) - ) - - private val pricingPlanBasedDefaults = mapOf( - PricingPlan.ANONYMOUS to mapOf( - FeatureType.AGENT to ModelSelection( - ServiceType.PROXYAI, - PROXYAI_AUTO, - "Auto" - ), - FeatureType.CHAT to ModelSelection( - ServiceType.PROXYAI, - GEMINI_FLASH_2_5, - "Gemini Flash 2.5" - ), - FeatureType.AUTO_APPLY to ModelSelection( - ServiceType.PROXYAI, - MERCURY_CODER, - "Mercury Coder" - ), - FeatureType.COMMIT_MESSAGE to ModelSelection( - ServiceType.PROXYAI, - GPT_5_MINI, - "GPT-5 Mini" - ), - FeatureType.INLINE_EDIT to ModelSelection( - ServiceType.PROXYAI, - GEMINI_FLASH_2_5, - "Gemini Flash 2.5" - ), - FeatureType.LOOKUP to ModelSelection( - ServiceType.PROXYAI, - GPT_5_MINI, - "GPT-5 Mini" - ), - FeatureType.CODE_COMPLETION to ModelSelection( - ServiceType.PROXYAI, - MERCURY_CODER, - "Mercury Coder" - ), - FeatureType.NEXT_EDIT to ModelSelection( - ServiceType.PROXYAI, - MERCURY_CODER, - "Mercury Coder" - ) - ), - PricingPlan.FREE to mapOf( - FeatureType.AGENT to ModelSelection( - ServiceType.PROXYAI, - PROXYAI_AUTO, - "Auto" - ), - FeatureType.CHAT to ModelSelection(ServiceType.PROXYAI, QWEN3_CODER, "Qwen3 Coder"), - FeatureType.AUTO_APPLY to ModelSelection( - ServiceType.PROXYAI, - MERCURY_CODER, - "Mercury Coder" - ), - FeatureType.COMMIT_MESSAGE to ModelSelection( - ServiceType.PROXYAI, - QWEN3_CODER, - "Qwen3 Coder" - ), - FeatureType.INLINE_EDIT to ModelSelection( - ServiceType.PROXYAI, - QWEN3_CODER, - "Qwen3 Coder" - ), - FeatureType.LOOKUP to ModelSelection(ServiceType.PROXYAI, QWEN3_CODER, "Qwen3 Coder"), - FeatureType.CODE_COMPLETION to ModelSelection( - ServiceType.PROXYAI, - MERCURY_CODER, - "Mercury Coder" - ), - FeatureType.NEXT_EDIT to ModelSelection( - ServiceType.PROXYAI, - MERCURY_CODER, - "Mercury Coder" - ) - ), - PricingPlan.INDIVIDUAL to mapOf( - FeatureType.AGENT to ModelSelection( - ServiceType.PROXYAI, - PROXYAI_AUTO, - "Auto" - ), - FeatureType.CHAT to ModelSelection( - ServiceType.PROXYAI, - CLAUDE_4_5_SONNET_THINKING, - "Claude 4.5 Sonnet Thinking" - ), - FeatureType.AUTO_APPLY to ModelSelection( - ServiceType.PROXYAI, - MERCURY_CODER, - "Mercury Coder" - ), - FeatureType.COMMIT_MESSAGE to ModelSelection( - ServiceType.PROXYAI, - GPT_5_CODEX, - "GPT-5 Codex" - ), - FeatureType.INLINE_EDIT to ModelSelection( - ServiceType.PROXYAI, - CLAUDE_4_5_SONNET, - "Claude 4 Sonnet" - ), - FeatureType.LOOKUP to ModelSelection(ServiceType.PROXYAI, GPT_5_CODEX, "GPT-5 Codex"), - FeatureType.CODE_COMPLETION to ModelSelection( - ServiceType.PROXYAI, - MERCURY_CODER, - "Mercury Coder" - ), - FeatureType.NEXT_EDIT to ModelSelection( - ServiceType.PROXYAI, - MERCURY_CODER, - "Mercury Coder" - ) - ) - ) - - private val fallbackDefaults = mapOf( - FeatureType.AGENT to ModelSelection( - ServiceType.PROXYAI, - PROXYAI_AUTO, - "Auto" - ), - FeatureType.CHAT to ModelSelection( - ServiceType.PROXYAI, - GEMINI_FLASH_2_5, - "Gemini Flash 2.5" - ), - FeatureType.AUTO_APPLY to ModelSelection( - ServiceType.PROXYAI, - MERCURY_CODER, - "Mercury Coder" - ), - FeatureType.COMMIT_MESSAGE to ModelSelection( - ServiceType.PROXYAI, - GPT_5_MINI, - "GPT-5 Mini" - ), - FeatureType.INLINE_EDIT to ModelSelection( - ServiceType.PROXYAI, - GEMINI_FLASH_2_5, - "Gemini Flash 2.5" - ), - FeatureType.LOOKUP to ModelSelection(ServiceType.PROXYAI, GPT_5_MINI, "GPT-5 Mini"), - FeatureType.CODE_COMPLETION to ModelSelection( - ServiceType.PROXYAI, - MERCURY_CODER, - "Mercury Coder" - ), - FeatureType.NEXT_EDIT to ModelSelection(ServiceType.PROXYAI, MERCURY_CODER, "Mercury Coder") - ) - - fun getAllModelsForFeature(featureType: FeatureType): List { - return when (featureType) { - FeatureType.AGENT -> getAllAgentModels() - FeatureType.CHAT, FeatureType.COMMIT_MESSAGE, - FeatureType.INLINE_EDIT, FeatureType.LOOKUP -> getAllChatModels() - - FeatureType.AUTO_APPLY -> getAllApplyModels() - FeatureType.CODE_COMPLETION -> getAllCodeModels() - FeatureType.NEXT_EDIT -> getNextEditModels() - } - } - - fun getDefaultModelForFeature( - featureType: FeatureType, - pricingPlan: PricingPlan? = null - ): ModelSelection { - val planBasedDefaults = pricingPlan?.let { pricingPlanBasedDefaults[it] } - return planBasedDefaults?.get(featureType) ?: fallbackDefaults[featureType]!! - } - - fun getProvidersForFeature(featureType: FeatureType): List { - return providerCapabilities.values - .filter { it.supportedFeatures.contains(featureType) } - .map { it.provider } - } - - fun isFeatureSupportedByProvider(featureType: FeatureType, provider: ServiceType): Boolean { - return providerCapabilities[provider]?.supportedFeatures?.contains(featureType) == true - } - - fun findModel(provider: ServiceType, modelCode: String): ModelSelection? { - return getAllModels() - .filter { it.provider == provider } - .find { it.id == modelCode || it.model == modelCode } - } - - fun getModelDisplayName(provider: ServiceType, modelCode: String?): String { - val code = modelCode?.takeIf { it.isNotBlank() } - return if (code != null) { - findModel(provider, code)?.displayName ?: code - } else { - // Fallback to the first known model for the provider, or provider name - getAllModels().firstOrNull { it.provider == provider }?.displayName ?: provider.name - } - } - - private fun getAllModels(): List { - return buildList { - addAll(getAllAgentModels()) - addAll(getAllApplyModels()) - addAll(getAllChatModels()) - addAll(getAllCodeModels()) - addAll(getNextEditModels()) - }.distinctBy { "${it.provider}:${it.model}" } - } - - fun getAgentModels(serviceType: ServiceType?): List { - return getAgentModels()[serviceType] ?: emptyList() - } - - data class LLMModelWrapper(val model: LLModel, val id: String = "", val name: String = "") - - private fun getProxyAIModelWrapper(id: String, name: String): LLMModelWrapper { - return LLMModelWrapper( - model = LLModel( - id = id, - provider = ProxyAILLMClient.ProxyAI, - capabilities = listOf( - LLMCapability.Temperature, - LLMCapability.Schema.JSON.Basic, - LLMCapability.Schema.JSON.Standard, - LLMCapability.Speculation, - LLMCapability.Tools, - LLMCapability.ToolChoice, - LLMCapability.Vision.Image, - LLMCapability.Document, - LLMCapability.Completion, - LLMCapability.MultipleChoices, - LLMCapability.OpenAIEndpoint.Completions, - ), - contextLength = 200_000, - maxOutputTokens = 32_768, - ), - name = name - ) - } - - fun getAgentModels(): Map> { - return mapOf( - ServiceType.PROXYAI to listOf( - getProxyAIModelWrapper(PROXYAI_AUTO, "Auto"), - getProxyAIModelWrapper(GPT5_2.id, "GPT-5.2"), - getProxyAIModelWrapper(GPT5_2_Codex.id, "GPT-5.2 Codex"), - getProxyAIModelWrapper(GPT5Mini.id, "GPT-5 Mini"), - getProxyAIModelWrapper(Opus_4_5.id, "Claude Opus 4.5"), - getProxyAIModelWrapper(Sonnet_4_5.id, "Claude Sonnet 4.5"), - getProxyAIModelWrapper(Haiku_4_5.id, "Claude Haiku 4.5"), - ), - ServiceType.OPENAI to listOf( - LLMModelWrapper(GPT5_2, name = "GPT-5.2"), - LLMModelWrapper(GPT5_2_Codex, name = "GPT-5.2 Codex"), - LLMModelWrapper(GPT5_1, name = "GPT-5.1"), - LLMModelWrapper(GPT5_1Codex, name = "GPT-5.1 Codex"), - LLMModelWrapper(GPT5Mini, name = "GPT-5 Mini"), - LLMModelWrapper(GPT4_1, name = "GPT-4.1"), - LLMModelWrapper(GPT4_1Mini, name = "GPT-4.1 Mini") - ), - ServiceType.ANTHROPIC to listOf( - LLMModelWrapper(Opus_4_5, name = "Claude Opus 4.5"), - LLMModelWrapper(Sonnet_4_5, name = "Claude Sonnet 4.5"), - LLMModelWrapper(Haiku_4_5, name = "Claude Haiku 4.5") - ), - ServiceType.CUSTOM_OPENAI to getCustomOpenAIModels().map { model -> - LLMModelWrapper( - model = LLModel( - id = model.model, - provider = CustomOpenAILLMClient.CustomOpenAI, - capabilities = listOf( - LLMCapability.Temperature, - LLMCapability.Schema.JSON.Basic, - LLMCapability.Schema.JSON.Standard, - LLMCapability.Speculation, - LLMCapability.Tools, - LLMCapability.ToolChoice, - LLMCapability.Vision.Image, - LLMCapability.Document, - LLMCapability.Completion, - LLMCapability.MultipleChoices, - LLMCapability.OpenAIEndpoint.Completions, - ), - contextLength = 200_000, - maxOutputTokens = 32_768, - ), - id = model.id ?: "", - name = model.displayName - ) - }, - ServiceType.GOOGLE to listOf( - LLMModelWrapper(Gemini3_Pro_Preview, name = "Gemini 3 Pro Preview"), - LLMModelWrapper(Gemini3_Flash_Preview, name = "Gemini 3 Flash Preview"), - LLMModelWrapper(Gemini2_5Pro, name = "Gemini 2.5 Pro"), - LLMModelWrapper(Gemini2_5Flash, name = "Gemini 2.5 Flash") - ), - ServiceType.MISTRAL to listOf( - LLMModelWrapper(Devstral2, name = "Devstral 2"), - LLMModelWrapper(DevstralMedium, name = "Devstral Medium"), - ), - ServiceType.INCEPTION to listOf( - LLMModelWrapper(model = Mercury, name = "Mercury"), - ), - ServiceType.OLLAMA to getOllamaModels().map { model -> - LLMModelWrapper( - LLModel( - id = model.model, - provider = LLMProvider.Ollama, - capabilities = listOf( - LLMCapability.Temperature, - LLMCapability.Schema.JSON.Basic, - LLMCapability.Schema.JSON.Standard, - LLMCapability.Speculation, - LLMCapability.Tools, - LLMCapability.ToolChoice, - LLMCapability.Vision.Image, - LLMCapability.Document, - LLMCapability.Completion, - LLMCapability.MultipleChoices, - LLMCapability.OpenAIEndpoint.Completions, - ), - contextLength = 200_000, - maxOutputTokens = 32_768, - ), - name = model.displayName - ) - } - ) - } - - private fun getAllAgentModels(): List { - return buildList { - addAll(getProxyAIAgentModels()) - addAll(getAnthropicAgentModels()) - addAll(getOpenAIAgentModels()) - addAll(getCustomOpenAIModels()) - addAll(getGoogleAgentModels()) - addAll(getMistralAgentModels()) - addAll(getOllamaModels()) - addAll(getInceptionAgentModels()) - } - } - - fun getOpenAIAgentModels(): List { - return listOf( - ModelSelection(ServiceType.OPENAI, GPT5_2.id, "GPT-5.2", Icons.OpenAI), - ModelSelection(ServiceType.OPENAI, GPT5_2_Codex.id, "GPT-5.2 Codex", Icons.OpenAI), - ModelSelection(ServiceType.OPENAI, GPT5_1.id, "GPT-5.1", Icons.OpenAI), - ModelSelection(ServiceType.OPENAI, GPT5_1Codex.id, "GPT-5.1 Codex", Icons.OpenAI), - ModelSelection(ServiceType.OPENAI, GPT5Mini.id, "GPT-5 Mini", Icons.OpenAI), - ) - } - - private fun getProxyAIAgentModels(): List { - return listOf( - ModelSelection(ServiceType.PROXYAI, PROXYAI_AUTO, "Auto", Icons.DefaultSmall), - ModelSelection(ServiceType.PROXYAI, GPT5_2.id, "GPT-5.2", Icons.OpenAI), - ModelSelection(ServiceType.PROXYAI, GPT5_2_Codex.id, "GPT-5.2 Codex", Icons.OpenAI), - ModelSelection(ServiceType.PROXYAI, GPT5Mini.id, "GPT-5 Mini", Icons.OpenAI), - ModelSelection(ServiceType.PROXYAI, Opus_4_5.id, "Claude Opus 4.5", Icons.Anthropic), - ModelSelection( - ServiceType.PROXYAI, - Sonnet_4_5.id, - "Claude Sonnet 4.5", - Icons.Anthropic - ), - ModelSelection( - ServiceType.PROXYAI, - Haiku_4_5.id, - "Claude Haiku 4.5", - Icons.Anthropic - ), - ) - } - - fun getAnthropicAgentModels(): List { - return listOf( - ModelSelection(ServiceType.ANTHROPIC, Opus_4_5.id, "Claude Opus 4.5", Icons.Anthropic), - ModelSelection( - ServiceType.ANTHROPIC, - Sonnet_4_5.id, - "Claude Sonnet 4.5", - Icons.Anthropic - ), - ModelSelection( - ServiceType.ANTHROPIC, - Haiku_4_5.id, - "Claude Haiku 4.5", - Icons.Anthropic - ), - ) - } - - fun getGoogleAgentModels(): List { - return listOf( - ModelSelection( - ServiceType.GOOGLE, - Gemini3_Pro_Preview.id, - "Gemini 3 Pro Preview", - Icons.Google - ), - ModelSelection( - ServiceType.GOOGLE, - Gemini3_Flash_Preview.id, - "Gemini 3 Flash Preview", - Icons.Google - ), - ModelSelection(ServiceType.GOOGLE, Gemini2_5Pro.id, "Gemini 2.5 Pro", Icons.Google), - ModelSelection( - ServiceType.GOOGLE, - Gemini2_5Flash.id, - "Gemini 2.5 Flash", - Icons.Google - ), - ) - } - - fun getMistralAgentModels(): List { - return listOf( - ModelSelection( - ServiceType.MISTRAL, - Devstral2.id, - "Devstral 2", - Icons.Mistral - ), - ModelSelection( - ServiceType.MISTRAL, - DevstralMedium.id, - "Devstral Medium", - Icons.Mistral - ), - ) - } - - private fun getAllChatModels(): List { - return buildList { - addAll(getProxyAIChatModels()) - addAll(getOpenAIChatModels()) - addAll(getAnthropicModels()) - addAll(getGoogleModels()) - addAll(getMistralModels()) - addAll(getLlamaModels()) - addAll(getOllamaModels()) - addAll(getCustomOpenAIModels()) - addAll(getInceptionModels()) - } - } - - private fun getAllApplyModels(): List { - return buildList { - addAll(getProxyAIApplyModels()) - addAll(getOpenAIChatModels()) - addAll(getAnthropicModels()) - addAll(getGoogleModels()) - addAll(getMistralModels()) - addAll(getLlamaModels()) - addAll(getOllamaModels()) - addAll(getCustomOpenAIModels()) - addAll(getInceptionModels()) - } - } - - private fun getAllCodeModels(): List { - return buildList { - addAll(getProxyAICodeModels()) - add(getOpenAICodeModel()) - addAll(getMistralCodeModels()) - addAll(getLlamaModels()) - addAll(getCustomOpenAICodeModels()) - addAll(getOllamaModels()) - addAll(getInceptionModels()) - } - } - - private fun getCustomOpenAICodeModels(): List { - return try { - val customServicesSettings = service() - customServicesSettings.state.services.mapNotNull { service -> - val serviceId = service.id ?: return@mapNotNull null - val serviceName = service.name ?: "" - val modelFromBody = service.codeCompletionSettings.body["model"] - val modelName = (modelFromBody as? String) - val displayName = if (!modelName.isNullOrEmpty()) { - if (modelName.length > 20) "$serviceName (...${modelName.takeLast(20)})" else "$serviceName ($modelName)" - } else serviceName - - ModelSelection(ServiceType.CUSTOM_OPENAI, serviceId, displayName) - } - } catch (e: Exception) { - logger.error("Failed to get Custom OpenAI code models", e) - emptyList() - } - } - - private fun getNextEditModels(): List { - return listOf( - ModelSelection(ServiceType.PROXYAI, MERCURY_CODER, "Mercury Coder"), - ModelSelection(ServiceType.INCEPTION, MERCURY_CODER, "Mercury Coder") - ) - } - - private fun getInceptionModels(): List { - return listOf( - ModelSelection(ServiceType.INCEPTION, MERCURY, "Mercury"), - ModelSelection(ServiceType.INCEPTION, MERCURY_CODER, "Mercury Coder") - ) - } - - fun getInceptionAgentModels(): List { - return listOf( - ModelSelection(ServiceType.INCEPTION, MERCURY, "Mercury"), - ) - } - - fun getProxyAIChatModels(): List { - return listOf( - ModelSelection( - ServiceType.PROXYAI, - GPT_5_CODEX, - "GPT-5 Codex", - Icons.OpenAI, - PricingPlan.INDIVIDUAL - ), - ModelSelection( - ServiceType.PROXYAI, - GPT_5_MINI, - "GPT-5 Mini", - Icons.OpenAI, - PricingPlan.ANONYMOUS - ), - ModelSelection( - ServiceType.PROXYAI, - CLAUDE_4_5_SONNET_THINKING, - "Claude Sonnet 4.5 (thinking)", - Icons.Anthropic, - PricingPlan.INDIVIDUAL - ), - ModelSelection( - ServiceType.PROXYAI, - CLAUDE_4_5_SONNET, - "Claude Sonnet 4.5", - Icons.Anthropic, - PricingPlan.INDIVIDUAL - ), - ModelSelection( - ServiceType.PROXYAI, - GEMINI_PRO_2_5, - "Gemini 3.0 Pro Preview", - Icons.Google, - PricingPlan.INDIVIDUAL - ), - ModelSelection( - ServiceType.PROXYAI, - GEMINI_FLASH_2_5, - "Gemini 2.5 Flash", - Icons.Google, - PricingPlan.ANONYMOUS - ), - ModelSelection( - ServiceType.PROXYAI, - DEEPSEEK_R1, - "DeepSeek R1", - Icons.DeepSeek, - PricingPlan.INDIVIDUAL - ), - ModelSelection( - ServiceType.PROXYAI, - DEEPSEEK_V3, - "DeepSeek V3", - Icons.DeepSeek, - PricingPlan.FREE - ), - ModelSelection( - ServiceType.PROXYAI, - QWEN3_CODER, - "Qwen3 Coder", - Icons.Qwen, - PricingPlan.FREE - ) - ) - } - - fun getProxyAIApplyModels(): List { - return listOf( - ModelSelection( - ServiceType.PROXYAI, - MERCURY_CODER, - "Mercury Coder", - Icons.Inception, - PricingPlan.FREE - ), - ModelSelection(ServiceType.PROXYAI, MORPH, "Morph", Icons.Morph, PricingPlan.FREE), - ModelSelection(ServiceType.PROXYAI, RELACE, "Relace", Icons.Relace, PricingPlan.FREE), - ) - } - - fun getProxyAIChatModelsForPricingPlan(userPricingPlan: PricingPlan?): List { - val allModels = getProxyAIChatModels() - return when (userPricingPlan) { - null, PricingPlan.ANONYMOUS -> allModels.filter { - it.pricingPlan == PricingPlan.ANONYMOUS || it.pricingPlan == PricingPlan.FREE || it.pricingPlan == PricingPlan.INDIVIDUAL - } - - PricingPlan.FREE -> allModels.filter { - it.pricingPlan != PricingPlan.INDIVIDUAL - } - - PricingPlan.INDIVIDUAL -> allModels - } - } - - private fun getProxyAICodeModels(): List { - return listOf( - ModelSelection( - ServiceType.PROXYAI, - MERCURY_CODER, - "Mercury Coder", - Icons.Inception, - PricingPlan.FREE - ), - ) - } - - private fun getOpenAIChatModels(): List { - val openAIModels = listOf( - GPT_5, - GPT_5_MINI, - O4_MINI, - O3_PRO, - O3, - O3_MINI, - GPT_4_1, - GPT_4_1_MINI, - GPT_4_1_NANO, - O1_PREVIEW, - O1_MINI, - GPT_5_MINI, - GPT_4O, - GPT_4_0125_PREVIEW, - GPT_3_5_TURBO_INSTRUCT, - GPT_4_VISION_PREVIEW, - GPT_5_CODEX, - ) - - return openAIModels.mapNotNull { modelId -> - OpenAIChatCompletionModel.entries.find { it.code == modelId }?.let { model -> - ModelSelection(ServiceType.OPENAI, model.code, model.description) - } - } - } - - private fun getOpenAICodeModel(): ModelSelection { - return ModelSelection(ServiceType.OPENAI, GPT_3_5_TURBO_INSTRUCT, "GPT-3.5 Turbo Instruct") - } - - private fun getAnthropicModels(): List { - return listOf( - ModelSelection(ServiceType.ANTHROPIC, CLAUDE_OPUS_4_20250514, "Claude Opus 4"), - ModelSelection(ServiceType.ANTHROPIC, CLAUDE_SONNET_4_20250514, "Claude Sonnet 4") - ) - } - - private fun getGoogleModels(): List { - return listOf( - ModelSelection( - ServiceType.GOOGLE, - GoogleModel.GEMINI_2_5_PRO_PREVIEW.code, - GoogleModel.GEMINI_2_5_PRO_PREVIEW.description - ), - ModelSelection( - ServiceType.GOOGLE, - GoogleModel.GEMINI_2_5_FLASH_PREVIEW.code, - GoogleModel.GEMINI_2_5_FLASH_PREVIEW.description - ), - ModelSelection( - ServiceType.GOOGLE, - GoogleModel.GEMINI_2_5_PRO.code, - GoogleModel.GEMINI_2_5_PRO.description - ), - ModelSelection( - ServiceType.GOOGLE, - GoogleModel.GEMINI_2_0_PRO_EXP.code, - GoogleModel.GEMINI_2_0_PRO_EXP.description - ), - ModelSelection( - ServiceType.GOOGLE, - GoogleModel.GEMINI_2_0_FLASH_THINKING_EXP.code, - GoogleModel.GEMINI_2_0_FLASH_THINKING_EXP.description - ), - ModelSelection( - ServiceType.GOOGLE, - GoogleModel.GEMINI_2_0_FLASH.code, - GoogleModel.GEMINI_2_0_FLASH.description - ), - ModelSelection( - ServiceType.GOOGLE, - GoogleModel.GEMINI_1_5_PRO.code, - GoogleModel.GEMINI_1_5_PRO.description - ) - ) - } - - private fun getMistralModels(): List { - return listOf( - ModelSelection(ServiceType.MISTRAL, Devstral2.id, "Devstral 2"), - ModelSelection(ServiceType.MISTRAL, DEVSTRAL_MEDIUM_2507, "Devstral Medium"), - ModelSelection(ServiceType.MISTRAL, MISTRAL_LARGE_2411, "Mistral Large"), - ModelSelection(ServiceType.MISTRAL, CODESTRAL_LATEST, "Codestral"), - ) - } - - private fun getMistralCodeModels(): List { - return listOf(ModelSelection(ServiceType.MISTRAL, CODESTRAL_LATEST, "Codestral")) - } - - private fun getOllamaModels(): List { - return try { - val ollamaSettings = service() - ollamaSettings.state.availableModels.map { model -> - ModelSelection(ServiceType.OLLAMA, model, model) - } - } catch (e: Exception) { - logger.error("Failed to get Ollama models", e) - emptyList() - } - } - - fun getCustomOpenAIModels(): List { - return try { - service().state.services.mapNotNull { service -> - val serviceId = service.id ?: return@mapNotNull null - val serviceName = service.name ?: "" - val model = service.chatCompletionSettings.body["model"] as? String ?: "" - val displayName = if (model.isNotEmpty()) { - if (model.length > 20) "$serviceName (...${model.takeLast(20)})" else "$serviceName ($model)" - } else serviceName - - ModelSelection( - id = serviceId, - provider = ServiceType.CUSTOM_OPENAI, - model = model, - displayName = displayName - ) - } - } catch (e: Exception) { - logger.error("Failed to get Custom OpenAI models", e) - emptyList() - } - } - - private fun getLlamaModels(): List { - return try { - LlamaModel.entries.flatMap { llamaModel -> - llamaModel.huggingFaceModels.map { hfModel -> - val displayName = - "${llamaModel.label} (${hfModel.parameterSize}B) / Q${hfModel.quantization}" - ModelSelection(ServiceType.LLAMA_CPP, hfModel.name, displayName) - } - } - } catch (e: Exception) { - logger.error("Failed to get llama.cpp models", e) - emptyList() - } - } - - companion object { - // ProxyAI Models - const val PROXYAI_AUTO = "auto" - const val GEMINI_PRO_2_5 = "gemini-pro-2.5" - const val GEMINI_FLASH_2_5 = "gemini-flash-2.5" - const val CLAUDE_4_SONNET = "claude-4-sonnet" - const val CLAUDE_4_SONNET_THINKING = "claude-4-sonnet-thinking" - const val CLAUDE_4_5_SONNET = "claude-sonnet-4-5" - const val CLAUDE_4_5_SONNET_THINKING = "claude-sonnet-4-5-thinking" - const val DEEPSEEK_R1 = "deepseek-r1" - const val DEEPSEEK_V3 = "deepseek-v3" - const val QWEN_2_5_32B_CODE = "qwen-2.5-32b-code" - const val QWEN3_CODER = "qwen3-coder" - const val RELACE = "relace" - const val MORPH = "morph" - - // OpenAI Models - const val GPT_3_5_TURBO_INSTRUCT = "gpt-3.5-turbo-instruct" - const val O4_MINI = "o4-mini" - const val O3_PRO = "o3-pro" - const val O3 = "o3" - const val O3_MINI = "o3-mini" - const val O1_PREVIEW = "o1-preview" - const val O1_MINI = "o1-mini" - const val GPT_4_1 = "gpt-4.1" - const val GPT_4_1_MINI = "gpt-4.1-mini" - const val GPT_4_1_NANO = "gpt-4.1-nano" - const val GPT_4O = "gpt-4o" - const val GPT_4O_MINI = "gpt-4o-mini" - const val GPT_4_0125_PREVIEW = "gpt-4-0125-preview" - const val GPT_4_VISION_PREVIEW = "gpt-4-vision-preview" - const val GPT_5 = "gpt-5" - const val GPT_5_MINI = "gpt-5-mini" - const val GPT_5_CODEX = "gpt-5-codex" - const val GPT_5_1_CODEX = "gpt-5.1-codex" - - // Anthropic Models - const val CLAUDE_OPUS_4_20250514 = "claude-opus-4-20250514" - const val CLAUDE_SONNET_4_20250514 = "claude-sonnet-4-20250514" - - // Google Models - const val GEMINI_2_0_FLASH = "gemini-2.0-flash" - - // Mistral Models - const val MISTRAL_LARGE_2411 = "mistral-large-2411" - const val DEVSTRAL_MEDIUM_2507 = "devstral-medium-2507" - const val CODESTRAL_LATEST = "codestral-latest" - - // Ollama default models - const val LLAMA_3_2 = "llama3.2" - - // Llama.cpp default models - const val LLAMA_3_2_3B_INSTRUCT = "llama-3.2-3b-instruct" - - const val MERCURY = "mercury" - const val MERCURY_CODER = "mercury-coder" - - @JvmStatic - fun getInstance(): ModelRegistry { - return ApplicationManager.getApplication().getService(ModelRegistry::class.java) - } - } -} - -public val Gemini3_Flash_Preview: LLModel = LLModel( - provider = LLMProvider.Google, - id = "gemini-3-flash-preview", - capabilities = listOf( - LLMCapability.Temperature, - LLMCapability.Completion, - LLMCapability.MultipleChoices, - LLMCapability.Vision.Image, - LLMCapability.Vision.Video, - LLMCapability.Audio, - LLMCapability.Tools, - LLMCapability.ToolChoice, - LLMCapability.Schema.JSON.Basic, - LLMCapability.Schema.JSON.Standard, - ), - contextLength = 1_048_576, - maxOutputTokens = 65_536, -) - - -public val GPT5_2_Codex: LLModel = LLModel( - provider = LLMProvider.OpenAI, - id = "gpt-5.2-codex", - capabilities = listOf( - LLMCapability.Completion, - LLMCapability.Schema.JSON.Basic, - LLMCapability.Schema.JSON.Standard, - LLMCapability.Speculation, - LLMCapability.Tools, - LLMCapability.ToolChoice, - LLMCapability.Vision.Image, - LLMCapability.Document, - LLMCapability.MultipleChoices, - LLMCapability.OpenAIEndpoint.Completions, - LLMCapability.OpenAIEndpoint.Responses, - ), - contextLength = 400_000, - maxOutputTokens = 128_000, -) - -public val Devstral2: LLModel = LLModel( - provider = LLMProvider.MistralAI, - id = "devstral-2512", - capabilities = listOf( - LLMCapability.Temperature, - LLMCapability.Completion, - LLMCapability.Tools, - LLMCapability.ToolChoice, - LLMCapability.Schema.JSON.Basic, - LLMCapability.Schema.JSON.Standard - ), - contextLength = 128_000 -) - -public val Mercury: LLModel = LLModel( - id = MERCURY, - provider = InceptionAILLMClient.Inception, - capabilities = listOf( - LLMCapability.Temperature, - LLMCapability.Schema.JSON.Basic, - LLMCapability.Schema.JSON.Standard, - LLMCapability.Speculation, - LLMCapability.Tools, - LLMCapability.ToolChoice, - LLMCapability.Vision.Image, - LLMCapability.Document, - LLMCapability.Completion, - LLMCapability.MultipleChoices, - LLMCapability.OpenAIEndpoint.Completions, - ), - contextLength = 200_000, - maxOutputTokens = 32_768, -) diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/models/ModelSettings.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/models/ModelSettings.kt index 6af86466..fc73dcf3 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/settings/models/ModelSettings.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/models/ModelSettings.kt @@ -1,13 +1,14 @@ package ee.carlrobert.codegpt.settings.models +import ai.koog.prompt.llm.LLModel import com.intellij.openapi.application.ApplicationManager import com.intellij.openapi.components.* +import com.intellij.openapi.diagnostic.thisLogger import ee.carlrobert.codegpt.settings.service.FeatureType import ee.carlrobert.codegpt.settings.service.ModelChangeNotifier import ee.carlrobert.codegpt.settings.service.ServiceType import ee.carlrobert.codegpt.settings.service.ServiceType.PROXYAI import ee.carlrobert.codegpt.settings.service.custom.CustomServicesSettings -import java.time.Month @Service @State( @@ -55,8 +56,28 @@ class ModelSettings : SimplePersistentStateComponent(ModelSe return publisherMethods[featureType] ?: error("No publisher method for $featureType") } - private fun getModelsForFeatureType(featureType: FeatureType): List { - return ModelRegistry.getInstance().getAllModelsForFeature(featureType) + fun getAvailableModels(featureType: FeatureType): List { + return service().getAllModelsForFeature(featureType) + } + + fun getAvailableProviders(featureType: FeatureType): List { + return service().getProvidersForFeature(featureType) + } + + fun isFeatureSupportedByProvider(featureType: FeatureType, provider: ServiceType): Boolean { + return service().isFeatureSupportedByProvider(featureType, provider) + } + + fun findModel(provider: ServiceType, modelCode: String): ModelSelection? { + return service().findModel(provider, modelCode) + } + + fun getDefaultModelSelection(featureType: FeatureType): ModelSelection { + return service().getDefaultModelForFeature(featureType) + } + + fun getModelDisplayName(provider: ServiceType, modelCode: String?): String { + return service().getModelDisplayName(provider, modelCode) } override fun loadState(state: ModelSettingsState) { @@ -87,19 +108,73 @@ class ModelSettings : SimplePersistentStateComponent(ModelSe val details = getModelDetailsState(featureType) return details?.model?.let { model -> details.provider?.let { provider -> - ModelRegistry.getInstance().findModel(provider, model) + findModel(provider, model) } } } - fun getModelForFeature(featureType: FeatureType): String? { + fun getStoredModelForFeature(featureType: FeatureType): String? { return getModelDetailsState(featureType)?.model } - fun getProviderForFeature(featureType: FeatureType): ServiceType? { + fun getStoredProviderForFeature(featureType: FeatureType): ServiceType? { return getModelDetailsState(featureType)?.provider } + fun getAgentModel(): LLModel { + return getModelSelectionForFeature(FeatureType.AGENT).llmModel + } + + fun getModelSelectionForFeature(featureType: FeatureType): ModelSelection { + val modelDetailsState = state.getModelSelection(featureType) + if (modelDetailsState?.model != null && modelDetailsState.provider != null) { + val foundModel = findModel(modelDetailsState.provider!!, modelDetailsState.model!!) + if (foundModel != null) { + return foundModel + } + + logger.warn( + "Stored model selection not found for feature=$featureType provider=${modelDetailsState.provider} model=${modelDetailsState.model}; using default." + ) + } + + return getDefaultModelSelection(featureType) + } + + fun getServiceForFeature(featureType: FeatureType): ServiceType { + return getModelSelectionForFeature(featureType).provider + } + + fun getModelForFeature(featureType: FeatureType): String { + return getModelSelectionForFeature(featureType).model + } + + fun syncWithAvailableCustomOpenAIModels(preferredServiceId: String? = null) { + FeatureType.entries.forEach { featureType -> + if (!isFeatureSupportedByProvider(featureType, ServiceType.CUSTOM_OPENAI)) { + return@forEach + } + + val current = getModelSelection(featureType) + if (current?.provider != ServiceType.CUSTOM_OPENAI) { + return@forEach + } + + val available = getAvailableModels(featureType) + .filter { it.provider == ServiceType.CUSTOM_OPENAI } + val isCurrentValid = available.any { it.model == current.model } + + if (!isCurrentValid) { + val newId = when { + !preferredServiceId.isNullOrBlank() && available.any { it.model == preferredServiceId } -> preferredServiceId + available.isNotEmpty() -> available.first().model + else -> null + } + setModelWithProvider(featureType, newId, ServiceType.CUSTOM_OPENAI) + } + } + } + private fun notifyModelChange( featureType: FeatureType, model: String?, @@ -120,7 +195,7 @@ class ModelSettings : SimplePersistentStateComponent(ModelSe val newModel = getModelFromState(newState, featureType) if (oldModel != newModel) { - val service = getProviderForFeature(featureType) ?: return + val service = getStoredProviderForFeature(featureType) ?: return notifyModelChange(featureType, newModel, service) } } @@ -153,8 +228,8 @@ class ModelSettings : SimplePersistentStateComponent(ModelSe private fun migrateProxyAIModels() { listOf(FeatureType.AUTO_APPLY, FeatureType.CODE_COMPLETION, FeatureType.NEXT_EDIT).forEach { val modelSelection = state.getModelSelection(it) - if (modelSelection?.provider == PROXYAI && modelSelection.model != ModelRegistry.MERCURY_CODER) { - setModelWithProvider(it, ModelRegistry.MERCURY_CODER, PROXYAI) + if (modelSelection?.provider == PROXYAI && modelSelection.model != ModelCatalog.MERCURY_CODER) { + setModelWithProvider(it, ModelCatalog.MERCURY_CODER, PROXYAI) } } listOf( @@ -165,10 +240,10 @@ class ModelSettings : SimplePersistentStateComponent(ModelSe ).forEach { val modelSelection = state.getModelSelection(it) if (modelSelection?.provider == PROXYAI) { - if (modelSelection.model == ModelRegistry.CLAUDE_4_SONNET) { - setModelWithProvider(it, ModelRegistry.CLAUDE_4_5_SONNET, PROXYAI) - } else if (modelSelection.model == ModelRegistry.CLAUDE_4_SONNET_THINKING) { - setModelWithProvider(it, ModelRegistry.CLAUDE_4_5_SONNET_THINKING, PROXYAI) + if (modelSelection.model == ModelCatalog.CLAUDE_4_SONNET) { + setModelWithProvider(it, ModelCatalog.CLAUDE_4_5_SONNET, PROXYAI) + } else if (modelSelection.model == ModelCatalog.CLAUDE_4_SONNET_THINKING) { + setModelWithProvider(it, ModelCatalog.CLAUDE_4_5_SONNET_THINKING, PROXYAI) } } } @@ -178,13 +253,12 @@ class ModelSettings : SimplePersistentStateComponent(ModelSe featureType: FeatureType, modelCode: String ): ServiceType? { - val models = getModelsForFeatureType(featureType) - return models.find { it.model == modelCode }?.provider + val models = getAvailableModels(featureType) + return models.find { it.selectionId == modelCode || it.modelId == modelCode }?.provider } private fun migrateCustomOpenAIModelCodesToIds() { val servicesByName: Map> = try { - CustomServicesSettings::class.java service().state.services .groupBy({ it.name ?: "" }, { it.id ?: "" }) .filterKeys { it.isNotEmpty() } @@ -210,6 +284,9 @@ class ModelSettings : SimplePersistentStateComponent(ModelSe } companion object { + private val logger = thisLogger() + + @JvmStatic fun getInstance(): ModelSettings = service() } } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/models/ModelSettingsForm.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/models/ModelSettingsForm.kt index 80d579b1..2496fbb4 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/settings/models/ModelSettingsForm.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/models/ModelSettingsForm.kt @@ -194,7 +194,7 @@ class ModelSettingsForm( } else -> { - ModelRegistry.getInstance().getDefaultModelForFeature(featureType) + settings.getDefaultModelSelection(featureType) } } @@ -292,7 +292,7 @@ class ModelSettingsForm( } private fun getModelsForFeature(featureType: FeatureType): List { - return ModelRegistry.getInstance().getAllModelsForFeature(featureType) + return settings.getAvailableModels(featureType) } private fun addPreviousModelRow( diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/models/ModelSettingsState.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/models/ModelSettingsState.kt index e789c47a..aaaab75a 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/settings/models/ModelSettingsState.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/models/ModelSettingsState.kt @@ -1,6 +1,7 @@ package ee.carlrobert.codegpt.settings.models import com.intellij.openapi.components.BaseState +import com.intellij.openapi.components.service import ee.carlrobert.codegpt.settings.service.FeatureType import ee.carlrobert.codegpt.settings.service.ServiceType @@ -14,7 +15,7 @@ class ModelSettingsState : BaseState() { } private fun initializeDefaults() { - val registry = ModelRegistry.getInstance() + val registry = service() FeatureType.entries.forEach { featureType -> val defaultModel = registry.getDefaultModelForFeature(featureType) setModelSelection(featureType, defaultModel.model, defaultModel.provider) diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/models/SettingsModelComboBoxAction.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/models/SettingsModelComboBoxAction.kt index 98f905c2..38881cae 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/settings/models/SettingsModelComboBoxAction.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/models/SettingsModelComboBoxAction.kt @@ -19,6 +19,7 @@ import ee.carlrobert.codegpt.settings.service.ServiceType import ee.carlrobert.codegpt.settings.service.llama.LlamaSettings import ee.carlrobert.codegpt.toolwindow.ui.ModelListPopup import java.awt.Color +import javax.swing.Icon import javax.swing.JComponent class SettingsModelComboBoxAction( @@ -36,7 +37,7 @@ class SettingsModelComboBoxAction( fun updateTemplatePresentation(model: ModelSelection?) { templatePresentation.text = model?.fullDisplayName ?: CodeGPTBundle.get("settings.models.selectModel") - templatePresentation.icon = model?.let { ModelIcons.getIconForModel(it) } + templatePresentation.icon = model?.let(::getIconForModel) } fun createCustomComponent(place: String): JComponent { @@ -90,8 +91,7 @@ class SettingsModelComboBoxAction( ServiceType.ANTHROPIC -> { val group = DefaultActionGroup.createPopupGroup { provider.label } - group.templatePresentation.icon = - ModelIcons.getIconForProvider(provider) + group.templatePresentation.icon = getIconForProvider(provider) providerModels.forEach { model -> group.add(createModelAction(model)) } @@ -100,8 +100,7 @@ class SettingsModelComboBoxAction( else -> { val group = DefaultActionGroup.createPopupGroup { provider.label } - group.templatePresentation.icon = - ModelIcons.getIconForProvider(provider) + group.templatePresentation.icon = getIconForProvider(provider) providerModels.forEach { model -> group.add(createModelAction(model)) } @@ -126,7 +125,7 @@ class SettingsModelComboBoxAction( val llamaAction = object : DumbAwareAction( llamaDisplayName, "", - ModelIcons.getIconForProvider(ServiceType.LLAMA_CPP) + getIconForProvider(ServiceType.LLAMA_CPP) ) { override fun update(e: AnActionEvent) { e.presentation.isEnabled = @@ -156,8 +155,7 @@ class SettingsModelComboBoxAction( else -> { val group = DefaultActionGroup.createPopupGroup { provider.label } - group.templatePresentation.icon = - ModelIcons.getIconForProvider(provider) + group.templatePresentation.icon = getIconForProvider(provider) providerModels.forEach { model -> group.add(createModelAction(model)) } @@ -206,7 +204,7 @@ class SettingsModelComboBoxAction( } private fun getModelsForFeature(featureType: FeatureType): List { - val allModels = ModelRegistry.getInstance().getAllModelsForFeature(featureType) + val allModels = service().getAvailableModels(featureType) return if (serviceType != null) { allModels.filter { it.provider == serviceType } @@ -219,7 +217,7 @@ class SettingsModelComboBoxAction( model: ModelSelection, useFullDisplayName: Boolean = false ): AnAction { - val icon = ModelIcons.getIconForModel(model) + val icon = getIconForModel(model) val displayText = if (useFullDisplayName) model.fullDisplayName else model.displayName return object : DumbAwareAction(displayText, "", icon) { @@ -240,10 +238,28 @@ class SettingsModelComboBoxAction( } } + private fun getIconForModel(model: ModelSelection): Icon? { + return model.icon ?: getIconForProvider(model.provider) + } + private fun getLlamaDisplayName(): String { val llamaSettings = ApplicationManager.getApplication().service().state val huggingFaceModel = llamaSettings.huggingFaceModel val llamaModel = LlamaModel.findByHuggingFaceModel(huggingFaceModel) return "${llamaModel.label} (${huggingFaceModel.parameterSize}B) / Q${huggingFaceModel.quantization}" } + + private fun getIconForProvider(provider: ServiceType): Icon? { + return when (provider) { + ServiceType.PROXYAI -> Icons.DefaultSmall + ServiceType.OPENAI -> Icons.OpenAI + ServiceType.ANTHROPIC -> Icons.Anthropic + ServiceType.GOOGLE -> Icons.Google + ServiceType.MISTRAL -> Icons.Mistral + ServiceType.OLLAMA -> Icons.Ollama + ServiceType.CUSTOM_OPENAI -> Icons.OpenAI + ServiceType.LLAMA_CPP -> Icons.Llama + ServiceType.INCEPTION -> Icons.Inception + } + } } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/ModelReplacementDialog.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/ModelReplacementDialog.kt index 0bdde91d..85d486bf 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/ModelReplacementDialog.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/ModelReplacementDialog.kt @@ -4,8 +4,6 @@ import com.intellij.openapi.components.service import com.intellij.openapi.diagnostic.thisLogger import com.intellij.openapi.project.Project import com.intellij.openapi.ui.DialogWrapper -import ee.carlrobert.codegpt.CodeGPTKeys.CODEGPT_USER_DETAILS -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.ModelSettingsForm @@ -66,16 +64,15 @@ class ModelReplacementDialog( } private fun generateInitialModelSelections(serviceType: ServiceType): Map { - val registry = ModelRegistry.getInstance() + val modelSettings = service() return FeatureType.entries.associateWith { featureType -> when (serviceType) { ServiceType.PROXYAI -> { - val userDetails = project?.let { CODEGPT_USER_DETAILS.get(it) } - registry.getDefaultModelForFeature(featureType, userDetails?.pricingPlan) + modelSettings.getDefaultModelSelection(featureType) } ServiceType.CUSTOM_OPENAI -> { - val models = registry.getAllModelsForFeature(featureType) + val models = modelSettings.getAvailableModels(featureType) .filter { it.provider == serviceType } if (!preferredCustomServiceId.isNullOrBlank()) { @@ -87,7 +84,7 @@ class ModelReplacementDialog( } else -> { - registry.getAllModelsForFeature(featureType) + modelSettings.getAvailableModels(featureType) .firstOrNull { it.provider == serviceType } } } @@ -130,16 +127,14 @@ class ModelReplacementDialog( serviceType: ServiceType, preferredCustomServiceId: String? = null ): Boolean { + val modelSettings = service() return FeatureType.entries.any { featureType -> - val registry = service() - if (!registry.isFeatureSupportedByProvider( - featureType, - serviceType - ) - ) return@any false + if (!modelSettings.isFeatureSupportedByProvider(featureType, serviceType)) { + return@any false + } - val currentSelection = service().getModelSelection(featureType) - val availableModels = registry.getAllModelsForFeature(featureType) + val currentSelection = modelSettings.getModelSelection(featureType) + val availableModels = modelSettings.getAvailableModels(featureType) .filter { it.provider == serviceType } if (availableModels.isEmpty()) return@any false @@ -157,4 +152,4 @@ class ModelReplacementDialog( } } } -} \ No newline at end of file +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/ModelSelectionService.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/ModelSelectionService.kt deleted file mode 100644 index 4832cd20..00000000 --- a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/ModelSelectionService.kt +++ /dev/null @@ -1,116 +0,0 @@ -package ee.carlrobert.codegpt.settings.service - -import ai.koog.prompt.llm.LLModel -import com.intellij.openapi.application.ApplicationManager -import com.intellij.openapi.components.Service -import com.intellij.openapi.components.service -import com.intellij.openapi.diagnostic.thisLogger -import ee.carlrobert.codegpt.settings.models.ModelRegistry -import ee.carlrobert.codegpt.settings.models.ModelSelection -import ee.carlrobert.codegpt.settings.models.ModelSettings -import ee.carlrobert.llm.client.codegpt.PricingPlan - -@Service -class ModelSelectionService { - - fun getAgentModel(): LLModel { - val modelDetailsState = service().state.getModelSelection(FeatureType.AGENT) - val model = modelDetailsState?.model - val provider = modelDetailsState?.provider - val models = service().getAgentModels()[provider] - - return models?.find { it.id == model || it.model.id == model }?.model - ?: throw IllegalArgumentException("Model '$model' not found for provider $provider") - } - - fun getModelSelectionForFeature( - featureType: FeatureType, - pricingPlan: PricingPlan? = null - ): ModelSelection { - return try { - val modelSettings = service() - val modelDetailsState = modelSettings.state.getModelSelection(featureType) - - if (modelDetailsState != null && modelDetailsState.model != null && modelDetailsState.provider != null) { - val foundModel = - service().findModel( - modelDetailsState.provider!!, - modelDetailsState.model!! - ) - if (foundModel != null) { - return foundModel - } - } - - service().getDefaultModelForFeature(featureType, pricingPlan) - } catch (exception: Exception) { - logger.warn( - "Error getting model selection for feature: $featureType, using default", - exception - ) - service().getDefaultModelForFeature(featureType, pricingPlan) - } - } - - fun getServiceForFeature(featureType: FeatureType): ServiceType { - return getServiceForFeature(featureType, null) - } - - fun getServiceForFeature(featureType: FeatureType, pricingPlan: PricingPlan?): ServiceType { - return try { - getModelSelectionForFeature(featureType, pricingPlan).provider - } catch (exception: Exception) { - logger.warn("Error getting service for feature: $featureType, using default", exception) - ServiceType.PROXYAI - } - } - - fun getModelForFeature(featureType: FeatureType, pricingPlan: PricingPlan? = null): String { - return try { - getModelSelectionForFeature(featureType, pricingPlan).model - } catch (exception: Exception) { - logger.warn("Error getting model for feature: $featureType, using default", exception) - service().getDefaultModelForFeature(featureType, pricingPlan).model - } - } - - fun syncWithAvailableCustomOpenAIModels(preferredServiceId: String? = null) { - val registry = service() - val settings = service() - - FeatureType.entries.forEach { featureType -> - if (!registry.isFeatureSupportedByProvider( - featureType, - ServiceType.CUSTOM_OPENAI - ) - ) return@forEach - - val current = settings.getModelSelection(featureType) - if (current?.provider != ServiceType.CUSTOM_OPENAI) return@forEach - - val available = registry.getAllModelsForFeature(featureType) - .filter { it.provider == ServiceType.CUSTOM_OPENAI } - - val isCurrentValid = available.any { it.model == current.model } - - if (!isCurrentValid) { - val newId = when { - !preferredServiceId.isNullOrBlank() && available.any { it.model == preferredServiceId } -> preferredServiceId - available.isNotEmpty() -> available.first().model - else -> null - } - settings.setModelWithProvider(featureType, newId, ServiceType.CUSTOM_OPENAI) - } - } - } - - companion object { - - private val logger = thisLogger() - - @JvmStatic - fun getInstance(): ModelSelectionService { - return ApplicationManager.getApplication().getService(ModelSelectionService::class.java) - } - } -} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/codegpt/CodeGPTApiException.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/codegpt/CodeGPTApiException.kt new file mode 100644 index 00000000..9981c440 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/codegpt/CodeGPTApiException.kt @@ -0,0 +1,11 @@ +package ee.carlrobert.codegpt.settings.service.codegpt + +import com.fasterxml.jackson.annotation.JsonIgnoreProperties + +@JsonIgnoreProperties(ignoreUnknown = true) +class CodeGPTApiException( + var title: String? = null, + var status: Int = 0, + var detail: String? = null, + var instance: String? = null, +) : RuntimeException(detail) diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/codegpt/CodeGPTAvailableModels.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/codegpt/CodeGPTAvailableModels.kt deleted file mode 100644 index 0c581ba8..00000000 --- a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/codegpt/CodeGPTAvailableModels.kt +++ /dev/null @@ -1,43 +0,0 @@ -package ee.carlrobert.codegpt.settings.service.codegpt - -import ee.carlrobert.codegpt.Icons -import ee.carlrobert.llm.client.codegpt.PricingPlan -import ee.carlrobert.llm.client.codegpt.PricingPlan.* -import javax.swing.Icon - -object CodeGPTAvailableModels { - - val DEFAULT_CODE_MODEL = CodeGPTModel("Qwen 2.5 Coder", "qwen-2.5-32b-code", Icons.Qwen, ANONYMOUS) - - @JvmStatic - val ALL_CHAT_MODELS: List = listOf( - CodeGPTModel("o4-mini", "o4-mini", Icons.OpenAI, INDIVIDUAL), - CodeGPTModel("GPT-5", "gpt-5", Icons.OpenAI, INDIVIDUAL), - CodeGPTModel("GPT-5 Mini", "gpt-5-mini", Icons.OpenAI, ANONYMOUS), - CodeGPTModel("Claude Sonnet 4.5 (thinking)", "claude-sonnet-4-5-thinking", Icons.Anthropic, INDIVIDUAL), - CodeGPTModel("Claude Sonnet 4.5", "claude-sonnet-4-5", Icons.Anthropic, INDIVIDUAL), - CodeGPTModel("Gemini 2.5 Pro", "gemini-pro-2.5", Icons.Google, INDIVIDUAL), - CodeGPTModel("Gemini 2.5 Flash", "gemini-flash-2.5", Icons.Google, ANONYMOUS), - CodeGPTModel("DeepSeek R1", "deepseek-r1", Icons.DeepSeek, INDIVIDUAL), - CodeGPTModel("DeepSeek V3", "deepseek-v3", Icons.DeepSeek, FREE), - CodeGPTModel("Qwen3 Coder", "qwen3-coder", Icons.Qwen, FREE), - ) - - @JvmStatic - val ALL_CODE_MODELS: List = listOf( - DEFAULT_CODE_MODEL, - CodeGPTModel("GPT-3.5 Turbo Instruct", "gpt-3.5-turbo-instruct", Icons.OpenAI, FREE), - ) - - @JvmStatic - fun findByCode(code: String?): CodeGPTModel? { - return ALL_CHAT_MODELS.union(ALL_CODE_MODELS).firstOrNull { it.code == code } - } -} - -data class CodeGPTModel( - val name: String, - val code: String, - val icon: Icon, - val pricingPlan: PricingPlan -) diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/codegpt/CodeGPTService.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/codegpt/CodeGPTService.kt index cfad1e86..456d3af5 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/codegpt/CodeGPTService.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/codegpt/CodeGPTService.kt @@ -5,7 +5,7 @@ import com.intellij.openapi.components.service import com.intellij.openapi.diagnostic.thisLogger import com.intellij.openapi.project.Project import ee.carlrobert.codegpt.CodeGPTKeys.CODEGPT_USER_DETAILS -import ee.carlrobert.codegpt.completions.CompletionClientProvider +import ee.carlrobert.codegpt.agent.clients.ProxyAILLMClient import ee.carlrobert.codegpt.credentials.CredentialsStore.CredentialKey.CodeGptApiKey import ee.carlrobert.codegpt.credentials.CredentialsStore.getCredential import ee.carlrobert.codegpt.settings.GeneralSettings @@ -34,9 +34,9 @@ class CodeGPTService private constructor(val project: Project) { try { val userDetails = withContext(Dispatchers.IO) { if (apiKey.isNullOrEmpty()) null - else CompletionClientProvider.getCodeGPTClient().getUserDetails(apiKey) + else ProxyAILLMClient(apiKey).use { client -> client.getUserDetails() } } - if (userDetails != null && userDetails.pricingPlan != null) { + if (userDetails != null) { if (!userDetails.fullName.isNullOrEmpty()) { service().state.run { displayName = userDetails.fullName @@ -48,7 +48,7 @@ class CodeGPTService private constructor(val project: Project) { CODEGPT_USER_DETAILS.set(project, userDetails) if (showDialog) { - val modality = modalityState ?: ModalityState.NON_MODAL + val modality = modalityState ?: ModalityState.defaultModalityState() ApplicationManager.getApplication().invokeLater({ ModelReplacementDialog.showDialog(ServiceType.PROXYAI) }, modality) diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/codegpt/CodeGPTServiceForm.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/codegpt/CodeGPTServiceForm.kt index e418c990..2c205ac6 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/codegpt/CodeGPTServiceForm.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/codegpt/CodeGPTServiceForm.kt @@ -13,9 +13,6 @@ import ee.carlrobert.codegpt.ui.UIUtil import ee.carlrobert.codegpt.util.ApplicationUtil import kotlinx.coroutines.Dispatchers import kotlinx.coroutines.runBlocking -import java.awt.Component -import javax.swing.DefaultListCellRenderer -import javax.swing.JList import javax.swing.JPanel class CodeGPTServiceForm { @@ -91,21 +88,4 @@ class CodeGPTServiceForm { } apiKeyField.text = getCredential(CodeGptApiKey) } - - private class CustomComboBoxRenderer : DefaultListCellRenderer() { - override fun getListCellRendererComponent( - list: JList<*>, - value: Any?, - index: Int, - isSelected: Boolean, - cellHasFocus: Boolean - ): Component { - val component = - super.getListCellRendererComponent(list, value, index, isSelected, cellHasFocus) - if (value is CodeGPTModel) { - text = value.name - } - return component - } - } } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/codegpt/CodeGPTServiceSettings.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/codegpt/CodeGPTServiceSettings.kt index c16beddd..63cb6b2b 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/codegpt/CodeGPTServiceSettings.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/codegpt/CodeGPTServiceSettings.kt @@ -1,7 +1,7 @@ package ee.carlrobert.codegpt.settings.service.codegpt import com.intellij.openapi.components.* -import ee.carlrobert.codegpt.settings.models.ModelRegistry +import ee.carlrobert.codegpt.settings.models.ModelCatalog @Service @State( @@ -18,10 +18,10 @@ class CodeGPTServiceSettingsState : BaseState() { } class CodeGPTServiceChatCompletionSettingsState : BaseState() { - var model by string(ModelRegistry.GPT_5_MINI) + var model by string(ModelCatalog.GPT_5_MINI) } class CodeGPTServiceCodeCompletionSettingsState : BaseState() { var codeCompletionsEnabled by property(true) - var model by string(ModelRegistry.MERCURY_CODER) + var model by string(ModelCatalog.MERCURY_CODER) } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/codegpt/CodeGPTUserDetails.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/codegpt/CodeGPTUserDetails.kt new file mode 100644 index 00000000..be643e61 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/codegpt/CodeGPTUserDetails.kt @@ -0,0 +1,20 @@ +package ee.carlrobert.codegpt.settings.service.codegpt + +import com.fasterxml.jackson.annotation.JsonIgnoreProperties + +@JsonIgnoreProperties(ignoreUnknown = true) +data class CodeGPTUserDetails( + val fullName: String? = null, + val pricingPlan: String? = null, + val creditsUsed: Long? = null, + val creditsTotal: Long? = null, + val availableModels: List? = null, + val toolWindowModels: List? = null, + val avatarBase64: String? = null, +) + +@JsonIgnoreProperties(ignoreUnknown = true) +data class AvailableModel( + val code: String? = null, + val title: String? = null, +) diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/codegpt/CodeGPTUserDetailsNotifier.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/codegpt/CodeGPTUserDetailsNotifier.kt index 47f4b372..961b86b2 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/codegpt/CodeGPTUserDetailsNotifier.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/codegpt/CodeGPTUserDetailsNotifier.kt @@ -1,7 +1,6 @@ package ee.carlrobert.codegpt.settings.service.codegpt import com.intellij.util.messages.Topic -import ee.carlrobert.llm.client.codegpt.CodeGPTUserDetails interface CodeGPTUserDetailsNotifier { fun userDetailsObtained(userDetails: CodeGPTUserDetails?) @@ -11,4 +10,4 @@ interface CodeGPTUserDetailsNotifier { val CODEGPT_USER_DETAILS_TOPIC = Topic.create("codegptUserDetails", CodeGPTUserDetailsNotifier::class.java) } -} \ No newline at end of file +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/CustomServiceConfigurable.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/CustomServiceConfigurable.kt index 161f978a..39f0810d 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/CustomServiceConfigurable.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/CustomServiceConfigurable.kt @@ -2,7 +2,7 @@ package ee.carlrobert.codegpt.settings.service.custom import com.intellij.openapi.components.service import com.intellij.openapi.options.Configurable -import ee.carlrobert.codegpt.settings.service.ModelSelectionService +import ee.carlrobert.codegpt.settings.models.ModelSettings import ee.carlrobert.codegpt.settings.service.custom.form.CustomServiceForm import ee.carlrobert.codegpt.util.coroutines.EdtDispatchers import kotlinx.coroutines.CoroutineScope @@ -28,7 +28,7 @@ class CustomServiceConfigurable : Configurable { override fun apply() { component.applyChanges() - ModelSelectionService.getInstance() + ModelSettings.getInstance() .syncWithAvailableCustomOpenAIModels(component.getSelectedServiceId()) } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/CustomServiceSettings.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/CustomServiceSettings.kt index 8697af16..2f371172 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/CustomServiceSettings.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/CustomServiceSettings.kt @@ -7,7 +7,7 @@ import com.intellij.util.xmlb.annotations.OptionTag import ee.carlrobert.codegpt.codecompletions.InfillPromptTemplate import ee.carlrobert.codegpt.credentials.CredentialsStore 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.custom.template.CustomServiceChatCompletionTemplate import ee.carlrobert.codegpt.settings.service.custom.template.CustomServiceCodeCompletionTemplate @@ -16,7 +16,7 @@ import ee.carlrobert.codegpt.util.BaseConverter import java.util.UUID import ee.carlrobert.codegpt.util.MapConverter -private const val DEFAULT_SERVICE_SETTINGS_NANE = "Default" +private const val DEFAULT_SERVICE_SETTINGS_NAME = "Default" @Service @State( @@ -118,7 +118,7 @@ class CustomServicesSettings : } fun customServiceStateForFeatureType(featureType: FeatureType): CustomServiceSettingsState { - val modelSelection = service() + val modelSelection = service() val featureSelection = modelSelection.getModelSelectionForFeature(featureType) if (featureSelection.provider != ServiceType.CUSTOM_OPENAI) @@ -127,8 +127,8 @@ class CustomServicesSettings : "This function should not be called in this context!" ) - return this.state.services.find { it.id == featureSelection.id } - ?: throw IllegalStateException("Unable to find custom service with id '${featureSelection.model}'.") + return this.state.services.find { it.id == featureSelection.serviceId } + ?: throw IllegalStateException("Unable to find custom service with id '${featureSelection.serviceId}'.") } } @@ -151,7 +151,7 @@ class CustomServicesState( @JsonIgnoreProperties(ignoreUnknown = true) class CustomServiceSettingsState : BaseState() { var id by string(UUID.randomUUID().toString()) - var name by string(DEFAULT_SERVICE_SETTINGS_NANE) + var name by string(DEFAULT_SERVICE_SETTINGS_NAME) var template by enum(CustomServiceTemplate.OPENAI) var chatCompletionSettings by property(CustomServiceChatCompletionSettingsState()) var codeCompletionSettings by property(CustomServiceCodeCompletionSettingsState()) diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/form/CustomServiceChatCompletionForm.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/form/CustomServiceChatCompletionForm.kt index dd627136..8d249a34 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/form/CustomServiceChatCompletionForm.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/form/CustomServiceChatCompletionForm.kt @@ -3,16 +3,15 @@ package ee.carlrobert.codegpt.settings.service.custom.form import com.intellij.openapi.ui.MessageType import com.intellij.util.ui.FormBuilder import ee.carlrobert.codegpt.CodeGPTBundle +import ee.carlrobert.codegpt.completions.CancellableRequest +import ee.carlrobert.codegpt.completions.CompletionError +import ee.carlrobert.codegpt.completions.CompletionStreamEventListener import ee.carlrobert.codegpt.completions.CompletionRequestService -import ee.carlrobert.codegpt.completions.factory.CustomOpenAIRequestFactory import ee.carlrobert.codegpt.settings.service.custom.CustomServiceChatCompletionSettingsState import ee.carlrobert.codegpt.settings.service.custom.CustomServiceFormTabbedPane import ee.carlrobert.codegpt.settings.service.custom.form.model.CustomServiceChatCompletionSettingsData import ee.carlrobert.codegpt.ui.OverlayUtil import ee.carlrobert.codegpt.ui.URLTextField -import ee.carlrobert.llm.client.openai.completion.ErrorDetails -import ee.carlrobert.llm.completion.CompletionEventListener -import okhttp3.sse.EventSource import java.awt.BorderLayout import javax.swing.JButton import javax.swing.JPanel @@ -76,26 +75,34 @@ class CustomServiceChatCompletionForm( testConnectionButton.isEnabled = false testConnectionButton.text = "Testing..." - val request = CustomOpenAIRequestFactory.buildCustomOpenAICompletionRequest( - "Test", - urlField.text, - tabbedPane.headers, - tabbedPane.body, - getApiKey.invoke() + val settings = CustomServiceChatCompletionSettingsState().apply { + url = urlField.text + headers.clear() + headers.putAll(tabbedPane.headers) + body.clear() + body.putAll(tabbedPane.body) + } + val listener = TestConnectionEventListener() + val request = CompletionRequestService.testCustomServiceConnectionAsync( + settings = settings, + apiKey = getApiKey.invoke(), + modelId = settings.body["model"]?.toString(), + eventListener = listener ) - - CompletionRequestService.getInstance().getCustomOpenAIChatCompletionAsync( - request, - TestConnectionEventListener() - ) + listener.attachRequest(request) } - internal inner class TestConnectionEventListener : CompletionEventListener { + internal inner class TestConnectionEventListener : CompletionStreamEventListener { private var responseReceived = false + private var request: CancellableRequest? = null - override fun onMessage(value: String?, eventSource: EventSource) { - if (!responseReceived) { + fun attachRequest(request: CancellableRequest) { + this.request = request + } + + override fun onMessage(message: String) { + if (!responseReceived && message.isNotBlank()) { responseReceived = true testConnectionButton.isEnabled = true testConnectionButton.text = @@ -105,11 +112,11 @@ class CustomServiceChatCompletionForm( MessageType.INFO, testConnectionButton ) - eventSource.cancel() + request?.cancel() } } - override fun onError(error: ErrorDetails, ex: Throwable) { + override fun onError(error: CompletionError, ex: Throwable) { testConnectionButton.isEnabled = true testConnectionButton.text = CodeGPTBundle.get("settingsConfigurable.service.custom.openai.testConnection.label") @@ -139,4 +146,4 @@ class CustomServiceChatCompletionForm( CodeGPTBundle.get("settingsConfigurable.service.custom.openai.testConnection.label") } } -} \ No newline at end of file +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/form/CustomServiceCodeCompletionForm.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/form/CustomServiceCodeCompletionForm.kt index b8ae9c7d..baa59e6c 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/form/CustomServiceCodeCompletionForm.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/form/CustomServiceCodeCompletionForm.kt @@ -2,6 +2,7 @@ package ee.carlrobert.codegpt.settings.service.custom.form import com.intellij.icons.AllIcons.General import com.intellij.ide.HelpTooltip +import com.intellij.openapi.components.service import com.intellij.openapi.ui.ComboBox import com.intellij.openapi.ui.MessageType import com.intellij.openapi.ui.panel.ComponentPanelBuilder @@ -10,19 +11,18 @@ import com.intellij.ui.components.JBCheckBox import com.intellij.ui.components.JBLabel import com.intellij.util.ui.FormBuilder import ee.carlrobert.codegpt.CodeGPTBundle -import ee.carlrobert.codegpt.codecompletions.CodeCompletionRequestFactory +import ee.carlrobert.codegpt.codecompletions.CodeCompletionService import ee.carlrobert.codegpt.codecompletions.InfillPromptTemplate import ee.carlrobert.codegpt.codecompletions.InfillRequest -import ee.carlrobert.codegpt.completions.CompletionRequestService +import ee.carlrobert.codegpt.completions.CancellableRequest +import ee.carlrobert.codegpt.completions.CompletionError +import ee.carlrobert.codegpt.completions.CompletionStreamEventListener import ee.carlrobert.codegpt.settings.Placeholder import ee.carlrobert.codegpt.settings.service.custom.CustomServiceCodeCompletionSettingsState import ee.carlrobert.codegpt.settings.service.custom.CustomServiceFormTabbedPane import ee.carlrobert.codegpt.settings.service.custom.form.model.CustomServiceCodeCompletionSettingsData import ee.carlrobert.codegpt.ui.OverlayUtil import ee.carlrobert.codegpt.ui.URLTextField -import ee.carlrobert.llm.client.openai.completion.ErrorDetails -import ee.carlrobert.llm.completion.CompletionEventListener -import okhttp3.sse.EventSource import org.apache.commons.text.StringEscapeUtils import java.awt.BorderLayout import java.awt.FlowLayout @@ -164,43 +164,37 @@ class CustomServiceCodeCompletionForm( testConnectionButton.isEnabled = false testConnectionButton.text = "Testing..." - val selectedTemplate = promptTemplateComboBox.selectedItem as InfillPromptTemplate - val testRequest = InfillRequest.Builder("Hello", "!", 0).build() - - if (selectedTemplate == InfillPromptTemplate.CHAT_COMPLETION) { - // Use chat completion endpoint for testing - CompletionRequestService.getInstance().getCustomOpenAIChatCompletionAsync( - CodeCompletionRequestFactory.buildChatBasedFIMHttpRequest( - testRequest, - urlField.text, - tabbedPane.headers, - tabbedPane.body, - getApiKey.invoke() - ), - TestConnectionEventListener() - ) - } else { - // Use traditional completion endpoint for testing - CompletionRequestService.getInstance().getCustomOpenAICompletionAsync( - CodeCompletionRequestFactory.buildCustomRequest( - testRequest, - urlField.text, - tabbedPane.headers, - tabbedPane.body, - selectedTemplate, - getApiKey.invoke() - ), - TestConnectionEventListener() - ) + val settings = CustomServiceCodeCompletionSettingsState().apply { + codeCompletionsEnabled = true + parseResponseAsChatCompletions = this@CustomServiceCodeCompletionForm.parseResponseAsChatCompletions + infillTemplate = this@CustomServiceCodeCompletionForm.infillTemplate + url = urlField.text + headers.clear() + headers.putAll(tabbedPane.headers) + body.clear() + body.putAll(tabbedPane.body) } + + val listener = TestConnectionEventListener() + val request = service().testCustomOpenAIConnectionAsync( + settings = settings, + apiKey = getApiKey.invoke(), + eventListener = listener + ) + listener.attachRequest(request) } - internal inner class TestConnectionEventListener : CompletionEventListener { + internal inner class TestConnectionEventListener : CompletionStreamEventListener { private var responseReceived = false + private var request: CancellableRequest? = null - override fun onMessage(value: String?, eventSource: EventSource) { - if (!responseReceived) { + fun attachRequest(request: CancellableRequest) { + this.request = request + } + + override fun onMessage(message: String) { + if (!responseReceived && message.isNotBlank()) { responseReceived = true testConnectionButton.isEnabled = true testConnectionButton.text = @@ -210,11 +204,11 @@ class CustomServiceCodeCompletionForm( MessageType.INFO, testConnectionButton ) - eventSource.cancel() + request?.cancel() } } - override fun onError(error: ErrorDetails, ex: Throwable) { + override fun onError(error: CompletionError, ex: Throwable) { testConnectionButton.isEnabled = true testConnectionButton.text = CodeGPTBundle.get("settingsConfigurable.service.custom.openai.testConnection.label") diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/google/GoogleSettings.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/google/GoogleSettings.kt deleted file mode 100644 index 8348aa9c..00000000 --- a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/google/GoogleSettings.kt +++ /dev/null @@ -1,12 +0,0 @@ -package ee.carlrobert.codegpt.settings.service.google - -import com.intellij.openapi.components.* -import ee.carlrobert.llm.client.google.models.GoogleModel - -@Service -@State(name = "CodeGPT_GoogleSettings_210", storages = [Storage("CodeGPT_GoogleSettings_210.xml")]) -class GoogleSettings : SimplePersistentStateComponent(GoogleSettingsState()) - -class GoogleSettingsState : BaseState() { - var model by string(GoogleModel.GEMINI_2_5_PRO.code) -} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/google/GoogleSettingsForm.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/google/GoogleSettingsForm.kt index 55ff23a0..5a3f6fbe 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/google/GoogleSettingsForm.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/google/GoogleSettingsForm.kt @@ -39,7 +39,5 @@ class GoogleSettingsForm { apiKeyField.text = getCredential(GoogleApiKey) } - fun isModified(): Boolean = service().state.run { - getApiKey() != getCredential(GoogleApiKey) - } + fun isModified(): Boolean = getApiKey() != getCredential(GoogleApiKey) } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/ollama/OllamaSettingsForm.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/ollama/OllamaSettingsForm.kt index 005fdab4..f24dc761 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/ollama/OllamaSettingsForm.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/ollama/OllamaSettingsForm.kt @@ -18,12 +18,16 @@ import ee.carlrobert.codegpt.CodeGPTBundle import ee.carlrobert.codegpt.credentials.CredentialsStore.CredentialKey.OllamaApikey import ee.carlrobert.codegpt.credentials.CredentialsStore.getCredential import ee.carlrobert.codegpt.credentials.CredentialsStore.setCredential +import ee.carlrobert.codegpt.agent.clients.HttpClientProvider import ee.carlrobert.codegpt.settings.service.CodeCompletionConfigurationForm import ee.carlrobert.codegpt.settings.service.FeatureType import ee.carlrobert.codegpt.ui.OverlayUtil import ee.carlrobert.codegpt.ui.UIUtil import ee.carlrobert.codegpt.ui.URLTextField -import ee.carlrobert.llm.client.ollama.OllamaClient +import ee.carlrobert.codegpt.util.JsonMapper +import io.ktor.client.request.get +import io.ktor.client.request.header +import io.ktor.client.statement.bodyAsText import kotlinx.coroutines.Dispatchers import kotlinx.coroutines.runBlocking import java.awt.Dimension @@ -154,16 +158,7 @@ class OllamaSettingsForm { disableModelComboBoxWithPlaceholder(DefaultComboBoxModel(arrayOf("Loading"))) ReadAction.nonBlocking> { try { - OllamaClient.Builder() - .setHost(hostField.text) - .setApiKey(getApiKey()) - .build() - .modelTags - .models - ?.map { it.name } - ?.sortedWith(compareBy({ it.split(":").first() }, { - if (it.contains("latest")) 1 else 0 - })) + fetchOllamaModels(hostField.text, getApiKey()) } catch (t: Throwable) { handleModelLoadingError(t) throw t @@ -175,6 +170,27 @@ class OllamaSettingsForm { .submit(AppExecutorUtil.getAppExecutorService()) } + private fun fetchOllamaModels(host: String, apiKey: String?): List { + val normalizedHost = host.trim().trimEnd('/') + val endpoint = "$normalizedHost/api/tags" + val body = runBlocking { + HttpClientProvider.createHttpClient().use { client -> + val response = client.get(endpoint) { + if (!apiKey.isNullOrBlank()) { + header("Authorization", "Bearer $apiKey") + } + } + response.bodyAsText() + } + } + val models = JsonMapper.mapper.readTree(body) + .path("models") + .mapNotNull { node -> node.path("name").takeIf { !it.isMissingNode }?.asText() } + return models.sortedWith(compareBy({ it.split(":").first() }, { + if (it.contains("latest")) 1 else 0 + })) + } + private fun updateModelComboBoxState( models: List, currentModels: Map diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/AgentEventHandler.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/AgentEventHandler.kt index 7055e192..4408bd08 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/AgentEventHandler.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/AgentEventHandler.kt @@ -1,12 +1,13 @@ package ee.carlrobert.codegpt.toolwindow.agent +import ai.koog.agents.core.agent.exception.AIAgentStuckInTheNodeException import ai.koog.http.client.KoogHttpClientException -import ai.koog.prompt.executor.clients.LLMClientException import com.intellij.openapi.Disposable import com.intellij.openapi.application.ApplicationManager import com.intellij.openapi.application.runInEdt import com.intellij.openapi.application.runReadAction import com.intellij.openapi.components.service +import com.intellij.openapi.diagnostic.Logger import com.intellij.openapi.diagnostic.thisLogger import com.intellij.openapi.fileEditor.FileDocumentManager import com.intellij.openapi.project.Project @@ -16,6 +17,7 @@ import ee.carlrobert.codegpt.agent.* import ee.carlrobert.codegpt.agent.history.CheckpointRef import ee.carlrobert.codegpt.agent.rollback.RollbackService import ee.carlrobert.codegpt.agent.tools.* +import ee.carlrobert.codegpt.settings.service.codegpt.CodeGPTApiException import ee.carlrobert.codegpt.settings.service.ServiceType import ee.carlrobert.codegpt.toolwindow.agent.ui.* import ee.carlrobert.codegpt.toolwindow.agent.ui.approval.* @@ -48,7 +50,7 @@ class AgentEventHandler( ) : AgentEvents, Disposable { companion object { - val logger = thisLogger() + val logger: Logger = thisLogger() } private val mainToolCards = ConcurrentHashMap() @@ -182,38 +184,19 @@ class AgentEventHandler( handleDone() } - override fun onClientException(provider: ServiceType, ex: LLMClientException) { - val cause = ex.cause - if (cause is KoogHttpClientException) { - if (provider == ServiceType.PROXYAI) { - handleProxyAIException(cause) - } else { - handleException(cause) - } + override fun onAgentException(provider: ServiceType, throwable: Throwable) { + val statusCode = extractStatusCode(throwable) + if (statusCode != null) { + handleHttpException(provider, statusCode, throwable) + return } + + currentResponseBody?.displayError(extractDisplayMessage(throwable)) + handleDone() } - private fun handleException(ex: KoogHttpClientException) { - when (ex.statusCode) { - 401 -> { - currentResponseBody?.displayMissingCredential() - handleDone() - } - - 403 -> { - currentResponseBody?.displayForbidden() - handleDone() - } - - else -> { - currentResponseBody?.displayError(ex.message) - handleDone() - } - } - } - - private fun handleProxyAIException(ex: KoogHttpClientException) { - when (ex.statusCode) { + private fun handleHttpException(provider: ServiceType, statusCode: Int, throwable: Throwable) { + when (statusCode) { 401 -> { currentResponseBody?.displayMissingCredential() handleDone() @@ -225,17 +208,54 @@ class AgentEventHandler( } 429 -> { - currentResponseBody?.displayCreditsExhausted() + if (provider == ServiceType.PROXYAI) { + currentResponseBody?.displayCreditsExhausted() + } else { + currentResponseBody?.displayError(extractDisplayMessage(throwable)) + } handleDone() } else -> { - currentResponseBody?.displayError("Something went wrong. Please try again.") + if (provider == ServiceType.PROXYAI) { + currentResponseBody?.displayError(CodeGPTBundle.get("toolwindow.agent.error.generic")) + } else { + currentResponseBody?.displayError(extractDisplayMessage(throwable)) + } handleDone() } } } + private fun extractStatusCode(throwable: Throwable): Int? { + return generateSequence(throwable) { it.cause } + .take(10) + .firstNotNullOfOrNull { cause -> + when (cause) { + is KoogHttpClientException -> cause.statusCode + is CodeGPTApiException -> cause.status.takeIf { it > 0 } + else -> null + } + } + } + + private fun extractDisplayMessage(throwable: Throwable): String { + val causes = generateSequence(throwable) { it.cause }.take(10).toList() + if (causes.any { it is AIAgentStuckInTheNodeException }) { + return CodeGPTBundle.get("toolwindow.agent.error.executionFailed") + } + + val message = causes.firstNotNullOfOrNull { cause -> + when (cause) { + is CodeGPTApiException -> cause.detail?.takeIf { it.isNotBlank() } ?: cause.title + is KoogHttpClientException -> cause.errorBody?.takeIf { it.isNotBlank() } ?: cause.message + else -> cause.message + }?.trim()?.takeIf { it.isNotEmpty() } + } ?: CodeGPTBundle.get("toolwindow.agent.error.generic") + + return if (message.length > 400) "${message.take(400)}..." else message + } + private fun handleDone() { runInEdt { currentRollbackRunId?.let { runId -> @@ -317,9 +337,7 @@ class AgentEventHandler( override fun onTextReceived(text: String) { runInEdt { - val cleanedText = - text.replace(Regex(".*?", RegexOption.DOT_MATCHES_ALL), "") - currentResponseBody?.updateMessage(cleanedText) + currentResponseBody?.updateMessage(text) scrollablePanel.update() scrollablePanel.scrollToBottom() } @@ -518,14 +536,9 @@ class AgentEventHandler( .onCreditsChanged(event) } - override fun onRetry(attempt: Int, maxAttempts: Int, reason: String?) { - val suffix = "($attempt/$maxAttempts)" - val base = when { - reason?.contains("timeout", ignoreCase = true) == true -> "Request timed out, retrying" - else -> "Retrying" - } + override fun onRetry(attempt: Int, maxAttempts: Int) { runInEdt { - onShowLoading("$base $suffix") + onShowLoading(CodeGPTBundle.get("toolwindow.agent.retrying", attempt, maxAttempts)) } } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/AgentSessionTimelineController.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/AgentSessionTimelineController.kt index a0511b9b..3e3029a4 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/AgentSessionTimelineController.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/AgentSessionTimelineController.kt @@ -39,6 +39,7 @@ import javax.swing.Action import javax.swing.JButton import javax.swing.JComponent import javax.swing.JPanel +import kotlin.time.ExperimentalTime import ai.koog.prompt.message.Message as PromptMessage internal data class AgentTimelineRunState( @@ -46,6 +47,7 @@ internal data class AgentTimelineRunState( val sourceMessage: Message? ) +@OptIn(ExperimentalTime::class) internal class AgentSessionTimelineController( private val project: Project, private val agentSession: AgentSession, diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/AgentSessionTimelineDialogBuilder.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/AgentSessionTimelineDialogBuilder.kt index 2ff6a9a3..c09182bc 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/AgentSessionTimelineDialogBuilder.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/AgentSessionTimelineDialogBuilder.kt @@ -224,7 +224,7 @@ internal object AgentSessionTimelineDialogBuilder { object : CheckboxTree.CheckboxTreeCellRenderer() { override fun customizeRenderer( tree: JTree, - value: Any?, + value: Any, selected: Boolean, expanded: Boolean, leaf: Boolean, diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/AgentToolWindowTabPanel.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/AgentToolWindowTabPanel.kt index 1e375b66..25d6e4a7 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/AgentToolWindowTabPanel.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/AgentToolWindowTabPanel.kt @@ -25,7 +25,7 @@ import ee.carlrobert.codegpt.conversations.message.Message import ee.carlrobert.codegpt.conversations.message.QueuedMessage import ee.carlrobert.codegpt.psistructure.PsiStructureProvider 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.toolwindow.agent.ui.AgentToolWindowLandingPanel import ee.carlrobert.codegpt.toolwindow.agent.ui.RollbackPanel import ee.carlrobert.codegpt.toolwindow.agent.ui.TodoListPanel @@ -277,7 +277,7 @@ class AgentToolWindowTabPanel( disposeLandingPanelIfPresent() scrollablePanel.clearLandingViewIfVisible() agentSession.serviceType = - ModelSelectionService.getInstance().getServiceForFeature(FeatureType.AGENT) + ModelSettings.getInstance().getServiceForFeature(FeatureType.AGENT) val agentService = project.service() diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/ui/AgentCreditsToolbarLabel.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/ui/AgentCreditsToolbarLabel.kt index b46f79a9..51abd181 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/ui/AgentCreditsToolbarLabel.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/ui/AgentCreditsToolbarLabel.kt @@ -8,12 +8,13 @@ import com.intellij.util.messages.MessageBusConnection import com.intellij.util.ui.JBFont import com.intellij.util.ui.JBUI import ee.carlrobert.codegpt.CodeGPTKeys +import ee.carlrobert.codegpt.settings.models.ModelSettings import ee.carlrobert.codegpt.settings.service.* import ee.carlrobert.codegpt.settings.service.codegpt.CodeGPTService +import ee.carlrobert.codegpt.settings.service.codegpt.CodeGPTUserDetails import ee.carlrobert.codegpt.settings.service.codegpt.CodeGPTUserDetailsNotifier import ee.carlrobert.codegpt.toolwindow.agent.AgentCreditsEvent import ee.carlrobert.codegpt.toolwindow.agent.AgentCreditsListener -import ee.carlrobert.llm.client.codegpt.CodeGPTUserDetails import java.text.NumberFormat import java.util.* @@ -69,7 +70,7 @@ class AgentCreditsToolbarLabel( private fun updateDisplay() { ApplicationManager.getApplication().invokeLater { - val provider = ModelSelectionService.getInstance() + val provider = ModelSettings.getInstance() .getServiceForFeature(FeatureType.AGENT) if (provider != ServiceType.PROXYAI) { isVisible = false diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/ui/AgentToolWindowLandingPanel.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/ui/AgentToolWindowLandingPanel.kt index 0fbc2119..275d6778 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/ui/AgentToolWindowLandingPanel.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/ui/AgentToolWindowLandingPanel.kt @@ -11,11 +11,14 @@ import com.intellij.openapi.vfs.LocalFileSystem import com.intellij.openapi.vfs.VfsUtilCore import com.intellij.openapi.vfs.VirtualFile import com.intellij.ui.SimpleTextAttributes +import com.intellij.ui.JBColor import com.intellij.ui.components.ActionLink import com.intellij.ui.components.JBLabel import com.intellij.util.ui.JBUI import com.intellij.util.ui.components.BorderLayoutPanel import ee.carlrobert.codegpt.agent.ProxyAIAgent.loadProjectInstructions +import ee.carlrobert.codegpt.credentials.CredentialsStore +import ee.carlrobert.codegpt.credentials.CredentialsStore.CredentialKey.CodeGptApiKey import ee.carlrobert.codegpt.agent.history.AgentCheckpointConversationMapper import ee.carlrobert.codegpt.agent.history.AgentCheckpointHistoryService import ee.carlrobert.codegpt.agent.history.AgentHistoryThreadSummary @@ -25,15 +28,19 @@ import ee.carlrobert.codegpt.settings.ProxyAISubagent import ee.carlrobert.codegpt.settings.agents.SubagentsConfigurable import ee.carlrobert.codegpt.settings.hooks.HookConfig import ee.carlrobert.codegpt.settings.hooks.HookConfiguration +import ee.carlrobert.codegpt.settings.models.ModelSettings +import ee.carlrobert.codegpt.settings.service.FeatureType +import ee.carlrobert.codegpt.settings.service.ServiceType +import ee.carlrobert.codegpt.settings.service.codegpt.CodeGPTServiceConfigurable import ee.carlrobert.codegpt.settings.skills.SkillDescriptor import ee.carlrobert.codegpt.settings.skills.SkillDiscoveryService import ee.carlrobert.codegpt.tokens.TokenComputationService import ee.carlrobert.codegpt.toolwindow.agent.AgentToolWindowContentManager import ee.carlrobert.codegpt.toolwindow.agent.history.AgentHistoryListPanel import ee.carlrobert.codegpt.toolwindow.ui.ResponseMessagePanel +import ee.carlrobert.codegpt.ui.UIUtil import ee.carlrobert.codegpt.ui.UIUtil.createTextPane import kotlinx.coroutines.Dispatchers -import kotlinx.coroutines.launch import ee.carlrobert.codegpt.util.coroutines.DisposableCoroutineScope import com.intellij.openapi.Disposable import java.awt.BorderLayout @@ -80,11 +87,56 @@ class AgentToolWindowLandingPanel(private val project: Project) : ResponseMessag private fun buildContent(): JPanel { return BorderLayoutPanel().apply { border = JBUI.Borders.empty(0) - add(createTextPane(welcomeMessage(), false), BorderLayout.NORTH) + add(topPanel(), BorderLayout.NORTH) add(centerPanel(), BorderLayout.CENTER) } } + private fun topPanel(): JPanel { + return BorderLayoutPanel().apply { + isOpaque = false + apiKeyPanel()?.let { addToTop(it) } + addToCenter(createTextPane(welcomeMessage(), false)) + } + } + + private fun apiKeyPanel(): JPanel? { + val provider = ModelSettings.getInstance().getServiceForFeature(FeatureType.AGENT) + if (provider != ServiceType.PROXYAI || CredentialsStore.isCredentialSet(CodeGptApiKey)) { + return null + } + + return ResponseMessagePanel().apply { + addContent( + createTextPane( + """ + +

+ It looks like you haven't configured your API key yet. Visit ProxyAI settings to do so. +

+

+ Don't have an account? Sign up to get started. +

+ + """.trimIndent(), + false + ) { event -> + if (event.eventType == javax.swing.event.HyperlinkEvent.EventType.ACTIVATED && + event.description == "#OPEN_SETTINGS" + ) { + ShowSettingsUtil.getInstance().showSettingsDialog( + project, + CodeGPTServiceConfigurable::class.java + ) + } else { + UIUtil.handleHyperlinkClicked(event) + } + } + ) + border = JBUI.Borders.customLine(JBColor.border(), 1, 0, 0, 0) + } + } + private fun centerPanel(): JPanel { val panel = BorderLayoutPanel() panel.addToTop(topSectionsPanel()) diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/ui/RollbackPanel.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/ui/RollbackPanel.kt index 1258548e..6a0a560d 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/ui/RollbackPanel.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/ui/RollbackPanel.kt @@ -64,6 +64,10 @@ class RollbackPanel( } private fun setupUI() { + isVisible = false + rollbackAllLink.isVisible = false + keepAllLink.isVisible = false + val headerPanel = JPanel(FlowLayout(FlowLayout.LEFT, 4, 2)).apply { isOpaque = false add(titleLabel.apply { diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/AutoApplyListener.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/AutoApplyListener.kt index d9051ad6..68f58e21 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/AutoApplyListener.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/AutoApplyListener.kt @@ -4,15 +4,14 @@ import com.intellij.openapi.diagnostic.logger import com.intellij.openapi.editor.ex.EditorEx import com.intellij.openapi.project.Project import com.intellij.openapi.vfs.VirtualFile -import com.intellij.openapi.vfs.readText +import ee.carlrobert.codegpt.completions.CancellableRequest +import ee.carlrobert.codegpt.completions.CompletionError +import ee.carlrobert.codegpt.completions.CompletionStreamEventListener import ee.carlrobert.codegpt.util.EditorUtil import ee.carlrobert.codegpt.codecompletions.CompletionProgressNotifier import ee.carlrobert.codegpt.toolwindow.chat.editor.header.DiffHeaderPanel import ee.carlrobert.codegpt.toolwindow.chat.editor.state.EditorStateManager import ee.carlrobert.codegpt.toolwindow.chat.parser.* -import ee.carlrobert.llm.client.openai.completion.ErrorDetails -import ee.carlrobert.llm.completion.CompletionEventListener -import okhttp3.sse.EventSource class AutoApplyListener( private val project: Project, @@ -20,35 +19,36 @@ class AutoApplyListener( private val virtualFile: VirtualFile, private val originalSuggestion: String, private val onEditorReplaced: (EditorEx, EditorEx) -> Unit -) : CompletionEventListener { +) : CompletionStreamEventListener { private val logger = logger() private var editorReplaced: Boolean = false private val messageParser = SseMessageParser() - private var eventSource: EventSource? = null + private var request: CancellableRequest? = null + + fun attachRequest(request: CancellableRequest) { + this.request = request + } override fun onOpen() { CompletionProgressNotifier.update(project, true) } - override fun onMessage(message: String, eventSource: EventSource?) { - if (this.eventSource == null && eventSource != null) { - this.eventSource = eventSource - } + override fun onMessage(message: String) { processMessageSegments(message) } - override fun onError(error: ErrorDetails?, ex: Throwable?) { + override fun onError(error: CompletionError, ex: Throwable) { logger.error("Something went wrong while applying the changes", ex) ErrorHandler.handleError(error, ex) handleComplete() } - override fun onCancelled(messageBuilder: java.lang.StringBuilder?) { + override fun onCancelled(messageBuilder: StringBuilder) { handleComplete() } - override fun onComplete(messageBuilder: StringBuilder?) { + override fun onComplete(messageBuilder: StringBuilder) { handleComplete() } @@ -89,7 +89,7 @@ class AutoApplyListener( val containsText = currentText.contains(segment.search.trim()) val newState = if (containsText) { - stateManager.createFromSegment(segment, false, eventSource, originalSuggestion) + stateManager.createFromSegment(segment, false, request, originalSuggestion) } else { stateManager.transitionToFailedDiffState( segment.search, @@ -105,7 +105,7 @@ class AutoApplyListener( val editor = stateManager.getCurrentState()?.editor ?: return (editor.permanentHeaderComponent as? DiffHeaderPanel)?.handleDone() CompletionProgressNotifier.update(project, false) - eventSource = null + request = null } } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/ErrorHandler.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/ErrorHandler.kt index 7a78c9df..95c3e9df 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/ErrorHandler.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/ErrorHandler.kt @@ -1,21 +1,21 @@ package ee.carlrobert.codegpt.toolwindow.chat.editor import com.intellij.notification.NotificationType +import ee.carlrobert.codegpt.completions.CompletionError import ee.carlrobert.codegpt.ui.OverlayUtil -import ee.carlrobert.llm.client.openai.completion.ErrorDetails object ErrorHandler { - fun handleError(error: ErrorDetails?, ex: Throwable?) { + fun handleError(error: CompletionError?, ex: Throwable?) { val errorMessage = formatErrorMessage(error, ex) OverlayUtil.showNotification(errorMessage, NotificationType.ERROR) } - fun formatErrorMessage(error: ErrorDetails?, ex: Throwable?): String { + fun formatErrorMessage(error: CompletionError?, ex: Throwable?): String { return when { error?.code == "insufficient_quota" -> "You exceeded your current quota, please check your plan and billing details." ex?.message != null -> "Error: ${ex.message}" else -> "An unknown error occurred while applying changes." } } -} \ No newline at end of file +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/ResponseEditorPanel.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/ResponseEditorPanel.kt index 0c060b70..0b5f9bdc 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/ResponseEditorPanel.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/ResponseEditorPanel.kt @@ -23,14 +23,8 @@ import com.intellij.util.ui.JBUI import com.intellij.util.ui.components.BorderLayoutPanel import ee.carlrobert.codegpt.codecompletions.CompletionProgressNotifier import ee.carlrobert.codegpt.completions.AutoApplyParameters -import ee.carlrobert.codegpt.completions.CompletionClientProvider -import ee.carlrobert.codegpt.completions.CompletionRequestService -import ee.carlrobert.codegpt.completions.factory.InceptionRequestFactory -import ee.carlrobert.codegpt.settings.models.ModelSelection -import ee.carlrobert.codegpt.settings.service.FeatureType -import ee.carlrobert.codegpt.settings.service.ModelSelectionService -import ee.carlrobert.codegpt.settings.service.ServiceType.INCEPTION -import ee.carlrobert.codegpt.settings.service.ServiceType.PROXYAI +import ee.carlrobert.codegpt.completions.AutoApplyService +import ee.carlrobert.codegpt.completions.CompletionService import ee.carlrobert.codegpt.toolwindow.chat.editor.diff.DiffSyncManager import ee.carlrobert.codegpt.toolwindow.chat.editor.factory.ComponentFactory import ee.carlrobert.codegpt.toolwindow.chat.editor.factory.ComponentFactory.EXPANDED_KEY @@ -44,9 +38,7 @@ import ee.carlrobert.codegpt.toolwindow.chat.parser.SearchReplace import ee.carlrobert.codegpt.toolwindow.chat.parser.Segment import ee.carlrobert.codegpt.ui.OverlayUtil import ee.carlrobert.codegpt.util.EditorUtil -import ee.carlrobert.llm.client.codegpt.request.AutoApplyRequest -import ee.carlrobert.llm.client.codegpt.response.CodeGPTException -import java.util.regex.Pattern +import ee.carlrobert.codegpt.settings.service.codegpt.CodeGPTApiException import javax.swing.BorderFactory class ResponseEditorPanel( @@ -72,7 +64,8 @@ class ResponseEditorPanel( init { isOpaque = false if (compact) { - val visibleBorderColor = JBColor.namedColor("Component.borderColor", JBColor(0xC4C9D0, 0x44484F)) + val visibleBorderColor = + JBColor.namedColor("Component.borderColor", JBColor(0xC4C9D0, 0x44484F)) border = BorderFactory.createCompoundBorder( JBUI.Borders.customLine(visibleBorderColor, 1), JBUI.Borders.empty(4, 6) @@ -139,7 +132,6 @@ class ResponseEditorPanel( } fun applyCode( - modelSelection: ModelSelection, params: AutoApplyParameters, headerPanel: DefaultHeaderPanel ) { @@ -147,25 +139,10 @@ class ResponseEditorPanel( CompletionProgressNotifier.update(project, true) application.executeOnPooledThread { - val model = service().getModelForFeature(FeatureType.AUTO_APPLY) val originalCode = EditorUtil.getFileContent(params.destination) try { - val response = if (modelSelection.provider == INCEPTION) { - val request = InceptionRequestFactory().createAutoApplyRequest(params) - val responseContent = CompletionClientProvider.getInceptionClient() - .getApplyEditCompletion(request) - .choices[0] - .message - .content - extractUpdatedCode(responseContent) - } else if (modelSelection.provider == PROXYAI) { - val request = AutoApplyRequest(model, originalCode, params.source) - CompletionClientProvider.getCodeGPTClient().applyChanges(request).mergedCode - } else { - null - } - - if (!response.isNullOrBlank()) { + val response = service().applyCode(params, originalCode) + if (response.isNotBlank()) { stateManager.transitionToDiffState( originalCode, response, @@ -174,10 +151,13 @@ class ResponseEditorPanel( ) } } catch (e: Exception) { - logger.error("Failed to apply changes", e) + logger.error("Failed to apply changes", e) ApplicationManager.getApplication().invokeLater { - if (e is CodeGPTException) { - OverlayUtil.showNotification(e.detail, NotificationType.ERROR) + if (e is CodeGPTApiException) { + OverlayUtil.showNotification( + e.detail ?: e.message ?: "Failed to apply code", + NotificationType.ERROR + ) } if (!project.isDisposed) { @@ -200,15 +180,18 @@ class ResponseEditorPanel( editor: EditorEx, headerPanel: DefaultHeaderPanel ) { - val eventSource = CompletionRequestService.getInstance().autoApplyAsync( + val listener = AutoApplyListener(project, stateManager, virtualFile, content) { oldEditor, newEditor -> + val responseEditorPanel = editor.component.parent as? ResponseEditorPanel + ?: throw IllegalStateException("Expected parent to be ResponseEditorPanel") + responseEditorPanel.replaceEditor(oldEditor, newEditor) + } + val request = service().autoApply( AutoApplyParameters(content, virtualFile), - AutoApplyListener(project, stateManager, virtualFile, content) { oldEditor, newEditor -> - val responseEditorPanel = editor.component.parent as? ResponseEditorPanel - ?: throw IllegalStateException("Expected parent to be ResponseEditorPanel") - responseEditorPanel.replaceEditor(oldEditor, newEditor) - }) + listener + ) + listener.attachRequest(request) - headerPanel.setLoading(eventSource) + headerPanel.setLoading(request) } internal fun createReplaceWaitingSegment( @@ -253,7 +236,7 @@ class ResponseEditorPanel( stateManager.createFromSegment( finalSegment, readOnly = false, - eventSource = null, + request = null, originalSuggestion = replaceContent ) } else { @@ -390,15 +373,4 @@ class ResponseEditorPanel( ) } } - - private fun extractUpdatedCode(content: String): String { - val pattern = - Pattern.compile("<\\|updated_code\\|>(.*?)<\\|/updated_code\\|>", Pattern.DOTALL) - val matcher = pattern.matcher(content) - return if (matcher.find()) { - matcher.group(1).trim() - } else { - content - } - } } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/actions/AutoApplyAction.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/actions/AutoApplyAction.kt index d307316b..4a72b4d0 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/actions/AutoApplyAction.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/actions/AutoApplyAction.kt @@ -14,7 +14,7 @@ import com.intellij.ui.components.AnActionLink import com.intellij.util.ui.JBUI import ee.carlrobert.codegpt.CodeGPTBundle 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.INCEPTION import ee.carlrobert.codegpt.util.EditorUtil import javax.swing.JComponent diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/header/DefaultHeaderPanel.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/header/DefaultHeaderPanel.kt index 1544cd6e..c1017cb4 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/header/DefaultHeaderPanel.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/header/DefaultHeaderPanel.kt @@ -4,6 +4,7 @@ import com.intellij.icons.AllIcons import com.intellij.openapi.actionSystem.* import com.intellij.openapi.actionSystem.toolbarLayout.ToolbarLayoutStrategy import com.intellij.openapi.application.runInEdt +import com.intellij.openapi.components.service import com.intellij.openapi.diagnostic.thisLogger import com.intellij.openapi.editor.ex.EditorEx import com.intellij.openapi.project.Project @@ -12,15 +13,15 @@ import com.intellij.openapi.ui.JBPopupMenu import com.intellij.util.ui.JBUI import ee.carlrobert.codegpt.CodeGPTBundle import ee.carlrobert.codegpt.completions.AutoApplyParameters +import ee.carlrobert.codegpt.completions.CancellableRequest +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.INCEPTION import ee.carlrobert.codegpt.settings.service.ServiceType.PROXYAI import ee.carlrobert.codegpt.toolwindow.chat.editor.ResponseEditorPanel import ee.carlrobert.codegpt.toolwindow.chat.editor.actions.* import ee.carlrobert.codegpt.util.EditorUtil import ee.carlrobert.codegpt.util.StringUtil -import okhttp3.sse.EventSource import javax.swing.JPanel class DefaultHeaderPanel(config: HeaderConfig) : HeaderPanel(config) { @@ -29,7 +30,6 @@ class DefaultHeaderPanel(config: HeaderConfig) : HeaderPanel(config) { private val logger = thisLogger() } - private var currentEventSource: EventSource? = null private val loadingPanel = LoadingPanel( CodeGPTBundle.get("toolwindow.chat.editor.diff.thinking") ) { @@ -49,12 +49,11 @@ class DefaultHeaderPanel(config: HeaderConfig) : HeaderPanel(config) { } fun setLoading( - eventSource: EventSource? = null, + request: CancellableRequest? = null, label: String = CodeGPTBundle.get("toolwindow.chat.editor.diff.applying") ) { - currentEventSource = eventSource loadingPanel.setText(label) - loadingPanel.setEventSource(eventSource) + loadingPanel.setRequest(request) runInEdt { setRightPanelComponent(loadingPanel) @@ -63,7 +62,6 @@ class DefaultHeaderPanel(config: HeaderConfig) : HeaderPanel(config) { fun handleDone() { runInEdt { - currentEventSource = null setRightPanelComponent(createHeaderActions().component) } } @@ -105,12 +103,11 @@ class DefaultHeaderPanel(config: HeaderConfig) : HeaderPanel(config) { return } - val modelSelection = - ModelSelectionService.getInstance() - .getModelSelectionForFeature(FeatureType.AUTO_APPLY); + val modelSelection = service() + .getModelSelectionForFeature(FeatureType.AUTO_APPLY) val params = AutoApplyParameters(editor.document.text, file) if (listOf(PROXYAI, INCEPTION).any { it == modelSelection.provider }) { - responseEditorPanel.applyCode(modelSelection, params, this) + responseEditorPanel.applyCode(params, this) } else { responseEditorPanel.applyCodeAsync(editor.document.text, file, editor, this) } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/header/DiffHeaderPanel.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/header/DiffHeaderPanel.kt index a6bfeaf5..d94a18d8 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/header/DiffHeaderPanel.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/header/DiffHeaderPanel.kt @@ -8,10 +8,10 @@ import com.intellij.openapi.actionSystem.AnActionEvent import com.intellij.openapi.application.runInEdt import com.intellij.ui.components.ActionLink import ee.carlrobert.codegpt.CodeGPTBundle +import ee.carlrobert.codegpt.completions.CancellableRequest import ee.carlrobert.codegpt.toolwindow.chat.editor.ResponseEditorPanel import ee.carlrobert.codegpt.toolwindow.chat.editor.diff.DiffAcceptedPanel import ee.carlrobert.codegpt.ui.IconActionButton -import okhttp3.sse.EventSource import java.awt.BorderLayout import javax.swing.* @@ -25,16 +25,16 @@ class DiffHeaderPanel( config: HeaderConfig, retry: Boolean, private val actions: DiffHeaderActions, - private var eventSource: EventSource? = null + private var request: CancellableRequest? = null ) : HeaderPanel(config) { private val loadingPanel = LoadingPanel( when { retry -> CodeGPTBundle.get("toolwindow.chat.editor.diff.retrying") - eventSource != null -> CodeGPTBundle.get("toolwindow.chat.editor.diff.applying") + request != null -> CodeGPTBundle.get("toolwindow.chat.editor.diff.applying") else -> CodeGPTBundle.get("toolwindow.chat.editor.diff.editing") }, - eventSource + request ) { handleDone() } @@ -60,7 +60,7 @@ class DiffHeaderPanel( setupUI() runInEdt { loadingPanel.isVisible = true - loadingPanel.setEventSource(eventSource) + loadingPanel.setRequest(request) } } @@ -74,7 +74,7 @@ class DiffHeaderPanel( } fun handleDone() { - eventSource = null + request = null runInEdt { actionLinksPanel.isVisible = true loadingPanel.isVisible = false @@ -84,7 +84,7 @@ class DiffHeaderPanel( } fun handleChangesApplied(before: String, after: String, patches: List) { - eventSource = null + request = null actionLinksPanel.isVisible = false loadingPanel.isVisible = false @@ -106,7 +106,7 @@ class DiffHeaderPanel( runInEdt { loadingPanel.setText(CodeGPTBundle.get("toolwindow.chat.editor.diff.editing")) loadingPanel.isVisible = true - loadingPanel.showStopButton(eventSource != null) + loadingPanel.showStopButton(request != null) } } } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/header/LoadingPanel.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/header/LoadingPanel.kt index 37ea0b69..b309b3c6 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/header/LoadingPanel.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/header/LoadingPanel.kt @@ -5,14 +5,14 @@ import com.intellij.openapi.actionSystem.AnAction import com.intellij.openapi.actionSystem.AnActionEvent import com.intellij.ui.* import com.intellij.ui.components.JBLabel +import ee.carlrobert.codegpt.completions.CancellableRequest import ee.carlrobert.codegpt.ui.IconActionButton -import okhttp3.sse.EventSource import javax.swing.BoxLayout import javax.swing.JPanel class LoadingPanel( initialText: String, - private var eventSource: EventSource? = null, + private var request: CancellableRequest? = null, private val onCancel: (() -> Unit)? = null ) : JPanel() { @@ -21,13 +21,13 @@ class LoadingPanel( private val stopButton = IconActionButton( object : AnAction("Stop", "Stop the current operation", AllIcons.Actions.Suspend) { override fun actionPerformed(e: AnActionEvent) { - eventSource?.cancel() + request?.cancel() onCancel?.invoke() } }, "stop-operation" ).apply { - isVisible = eventSource != null + isVisible = request != null } init { @@ -47,14 +47,14 @@ class LoadingPanel( loadingLabel.text = text } - fun setEventSource(source: EventSource?) { - eventSource = source - stopButton.isVisible = source != null + fun setRequest(request: CancellableRequest?) { + this.request = request + stopButton.isVisible = request != null revalidate() repaint() } fun showStopButton(show: Boolean) { - stopButton.isVisible = show && eventSource != null + stopButton.isVisible = show && request != null } -} \ No newline at end of file +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/state/DiffEditorState.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/state/DiffEditorState.kt index 88eac704..8e880493 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/state/DiffEditorState.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/state/DiffEditorState.kt @@ -24,12 +24,12 @@ import com.intellij.ui.JBColor import com.intellij.util.ui.JBUI import ee.carlrobert.codegpt.CodeGPTBundle import ee.carlrobert.codegpt.Icons +import ee.carlrobert.codegpt.completions.CancellableRequest import ee.carlrobert.codegpt.toolwindow.chat.editor.header.DiffHeaderActions import ee.carlrobert.codegpt.toolwindow.chat.editor.header.DiffHeaderPanel import ee.carlrobert.codegpt.toolwindow.chat.editor.header.HeaderConfig import ee.carlrobert.codegpt.toolwindow.chat.parser.Segment import ee.carlrobert.codegpt.util.file.FileUtil -import okhttp3.sse.EventSource import java.util.* import javax.swing.Icon import javax.swing.JButton @@ -41,7 +41,7 @@ abstract class DiffEditorState( override val project: Project, val diffViewer: UnifiedDiffViewer?, val virtualFile: VirtualFile?, - private val eventSource: EventSource? = null + private val request: CancellableRequest? = null ) : EditorState { companion object { @@ -109,7 +109,7 @@ abstract class DiffEditorState( ), readOnly, actions, - eventSource + request ) } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/state/EditorStateManager.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/state/EditorStateManager.kt index ecba725a..7ebe9392 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/state/EditorStateManager.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/state/EditorStateManager.kt @@ -6,6 +6,7 @@ import com.intellij.openapi.editor.ex.EditorEx import com.intellij.openapi.project.Project import com.intellij.openapi.vfs.LocalFileSystem import com.intellij.openapi.vfs.VirtualFile +import ee.carlrobert.codegpt.completions.CancellableRequest import ee.carlrobert.codegpt.toolwindow.chat.editor.ResponseEditorPanel import ee.carlrobert.codegpt.toolwindow.chat.editor.ResponseEditorPanel.Companion.RESPONSE_EDITOR_STATE_KEY import ee.carlrobert.codegpt.toolwindow.chat.editor.ResponseEditorPanel.Companion.RESPONSE_EDITOR_DIFF_VIEWER_KEY @@ -18,7 +19,6 @@ import ee.carlrobert.codegpt.toolwindow.chat.editor.header.DiffHeaderPanel import com.intellij.diff.util.DiffUtil import com.intellij.diff.util.Side import com.intellij.openapi.util.text.StringUtil -import okhttp3.sse.EventSource class EditorStateManager(private val project: Project) { @@ -28,13 +28,13 @@ class EditorStateManager(private val project: Project) { fun createFromSegment( segment: Segment, readOnly: Boolean = false, - eventSource: EventSource? = null, + request: CancellableRequest? = null, originalSuggestion: String? = null, compact: Boolean = false ): EditorState { val editor = EditorFactory.createEditor(project, segment) val state = if (editor.editorKind == EditorKind.DIFF) { - createDiffState(editor, segment, eventSource, originalSuggestion) + createDiffState(editor, segment, request, originalSuggestion) } else { RegularEditorState(editor, segment, project) } @@ -135,11 +135,13 @@ class EditorStateManager(private val project: Project) { private fun createDiffState( editor: EditorEx, segment: Segment, - eventSource: EventSource? = null, + request: CancellableRequest? = null, originalSuggestion: String? = null ): EditorState { val virtualFile = getVirtualFile(segment.filePath) - val diffViewer = RESPONSE_EDITOR_DIFF_VIEWER_KEY.get(editor) + val diffViewer = requireNotNull(RESPONSE_EDITOR_DIFF_VIEWER_KEY.get(editor)) { + "Diff viewer is missing for diff editor state" + } val diffEditorManager = DiffEditorManager(project, diffViewer) this.diffEditorManager = diffEditorManager @@ -150,7 +152,7 @@ class EditorStateManager(private val project: Project) { diffViewer, virtualFile, diffEditorManager, - eventSource, + request, originalSuggestion ) return state diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/state/StandardDiffEditorState.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/state/StandardDiffEditorState.kt index 8999c1a5..7e0ffeb8 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/state/StandardDiffEditorState.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/editor/state/StandardDiffEditorState.kt @@ -9,6 +9,7 @@ import com.intellij.openapi.editor.ex.EditorEx import com.intellij.openapi.project.Project import com.intellij.openapi.vfs.VirtualFile import com.intellij.util.application +import ee.carlrobert.codegpt.completions.CancellableRequest import ee.carlrobert.codegpt.toolwindow.chat.editor.ResponseEditorPanel import ee.carlrobert.codegpt.toolwindow.chat.editor.diff.DiffEditorManager import ee.carlrobert.codegpt.toolwindow.chat.editor.header.DiffHeaderPanel @@ -16,7 +17,6 @@ import ee.carlrobert.codegpt.toolwindow.chat.parser.Code import ee.carlrobert.codegpt.toolwindow.chat.parser.ReplaceWaiting import ee.carlrobert.codegpt.toolwindow.chat.parser.SearchReplace import ee.carlrobert.codegpt.toolwindow.chat.parser.Segment -import okhttp3.sse.EventSource class StandardDiffEditorState( editor: EditorEx, @@ -25,9 +25,9 @@ class StandardDiffEditorState( diffViewer: UnifiedDiffViewer?, virtualFile: VirtualFile?, private val diffEditorManager: DiffEditorManager, - eventSource: EventSource? = null, + request: CancellableRequest? = null, private val originalSuggestion: String? = null -) : DiffEditorState(editor, segment, project, diffViewer, virtualFile, eventSource) { +) : DiffEditorState(editor, segment, project, diffViewer, virtualFile, request) { private var latestSuggestion: String? = when (segment) { is SearchReplace -> segment.replace.takeIf { it.isNotBlank() } is ReplaceWaiting -> segment.replace.takeIf { it.isNotBlank() } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/ui/textarea/TotalTokensPanel.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/ui/textarea/TotalTokensPanel.kt index c55265d6..516c332e 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/ui/textarea/TotalTokensPanel.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/ui/textarea/TotalTokensPanel.kt @@ -18,7 +18,7 @@ import ee.carlrobert.codegpt.psistructure.ClassStructureSerializer import ee.carlrobert.codegpt.settings.configuration.ConfigurationSettings 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.settings.models.ModelSettings import ee.carlrobert.codegpt.settings.service.ServiceType import ee.carlrobert.codegpt.toolwindow.chat.structure.data.PsiStructureRepository import ee.carlrobert.codegpt.tokens.TokenComputationService @@ -46,6 +46,7 @@ class TotalTokensPanel( private val encodingManager = EncodingManager.getInstance() private val totalTokensDetails = createTokenDetails(conversation, highlightedText) private val label = JBLabel(getLabelHtml(totalTokensDetails.total)) + private val totalTokensListeners = mutableListOf<(Int) -> Unit>() private val dispatchers = CoroutineDispatchers() private val scope = DisposableCoroutineScope(dispatchers.default()) @@ -113,11 +114,19 @@ class TotalTokensPanel( fun getTokenDetails(): TotalTokensDetails = totalTokensDetails + fun addTotalTokensListener(listener: (Int) -> Unit) { + totalTokensListeners += listener + listener(totalTokensDetails.total) + } + fun update() = update(totalTokensDetails.total) fun update(total: Int) { scope.launch { - withEdt { label.text = getLabelHtml(total) } + withEdt { + label.text = getLabelHtml(total) + totalTokensListeners.forEach { it(total) } + } } } @@ -201,7 +210,7 @@ class TotalTokensPanel( } private fun getIconToolTipText(html: String): String { - return if (ModelSelectionService.getInstance().getServiceForFeature(FeatureType.CHAT) != ServiceType.OPENAI) { + return if (ModelSettings.getInstance().getServiceForFeature(FeatureType.CHAT) != ServiceType.OPENAI) { """ diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/ui/ModelListPopup.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/ui/ModelListPopup.kt index d0b3d07e..696e3383 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/ui/ModelListPopup.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/ui/ModelListPopup.kt @@ -11,8 +11,6 @@ import com.intellij.ui.popup.list.PopupListElementRenderer import com.intellij.util.ui.JBUI import com.intellij.util.ui.components.BorderLayoutPanel import ee.carlrobert.codegpt.Icons -import ee.carlrobert.codegpt.settings.service.codegpt.CodeGPTModel -import ee.carlrobert.llm.client.codegpt.PricingPlan import javax.swing.Box import javax.swing.Icon import javax.swing.JComponent @@ -60,34 +58,21 @@ class ModelListPopup( } class CodeGPTModelsListPopupAction : DumbAwareAction { - - val model: CodeGPTModel? + + val modelCode: String private val locked: Boolean private val selected: Boolean private val onModelChanged: Runnable? - - constructor( - model: CodeGPTModel, - locked: Boolean = false, - selected: Boolean = false, - onModelChanged: Runnable? - ) : super(model.name, "", if (locked) Icons.Locked else model.icon) { - this.model = model - this.locked = locked - this.selected = selected - this.onModelChanged = onModelChanged - } - + constructor( name: String, code: String, icon: Icon, - pricingPlan: PricingPlan, locked: Boolean = false, selected: Boolean = false, onModelChanged: Runnable? ) : super(name, "", if (locked) Icons.Locked else icon) { - this.model = CodeGPTModel(name, code, icon, pricingPlan) + this.modelCode = code this.locked = locked this.selected = selected this.onModelChanged = onModelChanged diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/ui/mcp/CompactToolPanel.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/ui/mcp/CompactToolPanel.kt deleted file mode 100644 index 1cd00924..00000000 --- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/ui/mcp/CompactToolPanel.kt +++ /dev/null @@ -1,427 +0,0 @@ -package ee.carlrobert.codegpt.toolwindow.ui.mcp - -import com.fasterxml.jackson.databind.ObjectMapper -import com.intellij.icons.AllIcons -import com.intellij.ui.JBColor -import com.intellij.ui.components.ActionLink -import com.intellij.ui.components.JBLabel -import com.intellij.ui.components.JBPanel -import com.intellij.ui.components.JBTextArea -import com.intellij.util.ui.JBUI -import com.intellij.util.ui.UIUtil -import ee.carlrobert.codegpt.Icons -import ee.carlrobert.llm.client.openai.completion.response.ToolCall -import java.awt.* -import java.awt.event.ActionEvent -import java.awt.event.KeyEvent -import java.awt.event.MouseAdapter -import java.awt.event.MouseEvent -import javax.swing.* -import kotlin.math.min - -class CompactToolPanel( - private val toolCall: ToolCall, - private val serverName: String, - private val onApprove: (autoApprove: Boolean) -> Unit, - private val onReject: () -> Unit, -) : JBPanel() { - - private var isExpanded = false - var isCompleted = false - private set - private var status: CompletionStatus? = null - - val toolCallId: String? - get() = toolCall.id - - private val parameterPreview: String - private val fullParameters: Map? - - private val mainPanel = JBPanel>() - private val detailsPanel = JBPanel>() - private val statusLabel = JBLabel() - - enum class CompletionStatus { - APPROVED, REJECTED, EXECUTING, COMPLETED, FAILED - } - - companion object { - private const val SEPARATOR_TEXT = " | " - private const val AUTO_TOOLTIP_TEXT = "Auto-approve for this session" - } - - init { - val (preview, full) = parseParameters() - parameterPreview = preview - fullParameters = full - - layout = BorderLayout() - isOpaque = false - border = JBUI.Borders.empty() - - setupUI() - setupKeyboardShortcuts() - } - - private fun parseParameters(): Pair?> { - return try { - val mapper = ObjectMapper() - val args = - mapper.readValue(toolCall.function.arguments, Map::class.java) as Map - - val name = toolCall.function.name.lowercase() - val preview = when { - args.isEmpty() -> "" - name.contains("bash") -> { - val cmd = (args["command"] as? String)?.trim() - if (!cmd.isNullOrEmpty()) "command: \"${cmd.take(80)}${if (cmd.length > 80) "..." else ""}\"" else - args.entries.take(2).joinToString(", ") { (k, v) -> "$k: ${formatPreviewValue(v)}" } - } - name.startsWith("read") -> { - val path = (args["file_path"] as? String) ?: (args["path"] as? String) - val off = (args["offset"] as? Number)?.toInt() - val lim = (args["limit"] as? Number)?.toInt() - buildString { - path?.let { append("file: ${it}") } - if (off != null || lim != null) { - if (isNotEmpty()) append(", ") - append("range: ") - append(off?.toString() ?: "1") - append("..") - append(lim?.toString() ?: "") - } - }.ifEmpty { - args.entries.take(2).joinToString(", ") { (k, v) -> "$k: ${formatPreviewValue(v)}" } - } - } - name.contains("grep") || name.contains("search") -> { - val pat = (args["pattern"] as? String) - val path = (args["path"] as? String) - buildString { - pat?.let { append("pattern: \"${it.take(40)}${if (it.length > 40) "..." else ""}\"") } - if (!path.isNullOrEmpty()) { - if (isNotEmpty()) append(", ") - append("in: $path") - } - }.ifEmpty { - args.entries.take(2).joinToString(", ") { (k, v) -> "$k: ${formatPreviewValue(v)}" } - } - } - name.startsWith("write") || name.contains("edit") -> { - val path = (args["file_path"] as? String) ?: (args["path"] as? String) - val content = (args["content"] as? String) - buildString { - path?.let { append("file: $it") } - if (!content.isNullOrEmpty()) { - if (isNotEmpty()) append(", ") - append("content: ${content.length} chars") - } - }.ifEmpty { - args.entries.take(2).joinToString(", ") { (k, v) -> "$k: ${formatPreviewValue(v)}" } - } - } - else -> { - args.entries.take(2).joinToString(", ") { (k, v) -> - "$k: ${formatPreviewValue(v)}" - } + if (args.size > 2) ", ..." else "" - } - } - - preview to args - } catch (_: Exception) { - val raw = toolCall.function.arguments - val preview = if (raw.isNullOrEmpty() || raw == "{}") "" else raw.take(50) + "..." - preview to null - } - } - - private fun setupUI() { - mainPanel.apply { - layout = BorderLayout() - isOpaque = false - border = JBUI.Borders.empty(4, 8, 4, 0) - - setupMainContent() - - addMouseListener(object : MouseAdapter() { - override fun mouseClicked(e: MouseEvent) { - if (!isCompleted && e.clickCount == 1) { - toggleExpanded() - } - } - - override fun mouseEntered(e: MouseEvent) { - if (!isCompleted) { - cursor = Cursor.getPredefinedCursor(Cursor.HAND_CURSOR) - background = JBUI.CurrentTheme.ActionButton.hoverBackground() - isOpaque = true - repaint() - } - } - - override fun mouseExited(e: MouseEvent) { - cursor = Cursor.getDefaultCursor() - isOpaque = false - repaint() - } - }) - } - - detailsPanel.apply { - isVisible = false - layout = BorderLayout() - isOpaque = false - border = JBUI.Borders.empty(0, 32, 8, 0) - - setupDetailsContent() - } - - add(mainPanel, BorderLayout.NORTH) - add(detailsPanel, BorderLayout.CENTER) - } - - private fun setupMainContent() { - mainPanel.layout = BorderLayout() - - val infoPanel = JBPanel>().apply { - layout = FlowLayout(FlowLayout.LEFT, 0, 0) - isOpaque = false - } - - infoPanel.add(Box.createHorizontalStrut(0)) - - val toolNameLabel = JBLabel(toolCall.function.name).apply { - font = JBUI.Fonts.label().asBold() - } - infoPanel.add(toolNameLabel) - - infoPanel.add(JBLabel(" • ").apply { foreground = JBColor.gray }) - - infoPanel.add(JBLabel(serverName).apply { - foreground = JBColor.gray - }) - - if (parameterPreview.isNotEmpty()) { - infoPanel.add(JBLabel(" • ").apply { foreground = JBColor.gray }) - val paramLabel = JBLabel(parameterPreview).apply { - foreground = JBColor.gray - toolTipText = McpParameterPreview.generateTooltip(toolCall.function.arguments) - } - infoPanel.add(paramLabel) - } - - if (fullParameters != null && fullParameters.isNotEmpty()) { - infoPanel.add(Box.createHorizontalStrut(4)) - infoPanel.add(JBLabel(AllIcons.General.ArrowDown).apply { - foreground = JBColor.gray - }) - } - - mainPanel.add(infoPanel, BorderLayout.CENTER) - - if (!isCompleted) { - val actionsPanel = setupActionLinks() - mainPanel.add(actionsPanel, BorderLayout.EAST) - } else { - statusLabel.font = JBUI.Fonts.smallFont() - statusLabel.foreground = when (status) { - CompletionStatus.COMPLETED -> JBColor.BLUE - CompletionStatus.FAILED -> JBColor.RED - CompletionStatus.APPROVED -> JBColor.GREEN - CompletionStatus.REJECTED -> JBColor.RED - CompletionStatus.EXECUTING -> JBColor.GRAY - else -> JBColor.GRAY - } - mainPanel.add(statusLabel, BorderLayout.EAST) - } - } - - private fun setupActionLinks(): JPanel { - val actionsPanel = JBPanel>().apply { - layout = FlowLayout(FlowLayout.RIGHT, 4, 0) - isOpaque = false - } - - actionsPanel.add(ActionLink("Yes") { - handleApproval(false) - }) - - actionsPanel.add(JBLabel(SEPARATOR_TEXT).apply { foreground = JBColor.gray }) - - actionsPanel.add(ActionLink("No") { - handleRejection() - }) - - return actionsPanel - } - - private fun setupDetailsContent() { - if (fullParameters == null || fullParameters.isEmpty()) return - - val paramsPanel = JBPanel>().apply { - layout = GridBagLayout() - isOpaque = false - } - - val gbc = GridBagConstraints().apply { - anchor = GridBagConstraints.NORTHWEST - insets = JBUI.insets(2, 0) - gridx = 0 - gridy = 0 - } - - fullParameters.forEach { (key, value) -> - gbc.gridx = 0 - gbc.weightx = 0.0 - paramsPanel.add(JBLabel("$key:").apply { - foreground = JBColor.gray - font = JBUI.Fonts.smallFont() - }, gbc) - - gbc.gridx = 1 - gbc.weightx = 1.0 - gbc.fill = GridBagConstraints.HORIZONTAL - gbc.insets = JBUI.insets(2, 8, 2, 0) - - val valueLabel = createValueLabel(value) - paramsPanel.add(valueLabel, gbc) - - gbc.gridy++ - gbc.insets = JBUI.insets(2, 0) - } - - detailsPanel.add(paramsPanel, BorderLayout.NORTH) - } - - private fun createValueLabel(value: Any): JComponent { - val text = formatDetailValue(value) - - return if (text.length > 100 || text.contains('\n')) { - JBTextArea(text).apply { - isEditable = false - font = JBUI.Fonts.smallFont() - background = UIUtil.getPanelBackground() - foreground = JBColor.foreground() - lineWrap = true - wrapStyleWord = true - border = JBUI.Borders.empty() - rows = min(3, text.lines().size) - } - } else { - JBLabel(text).apply { - font = JBUI.Fonts.smallFont() - } - } - } - - private fun formatPreviewValue(value: Any): String { - return when (value) { - is String -> if (value.length > 20) "\"${value.take(20)}...\"" else "\"$value\"" - is List<*> -> "[${value.size}]" - is Map<*, *> -> "{${value.size}}" - else -> value.toString().take(20) - } - } - - private fun formatDetailValue(value: Any): String { - return when (value) { - is String -> value - is List<*> -> value.joinToString("\n") { "• ${it.toString()}" } - is Map<*, *> -> value.entries.joinToString("\n") { (k, v) -> "$k: $v" } - else -> value.toString() - } - } - - private fun toggleExpanded() { - isExpanded = !isExpanded - - val infoPanel = (mainPanel.getComponent(0) as? JPanel) ?: return - infoPanel.components - .filterIsInstance() - .lastOrNull { it.icon == AllIcons.General.ArrowDown || it.icon == AllIcons.General.ArrowUp } - ?.icon = if (isExpanded) AllIcons.General.ArrowUp else AllIcons.General.ArrowDown - - if (isExpanded) { - detailsPanel.isVisible = true - val targetHeight = detailsPanel.preferredSize.height - detailsPanel.preferredSize = detailsPanel.preferredSize.apply { height = 0 } - McpAnimations.animateHeight(detailsPanel, targetHeight, duration = 150) { - revalidate() - } - } else { - McpAnimations.animateHeight(detailsPanel, 0, duration = 150) { - detailsPanel.isVisible = false - revalidate() - } - } - } - - private fun handleApproval(autoApprove: Boolean) { - status = CompletionStatus.APPROVED - statusLabel.text = "Approved${if (autoApprove) " (auto)" else ""}" - isCompleted = true - setupMainContent() - onApprove(autoApprove) - } - - private fun handleRejection() { - status = CompletionStatus.REJECTED - statusLabel.text = "Rejected" - setupMainContent() - onReject() - } - - fun showExecutionStatus() { - status = CompletionStatus.EXECUTING - statusLabel.text = "Executing..." - isCompleted = true - setupMainContent() - } - - fun showCompletionResult(result: String, isSuccess: Boolean) { - status = if (isSuccess) CompletionStatus.COMPLETED else CompletionStatus.FAILED - statusLabel.text = if (isSuccess) "Completed" else "Failed" - isCompleted = true - setupMainContent() - } - - private fun setupKeyboardShortcuts() { - if (isCompleted) return - - isFocusable = true - - val inputMap = getInputMap(WHEN_IN_FOCUSED_WINDOW) - val actionMap = getActionMap() - - inputMap.put(KeyStroke.getKeyStroke(KeyEvent.VK_Y, 0), "approve") - actionMap.put("approve", object : AbstractAction() { - override fun actionPerformed(e: ActionEvent) { - if (!isCompleted) handleApproval(false) - } - }) - - inputMap.put(KeyStroke.getKeyStroke(KeyEvent.VK_A, 0), "autoApprove") - actionMap.put("autoApprove", object : AbstractAction() { - override fun actionPerformed(e: ActionEvent) { - if (!isCompleted) handleApproval(true) - } - }) - - inputMap.put(KeyStroke.getKeyStroke(KeyEvent.VK_N, 0), "reject") - actionMap.put("reject", object : AbstractAction() { - override fun actionPerformed(e: ActionEvent) { - if (!isCompleted) handleRejection() - } - }) - - inputMap.put(KeyStroke.getKeyStroke(KeyEvent.VK_SPACE, 0), "toggle") - inputMap.put(KeyStroke.getKeyStroke(KeyEvent.VK_ENTER, 0), "toggle") - actionMap.put("toggle", object : AbstractAction() { - override fun actionPerformed(e: ActionEvent) { - if (!isCompleted && fullParameters != null && fullParameters.isNotEmpty()) { - toggleExpanded() - } - } - }) - } -} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/ui/mcp/McpAnimations.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/ui/mcp/McpAnimations.kt deleted file mode 100644 index bddf8f85..00000000 --- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/ui/mcp/McpAnimations.kt +++ /dev/null @@ -1,69 +0,0 @@ -package ee.carlrobert.codegpt.toolwindow.ui.mcp - -import java.awt.AlphaComposite -import java.awt.Graphics -import java.awt.Graphics2D -import javax.swing.JComponent -import javax.swing.Timer -import kotlin.math.pow - -object McpAnimations { - - fun animateHeight( - component: JComponent, - targetHeight: Int, - duration: Int = 200, - onComplete: (() -> Unit)? = null - ) { - val startHeight = component.height - animate(0f, 1f, duration) { progress -> - val currentHeight = (startHeight + (targetHeight - startHeight) * progress).toInt() - component.preferredSize = component.preferredSize.apply { height = currentHeight } - component.revalidate() - component.repaint() - if (progress >= 1f) { - onComplete?.invoke() - } - } - } - - private fun animate(start: Float, end: Float, duration: Int, onUpdate: (Float) -> Unit) { - val startTime = System.currentTimeMillis() - val timer = Timer(16) { _ -> - val elapsed = System.currentTimeMillis() - startTime - val progress = (elapsed.toFloat() / duration).coerceIn(0f, 1f) - val value = start + (end - start) * easeInOutCubic(progress) - onUpdate(value) - } - - timer.addActionListener { - val elapsed = System.currentTimeMillis() - startTime - if (elapsed >= duration) { - timer.stop() - onUpdate(end) - } - } - - timer.start() - } - - private fun easeInOutCubic(t: Float): Float { - return if (t < 0.5f) { - 4 * t * t * t - } else { - 1 - (-2.0 * t + 2.0).pow(3.0).toFloat() / 2 - } - } -} - -fun JComponent.paintWithAlpha(g: Graphics, paintAction: () -> Unit) { - val alpha = getClientProperty("alpha") as? Float ?: 1f - if (alpha < 1f) { - val g2d = g.create() as Graphics2D - g2d.composite = AlphaComposite.getInstance(AlphaComposite.SRC_OVER, alpha) - paintAction() - g2d.dispose() - } else { - paintAction() - } -} \ No newline at end of file diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/ui/mcp/McpApprovalPanel.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/ui/mcp/McpApprovalPanel.kt deleted file mode 100644 index 169b2fd4..00000000 --- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/ui/mcp/McpApprovalPanel.kt +++ /dev/null @@ -1,234 +0,0 @@ -package ee.carlrobert.codegpt.toolwindow.ui.mcp - -import com.intellij.ui.components.ActionLink -import com.intellij.ui.components.JBLabel -import com.intellij.ui.components.JBPanel -import com.intellij.ui.components.JBTextArea -import com.intellij.util.ui.JBUI -import ee.carlrobert.llm.client.openai.completion.response.ToolCall -import java.awt.BorderLayout -import java.awt.FlowLayout -import java.awt.Graphics -import javax.swing.BorderFactory -import javax.swing.JPanel - -/** - * Modern approval panel following JetBrains UI/UX guidelines. - * Features a clean card-like design with proper spacing and visual hierarchy. - */ -class McpApprovalPanel( - private val toolCall: ToolCall, - private val serverName: String, - private val onApprove: (autoApprove: Boolean) -> Unit, - private val onReject: () -> Unit, -) : JBPanel() { - - private var isCompleted = false - private val fullParameters = McpParameterPreview.formatDetailed(toolCall.function.arguments) - - init { - layout = BorderLayout() - background = JBUI.CurrentTheme.ToolWindow.background() - border = BorderFactory.createCompoundBorder( - JBUI.Borders.empty(2, 0), - BorderFactory.createCompoundBorder( - BorderFactory.createLineBorder( - JBUI.CurrentTheme.CustomFrameDecorations.separatorForeground(), - 1 - ), - JBUI.Borders.empty(1) - ) - ) - - setupUI() - - toolTipText = createParametersTooltip() - } - - private fun setupUI() { - val contentPanel = JBPanel>().apply { - layout = BorderLayout() - isOpaque = false - border = JBUI.Borders.empty(8, 12) - } - - contentPanel.add(createHeaderSection(), BorderLayout.NORTH) - contentPanel.add(createActionsSection(), BorderLayout.CENTER) - - add(contentPanel, BorderLayout.CENTER) - } - - private fun createHeaderSection(): JPanel { - return JBPanel>().apply { - layout = BorderLayout() - isOpaque = false - border = JBUI.Borders.emptyBottom(8) - - val approvalLabel = JBLabel("Allow ProxyAI to execute this tool?").apply { - font = JBUI.Fonts.label() - foreground = JBUI.CurrentTheme.Label.foreground() - } - add(approvalLabel, BorderLayout.NORTH) - - val toolInfoPanel = JBPanel>().apply { - layout = FlowLayout(FlowLayout.LEFT, 0, 0) - isOpaque = false - border = JBUI.Borders.emptyTop(4) - } - - toolInfoPanel.add(JBLabel(toolCall.function.name).apply { - font = JBUI.Fonts.label().asBold() - toolTipText = createParametersTooltip() - }) - - toolInfoPanel.add(JBLabel(" • ").apply { - foreground = JBUI.CurrentTheme.Label.disabledForeground() - }) - - toolInfoPanel.add(JBLabel(serverName).apply { - font = JBUI.Fonts.label() - foreground = JBUI.CurrentTheme.Label.disabledForeground() - }) - - add(toolInfoPanel, BorderLayout.CENTER) - } - } - - - private fun createActionsSection(): JPanel { - return JBPanel>().apply { - layout = BorderLayout() - isOpaque = false - - if (fullParameters != null && fullParameters.isNotEmpty()) { - val paramsPanel = JBPanel>().apply { - layout = BorderLayout() - isOpaque = false - border = JBUI.Borders.emptyBottom(8) - } - - val paramsText = fullParameters.entries.joinToString("\n") { (key, value) -> - formatParameterEntry(key, value) - } - - paramsPanel.add(createWrappedParamsArea(paramsText), BorderLayout.CENTER) - add(paramsPanel, BorderLayout.NORTH) - } - - val linksPanel = JBPanel>().apply { - layout = FlowLayout(FlowLayout.LEFT, 0, 0) - isOpaque = false - } - - linksPanel.add(ActionLink("Yes") { - handleApproval(false) - }) - - linksPanel.add(JBLabel(" | ").apply { - foreground = JBUI.CurrentTheme.Label.disabledForeground() - }) - - linksPanel.add(ActionLink("Always accept for this session") { - handleApproval(true) - }) - - linksPanel.add(JBLabel(" | ").apply { - foreground = JBUI.CurrentTheme.Label.disabledForeground() - }) - - linksPanel.add(ActionLink("No") { - handleRejection() - }) - - add(linksPanel, BorderLayout.CENTER) - } - } - - private fun handleApproval(autoApprove: Boolean) { - isCompleted = true - isVisible = false - parent?.remove(this) - parent?.revalidate() - parent?.repaint() - onApprove(autoApprove) - } - - private fun handleRejection() { - isCompleted = true - transformToCompactStatus("✗ Rejected", false) - onReject() - } - - private fun transformToCompactStatus(statusText: String, success: Boolean) { - removeAll() - - layout = FlowLayout(FlowLayout.LEFT, 0, 0) - border = JBUI.Borders.empty(8, 12) - background = if (success) { - JBUI.CurrentTheme.NotificationInfo.backgroundColor() - } else { - JBUI.CurrentTheme.NotificationError.backgroundColor() - } - - // Removed MCP icon - add(JBLabel(statusText).apply { - foreground = if (success) { - JBUI.CurrentTheme.NotificationInfo.foregroundColor() - } else { - JBUI.CurrentTheme.NotificationError.foregroundColor() - } - }) - - revalidate() - repaint() - } - - private fun createParametersTooltip(): String { - val params = fullParameters ?: return toolCall.function.name - - val paramHtml = if (params.isNotEmpty()) { - params.entries.joinToString("
") { (key, value) -> - "$key: ${value.toString().replace("\n", "
").replace(" ", " ")}" - } - } else { - "No parameters" - } - - return "" + - "Tool: ${toolCall.function.name}
" + - "Server: $serverName

" + - "Parameters:
" + - paramHtml + - "" - } - - override fun paintComponent(g: Graphics) { - paintWithAlpha(g) { - super.paintComponent(g) - } - } - - private fun formatParameterEntry(key: String, value: String): String { - return "$key: ${value.replace("\n", "\n ")}" - } - - private fun createWrappedParamsArea(text: String): JBTextArea { - return JBTextArea(text).apply { - isEditable = false - isFocusable = false - lineWrap = true - wrapStyleWord = false - font = JBUI.Fonts.smallFont() - foreground = JBUI.CurrentTheme.Label.foreground() - border = BorderFactory.createCompoundBorder( - BorderFactory.createLineBorder( - JBUI.CurrentTheme.CustomFrameDecorations.separatorForeground(), - 1 - ), - JBUI.Borders.empty(4, 8) - ) - background = JBUI.CurrentTheme.CustomFrameDecorations.paneBackground() - isOpaque = true - } - } -} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/ui/mcp/McpMessageListItem.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/ui/mcp/McpMessageListItem.kt deleted file mode 100644 index a68a58a3..00000000 --- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/ui/mcp/McpMessageListItem.kt +++ /dev/null @@ -1,47 +0,0 @@ -package ee.carlrobert.codegpt.toolwindow.ui.mcp - -import com.intellij.ui.components.JBLabel -import com.intellij.util.ui.JBUI -import ee.carlrobert.codegpt.Icons -import ee.carlrobert.codegpt.toolwindow.ui.BaseMessagePanel -import java.awt.Graphics -import javax.swing.JPanel -import javax.swing.SwingConstants - -class McpMessageListItem( - content: JPanel -) : BaseMessagePanel() { - - init { - addContent(content) - border = JBUI.Borders.empty() - } - - override fun createDisplayNameLabel(): JBLabel { - return JBLabel( - "Tool", - null, - SwingConstants.LEADING - ).apply { - iconTextGap = 6 - font = JBUI.Fonts.label() - foreground = JBUI.CurrentTheme.Label.disabledForeground() - } - } - - companion object { - fun fromToolPanel(toolPanel: CompactToolPanel): McpMessageListItem { - return McpMessageListItem(toolPanel) - } - - fun fromGroupPanel(groupPanel: ToolGroupPanel): McpMessageListItem { - return McpMessageListItem(groupPanel) - } - } - - override fun paintComponent(g: Graphics) { - paintWithAlpha(g) { - super.paintComponent(g) - } - } -} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/ui/mcp/McpParameterPreview.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/ui/mcp/McpParameterPreview.kt deleted file mode 100644 index 7bb5e028..00000000 --- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/ui/mcp/McpParameterPreview.kt +++ /dev/null @@ -1,90 +0,0 @@ -package ee.carlrobert.codegpt.toolwindow.ui.mcp - -import com.fasterxml.jackson.databind.ObjectMapper -import com.intellij.openapi.util.text.StringUtil - -object McpParameterPreview { - - private val objectMapper = ObjectMapper() - - fun formatDetailed(jsonArgs: String?): Map? { - if (jsonArgs.isNullOrEmpty() || jsonArgs == "{}") return null - - return try { - val args = objectMapper.readValue(jsonArgs, Map::class.java) as Map - args.mapValues { (_, value) -> formatValueDetailed(value) } - } catch (e: Exception) { - null - } - } - - private fun formatValueDetailed(value: Any?): String { - return when (value) { - null -> "null" - is String -> value - is Number -> value.toString() - is Boolean -> value.toString() - is List<*> -> formatListDetailed(value) - is Map<*, *> -> formatMapDetailed(value) - is Array<*> -> formatListDetailed(value.toList()) - else -> value.toString() - } - } - - private fun formatListDetailed(list: List<*>): String { - return if (list.isEmpty()) { - "[]" - } else if (list.size <= 3) { - list.joinToString("\n") { "• ${formatValueDetailed(it)}" } - } else { - val preview = list.take(3).joinToString("\n") { "• ${formatValueDetailed(it)}" } - "$preview\n• ... (${list.size - 3} more items)" - } - } - - private fun formatMapDetailed(map: Map<*, *>): String { - return if (map.isEmpty()) { - "{}" - } else { - map.entries.joinToString("\n") { (k, v) -> - "${k}: ${formatValueDetailed(v)}" - } - } - } - - fun generateTooltip(jsonArgs: String?): String? { - if (jsonArgs.isNullOrEmpty() || jsonArgs == "{}") return null - - return try { - val args = objectMapper.readValue(jsonArgs, Map::class.java) as Map - if (args.isEmpty()) return null - - "" + - args.entries.joinToString("
") { (key, value) -> - "$key: ${formatTooltipValue(value)}" - } + - "" - } catch (e: Exception) { - null - } - } - - private fun formatTooltipValue(value: Any?): String { - val formatted = when (value) { - null -> "null" - is String -> if (value.length > 100) { - StringUtil.shortenTextWithEllipsis(value, 100, 10) - } else value - - is List<*> -> "[${value.size} items]" - is Map<*, *> -> "{${value.size} properties}" - else -> value.toString() - } - - return formatted - .replace("&", "&") - .replace("<", "<") - .replace(">", ">") - .replace("\"", """) - } -} \ No newline at end of file diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/ui/mcp/McpStatusPanel.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/ui/mcp/McpStatusPanel.kt deleted file mode 100644 index ac5750db..00000000 --- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/ui/mcp/McpStatusPanel.kt +++ /dev/null @@ -1,430 +0,0 @@ -package ee.carlrobert.codegpt.toolwindow.ui.mcp - -import com.intellij.icons.AllIcons -import com.intellij.ui.JBColor -import com.intellij.ui.components.ActionLink -import com.intellij.ui.components.JBLabel -import com.intellij.ui.components.JBPanel -import com.intellij.util.ui.JBUI -import com.intellij.util.ui.components.BorderLayoutPanel -import ee.carlrobert.codegpt.Icons -import java.awt.BorderLayout -import java.awt.FlowLayout -import java.awt.Cursor -import java.awt.Graphics -import java.awt.Toolkit -import java.awt.datatransfer.StringSelection -import java.time.Duration -import java.time.LocalDateTime -import javax.swing.* - -class McpStatusPanel( - private val toolName: String, - private val serverName: String, - initialStatus: String = "Executing...", - private val args: Map? = null, -) : JBPanel() { - - private val startTime = LocalDateTime.now() - private val statusLabel = JBLabel(initialStatus) - private val spinner = JProgressBar().apply { - isIndeterminate = true - preferredSize = JBUI.size(60, 8) - border = JBUI.Borders.empty() - } - - private var resultPreviewLabel: JBLabel? = null - private var resultActionsPanel: JPanel? = null - private var currentResult: String? = null - private var currentError: String? = null - private var isCompleted = false - private var isExpanded = false - - private lateinit var contentContainer: BorderLayoutPanel - private val expandIcon = JBLabel(AllIcons.General.ArrowRight).apply { - border = JBUI.Borders.emptyRight(6) - } - - init { - layout = BorderLayout() - isOpaque = false - border = JBUI.Borders.empty(8, 0) - - setupUI() - } - - private fun setupUI() { - layout = BorderLayout() - isOpaque = false - border = JBUI.Borders.empty(8, 0) - - val headerPanel = JBPanel>().apply { - layout = FlowLayout(FlowLayout.LEFT, 0, 0) - isOpaque = false - border = JBUI.Borders.empty(0, 0, 4, 0) - cursor = Cursor.getPredefinedCursor(Cursor.HAND_CURSOR) - addMouseListener(object : java.awt.event.MouseAdapter() { - override fun mouseClicked(e: java.awt.event.MouseEvent?) { - toggleExpanded() - } - }) - } - - headerPanel.add(expandIcon) - - val infoLabel = JBLabel("$toolName • $serverName").apply { - font = JBUI.Fonts.label() - foreground = JBUI.CurrentTheme.Label.disabledForeground() - } - headerPanel.add(infoLabel) - headerPanel.add(Box.createHorizontalStrut(12)) - - statusLabel.font = JBUI.Fonts.smallFont() - statusLabel.foreground = JBUI.CurrentTheme.Label.disabledForeground() - headerPanel.add(statusLabel) - headerPanel.add(Box.createHorizontalStrut(8)) - - headerPanel.add(spinner) - - add(headerPanel, BorderLayout.NORTH) - contentContainer = BorderLayoutPanel() - contentContainer.isVisible = false - add(contentContainer, BorderLayout.CENTER) - } - - fun complete(success: Boolean, message: String? = null, result: String? = null) { - SwingUtilities.invokeLater { - spinner.isVisible = false - isCompleted = true - - val duration = Duration.between(startTime, LocalDateTime.now()) - val durationText = formatDuration(duration) - - if (success) { - val lowerName = toolName.lowercase() - if (lowerName.contains("write") || lowerName.contains("edit")) { - val filePath = args?.get("file_path") as? String ?: args?.get("path") as? String - val created = (result ?: "").lowercase().contains("created") - statusLabel.text = "${if (created) "Created" else "Updated"}${filePath?.let { ": $it" } ?: ""} • $durationText" - statusLabel.foreground = JBUI.CurrentTheme.Label.foreground() - val content = args?.get("content") as? String - currentResult = content ?: result - showResultPreview(currentResult ?: "") - } else { - currentResult = result - statusLabel.text = "${message ?: "Completed"} • $durationText" - statusLabel.foreground = JBUI.CurrentTheme.Label.foreground() - result?.let { showResultPreview(it) } - } - } else { - currentError = message ?: result ?: "Failed" - statusLabel.text = "Failed • $durationText" - statusLabel.foreground = JBColor.RED - - currentError?.let { showErrorPreview(it) } - } - - updateExpansionUi() - revalidate() - repaint() - } - } - - private fun showResultPreview(result: String) { - if (result.isBlank()) return - - val previewText = buildToolSpecificPreview(result).ifEmpty { formatResultPreview(result) } - if (previewText.isNotEmpty()) { - removeResultComponents() - - val resultPanel = JBPanel>().apply { - layout = BorderLayout() - isOpaque = false - border = JBUI.Borders.empty(8, 24, 4, 0) - } - - resultPreviewLabel = JBLabel(previewText).apply { - font = JBUI.Fonts.smallFont() - foreground = JBUI.CurrentTheme.Label.disabledForeground() - } - - resultActionsPanel = JPanel(FlowLayout(FlowLayout.LEFT, 0, 2)).apply { - isOpaque = false - border = JBUI.Borders.empty(4, 0, 0, 0) - - add(createShowMoreLink()) - add(Box.createHorizontalStrut(12)) - add(createCopyLink()) - } - - resultPanel.add(resultPreviewLabel!!, BorderLayout.CENTER) - resultPanel.add(resultActionsPanel!!, BorderLayout.SOUTH) - - contentContainer.removeAll() - contentContainer.add(resultPanel, BorderLayout.CENTER) - updateExpansionUi() - } - } - - private fun showErrorPreview(error: String) { - val previewText = error.trim().take(100).let { - if (it.length < error.trim().length) "$it..." else it - } - - removeResultComponents() - - val errorPanel = JBPanel>().apply { - layout = BorderLayout() - isOpaque = false - border = JBUI.Borders.empty(8, 24, 4, 0) - } - - resultPreviewLabel = JBLabel("Error: $previewText").apply { - font = JBUI.Fonts.smallFont() - foreground = JBColor.RED - } - - resultActionsPanel = JPanel(FlowLayout(FlowLayout.LEFT, 0, 2)).apply { - isOpaque = false - border = JBUI.Borders.empty(4, 0, 0, 0) - - add(createShowErrorDetailsLink()) - add(Box.createHorizontalStrut(12)) - add(createCopyLink()) - } - - errorPanel.add(resultPreviewLabel!!, BorderLayout.CENTER) - errorPanel.add(resultActionsPanel!!, BorderLayout.SOUTH) - - contentContainer.removeAll() - contentContainer.add(errorPanel, BorderLayout.CENTER) - updateExpansionUi() - } - - private fun removeResultComponents() { - contentContainer.removeAll() - } - - private fun createShowMoreLink(): ActionLink { - return ActionLink("Show details") { - showFullContent(false) - }.apply { - font = JBUI.Fonts.smallFont() - } - } - - private fun createShowErrorDetailsLink(): ActionLink { - return ActionLink("Show details") { - showFullContent(true) - }.apply { - font = JBUI.Fonts.smallFont() - } - } - - private fun createCopyLink(): ActionLink { - return ActionLink("Copy") { - val content = currentResult ?: currentError ?: return@ActionLink - val selection = StringSelection(content) - Toolkit.getDefaultToolkit().systemClipboard.setContents(selection, null) - }.apply { - font = JBUI.Fonts.smallFont() - } - } - - private fun showFullContent(isError: Boolean) { - val content = currentResult ?: currentError ?: return - val dialog = McpResultDialog(toolName, serverName, content, isError) - dialog.isVisible = true - } - - private fun toggleExpanded() { - isExpanded = !isExpanded - updateExpansionUi() - } - - private fun updateExpansionUi() { - val hasContent = (currentResult?.isNotBlank() == true) || (currentError?.isNotBlank() == true) - contentContainer.isVisible = hasContent && isExpanded - expandIcon.icon = if (hasContent && isExpanded) AllIcons.General.ArrowDown else AllIcons.General.ArrowRight - } - - private fun formatResultPreview(result: String): String { - val cleaned = result - .replace("\n", " ") - .replace("\r", " ") - .replace(" ", " ") - .trim() - - return when { - cleaned.contains("search results") || cleaned.contains("results found") -> { - val countMatch = - Regex("(\\d+)\\s*(?:search\\s*)?results?").find(cleaned.lowercase()) - if (countMatch != null) { - "Found ${countMatch.groupValues[1]} results" - } else { - "Search completed" - } - } - - cleaned.contains("SELECT") || cleaned.contains("INSERT") || cleaned.contains("UPDATE") -> { - val rowMatch = Regex("(\\d+)\\s*rows?").find(cleaned) - if (rowMatch != null) { - "Query returned ${rowMatch.groupValues[1]} rows" - } else { - "Query executed successfully" - } - } - - cleaned.length > 150 -> { - cleaned.take(150) + "..." - } - - else -> cleaned - } - } - - private fun buildToolSpecificPreview(result: String): String { - val name = toolName.lowercase() - - return when { - name.contains("bash") -> { - val cmd = (args?.get("command") as? String) - ?: Regex("^Command:\\s*(.*)$", RegexOption.MULTILINE).find(result)?.groupValues?.getOrNull(1) - cmd?.let { "Command: ${it.trim()}" } ?: "" - } - name.startsWith("read") -> { - val filePath = args?.get("file_path") as? String ?: args?.get("path") as? String - val count = Regex("^\\d+\\t", RegexOption.MULTILINE).findAll(result).count() - if (filePath != null && count > 0) { - "Read ${filePath} • ${count} lines" - } else "" - } - name.contains("grep") || name.contains("search") -> { - val total = Regex("Total matches:\\s*(\\d+)").find(result)?.groupValues?.getOrNull(1) - val pattern = args?.get("pattern") as? String - val path = args?.get("path") as? String - buildString { - if (!pattern.isNullOrBlank()) append("pattern: \"${pattern.take(40)}${if (pattern.length>40) "..." else ""}\"") - if (!path.isNullOrBlank()) { - if (isNotEmpty()) append(" • ") - append("in: $path") - } - if (!total.isNullOrBlank()) { - if (isNotEmpty()) append(" • ") - append("${total} matches") - } - } - } - name.contains("write") || name.contains("edit") -> { - val filePath = args?.get("file_path") as? String ?: args?.get("path") as? String - val content = args?.get("content") as? String - val contentPreview = content?.let { it.replace("\n", " ").take(150) + if (it.length>150) "..." else "" } - buildString { - if (filePath != null) append("Wrote ${filePath}") - if (!contentPreview.isNullOrBlank()) { - if (isNotEmpty()) append(" • ") - append(contentPreview) - } - } - } - else -> "" - } - } - - private fun formatDuration(duration: Duration): String { - val seconds = duration.toMillis() / 1000.0 - return when { - seconds < 1 -> "${duration.toMillis()}ms" - seconds < 60 -> "%.1fs".format(seconds) - else -> "${duration.toSeconds()}s" - } - } - - override fun paintComponent(g: Graphics) { - paintWithAlpha(g) { - super.paintComponent(g) - } - } -} - -class McpResultDialog( - private val toolName: String, - private val serverName: String, - private val content: String, - private val isError: Boolean -) : JDialog() { - - init { - title = if (isError) "Tool Error: $toolName" else "Tool Result: $toolName" - isModal = true - setupUI() - pack() - setLocationRelativeTo(null) - } - - private fun setupUI() { - layout = BorderLayout() - - val headerPanel = JPanel(FlowLayout(FlowLayout.LEFT)).apply { - border = JBUI.Borders.empty() - isOpaque = false - - add(JLabel(if (isError) AllIcons.General.Error else null)) - add(Box.createHorizontalStrut(8)) - add(JLabel("Tool: $toolName").apply { - font = JBUI.Fonts.label().asBold() - }) - add(JLabel(" • ").apply { - foreground = JBUI.CurrentTheme.Label.disabledForeground() - }) - add(JLabel(serverName).apply { - foreground = JBUI.CurrentTheme.Label.disabledForeground() - }) - } - - val textArea = JTextArea(content).apply { - isEditable = false - font = JBUI.Fonts.label() - lineWrap = true - wrapStyleWord = true - caretPosition = 0 - border = JBUI.Borders.empty() - } - - val scrollPane = JScrollPane(textArea).apply { - preferredSize = JBUI.size(600, 400) - border = - JBUI.Borders.customLine(JBUI.CurrentTheme.CustomFrameDecorations.separatorForeground()) - } - - val buttonPanel = JPanel(FlowLayout(FlowLayout.RIGHT)).apply { - border = JBUI.Borders.empty(8, 0, 0, 0) - isOpaque = false - - val copyButton = JButton("Copy to Clipboard").apply { - addActionListener { - val selection = StringSelection(content) - Toolkit.getDefaultToolkit().systemClipboard.setContents(selection, null) - } - } - - val closeButton = JButton("Close").apply { - addActionListener { dispose() } - } - - add(copyButton) - add(Box.createHorizontalStrut(8)) - add(closeButton) - } - - add(headerPanel, BorderLayout.NORTH) - add(scrollPane, BorderLayout.CENTER) - add(buttonPanel, BorderLayout.SOUTH) - - rootPane.registerKeyboardAction( - { dispose() }, - KeyStroke.getKeyStroke("ESCAPE"), - JComponent.WHEN_IN_FOCUSED_WINDOW - ) - } -} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/ui/mcp/McpToolResultPanel.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/ui/mcp/McpToolResultPanel.kt deleted file mode 100644 index b1a6909e..00000000 --- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/ui/mcp/McpToolResultPanel.kt +++ /dev/null @@ -1,197 +0,0 @@ -package ee.carlrobert.codegpt.toolwindow.ui.mcp - -import com.intellij.icons.AllIcons -import com.intellij.openapi.ui.DialogWrapper -import com.intellij.ui.JBColor -import com.intellij.ui.components.ActionLink -import com.intellij.ui.components.JBLabel -import com.intellij.ui.components.JBPanel -import com.intellij.ui.components.JBScrollPane -import com.intellij.util.ui.JBUI -import com.intellij.util.ui.UIUtil -import com.intellij.util.ui.components.BorderLayoutPanel -import java.awt.BorderLayout -import java.awt.FlowLayout -import java.awt.Graphics -import javax.swing.* - -class McpToolResultPanel( - private val toolName: String, - private val result: String, - private val serverName: String = "", - private val isError: Boolean = false -) : JBPanel() { - - init { - layout = BorderLayout() - background = if (isError) { - UIUtil.getPanelBackground() - } else { - UIUtil.getListBackground() - } - border = BorderFactory.createCompoundBorder( - JBUI.Borders.empty(4, 0), - BorderFactory.createCompoundBorder( - BorderFactory.createLineBorder( - if (isError) JBColor.RED.darker() else JBUI.CurrentTheme.CustomFrameDecorations.separatorForeground(), - 1 - ), - JBUI.Borders.empty(8) - ) - ) - - setupUI() - } - - private fun setupUI() { - val headerPanel = JBPanel>().apply { - layout = FlowLayout(FlowLayout.LEFT, 0, 0) - isOpaque = false - - add(JBLabel(if (isError) AllIcons.General.Error else null)) - add(Box.createHorizontalStrut(8)) - add(JBLabel("Tool: $toolName").apply { - font = JBUI.Fonts.label().asBold() - foreground = if (isError) JBColor.RED else JBUI.CurrentTheme.Label.foreground() - }) - } - - add(headerPanel, BorderLayout.NORTH) - - val resultPreview = formatResultPreview(result) - if (resultPreview.isNotEmpty()) { - val previewPanel = JBPanel>().apply { - layout = BorderLayout() - isOpaque = false - border = JBUI.Borders.emptyTop(4) - } - - val resultLabel = JBLabel(resultPreview).apply { - font = JBUI.Fonts.smallFont() - foreground = JBUI.CurrentTheme.Label.disabledForeground() - - toolTipText = formatTooltip(result) - } - previewPanel.add(resultLabel, BorderLayout.CENTER) - - if (resultPreview.endsWith("...") || result.length > 150) { - val showMoreLink = ActionLink("Show more") { - showFullContent() - }.apply { - font = JBUI.Fonts.smallFont() - border = JBUI.Borders.emptyLeft(8) - } - previewPanel.add(showMoreLink, BorderLayout.EAST) - } - - add(previewPanel, BorderLayout.CENTER) - } - - toolTipText = formatTooltip(result) - } - - private fun formatResultPreview(result: String): String { - val cleaned = result - .replace("\n", " ") - .replace("\r", " ") - .replace(" ", " ") - .trim() - - return when { - isError -> { - if (cleaned.length > 100) { - cleaned.take(100) + "..." - } else { - cleaned - } - } - - cleaned.contains("search results") || cleaned.contains("results found") -> { - val countMatch = - Regex("(\\d+)\\s*(?:search\\s*)?results?").find(cleaned.lowercase()) - if (countMatch != null) { - "Found ${countMatch.groupValues[1]} results" - } else { - "Search completed" - } - } - - cleaned.contains("SELECT") || cleaned.contains("INSERT") || cleaned.contains("UPDATE") -> { - val rowMatch = Regex("(\\d+)\\s*rows?").find(cleaned) - if (rowMatch != null) { - "Query returned ${rowMatch.groupValues[1]} rows" - } else { - "Query executed successfully" - } - } - - cleaned.length > 150 -> { - cleaned.take(150) + "..." - } - - else -> cleaned - } - } - - private fun formatTooltip(result: String): String { - return "" + - "Tool: $toolName

" + - result.replace("\n", "
") - .replace("\t", "    ") - .replace(" ", " ") + - "" - } - - private fun showFullContent() { - val dialog = object : DialogWrapper(false) { - init { - title = "Tool Result: $toolName" - init() - } - - override fun createCenterPanel(): JComponent { - val panel = BorderLayoutPanel().apply { - border = JBUI.Borders.empty() - } - - val headerPanel = JPanel(FlowLayout(FlowLayout.LEFT)).apply { - add(JLabel(if (isError) AllIcons.General.Error else null)) - add(Box.createHorizontalStrut(8)) - add(JLabel("Tool: $toolName").apply { - font = JBUI.Fonts.label().asBold() - }) - add(JLabel(" • ").apply { - foreground = JBUI.CurrentTheme.Label.disabledForeground() - }) - add(JLabel(serverName).apply { - foreground = JBUI.CurrentTheme.Label.disabledForeground() - }) - } - panel.add(headerPanel, BorderLayout.NORTH) - - val textArea = JTextArea(result).apply { - isEditable = false - font = JBUI.Fonts.label() - lineWrap = true - wrapStyleWord = true - caretPosition = 0 - } - - val scrollPane = JBScrollPane(textArea).apply { - preferredSize = JBUI.size(600, 400) - } - panel.add(scrollPane, BorderLayout.CENTER) - - return panel - } - } - - dialog.show() - } - - override fun paintComponent(g: Graphics) { - paintWithAlpha(g) { - super.paintComponent(g) - } - } -} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/ui/mcp/ToolGroupPanel.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/ui/mcp/ToolGroupPanel.kt deleted file mode 100644 index 77ed1f7e..00000000 --- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/ui/mcp/ToolGroupPanel.kt +++ /dev/null @@ -1,218 +0,0 @@ -package ee.carlrobert.codegpt.toolwindow.ui.mcp - -import com.intellij.openapi.ui.VerticalFlowLayout -import com.intellij.ui.JBColor -import com.intellij.ui.components.ActionLink -import com.intellij.ui.components.JBLabel -import com.intellij.ui.components.JBPanel -import com.intellij.util.ui.JBUI -import ee.carlrobert.llm.client.openai.completion.response.ToolCall -import java.awt.BorderLayout -import java.awt.FlowLayout -import javax.swing.BorderFactory -import javax.swing.Box -import javax.swing.JComponent -import javax.swing.JPanel - -class ToolGroupPanel( - private val toolCalls: List>, - private val onApprove: (ToolCall, CompactToolPanel, Boolean) -> Unit, - private val onReject: (ToolCall, CompactToolPanel) -> Unit, - private val onBatchApprove: (List, autoApprove: Boolean) -> Unit, - private val onBatchReject: (List) -> Unit -) : JBPanel() { - - private val toolPanels = mutableListOf() - private val containerPanel = JBPanel>() - - companion object { - private const val TOOL_REQUESTS_TEXT = "tool requests from ProxyAI" - private const val APPROVE_ALL_TEXT = "Approve all" - private const val REJECT_ALL_TEXT = "Reject all" - private const val SEPARATOR_TEXT = " | " - } - - init { - layout = BorderLayout() - isOpaque = false - border = JBUI.Borders.empty() - - setupUI() - } - - private fun removeToolPanel(toolCall: ToolCall) { - toolPanels.find { it.toolCallId == toolCall.id }?.let { panel -> - containerPanel.components - .filterIsInstance() - .find { wrapper -> wrapper.components.contains(panel) } - ?.let { wrapperPanel -> - containerPanel.remove(wrapperPanel) - containerPanel.revalidate() - containerPanel.repaint() - } - } - } - - private fun cleanupGroupPanel() { - toolPanels.forEach { it.isVisible = false } - containerPanel.removeAll() - containerPanel.revalidate() - containerPanel.repaint() - - if (componentCount > 0) { - getComponent(0).isVisible = false - } - } - - private fun removeSelfFromParent() { - isVisible = false - parent?.apply { - remove(this@ToolGroupPanel) - revalidate() - repaint() - } - } - - private fun setupUI() { - if (toolCalls.size > 1) { - add(createBatchActionsHeader(), BorderLayout.NORTH) - } - - containerPanel.apply { - layout = VerticalFlowLayout(VerticalFlowLayout.TOP, 0, 4, true, false) - isOpaque = false - border = JBUI.Borders.empty() - } - - toolCalls.forEach { (toolCall, serverName) -> - lateinit var panel: CompactToolPanel - - panel = CompactToolPanel( - toolCall = toolCall, - serverName = serverName, - onApprove = { autoApprove -> - handleIndividualApproval(toolCall, panel, autoApprove) - }, - onReject = { - handleIndividualRejection(toolCall, panel) - }, - ) - - toolPanels.add(panel) - containerPanel.add(wrapInBorder(panel)) - } - - add(containerPanel, BorderLayout.CENTER) - } - - private fun createBatchActionsHeader(): JPanel { - return JBPanel>().apply { - layout = BorderLayout() - isOpaque = false - border = JBUI.Borders.empty(0, 0, 8, 0) - - val contentPanel = JBPanel>().apply { - layout = FlowLayout(FlowLayout.LEFT, 0, 0) - isOpaque = false - } - - contentPanel.add(JBLabel("${toolCalls.size} $TOOL_REQUESTS_TEXT").apply { - font = JBUI.Fonts.label() - foreground = JBUI.CurrentTheme.Label.foreground() - }) - - contentPanel.add(Box.createHorizontalStrut(12)) - - actions { - action(APPROVE_ALL_TEXT) { handleBatchApproval() } - separator() - action(REJECT_ALL_TEXT) { handleBatchRejection() } - }.forEach { contentPanel.add(it) } - - add(contentPanel, BorderLayout.CENTER) - } - } - - private fun actions(block: ActionLinksBuilder.() -> Unit): List { - return ActionLinksBuilder().apply(block).build() - } - - private class ActionLinksBuilder { - private val components = mutableListOf() - - fun action(text: String, action: () -> Unit) { - components.add(ActionLink(text) { action() }) - } - - fun separator() { - components.add(JBLabel(SEPARATOR_TEXT).apply { - foreground = JBUI.CurrentTheme.Label.disabledForeground() - }) - } - - fun build(): List = components - } - - private fun wrapInBorder(panel: CompactToolPanel): JPanel { - return JBPanel>().apply { - layout = BorderLayout() - isOpaque = false - border = BorderFactory.createCompoundBorder( - JBUI.Borders.customLine(JBColor.border(), 0, 0, 1, 0), - JBUI.Borders.empty(4, 0) - ) - add(panel, BorderLayout.CENTER) - } - } - - private fun handleIndividualApproval(toolCall: ToolCall, panel: CompactToolPanel, autoApprove: Boolean) { - onApprove(toolCall, panel, autoApprove) - updateHeader() - - if (toolPanels.all { it.isCompleted }) { - removeSelfFromParent() - } - } - - private fun handleIndividualRejection(toolCall: ToolCall, panel: CompactToolPanel) { - onReject(toolCall, panel) - removeToolPanel(toolCall) - updateHeader() - - if (toolPanels.all { it.isCompleted }) { - removeSelfFromParent() - } - } - - private fun handleBatchApproval() { - onBatchApprove(toolPanels, false) - updateHeader() - - if (toolPanels.all { it.isCompleted }) { - removeSelfFromParent() - } - } - - private fun handleBatchRejection() { - cleanupGroupPanel() - onBatchReject(toolPanels) - updateHeader() - removeSelfFromParent() - } - - private fun updateHeader() { - if (componentCount == 0) return - - (getComponent(0) as? JPanel)?.let { headerPanel -> - val remainingCount = toolPanels.count { !it.isCompleted } - - headerPanel.isVisible = remainingCount > 1 - - if (remainingCount > 1) { - (headerPanel.getComponent(0) as? JPanel)?.let { contentPanel -> - (contentPanel.getComponent(0) as? JBLabel)?.text = "$remainingCount $TOOL_REQUESTS_TEXT" - } - } - } - } -} \ No newline at end of file diff --git a/src/main/kotlin/ee/carlrobert/codegpt/ui/components/TokenUsageCounterPanel.kt b/src/main/kotlin/ee/carlrobert/codegpt/ui/components/TokenUsageCounterPanel.kt index c39f701d..b480ed5a 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/ui/components/TokenUsageCounterPanel.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/ui/components/TokenUsageCounterPanel.kt @@ -13,6 +13,7 @@ import com.intellij.util.ui.JBFont import com.intellij.util.ui.JBUI import ee.carlrobert.codegpt.agent.AgentService import ee.carlrobert.codegpt.settings.configuration.ConfigurationSettings +import ee.carlrobert.codegpt.settings.models.ModelSettings import ee.carlrobert.codegpt.toolwindow.agent.TokenUsageEvent import ee.carlrobert.codegpt.toolwindow.agent.TokenUsageListener import ee.carlrobert.codegpt.util.coroutines.DisposableCoroutineScope @@ -37,7 +38,8 @@ class TokenUsageCounterPanel( isGroupingUsed = true } private var messageBusConnection: MessageBusConnection? = null - private var currentSessionId: String = sessionId + private var lastEstimatedPromptTokens: Long = 0L + private var calibrationOffset: Long = 0L init { Disposer.register(ApplicationManager.getApplication(), scope) @@ -45,6 +47,7 @@ class TokenUsageCounterPanel( isOpaque = false font = JBFont.small() border = JBUI.Borders.empty(5, 7) + text = "100% context left" setupMessageBusConnection() } @@ -56,20 +59,29 @@ class TokenUsageCounterPanel( object : TokenUsageListener { override fun onTokenUsageChanged(event: TokenUsageEvent) { if (event.sessionId == sessionId) { - currentSessionId = event.sessionId - val model = getAgentModelForSession(event.sessionId) ?: return - updateDisplay(event, model) - updateTooltipText(event, model) + val model = getAgentModelForSession(event.sessionId) + ?: getSelectedAgentModel() + ?: return + calibrationOffset = event.totalTokens - lastEstimatedPromptTokens + updateDisplay(event.totalTokens, model) } } }) } } - private fun updateDisplay(event: TokenUsageEvent, model: LLModel) { + fun updateFromTotalTokens(totalTokens: Long) { + lastEstimatedPromptTokens = totalTokens + val model = getSelectedAgentModel() ?: return + val calibratedTotal = (totalTokens + calibrationOffset).coerceAtLeast(0L) + updateDisplay(calibratedTotal, model) + } + + private fun updateDisplay(totalTokens: Long, model: LLModel) { scope.launch { withEdt { - updateColorAndText(event.totalTokens, model) + updateColorAndText(totalTokens, model) + updateTooltipText(totalTokens, model) revalidate() repaint() } @@ -93,12 +105,12 @@ class TokenUsageCounterPanel( text = "${percentageLeft.toInt()}% context left" } - fun updateTooltipText(event: TokenUsageEvent, model: LLModel) { + private fun updateTooltipText(totalTokens: Long, model: LLModel) { val budget = computeBudget(model) toolTipText = buildString { append("") append("Usage Details
") - append("Input size: ${numberFormat.format(event.totalTokens)} tokens
") + append("Input size: ${numberFormat.format(totalTokens)} tokens
") append("Max output size: ${numberFormat.format(budget.reservedOutput)} tokens
") append("Max context size: ${numberFormat.format(budget.contextLength)} tokens
") append("") @@ -127,13 +139,17 @@ class TokenUsageCounterPanel( ) private fun computeBudget(model: LLModel): Budget { - val contextLength = model.contextLength + val contextLength = model.contextLength ?: 128_000L val configured = - (model.maxOutputTokens ?: ConfigurationSettings.getState().maxTokens).toLong() + model.maxOutputTokens ?: ConfigurationSettings.getState().maxTokens.toLong() val reserved = configured .coerceAtLeast(0L) .coerceAtMost(contextLength) val inputBudget = (contextLength - reserved).coerceAtLeast(1L) return Budget(contextLength, reserved, inputBudget) } + + private fun getSelectedAgentModel(): LLModel? { + return runCatching { ModelSettings.getInstance().getAgentModel() }.getOrNull() + } } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/UserInputPanel.kt b/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/UserInputPanel.kt index 720b1013..437ee1b6 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/UserInputPanel.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/UserInputPanel.kt @@ -25,15 +25,15 @@ import com.intellij.ui.dsl.builder.panel import com.intellij.util.IconUtil import com.intellij.util.ui.AsyncProcessIcon import com.intellij.util.ui.JBUI +import ai.koog.prompt.llm.LLMCapability import com.intellij.util.ui.components.BorderLayoutPanel import ee.carlrobert.codegpt.CodeGPTBundle import ee.carlrobert.codegpt.Icons import ee.carlrobert.codegpt.agent.PromptEnhancer import ee.carlrobert.codegpt.settings.ProxyAISettingsService import ee.carlrobert.codegpt.settings.configuration.ChatMode -import ee.carlrobert.codegpt.settings.models.ModelRegistry 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.toolwindow.chat.ui.textarea.ModelComboBoxAction import ee.carlrobert.codegpt.toolwindow.chat.ui.textarea.TotalTokensPanel @@ -41,6 +41,7 @@ import ee.carlrobert.codegpt.ui.IconActionButton import ee.carlrobert.codegpt.ui.OverlayUtil import ee.carlrobert.codegpt.ui.components.BadgeChip import ee.carlrobert.codegpt.ui.components.InlineEditChips +import ee.carlrobert.codegpt.ui.components.TokenUsageCounterPanel import ee.carlrobert.codegpt.ui.dnd.FileDragAndDrop import ee.carlrobert.codegpt.ui.textarea.header.UserInputHeaderPanel import ee.carlrobert.codegpt.ui.textarea.header.tag.* @@ -144,6 +145,7 @@ class UserInputPanel @JvmOverloads constructor( ) private var footerPanelRef: JPanel? = null + private val tokenUsageCounterPanel = agentTokenCounterPanel as? TokenUsageCounterPanel private val applyChip = onApply?.let { BadgeChip(CodeGPTBundle.get("shared.apply"), InlineEditChips.GREEN, it) } @@ -264,6 +266,7 @@ class UserInputPanel @JvmOverloads constructor( init { setupDisposables(parentDisposable) setupLayout(featureType) + bindTokenCounterToInputTokens() addSelectedEditorContent() if (featureType == FeatureType.INLINE_EDIT) { setupTextChangeListener() @@ -320,13 +323,20 @@ class UserInputPanel @JvmOverloads constructor( } } + private fun bindTokenCounterToInputTokens() { + val tokenCounter = tokenUsageCounterPanel ?: return + totalTokensPanel.addTotalTokensListener { total -> + tokenCounter.updateFromTotalTokens(total.toLong()) + } + } + private fun addSelectedEditorContent() { runInEdt { EditorUtil.getSelectedEditor(project)?.let { editor -> if (EditorUtil.hasSelection(editor)) { - addTag( - EditorSelectionTagDetails(editor.virtualFile, editor.selectionModel) - ) + editor.virtualFile?.let { virtualFile -> + addTag(EditorSelectionTagDetails(virtualFile, editor.selectionModel)) + } } } } @@ -581,11 +591,10 @@ class UserInputPanel @JvmOverloads constructor( } private fun createFooterPanel(featureType: FeatureType): JPanel { - val currentService = - ModelSelectionService.getInstance().getServiceForFeature(featureType) - val availableProviders = ModelRegistry.getInstance().getProvidersForFeature(featureType) + val modelSettings = ModelSettings.getInstance() + val currentService = modelSettings.getServiceForFeature(featureType) + val availableProviders = modelSettings.getAvailableProviders(featureType) val modelComboBox = ModelComboBoxAction( - project, { imageActionSupported.set(isImageActionSupported()) }, currentService, availableProviders, @@ -653,9 +662,9 @@ class UserInputPanel @JvmOverloads constructor( } private fun isImageActionSupported(): Boolean { - val currentModel = ModelSelectionService.getInstance().getModelForFeature(featureType) + val currentModel = ModelSettings.getInstance().getModelForFeature(featureType) val currentService = - ModelSelectionService.getInstance().getServiceForFeature(featureType) + ModelSettings.getInstance().getServiceForFeature(featureType) return when (currentService) { ServiceType.CUSTOM_OPENAI, @@ -670,14 +679,8 @@ class UserInputPanel @JvmOverloads constructor( } private fun isCodeGPTModelSupported(modelCode: String): Boolean { - return modelCode in setOf( - ModelRegistry.GPT_4_1, - ModelRegistry.GPT_4_1_MINI, - ModelRegistry.GEMINI_PRO_2_5, - ModelRegistry.GEMINI_FLASH_2_5, - ModelRegistry.CLAUDE_4_5_SONNET, - ModelRegistry.CLAUDE_4_5_SONNET_THINKING - ) + val model = ModelSettings.getInstance().findModel(ServiceType.PROXYAI, modelCode) + return model?.llmModel?.capabilities?.any { it is LLMCapability.Vision.Image } == true } private fun updatePreferredSizeFromChildren() { diff --git a/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/header/UserInputHeaderPanel.kt b/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/header/UserInputHeaderPanel.kt index ddd5ee50..77f9050a 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/header/UserInputHeaderPanel.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/header/UserInputHeaderPanel.kt @@ -439,7 +439,7 @@ class UserInputHeaderPanel( private inner class EditorReleasedListener : EditorNotifier.Released { override fun editorReleased(editor: Editor) { if (editor.editorKind == EditorKind.MAIN_EDITOR && !editor.isDisposed && editor.virtualFile != null) { - tagManager.remove(EditorTagDetails(editor.virtualFile)) + editor.virtualFile?.let { tagManager.remove(EditorTagDetails(it)) } } } } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/lookup/action/WebActionItem.kt b/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/lookup/action/WebActionItem.kt index 6f49e059..532da685 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/lookup/action/WebActionItem.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/lookup/action/WebActionItem.kt @@ -4,7 +4,7 @@ import com.intellij.icons.AllIcons import com.intellij.openapi.project.Project import ee.carlrobert.codegpt.CodeGPTBundle 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.ui.textarea.UserInputPanel import ee.carlrobert.codegpt.ui.textarea.header.tag.TagManager @@ -19,7 +19,7 @@ class WebActionItem(private val tagManager: TagManager) : AbstractLookupActionIt get() = enabled() fun enabled(): Boolean { - if (ModelSelectionService.getInstance().getServiceForFeature(FeatureType.CHAT) != ServiceType.PROXYAI) { + if (ModelSettings.getInstance().getServiceForFeature(FeatureType.CHAT) != ServiceType.PROXYAI) { return false } return tagManager.getTags().none { it is WebTagDetails } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/lookup/group/MCPGroupItem.kt b/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/lookup/group/MCPGroupItem.kt index d817e4f7..36782f00 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/lookup/group/MCPGroupItem.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/lookup/group/MCPGroupItem.kt @@ -7,7 +7,7 @@ import ee.carlrobert.codegpt.CodeGPTBundle import ee.carlrobert.codegpt.Icons import ee.carlrobert.codegpt.settings.mcp.McpSettings 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.CUSTOM_OPENAI import ee.carlrobert.codegpt.settings.service.ServiceType.OPENAI import ee.carlrobert.codegpt.ui.textarea.header.tag.McpTagDetails @@ -30,7 +30,7 @@ class MCPGroupItem( if (featureType == FeatureType.AGENT) { return true } - val serviceType = service().getServiceForFeature(featureType) + val serviceType = service().getServiceForFeature(featureType) return serviceType == OPENAI || serviceType == CUSTOM_OPENAI } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/util/EditorDiffUtil.kt b/src/main/kotlin/ee/carlrobert/codegpt/util/EditorDiffUtil.kt index a54c5106..57579034 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/util/EditorDiffUtil.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/util/EditorDiffUtil.kt @@ -29,7 +29,7 @@ object EditorDiffUtil { createTempDiffContent(mainEditor, toolwindowEditor, highlightedText) ) DiffManager.getInstance() - .showDiff(project, createDiffRequest(project, tempFile, mainEditor.virtualFile)) + .showDiff(project, createDiffRequest(project, tempFile, mainEditorFile)) } private fun createTempDiffContent( diff --git a/src/main/kotlin/ee/carlrobert/codegpt/util/EditorUtil.kt b/src/main/kotlin/ee/carlrobert/codegpt/util/EditorUtil.kt index 82738366..05f288be 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/util/EditorUtil.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/util/EditorUtil.kt @@ -24,7 +24,6 @@ import com.intellij.psi.PsiDocumentManager import com.intellij.psi.codeStyle.CodeStyleManager import com.intellij.testFramework.LightVirtualFile import ee.carlrobert.codegpt.settings.configuration.ConfigurationSettings -import ee.carlrobert.llm.client.codegpt.request.prediction.FileDetails import java.time.LocalDateTime import java.time.format.DateTimeFormatter diff --git a/src/main/kotlin/ee/carlrobert/codegpt/util/FileDetails.kt b/src/main/kotlin/ee/carlrobert/codegpt/util/FileDetails.kt new file mode 100644 index 00000000..45e9e14b --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/util/FileDetails.kt @@ -0,0 +1,7 @@ +package ee.carlrobert.codegpt.util + +class FileDetails { + var name: String? = null + var content: String? = null + var modificationStamp: Long? = null +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/util/JsonMapper.kt b/src/main/kotlin/ee/carlrobert/codegpt/util/JsonMapper.kt new file mode 100644 index 00000000..a40ca198 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/util/JsonMapper.kt @@ -0,0 +1,8 @@ +package ee.carlrobert.codegpt.util + +import com.fasterxml.jackson.databind.ObjectMapper +import com.fasterxml.jackson.module.kotlin.registerKotlinModule + +object JsonMapper { + val mapper: ObjectMapper = ObjectMapper().registerKotlinModule() +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/util/ReasoningFrameTextAdapter.kt b/src/main/kotlin/ee/carlrobert/codegpt/util/ReasoningFrameTextAdapter.kt new file mode 100644 index 00000000..12227dd7 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/util/ReasoningFrameTextAdapter.kt @@ -0,0 +1,119 @@ +package ee.carlrobert.codegpt.util + +import ai.koog.prompt.streaming.StreamFrame + +internal class ReasoningFrameTextAdapter { + + private var thinkingOpen = false + private val textDeltaIndexes = mutableSetOf() + private val reasoningDeltaIndexes = mutableSetOf() + + fun consume(frame: StreamFrame): List { + return when (frame) { + is StreamFrame.TextDelta -> emitTextDelta(frame) + is StreamFrame.TextComplete -> emitTextComplete(frame) + is StreamFrame.ReasoningDelta -> emitReasoning(frame.text, frame.summary, frame.index) + is StreamFrame.ReasoningComplete -> emitReasoningComplete(frame) + is StreamFrame.End -> { + textDeltaIndexes.clear() + reasoningDeltaIndexes.clear() + closeThinking() + } + else -> emptyList() + } + } + + private fun emitTextDelta(frame: StreamFrame.TextDelta): List { + if (frame.text.isEmpty()) return emptyList() + textDeltaIndexes.add(frame.index) + return emitText(frame.text) + } + + private fun emitTextComplete(frame: StreamFrame.TextComplete): List { + val hadDeltaForThisText = textDeltaIndexes.remove(frame.index) + return if (hadDeltaForThisText) emptyList() else emitText(frame.text) + } + + private fun emitText(text: String): List { + if (text.isEmpty()) return emptyList() + + val parts = mutableListOf() + if (thinkingOpen) { + parts.add(THINK_END) + thinkingOpen = false + } + parts.add(text) + return parts + } + + private fun emitReasoning(text: String?, summary: String?, index: Int?): List { + val parts = mutableListOf() + var emitted = false + + if (!text.isNullOrBlank()) { + if (!thinkingOpen) { + parts.add(THINK_START) + thinkingOpen = true + } + parts.add(text) + emitted = true + } + + if (!summary.isNullOrBlank()) { + if (!thinkingOpen) { + parts.add(THINK_START) + thinkingOpen = true + } + parts.add(summary) + emitted = true + } + + if (emitted) { + reasoningDeltaIndexes.add(index) + } + + return parts + } + + private fun emitReasoningComplete(frame: StreamFrame.ReasoningComplete): List { + val hadDeltaForThisReasoning = reasoningDeltaIndexes.remove(frame.index) + if (hadDeltaForThisReasoning) { + return closeThinking() + } + + val reasoningText = frame.text.joinToString("\n").trim() + val summaryText = frame.summary?.joinToString("\n")?.trim().orEmpty() + + val parts = mutableListOf() + if (reasoningText.isNotEmpty() || summaryText.isNotEmpty()) { + if (!thinkingOpen) { + parts.add(THINK_START) + thinkingOpen = true + } + + if (reasoningText.isNotEmpty()) { + parts.add(reasoningText) + } + if (summaryText.isNotEmpty()) { + parts.add(summaryText) + } + } + + parts.addAll(closeThinking()) + return parts + } + + private fun closeThinking(): List { + return if (thinkingOpen) { + thinkingOpen = false + listOf(THINK_END) + } else { + emptyList() + } + } + + private companion object { + private const val THINK_START = "" + private const val THINK_END = "" + } +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/util/StringUtil.kt b/src/main/kotlin/ee/carlrobert/codegpt/util/StringUtil.kt index 45d4284c..86c3b277 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/util/StringUtil.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/util/StringUtil.kt @@ -1,7 +1,5 @@ package ee.carlrobert.codegpt.util -import ai.grazie.nlp.utils.takeWhitespaces - object StringUtil { private val thinkBlockRegex = Regex("(?s).*?\\s*") @@ -15,14 +13,14 @@ object StringUtil { completionLine: String, editorLine: String ): String { - val editorWhitespaces = editorLine.takeWhitespaces() + val editorWhitespaces = editorLine.leadingWhitespaces() if (completionLine.isNotEmpty() && editorWhitespaces.isNotEmpty()) { if (completionLine.startsWith(editorWhitespaces)) { return completionLine.substring(editorWhitespaces.length) } if (editorLine.isBlank()) { - val completionWhitespaces = completionLine.takeWhitespaces() + val completionWhitespaces = completionLine.leadingWhitespaces() return completionLine.substring(completionWhitespaces.length) } } @@ -40,4 +38,6 @@ object StringUtil { val intersection = bigrams1.intersect(bigrams2).size return (2.0 * intersection) / (bigrams1.size + bigrams2.size) } + + private fun String.leadingWhitespaces(): String = takeWhile { it.isWhitespace() } } diff --git a/src/main/resources/META-INF/plugin-python.xml b/src/main/resources/META-INF/plugin-python.xml deleted file mode 100644 index 36ff30f4..00000000 --- a/src/main/resources/META-INF/plugin-python.xml +++ /dev/null @@ -1,6 +0,0 @@ - - - - - diff --git a/src/main/resources/META-INF/plugin.xml b/src/main/resources/META-INF/plugin.xml index d51b2d23..12d53b68 100644 --- a/src/main/resources/META-INF/plugin.xml +++ b/src/main/resources/META-INF/plugin.xml @@ -6,7 +6,6 @@ com.intellij.modules.lang org.jetbrains.kotlin com.intellij.modules.java - com.intellij.modules.python JavaScript com.intellij.modules.cpp diff --git a/src/main/resources/icons/morph.svg b/src/main/resources/icons/morph.svg deleted file mode 100644 index 5e5265fd..00000000 --- a/src/main/resources/icons/morph.svg +++ /dev/null @@ -1,3 +0,0 @@ - - - \ No newline at end of file diff --git a/src/main/resources/icons/relace.svg b/src/main/resources/icons/relace.svg deleted file mode 100644 index 6945d7d0..00000000 --- a/src/main/resources/icons/relace.svg +++ /dev/null @@ -1,4 +0,0 @@ - - - - \ No newline at end of file diff --git a/src/main/resources/icons/relace_dark.svg b/src/main/resources/icons/relace_dark.svg deleted file mode 100644 index 5fb27a5d..00000000 --- a/src/main/resources/icons/relace_dark.svg +++ /dev/null @@ -1,4 +0,0 @@ - - - - \ No newline at end of file diff --git a/src/main/resources/messages/codegpt.properties b/src/main/resources/messages/codegpt.properties index 2a89c63c..094aa459 100644 --- a/src/main/resources/messages/codegpt.properties +++ b/src/main/resources/messages/codegpt.properties @@ -368,6 +368,9 @@ tagPopupMenuItem.closeTagsToLeft=Close Tags to the Left tagPopupMenuItem.closeTagsToRight=Close Tags to the Right toolwindow.chat.loading=Thinking... toolwindow.chat.compressingHistory=Compressing history... +toolwindow.agent.retrying=Retrying ({0}/{1}) +toolwindow.agent.error.generic=Something went wrong. Please try again. +toolwindow.agent.error.executionFailed=Agent execution failed unexpectedly. Please try again. headerPanel.error.searchBlockNotMapped.title=Failed to Locate Search Block llama.build.cmake.setup=Setting up CMake... llama.build.cmake.build=Building project... diff --git a/src/main/resources/messages/codegpt_zh.properties b/src/main/resources/messages/codegpt_zh.properties index 057dab95..c5a50bb5 100644 --- a/src/main/resources/messages/codegpt_zh.properties +++ b/src/main/resources/messages/codegpt_zh.properties @@ -355,6 +355,9 @@ tagPopupMenuItem.closeTagsToLeft=\u5173\u95ED\u5DE6\u4FA7\u6807\u7B7E tagPopupMenuItem.closeTagsToRight=\u5173\u95ED\u53F3\u4FA7\u6807\u7B7E toolwindow.chat.loading=\u6B63\u5728\u751F\u6210\u54CD\u5E94... toolwindow.chat.compressingHistory=\u6B63\u5728\u538B\u7F29\u5386\u53F2\u8BB0\u5F55... +toolwindow.agent.retrying=\u6B63\u5728\u91CD\u8BD5\uFF08{0}/{1}\uFF09 +toolwindow.agent.error.generic=\u51FA\u73B0\u9519\u8BEF\uFF0C\u8BF7\u91CD\u8BD5\u3002 +toolwindow.agent.error.executionFailed=Agent \u6267\u884C\u5931\u8D25\uFF0C\u8BF7\u91CD\u8BD5\u3002 headerPanel.error.searchBlockNotMapped.title=\u65E0\u6CD5\u5B9A\u4F4D\u641C\u7D22\u5757 llama.build.cmake.setup=\u6B63\u5728\u8BBE\u7F6ECMake... llama.build.cmake.build=\u6B63\u5728\u6784\u5EFA\u9879\u76EE... diff --git a/src/test/kotlin/ee/carlrobert/codegpt/agent/AgentProviderIntegrationTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/agent/AgentProviderIntegrationTest.kt new file mode 100644 index 00000000..3c1b6063 --- /dev/null +++ b/src/test/kotlin/ee/carlrobert/codegpt/agent/AgentProviderIntegrationTest.kt @@ -0,0 +1,649 @@ +package ee.carlrobert.codegpt.agent + +import ai.koog.agents.snapshot.providers.file.JVMFilePersistenceStorageProvider +import com.intellij.openapi.components.service +import ee.carlrobert.codegpt.credentials.CredentialsStore.CredentialKey.* +import ee.carlrobert.codegpt.credentials.CredentialsStore.setCredential +import ee.carlrobert.codegpt.settings.models.ModelCatalog +import ee.carlrobert.codegpt.settings.models.ModelSettings +import ee.carlrobert.codegpt.settings.service.FeatureType +import ee.carlrobert.codegpt.settings.service.ServiceType +import ee.carlrobert.codegpt.settings.service.custom.CustomServiceSettingsState +import ee.carlrobert.codegpt.settings.service.custom.CustomServicesSettings +import kotlinx.coroutines.runBlocking +import org.assertj.core.api.Assertions.assertThat +import testsupport.IntegrationTest +import testsupport.http.RequestEntity +import testsupport.http.ResponseEntity +import testsupport.http.exchange.BasicHttpExchange +import testsupport.http.exchange.StreamHttpExchange +import testsupport.json.JSONUtil.e +import testsupport.json.JSONUtil.jsonArray +import testsupport.json.JSONUtil.jsonMap +import testsupport.json.JSONUtil.jsonMapResponse +import java.util.* +import java.util.concurrent.ConcurrentHashMap +import kotlin.io.path.createTempDirectory +import kotlin.io.path.createTempFile +import kotlin.io.path.writeText + +class AgentProviderIntegrationTest : IntegrationTest() { + + fun testOpenAIAgentUsesMockHarnessThroughRealExecutor() { + setCredential(OpenaiApiKey, "TEST_API_KEY") + service().setModel(FeatureType.AGENT, "gpt-4.1-mini", ServiceType.OPENAI) + expectOpenAI(StreamHttpExchange { request -> + assertThat(request.method).isEqualTo("POST") + assertThat(extractPromptText(request)).contains("Say hello from OpenAI") + when (request.uri.path) { + "/v1/chat/completions" -> chatCompletionChunks("gpt-4.1-mini", "Hello from OpenAI") + "/v1/responses" -> openAiResponsesChunks() + else -> error("Unexpected OpenAI path: ${request.uri.path}") + } + }) + + val result = runAgent(ServiceType.OPENAI, "Say hello from OpenAI") + + assertThat(result.output).isEqualTo("Hello from OpenAI") + assertThat(result.events.text.toString()).isEqualTo("Hello from OpenAI") + } + + fun testProxyAIAgentUsesMockHarnessThroughRealExecutor() { + setCredential(CodeGptApiKey, "TEST_API_KEY") + service().setModel( + FeatureType.AGENT, + ModelCatalog.PROXYAI_AUTO, + ServiceType.PROXYAI + ) + val fixture = createReadFixture("ProxyAI fixture") + expectCodeGPT { request -> + assertThat(request.uri.path).isEqualTo("/v2/chat/completions") + assertThat(request.method).isEqualTo("POST") + assertThat(extractPromptText(request)).contains("Read the fixture and repeat its contents") + proxyAiChatCompletionToolCallChunks("""{"file_path":"${fixture.path}"}""") + } + expectCodeGPT { request -> + assertThat(request.uri.path).isEqualTo("/v2/chat/completions") + assertThat(request.method).isEqualTo("POST") + assertThat(request.body.toString()).contains(fixture.renderedResult) + proxyAiChatCompletionChunks("ProxyAI read: ${fixture.contents}") + } + + val result = runAgent(ServiceType.PROXYAI, "Read the fixture and repeat its contents") + + assertThat(result.output).isEqualTo("ProxyAI read: ${fixture.contents}") + assertThat(result.events.text.toString()).isEqualTo("ProxyAI read: ${fixture.contents}") + } + + fun testAnthropicAgentCompletesReadToolLoopThroughRealExecutor() { + setCredential(AnthropicApiKey, "TEST_API_KEY") + service().setModel( + FeatureType.AGENT, + "claude-haiku-4-5", + ServiceType.ANTHROPIC + ) + val fixture = createReadFixture("Anthropic fixture") + expectAnthropic { request -> + assertThat(request.uri.path).isEqualTo("/v1/messages") + assertThat(request.method).isEqualTo("POST") + assertThat(extractPromptText(request)).contains("Read the fixture and repeat its contents") + anthropicToolUseChunks("""{"file_path":"${fixture.path}"}""") + } + expectAnthropic { request -> + assertThat(request.uri.path).isEqualTo("/v1/messages") + assertThat(request.method).isEqualTo("POST") + assertThat(request.body.toString()).contains(fixture.renderedResult) + anthropicTextChunks("Anthropic read: ${fixture.contents}") + } + + val result = runAgent(ServiceType.ANTHROPIC, "Read the fixture and repeat its contents") + + assertThat(result.output).isEqualTo("Anthropic read: ${fixture.contents}") + assertThat(result.events.text.toString()).isEqualTo("Anthropic read: ${fixture.contents}") + } + + fun testGoogleAgentUsesMockHarnessThroughRealExecutor() { + setCredential(GoogleApiKey, "TEST_API_KEY") + service().setModel( + FeatureType.AGENT, + "gemini-3-flash-preview", + ServiceType.GOOGLE + ) + expectGoogle { request -> + assertThat(request.uri.path).isEqualTo("/v1beta/models/gemini-3-flash-preview:generateContent") + assertThat(request.method).isEqualTo("POST") + assertThat(extractGooglePromptText(request)).contains("Say hello from Google") + ResponseEntity( + googleTextResponse() + ) + } + + val result = runAgent(ServiceType.GOOGLE, "Say hello from Google") + + assertThat(result.output).isEqualTo("Hello from Google") + assertThat(result.events.text.toString()).isEqualTo("Hello from Google") + } + + fun testMistralAgentUsesMockHarnessThroughRealExecutor() { + setCredential(MistralApiKey, "TEST_API_KEY") + service().setModel(FeatureType.AGENT, "devstral-2512", ServiceType.MISTRAL) + expectMistral(StreamHttpExchange { request: RequestEntity -> + assertThat(request.uri.path).isEqualTo("/v1/chat/completions") + assertThat(request.method).isEqualTo("POST") + assertThat(extractPromptText(request)).contains("Say hello from Mistral") + chatCompletionChunks("devstral-2512", "Hello from Mistral") + }) + + val result = runAgent(ServiceType.MISTRAL, "Say hello from Mistral") + + assertThat(result.output).isEqualTo("Hello from Mistral") + assertThat(result.events.text.toString()).isEqualTo("Hello from Mistral") + } + + fun testInceptionAgentUsesMockHarnessThroughRealExecutor() { + setCredential(InceptionApiKey, "TEST_API_KEY") + service().setModel(FeatureType.AGENT, "mercury", ServiceType.INCEPTION) + expectInception(StreamHttpExchange { request: RequestEntity -> + assertThat(request.uri.path).isEqualTo("/v1/chat/completions") + assertThat(request.method).isEqualTo("POST") + assertThat(extractPromptText(request)).contains("Say hello from Inception") + chatCompletionChunks("mercury", "Hello from Inception") + }) + + val result = runAgent(ServiceType.INCEPTION, "Say hello from Inception") + + assertThat(result.output).isEqualTo("Hello from Inception") + assertThat(result.events.text.toString()).isEqualTo("Hello from Inception") + } + + fun testCustomOpenAIAgentUsesMockHarnessThroughRealExecutor() { + val customService = configureCustomOpenAIService() + expectCustomOpenAI(BasicHttpExchange { request -> + assertThat(request.uri.path).isEqualTo("/v1/chat/completions") + assertThat(request.method).isEqualTo("POST") + assertThat(extractPromptText(request)).contains("Say hello from Custom OpenAI") + ResponseEntity( + customOpenAiResponse() + ) + }) + + val result = runAgent(ServiceType.CUSTOM_OPENAI, "Say hello from Custom OpenAI") + + assertThat(customService.id).isNotBlank() + assertThat(result.output).isEqualTo("Hello from Custom OpenAI") + assertThat(result.events.text.toString()).isEqualTo("Hello from Custom OpenAI") + } + + private fun runAgent( + provider: ServiceType, + userMessage: String + ): AgentRunResult { + val events = RecordingAgentEvents() + val checkpointStorage = + JVMFilePersistenceStorageProvider(createTempDirectory("agent-provider-test")) + val service = ProxyAIAgent.createService( + project = project, + checkpointStorage = checkpointStorage, + provider = provider, + events = events, + sessionId = "session-${UUID.randomUUID()}", + pendingMessages = ConcurrentHashMap() + ) + + return try { + val output = runBlocking { + val agent = service.createAgent() + agent.run(MessageWithContext(userMessage)) + } + AgentRunResult(output, events) + } finally { + runBlocking { service.closeAll() } + } + } + + private fun chatCompletionChunks(model: String, text: String): List { + val chunks = text.chunked(4) + return chunks.mapIndexed { index, chunk -> + jsonMapResponse( + e("id", "chatcmpl-test"), + e("object", "chat.completion.chunk"), + e("created", 1), + e("model", model), + e( + "choices", + jsonArray( + jsonMap( + e("index", 0), + e( + "delta", + jsonMap( + *buildList { + if (index == 0) add(e("role", "assistant")) + add(e("content", chunk)) + }.toTypedArray() + ) + ), + e("finishReason", null) + ) + ) + ) + ) + } + jsonMapResponse( + e("id", "chatcmpl-test"), + e("object", "chat.completion.chunk"), + e("created", 1), + e("model", model), + e( + "choices", + jsonArray( + jsonMap( + e("index", 0), + e("delta", jsonMap()), + e("finishReason", "stop") + ) + ) + ) + ) + } + + private fun openAiResponsesChunks(): List { + val model = "gpt-4.1-mini" + val text = "Hello from OpenAI" + val chunks = text.chunked(4) + return chunks.mapIndexed { index, chunk -> + jsonMapResponse( + e("type", "response.output_text.delta"), + e("item_id", "msg_1"), + e("output_index", 0), + e("content_index", 0), + e("delta", chunk), + e("sequence_number", index + 1) + ) + } + jsonMapResponse( + e("type", "response.output_item.done"), + e("item", text), + e("output_index", 0), + e("sequence_number", chunks.size + 1) + ) + jsonMapResponse( + e("type", "response.completed"), + e( + "response", + jsonMap( + e("id", "resp-openai-test"), + e("object", "response"), + e("created_at", 1), + e("model", model), + e( + "output", + jsonArray( + jsonMap( + e("type", "message"), + e("id", "msg_1"), + e("role", "assistant"), + e("status", "completed"), + e( + "content", + jsonArray( + jsonMap( + e("type", "output_text"), + e("text", text), + e("annotations", jsonArray()) + ) + ) + ) + ) + ) + ), + e("parallel_tool_calls", true), + e("status", "completed"), + e("text", jsonMap()), + e( + "usage", + jsonMap( + e("input_tokens", 1), + e("input_tokens_details", jsonMap("cached_tokens", 0)), + e("output_tokens", 1), + e("output_tokens_details", jsonMap("reasoning_tokens", 0)), + e("total_tokens", 2) + ) + ) + ) + ), + e("sequence_number", chunks.size + 2) + ) + } + + private fun anthropicTextChunks(text: String): List { + return listOf( + jsonMapResponse( + e("type", "content_block_start"), + e("index", 0), + e( + "content_block", + jsonMap( + e("type", "text"), + e("text", text) + ) + ) + ), + jsonMapResponse( + e("type", "message_delta"), + e( + "delta", + jsonMap( + e("stop_reason", "end_turn") + ) + ) + ) + ) + } + + private fun anthropicToolUseChunks(arguments: String): List { + val toolId = "toolu_anthropic_read" + val toolName = "Read" + return listOf( + jsonMapResponse( + e("type", "content_block_start"), + e("index", 0), + e( + "content_block", + jsonMap( + e("type", "tool_use"), + e("id", toolId), + e("name", toolName), + e("input", jsonMap()) + ) + ) + ), + jsonMapResponse( + e("type", "content_block_delta"), + e("index", 0), + e( + "delta", + jsonMap( + e("type", "input_json_delta"), + e("partial_json", arguments) + ) + ) + ), + jsonMapResponse( + e("type", "content_block_stop"), + e("index", 0), + e( + "delta", + jsonMap( + e("type", "input_json_delta") + ) + ) + ), + jsonMapResponse( + e("type", "message_delta"), + e( + "delta", + jsonMap( + e("stop_reason", "tool_use") + ) + ) + ) + ) + } + + private fun proxyAiChatCompletionChunks(text: String): List { + val model = "auto" + val chunks = text.chunked(4) + return chunks.mapIndexed { index, chunk -> + jsonMapResponse( + e("id", "proxyai-test"), + e("object", "chat.completion.chunk"), + e("created", 1), + e("model", model), + e( + "choices", + jsonArray( + jsonMap( + e( + "delta", + jsonMap( + *buildList { + if (index == 0) add(e("role", "assistant")) + add(e("content", chunk)) + }.toTypedArray() + ) + ), + e("finishReason", null), + e("nativeFinishReason", null) + ) + ) + ) + ) + } + jsonMapResponse( + e("id", "proxyai-test"), + e("object", "chat.completion.chunk"), + e("created", 1), + e("model", model), + e( + "choices", + jsonArray( + jsonMap( + e("delta", jsonMap()), + e("finishReason", "stop"), + e("nativeFinishReason", "stop") + ) + ) + ) + ) + } + + private fun proxyAiChatCompletionToolCallChunks(arguments: String): List { + val model = "auto" + val toolId = "call_proxyai_read" + val toolName = "Read" + return listOf( + jsonMapResponse( + e("id", "proxyai-tool-call"), + e("object", "chat.completion.chunk"), + e("created", 1), + e("model", model), + e( + "choices", + jsonArray( + jsonMap( + e( + "delta", + jsonMap( + e("role", "assistant"), + e( + "tool_calls", + jsonArray( + jsonMap( + e("index", 0), + e("id", toolId), + e( + "function", + jsonMap( + e("name", toolName), + e("arguments", arguments) + ) + ) + ) + ) + ) + ) + ), + e("finishReason", null), + e("nativeFinishReason", null) + ) + ) + ) + ), + jsonMapResponse( + e("id", "proxyai-tool-call"), + e("object", "chat.completion.chunk"), + e("created", 1), + e("model", model), + e( + "choices", + jsonArray( + jsonMap( + e("delta", jsonMap()), + e("finishReason", "tool_calls"), + e("nativeFinishReason", "tool_calls") + ) + ) + ) + ) + ) + } + + private fun googleTextResponse(): String { + val text = "Hello from Google" + return jsonMapResponse( + e( + "candidates", + jsonArray( + jsonMap( + e( + "content", + jsonMap( + e("role", "model"), + e( + "parts", + jsonArray( + jsonMap(e("text", text)) + ) + ) + ) + ), + e("finishReason", "STOP") + ) + ) + ), + e( + "usageMetadata", + jsonMap( + e("promptTokenCount", 1), + e("candidatesTokenCount", 1), + e("totalTokenCount", 2) + ) + ) + ) + } + + private fun customOpenAiResponse(): String { + val model = "custom-agent-model" + val text = "Hello from Custom OpenAI" + return jsonMapResponse( + e( + "choices", + jsonArray( + jsonMap( + e("finishReason", "stop"), + e( + "message", + jsonMap( + e("role", "assistant"), + e("content", text) + ) + ) + ) + ) + ), + e("created", 1), + e("id", "custom-openai-test"), + e("model", model), + e("object", "chat.completion"), + e( + "usage", + jsonMap( + e("prompt_tokens", 1), + e("completion_tokens", 1), + e("total_tokens", 2) + ) + ) + ) + } + + private fun createReadFixture(contents: String): ReadFixture { + val path = createTempFile("agent-provider-read-", ".txt") + path.writeText(contents) + return ReadFixture( + path = path.toAbsolutePath().toString(), + contents = contents, + renderedResult = "1\t$contents" + ) + } + + private fun configureCustomOpenAIService(): CustomServiceSettingsState { + val settings = service() + val serviceState = CustomServiceSettingsState().apply { + name = "Agent Test Custom Service" + chatCompletionSettings.url = + System.getProperty("customOpenAI.baseUrl") + "/v1/chat/completions" + chatCompletionSettings.headers.clear() + chatCompletionSettings.body.clear() + chatCompletionSettings.body["model"] = "custom-agent-model" + } + settings.state.services.clear() + settings.state.services.add(serviceState) + setCredential( + CustomServiceApiKeyById(requireNotNull(serviceState.id)), + "TEST_API_KEY" + ) + service().setModel( + FeatureType.AGENT, + serviceState.id, + ServiceType.CUSTOM_OPENAI + ) + return serviceState + } + + private fun extractPromptText(request: RequestEntity): String { + val messages = + request.body["messages"] as? List<*> ?: request.body["input"] as? List<*> ?: return "" + return messages.joinToString("\n") { message -> + when (val content = (message as? Map<*, *>)?.get("content")) { + is String -> content + is List<*> -> content.joinToString("\n") { part -> + val partMap = part as? Map<*, *> + (partMap?.get("text") as? String) + ?: (partMap?.get("value") as? String) + ?: partMap.toString() + } + + null -> "" + else -> content.toString() + } + } + } + + private fun extractGooglePromptText(request: RequestEntity): String { + val contents = request.body["contents"] as? List<*> ?: return "" + return contents.joinToString("\n") { content -> + val parts = (content as? Map<*, *>)?.get("parts") as? List<*> ?: emptyList() + parts.joinToString("\n") { part -> + ((part as? Map<*, *>)?.get("text") as? String).orEmpty() + } + } + } + + private data class AgentRunResult( + val output: String, + val events: RecordingAgentEvents + ) + + private data class ReadFixture( + val path: String, + val contents: String, + val renderedResult: String + ) + + private class RecordingAgentEvents : AgentEvents { + val text = StringBuilder() + + override fun onTextReceived(text: String) { + this.text.append(text) + } + + override fun onQueuedMessagesResolved() = Unit + } +} diff --git a/src/test/kotlin/ee/carlrobert/codegpt/agent/AgentServiceIntegrationTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/agent/AgentServiceIntegrationTest.kt index c0147790..d022fc41 100644 --- a/src/test/kotlin/ee/carlrobert/codegpt/agent/AgentServiceIntegrationTest.kt +++ b/src/test/kotlin/ee/carlrobert/codegpt/agent/AgentServiceIntegrationTest.kt @@ -18,14 +18,13 @@ import com.intellij.openapi.components.service import com.intellij.openapi.project.Project import ee.carlrobert.codegpt.agent.history.CheckpointRef import ee.carlrobert.codegpt.conversations.Conversation -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.toolwindow.agent.AgentSession import ee.carlrobert.codegpt.toolwindow.agent.AgentToolWindowContentManager import ee.carlrobert.codegpt.ui.textarea.header.tag.McpTagDetails import kotlinx.coroutines.flow.Flow import kotlinx.coroutines.runBlocking -import kotlinx.datetime.Clock import kotlinx.serialization.json.JsonNull import org.assertj.core.api.Assertions.assertThat import org.assertj.core.groups.Tuple.tuple @@ -33,6 +32,7 @@ import testsupport.IntegrationTest import java.util.* import java.util.concurrent.ConcurrentHashMap import kotlin.io.path.Path +import kotlin.time.Clock class AgentServiceIntegrationTest : IntegrationTest() { @@ -138,7 +138,8 @@ class AgentServiceIntegrationTest : IntegrationTest() { val sessionId = createSession("agent-checkpoint-fallback") val runtimeAgentId = "runtime-agent-${UUID.randomUUID()}" val checkpointId = "checkpoint-${UUID.randomUUID()}" - val checkpointStorage = JVMFilePersistenceStorageProvider(Path(project.basePath ?: "", ".proxyai")) + val checkpointStorage = + JVMFilePersistenceStorageProvider(Path(project.basePath ?: "", ".proxyai")) val staleRef = CheckpointRef(agentId = "missing-agent", checkpointId = "missing-checkpoint") runBlocking { checkpointStorage.saveCheckpoint( @@ -205,7 +206,7 @@ class AgentServiceIntegrationTest : IntegrationTest() { } private fun agentModel(): LLModel { - return project.service().getAgentModel() + return project.service().getAgentModel() } private data class SessionSnapshot( diff --git a/src/test/kotlin/ee/carlrobert/codegpt/agent/HooksTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/agent/HooksTest.kt index 7d9d5ca3..97d44003 100644 --- a/src/test/kotlin/ee/carlrobert/codegpt/agent/HooksTest.kt +++ b/src/test/kotlin/ee/carlrobert/codegpt/agent/HooksTest.kt @@ -29,7 +29,7 @@ class HooksTest : IntegrationTest() { val tool = BashTool( project, { ShellCommandConfirmation.Approved }, - "test-pretool-deny", + "test-session-id", HookManager(project) ) ToolRunContext.set("test-pretool-deny", "tool-pretool-deny") @@ -64,10 +64,10 @@ class HooksTest : IntegrationTest() { val tool = BashTool( project, { ShellCommandConfirmation.Approved }, - "test", + "test-session-id", HookManager(project) ) - ToolRunContext.set("test", "tool-1") + ToolRunContext.set("test-session-id", "tool-1") val result = runBlocking { tool.execute(BashTool.Args(command = "echo ORIGINAL", description = "test")) @@ -105,10 +105,10 @@ class HooksTest : IntegrationTest() { val tool = BashTool( project, { ShellCommandConfirmation.Approved }, - "test", + "test-session-id", HookManager(project) ) - ToolRunContext.set("test", "tool-1") + ToolRunContext.set("test-session-id", "tool-1") val result = runBlocking { tool.execute(BashTool.Args(command = "echo ORIGINAL", description = "test")) @@ -137,7 +137,7 @@ class HooksTest : IntegrationTest() { val tool = BashTool( project, { ShellCommandConfirmation.Approved }, - "test", + "test-session-id", HookManager(project) ) ToolRunContext.set("test", "tool-1") @@ -168,26 +168,24 @@ class HooksTest : IntegrationTest() { val settings = """{"beforeShellExecution":[{"command":".proxyai/hooks/${hookScript.name}","matcher":"gradlew"}]}""" writeSettings(settings) - val tool = BashTool( project, { ShellCommandConfirmation.Approved }, - "test-bash-gradlew", + "test-session-id", HookManager(project) ) - ToolRunContext.set("test-bash-gradlew", "tool-bash-gradlew") + ToolRunContext.set("test-session-id", "tool-bash-gradlew") val blocked = runBlocking { tool.execute(BashTool.Args(command = "bash ./gradlew test", description = "test")) } + assertThat(blocked.exitCode).isNull() assertThat(blocked.output).isEqualTo("Blocked bash gradlew") - val map = ObjectMapper().registerKotlinModule() .readValue(logFile.readText(), Map::class.java) assertThat(map["hook_event_name"]).isEqualTo("beforeShellExecution") assertThat(map["command"]).isEqualTo("bash ./gradlew test") - val allowed = runBlocking { tool.execute(BashTool.Args(command = "echo should_run", description = "test")) } @@ -212,7 +210,7 @@ class HooksTest : IntegrationTest() { writeSettings(settings) val result = runBlocking { - EditTool(project, HookManager(project), "test") + EditTool(project, "test-session-id", HookManager(project)) .execute( EditTool.Args( target.absolutePath, @@ -243,7 +241,7 @@ class HooksTest : IntegrationTest() { writeSettings(settings) val result = runBlocking { - ReadTool(project, HookManager(project), "test") + ReadTool(project, "test-session-id", HookManager(project)) .execute(ReadTool.Args(target.absolutePath)) } @@ -268,7 +266,7 @@ class HooksTest : IntegrationTest() { writeSettings(settings) val result = runBlocking { - EditTool(project, HookManager(project), "test") + EditTool(project, "test-session-id", HookManager(project)) .execute( EditTool.Args( target.absolutePath, diff --git a/src/test/kotlin/ee/carlrobert/codegpt/agent/IgnoreRulesTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/agent/IgnoreRulesTest.kt index 88d69a44..d204b84f 100644 --- a/src/test/kotlin/ee/carlrobert/codegpt/agent/IgnoreRulesTest.kt +++ b/src/test/kotlin/ee/carlrobert/codegpt/agent/IgnoreRulesTest.kt @@ -19,7 +19,7 @@ class IgnoreRulesTest : IntegrationTest() { writeSettings(ignoreEntries = listOf(".env")) val result = runBlocking { - ReadTool(project, HookManager(project), "ignore-test") + ReadTool(project, "test-session-id", HookManager(project)) .execute(ReadTool.Args(envFile.absolutePath)) } @@ -33,7 +33,7 @@ class IgnoreRulesTest : IntegrationTest() { writeSettings(ignoreEntries = listOf("**/*.pem")) val result = runBlocking { - WriteTool(project, HookManager(project)) + WriteTool(project, "test-session-id", HookManager(project)) .execute(WriteTool.Args(pemFile.absolutePath, "PRIVATE KEY")) } @@ -52,7 +52,7 @@ class IgnoreRulesTest : IntegrationTest() { val tool = BashTool( project = project, confirmationHandler = { ShellCommandConfirmation.Approved }, - sessionId = "ignore-bash", + sessionId = "test-session-id", hookManager = HookManager(project) ) val result = runBlocking { @@ -76,7 +76,7 @@ class IgnoreRulesTest : IntegrationTest() { writeSettings(ignoreEntries = listOf("secrets/**")) val result = runBlocking { - ReadTool(project, HookManager(project), "ignore-test") + ReadTool(project, "test-session-id", HookManager(project)) .execute(ReadTool.Args(file.absolutePath)) } @@ -91,7 +91,7 @@ class IgnoreRulesTest : IntegrationTest() { writeSettingsRaw("""{"ignore":["app/src/main/**"]}""") val result = runBlocking { - ReadTool(project, HookManager(project), "ignore-test") + ReadTool(project, "test-session-id", HookManager(project)) .execute(ReadTool.Args(file.absolutePath)) } @@ -108,7 +108,7 @@ class IgnoreRulesTest : IntegrationTest() { writeSettings(ignoreEntries = listOf("build/")) val result = runBlocking { - ReadTool(project, HookManager(project), "ignore-test") + ReadTool(project, "test-session-id", HookManager(project)) .execute(ReadTool.Args(file.absolutePath)) } @@ -122,7 +122,7 @@ class IgnoreRulesTest : IntegrationTest() { writeSettings(ignoreEntries = listOf("*.pem")) val result = runBlocking { - WriteTool(project, HookManager(project)) + WriteTool(project, "test-session-id", HookManager(project)) .execute(WriteTool.Args(file.absolutePath, "PRIVATE KEY")) } @@ -136,7 +136,10 @@ class IgnoreRulesTest : IntegrationTest() { val file = File(project.basePath, ".proxyai/settings.json") file.parentFile.mkdirs() file.writeText("""{"ignore":[$ignoreJson],"permissions":{"allow":[],"ask":[],"deny":[]},"hooks":{}}""") - Files.setLastModifiedTime(file.toPath(), FileTime.fromMillis(System.currentTimeMillis() + 1000)) + Files.setLastModifiedTime( + file.toPath(), + FileTime.fromMillis(System.currentTimeMillis() + 1000) + ) return file } @@ -144,7 +147,10 @@ class IgnoreRulesTest : IntegrationTest() { val file = File(project.basePath, ".proxyai/settings.json") file.parentFile.mkdirs() file.writeText(rawJson) - Files.setLastModifiedTime(file.toPath(), FileTime.fromMillis(System.currentTimeMillis() + 1000)) + Files.setLastModifiedTime( + file.toPath(), + FileTime.fromMillis(System.currentTimeMillis() + 1000) + ) return file } } diff --git a/src/test/kotlin/ee/carlrobert/codegpt/agent/PermissionsTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/agent/PermissionsTest.kt index 084e1f08..3cdc1fb3 100644 --- a/src/test/kotlin/ee/carlrobert/codegpt/agent/PermissionsTest.kt +++ b/src/test/kotlin/ee/carlrobert/codegpt/agent/PermissionsTest.kt @@ -20,8 +20,8 @@ class PermissionsTest : IntegrationTest() { ) val result = runBlocking { - seedToolContext("perm-test") - ReadTool(project, HookManager(project), "perm-test") + seedToolContext() + ReadTool(project, "test-session-id", HookManager(project)) .execute(ReadTool.Args(file.absolutePath)) } @@ -40,8 +40,8 @@ class PermissionsTest : IntegrationTest() { ) val result = runBlocking { - seedToolContext("perm-test") - ReadTool(project, HookManager(project), "perm-test") + seedToolContext() + ReadTool(project, "test-session-id", HookManager(project)) .execute(ReadTool.Args(file.absolutePath)) } @@ -60,8 +60,8 @@ class PermissionsTest : IntegrationTest() { ) val result = runBlocking { - seedToolContext("perm-test") - ReadTool(project, HookManager(project), "perm-test") + seedToolContext() + ReadTool(project, "test-session-id", HookManager(project)) .execute(ReadTool.Args(file.absolutePath)) } @@ -73,9 +73,9 @@ class PermissionsTest : IntegrationTest() { denyEntries = listOf("Bash(git push *)") ) - val tool = createBashTool("perm-bash") { ShellCommandConfirmation.Approved } + val tool = createBashTool("test-session-id") { ShellCommandConfirmation.Approved } val result = runBlocking { - seedToolContext("perm-bash") + seedToolContext() tool.execute(BashTool.Args(command = "git push origin main", description = "push")) } @@ -88,11 +88,11 @@ class PermissionsTest : IntegrationTest() { allowEntries = listOf("Bash(echo *)") ) - val tool = createBashTool("perm-bash-allow") { + val tool = createBashTool("test-session-id") { ShellCommandConfirmation.Denied("should not be called when allow matches") } val result = runBlocking { - seedToolContext("perm-bash-allow") + seedToolContext() tool.execute(BashTool.Args(command = "echo OK", description = "echo")) } @@ -105,11 +105,11 @@ class PermissionsTest : IntegrationTest() { allowEntries = listOf("Bash(echo *)") ) - val tool = createBashTool("perm-bash-deny") { + val tool = createBashTool("test-session-id") { ShellCommandConfirmation.Denied("explicit deny") } val result = runBlocking { - seedToolContext("perm-bash-deny") + seedToolContext() tool.execute(BashTool.Args(command = "git status", description = "status")) } @@ -147,10 +147,10 @@ class PermissionsTest : IntegrationTest() { return file } - private fun seedToolContext(sessionId: String) { + private fun seedToolContext() { ToolRunContext.set( - sessionId = sessionId, - toolId = "test-tool-$sessionId" + sessionId = "test-session-id", + toolId = "test-tool-id" ) } } diff --git a/src/test/kotlin/ee/carlrobert/codegpt/agent/WebFetchToolTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/agent/WebFetchToolTest.kt index bdff599e..28e82183 100644 --- a/src/test/kotlin/ee/carlrobert/codegpt/agent/WebFetchToolTest.kt +++ b/src/test/kotlin/ee/carlrobert/codegpt/agent/WebFetchToolTest.kt @@ -32,6 +32,7 @@ class WebFetchToolTest : IntegrationTest() { val result = runBlocking { WebFetchTool( workingDirectory = project.basePath ?: "", + sessionId = "test-session-id", hookManager = HookManager(project) ).execute(WebFetchTool.Args(url = "$baseUrl/docs")) } @@ -69,6 +70,7 @@ class WebFetchToolTest : IntegrationTest() { val result = runBlocking { WebFetchTool( workingDirectory = project.basePath ?: "", + sessionId = "test-session-id", hookManager = HookManager(project) ).execute( WebFetchTool.Args( @@ -97,6 +99,7 @@ class WebFetchToolTest : IntegrationTest() { val result = runBlocking { WebFetchTool( workingDirectory = project.basePath ?: "", + sessionId = "test-session-id", hookManager = HookManager(project) ).execute(WebFetchTool.Args(url = "$baseUrl/api")) } @@ -161,6 +164,7 @@ class WebFetchToolTest : IntegrationTest() { private fun createTool(): WebFetchTool { return WebFetchTool( workingDirectory = project.basePath ?: "", + sessionId = "test-session-id", hookManager = HookManager(project) ) } diff --git a/src/test/kotlin/ee/carlrobert/codegpt/agent/WriteToolPathResolutionTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/agent/WriteToolPathResolutionTest.kt index abcab983..8dbd27f5 100644 --- a/src/test/kotlin/ee/carlrobert/codegpt/agent/WriteToolPathResolutionTest.kt +++ b/src/test/kotlin/ee/carlrobert/codegpt/agent/WriteToolPathResolutionTest.kt @@ -13,7 +13,7 @@ class WriteToolPathResolutionTest : IntegrationTest() { val target = File(project.basePath, "app/components/NewLanding.tsx") val result = runBlocking { - WriteTool(project, HookManager(project)) + WriteTool(project, "test-session-id", HookManager(project)) .execute( WriteTool.Args( target.absolutePath, @@ -32,7 +32,7 @@ class WriteToolPathResolutionTest : IntegrationTest() { target.parentFile.mkdirs() val result = runBlocking { - WriteTool(project, HookManager(project)) + WriteTool(project, "test-session-id", HookManager(project)) .execute( WriteTool.Args( "/$pseudoRootDir/components/NewLanding.tsx", diff --git a/src/test/kotlin/ee/carlrobert/codegpt/agent/clients/CustomOpenAILLMClientBodyTransformTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/agent/clients/CustomOpenAILLMClientBodyTransformTest.kt new file mode 100644 index 00000000..93b1f194 --- /dev/null +++ b/src/test/kotlin/ee/carlrobert/codegpt/agent/clients/CustomOpenAILLMClientBodyTransformTest.kt @@ -0,0 +1,170 @@ +package ee.carlrobert.codegpt.agent.clients + +import ai.koog.prompt.executor.clients.openai.base.models.Content +import ai.koog.prompt.executor.clients.openai.base.models.OpenAIMessage +import ee.carlrobert.codegpt.settings.service.custom.CustomServiceChatCompletionSettingsState +import kotlinx.serialization.json.* +import org.assertj.core.api.Assertions.assertThat +import org.junit.Test + +class CustomOpenAILLMClientBodyTransformTest { + + private val json = Json { + ignoreUnknownKeys = true + encodeDefaults = true + explicitNulls = false + } + + @Test + fun `should transform nested placeholder values in custom body`() { + val body = mapOf( + "model" to "gpt-test", + "payload" to mapOf( + "prompt_alias" to $$"$PROMPT", + "messages_alias" to $$"$OPENAI_MESSAGES", + "items" to listOf( + mapOf("kind" to "prompt", "value" to $$"$PROMPT"), + mapOf("kind" to "messages", "value" to $$"$OPENAI_MESSAGES"), + ) + ), + "token" to $$"Bearer $CUSTOM_SERVICE_API_KEY", + ) + + val properties = buildCustomOpenAIAdditionalProperties( + body = body, + messages = listOf( + OpenAIMessage.System(content = Content.Text("System instructions")), + OpenAIMessage.User(content = Content.Text("Write a test")), + ), + streamRequest = true, + credential = "secret-key", + json = json + ) ?: emptyMap() + + val payload = properties.getValue("payload").jsonObject + val payloadItems = payload.getValue("items").jsonArray + assertThat(properties.getValue("token").jsonPrimitive.content).isEqualTo("Bearer secret-key") + assertThat(payload.getValue("prompt_alias").jsonPrimitive.content) + .isEqualTo("System instructions\n\nWrite a test") + assertThat( + payload.getValue("messages_alias").jsonArray + .map { message -> + val msg = message.jsonObject + msg.getValue("role").jsonPrimitive.content to msg.getValue("content").jsonPrimitive.content + } + ).containsExactly( + "system" to "System instructions", + "user" to "Write a test" + ) + assertThat( + payloadItems.map { item -> + val node = item.jsonObject + Triple( + node.getValue("kind").jsonPrimitive.content, + node.getValue("value").toString().startsWith("["), + node.getValue("value").toString().startsWith("\"") + ) + } + ).containsExactly( + Triple("prompt", false, true), + Triple("messages", true, false) + ) + assertThat(payloadItems[0].jsonObject.getValue("value").jsonPrimitive.content) + .isEqualTo("System instructions\n\nWrite a test") + assertThat( + payloadItems[1].jsonObject.getValue("value").jsonArray + .map { it.jsonObject.getValue("role").jsonPrimitive.content } + ).containsExactly("system", "user") + } + + @Test + fun `should preserve nested arrays and objects in custom body`() { + val body = mapOf( + "model" to "gpt-test", + "config" to mapOf( + "temperature" to 0.2, + "enabled" to true, + "tags" to listOf("alpha", 2, false), + "nested" to mapOf("arr" to listOf(mapOf("k" to "v"))) + ) + ) + + val properties = buildCustomOpenAIAdditionalProperties( + body = body, + messages = listOf(OpenAIMessage.User(content = Content.Text("hello"))), + streamRequest = true, + credential = "secret-key", + json = json + ) ?: emptyMap() + + val config = properties.getValue("config").jsonObject + assertThat(config.getValue("temperature").jsonPrimitive.double).isEqualTo(0.2) + assertThat(config.getValue("enabled").jsonPrimitive.boolean).isTrue() + assertThat( + config.getValue("tags").jsonArray + .map { it.jsonPrimitive.content }).containsExactly("alpha", "2", "false") + assertThat( + config.getValue("nested").jsonObject.getValue("arr").jsonArray + .map { it.jsonObject.getValue("k").jsonPrimitive.content } + ).containsExactly("v") + } + + @Test + fun `should replace nested stream flags and api key placeholders when stream is disabled`() { + val body = mapOf( + "model" to "gpt-test", + "stream" to "true", + "config" to mapOf( + "stream" to "true", + "authorization" to $$"Bearer $CUSTOM_SERVICE_API_KEY" + ), + "items" to listOf( + mapOf("stream" to "true"), + $$"token:$CUSTOM_SERVICE_API_KEY" + ) + ) + + val properties = buildCustomOpenAIAdditionalProperties( + body = body, + messages = listOf(OpenAIMessage.User(content = Content.Text("hello"))), + streamRequest = false, + credential = "secret-key", + json = json + ) ?: emptyMap() + + assertThat(properties).doesNotContainKey("stream") + assertThat( + listOf( + properties.getValue("config").jsonObject.getValue("stream").jsonPrimitive.boolean, + properties.getValue("items").jsonArray[0].jsonObject.getValue("stream").jsonPrimitive.boolean + ) + ).containsExactly(false, false) + assertThat( + listOf( + properties.getValue("config").jsonObject.getValue("authorization").jsonPrimitive.content, + properties.getValue("items").jsonArray[1].jsonPrimitive.content + ) + ).containsExactly("Bearer secret-key", "token:secret-key") + } + + @Test + fun `should enable custom openai streaming only when body stream is explicitly true`() { + val defaultState = CustomServiceChatCompletionSettingsState().apply { + body.remove("stream") + } + val disabledState = CustomServiceChatCompletionSettingsState().apply { + body["stream"] = false + } + val enabledState = CustomServiceChatCompletionSettingsState().apply { + body["stream"] = true + } + val enabledStringState = CustomServiceChatCompletionSettingsState().apply { + body["stream"] = "true" + } + + assertThat(defaultState.shouldStream()).isFalse() + assertThat(disabledState.shouldStream()).isFalse() + assertThat(enabledState.shouldStream()).isTrue() + assertThat(enabledStringState.shouldStream()).isTrue() + } +} diff --git a/src/test/kotlin/ee/carlrobert/codegpt/agent/clients/RetryingPromptExecutorRetryabilityTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/agent/clients/RetryingPromptExecutorRetryabilityTest.kt new file mode 100644 index 00000000..e548bec6 --- /dev/null +++ b/src/test/kotlin/ee/carlrobert/codegpt/agent/clients/RetryingPromptExecutorRetryabilityTest.kt @@ -0,0 +1,45 @@ +package ee.carlrobert.codegpt.agent.clients + +import ai.koog.http.client.KoogHttpClientException +import org.assertj.core.api.Assertions.assertThat +import org.junit.Test + +class RetryingPromptExecutorRetryabilityTest { + + @Test + fun `should retry for transient HTTP status codes`() { + val tooManyRequests = KoogHttpClientException("rate limited", 429, null, null, null) + val unavailable = KoogHttpClientException("service unavailable", 503, null, null, null) + + assertThat(RetryingPromptExecutor.isRetryableFailure(tooManyRequests)).isTrue() + assertThat(RetryingPromptExecutor.isRetryableFailure(unavailable)).isTrue() + } + + @Test + fun `should not retry for non-transient HTTP status codes`() { + val badRequest = KoogHttpClientException("bad request", 400, null, null, null) + val unauthorized = KoogHttpClientException("unauthorized", 401, null, null, null) + + assertThat(RetryingPromptExecutor.isRetryableFailure(badRequest)).isFalse() + assertThat(RetryingPromptExecutor.isRetryableFailure(unauthorized)).isFalse() + } + + @Test + fun `should retry when transient status is in nested cause chain`() { + val nested = RuntimeException( + "outer", + KoogHttpClientException("service unavailable", 503, null, null, null) + ) + assertThat(RetryingPromptExecutor.isRetryableFailure(nested)).isTrue() + } + + @Test + fun `should not retry when no HTTP status is present`() { + assertThat(RetryingPromptExecutor.isRetryableFailure(RuntimeException("connection reset"))).isFalse() + } + + @Test + fun `should not retry for non-network application errors`() { + assertThat(RetryingPromptExecutor.isRetryableFailure(IllegalArgumentException("invalid prompt"))).isFalse() + } +} diff --git a/src/test/kotlin/ee/carlrobert/codegpt/agent/history/AgentCheckpointConversationMapperTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/agent/history/AgentCheckpointConversationMapperTest.kt index ac1715ba..a25fcd56 100644 --- a/src/test/kotlin/ee/carlrobert/codegpt/agent/history/AgentCheckpointConversationMapperTest.kt +++ b/src/test/kotlin/ee/carlrobert/codegpt/agent/history/AgentCheckpointConversationMapperTest.kt @@ -4,13 +4,13 @@ import ai.koog.agents.snapshot.feature.AgentCheckpointData import ai.koog.prompt.message.Message import ai.koog.prompt.message.RequestMetaInfo import ai.koog.prompt.message.ResponseMetaInfo -import kotlinx.datetime.Clock -import kotlinx.datetime.Instant import kotlinx.serialization.json.JsonNull import kotlinx.serialization.json.JsonObject import kotlinx.serialization.json.JsonPrimitive import org.assertj.core.api.Assertions.assertThat import org.assertj.core.groups.Tuple.tuple +import kotlin.time.Clock +import kotlin.time.Instant import kotlin.test.Test class AgentCheckpointConversationMapperTest { diff --git a/src/test/kotlin/ee/carlrobert/codegpt/agent/history/AgentCheckpointHistoryServiceTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/agent/history/AgentCheckpointHistoryServiceTest.kt index e9cae785..5214d77f 100644 --- a/src/test/kotlin/ee/carlrobert/codegpt/agent/history/AgentCheckpointHistoryServiceTest.kt +++ b/src/test/kotlin/ee/carlrobert/codegpt/agent/history/AgentCheckpointHistoryServiceTest.kt @@ -7,8 +7,6 @@ import ai.koog.prompt.message.RequestMetaInfo import ai.koog.prompt.message.ResponseMetaInfo import com.intellij.openapi.components.service import kotlinx.coroutines.runBlocking -import kotlinx.datetime.Clock -import kotlinx.datetime.Instant import kotlinx.serialization.json.JsonNull import kotlinx.serialization.json.JsonObject import kotlinx.serialization.json.JsonPrimitive @@ -16,6 +14,8 @@ import org.assertj.core.api.Assertions.assertThat import org.assertj.core.groups.Tuple.tuple import testsupport.IntegrationTest import kotlin.io.path.Path +import kotlin.time.Clock +import kotlin.time.Instant import java.util.UUID class AgentCheckpointHistoryServiceTest : IntegrationTest() { diff --git a/src/test/kotlin/ee/carlrobert/codegpt/agent/history/AgentCheckpointTurnSequencerTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/agent/history/AgentCheckpointTurnSequencerTest.kt index 92e517b3..034b22ab 100644 --- a/src/test/kotlin/ee/carlrobert/codegpt/agent/history/AgentCheckpointTurnSequencerTest.kt +++ b/src/test/kotlin/ee/carlrobert/codegpt/agent/history/AgentCheckpointTurnSequencerTest.kt @@ -4,12 +4,12 @@ import ai.koog.agents.snapshot.feature.AgentCheckpointData import ai.koog.prompt.message.Message import ai.koog.prompt.message.RequestMetaInfo import ai.koog.prompt.message.ResponseMetaInfo -import kotlinx.datetime.Clock -import kotlinx.datetime.Instant import kotlinx.serialization.json.JsonNull import kotlinx.serialization.json.JsonObject import kotlinx.serialization.json.JsonPrimitive import org.assertj.core.api.Assertions.assertThat +import kotlin.time.Clock +import kotlin.time.Instant import kotlin.test.Test class AgentCheckpointTurnSequencerTest { diff --git a/src/test/kotlin/ee/carlrobert/codegpt/agent/strategy/SingleRunStrategyProviderTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/agent/strategy/SingleRunStrategyProviderTest.kt new file mode 100644 index 00000000..b6037155 --- /dev/null +++ b/src/test/kotlin/ee/carlrobert/codegpt/agent/strategy/SingleRunStrategyProviderTest.kt @@ -0,0 +1,97 @@ +package ee.carlrobert.codegpt.agent.strategy + +import ai.koog.prompt.llm.LLMProvider +import ai.koog.prompt.message.Message +import ai.koog.prompt.message.ResponseMetaInfo +import org.assertj.core.api.Assertions.assertThat +import kotlin.test.Test + +class SingleRunStrategyProviderTest { + + @Test + fun `extract assistant output prefers assistant text when available`() { + val responses = listOf( + Message.Reasoning(content = "Reasoning text", metaInfo = ResponseMetaInfo.Empty), + Message.Assistant("Final answer", ResponseMetaInfo.Empty) + ) + + val actual = extractAssistantOutput(responses) + + assertThat(actual).isEqualTo("Final answer") + } + + @Test + fun `extract assistant output falls back to reasoning when assistant text is blank`() { + val responses = listOf( + Message.Reasoning(content = "Reasoning text", metaInfo = ResponseMetaInfo.Empty), + Message.Assistant("", ResponseMetaInfo.Empty) + ) + + val actual = extractAssistantOutput(responses) + + assertThat(actual).isEqualTo("Reasoning text") + } + + @Test + fun `appendable responses keeps reasoning fallback for non google when assistant is blank`() { + val responses = listOf( + Message.Reasoning(content = "Reasoning text", metaInfo = ResponseMetaInfo.Empty), + Message.Assistant("", ResponseMetaInfo.Empty) + ) + + val actual = appendableResponses(responses, LLMProvider.OpenAI) + + assertThat(actual) + .singleElement() + .isInstanceOf(Message.Reasoning::class.java) + } + + @Test + fun `appendable responses removes reasoning for non google when assistant has text`() { + val responses = listOf( + Message.Reasoning(content = "Reasoning text", metaInfo = ResponseMetaInfo.Empty), + Message.Assistant("Final answer", ResponseMetaInfo.Empty) + ) + + val actual = appendableResponses(responses, LLMProvider.OpenAI) + + assertThat(actual) + .singleElement() + .isInstanceOf(Message.Assistant::class.java) + } + + @Test + fun `appendable responses keeps reasoning for google`() { + val responses = listOf( + Message.Reasoning(content = "Reasoning text", metaInfo = ResponseMetaInfo.Empty), + Message.Assistant("Final answer", ResponseMetaInfo.Empty) + ) + + val actual = appendableResponses(responses, LLMProvider.Google) + + assertThat(actual.map { it::class.simpleName }) + .containsExactly("Assistant", "Reasoning") + } + + @Test + fun `appendable responses drops assistant text for google tool turns`() { + val responses = listOf( + Message.Reasoning(content = "Reasoning text", metaInfo = ResponseMetaInfo.Empty), + Message.Assistant("I will inspect the file.", ResponseMetaInfo.Empty), + Message.Tool.Call( + id = "tool-1", + tool = "Read", + content = """{"path":"build.gradle.kts"}""", + metaInfo = ResponseMetaInfo.Empty + ) + ) + + val actual = appendableResponses(responses, LLMProvider.Google) + val toolCall = actual.filterIsInstance().single() + + assertThat(actual.map { it::class.simpleName }) + .containsExactly("Reasoning", "Call") + assertThat(toolCall.tool).isEqualTo("Read") + assertThat(toolCall.content).isEqualTo("""{"path":"build.gradle.kts"}""") + } +} diff --git a/src/test/kotlin/ee/carlrobert/codegpt/codecompletions/ChatBasedFIMTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/codecompletions/ChatBasedFIMTest.kt deleted file mode 100644 index 5bf594d7..00000000 --- a/src/test/kotlin/ee/carlrobert/codegpt/codecompletions/ChatBasedFIMTest.kt +++ /dev/null @@ -1,99 +0,0 @@ -package ee.carlrobert.codegpt.codecompletions - -import ee.carlrobert.codegpt.settings.service.FeatureType -import ee.carlrobert.codegpt.settings.service.ModelSelectionService -import ee.carlrobert.llm.client.openai.completion.request.OpenAIChatCompletionStandardMessage -import org.junit.jupiter.api.Test -import org.junit.jupiter.api.Assertions.* - -class ChatBasedFIMTest { - - @Test - fun `test chat completion template is available`() { - val templates = InfillPromptTemplate.values() - val chatTemplate = templates.find { it == InfillPromptTemplate.CHAT_COMPLETION } - - assertNotNull(chatTemplate) - assertEquals("Chat-based FIM", chatTemplate?.label) - assertEquals(listOf("\n\n", "```"), chatTemplate?.stopTokens) - } - - @Test - fun `test chat completion template builds placeholder`() { - val template = InfillPromptTemplate.CHAT_COMPLETION - val infillRequest = InfillRequest.Builder("function test() {", "}", 0).build() - - val result = template.buildPrompt(infillRequest) - assertEquals("CHAT_FIM_PLACEHOLDER", result) - } - - @Test - fun `test chat based FIM request creation`() { - val infillRequest = InfillRequest.Builder( - "function calculateSum(a, b) {", - "return result;\n}", - 0 - ).build() - - val chatRequest = CodeCompletionRequestFactory.buildChatBasedFIMRequest(infillRequest) - - assertNotNull(chatRequest) - assertEquals(2, chatRequest.messages.size) - - val systemMessage = chatRequest.messages[0] as OpenAIChatCompletionStandardMessage - assertEquals("system", systemMessage.role) - assertTrue(systemMessage.content.contains("expert coding assistant")) - - val userMessage = chatRequest.messages[1] as OpenAIChatCompletionStandardMessage - assertEquals("user", userMessage.role) - assertTrue(userMessage.content.contains("PREFIX:")) - assertTrue(userMessage.content.contains("SUFFIX:")) - assertTrue(userMessage.content.contains("function calculateSum(a, b) {")) - assertTrue(userMessage.content.contains("return result;")) - - assertTrue(chatRequest.isStream) - assertEquals(128, chatRequest.maxTokens) - assertEquals(0.0, chatRequest.temperature) - } - - @Test - fun `test custom request throws exception for chat completion template`() { - val infillRequest = InfillRequest.Builder("test", "test", 0).build() - - assertThrows(IllegalArgumentException::class.java) { - CodeCompletionRequestFactory.buildCustomRequest( - infillRequest, - "http://test.com", - emptyMap(), - emptyMap(), - InfillPromptTemplate.CHAT_COMPLETION, - null - ) - } - } - - @Test - fun `test chat based FIM HTTP request creation`() { - val infillRequest = InfillRequest.Builder("test prefix", "test suffix", 0).build() - val headers = mapOf("Authorization" to "Bearer \$CUSTOM_SERVICE_API_KEY") - val body = mapOf( - "model" to "gpt-4.1", - "temperature" to 0.2, - "max_tokens" to 24, - "stream" to false - ) - - val httpRequest = CodeCompletionRequestFactory.buildChatBasedFIMHttpRequest( - infillRequest, - "https://api.openai.com/v1/chat/completions", - headers, - body, - "test-api-key" - ) - - assertNotNull(httpRequest) - assertEquals("https://api.openai.com/v1/chat/completions", httpRequest.url.toString()) - assertEquals("Bearer test-api-key", httpRequest.header("Authorization")) - assertEquals("application/json", httpRequest.body?.contentType()?.toString()) - } -} diff --git a/src/test/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionPromptFactoryTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionPromptFactoryTest.kt new file mode 100644 index 00000000..c4b98c46 --- /dev/null +++ b/src/test/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionPromptFactoryTest.kt @@ -0,0 +1,46 @@ +package ee.carlrobert.codegpt.codecompletions + +import ai.koog.prompt.message.Message +import ee.carlrobert.codegpt.EncodingManager +import org.assertj.core.api.Assertions.assertThat +import junit.framework.TestCase + +class CodeCompletionPromptFactoryTest : TestCase() { + + fun testBuildsChatFimPromptWithDeterministicCompletionParams() { + val infillRequest = InfillRequest.Builder( + prefix = "fun sum(a: Int, b: Int) = ", + suffix = " + c", + caretOffset = 0 + ).build() + + val prompt = CodeCompletionPromptFactory.createPrompt(infillRequest) + val systemMessage = prompt.messages[0] as Message.System + val userMessage = prompt.messages[1] as Message.User + + assertThat(prompt.params.maxTokens).isEqualTo(128) + assertThat(prompt.params.temperature).isEqualTo(0.0) + assertThat(systemMessage.content).contains("code completion assistant") + assertThat(userMessage.content).contains("") + assertThat(userMessage.content).contains("") + assertThat(userMessage.content).contains("") + assertThat(userMessage.content).contains("") + assertThat(userMessage.content).contains("fun sum(a: Int, b: Int) = ") + assertThat(userMessage.content).contains(" + c") + } + + fun testBuildTruncatesGitDiffAugmentedPrefixBackToPromptBudget() { + val gitDiff = (1..4_000).joinToString("\n") { "diff line $it" } + val originalPrefix = "fun main() {\n pri" + val infillRequest = InfillRequest.Builder( + prefix = originalPrefix, + suffix = "", + caretOffset = 0 + ).gitDiff(gitDiff).build() + + assertThat(EncodingManager.getInstance().countTokens(infillRequest.prefix)) + .isLessThanOrEqualTo(MAX_PROMPT_TOKENS) + assertThat(infillRequest.prefix).contains(originalPrefix) + assertThat(infillRequest.prefix).doesNotContain("diff line 1") + } +} diff --git a/src/test/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionServiceTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionServiceTest.kt index 729343ea..a052c8c3 100644 --- a/src/test/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionServiceTest.kt +++ b/src/test/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionServiceTest.kt @@ -7,17 +7,46 @@ import com.intellij.openapi.util.TextRange import com.intellij.testFramework.PlatformTestUtil import ee.carlrobert.codegpt.CodeGPTKeys.REMAINING_EDITOR_COMPLETION import ee.carlrobert.codegpt.completions.HuggingFaceModel +import ee.carlrobert.codegpt.completions.llama.LlamaModel +import ee.carlrobert.codegpt.credentials.CredentialsStore.CredentialKey.CustomServiceApiKeyById +import ee.carlrobert.codegpt.credentials.CredentialsStore.setCredential +import ee.carlrobert.codegpt.settings.models.ModelSettings import ee.carlrobert.codegpt.settings.service.FeatureType +import ee.carlrobert.codegpt.settings.service.ServiceType import ee.carlrobert.codegpt.settings.service.codegpt.CodeGPTServiceSettings +import ee.carlrobert.codegpt.settings.service.custom.CustomServiceSettingsState +import ee.carlrobert.codegpt.settings.service.custom.CustomServicesSettings import ee.carlrobert.codegpt.util.file.FileUtil -import ee.carlrobert.llm.client.http.RequestEntity -import ee.carlrobert.llm.client.http.exchange.StreamHttpExchange -import ee.carlrobert.llm.client.util.JSONUtil.* import org.assertj.core.api.Assertions.assertThat import testsupport.IntegrationTest +import testsupport.http.RequestEntity +import testsupport.http.ResponseEntity +import testsupport.http.exchange.BasicHttpExchange +import testsupport.http.exchange.StreamHttpExchange +import testsupport.json.JSONUtil.e +import testsupport.json.JSONUtil.jsonArray +import testsupport.json.JSONUtil.jsonMap +import testsupport.json.JSONUtil.jsonMapResponse +import ee.carlrobert.service.CodeCompletionServiceImplGrpc +import ee.carlrobert.service.GrpcCodeCompletionRequest +import ee.carlrobert.service.PartialCodeCompletionResponse +import io.grpc.Server +import io.grpc.ServerBuilder +import io.grpc.stub.StreamObserver class CodeCompletionServiceTest : IntegrationTest() { + private var proxyAiGrpcServer: Server? = null + + override fun tearDown() { + proxyAiGrpcServer?.shutdownNow() + proxyAiGrpcServer = null + System.clearProperty("proxyai.grpc.host") + System.clearProperty("proxyai.grpc.port") + System.clearProperty("proxyai.grpc.plaintext") + super.tearDown() + } + fun `test code completion with OpenAI provider`() { useOpenAIService("gpt-4", FeatureType.CODE_COMPLETION) service().state.nextEditsEnabled = false @@ -39,16 +68,17 @@ class CodeCompletionServiceTest : IntegrationTest() { zzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzz xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx """.trimIndent() - expectOpenAI(StreamHttpExchange { request: RequestEntity -> + expectOpenAI(BasicHttpExchange { request: RequestEntity -> assertThat(request.uri.path).isEqualTo("/v1/completions") assertThat(request.method).isEqualTo("POST") assertThat(request.body) - .extracting("model", "prompt", "suffix", "max_tokens") - .containsExactly("gpt-3.5-turbo-instruct", prefix, suffix, 128) - listOf( - jsonMapResponse("choices", jsonArray(jsonMap("text", "ublic "))), - jsonMapResponse("choices", jsonArray(jsonMap("text", "void"))), - jsonMapResponse("choices", jsonArray(jsonMap("text", " main"))) + .extracting("model", "prompt", "suffix", "max_tokens", "stream") + .containsExactly("gpt-3.5-turbo-instruct", prefix, suffix, 128, false) + ResponseEntity( + jsonMapResponse( + "choices", + jsonArray(jsonMap("text", "ublic void main")) + ) ) }) @@ -79,19 +109,19 @@ class CodeCompletionServiceTest : IntegrationTest() { zzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzz xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx """.trimIndent() - expectOllama(StreamHttpExchange { request: RequestEntity -> + expectOllama(BasicHttpExchange { request: RequestEntity -> assertThat(request.uri.path).isEqualTo("/api/generate") assertThat(request.method).isEqualTo("POST") assertThat(request.body) - .extracting("model", "prompt", "suffix") - .containsExactly(HuggingFaceModel.CODE_QWEN_2_5_3B_Q4_K_M.code, prefix, suffix) - listOf( + .extracting("model", "prompt", "suffix", "stream") + .containsExactly(HuggingFaceModel.CODE_QWEN_2_5_3B_Q4_K_M.code, prefix, suffix, false) + ResponseEntity( jsonMapResponse( e("model", HuggingFaceModel.CODE_QWEN_2_5_3B_Q4_K_M.code), e("created_at", "2023-08-04T08:52:19.385406455-07:00"), e("response", "rivate void main"), e("done", true), - ), + ) ) }) @@ -102,6 +132,160 @@ class CodeCompletionServiceTest : IntegrationTest() { } } + fun `test code completion with custom openai provider uses completion endpoint`() { + val customService = CustomServiceSettingsState().apply { + name = "Completion Test Custom Service" + codeCompletionSettings.url = System.getProperty("customOpenAI.baseUrl") + "/v1/completions" + codeCompletionSettings.headers.clear() + codeCompletionSettings.headers["Authorization"] = "Bearer \$CUSTOM_SERVICE_API_KEY" + codeCompletionSettings.body.clear() + codeCompletionSettings.body["stream"] = true + codeCompletionSettings.body["prompt"] = "\$PREFIX" + codeCompletionSettings.body["suffix"] = "\$SUFFIX" + codeCompletionSettings.body["model"] = "custom-code-model" + codeCompletionSettings.body["max_tokens"] = 128 + } + service().state.services.clear() + service().state.services.add(customService) + setCredential(CustomServiceApiKeyById(requireNotNull(customService.id)), "TEST_API_KEY") + service().setModel( + FeatureType.CODE_COMPLETION, + customService.id, + ServiceType.CUSTOM_OPENAI + ) + + myFixture.configureByText( + "CompletionTest.txt", + FileUtil.getResourceContent("/codecompletions/code-completion-file.txt") + ) + myFixture.editor.caretModel.moveToVisualPosition(VisualPosition(3, 0)) + project.service().clear() + val prefix = """ + xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx + zzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzz + [INPUT] + p + """.trimIndent() + val suffix = """ + + [\INPUT] + zzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzz + xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx + """.trimIndent() + expectCustomOpenAI(BasicHttpExchange { request: RequestEntity -> + assertThat(request.uri.path).isEqualTo("/v1/completions") + assertThat(request.method).isEqualTo("POST") + assertThat(request.body) + .extracting("model", "prompt", "suffix", "max_tokens", "stream") + .containsExactly("custom-code-model", prefix, suffix, 128, false) + ResponseEntity( + jsonMapResponse( + "choices", + jsonArray(jsonMap("text", "ublic void main")) + ) + ) + }) + + myFixture.type('p') + + assertInlineSuggestion("Failed to display initial inline suggestion.") { + "ublic void main" == it + } + } + + fun `test code completion with Mistral provider`() { + useMistralService(FeatureType.CODE_COMPLETION) + val fixture = prepareStandardCompletionFixture() + expectMistral(BasicHttpExchange { request: RequestEntity -> + assertThat(request.uri.path).isEqualTo("/v1/fim/completions") + assertThat(request.method).isEqualTo("POST") + assertThat(request.body) + .extracting("model", "prompt", "suffix", "max_tokens", "stream") + .containsExactly("codestral-latest", fixture.prefix, fixture.suffix, 128, false) + ResponseEntity( + jsonMapResponse( + "choices", + jsonArray( + jsonMap( + "message", + jsonMap("content", "ublic void main") + ) + ) + ) + ) + }) + + myFixture.type('p') + + assertInlineSuggestion("Failed to display initial inline suggestion.") { + "ublic void main" == it + } + } + + fun `test code completion with Inception provider`() { + useInceptionService(FeatureType.CODE_COMPLETION) + val fixture = prepareStandardCompletionFixture() + expectInception(BasicHttpExchange { request: RequestEntity -> + assertThat(request.uri.path).isEqualTo("/v1/fim/completions") + assertThat(request.method).isEqualTo("POST") + assertThat(request.body) + .extracting("model", "prompt", "suffix", "stream") + .containsExactly("mercury-coder", fixture.prefix, fixture.suffix, false) + ResponseEntity( + jsonMapResponse( + "choices", + jsonArray(jsonMap("text", "ublic void main")) + ) + ) + }) + + myFixture.type('p') + + assertInlineSuggestion("Failed to display initial inline suggestion.") { + "ublic void main" == it + } + } + + fun `test code completion with llama cpp provider`() { + useLlamaService(codeCompletionsEnabled = true, role = FeatureType.CODE_COMPLETION) + val fixture = prepareStandardCompletionFixture() + val prompt = LlamaModel.findByHuggingFaceModel(HuggingFaceModel.CODE_LLAMA_7B_Q4) + .infillPromptTemplate + .buildPrompt(InfillRequest.Builder(fixture.prefix, fixture.suffix, 0).build()) + expectLlama(BasicHttpExchange { request: RequestEntity -> + assertThat(request.uri.path).isEqualTo("/completion") + assertThat(request.method).isEqualTo("POST") + assertThat(request.body) + .extracting("prompt", "stream", "n_predict") + .containsExactly(prompt, false, 128) + ResponseEntity( + jsonMapResponse( + "content", + "ublic void main" + ) + ) + }) + + myFixture.type('p') + + assertInlineSuggestion("Failed to display initial inline suggestion.") { + "ublic void main" == it + } + } + + fun `test code completion with ProxyAI provider`() { + useCodeGPTService(FeatureType.CODE_COMPLETION) + service().state.nextEditsEnabled = false + val fixture = prepareStandardCompletionFixture() + startProxyAiGrpcServer(fixture) + + myFixture.type('p') + + assertInlineSuggestion("Failed to display initial inline suggestion.") { + "ublic void main" == it + } + } + fun `_test apply inline suggestions without initial following text`() { useCodeGPTService(FeatureType.CODE_COMPLETION) service().state.nextEditsEnabled = false @@ -360,4 +544,62 @@ class CodeCompletionServiceTest : IntegrationTest() { 5 ) } + + private fun prepareStandardCompletionFixture(): CompletionFixture { + service().state.nextEditsEnabled = false + myFixture.configureByText( + "CompletionTest.txt", + FileUtil.getResourceContent("/codecompletions/code-completion-file.txt") + ) + myFixture.editor.caretModel.moveToVisualPosition(VisualPosition(3, 0)) + project.service().clear() + return CompletionFixture( + prefix = """ + xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx + zzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzz + [INPUT] + p + """.trimIndent(), + suffix = """ + + [\INPUT] + zzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzz + xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx + """.trimIndent() + ) + } + + private data class CompletionFixture( + val prefix: String, + val suffix: String + ) + + private fun startProxyAiGrpcServer(fixture: CompletionFixture) { + val service = object : CodeCompletionServiceImplGrpc.CodeCompletionServiceImplImplBase() { + override fun getCodeCompletion( + request: GrpcCodeCompletionRequest, + responseObserver: StreamObserver + ) { + assertThat(request.model).isEqualTo("mercury-coder") + assertThat(request.fileContent).contains(fixture.prefix) + assertThat(request.cursorPosition).isGreaterThan(0) + responseObserver.onNext( + PartialCodeCompletionResponse.newBuilder() + .setId("test-response-id") + .setPartialCompletion("ublic void main") + .build() + ) + responseObserver.onCompleted() + } + } + + proxyAiGrpcServer = ServerBuilder.forPort(0) + .addService(service) + .build() + .start() + + System.setProperty("proxyai.grpc.host", "127.0.0.1") + System.setProperty("proxyai.grpc.port", proxyAiGrpcServer!!.port.toString()) + System.setProperty("proxyai.grpc.plaintext", "true") + } } diff --git a/src/test/kotlin/ee/carlrobert/codegpt/codecompletions/psi/PythonContextFinderTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/codecompletions/psi/PythonContextFinderTest.kt deleted file mode 100644 index bbb9369c..00000000 --- a/src/test/kotlin/ee/carlrobert/codegpt/codecompletions/psi/PythonContextFinderTest.kt +++ /dev/null @@ -1,68 +0,0 @@ -package ee.carlrobert.codegpt.codecompletions.psi - -import com.intellij.openapi.editor.VisualPosition -import com.intellij.testFramework.fixtures.BasePlatformTestCase -import com.jetbrains.python.psi.PyClass -import com.jetbrains.python.psi.PyFile -import com.jetbrains.python.psi.PyFunction -import ee.carlrobert.codegpt.util.file.FileUtil.getResourceContent -import org.assertj.core.api.Assertions.assertThat -import org.junit.Ignore - -@Ignore -class PythonContextFinderTest : BasePlatformTestCase() { - - private val contextFinder = PythonContextFinder() - - fun testFindEnclosingElementMethod() { - val file = myFixture.configureByText( - "util.py", - getResourceContent("/codecompletions/psi/python.txt") - ) - val psiElement = - file.findElementAt(myFixture.editor.visualPositionToOffset(VisualPosition(11, 1))) - val enclosingElement = contextFinder.findEnclosingElement(psiElement!!) - assertThat(enclosingElement) - .isInstanceOf(PyFunction::class.java) - .extracting("name") - .isEqualTo("randomStrings") - } - - fun testFindEnclosingElementClass() { - val file = myFixture.configureByText( - "util.py", - getResourceContent("/codecompletions/psi/python.txt") - ) - val psiElement = - file.findElementAt(myFixture.editor.visualPositionToOffset(VisualPosition(6, 2))) - val contextFinder = contextFinder - val enclosingElement = contextFinder.findEnclosingElement(psiElement!!) - assertThat(enclosingElement) - .isInstanceOf(PyClass::class.java) - .extracting("name") - .isEqualTo("Util") - } - - fun testFindRelevantElements() { - val file = myFixture.configureByText( - "util.py", - getResourceContent("/codecompletions/psi/python.txt") - ) - val psiMethod = - (file as PyFile).topLevelClasses[0].methods[1] - val relevantElements = contextFinder.findRelevantElements(psiMethod, psiMethod) - assertThat(relevantElements) - .hasSize(7) - .extracting("name") - .containsExactly( - "test", - "int", - "List", - "str", - "randint", - "range", - "n", - ) - } - -} diff --git a/src/test/kotlin/ee/carlrobert/codegpt/completions/CompletionRequestProviderTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/completions/CompletionRequestProviderTest.kt index c2fabc48..1bee3228 100644 --- a/src/test/kotlin/ee/carlrobert/codegpt/completions/CompletionRequestProviderTest.kt +++ b/src/test/kotlin/ee/carlrobert/codegpt/completions/CompletionRequestProviderTest.kt @@ -1,22 +1,20 @@ package ee.carlrobert.codegpt.completions +import ai.koog.prompt.message.Message as KoogMessage import com.intellij.openapi.components.service import ee.carlrobert.codegpt.completions.factory.OpenAIRequestFactory import ee.carlrobert.codegpt.conversations.ConversationService import ee.carlrobert.codegpt.conversations.message.Message import ee.carlrobert.codegpt.settings.prompts.PersonaPromptDetailsState import ee.carlrobert.codegpt.settings.prompts.PromptsSettings -import ee.carlrobert.llm.client.openai.completion.OpenAIChatCompletionModel -import ee.carlrobert.codegpt.util.file.FileUtil.getResourceContent import org.assertj.core.api.Assertions.assertThat -import org.assertj.core.groups.Tuple -import org.junit.jupiter.api.Assertions.assertThrows +import org.assertj.core.api.Assertions.tuple import testsupport.IntegrationTest class CompletionRequestProviderTest : IntegrationTest() { fun testChatCompletionRequestWithSystemPromptOverride() { - useOpenAIService(OpenAIChatCompletionModel.GPT_4_O.code) + useOpenAIService("gpt-4o") val customPersona = PersonaPromptDetailsState().apply { id = 999L name = "Test Persona" @@ -32,24 +30,23 @@ class CompletionRequestProviderTest : IntegrationTest() { .builder(conversation, Message("TEST_CHAT_COMPLETION_PROMPT")) .build() - val request = OpenAIRequestFactory().createChatRequest(callParameters) + val prompt = OpenAIRequestFactory().createChatCompletionPrompt(callParameters) + val normalized = normalize(prompt.messages) - val guidelines = getResourceContent("/prompts/persona/psi-navigation-guidelines.txt") - val expectedSystem = "TEST_SYSTEM_PROMPT\n$guidelines" - assertThat(request.messages) + assertThat(normalized[0]["content"]).contains("TEST_SYSTEM_PROMPT") + assertThat(normalized.drop(1)) .extracting("role", "content") .containsExactly( - Tuple.tuple("system", expectedSystem), - Tuple.tuple("user", "TEST_PROMPT"), - Tuple.tuple("assistant", firstMessage.response), - Tuple.tuple("user", "TEST_PROMPT"), - Tuple.tuple("assistant", secondMessage.response), - Tuple.tuple("user", "TEST_CHAT_COMPLETION_PROMPT") + tuple("user", "TEST_PROMPT"), + tuple("assistant", firstMessage.response), + tuple("user", "TEST_PROMPT"), + tuple("assistant", secondMessage.response), + tuple("user", "TEST_CHAT_COMPLETION_PROMPT") ) } fun testChatCompletionRequestRetry() { - useOpenAIService(OpenAIChatCompletionModel.GPT_4_O.code) + useOpenAIService("gpt-4o") val customPersona = PersonaPromptDetailsState().apply { id = 999L name = "Test Persona" @@ -65,17 +62,16 @@ class CompletionRequestProviderTest : IntegrationTest() { .retry(true) .build() - val request = OpenAIRequestFactory().createChatRequest(callParameters) + val prompt = OpenAIRequestFactory().createChatCompletionPrompt(callParameters) + val normalized = normalize(prompt.messages) - val guidelines = getResourceContent("/prompts/persona/psi-navigation-guidelines.txt") - val expectedSystem = "TEST_SYSTEM_PROMPT\n$guidelines" - assertThat(request.messages) + assertThat(normalized[0]["content"]).contains("TEST_SYSTEM_PROMPT") + assertThat(normalized.drop(1)) .extracting("role", "content") .containsExactly( - Tuple.tuple("system", expectedSystem), - Tuple.tuple("user", "FIRST_TEST_PROMPT"), - Tuple.tuple("assistant", firstMessage.response), - Tuple.tuple("user", "SECOND_TEST_PROMPT") + tuple("user", "FIRST_TEST_PROMPT"), + tuple("assistant", firstMessage.response), + tuple("user", "SECOND_TEST_PROMPT") ) } @@ -88,4 +84,15 @@ class CompletionRequestProviderTest : IntegrationTest() { message.response = "zz".repeat((tokenSize) - 6 - 7) return message } + + private fun normalize(messages: List): List> { + return messages.mapNotNull { msg -> + when (msg) { + is KoogMessage.System -> mapOf("role" to "system", "content" to msg.content) + is KoogMessage.User -> mapOf("role" to "user", "content" to msg.content) + is KoogMessage.Assistant -> mapOf("role" to "assistant", "content" to msg.content) + else -> null + } + } + } } diff --git a/src/test/kotlin/ee/carlrobert/codegpt/completions/CompletionRequestServiceKoogIntegrationTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/completions/CompletionRequestServiceKoogIntegrationTest.kt new file mode 100644 index 00000000..a2e2cbb3 --- /dev/null +++ b/src/test/kotlin/ee/carlrobert/codegpt/completions/CompletionRequestServiceKoogIntegrationTest.kt @@ -0,0 +1,285 @@ +package ee.carlrobert.codegpt.completions + +import com.intellij.openapi.components.service +import ee.carlrobert.codegpt.conversations.ConversationService +import ee.carlrobert.codegpt.conversations.message.Message +import ee.carlrobert.codegpt.settings.models.ModelSettings +import ee.carlrobert.codegpt.settings.service.FeatureType +import ee.carlrobert.codegpt.settings.service.ServiceType +import org.assertj.core.api.Assertions.assertThat +import testsupport.IntegrationTest +import testsupport.http.RequestEntity +import testsupport.http.ResponseEntity +import testsupport.http.exchange.BasicHttpExchange +import testsupport.http.exchange.StreamHttpExchange +import testsupport.json.JSONUtil.e +import testsupport.json.JSONUtil.jsonArray +import testsupport.json.JSONUtil.jsonMap +import testsupport.json.JSONUtil.jsonMapResponse + +class CompletionServiceKoogIntegrationTest : IntegrationTest() { + + fun testLookupCompletionUsesKoogPromptExecution() { + useOpenAIService("gpt-4o", FeatureType.LOOKUP) + useStableOpenAIModel(FeatureType.LOOKUP) + expectOpenAISyncResponse("lookup_1,lookup_2") { promptText -> + assertThat(promptText).contains("generate names for InvoiceService") + } + + val result = service() + .getLookupCompletion(LookupCompletionParameters("generate names for InvoiceService")) + + assertThat(result).isEqualTo("lookup_1,lookup_2") + } + + fun testCommitMessageFlowUsesKoogStreamingExecution() { + useOpenAIService("gpt-4o", FeatureType.COMMIT_MESSAGE) + useStableOpenAIModel(FeatureType.COMMIT_MESSAGE) + expectOpenAIStreamingHello { promptText -> + assertThat(promptText).contains("Generate a concise commit message") + assertThat(promptText).contains("diff --git a/A.kt b/A.kt") + } + + val listener = RecordingListener() + + service().getCommitMessage( + CommitMessageCompletionParameters( + gitDiff = "diff --git a/A.kt b/A.kt", + systemPrompt = "Generate a concise commit message" + ), + listener + ) + waitExpecting { listener.completed == "Hello!" } + + assertThat(listener.opened).isTrue() + assertThat(listener.error).isNull() + assertThat(listener.chunks).containsExactly("Hel", "lo", "!") + } + + fun testInlineEditFlowUsesKoogStreamingExecution() { + useOpenAIService("gpt-4o", FeatureType.INLINE_EDIT) + useStableOpenAIModel(FeatureType.INLINE_EDIT) + expectOpenAIStreamingHello { promptText -> + assertThat(promptText).contains("Rename this function") + assertThat(promptText).contains("Sure, renaming now") + assertThat(promptText).contains("Implement.") + } + + val conversation = ConversationService.getInstance().startConversation(project) + conversation.addMessage(Message("Rename this function", "Sure, renaming now")) + + val listener = RecordingListener() + + service().getInlineEditCompletion( + InlineEditCompletionParameters( + selectedText = "fun oldName() = 1", + fileExtension = "kt", + projectBasePath = project.basePath, + conversation = conversation + ), + listener + ) + waitExpecting { listener.completed == "Hello!" } + + assertThat(listener.opened).isTrue() + assertThat(listener.error).isNull() + assertThat(listener.chunks).containsExactly("Hel", "lo", "!") + } + + fun testAutoApplyFlowUsesKoogStreamingExecution() { + useOpenAIService("gpt-4o", FeatureType.AUTO_APPLY) + useStableOpenAIModel(FeatureType.AUTO_APPLY) + expectOpenAIStreamingHello { promptText -> + assertThat(promptText).contains("fun value() = 1") + assertThat(promptText).contains("fun value() = 2") + } + + val destination = myFixture.configureByText( + "Destination.kt", + "fun value() = 1" + ).virtualFile + val listener = RecordingListener() + + service().autoApply( + AutoApplyParameters( + source = "fun value() = 2", + destination = destination + ), + listener + ) + waitExpecting { listener.completed == "Hello!" } + + assertThat(listener.opened).isTrue() + assertThat(listener.error).isNull() + assertThat(listener.chunks).containsExactly("Hel", "lo", "!") + } + + fun testAutoApplySyncFlowUsesKoogExecution() { + useOpenAIService("gpt-4o", FeatureType.AUTO_APPLY) + useStableOpenAIModel(FeatureType.AUTO_APPLY) + expectOpenAISyncResponse("fun value() = 2") { promptText -> + assertThat(promptText).contains("fun value() = 1") + assertThat(promptText).contains("fun value() = 2") + } + + val destination = myFixture.configureByText( + "SyncDestination.kt", + "fun value() = 1" + ).virtualFile + + val result = service().autoApply( + AutoApplyParameters( + source = "fun value() = 2", + destination = destination + ) + ) + + assertThat(result).isEqualTo("fun value() = 2") + } + + private class RecordingListener : CompletionStreamEventListener { + val chunks = mutableListOf() + var completed: String? = null + var opened: Boolean = false + var error: Throwable? = null + + override fun onOpen() { + opened = true + } + + override fun onMessage(message: String) { + if (message.isNotEmpty()) { + chunks.add(message) + } + } + + override fun onComplete(messageBuilder: StringBuilder) { + completed = messageBuilder.toString() + } + + override fun onError(error: CompletionError, ex: Throwable) { + this.error = ex + } + + override fun onCancelled(messageBuilder: StringBuilder) { + completed = messageBuilder.toString() + } + } + + private fun expectOpenAIStreamingHello(assertPrompt: (String) -> Unit) { + expectOpenAI(StreamHttpExchange { request: RequestEntity -> + assertThat(request.uri.path).isEqualTo("/v1/responses") + assertThat(request.method).isEqualTo("POST") + assertPrompt(extractPromptText(request)) + listOf( + streamingChunk("Hel", 1), + streamingChunk("lo", 2), + streamingChunk("!", 3) + ) + }) + } + + private fun expectOpenAISyncResponse( + responseText: String, + assertPrompt: (String) -> Unit + ) { + expectOpenAI(BasicHttpExchange { request: RequestEntity -> + assertThat(request.uri.path).isEqualTo("/v1/responses") + assertThat(request.method).isEqualTo("POST") + assertPrompt(extractPromptText(request)) + ResponseEntity( + jsonMapResponse( + e("id", "resp-test"), + e("object", "response"), + e("created_at", 1), + e("model", "gpt-5-mini"), + e( + "output", + jsonArray( + jsonMap( + e("type", "message"), + e("id", "msg_1"), + e("role", "assistant"), + e("status", "completed"), + e( + "content", + jsonArray( + jsonMap( + e("type", "output_text"), + e("text", responseText), + e("annotations", jsonArray()) + ) + ) + ) + ) + ) + ), + e("parallel_tool_calls", true), + e("status", "completed"), + e("text", jsonMap()), + e( + "usage", + jsonMap( + e("input_tokens", 1), + e("input_tokens_details", jsonMap("cached_tokens", 0)), + e("output_tokens", 1), + e("output_tokens_details", jsonMap("reasoning_tokens", 0)), + e("total_tokens", 2) + ) + ) + ) + ) + }) + } + + private fun streamingChunk(content: String, sequenceNumber: Int): String { + return jsonMapResponse( + e("type", "response.output_text.delta"), + e("item_id", "msg_1"), + e("output_index", 0), + e("content_index", 0), + e("delta", content), + e("sequence_number", sequenceNumber) + ) + } + + private fun extractPromptText(request: RequestEntity): String { + val messages = request.body["messages"] as? List<*> + if (messages != null) { + return messages.joinToString("\n") { message -> + val content = (message as? Map<*, *>)?.get("content") + when (content) { + is String -> content + is List<*> -> content.joinToString("\n") { part -> + val partMap = part as? Map<*, *> + (partMap?.get("text") as? String) ?: partMap.toString() + } + + null -> "" + else -> content.toString() + } + } + } + + val input = request.body["input"] as? List<*> ?: return "" + return input.joinToString("\n") { item -> + val content = (item as? Map<*, *>)?.get("content") + when (content) { + is String -> content + is List<*> -> content.joinToString("\n") { part -> + val partMap = part as? Map<*, *> + (partMap?.get("text") as? String) + ?: (partMap?.get("value") as? String) + ?: partMap.toString() + } + + null -> "" + else -> content.toString() + } + } + } + + private fun useStableOpenAIModel(featureType: FeatureType) { + service().setModel(featureType, "gpt-5-mini", ServiceType.OPENAI) + } +} diff --git a/src/test/kotlin/ee/carlrobert/codegpt/completions/DefaultToolwindowChatCompletionRequestHandlerTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/completions/DefaultToolwindowChatCompletionRequestHandlerTest.kt deleted file mode 100644 index 35e3fbef..00000000 --- a/src/test/kotlin/ee/carlrobert/codegpt/completions/DefaultToolwindowChatCompletionRequestHandlerTest.kt +++ /dev/null @@ -1,268 +0,0 @@ -package ee.carlrobert.codegpt.completions - -import com.intellij.openapi.components.service -import ee.carlrobert.codegpt.completions.llama.PromptTemplate.LLAMA -import ee.carlrobert.codegpt.conversations.ConversationService -import ee.carlrobert.codegpt.conversations.message.Message -import ee.carlrobert.codegpt.settings.configuration.ConfigurationSettings -import ee.carlrobert.codegpt.settings.prompts.PersonaPromptDetailsState -import ee.carlrobert.codegpt.settings.prompts.PromptsSettings -import ee.carlrobert.llm.client.http.RequestEntity -import ee.carlrobert.llm.client.http.exchange.NdJsonStreamHttpExchange -import ee.carlrobert.llm.client.http.exchange.StreamHttpExchange -import ee.carlrobert.llm.client.util.JSONUtil.* -import org.apache.http.HttpHeaders -import org.assertj.core.api.Assertions.assertThat -import testsupport.IntegrationTest -import ee.carlrobert.codegpt.util.file.FileUtil.getResourceContent - -class DefaultToolwindowChatCompletionRequestHandlerTest : IntegrationTest() { - - fun testOpenAIChatCompletionCall() { - useOpenAIService() - val customPersona = PersonaPromptDetailsState().apply { - id = 999L - name = "Test Persona" - instructions = "TEST_SYSTEM_PROMPT" - } - service().state.personas.selectedPersona = customPersona - val message = Message("TEST_PROMPT") - val conversation = ConversationService.getInstance().startConversation(project) - expectOpenAI(StreamHttpExchange { request: RequestEntity -> - assertThat(request.uri.path).isEqualTo("/v1/chat/completions") - assertThat(request.method).isEqualTo("POST") - assertThat(request.headers[HttpHeaders.AUTHORIZATION]!![0]).isEqualTo("Bearer TEST_API_KEY") - val guidelines = getResourceContent("/prompts/persona/psi-navigation-guidelines.txt") - val expectedSystem = "TEST_SYSTEM_PROMPT\n$guidelines" - assertThat(request.body) - .extracting( - "model", - "messages" - ) - .containsExactly( - "gpt-4o", - listOf( - mapOf("role" to "system", "content" to expectedSystem), - mapOf("role" to "user", "content" to "TEST_PROMPT") - ) - ) - listOf( - jsonMapResponse( - "choices", - jsonArray(jsonMap("delta", jsonMap("role", "assistant"))) - ), - jsonMapResponse("choices", jsonArray(jsonMap("delta", jsonMap("content", "Hel")))), - jsonMapResponse("choices", jsonArray(jsonMap("delta", jsonMap("content", "lo")))), - jsonMapResponse("choices", jsonArray(jsonMap("delta", jsonMap("content", "!")))) - ) - }) - val requestHandler = - ToolwindowChatCompletionRequestHandler(project, getRequestEventListener(message), null) - - requestHandler.call(ChatCompletionParameters.builder(conversation, message).build()) - - waitExpecting { "Hello!" == message.response } - } - - fun testLlamaChatCompletionCall() { - useLlamaService() - service().state.maxTokens = 99 - val customPersona = PersonaPromptDetailsState().apply { - id = 999L - name = "Test Persona" - instructions = "TEST_SYSTEM_PROMPT" - } - service().state.personas.selectedPersona = customPersona - val message = Message("TEST_PROMPT") - val conversation = ConversationService.getInstance().startConversation(project) - conversation.addMessage(Message("Ping", "Pong")) - expectLlama(StreamHttpExchange { request: RequestEntity -> - assertThat(request.uri.path).isEqualTo("/v1/chat/completions") - val guidelines = getResourceContent("/prompts/persona/psi-navigation-guidelines.txt") - val expectedSystem = "TEST_SYSTEM_PROMPT\n$guidelines" - assertThat(request.body) - .extracting( - "model", - "messages" - ) - .containsExactly( - HuggingFaceModel.CODE_LLAMA_7B_Q4.code, - listOf( - mapOf("role" to "system", "content" to expectedSystem), - mapOf("role" to "user", "content" to "Ping"), - mapOf("role" to "assistant", "content" to "Pong"), - mapOf("role" to "user", "content" to "TEST_PROMPT") - ) - ) - listOf( - jsonMapResponse( - "choices", - jsonArray(jsonMap("delta", jsonMap("role", "assistant"))) - ), - jsonMapResponse("choices", jsonArray(jsonMap("delta", jsonMap("content", "Hel")))), - jsonMapResponse("choices", jsonArray(jsonMap("delta", jsonMap("content", "lo")))), - jsonMapResponse("choices", jsonArray(jsonMap("delta", jsonMap("content", "!")))) - ) - }) - val requestHandler = - ToolwindowChatCompletionRequestHandler(project, getRequestEventListener(message), null) - - requestHandler.call(ChatCompletionParameters.builder(conversation, message).build()) - - waitExpecting { "Hello!" == message.response } - } - - fun testOllamaChatCompletionCall() { - useOllamaService() - service().state.maxTokens = 99 - val customPersona = PersonaPromptDetailsState().apply { - id = 999L - name = "Test Persona" - instructions = "TEST_SYSTEM_PROMPT" - } - service().state.personas.selectedPersona = customPersona - val message = Message("TEST_PROMPT") - val conversation = ConversationService.getInstance().startConversation(project) - expectOllama(NdJsonStreamHttpExchange { request: RequestEntity -> - assertThat(request.uri.path).isEqualTo("/v1/chat/completions") - assertThat(request.method).isEqualTo("POST") - assertThat(request.headers[HttpHeaders.AUTHORIZATION]!![0]).isEqualTo("Bearer TEST_API_KEY") - val guidelines = getResourceContent("/prompts/persona/psi-navigation-guidelines.txt") - val expectedSystem = "TEST_SYSTEM_PROMPT\n$guidelines" - assertThat(request.body) - .extracting( - "model", - "messages" - ) - .containsExactly( - HuggingFaceModel.LLAMA_3_8B_Q6_K.code, - listOf( - mapOf("role" to "system", "content" to expectedSystem), - mapOf("role" to "user", "content" to "TEST_PROMPT") - ) - ) - listOf( - jsonMapResponse( - "choices", - jsonArray(jsonMap("delta", jsonMap("role", "assistant"))) - ), - jsonMapResponse("choices", jsonArray(jsonMap("delta", jsonMap("content", "Hel")))), - jsonMapResponse("choices", jsonArray(jsonMap("delta", jsonMap("content", "lo")))), - jsonMapResponse("choices", jsonArray(jsonMap("delta", jsonMap("content", "!")))) - ) - }) - val requestHandler = - ToolwindowChatCompletionRequestHandler(project, getRequestEventListener(message), null) - - requestHandler.call(ChatCompletionParameters.builder(conversation, message).build()) - - waitExpecting { "Hello!" == message.response } - } - - fun testGoogleChatCompletionCall() { - useGoogleService() - service().state - .chatCompletionSettings.clickableLinksEnabled = true - val customPersona = PersonaPromptDetailsState().apply { - id = 999L - name = "Test Persona" - instructions = "TEST_SYSTEM_PROMPT" - } - service().state.personas.selectedPersona = customPersona - val message = Message("TEST_PROMPT") - val conversation = ConversationService.getInstance().startConversation(project) - expectGoogle(StreamHttpExchange { request: RequestEntity -> - assertThat(request.uri.path).isEqualTo("/v1/models/gemini-2.0-flash:streamGenerateContent") - assertThat(request.method).isEqualTo("POST") - assertThat(request.uri.query).isEqualTo("key=TEST_API_KEY&alt=sse") - val guidelines = getResourceContent("/prompts/persona/psi-navigation-guidelines.txt") - val expectedSystem = "TEST_SYSTEM_PROMPT\n$guidelines" - assertThat(request.body) - .extracting("contents", "systemInstruction") - .containsExactly( - listOf( - mapOf("parts" to listOf(mapOf("text" to "TEST_PROMPT")), "role" to "user"), - ), - mapOf("parts" to listOf(mapOf("text" to expectedSystem))) - ) - listOf( - jsonMapResponse( - "candidates", - jsonArray( - jsonMap( - "content", - jsonMap("parts", jsonArray(jsonMap("text", "Hello"))) - ) - ) - ), - jsonMapResponse( - "candidates", - jsonArray(jsonMap("content", jsonMap("parts", jsonArray(jsonMap("text", "!"))))) - ) - ) - }) - val requestHandler = - ToolwindowChatCompletionRequestHandler(project, getRequestEventListener(message), null) - - requestHandler.call(ChatCompletionParameters.builder(conversation, message).build()) - - waitExpecting { "Hello!" == message.response } - } - - fun testCodeGPTServiceChatCompletionCall() { - useCodeGPTService() - val customPersona = PersonaPromptDetailsState().apply { - id = 999L - name = "Test Persona" - instructions = "TEST_SYSTEM_PROMPT" - } - service().state.personas.selectedPersona = customPersona - val message = Message("TEST_PROMPT") - val conversation = ConversationService.getInstance().startConversation(project) - expectCodeGPT(StreamHttpExchange { request: RequestEntity -> - assertThat(request.uri.path).isEqualTo("/v1/chat/completions") - assertThat(request.method).isEqualTo("POST") - assertThat(request.headers[HttpHeaders.AUTHORIZATION]!![0]).isEqualTo("Bearer TEST_API_KEY") - val guidelines = getResourceContent("/prompts/persona/psi-navigation-guidelines.txt") - val expectedSystem = "TEST_SYSTEM_PROMPT\n$guidelines" - assertThat(request.body) - .extracting( - "model", - "messages" - ) - .containsExactly( - "gpt-5-mini", - listOf( - mapOf("role" to "user", "content" to expectedSystem), - mapOf("role" to "user", "content" to "TEST_PROMPT") - ) - ) - listOf( - jsonMapResponse( - "choices", - jsonArray(jsonMap("delta", jsonMap("role", "assistant"))) - ), - jsonMapResponse("choices", jsonArray(jsonMap("delta", jsonMap("content", "Hel")))), - jsonMapResponse("choices", jsonArray(jsonMap("delta", jsonMap("content", "lo")))), - jsonMapResponse("choices", jsonArray(jsonMap("delta", jsonMap("content", "!")))) - ) - }) - val requestHandler = - ToolwindowChatCompletionRequestHandler(project, getRequestEventListener(message), null) - - requestHandler.call(ChatCompletionParameters.builder(conversation, message).build()) - - waitExpecting { "Hello!" == message.response } - } - - private fun getRequestEventListener(message: Message): CompletionResponseEventListener { - return object : CompletionResponseEventListener { - override fun handleCompleted( - fullMessage: String, - callParameters: ChatCompletionParameters - ) { - message.response = fullMessage - } - } - } -} diff --git a/src/test/kotlin/ee/carlrobert/codegpt/completions/OpenAIRequestFactoryIntegrationTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/completions/OpenAIRequestFactoryIntegrationTest.kt index 198244fd..7a9eabbb 100644 --- a/src/test/kotlin/ee/carlrobert/codegpt/completions/OpenAIRequestFactoryIntegrationTest.kt +++ b/src/test/kotlin/ee/carlrobert/codegpt/completions/OpenAIRequestFactoryIntegrationTest.kt @@ -1,5 +1,6 @@ package ee.carlrobert.codegpt.completions +import ai.koog.prompt.message.Message as KoogMessage import com.intellij.openapi.components.service import ee.carlrobert.codegpt.ReferencedFile import ee.carlrobert.codegpt.completions.factory.OpenAIRequestFactory @@ -10,13 +11,7 @@ import ee.carlrobert.codegpt.settings.configuration.ChatMode import ee.carlrobert.codegpt.settings.prompts.* import ee.carlrobert.codegpt.settings.service.FeatureType import ee.carlrobert.codegpt.util.file.FileUtil.getResourceContent -import ee.carlrobert.llm.client.openai.completion.OpenAIChatCompletionModel -import ee.carlrobert.llm.client.openai.completion.request.OpenAIChatCompletionAssistantMessage -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.OpenAIChatCompletionToolMessage -import ee.carlrobert.llm.client.openai.completion.response.ToolCall -import ee.carlrobert.llm.client.openai.completion.response.ToolFunctionResponse + import org.assertj.core.api.Assertions.assertThat import testsupport.IntegrationTest import java.io.File @@ -24,7 +19,7 @@ import java.io.File class OpenAIRequestFactoryIntegrationTest : IntegrationTest() { fun testDefaultPersonaUsesEditModePromptWhenEnabled() { - useOpenAIService(OpenAIChatCompletionModel.GPT_4_O.code) + useOpenAIService("gpt-4o") service().state.personas.selectedPersona = PersonasState.DEFAULT_PERSONA val conversation = ConversationService.getInstance().startConversation(project) val message = Message("Please refactor this code") @@ -33,11 +28,10 @@ class OpenAIRequestFactoryIntegrationTest : IntegrationTest() { .chatMode(ChatMode.EDIT) .build() - val request = OpenAIRequestFactory().createChatRequest(callParameters) + val prompt = OpenAIRequestFactory().createChatCompletionPrompt(callParameters) - val systemMessages = request.messages - .filterIsInstance() - .filter { it.role == "system" } + val systemMessages = prompt.messages + .filterIsInstance() .map { it.content } assertThat(systemMessages).isNotEmpty() val systemContent = systemMessages.first() @@ -50,54 +44,49 @@ class OpenAIRequestFactoryIntegrationTest : IntegrationTest() { } fun testToolCallRequestBuildsAssistantToolCallsMessage() { - useOpenAIService(OpenAIChatCompletionModel.GPT_4_O.code) + useOpenAIService("gpt-4o") service().state.personas.selectedPersona = PersonasState.DEFAULT_PERSONA val conversation = ConversationService.getInstance().startConversation(project) val message = Message("Run tool").apply { response = "Calling tool..." toolCalls = - listOf(ToolCall(null, "tc_1", "function", ToolFunctionResponse("search", "{}"))) + listOf(ChatToolCall(null, "tc_1", "function", ChatToolFunction("search", "{}"))) } val callParameters = ChatCompletionParameters.builder(conversation, message) .requestType(RequestType.TOOL_CALL_REQUEST).build() val filtered = service().getFilteredPersonaPrompt(ChatMode.ASK) .addProjectPath() - val expectedSystem = service().applyClickableLinks(filtered) val expected = listOf( - mapOf("role" to "system", "content" to expectedSystem), + mapOf("role" to "system"), mapOf("role" to "user", "content" to "Run tool"), mapOf( "role" to "assistant", - "content" to "Calling tool...", - "toolCalls" to listOf( - mapOf( - "id" to "tc_1", - "type" to "function", - "fn" to "search", - "args" to "{}" - ) - ) - ) + "content" to "Calling tool..." + ), + mapOf("role" to "tool_call", "id" to "tc_1", "tool" to "search", "args" to "{}") ) - val request = OpenAIRequestFactory().createChatRequest(callParameters) + val prompt = OpenAIRequestFactory().createChatCompletionPrompt(callParameters) - assertThat(normalize(request.messages)).isEqualTo(expected) + val actual = normalize(prompt.messages) + assertThat(actual.drop(1)).isEqualTo(expected.drop(1)) + assertThat(actual.first()["role"]).isEqualTo("system") + assertThat(actual.first()["content"].toString()).contains("JetBrains Navigation Links (MANDATORY)") } fun testToolCallContinuationIncludesResultsAndHistory() { - useOpenAIService(OpenAIChatCompletionModel.GPT_4_O.code) + useOpenAIService("gpt-4o") service().state.personas.selectedPersona = PersonasState.DEFAULT_PERSONA val conversation = ConversationService.getInstance().startConversation(project) val prev = Message("Prev run").apply { response = "Prev call" toolCalls = listOf( - ToolCall( + ChatToolCall( null, "tc_prev", "function", - ToolFunctionResponse("list", "{\"q\":1}") + ChatToolFunction("list", "{\"q\":1}") ) ) addToolCallResult("tc_prev", "prev_result") @@ -106,60 +95,46 @@ class OpenAIRequestFactoryIntegrationTest : IntegrationTest() { val current = Message("Run again").apply { response = "Curr call" toolCalls = listOf( - ToolCall( + ChatToolCall( null, "tc_curr", "function", - ToolFunctionResponse("list", "{\"q\":2}") + ChatToolFunction("list", "{\"q\":2}") ) ) addToolCallResult("tc_curr", "curr_result") } val callParameters = ChatCompletionParameters.builder(conversation, current) .requestType(RequestType.TOOL_CALL_CONTINUATION).build() - val filtered = service().getFilteredPersonaPrompt(ChatMode.ASK) - .addProjectPath() - val expectedSystem = service().applyClickableLinks(filtered) val expected = listOf( - mapOf("role" to "system", "content" to expectedSystem), + mapOf("role" to "system"), mapOf("role" to "user", "content" to "Prev run"), mapOf( "role" to "assistant", - "content" to "Prev call", - "toolCalls" to listOf( - mapOf( - "id" to "tc_prev", - "type" to "function", - "fn" to "list", - "args" to "{\"q\":1}" - ) - ) + "content" to "Prev call" ), - mapOf("role" to "tool", "callId" to "tc_prev", "content" to "prev_result"), + mapOf("role" to "tool_call", "id" to "tc_prev", "tool" to "list", "args" to "{\"q\":1}"), + mapOf("role" to "tool_result", "id" to "tc_prev", "tool" to "list", "content" to "prev_result"), mapOf("role" to "user", "content" to "Run again"), mapOf( "role" to "assistant", - "content" to "Curr call", - "toolCalls" to listOf( - mapOf( - "id" to "tc_curr", - "type" to "function", - "fn" to "list", - "args" to "{\"q\":2}" - ) - ) + "content" to "Curr call" ), - mapOf("role" to "tool", "callId" to "tc_curr", "content" to "curr_result"), + mapOf("role" to "tool_call", "id" to "tc_curr", "tool" to "list", "args" to "{\"q\":2}"), + mapOf("role" to "tool_result", "id" to "tc_curr", "tool" to "list", "content" to "curr_result"), mapOf("role" to "assistant", "content" to "Curr call") ) - val request = OpenAIRequestFactory().createChatRequest(callParameters) + val prompt = OpenAIRequestFactory().createChatCompletionPrompt(callParameters) - assertThat(normalize(request.messages)).isEqualTo(expected) + val actual = normalize(prompt.messages) + assertThat(actual.drop(1)).isEqualTo(expected.drop(1)) + assertThat(actual.first()["role"]).isEqualTo("system") + assertThat(actual.first()["content"].toString()).contains("JetBrains Navigation Links (MANDATORY)") } fun testChatRequestEmbedsReferencedFilesInSystemMessage() { - useOpenAIService(OpenAIChatCompletionModel.GPT_4_O.code) + useOpenAIService("gpt-4o") service().state.personas.selectedPersona = PersonasState.DEFAULT_PERSONA val conversation = ConversationService.getInstance().startConversation(project) val message = Message("What does the code do?") @@ -172,28 +147,22 @@ class OpenAIRequestFactoryIntegrationTest : IntegrationTest() { val filtered = service().getFilteredPersonaPrompt(ChatMode.ASK) .addProjectPath() - val withGuidelines = service().applyClickableLinks(filtered) - val filesBlock = "\n\n\n" + CompletionRequestUtil.formatCode( - "class A {}", - "/path/A.java" - ) + "\n" + CompletionRequestUtil.formatCode( - "class B", - "/path/B.kt" - ) + "\n" - val expectedSystem = withGuidelines + filesBlock + val expectedSystem = service().applyClickableLinks(filtered) - val request = OpenAIRequestFactory().createChatRequest(callParameters) + val prompt = OpenAIRequestFactory().createChatCompletionPrompt(callParameters) - val msgs = request.messages + val msgs = prompt.messages assertThat(msgs).hasSize(2) - assertThat((msgs[0] as OpenAIChatCompletionStandardMessage).content).isEqualTo( - expectedSystem - ) - assertThat((msgs[1] as OpenAIChatCompletionStandardMessage).content).isEqualTo("What does the code do?") + assertThat((msgs[0] as KoogMessage.System).content).contains(expectedSystem.trim()) + assertThat((msgs[1] as KoogMessage.User).content).contains("What does the code do?") + assertThat((msgs[1] as KoogMessage.User).content).contains("/path/A.java") + assertThat((msgs[1] as KoogMessage.User).content).contains("class A {}") + assertThat((msgs[1] as KoogMessage.User).content).contains("/path/B.kt") + assertThat((msgs[1] as KoogMessage.User).content).contains("class B") } fun testDefaultPersonaIsFilteredInAskMode() { - useOpenAIService(OpenAIChatCompletionModel.GPT_4_O.code) + useOpenAIService("gpt-4o") service().state.personas.selectedPersona = PersonasState.DEFAULT_PERSONA val conversation = ConversationService.getInstance().startConversation(project) val message = Message("Please refactor this code") @@ -204,16 +173,14 @@ class OpenAIRequestFactoryIntegrationTest : IntegrationTest() { .addProjectPath() val expectedSystem = service().applyClickableLinks(filtered) - val request = OpenAIRequestFactory().createChatRequest(callParameters) + val prompt = OpenAIRequestFactory().createChatCompletionPrompt(callParameters) - val systemMessages = - request.messages.filterIsInstance() - .filter { it.role == "system" }.map { it.content } - assertThat(systemMessages.first()).isEqualTo(expectedSystem) + val systemMessages = prompt.messages.filterIsInstance().map { it.content } + assertThat(systemMessages.first()).contains(expectedSystem.trim()) } fun testChatRequestUsesFilteredPersonaPromptInAskMode() { - useOpenAIService(OpenAIChatCompletionModel.GPT_4_O.code) + useOpenAIService("gpt-4o") val personaPromptWithSearchReplace = """ You are a helpful assistant. For refactoring or editing an existing file, always generate a SEARCH/REPLACE block. @@ -242,11 +209,10 @@ class OpenAIRequestFactoryIntegrationTest : IntegrationTest() { .chatMode(ChatMode.ASK) .build() - val request = OpenAIRequestFactory().createChatRequest(callParameters) + val prompt = OpenAIRequestFactory().createChatCompletionPrompt(callParameters) - val systemMessages = request.messages - .filterIsInstance() - .filter { it.role == "system" } + val systemMessages = prompt.messages + .filterIsInstance() .map { it.content } assertThat(systemMessages).isNotEmpty() val systemContent = systemMessages.first() @@ -257,13 +223,16 @@ class OpenAIRequestFactoryIntegrationTest : IntegrationTest() { } fun testChatRequestKeepsOriginalPersonaPromptInEditMode() { - useOpenAIService(OpenAIChatCompletionModel.GPT_4_O.code) + useOpenAIService("gpt-4o") val personaPromptWithSearchReplace = """ You are a helpful assistant. For refactoring or editing an existing file, always generate a SEARCH/REPLACE block. """.trimIndent() - service().state.personas.selectedPersona.instructions = - personaPromptWithSearchReplace + service().state.personas.selectedPersona = PersonaPromptDetailsState().apply { + id = 999L + name = "Custom Test Persona" + instructions = personaPromptWithSearchReplace + } val conversation = ConversationService.getInstance().startConversation(project) val message = Message("Please refactor this code") val callParameters = ChatCompletionParameters @@ -271,11 +240,10 @@ class OpenAIRequestFactoryIntegrationTest : IntegrationTest() { .chatMode(ChatMode.EDIT) .build() - val request = OpenAIRequestFactory().createChatRequest(callParameters) + val prompt = OpenAIRequestFactory().createChatCompletionPrompt(callParameters) - val systemMessages = request.messages - .filterIsInstance() - .filter { it.role == "system" } + val systemMessages = prompt.messages + .filterIsInstance() .map { it.content } assertThat(systemMessages).isNotEmpty() val systemContent = systemMessages.first() @@ -288,7 +256,7 @@ class OpenAIRequestFactoryIntegrationTest : IntegrationTest() { } fun testInlineEditSingleRequestNoHistory() { - useOpenAIService(OpenAIChatCompletionModel.GPT_4_1.code, FeatureType.INLINE_EDIT) + useOpenAIService("gpt-4.1", FeatureType.INLINE_EDIT) val testFileContent = getResourceContent("/inline/TestClass.java") val tempFile = File.createTempFile("TestClass", ".java") tempFile.writeText(testFileContent) @@ -305,10 +273,9 @@ class OpenAIRequestFactoryIntegrationTest : IntegrationTest() { diagnosticsInfo = null ) - val request = OpenAIRequestFactory().createInlineEditRequest(parameters) + val prompt = OpenAIRequestFactory().createInlineEditPrompt(parameters) - val systemMessage = request.messages[0] as OpenAIChatCompletionStandardMessage - assertThat(systemMessage.role).isEqualTo("system") + val systemMessage = prompt.messages[0] as KoogMessage.System val expectedSystem = buildExpectedInlineEditSystemPrompt( language = "java", filePath = tempFile.absolutePath, @@ -320,13 +287,12 @@ class OpenAIRequestFactoryIntegrationTest : IntegrationTest() { diagnosticsInfo = null, ) assertThat(systemMessage.content).isEqualTo(expectedSystem) - val userMessage = request.messages[1] as OpenAIChatCompletionStandardMessage - assertThat(userMessage.role).isEqualTo("user") + val userMessage = prompt.messages[1] as KoogMessage.User assertThat(userMessage.content).isEqualTo("Implement.") } fun testInlineEditFollowUpWithHistory() { - useOpenAIService(OpenAIChatCompletionModel.GPT_4_1.code, FeatureType.INLINE_EDIT) + useOpenAIService("gpt-4.1", FeatureType.INLINE_EDIT) val testFileContent = getResourceContent("/inline/TestClass.java") val tempFile = File.createTempFile("TestClass", ".java") tempFile.writeText(testFileContent) @@ -375,10 +341,9 @@ class OpenAIRequestFactoryIntegrationTest : IntegrationTest() { diagnosticsInfo = null ) - val request = OpenAIRequestFactory().createInlineEditRequest(parameters) + val prompt = OpenAIRequestFactory().createInlineEditPrompt(parameters) - val systemMessage = request.messages[0] as OpenAIChatCompletionStandardMessage - assertThat(systemMessage.role).isEqualTo("system") + val systemMessage = prompt.messages[0] as KoogMessage.System val expectedSystem = buildExpectedInlineEditSystemPrompt( language = "java", filePath = tempFile.absolutePath, @@ -390,14 +355,11 @@ class OpenAIRequestFactoryIntegrationTest : IntegrationTest() { diagnosticsInfo = parameters.diagnosticsInfo, ) assertThat(systemMessage.content).isEqualTo(expectedSystem) - val prevUserMessage = request.messages[1] as OpenAIChatCompletionStandardMessage - assertThat(prevUserMessage.role).isEqualTo("user") + val prevUserMessage = prompt.messages[1] as KoogMessage.User assertThat(prevUserMessage.content).isEqualTo("PREV_PROMPT") - val prevResponse = request.messages[2] as OpenAIChatCompletionStandardMessage - assertThat(prevResponse.role).isEqualTo("assistant") + val prevResponse = prompt.messages[2] as KoogMessage.Assistant assertThat(prevResponse.content).isEqualTo("PREV_RESPONSE") - val followUpMessage = request.messages[3] as OpenAIChatCompletionStandardMessage - assertThat(followUpMessage.role).isEqualTo("user") + val followUpMessage = prompt.messages[3] as KoogMessage.User assertThat(followUpMessage.content).isEqualTo("Implement.") } @@ -486,31 +448,36 @@ class OpenAIRequestFactoryIntegrationTest : IntegrationTest() { return systemPrompt } - private fun normalize(messages: List): List> { + private fun normalize(messages: List): List> { return messages.map { msg -> when (msg) { - is OpenAIChatCompletionStandardMessage -> mapOf( - "role" to msg.role, + is KoogMessage.System -> mapOf( + "role" to "system", "content" to msg.content ) - is OpenAIChatCompletionAssistantMessage -> mapOf( - "role" to "assistant", - "content" to (msg.content ?: ""), - "toolCalls" to (msg.toolCalls?.map { tc -> - mapOf( - "id" to (tc.id ?: ""), - "type" to (tc.type ?: ""), - "fn" to (tc.function?.name ?: ""), - "args" to (tc.function?.arguments ?: "") - ) - } ?: emptyList()) + is KoogMessage.User -> mapOf( + "role" to "user", + "content" to msg.content ) - is OpenAIChatCompletionToolMessage -> mapOf( - "role" to "tool", - "callId" to (msg.callId ?: ""), - "content" to (msg.content ?: "") + is KoogMessage.Assistant -> mapOf( + "role" to "assistant", + "content" to msg.content + ) + + is KoogMessage.Tool.Call -> mapOf( + "role" to "tool_call", + "id" to msg.id, + "tool" to msg.tool, + "args" to msg.content + ) + + is KoogMessage.Tool.Result -> mapOf( + "role" to "tool_result", + "id" to msg.id, + "tool" to msg.tool, + "content" to msg.content ) else -> mapOf("role" to "unknown") diff --git a/src/test/kotlin/ee/carlrobert/codegpt/completions/factory/CustomOpenAIRequestFactoryTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/completions/factory/CustomOpenAIRequestFactoryTest.kt deleted file mode 100644 index dab3052f..00000000 --- a/src/test/kotlin/ee/carlrobert/codegpt/completions/factory/CustomOpenAIRequestFactoryTest.kt +++ /dev/null @@ -1,164 +0,0 @@ -package ee.carlrobert.codegpt.completions.factory - -import com.fasterxml.jackson.databind.JsonNode -import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper -import ee.carlrobert.codegpt.settings.service.custom.CustomServiceChatCompletionSettingsState -import ee.carlrobert.llm.client.openai.completion.request.OpenAIChatCompletionMessage -import ee.carlrobert.llm.client.openai.completion.request.OpenAIChatCompletionStandardMessage -import okhttp3.Request -import okio.Buffer -import org.assertj.core.api.Assertions.assertThat -import org.assertj.core.api.Assertions.tuple -import org.junit.Test -import java.util.function.Function - -class CustomOpenAIRequestFactoryTest { - - @Test - fun `should transform nested placeholder values for chat request body`() { - val settings = CustomServiceChatCompletionSettingsState().apply { - url = "https://example.com/v1/chat/completions" - headers = mutableMapOf("Authorization" to $$"Bearer $CUSTOM_SERVICE_API_KEY") - body = mutableMapOf( - "model" to "gpt-test", - "payload" to mapOf( - "prompt_alias" to $$"$PROMPT", - "messages_alias" to $$"$OPENAI_MESSAGES", - "items" to listOf( - mapOf("kind" to "prompt", "value" to $$"$PROMPT"), - mapOf("kind" to "messages", "value" to $$"$OPENAI_MESSAGES") - ) - ) - ) - } - val messages = listOf( - OpenAIChatCompletionStandardMessage("system", "System instructions"), - OpenAIChatCompletionStandardMessage("user", "Write a test"), - ) - - val request = CustomOpenAIRequestFactory.buildCustomOpenAIChatCompletionRequest( - settings = settings, - messages = messages, - streamRequest = true, - credential = "secret-key", - ) - val body = request.bodyAsJson() - val payload = body["payload"] - val aliasedMessages = payload["messages_alias"].elements().asSequence().toList() - val payloadItems = payload["items"].elements().asSequence().toList() - val nestedMessageItemValues = payloadItems[1]["value"].elements().asSequence().toList() - - assertThat(request.header("Authorization")).isEqualTo("Bearer secret-key") - assertThat(payload["prompt_alias"].asText()) - .isEqualTo("System instructions\n\nWrite a test") - assertThat(aliasedMessages) - .extracting({ it["role"].asText() }, { it["content"].asText() }) - .containsExactly( - tuple("system", "System instructions"), - tuple("user", "Write a test") - ) - assertThat(payloadItems) - .extracting( - { it["kind"].asText() }, - { it["value"].isArray }, - { it["value"].isTextual } - ) - .containsExactly( - tuple("prompt", false, true), - tuple("messages", true, false) - ) - assertThat(payloadItems[0]["value"].asText()) - .isEqualTo("System instructions\n\nWrite a test") - assertThat(nestedMessageItemValues) - .extracting(Function { it["role"].asText() }) - .containsExactly("system", "user") - } - - @Test - fun `should preserve nested arrays and objects in chat request body`() { - val settings = CustomServiceChatCompletionSettingsState().apply { - url = "https://example.com/v1/chat/completions" - headers = mutableMapOf("X-Test" to "true") - body = mutableMapOf( - "model" to "gpt-test", - "config" to mapOf( - "temperature" to 0.2, - "enabled" to true, - "tags" to listOf("alpha", 2, false), - "nested" to mapOf("arr" to listOf(mapOf("k" to "v"))) - ) - ) - } - val messages = listOf( - OpenAIChatCompletionStandardMessage("user", "hello") - ) - - val request = CustomOpenAIRequestFactory.buildCustomOpenAIChatCompletionRequest( - settings = settings, - messages = messages, - streamRequest = true, - credential = null, - ) - val body = request.bodyAsJson() - val tags = body["config"]["tags"].elements().asSequence().toList() - val nestedArray = body["config"]["nested"]["arr"].elements().asSequence().toList() - - assertThat(body["config"]["temperature"].asDouble()).isEqualTo(0.2) - assertThat(body["config"]["enabled"].asBoolean()).isTrue() - assertThat(tags) - .extracting({ it.nodeType.name }, { it.asText() }) - .containsExactly( - tuple("STRING", "alpha"), - tuple("NUMBER", "2"), - tuple("BOOLEAN", "false") - ) - assertThat(nestedArray) - .extracting(Function { it["k"].asText() }) - .containsExactly("v") - } - - @Test - fun `should replace nested stream flags and nested api key placeholders when stream is disabled`() { - val settings = CustomServiceChatCompletionSettingsState().apply { - url = "https://example.com/v1/chat/completions" - headers = mutableMapOf("X-Test" to "true") - body = mutableMapOf( - "model" to "gpt-test", - "stream" to "true", - "config" to mapOf( - "stream" to "true", - "authorization" to $$"Bearer $CUSTOM_SERVICE_API_KEY" - ), - "items" to listOf( - mapOf("stream" to "true"), - $$"token:$CUSTOM_SERVICE_API_KEY" - ) - ) - } - val messages = listOf( - OpenAIChatCompletionStandardMessage("user", "hello") - ) - - val request = CustomOpenAIRequestFactory.buildCustomOpenAIChatCompletionRequest( - settings = settings, - messages = messages, - streamRequest = false, - credential = "secret-key", - ) - val body = request.bodyAsJson() - - assertThat(listOf(body["stream"], body["config"]["stream"], body["items"][0]["stream"])) - .extracting(Function { it.asBoolean() }) - .containsExactly(false, false, false) - assertThat(listOf(body["config"]["authorization"], body["items"][1])) - .extracting(Function { it.asText() }) - .containsExactly("Bearer secret-key", "token:secret-key") - } - - private fun Request.bodyAsJson(): JsonNode { - val requestBody = requireNotNull(body) { "Request body is null" } - val buffer = Buffer() - requestBody.writeTo(buffer) - return jacksonObjectMapper().readTree(buffer.readUtf8()) - } -} diff --git a/src/test/kotlin/ee/carlrobert/codegpt/conversations/ConversationsStateTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/conversations/ConversationsStateTest.kt index 09a04fa3..e6f68408 100644 --- a/src/test/kotlin/ee/carlrobert/codegpt/conversations/ConversationsStateTest.kt +++ b/src/test/kotlin/ee/carlrobert/codegpt/conversations/ConversationsStateTest.kt @@ -5,7 +5,6 @@ import ee.carlrobert.codegpt.conversations.message.Message import ee.carlrobert.codegpt.settings.GeneralSettings import ee.carlrobert.codegpt.settings.service.ServiceType import ee.carlrobert.codegpt.settings.service.openai.OpenAISettings -import ee.carlrobert.llm.client.openai.completion.OpenAIChatCompletionModel import org.assertj.core.api.Assertions.assertThat class ConversationsStateTest : BasePlatformTestCase() { diff --git a/src/test/kotlin/ee/carlrobert/codegpt/conversations/message/MessageThreadSafetyTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/conversations/message/MessageThreadSafetyTest.kt new file mode 100644 index 00000000..03822135 --- /dev/null +++ b/src/test/kotlin/ee/carlrobert/codegpt/conversations/message/MessageThreadSafetyTest.kt @@ -0,0 +1,40 @@ +package ee.carlrobert.codegpt.conversations.message + +import ee.carlrobert.codegpt.completions.ChatToolCall +import ee.carlrobert.codegpt.completions.ChatToolFunction +import org.assertj.core.api.Assertions.assertThat +import org.junit.Test +import java.util.concurrent.CountDownLatch +import java.util.concurrent.Executors +import java.util.concurrent.TimeUnit + +class MessageThreadSafetyTest { + + @Test + fun `tool call updates should remain consistent under parallel writes`() { + val message = Message("prompt") + val tasks = 200 + val done = CountDownLatch(tasks) + val pool = Executors.newFixedThreadPool(8) + + repeat(tasks) { index -> + pool.execute { + val callId = "call-$index" + message.addToolCall( + ChatToolCall( + id = callId, + function = ChatToolFunction(name = "tool-$index", arguments = "{}") + ) + ) + message.addToolCallResult(callId, "result-$index") + done.countDown() + } + } + + assertThat(done.await(30, TimeUnit.SECONDS)).isTrue() + pool.shutdownNow() + + assertThat(message.toolCalls).hasSize(tasks) + assertThat(message.toolCallResults).hasSize(tasks) + } +} diff --git a/src/test/kotlin/ee/carlrobert/codegpt/mcp/McpToolArgumentsParserTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/mcp/McpToolArgumentsParserTest.kt new file mode 100644 index 00000000..d788cccf --- /dev/null +++ b/src/test/kotlin/ee/carlrobert/codegpt/mcp/McpToolArgumentsParserTest.kt @@ -0,0 +1,26 @@ +package ee.carlrobert.codegpt.mcp + +import org.assertj.core.api.Assertions.assertThat +import org.junit.Test + +class McpToolArgumentsParserTest { + + @Test + fun `parse should preserve explicit null values and normalize integral doubles`() { + val result = McpToolArgumentsParser.parse("""{"path":null,"limit":10.0,"ratio":1.5}""") + + assertThat(result.error).isNull() + assertThat(result.arguments).containsEntry("path", null) + assertThat(result.arguments["limit"]).isEqualTo(10) + assertThat(result.arguments["ratio"]).isEqualTo(1.5) + } + + @Test + fun `parse should return explicit error on invalid json`() { + val result = McpToolArgumentsParser.parse("""{"path":""") + + assertThat(result.arguments).isEmpty() + assertThat(result.error).isNotBlank() + } +} + diff --git a/src/test/kotlin/ee/carlrobert/codegpt/settings/models/ModelRegistryTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/settings/models/ModelRegistryTest.kt deleted file mode 100644 index 1378ba11..00000000 --- a/src/test/kotlin/ee/carlrobert/codegpt/settings/models/ModelRegistryTest.kt +++ /dev/null @@ -1,187 +0,0 @@ -package ee.carlrobert.codegpt.settings.models - -import ee.carlrobert.codegpt.settings.service.FeatureType -import ee.carlrobert.codegpt.settings.service.ServiceType -import ee.carlrobert.llm.client.codegpt.PricingPlan -import org.assertj.core.api.Assertions.assertThat -import testsupport.IntegrationTest - -class ModelRegistryTest : IntegrationTest() { - - private lateinit var modelRegistry: ModelRegistry - - override fun setUp() { - super.setUp() - modelRegistry = ModelRegistry.getInstance() - } - - fun `test getDefaultModelForFeature with individual plan returns premium model`() { - val result = modelRegistry.getDefaultModelForFeature(FeatureType.CHAT, PricingPlan.INDIVIDUAL) - - assertThat(result.provider).isEqualTo(ServiceType.PROXYAI) - assertThat(result.model).isEqualTo("claude-sonnet-4-5-thinking") - assertThat(result.displayName).isEqualTo("Claude 4.5 Sonnet Thinking") - } - - fun `test getDefaultModelForFeature with free plan returns free model`() { - val result = modelRegistry.getDefaultModelForFeature(FeatureType.CHAT, PricingPlan.FREE) - - assertThat(result.provider).isEqualTo(ServiceType.PROXYAI) - assertThat(result.model).isEqualTo("qwen3-coder") - assertThat(result.displayName).isEqualTo("Qwen3 Coder") - } - - fun `test getDefaultModelForFeature with anonymous plan returns basic model`() { - val result = modelRegistry.getDefaultModelForFeature(FeatureType.CHAT, PricingPlan.ANONYMOUS) - - assertThat(result.provider).isEqualTo(ServiceType.PROXYAI) - assertThat(result.model).isEqualTo("gemini-flash-2.5") - assertThat(result.displayName).isEqualTo("Gemini Flash 2.5") - } - - fun `test getDefaultModelForFeature with null plan returns fallback model`() { - val result = modelRegistry.getDefaultModelForFeature(FeatureType.CHAT, null) - - assertThat(result.provider).isEqualTo(ServiceType.PROXYAI) - assertThat(result.model).isEqualTo("gemini-flash-2.5") - assertThat(result.displayName).isEqualTo("Gemini Flash 2.5") - } - - fun `test getDefaultModelForFeature with code completion returns code model`() { - val result = modelRegistry.getDefaultModelForFeature(FeatureType.CODE_COMPLETION, PricingPlan.INDIVIDUAL) - - assertThat(result.provider).isEqualTo(ServiceType.PROXYAI) - assertThat(result.model).isEqualTo("mercury-coder") - assertThat(result.displayName).isEqualTo("Mercury Coder") - } - - fun `test getDefaultModelForFeature with next edit returns zeta model`() { - val result = modelRegistry.getDefaultModelForFeature(FeatureType.NEXT_EDIT, PricingPlan.INDIVIDUAL) - - assertThat(result.provider).isEqualTo(ServiceType.PROXYAI) - assertThat(result.model).isEqualTo("mercury-coder") - assertThat(result.displayName).isEqualTo("Mercury Coder") - } - - fun `test getAllModelsForFeature with chat returns chat models only`() { - val result = modelRegistry.getAllModelsForFeature(FeatureType.CHAT) - - assertThat(result).anyMatch { it.provider == ServiceType.PROXYAI && it.model == "claude-sonnet-4-5-thinking" } - assertThat(result).anyMatch { it.provider == ServiceType.OPENAI && it.model == "gpt-4.1" } - assertThat(result).anyMatch { it.provider == ServiceType.ANTHROPIC && it.model == "claude-sonnet-4-20250514" } - assertThat(result).anyMatch { it.provider == ServiceType.GOOGLE && it.model.contains("gemini") } - assertThat(result).anyMatch { it.provider == ServiceType.MISTRAL && it.model == "codestral-latest" } - assertThat(result).noneMatch { it.provider == ServiceType.ANTHROPIC && it.model == "qwen-2.5-32b-code" } - } - - fun `test getAllModelsForFeature with code completion returns code models only`() { - val result = modelRegistry.getAllModelsForFeature(FeatureType.CODE_COMPLETION) - - assertThat(result).anyMatch { it.provider == ServiceType.PROXYAI && it.model == "mercury-coder" } - assertThat(result).anyMatch { it.provider == ServiceType.OPENAI && it.model == "gpt-3.5-turbo-instruct" } - assertThat(result).anyMatch { it.provider == ServiceType.MISTRAL && it.model == "codestral-latest" } - assertThat(result).noneMatch { it.provider == ServiceType.ANTHROPIC } - assertThat(result).noneMatch { it.provider == ServiceType.GOOGLE } - } - - fun `test getAllModelsForFeature with next edit returns next edit models only`() { - val result = modelRegistry.getAllModelsForFeature(FeatureType.NEXT_EDIT) - - assertThat(result.first().provider).isEqualTo(ServiceType.PROXYAI) - assertThat(result.first().model).isEqualTo("mercury-coder") - } - - fun `test getProvidersForFeature with chat returns all chat providers`() { - val result = modelRegistry.getProvidersForFeature(FeatureType.CHAT) - - assertThat(result).containsExactlyInAnyOrder( - ServiceType.PROXYAI, - ServiceType.OPENAI, - ServiceType.ANTHROPIC, - ServiceType.GOOGLE, - ServiceType.MISTRAL, - ServiceType.OLLAMA, - ServiceType.LLAMA_CPP, - ServiceType.CUSTOM_OPENAI, - ServiceType.INCEPTION, - ) - } - - fun `test getProvidersForFeature with code completion returns code providers only`() { - val result = modelRegistry.getProvidersForFeature(FeatureType.CODE_COMPLETION) - - assertThat(result).containsExactlyInAnyOrder( - ServiceType.PROXYAI, - ServiceType.OPENAI, - ServiceType.MISTRAL, - ServiceType.OLLAMA, - ServiceType.LLAMA_CPP, - ServiceType.CUSTOM_OPENAI, - ServiceType.INCEPTION - ) - assertThat(result).doesNotContain(ServiceType.ANTHROPIC, ServiceType.GOOGLE) - } - - fun `test getProvidersForFeature with next edit returns proxyai and inception`() { - val result = modelRegistry.getProvidersForFeature(FeatureType.NEXT_EDIT) - - assertThat(result).containsExactlyInAnyOrder(ServiceType.PROXYAI, ServiceType.INCEPTION) - } - - fun `test isFeatureSupportedByProvider with valid chat provider returns true`() { - val result = modelRegistry.isFeatureSupportedByProvider(FeatureType.CHAT, ServiceType.OPENAI) - - assertThat(result).isTrue - } - - fun `test isFeatureSupportedByProvider with invalid code provider returns false`() { - val result = modelRegistry.isFeatureSupportedByProvider(FeatureType.CODE_COMPLETION, ServiceType.ANTHROPIC) - - assertThat(result).isFalse - } - - fun `test isFeatureSupportedByProvider with next edit non-proxyai returns false`() { - val result = modelRegistry.isFeatureSupportedByProvider(FeatureType.NEXT_EDIT, ServiceType.OPENAI) - - assertThat(result).isFalse - } - - fun `test findModel with existing provider and model returns model selection`() { - val result = modelRegistry.findModel(ServiceType.PROXYAI, "gemini-flash-2.5") - - assertThat(result).isNotNull - assertThat(result!!.provider).isEqualTo(ServiceType.PROXYAI) - assertThat(result.model).isEqualTo("gemini-flash-2.5") - assertThat(result.displayName).isEqualTo("Gemini 2.5 Flash") - } - - fun `test findModel with existing anthropic model returns model selection`() { - val result = modelRegistry.findModel(ServiceType.ANTHROPIC, "claude-sonnet-4-20250514") - - assertThat(result).isNotNull - assertThat(result!!.provider).isEqualTo(ServiceType.ANTHROPIC) - assertThat(result.model).isEqualTo("claude-sonnet-4-20250514") - assertThat(result.displayName).isEqualTo("Claude Sonnet 4") - } - - fun `test findModel with existing mistral model returns model selection`() { - val result = modelRegistry.findModel(ServiceType.MISTRAL, "codestral-latest") - - assertThat(result).isNotNull - assertThat(result!!.provider).isEqualTo(ServiceType.MISTRAL) - assertThat(result.model).isEqualTo("codestral-latest") - assertThat(result.displayName).isEqualTo("Codestral") - } - - fun `test findModel with non-existing model returns null`() { - val result = modelRegistry.findModel(ServiceType.OPENAI, "non-existing-model") - - assertThat(result).isNull() - } - - fun `test findModel with wrong provider for model returns null`() { - val result = modelRegistry.findModel(ServiceType.ANTHROPIC, "gpt-4o") - - assertThat(result).isNull() - } -} diff --git a/src/test/kotlin/ee/carlrobert/codegpt/settings/models/ModelSelectionServiceTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/settings/models/ModelSelectionServiceTest.kt deleted file mode 100644 index 08b145ee..00000000 --- a/src/test/kotlin/ee/carlrobert/codegpt/settings/models/ModelSelectionServiceTest.kt +++ /dev/null @@ -1,80 +0,0 @@ -package ee.carlrobert.codegpt.settings.service - -import com.intellij.openapi.components.service -import ee.carlrobert.codegpt.settings.models.ModelRegistry -import ee.carlrobert.codegpt.settings.models.ModelSettings -import ee.carlrobert.codegpt.settings.models.ModelSettingsState -import ee.carlrobert.llm.client.codegpt.PricingPlan -import org.assertj.core.api.Assertions.assertThat -import testsupport.IntegrationTest - -class ModelSelectionServiceTest : IntegrationTest() { - - private lateinit var modelSelectionService: ModelSelectionService - private lateinit var modelSettings: ModelSettings - - override fun setUp() { - super.setUp() - modelSelectionService = service() - modelSettings = service() - - val cleanState = ModelSettingsState() - cleanState.modelSelections.clear() - modelSettings.loadState(cleanState) - } - - fun `test getModelSelectionForFeature with valid feature returns model`() { - val result = modelSelectionService.getModelSelectionForFeature(FeatureType.CHAT) - val expected = ModelRegistry.getInstance().getDefaultModelForFeature(FeatureType.CHAT) - - assertThat(result.provider).isEqualTo(expected.provider) - assertThat(result.model).isEqualTo(expected.model) - assertThat(result.displayName).isEqualTo(expected.displayName) - } - - fun `test getModelSelectionForFeature with pricing plan returns plan-specific model`() { - val individualResult = modelSelectionService.getModelSelectionForFeature( - FeatureType.CHAT, - PricingPlan.INDIVIDUAL - ) - val freeResult = - modelSelectionService.getModelSelectionForFeature(FeatureType.CHAT, PricingPlan.FREE) - - assertThat(individualResult.model).isEqualTo("claude-sonnet-4-5-thinking") - assertThat(freeResult.model).isEqualTo("qwen3-coder") - } - - fun `test getModelSelectionForFeature with code completion returns code model`() { - val result = modelSelectionService.getModelSelectionForFeature(FeatureType.CODE_COMPLETION) - - assertThat(result.provider).isEqualTo(ServiceType.PROXYAI) - assertThat(result.model).isEqualTo("mercury-coder") - } - - fun `test getServiceForFeature with valid feature returns correct provider`() { - val result = modelSelectionService.getServiceForFeature(FeatureType.CHAT) - - assertThat(result).isEqualTo(ServiceType.PROXYAI) - } - - fun `test getServiceForFeature with pricing plan returns plan-specific provider`() { - val result = - modelSelectionService.getServiceForFeature(FeatureType.CHAT, PricingPlan.INDIVIDUAL) - - assertThat(result).isEqualTo(ServiceType.PROXYAI) - } - - fun `test getModelForFeature with valid feature returns model string`() { - val result = modelSelectionService.getModelForFeature(FeatureType.CHAT) - val expected = ModelRegistry.getInstance().getDefaultModelForFeature(FeatureType.CHAT) - - assertThat(result).isEqualTo(expected.model) - } - - fun `test getModelForFeature with pricing plan returns plan-specific model`() { - val result = - modelSelectionService.getModelForFeature(FeatureType.CHAT, PricingPlan.INDIVIDUAL) - - assertThat(result).isEqualTo("claude-sonnet-4-5-thinking") - } -} \ No newline at end of file diff --git a/src/test/kotlin/ee/carlrobert/codegpt/settings/models/ModelSettingsTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/settings/models/ModelSettingsTest.kt index a8d694e6..246822d3 100644 --- a/src/test/kotlin/ee/carlrobert/codegpt/settings/models/ModelSettingsTest.kt +++ b/src/test/kotlin/ee/carlrobert/codegpt/settings/models/ModelSettingsTest.kt @@ -177,27 +177,27 @@ class ModelSettingsTest : IntegrationTest() { } fun `test getModelSelection with valid feature returns model selection`() { - modelSettings.setModelWithProvider(FeatureType.CHAT, "gpt-4o", ServiceType.OPENAI) + modelSettings.setModelWithProvider(FeatureType.CHAT, "gpt-4.1", ServiceType.OPENAI) val result = modelSettings.getModelSelection(FeatureType.CHAT) assertThat(result).isNotNull assertThat(result?.provider).isEqualTo(ServiceType.OPENAI) - assertThat(result?.model).isEqualTo("gpt-4o") + assertThat(result?.model).isEqualTo("gpt-4.1") } fun `test getModelForFeature returns stored model`() { - modelSettings.setModelWithProvider(FeatureType.CHAT, "gpt-4o", ServiceType.OPENAI) + modelSettings.setModelWithProvider(FeatureType.CHAT, "gpt-4.1", ServiceType.OPENAI) - val result = modelSettings.getModelForFeature(FeatureType.CHAT) + val result = modelSettings.getStoredModelForFeature(FeatureType.CHAT) - assertThat(result).isEqualTo("gpt-4o") + assertThat(result).isEqualTo("gpt-4.1") } fun `test getProviderForFeature returns stored provider`() { - modelSettings.setModelWithProvider(FeatureType.CHAT, "gpt-4o", ServiceType.OPENAI) + modelSettings.setModelWithProvider(FeatureType.CHAT, "gpt-4.1", ServiceType.OPENAI) - val result = modelSettings.getProviderForFeature(FeatureType.CHAT) + val result = modelSettings.getStoredProviderForFeature(FeatureType.CHAT) assertThat(result).isEqualTo(ServiceType.OPENAI) } @@ -205,13 +205,13 @@ class ModelSettingsTest : IntegrationTest() { fun `test migrateMissingProviderInformation updates missing providers`() { val state = ModelSettingsState() val detailsState = ModelDetailsState() - detailsState.model = "gpt-4o" + detailsState.model = "gpt-4.1" detailsState.provider = null state.modelSelections["CHAT"] = detailsState modelSettings.loadState(state) - val result = modelSettings.getProviderForFeature(FeatureType.CHAT) + val result = modelSettings.getStoredProviderForFeature(FeatureType.CHAT) assertThat(result).isEqualTo(ServiceType.OPENAI) } @@ -224,7 +224,7 @@ class ModelSettingsTest : IntegrationTest() { modelSettings.loadState(state) - assertThat(modelSettings.getModelForFeature(FeatureType.CHAT)).isEqualTo("unknown-model") - assertThat(modelSettings.getProviderForFeature(FeatureType.CHAT)).isNull() + assertThat(modelSettings.getStoredModelForFeature(FeatureType.CHAT)).isEqualTo("unknown-model") + assertThat(modelSettings.getStoredProviderForFeature(FeatureType.CHAT)).isNull() } -} \ No newline at end of file +} diff --git a/src/test/kotlin/ee/carlrobert/codegpt/toolwindow/chat/ChatToolWindowTabPanelTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/toolwindow/chat/ChatToolWindowTabPanelTest.kt index b5428b91..fb853688 100644 --- a/src/test/kotlin/ee/carlrobert/codegpt/toolwindow/chat/ChatToolWindowTabPanelTest.kt +++ b/src/test/kotlin/ee/carlrobert/codegpt/toolwindow/chat/ChatToolWindowTabPanelTest.kt @@ -3,410 +3,168 @@ package ee.carlrobert.codegpt.toolwindow.chat import com.intellij.openapi.components.service import com.intellij.testFramework.LightVirtualFile import ee.carlrobert.codegpt.CodeGPTKeys -import ee.carlrobert.codegpt.EncodingManager -import ee.carlrobert.codegpt.completions.CompletionRequestUtil import ee.carlrobert.codegpt.completions.ConversationType -import ee.carlrobert.codegpt.completions.HuggingFaceModel -import ee.carlrobert.codegpt.completions.llama.PromptTemplate.LLAMA import ee.carlrobert.codegpt.conversations.ConversationService import ee.carlrobert.codegpt.conversations.message.Message -import ee.carlrobert.codegpt.settings.configuration.ConfigurationSettings +import ee.carlrobert.codegpt.settings.models.ModelSettings import ee.carlrobert.codegpt.settings.prompts.PromptsSettings -import ee.carlrobert.codegpt.settings.service.llama.LlamaSettings -import ee.carlrobert.llm.client.http.RequestEntity -import ee.carlrobert.llm.client.http.exchange.StreamHttpExchange -import ee.carlrobert.llm.client.util.JSONUtil.* -import org.apache.http.HttpHeaders -import ee.carlrobert.codegpt.util.file.FileUtil.getResourceContent +import ee.carlrobert.codegpt.settings.service.FeatureType +import ee.carlrobert.codegpt.settings.service.ServiceType import org.assertj.core.api.Assertions.assertThat import testsupport.IntegrationTest -import java.io.IOException -import java.nio.file.Files -import java.nio.file.Path +import testsupport.http.RequestEntity +import testsupport.http.exchange.StreamHttpExchange +import testsupport.json.JSONUtil.e +import testsupport.json.JSONUtil.jsonMapResponse import java.util.* class ChatToolWindowTabPanelTest : IntegrationTest() { - fun testSendingOpenAIMessage() { + override fun setUp() { + super.setUp() useOpenAIService() + service().setModel(FeatureType.CHAT, "gpt-5-mini", ServiceType.OPENAI) + } + + fun testSendingMessagePersistsResponseAndTokenDetails() { service().state.personas.selectedPersona.instructions = "TEST_SYSTEM_PROMPT" + val message = Message("Hello!") val conversation = ConversationService.getInstance().startConversation(project) val panel = ChatToolWindowTabPanel(project, conversation) - expectOpenAI(StreamHttpExchange { request: RequestEntity -> - assertThat(request.uri.path).isEqualTo("/v1/chat/completions") - assertThat(request.method).isEqualTo("POST") - assertThat(request.headers[HttpHeaders.AUTHORIZATION]!![0]).isEqualTo("Bearer TEST_API_KEY") - val guidelines = getResourceContent("/prompts/persona/psi-navigation-guidelines.txt") - val expectedSystem = "TEST_SYSTEM_PROMPT\n$guidelines" - assertThat(request.body) - .extracting( - "model", - "messages" - ) - .containsExactly( - "gpt-4o", - listOf( - mapOf("role" to "system", "content" to expectedSystem), - mapOf("role" to "user", "content" to "Hello!") - ) - ) - listOf( - jsonMapResponse( - "choices", - jsonArray(jsonMap("delta", jsonMap("role", "assistant"))) - ), - jsonMapResponse("choices", jsonArray(jsonMap("delta", jsonMap("content", "Hel")))), - jsonMapResponse("choices", jsonArray(jsonMap("delta", jsonMap("content", "lo")))), - jsonMapResponse("choices", jsonArray(jsonMap("delta", jsonMap("content", "!")))) - ) - }) + expectOpenAIStreamingHello { promptText -> + assertThat(promptText).contains("TEST_SYSTEM_PROMPT") + assertThat(promptText).contains("Hello!") + } panel.sendMessage(message, ConversationType.DEFAULT) waitExpecting { - val messages = conversation.messages - messages.isNotEmpty() && "Hello!" == messages[0].response - && panel.tokenDetails.conversationTokens > 0 + conversation.messages.isNotEmpty() && "Hello!" == conversation.messages[0].response } - val encodingManager = EncodingManager.getInstance() - assertThat(panel.tokenDetails).extracting( - "systemPromptTokens", - "conversationTokens", - "userPromptTokens", - "highlightedTokens" - ) - .containsExactly( - encodingManager.countTokens("TEST_SYSTEM_PROMPT"), - encodingManager.countConversationTokens(conversation), - 0, - 0 - ) - assertThat(panel.conversation) - .isNotNull() - .extracting("id", "discardTokenLimit") - .containsExactly(conversation.id, false) - val messages = panel.conversation.messages - assertThat(messages).hasSize(1) - assertThat(messages[0]) - .extracting("id", "prompt", "response") - .containsExactly(message.id, message.prompt, message.response) + + assertThat(panel.conversation.messages).hasSize(1) + assertThat(panel.conversation.messages[0].response).isEqualTo("Hello!") } - fun testSendingOpenAIMessageWithReferencedContext() { - useOpenAIService() - service().state.personas.selectedPersona.instructions = - "TEST_SYSTEM_PROMPT" - val message = Message("TEST_MESSAGE") - val conversation = ConversationService.getInstance().startConversation(project) - val panel = ChatToolWindowTabPanel(project, conversation) - val vf1 = LightVirtualFile("TEST_FILE_NAME_1", "TEST_FILE_CONTENT_1") - val vf2 = LightVirtualFile("TEST_FILE_NAME_2", "TEST_FILE_CONTENT_2") - val vf3 = LightVirtualFile("TEST_FILE_NAME_3", "TEST_FILE_CONTENT_3") - panel.includeFiles(listOf(vf1, vf2, vf3)) - expectOpenAI(StreamHttpExchange { request: RequestEntity -> - assertThat(request.uri.path).isEqualTo("/v1/chat/completions") - assertThat(request.method).isEqualTo("POST") - assertThat(request.headers[HttpHeaders.AUTHORIZATION]!![0]).isEqualTo("Bearer TEST_API_KEY") - val guidelines = getResourceContent("/prompts/persona/psi-navigation-guidelines.txt") - val filesBlock = buildString { - append("\n\n") - append("\n") - append(CompletionRequestUtil.formatCode("TEST_FILE_CONTENT_1", vf1.path)) - append("\n") - append(CompletionRequestUtil.formatCode("TEST_FILE_CONTENT_2", vf2.path)) - append("\n") - append(CompletionRequestUtil.formatCode("TEST_FILE_CONTENT_3", vf3.path)) - append("\n") - } - val expectedSystem = "TEST_SYSTEM_PROMPT\n$guidelines$filesBlock" - assertThat(request.body) - .extracting( - "model", - "messages" - ) - .containsExactly( - "gpt-4o", - listOf( - mapOf("role" to "system", "content" to expectedSystem), - mapOf("role" to "user", "content" to "TEST_MESSAGE") - ) - ) - listOf( - jsonMapResponse( - "choices", - jsonArray(jsonMap("delta", jsonMap("role", "assistant"))) - ), - jsonMapResponse("choices", jsonArray(jsonMap("delta", jsonMap("content", "Hel")))), - jsonMapResponse("choices", jsonArray(jsonMap("delta", jsonMap("content", "lo")))), - jsonMapResponse("choices", jsonArray(jsonMap("delta", jsonMap("content", "!")))) - ) - }) - - panel.sendMessage(message, ConversationType.DEFAULT) - - waitExpecting { - val messages = conversation.messages - messages.isNotEmpty() && "Hello!" == messages[0].response - && panel.tokenDetails.conversationTokens > 0 - } - val encodingManager = EncodingManager.getInstance() - assertThat(panel.tokenDetails).extracting( - "systemPromptTokens", - "conversationTokens", - "userPromptTokens", - "highlightedTokens" - ) - .containsExactly( - encodingManager.countTokens("TEST_SYSTEM_PROMPT"), - encodingManager.countConversationTokens(conversation), - 0, - 0 - ) - assertThat(panel.conversation) - .isNotNull() - .extracting("id", "discardTokenLimit") - .containsExactly(conversation.id, false) - val messages = panel.conversation.messages - assertThat(messages).hasSize(1) - assertThat(messages[0]) - .extracting("id", "prompt", "response") - .containsExactly( - message.id, - message.prompt, - message.response - ) - } - - fun testSendingOpenAIMessageWithImageInSession() { - val testImagePath = - Objects.requireNonNull(javaClass.getResource("/images/test-image.png")).path - project.putUserData(CodeGPTKeys.IMAGE_ATTACHMENT_FILE_PATH, testImagePath) - useOpenAIService("gpt-4-vision-preview") - service().state.personas.selectedPersona.instructions = - "TEST_SYSTEM_PROMPT" - val message = Message("TEST_MESSAGE") - val conversation = ConversationService.getInstance().startConversation(project) - val panel = ChatToolWindowTabPanel(project, conversation) - expectOpenAI(StreamHttpExchange { request: RequestEntity -> - assertThat(request.uri.path).isEqualTo("/v1/chat/completions") - assertThat(request.method).isEqualTo("POST") - assertThat(request.headers[HttpHeaders.AUTHORIZATION]!![0]).isEqualTo("Bearer TEST_API_KEY") - try { - val guidelines = getResourceContent("/prompts/persona/psi-navigation-guidelines.txt") - val expectedSystem = "TEST_SYSTEM_PROMPT\n$guidelines" - val testImageUrl = ("data:image/png;base64," - + Base64.getEncoder() - .encodeToString(Files.readAllBytes(Path.of(testImagePath)))) - assertThat(request.body) - .extracting("model", "messages") - .containsExactly( - "gpt-4-vision-preview", - listOf( - mapOf("role" to "system", "content" to expectedSystem), - mapOf( - "role" to "user", "content" to listOf( - mapOf( - "type" to "image_url", - "image_url" to mapOf("url" to testImageUrl) - ), - mapOf("type" to "text", "text" to "TEST_MESSAGE") - ) - ) - ) - ) - } catch (e: IOException) { - throw RuntimeException(e) - } - listOf( - jsonMapResponse( - "choices", - jsonArray(jsonMap("delta", jsonMap("role", "assistant"))) - ), - jsonMapResponse("choices", jsonArray(jsonMap("delta", jsonMap("content", "Hel")))), - jsonMapResponse("choices", jsonArray(jsonMap("delta", jsonMap("content", "lo")))), - jsonMapResponse("choices", jsonArray(jsonMap("delta", jsonMap("content", "!")))) - ) - }) - - panel.sendMessage(message, ConversationType.DEFAULT) - - waitExpecting { - val messages = conversation.messages - messages.isNotEmpty() && "Hello!" == messages[0].response - && panel.tokenDetails.conversationTokens > 0 - } - val encodingManager = EncodingManager.getInstance() - assertThat(panel.tokenDetails).extracting( - "systemPromptTokens", - "conversationTokens", - "userPromptTokens", - "highlightedTokens" - ) - .containsExactly( - encodingManager.countTokens("TEST_SYSTEM_PROMPT"), - encodingManager.countConversationTokens(conversation), - 0, - 0 - ) - assertThat(panel.conversation) - .isNotNull() - .extracting("id", "discardTokenLimit") - .containsExactly(conversation.id, false) - val messages = panel.conversation.messages - assertThat(messages).hasSize(1) - assertThat(messages[0]) - .extracting("id", "prompt", "response", "imageFilePath") - .containsExactly( - message.id, - message.prompt, - message.response, - message.imageFilePath - ) - } - - fun testFixCompileErrorsWithOpenAIService() { - useOpenAIService() - service().state.personas.selectedPersona.instructions = "TEST_SYSTEM_PROMPT" - val message = Message("TEST_MESSAGE") + fun testSendingMessageWithReferencedFilesAddsFileContextToPrompt() { + val message = Message("Explain referenced files") val conversation = ConversationService.getInstance().startConversation(project) val panel = ChatToolWindowTabPanel(project, conversation) panel.includeFiles( listOf( - LightVirtualFile("TEST_FILE_NAME_1", "TEST_FILE_CONTENT_1"), - LightVirtualFile("TEST_FILE_NAME_2", "TEST_FILE_CONTENT_2"), - LightVirtualFile("TEST_FILE_NAME_3", "TEST_FILE_CONTENT_3"), + LightVirtualFile("A.kt", "fun a() = 1"), + LightVirtualFile("B.kt", "class B") ) ) - expectOpenAI(StreamHttpExchange { request: RequestEntity -> - assertThat(request.uri.path).isEqualTo("/v1/chat/completions") - assertThat(request.method).isEqualTo("POST") - assertThat(request.headers[HttpHeaders.AUTHORIZATION]!![0]).isEqualTo("Bearer TEST_API_KEY") - assertThat(request.body) - .extracting( - "model", - "messages" - ) - .containsExactly( - "gpt-4o", - listOf( - mapOf( - "role" to "system", - "content" to service().state.coreActions.fixCompileErrors.instructions - ), - mapOf( - "role" to "user", - "content" to "TEST_MESSAGE" - ) - ) - ) - listOf( - jsonMapResponse( - "choices", - jsonArray(jsonMap("delta", jsonMap("role", "assistant"))) - ), - jsonMapResponse("choices", jsonArray(jsonMap("delta", jsonMap("content", "Hel")))), - jsonMapResponse("choices", jsonArray(jsonMap("delta", jsonMap("content", "lo")))), - jsonMapResponse("choices", jsonArray(jsonMap("delta", jsonMap("content", "!")))) - ) - }) - - panel.sendMessage(message, ConversationType.FIX_COMPILE_ERRORS) - - waitExpecting { - val messages = conversation.messages - messages.isNotEmpty() && "Hello!" == messages[0].response - && panel.tokenDetails.conversationTokens > 0 + expectOpenAIStreamingHello { promptText -> + assertThat(promptText).contains("A.kt") + assertThat(promptText).contains("fun a() = 1") + assertThat(promptText).contains("B.kt") + assertThat(promptText).contains("class B") } - val encodingManager = EncodingManager.getInstance() - assertThat(panel.tokenDetails).extracting( - "systemPromptTokens", - "conversationTokens", - "userPromptTokens", - "highlightedTokens" - ) - .containsExactly( - encodingManager.countTokens("TEST_SYSTEM_PROMPT"), - encodingManager.countConversationTokens(conversation), - 0, - 0 - ) - assertThat(panel.conversation) - .isNotNull() - .extracting("id", "discardTokenLimit") - .containsExactly(conversation.id, false) - val messages = panel.conversation.messages - assertThat(messages).hasSize(1) - assertThat(messages[0]) - .extracting("id", "prompt", "response") - .containsExactly( - message.id, - message.prompt, - message.response - ) - } - - fun testSendingLlamaMessage() { - useLlamaService() - val configurationState = service().state - service().state.personas.selectedPersona.instructions = - "TEST_SYSTEM_PROMPT" - configurationState.maxTokens = 1000 - configurationState.temperature = 0.1f - val llamaSettings = LlamaSettings.getCurrentState() - llamaSettings.isUseCustomModel = false - llamaSettings.huggingFaceModel = HuggingFaceModel.CODE_LLAMA_7B_Q4 - llamaSettings.topK = 30 - llamaSettings.topP = 0.8 - llamaSettings.minP = 0.03 - llamaSettings.repeatPenalty = 1.3 - val message = Message("TEST_PROMPT") - val conversation = ConversationService.getInstance().startConversation(project) - val panel = ChatToolWindowTabPanel(project, conversation) - expectLlama(StreamHttpExchange { request: RequestEntity -> - assertThat(request.uri.path).isEqualTo("/v1/chat/completions") - val expectedSystem = getResourceContent("/prompts/persona/psi-navigation-guidelines.txt").let { g -> - "TEST_SYSTEM_PROMPT\n$g" - } - assertThat(request.body) - .extracting( - "model", - "messages" - ) - .containsExactly( - HuggingFaceModel.CODE_LLAMA_7B_Q4.code, - listOf( - mapOf("role" to "system", "content" to expectedSystem), - mapOf("role" to "user", "content" to "TEST_PROMPT") - ) - ) - listOf( - jsonMapResponse( - "choices", - jsonArray(jsonMap("delta", jsonMap("role", "assistant"))) - ), - jsonMapResponse("choices", jsonArray(jsonMap("delta", jsonMap("content", "Hel")))), - jsonMapResponse("choices", jsonArray(jsonMap("delta", jsonMap("content", "lo")))), - jsonMapResponse("choices", jsonArray(jsonMap("delta", jsonMap("content", "!")))) - ) - }) panel.sendMessage(message, ConversationType.DEFAULT) waitExpecting { - val messages = conversation.messages - messages.isNotEmpty() && "Hello!" == messages[0].response - && panel.tokenDetails.conversationTokens > 0 + conversation.messages.isNotEmpty() && "Hello!" == conversation.messages[0].response + } + } + + fun testFixCompileErrorsConversationUsesCoreActionPrompt() { + service().state.coreActions.fixCompileErrors.instructions = + "FIX_ERRORS_SYSTEM_PROMPT" + + val message = Message("Fix compile errors in this class") + val conversation = ConversationService.getInstance().startConversation(project) + val panel = ChatToolWindowTabPanel(project, conversation) + expectOpenAIStreamingHello { promptText -> + assertThat(promptText).contains("FIX_ERRORS_SYSTEM_PROMPT") + assertThat(promptText).contains("Fix compile errors in this class") + } + + panel.sendMessage(message, ConversationType.FIX_COMPILE_ERRORS) + + waitExpecting { + conversation.messages.isNotEmpty() && "Hello!" == conversation.messages[0].response + } + } + + fun testSendingMessageWithImageStillBuildsPromptAndResponse() { + val testImagePath = + Objects.requireNonNull(javaClass.getResource("/images/test-image.png")).path + project.putUserData(CodeGPTKeys.IMAGE_ATTACHMENT_FILE_PATH, testImagePath) + + val message = Message("What is in this image?") + val conversation = ConversationService.getInstance().startConversation(project) + val panel = ChatToolWindowTabPanel(project, conversation) + expectOpenAIStreamingHello { promptText -> + assertThat(promptText).contains("What is in this image?") + } + + panel.sendMessage(message, ConversationType.DEFAULT) + + waitExpecting { + conversation.messages.isNotEmpty() && "Hello!" == conversation.messages[0].response + } + } + + private fun expectOpenAIStreamingHello(assertPrompt: (String) -> Unit) { + expectOpenAI(StreamHttpExchange { request: RequestEntity -> + assertThat(request.uri.path).isEqualTo("/v1/responses") + assertThat(request.method).isEqualTo("POST") + assertPrompt(extractPromptText(request)) + listOf( + streamingChunk("Hel", 1), + streamingChunk("lo", 2), + streamingChunk("!", 3) + ) + }) + } + + private fun streamingChunk(content: String, sequenceNumber: Int): String { + return jsonMapResponse( + e("type", "response.output_text.delta"), + e("item_id", "msg_1"), + e("output_index", 0), + e("content_index", 0), + e("delta", content), + e("sequence_number", sequenceNumber) + ) + } + + private fun extractPromptText(request: RequestEntity): String { + val messages = request.body["messages"] as? List<*> + if (messages != null) { + return messages.joinToString("\n") { message -> + when (val content = (message as? Map<*, *>)?.get("content")) { + is String -> content + is List<*> -> content.joinToString("\n") { part -> + val partMap = part as? Map<*, *> + (partMap?.get("text") as? String) ?: partMap.toString() + } + + null -> "" + else -> content.toString() + } + } + } + + val input = request.body["input"] as? List<*> ?: return "" + return input.joinToString("\n") { item -> + when (val content = (item as? Map<*, *>)?.get("content")) { + is String -> content + is List<*> -> content.joinToString("\n") { part -> + val partMap = part as? Map<*, *> + (partMap?.get("text") as? String) + ?: (partMap?.get("value") as? String) + ?: partMap.toString() + } + + null -> "" + else -> content.toString() + } } - assertThat(panel.conversation) - .isNotNull() - .extracting("id", "discardTokenLimit") - .containsExactly(conversation.id, false) - val messages = panel.conversation.messages - assertThat(messages).hasSize(1) - assertThat(messages[0]) - .extracting("id", "prompt", "response") - .containsExactly(message.id, message.prompt, message.response) } } diff --git a/src/test/kotlin/ee/carlrobert/codegpt/util/ReasoningFrameTextAdapterTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/util/ReasoningFrameTextAdapterTest.kt new file mode 100644 index 00000000..ce4d223c --- /dev/null +++ b/src/test/kotlin/ee/carlrobert/codegpt/util/ReasoningFrameTextAdapterTest.kt @@ -0,0 +1,64 @@ +package ee.carlrobert.codegpt.util + +import ai.koog.prompt.streaming.StreamFrame +import org.assertj.core.api.Assertions.assertThat +import org.junit.Test + +class ReasoningFrameTextAdapterTest { + + @Test + fun `should not duplicate reasoning complete when deltas were already emitted`() { + val adapter = ReasoningFrameTextAdapter() + + assertThat(adapter.consume(StreamFrame.ReasoningDelta(text = "step 1", summary = null, index = 0))) + .containsExactly("", "step 1") + assertThat(adapter.consume(StreamFrame.ReasoningDelta(text = " step 2", summary = null, index = 0))) + .containsExactly(" step 2") + + assertThat( + adapter.consume( + StreamFrame.ReasoningComplete( + text = listOf("step 1 step 2"), + summary = emptyList(), + encrypted = null, + index = 0 + ) + ) + ).containsExactly("") + } + + @Test + fun `should emit full reasoning complete when there were no deltas`() { + val adapter = ReasoningFrameTextAdapter() + + assertThat( + adapter.consume( + StreamFrame.ReasoningComplete( + text = listOf("full reasoning"), + summary = emptyList(), + encrypted = null, + index = 1 + ) + ) + ).containsExactly("", "full reasoning", "") + } + + @Test + fun `should deduplicate for null reasoning index`() { + val adapter = ReasoningFrameTextAdapter() + + assertThat(adapter.consume(StreamFrame.ReasoningDelta(text = "delta", summary = null, index = null))) + .containsExactly("", "delta") + + assertThat( + adapter.consume( + StreamFrame.ReasoningComplete( + text = listOf("delta"), + summary = emptyList(), + encrypted = null, + index = null + ) + ) + ).containsExactly("") + } +} diff --git a/src/test/kotlin/testsupport/IntegrationTest.kt b/src/test/kotlin/testsupport/IntegrationTest.kt index d5ec501f..9a91541e 100644 --- a/src/test/kotlin/testsupport/IntegrationTest.kt +++ b/src/test/kotlin/testsupport/IntegrationTest.kt @@ -3,7 +3,7 @@ package testsupport import com.intellij.openapi.util.Key import com.intellij.testFramework.fixtures.BasePlatformTestCase import ee.carlrobert.codegpt.CodeGPTKeys -import ee.carlrobert.llm.client.mixin.ExternalServiceTestMixin +import testsupport.mixin.ExternalServiceTestMixin import testsupport.mixin.ShortcutsTestMixin open class IntegrationTest : BasePlatformTestCase(), ExternalServiceTestMixin, ShortcutsTestMixin { diff --git a/src/test/kotlin/testsupport/http/RequestEntity.kt b/src/test/kotlin/testsupport/http/RequestEntity.kt new file mode 100644 index 00000000..4fdda422 --- /dev/null +++ b/src/test/kotlin/testsupport/http/RequestEntity.kt @@ -0,0 +1,9 @@ +package testsupport.http + +import java.net.URI + +data class RequestEntity( + val method: String, + val uri: URI, + val body: Map, +) diff --git a/src/test/kotlin/testsupport/http/ResponseEntity.kt b/src/test/kotlin/testsupport/http/ResponseEntity.kt new file mode 100644 index 00000000..7155116e --- /dev/null +++ b/src/test/kotlin/testsupport/http/ResponseEntity.kt @@ -0,0 +1,7 @@ +package testsupport.http + +data class ResponseEntity( + val body: String, + val statusCode: Int = 200, + val contentType: String = "application/json", +) diff --git a/src/test/kotlin/testsupport/http/exchange/Exchanges.kt b/src/test/kotlin/testsupport/http/exchange/Exchanges.kt new file mode 100644 index 00000000..cb6f58c3 --- /dev/null +++ b/src/test/kotlin/testsupport/http/exchange/Exchanges.kt @@ -0,0 +1,16 @@ +package testsupport.http.exchange + +import testsupport.http.RequestEntity +import testsupport.http.ResponseEntity + +fun interface BasicHttpExchange { + fun exchange(request: RequestEntity): ResponseEntity +} + +fun interface StreamHttpExchange { + fun exchange(request: RequestEntity): List +} + +fun interface NdJsonStreamHttpExchange { + fun exchange(request: RequestEntity): List +} diff --git a/src/test/kotlin/testsupport/json/JSONUtil.kt b/src/test/kotlin/testsupport/json/JSONUtil.kt new file mode 100644 index 00000000..01e8e4e4 --- /dev/null +++ b/src/test/kotlin/testsupport/json/JSONUtil.kt @@ -0,0 +1,22 @@ +package testsupport.json + +import ee.carlrobert.codegpt.util.JsonMapper + +object JSONUtil { + + fun e(key: String, value: Any?): Pair = key to value + + fun jsonArray(vararg values: Any?): List = values.toList() + + fun jsonMap(vararg entries: Pair): Map = linkedMapOf(*entries) + + fun jsonMap(key: String, value: Any?): Map = linkedMapOf(key to value) + + fun jsonMapResponse(vararg entries: Pair): String { + return JsonMapper.mapper.writeValueAsString(jsonMap(*entries)) + } + + fun jsonMapResponse(key: String, value: Any?): String { + return JsonMapper.mapper.writeValueAsString(jsonMap(key, value)) + } +} diff --git a/src/test/kotlin/testsupport/mixin/ExternalServiceTestMixin.kt b/src/test/kotlin/testsupport/mixin/ExternalServiceTestMixin.kt new file mode 100644 index 00000000..b2c38aaa --- /dev/null +++ b/src/test/kotlin/testsupport/mixin/ExternalServiceTestMixin.kt @@ -0,0 +1,228 @@ +package testsupport.mixin + +import com.sun.net.httpserver.HttpExchange +import com.sun.net.httpserver.HttpServer +import com.fasterxml.jackson.core.type.TypeReference +import ee.carlrobert.codegpt.util.JsonMapper +import testsupport.http.RequestEntity +import testsupport.http.ResponseEntity +import testsupport.http.exchange.BasicHttpExchange +import testsupport.http.exchange.NdJsonStreamHttpExchange +import testsupport.http.exchange.StreamHttpExchange +import java.net.InetSocketAddress +import java.nio.charset.StandardCharsets +import java.util.concurrent.ConcurrentHashMap +import java.util.concurrent.ConcurrentLinkedQueue + +interface ExternalServiceTestMixin { + + fun expectOpenAI(exchange: StreamHttpExchange) { + ExternalServiceMockServer.expect(Service.OPENAI, exchange) + } + + fun expectOpenAI(exchange: BasicHttpExchange) { + ExternalServiceMockServer.expect(Service.OPENAI, exchange) + } + + fun expectCodeGPT(exchange: StreamHttpExchange) { + ExternalServiceMockServer.expect(Service.PROXYAI, exchange) + } + + fun expectOllama(exchange: NdJsonStreamHttpExchange) { + ExternalServiceMockServer.expect(Service.OLLAMA, exchange) + } + + fun expectOllama(exchange: BasicHttpExchange) { + ExternalServiceMockServer.expect(Service.OLLAMA, exchange) + } + + fun expectLlama(exchange: BasicHttpExchange) { + ExternalServiceMockServer.expect(Service.LLAMA_CPP, exchange) + } + + fun expectGoogle(exchange: BasicHttpExchange) { + ExternalServiceMockServer.expect(Service.GOOGLE, exchange) + } + + fun expectInception(exchange: StreamHttpExchange) { + ExternalServiceMockServer.expect(Service.INCEPTION, exchange) + } + + fun expectInception(exchange: BasicHttpExchange) { + ExternalServiceMockServer.expect(Service.INCEPTION, exchange) + } + + fun expectAnthropic(exchange: StreamHttpExchange) { + ExternalServiceMockServer.expect(Service.ANTHROPIC, exchange) + } + + fun expectMistral(exchange: StreamHttpExchange) { + ExternalServiceMockServer.expect(Service.MISTRAL, exchange) + } + + fun expectMistral(exchange: BasicHttpExchange) { + ExternalServiceMockServer.expect(Service.MISTRAL, exchange) + } + + fun expectCustomOpenAI(exchange: BasicHttpExchange) { + ExternalServiceMockServer.expect(Service.CUSTOM_OPENAI, exchange) + } + + fun expectCustomOpenAI(exchange: StreamHttpExchange) { + ExternalServiceMockServer.expect(Service.CUSTOM_OPENAI, exchange) + } + + companion object { + @JvmStatic + fun init() { + ExternalServiceMockServer.init() + } + + @JvmStatic + fun clearAll() { + ExternalServiceMockServer.clearAll() + } + } +} + +private enum class Service { + OPENAI, + PROXYAI, + OLLAMA, + LLAMA_CPP, + GOOGLE, + INCEPTION, + ANTHROPIC, + MISTRAL, + CUSTOM_OPENAI, +} + +private object ExternalServiceMockServer { + + private data class ServerState( + val server: HttpServer, + val expectations: ConcurrentLinkedQueue = ConcurrentLinkedQueue(), + ) + + private val servers = ConcurrentHashMap() + + fun init() { + if (servers.isNotEmpty()) { + return + } + Service.entries.forEach { service -> + val server = HttpServer.create(InetSocketAddress("127.0.0.1", 0), 0) + val state = ServerState(server) + server.createContext("/") { exchange -> + handleRequest(service, state, exchange) + } + server.start() + servers[service] = state + System.setProperty(baseUrlProperty(service), "http://127.0.0.1:${server.address.port}") + } + } + + fun clearAll() { + servers.values.forEach { it.expectations.clear() } + } + + fun expect(service: Service, exchange: Any) { + init() + servers.getValue(service).expectations.add(exchange) + } + + private fun handleRequest(service: Service, state: ServerState, exchange: HttpExchange) { + runCatching { + val expectation = state.expectations.poll() + ?: error("No expectation configured for $service ${exchange.requestMethod} ${exchange.requestURI.path}") + val request = toRequestEntity(exchange) + when (expectation) { + is BasicHttpExchange -> { + val response = expectation.exchange(request) + respondWithJson(exchange, response) + } + + is StreamHttpExchange -> { + val chunks = expectation.exchange(request) + respondAsSse(exchange, chunks) + } + + is NdJsonStreamHttpExchange -> { + val chunks = expectation.exchange(request) + respondAsNdJson(exchange, chunks) + } + + else -> error("Unsupported expectation type: ${expectation::class.java.name}") + } + }.onFailure { throwable -> + val body = (throwable.message ?: "Request handling failed") + .toByteArray(StandardCharsets.UTF_8) + exchange.responseHeaders["Content-Type"] = listOf("text/plain") + exchange.sendResponseHeaders(500, body.size.toLong()) + exchange.responseBody.use { out -> out.write(body) } + } + } + + private fun toRequestEntity(exchange: HttpExchange): RequestEntity { + val requestBody = exchange.requestBody.use { input -> + input.readBytes().toString(StandardCharsets.UTF_8) + } + val body = if (requestBody.isBlank()) { + emptyMap() + } else { + JsonMapper.mapper.readValue(requestBody, object : TypeReference>() {}) + } + return RequestEntity( + method = exchange.requestMethod, + uri = exchange.requestURI, + body = body, + ) + } + + private fun respondWithJson(exchange: HttpExchange, response: ResponseEntity) { + val bytes = response.body.toByteArray(StandardCharsets.UTF_8) + exchange.responseHeaders["Content-Type"] = listOf(response.contentType) + exchange.sendResponseHeaders(response.statusCode, bytes.size.toLong()) + exchange.responseBody.use { out -> out.write(bytes) } + } + + private fun respondAsSse(exchange: HttpExchange, chunks: List) { + exchange.responseHeaders["Content-Type"] = listOf("text/event-stream") + exchange.responseHeaders["Cache-Control"] = listOf("no-cache") + exchange.sendResponseHeaders(200, 0) + exchange.responseBody.bufferedWriter(StandardCharsets.UTF_8).use { writer -> + chunks.forEach { chunk -> + writer.write("data: ") + writer.write(chunk) + writer.write("\n\n") + writer.flush() + } + } + } + + private fun respondAsNdJson(exchange: HttpExchange, chunks: List) { + exchange.responseHeaders["Content-Type"] = listOf("application/x-ndjson") + exchange.sendResponseHeaders(200, 0) + exchange.responseBody.bufferedWriter(StandardCharsets.UTF_8).use { writer -> + chunks.forEach { chunk -> + writer.write(chunk) + writer.write("\n") + writer.flush() + } + } + } + + private fun baseUrlProperty(service: Service): String { + return when (service) { + Service.OPENAI -> "openai.baseUrl" + Service.PROXYAI -> "proxyai.baseUrl" + Service.OLLAMA -> "ollama.baseUrl" + Service.LLAMA_CPP -> "llama.baseUrl" + Service.GOOGLE -> "google.baseUrl" + Service.INCEPTION -> "inception.baseUrl" + Service.ANTHROPIC -> "anthropic.baseUrl" + Service.MISTRAL -> "mistral.baseUrl" + Service.CUSTOM_OPENAI -> "customOpenAI.baseUrl" + } + } +} diff --git a/src/test/kotlin/testsupport/mixin/ShortcutsTestMixin.kt b/src/test/kotlin/testsupport/mixin/ShortcutsTestMixin.kt index 3f74bd69..4d751de6 100644 --- a/src/test/kotlin/testsupport/mixin/ShortcutsTestMixin.kt +++ b/src/test/kotlin/testsupport/mixin/ShortcutsTestMixin.kt @@ -5,14 +5,16 @@ import com.intellij.testFramework.PlatformTestUtil import ee.carlrobert.codegpt.completions.HuggingFaceModel import ee.carlrobert.codegpt.credentials.CredentialsStore.CredentialKey.* import ee.carlrobert.codegpt.credentials.CredentialsStore.setCredential +import ee.carlrobert.codegpt.settings.models.ModelCatalog import ee.carlrobert.codegpt.settings.models.ModelSettings import ee.carlrobert.codegpt.settings.service.FeatureType import ee.carlrobert.codegpt.settings.service.ServiceType import ee.carlrobert.codegpt.settings.service.codegpt.CodeGPTServiceSettings +import ee.carlrobert.codegpt.settings.service.inception.InceptionSettings 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.google.models.GoogleModel +import ee.carlrobert.codegpt.settings.service.mistral.MistralSettings import java.util.function.BooleanSupplier interface ShortcutsTestMixin { @@ -27,20 +29,32 @@ interface ShortcutsTestMixin { } } - fun useOpenAIService(chatModel: String? = "gpt-4o", featureType: FeatureType = FeatureType.CHAT) { + fun useOpenAIService( + chatModel: String? = "gpt-4o", + featureType: FeatureType = FeatureType.CHAT + ) { setCredential(OpenaiApiKey, "TEST_API_KEY") val modelSettings = service() - + when (featureType) { FeatureType.CODE_COMPLETION -> { - modelSettings.setModel(FeatureType.CODE_COMPLETION, "gpt-3.5-turbo-instruct", ServiceType.OPENAI) + modelSettings.setModel( + FeatureType.CODE_COMPLETION, + ModelCatalog.GPT_3_5_TURBO_INSTRUCT, + ServiceType.OPENAI + ) } + else -> { modelSettings.setModel(featureType, chatModel ?: "gpt-4o", ServiceType.OPENAI) - modelSettings.setModel(FeatureType.CODE_COMPLETION, "gpt-3.5-turbo-instruct", ServiceType.OPENAI) + modelSettings.setModel( + FeatureType.CODE_COMPLETION, + ModelCatalog.GPT_3_5_TURBO_INSTRUCT, + ServiceType.OPENAI + ) } } - + service().state.run { isCodeCompletionsEnabled = true } @@ -53,18 +67,36 @@ interface ShortcutsTestMixin { LlamaSettings.getCurrentState().serverPort = null LlamaSettings.getCurrentState().isCodeCompletionsEnabled = codeCompletionsEnabled LlamaSettings.getCurrentState().huggingFaceModel = HuggingFaceModel.CODE_LLAMA_7B_Q4 - + val modelSettings = service() when (role) { FeatureType.CHAT -> { - modelSettings.setModel(FeatureType.CHAT, HuggingFaceModel.CODE_LLAMA_7B_Q4.code, ServiceType.LLAMA_CPP) + modelSettings.setModel( + FeatureType.CHAT, + HuggingFaceModel.CODE_LLAMA_7B_Q4.code, + ServiceType.LLAMA_CPP + ) } + FeatureType.CODE_COMPLETION -> { - modelSettings.setModel(FeatureType.CODE_COMPLETION, HuggingFaceModel.CODE_LLAMA_7B_Q4.code, ServiceType.LLAMA_CPP) + modelSettings.setModel( + FeatureType.CODE_COMPLETION, + HuggingFaceModel.CODE_LLAMA_7B_Q4.code, + ServiceType.LLAMA_CPP + ) } + else -> { - modelSettings.setModel(FeatureType.CHAT, HuggingFaceModel.CODE_LLAMA_7B_Q4.code, ServiceType.LLAMA_CPP) - modelSettings.setModel(FeatureType.CODE_COMPLETION, HuggingFaceModel.CODE_LLAMA_7B_Q4.code, ServiceType.LLAMA_CPP) + modelSettings.setModel( + FeatureType.CHAT, + HuggingFaceModel.CODE_LLAMA_7B_Q4.code, + ServiceType.LLAMA_CPP + ) + modelSettings.setModel( + FeatureType.CODE_COMPLETION, + HuggingFaceModel.CODE_LLAMA_7B_Q4.code, + ServiceType.LLAMA_CPP + ) } } } @@ -81,25 +113,97 @@ interface ShortcutsTestMixin { HuggingFaceModel.CODE_QWEN_2_5_3B_Q4_K_M.code ) } - + val modelSettings = service() when (role) { FeatureType.CHAT -> { - modelSettings.setModel(FeatureType.CHAT, HuggingFaceModel.LLAMA_3_8B_Q6_K.code, ServiceType.OLLAMA) + modelSettings.setModel( + FeatureType.CHAT, + HuggingFaceModel.LLAMA_3_8B_Q6_K.code, + ServiceType.OLLAMA + ) } + FeatureType.CODE_COMPLETION -> { - modelSettings.setModel(FeatureType.CODE_COMPLETION, HuggingFaceModel.CODE_QWEN_2_5_3B_Q4_K_M.code, ServiceType.OLLAMA) + modelSettings.setModel( + FeatureType.CODE_COMPLETION, + HuggingFaceModel.CODE_QWEN_2_5_3B_Q4_K_M.code, + ServiceType.OLLAMA + ) } + else -> { - modelSettings.setModel(FeatureType.CHAT, HuggingFaceModel.LLAMA_3_8B_Q6_K.code, ServiceType.OLLAMA) - modelSettings.setModel(FeatureType.CODE_COMPLETION, HuggingFaceModel.CODE_QWEN_2_5_3B_Q4_K_M.code, ServiceType.OLLAMA) + modelSettings.setModel( + FeatureType.CHAT, + HuggingFaceModel.LLAMA_3_8B_Q6_K.code, + ServiceType.OLLAMA + ) + modelSettings.setModel( + FeatureType.CODE_COMPLETION, + HuggingFaceModel.CODE_QWEN_2_5_3B_Q4_K_M.code, + ServiceType.OLLAMA + ) + } + } + } + + fun useMistralService(role: FeatureType = FeatureType.CHAT) { + setCredential(MistralApiKey, "TEST_API_KEY") + MistralSettings.getCurrentState().isCodeCompletionsEnabled = true + + val modelSettings = service() + when (role) { + FeatureType.CODE_COMPLETION -> { + modelSettings.setModel( + FeatureType.CODE_COMPLETION, + "codestral-latest", + ServiceType.MISTRAL + ) + } + + else -> { + modelSettings.setModel(role, "mistral-medium-latest", ServiceType.MISTRAL) + modelSettings.setModel( + FeatureType.CODE_COMPLETION, + "codestral-latest", + ServiceType.MISTRAL + ) + } + } + } + + fun useInceptionService(role: FeatureType = FeatureType.CHAT) { + setCredential(InceptionApiKey, "TEST_API_KEY") + service().state.codeCompletionsEnabled = true + + val modelSettings = service() + when (role) { + FeatureType.CODE_COMPLETION -> { + modelSettings.setModel( + FeatureType.CODE_COMPLETION, + ModelCatalog.MERCURY_CODER, + ServiceType.INCEPTION + ) + } + + else -> { + modelSettings.setModel(role, "mercury", ServiceType.INCEPTION) + modelSettings.setModel( + FeatureType.CODE_COMPLETION, + ModelCatalog.MERCURY_CODER, + ServiceType.INCEPTION + ) } } } fun useGoogleService(role: FeatureType = FeatureType.CHAT) { setCredential(GoogleApiKey, "TEST_API_KEY") - service().setModel(FeatureType.CHAT, GoogleModel.GEMINI_2_0_FLASH.code, ServiceType.GOOGLE) + service().setModel( + FeatureType.CHAT, + "gemini-2.5-flash", + ServiceType.GOOGLE + ) } fun waitExpecting(condition: BooleanSupplier?) { diff --git a/src/test/resources/codecompletions/psi/python.txt b/src/test/resources/codecompletions/psi/python.txt deleted file mode 100644 index 68333816..00000000 --- a/src/test/resources/codecompletions/psi/python.txt +++ /dev/null @@ -1,14 +0,0 @@ -import random -from typing import List - -class Util: - - test: str = "" - - def __init__(self, test: str): - self.test = test - - def randomStrings(n: int) -> List[str]: - return [str(random.randint(0, 100)) for _ in range(n)] - } -}