Skip to content

Commit 78505ed

Browse files
thecodacusthecodacus
authored
Merge pull request #769 from thecodacus/token-usage
feat: Show token usage on LLM call assistant message
2 parents dd24ccc + a155447 commit 78505ed

File tree

11 files changed

+166
-374
lines changed

11 files changed

+166
-374
lines changed

‎app/commit.json‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +1,2 @@
1-
{ "commit": "77073a5e7f759ae8e5752628131d0c56df6b5c34" , "version": "" }
1+
{ "commit": "77073a5e7f759ae8e5752628131d0c56df6b5c34" , "version": "0.0.1" }
2+

‎app/components/chat/AssistantMessage.tsx‎

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,30 @@
11
import { memo } from 'react';
22
import { Markdown } from './Markdown';
3+
import type { JSONValue } from 'ai';
34

45
interface AssistantMessageProps {
56
content: string;
7+
annotations?: JSONValue[];
68
}
79

8-
export const AssistantMessage = memo(({ content }: AssistantMessageProps) => {
10+
export const AssistantMessage = memo(({ content, annotations }: AssistantMessageProps) => {
11+
const filteredAnnotations = (annotations?.filter(
12+
(annotation: JSONValue) => annotation && typeof annotation === 'object' && Object.keys(annotation).includes('type'),
13+
) || []) as { type: string; value: any }[];
14+
15+
const usage: {
16+
completionTokens: number;
17+
promptTokens: number;
18+
totalTokens: number;
19+
} = filteredAnnotations.find((annotation) => annotation.type === 'usage')?.value;
20+
921
return (
1022
<div className="overflow-hidden w-full">
23+
{usage && (
24+
<div className="text-sm text-bolt-elements-textSecondary mb-2">
25+
Tokens: {usage.totalTokens} (prompt: {usage.promptTokens}, completion: {usage.completionTokens})
26+
</div>
27+
)}
1128
<Markdown html>{content}</Markdown>
1229
</div>
1330
);

‎app/components/chat/Chat.client.tsx‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -116,13 +116,22 @@ export const ChatImpl = memo(
116116
apiKeys,
117117
files,
118118
},
119+
sendExtraMessageFields: true,
119120
onError: (error) => {
120121
logger.error('Request failed\n\n', error);
121122
toast.error(
122123
'There was an error processing your request: ' + (error.message ? error.message : 'No details were returned'),
123124
);
124125
},
125-
onFinish: () => {
126+
onFinish: (message, response) => {
127+
const usage = response.usage;
128+
129+
if (usage) {
130+
console.log('Token usage:', usage);
131+
132+
// You can now use the usage data as needed
133+
}
134+
126135
logger.debug('Finished streaming');
127136
},
128137
initialMessages,

‎app/components/chat/Messages.client.tsx‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -65,7 +65,11 @@ export const Messages = React.forwardRef<HTMLDivElement, MessagesProps>((props:
6565
</div>
6666
)}
6767
<div className="grid grid-col-1 w-full">
68-
{isUserMessage ? <UserMessage content={content} /> : <AssistantMessage content={content} />}
68+
{isUserMessage ? (
69+
<UserMessage content={content} />
70+
) : (
71+
<AssistantMessage content={content} annotations={message.annotations} />
72+
)}
6973
</div>
7074
{!isUserMessage && (
7175
<div className="flex gap-2 flex-col lg:flex-row">

‎app/components/chat/UserMessage.tsx‎

Lines changed: 15 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -12,42 +12,36 @@ interface UserMessageProps {
1212
export function UserMessage({ content }: UserMessageProps) {
1313
if (Array.isArray(content)) {
1414
const textItem = content.find((item) => item.type === 'text');
15-
const textContent = sanitizeUserMessage(textItem?.text || '');
15+
const textContent = stripMetadata(textItem?.text || '');
1616
const images = content.filter((item) => item.type === 'image' && item.image);
1717

1818
return (
1919
<div className="overflow-hidden pt-[4px]">
20-
<div className="flex items-start gap-4">
21-
<div className="flex-1">
22-
<Markdown limitedMarkdown>{textContent}</Markdown>
23-
</div>
24-
{images.length > 0 && (
25-
<div className="flex-shrink-0 w-[160px]">
26-
{images.map((item, index) => (
27-
<div key={index} className="relative">
28-
<img
29-
src={item.image}
30-
alt={`Uploaded image ${index + 1}`}
31-
className="w-full h-[160px] rounded-lg object-cover border border-bolt-elements-borderColor"
32-
/>
33-
</div>
34-
))}
35-
</div>
36-
)}
20+
<div className="flex flex-col gap-4">
21+
{textContent && <Markdown html>{textContent}</Markdown>}
22+
{images.map((item, index) => (
23+
<img
24+
key={index}
25+
src={item.image}
26+
alt={`Image ${index + 1}`}
27+
className="max-w-full h-auto rounded-lg"
28+
style={{ maxHeight: '512px', objectFit: 'contain' }}
29+
/>
30+
))}
3731
</div>
3832
</div>
3933
);
4034
}
4135

42-
const textContent = sanitizeUserMessage(content);
36+
const textContent = stripMetadata(content);
4337

4438
return (
4539
<div className="overflow-hidden pt-[4px]">
46-
<Markdown limitedMarkdown>{textContent}</Markdown>
40+
<Markdown html>{textContent}</Markdown>
4741
</div>
4842
);
4943
}
5044

51-
function sanitizeUserMessage(content: string) {
45+
function stripMetadata(content: string) {
5246
return content.replace(MODEL_REGEX, '').replace(PROVIDER_REGEX, '');
5347
}

‎app/components/settings/features/FeaturesTab.tsx‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,8 @@ import { Switch } from '~/components/ui/Switch';
33
import { useSettings } from '~/lib/hooks/useSettings';
44

55
export default function FeaturesTab() {
6-
const { debug, enableDebugMode, isLocalModel, enableLocalModels, enableEventLogs, latestBranch, enableLatestBranch } =
7-
useSettings();
6+
7+
const { debug, enableDebugMode, isLocalModel, enableLocalModels, enableEventLogs, latestBranch, enableLatestBranch } = useSettings();
88

99
const handleToggle = (enabled: boolean) => {
1010
enableDebugMode(enabled);

‎app/lib/hooks/useSettings.tsx‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -99,7 +99,6 @@ export function useSettings() {
9999
if (checkCommit === undefined) {
100100
checkCommit = commit.commit;
101101
}
102-
103102
if (savedLatestBranch === undefined || checkCommit !== commit.commit) {
104103
// If setting hasn't been set by user, check version
105104
checkIsStableVersion().then((isStable) => {

‎app/routes/api.chat.ts‎

Lines changed: 38 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import { type ActionFunctionArgs } from '@remix-run/cloudflare';
2+
import { createDataStream } from 'ai';
23
import { MAX_RESPONSE_SEGMENTS, MAX_TOKENS } from '~/lib/.server/llm/constants';
34
import { CONTINUE_PROMPT } from '~/lib/.server/llm/prompts';
45
import { streamText, type Messages, type StreamingOptions } from '~/lib/.server/llm/stream-text';
@@ -9,17 +10,15 @@ export async function action(args: ActionFunctionArgs) {
910
return chatAction(args);
1011
}
1112

12-
function parseCookies(cookieHeader: string) {
13-
const cookies: any = {};
13+
function parseCookies(cookieHeader: string): Record<string, string> {
14+
const cookies: Record<string, string> = {};
1415

15-
// Split the cookie string by semicolons and spaces
1616
const items = cookieHeader.split(';').map((cookie) => cookie.trim());
1717

1818
items.forEach((item) => {
1919
const [name, ...rest] = item.split('=');
2020

2121
if (name && rest) {
22-
// Decode the name and value, and join value parts in case it contains '='
2322
const decodedName = decodeURIComponent(name.trim());
2423
const decodedValue = decodeURIComponent(rest.join('=').trim());
2524
cookies[decodedName] = decodedValue;
@@ -36,21 +35,49 @@ async function chatAction({ context, request }: ActionFunctionArgs) {
3635
}>();
3736

3837
const cookieHeader = request.headers.get('Cookie');
39-
40-
// Parse the cookie's value (returns an object or null if no cookie exists)
4138
const apiKeys = JSON.parse(parseCookies(cookieHeader || '').apiKeys || '{}');
4239
const providerSettings: Record<string, IProviderSetting> = JSON.parse(
4340
parseCookies(cookieHeader || '').providers || '{}',
4441
);
4542

4643
const stream = new SwitchableStream();
4744

45+
const cumulativeUsage = {
46+
completionTokens: 0,
47+
promptTokens: 0,
48+
totalTokens: 0,
49+
};
50+
4851
try {
4952
const options: StreamingOptions = {
5053
toolChoice: 'none',
51-
onFinish: async ({ text: content, finishReason }) => {
54+
onFinish: async ({ text: content, finishReason, usage }) => {
55+
console.log('usage', usage);
56+
57+
if (usage) {
58+
cumulativeUsage.completionTokens += usage.completionTokens || 0;
59+
cumulativeUsage.promptTokens += usage.promptTokens || 0;
60+
cumulativeUsage.totalTokens += usage.totalTokens || 0;
61+
}
62+
5263
if (finishReason !== 'length') {
53-
return stream.close();
64+
return stream
65+
.switchSource(
66+
createDataStream({
67+
async execute(dataStream) {
68+
dataStream.writeMessageAnnotation({
69+
type: 'usage',
70+
value: {
71+
completionTokens: cumulativeUsage.completionTokens,
72+
promptTokens: cumulativeUsage.promptTokens,
73+
totalTokens: cumulativeUsage.totalTokens,
74+
},
75+
});
76+
},
77+
onError: (error: any) => `Custom error: ${error.message}`,
78+
}),
79+
)
80+
.then(() => stream.close());
5481
}
5582

5683
if (stream.switches >= MAX_RESPONSE_SEGMENTS) {
@@ -73,7 +100,7 @@ async function chatAction({ context, request }: ActionFunctionArgs) {
73100
providerSettings,
74101
});
75102

76-
return stream.switchSource(result.toAIStream());
103+
return stream.switchSource(result.toDataStream());
77104
},
78105
};
79106

@@ -86,7 +113,7 @@ async function chatAction({ context, request }: ActionFunctionArgs) {
86113
providerSettings,
87114
});
88115

89-
stream.switchSource(result.toAIStream());
116+
stream.switchSource(result.toDataStream());
90117

91118
return new Response(stream.readable, {
92119
status: 200,
@@ -95,7 +122,7 @@ async function chatAction({ context, request }: ActionFunctionArgs) {
95122
},
96123
});
97124
} catch (error: any) {
98-
console.log(error);
125+
console.error(error);
99126

100127
if (error.message?.includes('API key')) {
101128
throw new Response('Invalid or missing API key', {

‎app/routes/api.enhancer.ts‎

Lines changed: 35 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import { type ActionFunctionArgs } from '@remix-run/cloudflare';
2-
import { StreamingTextResponse, parseStreamPart } from 'ai';
2+
3+
//import { StreamingTextResponse, parseStreamPart } from 'ai';
34
import { streamText } from '~/lib/.server/llm/stream-text';
45
import { stripIndents } from '~/utils/stripIndent';
56
import type { IProviderSetting, ProviderInfo } from '~/types/model';
@@ -73,32 +74,32 @@ async function enhancerAction({ context, request }: ActionFunctionArgs) {
7374
`[Model: ${model}]\n\n[Provider: ${providerName}]\n\n` +
7475
stripIndents`
7576
You are a professional prompt engineer specializing in crafting precise, effective prompts.
76-
Your task is to enhance prompts by making them more specific, actionable, and effective.
77-
78-
I want you to improve the user prompt that is wrapped in \`<original_prompt>\` tags.
79-
80-
For valid prompts:
81-
- Make instructions explicit and unambiguous
82-
- Add relevant context and constraints
83-
- Remove redundant information
84-
- Maintain the core intent
85-
- Ensure the prompt is self-contained
86-
- Use professional language
87-
88-
For invalid or unclear prompts:
89-
- Respond with a clear, professional guidance message
90-
- Keep responses concise and actionable
91-
- Maintain a helpful, constructive tone
92-
- Focus on what the user should provide
93-
- Use a standard template for consistency
94-
95-
IMPORTANT: Your response must ONLY contain the enhanced prompt text.
96-
Do not include any explanations, metadata, or wrapper tags.
97-
98-
<original_prompt>
99-
${message}
100-
</original_prompt>
101-
`,
77+
Your task is to enhance prompts by making them more specific, actionable, and effective.
78+
79+
I want you to improve the user prompt that is wrapped in \`<original_prompt>\` tags.
80+
81+
For valid prompts:
82+
- Make instructions explicit and unambiguous
83+
- Add relevant context and constraints
84+
- Remove redundant information
85+
- Maintain the core intent
86+
- Ensure the prompt is self-contained
87+
- Use professional language
88+
89+
For invalid or unclear prompts:
90+
- Respond with clear, professional guidance
91+
- Keep responses concise and actionable
92+
- Maintain a helpful, constructive tone
93+
- Focus on what the user should provide
94+
- Use a standard template for consistency
95+
96+
IMPORTANT: Your response must ONLY contain the enhanced prompt text.
97+
Do not include any explanations, metadata, or wrapper tags.
98+
99+
<original_prompt>
100+
${message}
101+
</original_prompt>
102+
`,
102103
},
103104
],
104105
env: context.cloudflare.env,
@@ -113,7 +114,7 @@ async function enhancerAction({ context, request }: ActionFunctionArgs) {
113114

114115
for (const line of lines) {
115116
try {
116-
const parsed = parseStreamPart(line);
117+
const parsed = JSON.parse(line);
117118

118119
if (parsed.type === 'text') {
119120
controller.enqueue(encoder.encode(parsed.value));
@@ -128,7 +129,12 @@ async function enhancerAction({ context, request }: ActionFunctionArgs) {
128129

129130
const transformedStream = result.toDataStream().pipeThrough(transformStream);
130131

131-
return new StreamingTextResponse(transformedStream);
132+
return new Response(transformedStream, {
133+
status: 200,
134+
headers: {
135+
'Content-Type': 'text/plain; charset=utf-8',
136+
},
137+
});
132138
} catch (error: unknown) {
133139
console.log(error);
134140

‎package.json‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,7 @@
7373
"@xterm/addon-fit": "^0.10.0",
7474
"@xterm/addon-web-links": "^0.11.0",
7575
"@xterm/xterm": "^5.5.0",
76-
"ai": "^3.4.33",
76+
"ai": "^4.0.13",
7777
"date-fns": "^3.6.0",
7878
"diff": "^5.2.0",
7979
"file-saver": "^2.0.5",

0 commit comments

Comments
 (0)