return usage stats from providers

Author: dlants

node/providers/anthropic.ts

@@ -6,6 +6,7 @@ import {
   type StopReason,
   type Provider,
   type ProviderMessage,
+  type Usage,
 } from "./provider.ts";
 import type { ToolRequestId } from "../tools/toolManager.ts";
 import { assertUnreachable } from "../utils/assertUnreachable.ts";
@@ -49,6 +50,7 @@ export class AnthropicProvider implements Provider {
   ): Promise<{
     toolRequests: Result<ToolManager.ToolRequest, { rawRequest: unknown }>[];
     stopReason: StopReason;
+    usage: Usage;
   }> {
     const buf: string[] = [];
     let flushInProgress: boolean = false;
@@ -225,9 +227,22 @@ export class AnthropicProvider implements Provider {
           return extendError(result, { rawRequest: req });
         });
 
-      this.nvim.logger?.debug("toolRequests: " + JSON.stringify(toolRequests));
-      this.nvim.logger?.debug("stopReason: " + response.stop_reason);
-      return { toolRequests, stopReason: response.stop_reason || "end_turn" };
+      const usage: Usage = {
+        inputTokens: response.usage.input_tokens,
+        outputTokens: response.usage.output_tokens,
+      };
+      if (response.usage.cache_read_input_tokens) {
+        usage.cacheHits = response.usage.cache_read_input_tokens;
+      }
+      if (response.usage.cache_creation_input_tokens) {
+        usage.cacheMisses = response.usage.cache_creation_input_tokens;
+      }
+
+      return {
+        toolRequests,
+        stopReason: response.stop_reason || "end_turn",
+        usage,
+      };
     } finally {
       this.request = undefined;
     }

node/providers/mock.ts

@@ -6,6 +6,7 @@ import {
   type Provider,
   type ProviderMessage,
   type StopReason,
+  type Usage,
 } from "./provider.ts";
 
 type MockRequest = {
@@ -15,6 +16,7 @@ type MockRequest = {
   defer: Defer<{
     toolRequests: Result<ToolRequest, { rawRequest: unknown }>[];
     stopReason: StopReason;
+    usage: Usage;
   }>;
 };
 
@@ -28,6 +30,10 @@ export class MockProvider implements Provider {
         lastRequest.defer.resolve({
           toolRequests: [],
           stopReason: "end_turn",
+          usage: {
+            inputTokens: 0,
+            outputTokens: 0,
+          },
         });
       }
     }
@@ -40,6 +46,7 @@ export class MockProvider implements Provider {
   ): Promise<{
     toolRequests: Result<ToolRequest, { rawRequest: unknown }>[];
     stopReason: StopReason;
+    usage: Usage;
   }> {
     const request: MockRequest = {
       messages,
@@ -79,6 +86,10 @@ export class MockProvider implements Provider {
     lastRequest.defer.resolve({
       toolRequests,
       stopReason,
+      usage: {
+        inputTokens: 0,
+        outputTokens: 0,
+      },
     });
   }
 }

node/providers/openai.ts

@@ -1,7 +1,12 @@
 import OpenAI from "openai";
 import * as ToolManager from "../tools/toolManager.ts";
 import { extendError, type Result } from "../utils/result.ts";
-import type { StopReason, Provider, ProviderMessage } from "./provider.ts";
+import type {
+  StopReason,
+  Provider,
+  ProviderMessage,
+  Usage,
+} from "./provider.ts";
 import { assertUnreachable } from "../utils/assertUnreachable.ts";
 import type { ToolName, ToolRequestId } from "../tools/toolManager.ts";
 import type { Nvim } from "nvim-node";
@@ -46,6 +51,7 @@ export class OpenAIProvider implements Provider {
   ): Promise<{
     toolRequests: Result<ToolManager.ToolRequest, { rawRequest: unknown }>[];
     stopReason: StopReason;
+    usage: Usage;
   }> {
     const openaiMessages: OpenAI.ChatCompletionMessageParam[] = [
       {
@@ -148,7 +154,9 @@ export class OpenAIProvider implements Provider {
 
       const toolRequests = [];
       let stopReason: StopReason | undefined;
+      let lastChunk: OpenAI.ChatCompletionChunk | undefined;
       for await (const chunk of stream) {
+        lastChunk = chunk;
         const choice = chunk.choices[0];
         if (choice.delta.content) {
           onText(choice.delta.content);
@@ -229,6 +237,15 @@ export class OpenAIProvider implements Provider {
             return extendError(result, { rawRequest: req });
           }),
         stopReason: stopReason || "end_turn",
+        usage: lastChunk?.usage
+          ? {
+              inputTokens: lastChunk.usage.prompt_tokens,
+              outputTokens: lastChunk.usage.completion_tokens,
+            }
+          : {
+              inputTokens: 0,
+              outputTokens: 0,
+            },
       };
     } finally {
       this.request = undefined;

node/providers/provider.ts

@@ -17,6 +17,13 @@ export type StopReason =
   | "content"
   | "stop_sequence";
 
+export type Usage = {
+  inputTokens: number;
+  outputTokens: number;
+  cacheHits?: number;
+  cacheMisses?: number;
+};
+
 export type ProviderMessage = {
   role: "user" | "assistant";
   content: string | Array<ProviderMessageContent>;
@@ -57,6 +64,7 @@ export interface Provider {
   ): Promise<{
     toolRequests: Result<ToolManager.ToolRequest, { rawRequest: unknown }>[];
     stopReason: StopReason;
+    usage: Usage;
   }>;
 
   abort(): void;