vrr/ProxyAI
2
0
Fork 0
mirror of https://github.com/carlrobertoh/ProxyAI.git synced 2026-05-22 11:38:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Display credit usage

Browse source
This commit is contained in:
Carl-Robert Linnupuu 2023-03-17 11:36:43 +00:00
parent ae47acfaa4
commit d236f2ca27
4 changed files with 95 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

29
src/main/java/ee/carlrobert/codegpt/client/Client.java
View file

@ -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)

28
src/main/java/ee/carlrobert/codegpt/client/CreditUsage.java Normal file
View file

@ -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;
}
}

8
src/main/java/ee/carlrobert/codegpt/ide/toolwindow/chat/ChatGptToolWindow.java
View file

@ -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());

33
src/main/java/ee/carlrobert/codegpt/ide/toolwindow/chat/actions/UsageToolbarLabelAction.java Normal file
View file

@ -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());
}
}