diff --git a/src/main/java/ee/carlrobert/codegpt/client/Client.java b/src/main/java/ee/carlrobert/codegpt/client/Client.java index e0ce69d4..1e6b17e1 100644 --- a/src/main/java/ee/carlrobert/codegpt/client/Client.java +++ b/src/main/java/ee/carlrobert/codegpt/client/Client.java @@ -4,7 +4,9 @@ import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.databind.ObjectMapper; import ee.carlrobert.codegpt.ide.conversations.Conversation; import ee.carlrobert.codegpt.ide.conversations.ConversationsState; +import ee.carlrobert.codegpt.ide.settings.SettingsState; import ee.carlrobert.codegpt.ide.settings.advanced.AdvancedSettingsState; +import java.io.IOException; import java.net.InetSocketAddress; import java.net.Proxy; import java.util.Map; @@ -19,9 +21,13 @@ import okhttp3.RequestBody; import okhttp3.sse.EventSource; import okhttp3.sse.EventSourceListener; import okhttp3.sse.EventSources; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; public abstract class Client { + private static final Logger LOG = LoggerFactory.getLogger(Client.class); + private final ObjectMapper objectMapper = new ObjectMapper(); private final ClientCode clientCode; private EventSource eventSource; @@ -58,6 +64,29 @@ public abstract class Client { getEventSourceListener(onMessageReceived, onComplete, onFailure)); } + public CreditUsage getCreditGrants() { + try { + var response = buildClient().newCall(new Request.Builder() + .url("https://api.openai.com/dashboard/billing/credit_grants") + .headers(Headers.of(Map.of( + "Content-Type", "application/json", + "Authorization", "Bearer " + SettingsState.getInstance().apiKey + ))) + .get() + .build()) + .execute(); + + if (response.body() == null) { + return null; + } + + return objectMapper.readValue(response.body().string(), CreditUsage.class); + } catch (IOException ex) { + LOG.error("Unable to retrieve credit info", ex); + return null; + } + } + public OkHttpClient buildClient() { OkHttpClient.Builder builder = new OkHttpClient.Builder() .connectTimeout(60, TimeUnit.SECONDS) diff --git a/src/main/java/ee/carlrobert/codegpt/client/CreditUsage.java b/src/main/java/ee/carlrobert/codegpt/client/CreditUsage.java new file mode 100644 index 00000000..a9369993 --- /dev/null +++ b/src/main/java/ee/carlrobert/codegpt/client/CreditUsage.java @@ -0,0 +1,28 @@ +package ee.carlrobert.codegpt.client; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonIgnoreProperties; +import com.fasterxml.jackson.annotation.JsonProperty; + +@JsonIgnoreProperties(ignoreUnknown = true) +public class CreditUsage { + + private final Double totalGranted; + private final Double totalUsed; + + @JsonCreator(mode = JsonCreator.Mode.PROPERTIES) + public CreditUsage( + @JsonProperty("total_granted") Double totalGranted, + @JsonProperty("total_used") Double totalUsed) { + this.totalGranted = totalGranted; + this.totalUsed = totalUsed; + } + + public Double getTotalGranted() { + return totalGranted; + } + + public Double getTotalUsed() { + return totalUsed; + } +} diff --git a/src/main/java/ee/carlrobert/codegpt/ide/toolwindow/chat/ChatGptToolWindow.java b/src/main/java/ee/carlrobert/codegpt/ide/toolwindow/chat/ChatGptToolWindow.java index 3a13719c..3a6dbd32 100644 --- a/src/main/java/ee/carlrobert/codegpt/ide/toolwindow/chat/ChatGptToolWindow.java +++ b/src/main/java/ee/carlrobert/codegpt/ide/toolwindow/chat/ChatGptToolWindow.java @@ -20,13 +20,14 @@ import ee.carlrobert.codegpt.ide.conversations.ConversationsState; import ee.carlrobert.codegpt.ide.settings.SettingsConfigurable; import ee.carlrobert.codegpt.ide.settings.SettingsState; import ee.carlrobert.codegpt.ide.toolwindow.ToolWindowService; +import ee.carlrobert.codegpt.ide.toolwindow.chat.actions.CreateNewConversationAction; +import ee.carlrobert.codegpt.ide.toolwindow.chat.actions.OpenInEditorAction; +import ee.carlrobert.codegpt.ide.toolwindow.chat.actions.UsageToolbarLabelAction; import ee.carlrobert.codegpt.ide.toolwindow.components.GenerateButton; import ee.carlrobert.codegpt.ide.toolwindow.components.LandingView; import ee.carlrobert.codegpt.ide.toolwindow.components.ScrollPane; import ee.carlrobert.codegpt.ide.toolwindow.components.SyntaxTextArea; import ee.carlrobert.codegpt.ide.toolwindow.components.TextArea; -import ee.carlrobert.codegpt.ide.toolwindow.chat.actions.CreateNewConversationAction; -import ee.carlrobert.codegpt.ide.toolwindow.chat.actions.OpenInEditorAction; import icons.Icons; import java.awt.Cursor; import java.awt.Dimension; @@ -69,8 +70,9 @@ public class ChatGptToolWindow { var actionGroup = new DefaultActionGroup("TOOLBAR_ACTION_GROUP", false); actionGroup.add(new CreateNewConversationAction()); - actionGroup.addSeparator(); actionGroup.add(new OpenInEditorAction()); + actionGroup.addSeparator(); + actionGroup.add(new UsageToolbarLabelAction()); // TODO: Data usage not enabled in stream mode https://community.openai.com/t/usage-info-in-api-responses/18862/11 // actionGroup.add(new TokenToolbarLabelAction()); diff --git a/src/main/java/ee/carlrobert/codegpt/ide/toolwindow/chat/actions/UsageToolbarLabelAction.java b/src/main/java/ee/carlrobert/codegpt/ide/toolwindow/chat/actions/UsageToolbarLabelAction.java new file mode 100644 index 00000000..b9fa995b --- /dev/null +++ b/src/main/java/ee/carlrobert/codegpt/ide/toolwindow/chat/actions/UsageToolbarLabelAction.java @@ -0,0 +1,33 @@ +package ee.carlrobert.codegpt.ide.toolwindow.chat.actions; + +import com.intellij.openapi.actionSystem.AnActionEvent; +import com.intellij.openapi.actionSystem.Presentation; +import com.intellij.openapi.actionSystem.ex.ToolbarLabelAction; +import com.intellij.util.ui.JBUI; +import ee.carlrobert.codegpt.client.ClientFactory; +import ee.carlrobert.codegpt.ide.settings.SettingsState; +import javax.swing.JComponent; +import org.jetbrains.annotations.NotNull; + +public class UsageToolbarLabelAction extends ToolbarLabelAction { + + @Override + public @NotNull JComponent createCustomComponent( + @NotNull Presentation presentation, + @NotNull String place) { + var client = new ClientFactory().getClient(); + var creditGrants = client.getCreditGrants(); + JComponent component = super.createCustomComponent(presentation, place); + component.setBorder(JBUI.Borders.empty(0, 2)); + if (creditGrants != null) { + presentation.setText(String.format("Credit used: %.2f/%.2f USD", creditGrants.getTotalUsed(), creditGrants.getTotalGranted())); + } + return component; + } + + @Override + public void update(@NotNull AnActionEvent event) { + super.update(event); + event.getPresentation().setVisible(!SettingsState.getInstance().apiKey.isEmpty()); + } +}