Skip to content

Commit

Permalink
Add token counting and display functionality using js-tiktoken
Browse files Browse the repository at this point in the history
- Add js-tiktoken library to count tokens in response
- Update ResultView component to display token count instead of character count
  • Loading branch information
曾凯 committed Dec 20, 2024
1 parent bef1fdd commit 94c5feb
Show file tree
Hide file tree
Showing 3 changed files with 26 additions and 2 deletions.
10 changes: 10 additions & 0 deletions package-lock.json

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions package.json
Original file line number Diff line number Diff line change
Expand Up @@ -209,6 +209,7 @@
"@types/hjson": "^2.4.6",
"glob": "^11.0.0",
"hjson": "^3.2.2",
"js-tiktoken": "^1.0.16",
"lru-cache": "^11.0.1",
"markdown-it": "^14.1.0",
"md5": "^2.3.0",
Expand Down
17 changes: 15 additions & 2 deletions src/components/ResultView.tsx
Original file line number Diff line number Diff line change
@@ -1,10 +1,23 @@
import { Detail, ActionPanel, Action, Icon, Clipboard, closeMainWindow, showHUD } from "@raycast/api";
import React, { useMemo } from "react";
import MarkdownIt from "markdown-it";
import { getEncoding } from "js-tiktoken";

// 初始化 Markdown 解析器
const md = new MarkdownIt();

// 修改 token 计算函数
const countTokens = (text: string): number => {
try {
const encoder = getEncoding("cl100k_base");
const tokens = encoder.encode(text);
return tokens.length;
} catch (error) {
console.error('Token counting error:', error);
return 0;
}
};

// 提取代码块的辅助函数,限制最大代码块数量
const extractCodeBlocks = (text: string, maxBlocks: number = 10): string[] => {
const tokens = md.parse(text, {});
Expand Down Expand Up @@ -147,9 +160,9 @@ export function ResultView({
<Detail.Metadata.Label title="Top P" text={topP.toFixed(2)} />
<Detail.Metadata.Separator />
<Detail.Metadata.Label title="Duration" text={`${duration}s`} />
<Detail.Metadata.Label title="Response Length" text={`${response.length} chars`} />
<Detail.Metadata.Label title="Response Tokens" text={`${countTokens(response)} tokens`} />
</Detail.Metadata>
), [model, temperature, maxTokens, topP, duration, response.length]);
), [model, temperature, maxTokens, topP, duration, response]);

return (
<Detail
Expand Down

0 comments on commit 94c5feb

Please sign in to comment.