From ae38713d5bd2bdc0b75aba21f204fdb7bda1af6e Mon Sep 17 00:00:00 2001 From: Saoud Rizwan <7799382+saoudrizwan@users.noreply.github.com> Date: Sun, 25 Aug 2024 19:52:22 -0400 Subject: [PATCH] Use 'user credits' header to update balance and show user under task header --- src/ClaudeDev.ts | 7 ++++- src/api/anthropic.ts | 16 ++++++---- src/api/bedrock.ts | 7 +++-- src/api/index.ts | 7 ++++- src/api/kodu.ts | 8 +++-- src/api/openrouter.ts | 6 ++-- src/providers/ClaudeDevProvider.ts | 4 +++ src/shared/kodu.ts | 2 +- webview-ui/src/App.tsx | 1 + webview-ui/src/components/ApiOptions.tsx | 10 +++--- webview-ui/src/components/ChatView.tsx | 5 +++ webview-ui/src/components/TaskHeader.tsx | 40 +++++++++++++++++++++++- 12 files changed, 89 insertions(+), 24 deletions(-) diff --git a/src/ClaudeDev.ts b/src/ClaudeDev.ts index 7a866c5..517a69f 100644 --- a/src/ClaudeDev.ts +++ b/src/ClaudeDev.ts @@ -1251,7 +1251,12 @@ ${this.customInstructions.trim()} this.apiConversationHistory, tools ) - return await this.api.createMessage(systemPrompt, adjustedMessages, tools) + const { message, userCredits } = await this.api.createMessage(systemPrompt, adjustedMessages, tools) + if (userCredits !== undefined) { + console.log("Updating kodu credits", userCredits) + this.providerRef.deref()?.updateKoduCredits(userCredits) + } + return message } catch (error) { const { response } = await this.ask( "api_req_failed", diff --git a/src/api/anthropic.ts b/src/api/anthropic.ts index ff74608..f43496f 100644 --- a/src/api/anthropic.ts +++ b/src/api/anthropic.ts @@ -1,5 +1,5 @@ import { Anthropic } from "@anthropic-ai/sdk" -import { ApiHandler, withoutImageData } from "." +import { ApiHandler, ApiHandlerMessageResponse, withoutImageData } from "." import { anthropicDefaultModelId, AnthropicModelId, anthropicModels, ApiHandlerOptions, ModelInfo } from "../shared/api" export class AnthropicHandler implements ApiHandler { @@ -15,12 +15,12 @@ export class AnthropicHandler implements ApiHandler { systemPrompt: string, messages: Anthropic.Messages.MessageParam[], tools: Anthropic.Messages.Tool[] - ): Promise { + ): Promise { const modelId = this.getModel().id switch (modelId) { case "claude-3-5-sonnet-20240620": case "claude-3-opus-20240229": - case "claude-3-haiku-20240307": + case "claude-3-haiku-20240307": { /* The latest message will be the new user message, one before will be the assistant message from a previous request, and the user message before that will be a previously cached user message. So we need to mark the latest user message as ephemeral to cache it for the next request, and mark the second to last user message as ephemeral to let the server know the last message to retrieve from the cache for the current request.. */ @@ -30,7 +30,7 @@ export class AnthropicHandler implements ApiHandler { ) const lastUserMsgIndex = userMsgIndices[userMsgIndices.length - 1] ?? -1 const secondLastMsgUserIndex = userMsgIndices[userMsgIndices.length - 2] ?? -1 - return await this.client.beta.promptCaching.messages.create( + const message = await this.client.beta.promptCaching.messages.create( { model: modelId, max_tokens: this.getModel().info.maxTokens, @@ -80,8 +80,10 @@ export class AnthropicHandler implements ApiHandler { } })() ) - default: - return await this.client.messages.create({ + return { message } + } + default: { + const message = await this.client.messages.create({ model: modelId, max_tokens: this.getModel().info.maxTokens, system: [{ text: systemPrompt, type: "text" }], @@ -89,6 +91,8 @@ export class AnthropicHandler implements ApiHandler { tools, tool_choice: { type: "auto" }, }) + return { message } + } } } diff --git a/src/api/bedrock.ts b/src/api/bedrock.ts index 975ba3f..5a9db5f 100644 --- a/src/api/bedrock.ts +++ b/src/api/bedrock.ts @@ -1,6 +1,6 @@ import AnthropicBedrock from "@anthropic-ai/bedrock-sdk" import { Anthropic } from "@anthropic-ai/sdk" -import { ApiHandler, withoutImageData } from "." +import { ApiHandler, ApiHandlerMessageResponse, withoutImageData } from "." import { ApiHandlerOptions, bedrockDefaultModelId, BedrockModelId, bedrockModels, ModelInfo } from "../shared/api" // https://docs.anthropic.com/en/api/claude-on-amazon-bedrock @@ -26,8 +26,8 @@ export class AwsBedrockHandler implements ApiHandler { systemPrompt: string, messages: Anthropic.Messages.MessageParam[], tools: Anthropic.Messages.Tool[] - ): Promise { - return await this.client.messages.create({ + ): Promise { + const message = await this.client.messages.create({ model: this.getModel().id, max_tokens: this.getModel().info.maxTokens, system: systemPrompt, @@ -35,6 +35,7 @@ export class AwsBedrockHandler implements ApiHandler { tools, tool_choice: { type: "auto" }, }) + return { message } } createUserReadableRequest( diff --git a/src/api/index.ts b/src/api/index.ts index f50d43b..c287000 100644 --- a/src/api/index.ts +++ b/src/api/index.ts @@ -5,12 +5,17 @@ import { AwsBedrockHandler } from "./bedrock" import { OpenRouterHandler } from "./openrouter" import { KoduHandler } from "./kodu" +export interface ApiHandlerMessageResponse { + message: Anthropic.Messages.Message + userCredits?: number +} + export interface ApiHandler { createMessage( systemPrompt: string, messages: Anthropic.Messages.MessageParam[], tools: Anthropic.Messages.Tool[] - ): Promise + ): Promise createUserReadableRequest( userContent: Array< diff --git a/src/api/kodu.ts b/src/api/kodu.ts index ecba2ea..20af022 100644 --- a/src/api/kodu.ts +++ b/src/api/kodu.ts @@ -1,6 +1,6 @@ import { Anthropic } from "@anthropic-ai/sdk" import axios from "axios" -import { ApiHandler, withoutImageData } from "." +import { ApiHandler, ApiHandlerMessageResponse, withoutImageData } from "." import { ApiHandlerOptions, koduDefaultModelId, KoduModelId, koduModels, ModelInfo } from "../shared/api" import { getKoduCreditsUrl, getKoduInferenceUrl } from "../shared/kodu" @@ -24,7 +24,7 @@ export class KoduHandler implements ApiHandler { systemPrompt: string, messages: Anthropic.Messages.MessageParam[], tools: Anthropic.Messages.Tool[] - ): Promise { + ): Promise { const modelId = this.getModel().id let requestBody: Anthropic.Beta.PromptCaching.Messages.MessageCreateParamsNonStreaming switch (modelId) { @@ -82,7 +82,9 @@ export class KoduHandler implements ApiHandler { "x-api-key": this.options.koduApiKey, }, }) - return response.data + const message = response.data + const userCredits = response.headers["user-credits"] + return { message, userCredits: userCredits !== undefined ? parseFloat(userCredits) : undefined } } createUserReadableRequest( diff --git a/src/api/openrouter.ts b/src/api/openrouter.ts index cd74934..992181d 100644 --- a/src/api/openrouter.ts +++ b/src/api/openrouter.ts @@ -1,6 +1,6 @@ import { Anthropic } from "@anthropic-ai/sdk" import OpenAI from "openai" -import { ApiHandler, withoutImageData } from "." +import { ApiHandler, ApiHandlerMessageResponse, withoutImageData } from "." import { ApiHandlerOptions, ModelInfo, @@ -30,7 +30,7 @@ export class OpenRouterHandler implements ApiHandler { systemPrompt: string, messages: Anthropic.Messages.MessageParam[], tools: Anthropic.Messages.Tool[] - ): Promise { + ): Promise { // Convert Anthropic messages to OpenAI format const openAiMessages: OpenAI.Chat.ChatCompletionMessageParam[] = [ { role: "system", content: systemPrompt }, @@ -120,7 +120,7 @@ export class OpenRouterHandler implements ApiHandler { ) } - return anthropicMessage + return { message: anthropicMessage } } /* diff --git a/src/providers/ClaudeDevProvider.ts b/src/providers/ClaudeDevProvider.ts index 1a9a414..bea13ca 100644 --- a/src/providers/ClaudeDevProvider.ts +++ b/src/providers/ClaudeDevProvider.ts @@ -696,6 +696,10 @@ export class ClaudeDevProvider implements vscode.WebviewViewProvider { return history } + async updateKoduCredits(credits: number) { + await this.updateGlobalState("koduCredits", credits) + } + // global private async updateGlobalState(key: GlobalStateKey, value: any) { diff --git a/src/shared/kodu.ts b/src/shared/kodu.ts index b61ff03..cfb40bd 100644 --- a/src/shared/kodu.ts +++ b/src/shared/kodu.ts @@ -1,4 +1,4 @@ -const KODU_BASE_URL = "https://claude-dev.com" +const KODU_BASE_URL = "https://kodu.ai" export function getKoduSignInUrl(uriScheme?: string) { return `${KODU_BASE_URL}/auth/login?redirectTo=${uriScheme}://saoudrizwan.claude-dev&ext=1` diff --git a/webview-ui/src/App.tsx b/webview-ui/src/App.tsx index db14491..3fa1037 100644 --- a/webview-ui/src/App.tsx +++ b/webview-ui/src/App.tsx @@ -169,6 +169,7 @@ const App: React.FC = () => { apiConfiguration={apiConfiguration} vscodeUriScheme={vscodeUriScheme} shouldShowKoduPromo={shouldShowKoduPromo} + koduCredits={koduCredits} /> )} diff --git a/webview-ui/src/components/ApiOptions.tsx b/webview-ui/src/components/ApiOptions.tsx index dd063ad..d577aeb 100644 --- a/webview-ui/src/components/ApiOptions.tsx +++ b/webview-ui/src/components/ApiOptions.tsx @@ -36,7 +36,7 @@ const ApiOptions: React.FC = ({ apiErrorMessage, vscodeUriScheme, }) => { - const [didFetchKoduCredits, setDidFetchKoduCredits] = useState(false) + const [, setDidFetchKoduCredits] = useState(false) const handleInputChange = (field: keyof ApiConfiguration) => (event: any) => { setApiConfiguration((prev) => ({ ...prev, [field]: event.target.value })) } @@ -78,11 +78,11 @@ const ApiOptions: React.FC = ({ } useEffect(() => { - if (selectedProvider === "kodu" && apiConfiguration?.koduApiKey) { + if (selectedProvider === "kodu" && apiConfiguration?.koduApiKey && koduCredits === undefined) { setDidFetchKoduCredits(false) vscode.postMessage({ type: "fetchKoduCredits" }) } - }, [selectedProvider, apiConfiguration?.koduApiKey]) + }, [selectedProvider, apiConfiguration?.koduApiKey, koduCredits]) const handleMessage = useCallback((e: MessageEvent) => { const message: ExtensionMessage = e.data @@ -178,7 +178,7 @@ const ApiOptions: React.FC = ({
Credits remaining:{" "} - + {formatPrice(koduCredits || 0)}
@@ -307,7 +307,7 @@ const ApiOptions: React.FC = ({ ) } -const formatPrice = (price: number) => { +export const formatPrice = (price: number) => { return new Intl.NumberFormat("en-US", { style: "currency", currency: "USD", diff --git a/webview-ui/src/components/ChatView.tsx b/webview-ui/src/components/ChatView.tsx index 1d38958..a6e8414 100644 --- a/webview-ui/src/components/ChatView.tsx +++ b/webview-ui/src/components/ChatView.tsx @@ -33,6 +33,7 @@ interface ChatViewProps { apiConfiguration?: ApiConfiguration vscodeUriScheme?: string shouldShowKoduPromo: boolean + koduCredits?: number } const MAX_IMAGES_PER_MESSAGE = 20 // Anthropic limits to 20 images @@ -51,6 +52,7 @@ const ChatView = ({ apiConfiguration, vscodeUriScheme, shouldShowKoduPromo, + koduCredits, }: ChatViewProps) => { //const task = messages.length > 0 ? (messages[0].say === "task" ? messages[0] : undefined) : undefined) : undefined const task = messages.length > 0 ? messages[0] : undefined // leaving this less safe version here since if the first message is not a task, then the extension is in a bad state and needs to be debugged (see ClaudeDev.abort) @@ -487,6 +489,9 @@ const ChatView = ({ totalCost={apiMetrics.totalCost} onClose={handleTaskCloseButtonClick} isHidden={isHidden} + koduCredits={koduCredits} + vscodeUriScheme={vscodeUriScheme} + apiProvider={apiConfiguration?.apiProvider} /> ) : ( <> diff --git a/webview-ui/src/components/TaskHeader.tsx b/webview-ui/src/components/TaskHeader.tsx index bc1dd35..89d5bc1 100644 --- a/webview-ui/src/components/TaskHeader.tsx +++ b/webview-ui/src/components/TaskHeader.tsx @@ -1,9 +1,12 @@ -import { VSCodeButton } from "@vscode/webview-ui-toolkit/react" +import { VSCodeButton, VSCodeLink } from "@vscode/webview-ui-toolkit/react" import React, { useEffect, useRef, useState } from "react" import { useWindowSize } from "react-use" import { ClaudeMessage } from "../../../src/shared/ExtensionMessage" import { vscode } from "../utils/vscode" import Thumbnails from "./Thumbnails" +import { formatPrice } from "./ApiOptions" +import { getKoduAddCreditsUrl } from "../../../src/shared/kodu" +import { ApiProvider } from "../../../src/shared/api" interface TaskHeaderProps { task: ClaudeMessage @@ -15,6 +18,9 @@ interface TaskHeaderProps { totalCost: number onClose: () => void isHidden: boolean + koduCredits?: number + vscodeUriScheme?: string + apiProvider?: ApiProvider } const TaskHeader: React.FC = ({ @@ -27,6 +33,9 @@ const TaskHeader: React.FC = ({ totalCost, onClose, isHidden, + koduCredits, + vscodeUriScheme, + apiProvider, }) => { const [isExpanded, setIsExpanded] = useState(false) const [showSeeMore, setShowSeeMore] = useState(false) @@ -107,6 +116,7 @@ const TaskHeader: React.FC = ({ flexDirection: "column", gap: "8px", position: "relative", + zIndex: 1, }}>
= ({
+ {apiProvider === "kodu" && ( +
+
Credits Remaining:
+
+ {formatPrice(koduCredits || 0)} + {(koduCredits || 0) < 1 && ( + <> + {" "} + + (get more?) + + + )} +
+
+ )} ) }