From 225b55387669c9216f60896514e762515b8b1838 Mon Sep 17 00:00:00 2001
From: eduardruzga <wonderwhy.er@gmail.com>
Date: Mon, 16 Dec 2024 11:01:41 +0200
Subject: [PATCH] Another attempt to add toek usage info

---
 app/commit.json                          |  2 +-
 app/components/chat/AssistantMessage.tsx | 12 +++++++-
 app/components/chat/UserMessage.tsx      | 36 ++++++++++--------------
 app/lib/.server/llm/switchable-stream.ts |  2 +-
 app/routes/api.chat.ts                   | 25 ++++++++++++++++
 app/utils/constants.ts                   |  1 +
 6 files changed, 54 insertions(+), 24 deletions(-)

diff --git a/app/commit.json b/app/commit.json
index 46c9734..018af74 100644
--- a/app/commit.json
+++ b/app/commit.json
@@ -1 +1 @@
-{ "commit": "fcb61ba49990d1d0cc13d05a0958007b6e9c1c38" }
+{ "commit": "2e05270bab264d7ee83d7a1dd5408c969f648a68" }
diff --git a/app/components/chat/AssistantMessage.tsx b/app/components/chat/AssistantMessage.tsx
index a5698e9..7cdddda 100644
--- a/app/components/chat/AssistantMessage.tsx
+++ b/app/components/chat/AssistantMessage.tsx
@@ -1,14 +1,24 @@
 import { memo } from 'react';
 import { Markdown } from './Markdown';
+import { USAGE_REGEX } from '~/utils/constants';
 
 interface AssistantMessageProps {
   content: string;
 }
 
 export const AssistantMessage = memo(({ content }: AssistantMessageProps) => {
+  const match = content.match(USAGE_REGEX);
+  const usage = match ? JSON.parse(match[1]) : null;
+  const cleanContent = content.replace(USAGE_REGEX, '').trim();
+
   return (
     <div className="overflow-hidden w-full">
-      <Markdown html>{content}</Markdown>
+      {usage && (
+        <div className="text-sm text-bolt-elements-textSecondary mb-2">
+          Tokens: {usage.totalTokens} (prompt: {usage.promptTokens}, completion: {usage.completionTokens})
+        </div>
+      )}
+      <Markdown html>{cleanContent}</Markdown>
     </div>
   );
 });
