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

continuedev · May 18, 2024 · b8221eb · b8221eb
1 parent 430c6fe
commit b8221eb
Show file tree

Hide file tree

Showing 5 changed files with 48 additions and 31 deletions.
diff --git a/core/llm/countTokens.ts b/core/llm/countTokens.ts
@@ -57,7 +57,7 @@ function countTokens(
         : 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
@@ -217,14 +217,16 @@ ${settings}
 ${prompt}`;
   }
 
-  private _logTokensGenerated(model: string, completion: string) {
-    let 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,
-      tokens: tokens,
+      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> {
@@ -335,7 +337,7 @@ ${prompt}`;
       yield chunk;
     }
 
-    this._logTokensGenerated(completionOptions.model, completion);
+    this._logTokensGenerated(completionOptions.model, prompt, completion);
 
     if (log && this.writeLog) {
       await this.writeLog(`Completion:\n\n${completion}\n\n`);
@@ -370,7 +372,7 @@ ${prompt}`;
 
     const completion = await this._complete(prompt, completionOptions);
 
-    this._logTokensGenerated(completionOptions.model, completion);
+    this._logTokensGenerated(completionOptions.model, prompt, completion);
     if (log && this.writeLog) {
       await this.writeLog(`Completion:\n\n${completion}\n\n`);
     }
@@ -432,7 +434,7 @@ ${prompt}`;
       throw error;
     }
 
-    this._logTokensGenerated(completionOptions.model, completion);
+    this._logTokensGenerated(completionOptions.model, prompt, completion);
     if (log && this.writeLog) {
       await this.writeLog(`Completion:\n\n${completion}\n\n`);
     }
@@ -531,3 +533,4 @@ ${prompt}`;
     }
   }
 }
+
diff --git a/core/util/devdataSqlite.ts b/core/util/devdataSqlite.ts
@@ -12,41 +12,50 @@ export class DevDataSqliteDb {
             model TEXT NOT NULL,
             provider TEXT NOT NULL,
             tokens_generated INTEGER NOT NULL,
+            tokens_prompt INTEGER NOT NULL DEFAULT 0,
             timestamp DATETIME DEFAULT CURRENT_TIMESTAMP
         )`,
     );
+
+    // 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");
+    if (!columnExists) {
+      await db.exec(`ALTER TABLE tokens_generated ADD COLUMN tokens_prompt INTEGER NOT NULL DEFAULT 0;`);
+    }
   }
 
   public static async logTokensGenerated(
     model: string,
     provider: string,
-    tokens: number,
+    promptTokens: 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],
     );
   }
 
   public static async getTokensPerDay() {
     const db = await DevDataSqliteDb.get();
-    // Return a sum of tokens_generated column aggregated by day
     const result = await db?.all(
-      `SELECT date(timestamp) as day, sum(tokens_generated) as tokens
+      // Return a sum of tokens_generated and tokens_prompt columns aggregated by day
+      `SELECT date(timestamp) as day, sum(tokens_prompt) as promptTokens, sum(tokens_generated) as generatedTokens
         FROM tokens_generated
         GROUP BY date(timestamp)`,
-      // WHERE model = ? AND provider = ?
-      // [model, provider],
     );
     return result ?? [];
   }
 
   public static async getTokensPerModel() {
     const db = await DevDataSqliteDb.get();
-    // Return a sum of tokens_generated column aggregated by model
     const result = await db?.all(
-      `SELECT model, sum(tokens_generated) as tokens
+      // Return a sum of tokens_generated and tokens_prompt columns aggregated by model
+      `SELECT model, sum(tokens_prompt) as promptTokens, sum(tokens_generated) as generatedTokens
         FROM tokens_generated
         GROUP BY model`,
     );
@@ -70,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
@@ -37,18 +37,19 @@ function Stats() {
   useNavigationListener();
   const navigate = useNavigate();
 
-  const [days, setDays] = useState<{ day: string; tokens: number }[]>([]);
+  const [days, setDays] = useState<{ day: string; promptTokens: number; generatedTokens: number }[]>([]);
+  const [models, setModels] = useState<{ model: string; promptTokens: number; generatedTokens: number }[]>([]);
 
   useEffect(() => {
     ideRequest("stats/getTokensPerDay", undefined).then((days) => {
+      console.log("days", days);
       setDays(days);
     });
   }, []);
 
-  const [models, setModels] = useState<{ model: string; tokens: number }[]>([]);
-
   useEffect(() => {
     ideRequest("stats/getTokensPerModel", undefined).then((models) => {
+      console.log("models", models);
       setModels(models);
     });
   }, []);
@@ -75,8 +76,8 @@ function Stats() {
         <h2 className="ml-2">Tokens per Day</h2>
         <CopyButton
           text={generateTable(
-            ([["Day", "Tokens"]] as any).concat(
-              days.map((day) => [day.day, day.tokens]),
+            ([["Day", "Generated Tokens", "Prompt Tokens"]] as any).concat(
+              days.map((day) => [day.day, day.generatedTokens, day.promptTokens]),
             ),
           )}
         />
@@ -85,14 +86,16 @@ function Stats() {
         <thead>
           <Tr>
             <Th>Day</Th>
-            <Th>Tokens</Th>
+            <Th>Generated Tokens</Th>
+            <Th>Prompt Tokens</Th>
           </Tr>
         </thead>
         <tbody>
           {days.map((day) => (
             <Tr key={day.day} className="">
               <Td>{day.day}</Td>
-              <Td>{day.tokens}</Td>
+              <Td>{day.generatedTokens}</Td>
+              <Td>{day.promptTokens}</Td>
             </Tr>
           ))}
         </tbody>
@@ -102,8 +105,8 @@ function Stats() {
         <h2 className="ml-2">Tokens per Model</h2>
         <CopyButton
           text={generateTable(
-            ([["Model", "Tokens"]] as any).concat(
-              models.map((model) => [model.model, model.tokens]),
+            ([["Model", "Generated Tokens", "Prompt Tokens"]] as any).concat(
+              models.map((model) => [model.model, model.generatedTokens, model.promptTokens]),
             ),
           )}
         />
@@ -112,14 +115,16 @@ function Stats() {
         <thead>
           <Tr>
             <Th>Model</Th>
-            <Th>Tokens</Th>
+            <Th>Generated Tokens</Th>
+            <Th>Prompt Tokens</Th>
           </Tr>
         </thead>
         <tbody>
           {models.map((model) => (
             <Tr key={model.model} className="">
               <Td>{model.model}</Td>
-              <Td>{model.tokens}</Td>
+              <Td>{model.generatedTokens}</Td>
+              <Td>{model.promptTokens}</Td>
             </Tr>
           ))}
         </tbody>