From fb89f60470d684cee7ccf3e8abf57f4247e27a3e Mon Sep 17 00:00:00 2001
From: Dustin Healy <dustinhealy1@gmail.com>
Date: Sat, 9 Aug 2025 10:26:28 -0700
Subject: [PATCH] feat: add transaction costs endpoint

---
 api/server/routes/messages.js | 102 +++++++++++++++++++++++++++++++++-
 1 file changed, 101 insertions(+), 1 deletion(-)

diff --git a/api/server/routes/messages.js b/api/server/routes/messages.js
index 0a277a1bd6..8a063cd5ca 100644
--- a/api/server/routes/messages.js
+++ b/api/server/routes/messages.js
@@ -13,8 +13,8 @@ const { findAllArtifacts, replaceArtifactContent } = require('~/server/services/
 const { requireJwtAuth, validateMessageReq } = require('~/server/middleware');
 const { cleanUpPrimaryKeyValue } = require('~/lib/utils/misc');
 const { getConvosQueried } = require('~/models/Conversation');
+const { Message, Transaction } = require('~/db/models');
 const { countTokens } = require('~/server/utils');
-const { Message } = require('~/db/models');
 
 const router = express.Router();
 router.use(requireJwtAuth);
@@ -294,4 +294,104 @@ router.delete('/:conversationId/:messageId', validateMessageReq, async (req, res
   }
 });
 
+router.get('/:conversationId/costs', validateMessageReq, async (req, res) => {
+  try {
+    const user = req.user.id;
+    const { conversationId } = req.params;
+
+    const [transactions, messages] = await Promise.all([
+      Transaction.find({
+        conversationId,
+        user,
+        tokenType: { $in: ['prompt', 'completion'] },
+      })
+        .select('tokenType tokenValue createdAt')
+        .sort({ createdAt: 1 })
+        .lean(),
+      Message.find({ conversationId, user })
+        .select('messageId isCreatedByUser tokenCount createdAt')
+        .sort({ createdAt: 1 })
+        .lean(),
+    ]);
+
+    const userMsgs = messages.filter((m) => m.isCreatedByUser);
+    const aiMsgs = messages.filter((m) => !m.isCreatedByUser);
+
+    const perMessageMap = new Map();
+    for (const msg of messages) {
+      perMessageMap.set(msg.messageId, {
+        messageId: msg.messageId,
+        tokenType: msg.isCreatedByUser ? 'prompt' : 'completion',
+        tokenCount: msg.tokenCount ?? 0,
+        tokenValue: 0,
+        usd: 0,
+      });
+    }
+
+    let currentPrompt = 0;
+    let currentCompletion = 0;
+
+    let promptTokenValue = 0;
+    let completionTokenValue = 0;
+
+    for (const tx of transactions) {
+      const value = Math.abs(tx.tokenValue ?? 0);
+      if (tx.tokenType === 'prompt') {
+        promptTokenValue += value;
+        const target = userMsgs[currentPrompt] ?? userMsgs[userMsgs.length - 1];
+        if (target) {
+          const entry = perMessageMap.get(target.messageId);
+          entry.tokenValue += value;
+          perMessageMap.set(target.messageId, entry);
+          if (currentPrompt < userMsgs.length - 1) {
+            currentPrompt++;
+          }
+        }
+      } else if (tx.tokenType === 'completion') {
+        completionTokenValue += value;
+        const target = aiMsgs[currentCompletion] ?? aiMsgs[aiMsgs.length - 1];
+        if (target) {
+          const entry = perMessageMap.get(target.messageId);
+          entry.tokenValue += value;
+          perMessageMap.set(target.messageId, entry);
+          if (currentCompletion < aiMsgs.length - 1) {
+            currentCompletion++;
+          }
+        }
+      }
+    }
+
+    const perMessage = Array.from(perMessageMap.values()).map((entry) => ({
+      messageId: entry.messageId,
+      tokenType: entry.tokenType,
+      tokenCount: entry.tokenCount,
+      usd: entry.tokenValue / 1_000_000,
+    }));
+
+    const promptTokenCount = userMsgs.reduce((sum, m) => sum + (m.tokenCount ?? 0), 0);
+    const completionTokenCount = aiMsgs.reduce((sum, m) => sum + (m.tokenCount ?? 0), 0);
+    const totalTokenCount = promptTokenCount + completionTokenCount;
+
+    const totals = {
+      prompt: {
+        usd: promptTokenValue / 1_000_000,
+        tokenCount: promptTokenCount,
+      },
+      completion: {
+        usd: completionTokenValue / 1_000_000,
+        tokenCount: completionTokenCount,
+      },
+      total: {
+        usd: (promptTokenValue + completionTokenValue) / 1_000_000,
+        tokenCount: totalTokenCount,
+      },
+    };
+
+    res.status(200).json({ conversationId, totals, perMessage });
+  } catch (error) {
+    logger.error('Error fetching conversation costs:', error);
+    res.status(500).json({ error: 'Internal server error' });
+  }
+});
+
 module.exports = router;