From a92b7642633ff11952a8111e2f69f35eca16b07d Mon Sep 17 00:00:00 2001
From: skeptrune <nick.k@trieve.ai>
Date: Tue, 19 Nov 2024 18:40:32 -0800
Subject: [PATCH] feature: add token usage to pdf2md ui

---
 pdf2md/server/src/templates/demo-ui.html | 25 ++++++++++++++++++++++--
 pdf2md/server/static/pdf2md.js           | 17 ++++++++++++----
 2 files changed, 36 insertions(+), 6 deletions(-)

diff --git a/pdf2md/server/src/templates/demo-ui.html b/pdf2md/server/src/templates/demo-ui.html
index 75be9c76f2..fd463231d4 100644
--- a/pdf2md/server/src/templates/demo-ui.html
+++ b/pdf2md/server/src/templates/demo-ui.html
@@ -31,8 +31,7 @@
         OCR With Intelligence
       </h2>
       <p class="mt-2 text-pretty text-lg/8 text-gray-700">
-        Convert any PDF to LLM-ready Markdown using latest-gen vision models
-        like GPT-4o.
+        Convert any PDF to LLM-ready Markdown using vision models like GPT-4o.
       </p>
     </div>
     <button
@@ -175,6 +174,18 @@
             >
               Status
             </th>
+            <th
+              scope="col"
+              class="px-3 py-3.5 text-left text-sm font-semibold text-gray-900"
+            >
+              Prompt Tokens
+            </th>
+            <th
+              scope="col"
+              class="px-3 py-3.5 text-left text-sm font-semibold text-gray-900"
+            >
+              Completion Tokens
+            </th>
             <th scope="col" class="relative py-3.5 pl-3 pr-4 sm:pr-0">
               <span class="sr-only">View</span>
             </th>
@@ -197,6 +208,16 @@
             >
               lindsay.walton@example.com
             </td>
+            <td
+              class="task-prompt-tokens whitespace-nowrap px-3 py-4 text-sm text-gray-500"
+            >
+              10
+            </td>
+            <td
+              class="task-completion-tokens whitespace-nowrap px-3 py-4 text-sm text-gray-500"
+            >
+              10
+            </td>
             <td
               class="relative whitespace-nowrap py-4 pl-3 pr-4 text-right text-sm font-medium sm:pr-0"
             >
diff --git a/pdf2md/server/static/pdf2md.js b/pdf2md/server/static/pdf2md.js
index 26f2d951ad..9cb7b19948 100644
--- a/pdf2md/server/static/pdf2md.js
+++ b/pdf2md/server/static/pdf2md.js
@@ -62,7 +62,7 @@ const displayTask = (task) => {
   if (!pages) {
     return;
   }
-  const sortedPages = pages.sort((a, b) => a.metadata.page - b.metadata.page);
+  const sortedPages = pages.sort((a, b) => a.page_num - b.page_num);
 
   PDFObject.embed(task.file_url, "#my-pdf", {
     pdfOpenParams: {
@@ -126,8 +126,7 @@ const getTaskPages = async (taskId, taskIdToDisplay) => {
       }
     }
 
-    pages = pages.sort((a, b) => a.metadata.page - b.metadata.page);
-    console.log("final pages", taskId, pages);
+    pages = pages.sort((a, b) => a.page_num - b.page_num);
     task.pages = pages;
     upsertTaskToStorage(task);
     if (taskIdToDisplay === taskId) {
@@ -135,7 +134,7 @@ const getTaskPages = async (taskId, taskIdToDisplay) => {
     }
   } catch (e) {
     console.error(e);
-    Notyf.error({
+    notyf.error({
       message: `Error fetching task pages. Please try again later. ${e}`,
       dismissable: true,
       type: "error",
@@ -231,6 +230,13 @@ const updateTaskStatusTable = () => {
   const firstRow = tbody.querySelector("tr");
   tbody.innerHTML = "";
   const htmlRows = tasks.map((task) => {
+    let completionTokens = 0;
+    let promptTokens = 0;
+    task.pages?.map((page) => {
+      promptTokens += page?.usage?.prompt_tokens || 0;
+      completionTokens += page?.usage?.completion_tokens || 0;
+    });
+
     const row = firstRow
       ? firstRow.cloneNode(true)
       : defaultTableRow.cloneNode(true);
@@ -245,6 +251,9 @@ const updateTaskStatusTable = () => {
       .classList.add(
         `status-${task.status.split(" ").join("-").toLowerCase()}`
       );
+    row.querySelector(".task-prompt-tokens").innerText = promptTokens;
+    row.querySelector(".task-completion-tokens").innerText = completionTokens;
+
     row.querySelector("button").addEventListener("click", () => {
       const url = new URL(window.location);
       url.searchParams.set("taskId", task.id);