diff --git a/app/components/chat/UserMessage.tsx b/app/components/chat/UserMessage.tsx
index 3e6485b..6fd6fc0 100644
--- a/app/components/chat/UserMessage.tsx
+++ b/app/components/chat/UserMessage.tsx
@@ -12,42 +12,36 @@ interface UserMessageProps {
 export function UserMessage({ content }: UserMessageProps) {
   if (Array.isArray(content)) {
     const textItem = content.find((item) => item.type === 'text');
-    const textContent = sanitizeUserMessage(textItem?.text || '');
+    const textContent = stripMetadata(textItem?.text || '');
     const images = content.filter((item) => item.type === 'image' && item.image);
 
     return (
       <div className="overflow-hidden pt-[4px]">
-        <div className="flex items-start gap-4">
-          <div className="flex-1">
-            <Markdown limitedMarkdown>{textContent}</Markdown>
-          </div>
-          {images.length > 0 && (
-            <div className="flex-shrink-0 w-[160px]">
-              {images.map((item, index) => (
-                <div key={index} className="relative">
-                  <img
-                    src={item.image}
-                    alt={`Uploaded image ${index + 1}`}
-                    className="w-full h-[160px] rounded-lg object-cover border border-bolt-elements-borderColor"
-                  />
-                </div>
-              ))}
-            </div>
-          )}
+        <div className="flex flex-col gap-4">
+          {textContent && <Markdown html>{textContent}</Markdown>}
+          {images.map((item, index) => (
+            <img
+              key={index}
+              src={item.image}
+              alt={`Image ${index + 1}`}
+              className="max-w-full h-auto rounded-lg"
+              style={{ maxHeight: '512px', objectFit: 'contain' }}
+            />
+          ))}
         </div>
       </div>
     );
   }
 
-  const textContent = sanitizeUserMessage(content);
+  const textContent = stripMetadata(content);
 
   return (
     <div className="overflow-hidden pt-[4px]">
-      <Markdown limitedMarkdown>{textContent}</Markdown>
+      <Markdown html>{textContent}</Markdown>
     </div>
   );
 }
 
-function sanitizeUserMessage(content: string) {
+function stripMetadata(content: string) {
   return content.replace(MODEL_REGEX, '').replace(PROVIDER_REGEX, '');
 }
diff --git a/app/lib/.server/llm/switchable-stream.ts b/app/lib/.server/llm/switchable-stream.ts
index 638db7a..51f46ff 100644
--- a/app/lib/.server/llm/switchable-stream.ts
+++ b/app/lib/.server/llm/switchable-stream.ts
@@ -1,5 +1,5 @@
 export default class SwitchableStream extends TransformStream {
-  private _controller: TransformStreamDefaultController | null = null;
+  _controller: TransformStreamDefaultController | null = null;
   private _currentReader: ReadableStreamDefaultReader | null = null;
   private _switches = 0;
 
diff --git a/app/routes/api.chat.ts b/app/routes/api.chat.ts
index 34dea34..c7d6383 100644
--- a/app/routes/api.chat.ts
+++ b/app/routes/api.chat.ts
@@ -41,12 +41,36 @@ async function chatAction({ context, request }: ActionFunctionArgs) {
 
   const stream = new SwitchableStream();
 
+  const cumulativeUsage = {
+    completionTokens: 0,
+    promptTokens: 0,
+    totalTokens: 0,
+  };
+
   try {
     const options: StreamingOptions = {
       toolChoice: 'none',
       onFinish: async ({ text: content, finishReason, usage }) => {
         console.log('usage', usage);
 
+        if (usage && stream._controller) {
+          cumulativeUsage.completionTokens += usage.completionTokens || 0;
+          cumulativeUsage.promptTokens += usage.promptTokens || 0;
+          cumulativeUsage.totalTokens += usage.totalTokens || 0;
+
+          // Send usage info in message metadata for assistant messages
+          const usageMetadata = `0:"[Usage: ${JSON.stringify({
+            completionTokens: cumulativeUsage.completionTokens,
+            promptTokens: cumulativeUsage.promptTokens,
+            totalTokens: cumulativeUsage.totalTokens,
+          })}\n]"`;
+
+          console.log(usageMetadata);
+
+          const encodedData = new TextEncoder().encode(usageMetadata);
+          stream._controller.enqueue(encodedData);
+        }
+
         if (finishReason !== 'length') {
           return stream.close();
         }
@@ -83,6 +107,7 @@ async function chatAction({ context, request }: ActionFunctionArgs) {
       files,
       providerSettings,
     });
+
     stream.switchSource(result.toDataStream());
 
     return new Response(stream.readable, {
diff --git a/app/utils/constants.ts b/app/utils/constants.ts
index cc2aafc..2ab6cfb 100644
--- a/app/utils/constants.ts
+++ b/app/utils/constants.ts
@@ -9,6 +9,7 @@ export const WORK_DIR = `/home/${WORK_DIR_NAME}`;
 export const MODIFICATIONS_TAG_NAME = 'bolt_file_modifications';
 export const MODEL_REGEX = /^\[Model: (.*?)\]\n\n/;
 export const PROVIDER_REGEX = /\[Provider: (.*?)\]\n\n/;
+export const USAGE_REGEX = /\[Usage: ({.*?})\]/; // Keep this regex for assistant messages
 export const DEFAULT_MODEL = 'claude-3-5-sonnet-latest';
 export const PROMPT_COOKIE_KEY = 'cachedPrompt';