Преглед изворни кода

Merge pull request #608 from MuriloFP/feature/context-token-counter

Feature/context token counter
Matt Rubens пре 11 месеци
родитељ
комит
3862057a4d

+ 23 - 0
src/shared/getApiMetrics.ts

@@ -6,6 +6,7 @@ interface ApiMetrics {
 	totalCacheWrites?: number
 	totalCacheReads?: number
 	totalCost: number
+	contextTokens: number // Total tokens in conversation (last message's tokensIn + tokensOut)
 }
 
 /**
@@ -32,8 +33,22 @@ export function getApiMetrics(messages: ClineMessage[]): ApiMetrics {
 		totalCacheWrites: undefined,
 		totalCacheReads: undefined,
 		totalCost: 0,
+		contextTokens: 0,
 	}
 
+	// Find the last api_req_started message that has valid token information
+	const lastApiReq = [...messages].reverse().find((message) => {
+		if (message.type === "say" && message.say === "api_req_started" && message.text) {
+			try {
+				const parsedData = JSON.parse(message.text)
+				return typeof parsedData.tokensIn === "number" && typeof parsedData.tokensOut === "number"
+			} catch {
+				return false
+			}
+		}
+		return false
+	})
+
 	messages.forEach((message) => {
 		if (message.type === "say" && message.say === "api_req_started" && message.text) {
 			try {
@@ -55,6 +70,14 @@ export function getApiMetrics(messages: ClineMessage[]): ApiMetrics {
 				if (typeof cost === "number") {
 					result.totalCost += cost
 				}
+
+				// If this is the last api request, use its tokens for context size
+				if (message === lastApiReq) {
+					// Only update context tokens if both input and output tokens are non-zero
+					if (tokensIn > 0 && tokensOut > 0) {
+						result.contextTokens = tokensIn + tokensOut
+					}
+				}
 			} catch (error) {
 				console.error("Error parsing JSON:", error)
 			}

+ 1 - 0
webview-ui/src/components/chat/ChatView.tsx

@@ -915,6 +915,7 @@ const ChatView = ({ isHidden, showAnnouncement, hideAnnouncement, showHistoryVie
 					cacheWrites={apiMetrics.totalCacheWrites}
 					cacheReads={apiMetrics.totalCacheReads}
 					totalCost={apiMetrics.totalCost}
+					contextTokens={apiMetrics.contextTokens}
 					onClose={handleTaskCloseButtonClick}
 				/>
 			) : (

+ 9 - 0
webview-ui/src/components/chat/TaskHeader.tsx

@@ -16,6 +16,7 @@ interface TaskHeaderProps {
 	cacheWrites?: number
 	cacheReads?: number
 	totalCost: number
+	contextTokens: number
 	onClose: () => void
 }
 
@@ -27,6 +28,7 @@ const TaskHeader: React.FC<TaskHeaderProps> = ({
 	cacheWrites,
 	cacheReads,
 	totalCost,
+	contextTokens,
 	onClose,
 }) => {
 	const { apiConfiguration } = useExtensionState()
@@ -272,6 +274,13 @@ const TaskHeader: React.FC<TaskHeaderProps> = ({
 								{!isCostAvailable && <ExportButton />}
 							</div>
 
+							<div style={{ display: "flex", alignItems: "center", gap: "4px", flexWrap: "wrap" }}>
+								<span style={{ fontWeight: "bold" }}>Context:</span>
+								<span style={{ display: "flex", alignItems: "center", gap: "3px" }}>
+									{contextTokens ? formatLargeNumber(contextTokens) : '-'}
+								</span>
+							</div>
+
 							{shouldShowPromptCacheInfo && (cacheReads !== undefined || cacheWrites !== undefined) && (
 								<div style={{ display: "flex", alignItems: "center", gap: "4px", flexWrap: "wrap" }}>
 									<span style={{ fontWeight: "bold" }}>Cache:</span>