Show token bar at bottom of chat input (#33)

src/ipc/handlers/token_count_handlers.ts

@@ -0,0 +1,127 @@
+import { ipcMain } from "electron";
+import { db } from "../../db";
+import { chats, messages } from "../../db/schema";
+import { eq } from "drizzle-orm";
+import { SYSTEM_PROMPT } from "../../prompts/system_prompt";
+import {
+  SUPABASE_AVAILABLE_SYSTEM_PROMPT,
+  SUPABASE_NOT_AVAILABLE_SYSTEM_PROMPT,
+} from "../../prompts/supabase_prompt";
+import { getDyadAppPath } from "../../paths/paths";
+import log from "electron-log";
+import { extractCodebase } from "../../utils/codebase";
+import { getSupabaseContext } from "../../supabase_admin/supabase_context";
+import { readSettings } from "../../main/settings";
+import { MODEL_OPTIONS } from "../../constants/models";
+import { TokenCountParams } from "../ipc_types";
+import { TokenCountResult } from "../ipc_types";
+
+const logger = log.scope("token_count_handlers");
+
+// Estimate tokens (4 characters per token)
+const estimateTokens = (text: string): number => {
+  return Math.ceil(text.length / 4);
+};
+
+export function registerTokenCountHandlers() {
+  ipcMain.handle(
+    "chat:count-tokens",
+    async (event, req: TokenCountParams): Promise<TokenCountResult> => {
+      try {
+        // Get the chat with messages
+        const chat = await db.query.chats.findFirst({
+          where: eq(chats.id, req.chatId),
+          with: {
+            messages: {
+              orderBy: (messages, { asc }) => [asc(messages.createdAt)],
+            },
+            app: true,
+          },
+        });
+
+        if (!chat) {
+          throw new Error(`Chat not found: ${req.chatId}`);
+        }
+
+        // Prepare message history for token counting
+        const messageHistory = chat.messages
+          .map((message) => message.content)
+          .join("");
+        const messageHistoryTokens = estimateTokens(messageHistory);
+
+        // Count input tokens
+        const inputTokens = estimateTokens(req.input);
+
+        // Count system prompt tokens
+        let systemPrompt = SYSTEM_PROMPT;
+        let supabaseContext = "";
+
+        if (chat.app?.supabaseProjectId) {
+          systemPrompt += "\n\n" + SUPABASE_AVAILABLE_SYSTEM_PROMPT;
+          supabaseContext = await getSupabaseContext({
+            supabaseProjectId: chat.app.supabaseProjectId,
+          });
+        } else {
+          systemPrompt += "\n\n" + SUPABASE_NOT_AVAILABLE_SYSTEM_PROMPT;
+        }
+
+        const systemPromptTokens = estimateTokens(
+          systemPrompt + supabaseContext
+        );
+
+        // Extract codebase information if app is associated with the chat
+        let codebaseInfo = "";
+        let codebaseTokens = 0;
+
+        if (chat.app) {
+          const appPath = getDyadAppPath(chat.app.path);
+          try {
+            codebaseInfo = await extractCodebase(appPath);
+            codebaseTokens = estimateTokens(codebaseInfo);
+            logger.log(
+              `Extracted codebase information from ${appPath}, tokens: ${codebaseTokens}`
+            );
+          } catch (error) {
+            logger.error("Error extracting codebase:", error);
+          }
+        }
+
+        // Calculate total tokens
+        const totalTokens =
+          messageHistoryTokens +
+          inputTokens +
+          systemPromptTokens +
+          codebaseTokens;
+
+        return {
+          totalTokens,
+          messageHistoryTokens,
+          codebaseTokens,
+          inputTokens,
+          systemPromptTokens,
+          contextWindow: getContextWindow(),
+        };
+      } catch (error) {
+        logger.error("Error counting tokens:", error);
+        throw error;
+      }
+    }
+  );
+}
+
+const DEFAULT_CONTEXT_WINDOW = 128_000;
+
+function getContextWindow() {
+  const settings = readSettings();
+  const model = settings.selectedModel;
+  if (!MODEL_OPTIONS[model.provider as keyof typeof MODEL_OPTIONS]) {
+    logger.warn(
+      `Model provider ${model.provider} not found in MODEL_OPTIONS. Using default max tokens.`
+    );
+    return DEFAULT_CONTEXT_WINDOW;
+  }
+  const modelOption = MODEL_OPTIONS[
+    model.provider as keyof typeof MODEL_OPTIONS
+  ].find((m) => m.name === model.name);
+  return modelOption?.contextWindow || DEFAULT_CONTEXT_WINDOW;
+}