Log prompt-tokens to devdb and show in 'My Usage' view (#1309)

* Also log the number of prompt tokens to the dev-db * Show prompt tokens in 'My Usage' view
continuedev · Jun 23, 2024 · 7c5eb47 · 7c5eb47
1 parent 4d795e3
commit 7c5eb47
Show file tree

Hide file tree

Showing 5 changed files with 26 additions and 38 deletions.
diff --git a/core/llm/countTokens.ts b/core/llm/countTokens.ts
@@ -55,6 +55,8 @@ function countTokens(
         ? countImageTokens(part)
         : encoding.encode(part.text ?? "", "all", []).length;
     }, 0);
+  } else {
+    return encoding.encode(content ?? "", "all", []).length;
   }
   return encoding.encode(content, "all", []).length;
 }

diff --git a/core/llm/index.ts b/core/llm/index.ts
@@ -225,15 +225,16 @@ ${settings}
 ${prompt}`;
   }
 
-  private _logTokensGenerated(model: string, completion: string) {
-    const tokens = this.countTokens(completion);
+  private _logTokensGenerated(model: string, prompt: string, completion: string) {
+    let promptTokens = this.countTokens(prompt);
+    let generatedTokens = this.countTokens(completion);
     Telemetry.capture("tokens_generated", {
       model: model,
       provider: this.providerName,
       promptTokens: promptTokens,
       generatedTokens: generatedTokens,
     });
-    DevDataSqliteDb.logTokensGenerated(model, this.providerName, tokens);
+    DevDataSqliteDb.logTokensGenerated(model, this.providerName, promptTokens, generatedTokens);
   }
 
   fetch(url: RequestInfo | URL, init?: RequestInit): Promise<Response> {
@@ -594,3 +595,4 @@ ${prompt}`;
     return rendered;
   }
 }
+
diff --git a/core/util/devdataSqlite.ts b/core/util/devdataSqlite.ts
@@ -19,28 +19,24 @@ export class DevDataSqliteDb {
 
     // Add tokens_prompt column if it doesn't exist
     const columnCheckResult = await db.all(
-      `PRAGMA table_info(tokens_generated);`,
-    );
-    const columnExists = columnCheckResult.some(
-      (col: any) => col.name === "tokens_prompt",
+      `PRAGMA table_info(tokens_generated);`
     );
+    const columnExists = columnCheckResult.some((col: any) => col.name === "tokens_prompt");
     if (!columnExists) {
-      await db.exec(
-        `ALTER TABLE tokens_generated ADD COLUMN tokens_prompt INTEGER NOT NULL DEFAULT 0;`,
-      );
+      await db.exec(`ALTER TABLE tokens_generated ADD COLUMN tokens_prompt INTEGER NOT NULL DEFAULT 0;`);
     }
   }
 
   public static async logTokensGenerated(
     model: string,
     provider: string,
     promptTokens: number,
-    generatedTokens: number,
+    generatedTokens: number
   ) {
     const db = await DevDataSqliteDb.get();
     await db?.run(
-      "INSERT INTO tokens_generated (model, provider, tokens_generated) VALUES (?, ?, ?)",
-      [model, provider, tokens],
+      `INSERT INTO tokens_generated (model, provider, tokens_prompt, tokens_generated) VALUES (?, ?, ?, ?)`,
+      [model, provider, promptTokens, generatedTokens],
     );
   }
 
@@ -83,4 +79,4 @@ export class DevDataSqliteDb {
 
     return DevDataSqliteDb.db;
   }
-}
+}
diff --git a/core/web/webviewProtocol.ts b/core/web/webviewProtocol.ts
@@ -86,8 +86,8 @@ export type WebviewProtocol = Protocol &
     reloadWindow: [undefined, void];
     focusEditor: [undefined, void];
     toggleFullScreen: [undefined, void];
-    "stats/getTokensPerDay": [undefined, { day: string; tokens: number }[]];
-    "stats/getTokensPerModel": [undefined, { model: string; tokens: number }[]];
+    "stats/getTokensPerDay": [undefined, { day: string; promptTokens: number; generatedTokens: number }[]];
+    "stats/getTokensPerModel": [undefined, { model: string; promptTokens: number; generatedTokens: number }[]];
     insertAtCursor: [{ text: string }, void];
     copyText: [{ text: string }, void];
     "jetbrains/editorInsetHeight": [{ height: number }, void];

diff --git a/gui/src/pages/stats.tsx b/gui/src/pages/stats.tsx
@@ -38,25 +38,21 @@ function Stats() {
   const navigate = useNavigate();
   const ideMessenger = useContext(IdeMessengerContext);
 
-  const [days, setDays] = useState<
-    { day: string; promptTokens: number; generatedTokens: number }[]
-  >([]);
-  const [models, setModels] = useState<
-    { model: string; promptTokens: number; generatedTokens: number }[]
-  >([]);
+  const [days, setDays] = useState<{ day: string; promptTokens: number; generatedTokens: number }[]>([]);
+  const [models, setModels] = useState<{ model: string; promptTokens: number; generatedTokens: number }[]>([]);
 
   useEffect(() => {
-    ideMessenger.request("stats/getTokensPerDay", undefined).then((days) => {
+    ideRequest("stats/getTokensPerDay", undefined).then((days) => {
+      console.log("days", days);
       setDays(days);
     });
   }, []);
 
   useEffect(() => {
-    ideMessenger
-      .request("stats/getTokensPerModel", undefined)
-      .then((models) => {
-        setModels(models);
-      });
+    ideRequest("stats/getTokensPerModel", undefined).then((models) => {
+      console.log("models", models);
+      setModels(models);
+    });
   }, []);
 
   return (
@@ -82,11 +78,7 @@ function Stats() {
         <CopyButton
           text={generateTable(
             ([["Day", "Generated Tokens", "Prompt Tokens"]] as any).concat(
-              days.map((day) => [
-                day.day,
-                day.generatedTokens,
-                day.promptTokens,
-              ]),
+              days.map((day) => [day.day, day.generatedTokens, day.promptTokens]),
             ),
           )}
         />
@@ -115,11 +107,7 @@ function Stats() {
         <CopyButton
           text={generateTable(
             ([["Model", "Generated Tokens", "Prompt Tokens"]] as any).concat(
-              models.map((model) => [
-                model.model,
-                model.generatedTokens,
-                model.promptTokens,
-              ]),
+              models.map((model) => [model.model, model.generatedTokens, model.promptTokens]),
             ),
           )}
         />