Tokens accounting

docker-compose.yml

@@ -33,6 +33,8 @@ services:
         volumes:
             - mongo-data:/data/db
         restart: unless-stopped
+        logging:
+            driver: none
 
 volumes:
     mongo-data:

server/auth.js

@@ -1,24 +1,6 @@
 import bcrypt from "bcrypt";
 import jwt from "jsonwebtoken";
-import mongoose from "mongoose";
-
-export const MONGODB_URI =
-    process.env.NODE_ENV === "production" ? "mongodb://mongodb:27017/allchat" : "mongodb://localhost:27017/allchat";
-
-// MongoDB connection
-mongoose
-    .connect(MONGODB_URI)
-    .then(() => console.log("MongoDB connected"))
-    .catch((err) => console.error("MongoDB connection error:", err));
-
-// User schema
-const userSchema = new mongoose.Schema({
-    email: { type: String, required: true, unique: true },
-    password: { type: String, required: true },
-});
-
-// User model
-const User = mongoose.model("User", userSchema);
+import { User } from "./model/User.js";
 
 // Register a new user
 export const registerUser = async (email, password) => {
@@ -61,10 +43,9 @@ export const verifyToken = (req, res, next) => {
     if (!token) {
         return res.status(401).json({ error: "Unauthorized" });
     }
-
     try {
         const decoded = jwt.verify(token, process.env.JWT_TOKEN);
-        req.userId = decoded.userId;
+        req.user = { id: decoded.userId }; // Set req.user.id instead of req.userId
         next();
     } catch (error) {
         return res.status(403).json({ error: "Invalid token" });

server/index.js

@@ -11,6 +11,8 @@ import * as xlsx from "xlsx";
 import { getTextClaude } from "./claude.js";
 import promBundle from "express-prom-bundle";
 import { authenticateUser, registerUser, verifyToken } from "./auth.js";
+import mongoose from "mongoose";
+import { countCharacters, countTokens, storeUsageStats } from "./model/User.js";
 
 const MAX_CONTEXT_LENGTH = 8000;
 const systemPrompt = `You are an AI assistant that interacts with the Gemini Pro and Claude Haiku language models. Your capabilities include:
@@ -113,18 +115,39 @@ app.post("/interact", verifyToken, async (req, res) => {
             .join("\n")}\n\nHuman: ${userInput}\nAssistant:`.slice(-MAX_CONTEXT_LENGTH);
 
         let textResponse;
+        let inputTokens = 0;
+        let outputTokens = 0;
+        let inputCharacters = 0;
+        let outputCharacters = 0;
+        let imagesGenerated = 0;
+
         if (model === "gemini") {
+            inputCharacters = countCharacters(contextPrompt);
             textResponse = await getTextGemini(contextPrompt, temperature);
+            outputCharacters = countCharacters(textResponse);
         } else if (model === "claude") {
+            inputTokens = countTokens(contextPrompt);
             textResponse = await getTextClaude(contextPrompt, "claude-3-haiku-20240307", temperature);
+            outputTokens = countTokens(textResponse);
         }
 
         userInput = userInput?.toLowerCase();
         let imageResponse;
         if (hasPaintWord(userInput)) {
             imageResponse = await getImageTitan(userInput?.substr(0, 200) + textResponse?.trim()?.substr(0, 200));
+            imagesGenerated = 1;
         }
 
+        storeUsageStats(
+            req.user.id,
+            model,
+            inputTokens,
+            outputTokens,
+            inputCharacters,
+            outputCharacters,
+            imagesGenerated
+        );
+
         res.json({ textResponse: textResponse?.trim(), imageResponse });
     } catch (error) {
         console.error(error);
@@ -157,3 +180,12 @@ app.post("/login", async (req, res) => {
 app.listen(5000, () => {
     console.log(`🚀 Server started on port 5000`);
 });
+
+export const MONGODB_URI =
+    process.env.NODE_ENV === "production" ? "mongodb://mongodb:27017/allchat" : "mongodb://localhost:27017/allchat";
+
+// MongoDB connection
+mongoose
+    .connect(MONGODB_URI)
+    .then(() => console.log("🚀 MongoDB connected"))
+    .catch((err) => console.error("MongoDB connection error:", err));

server/model/User.js

@@ -0,0 +1,80 @@
+import mongoose from "mongoose";
+
+// User schema
+const userSchema = new mongoose.Schema({
+    email: { type: String, required: true, unique: true },
+    password: { type: String, required: true },
+    usageStats: {
+        gemini: {
+            inputCharacters: { type: Number, default: 0 },
+            outputCharacters: { type: Number, default: 0 },
+            imagesGenerated: { type: Number, default: 0 },
+            moneyConsumed: { type: Number, default: 0 },
+        },
+        claude: {
+            inputTokens: { type: Number, default: 0 },
+            outputTokens: { type: Number, default: 0 },
+            moneyConsumed: { type: Number, default: 0 },
+        },
+    },
+});
+
+// User model
+export const User = mongoose.model("User", userSchema);
+
+export function countCharacters(text) {
+    return text.length;
+}
+
+export function countTokens(text) {
+    let tokenCount = 0;
+    for (let i = 0; i < text.length; i++) {
+        const char = text[i];
+        if (/[a-zA-Z]/.test(char)) {
+            tokenCount += 0.25;
+        } else {
+            tokenCount += 1;
+        }
+    }
+    return Math.round(tokenCount);
+}
+
+export function storeUsageStats(
+    userId,
+    model,
+    inputTokens,
+    outputTokens,
+    inputCharacters,
+    outputCharacters,
+    imagesGenerated
+) {
+    let moneyConsumed = 0;
+
+    if (model === "gemini") {
+        const inputCharactersCost = (inputCharacters / 1000) * 0.000125;
+        const outputCharactersCost = (outputCharacters / 1000) * 0.000375;
+        const imagesGeneratedCost = imagesGenerated * 0.0025;
+        moneyConsumed = inputCharactersCost + outputCharactersCost + imagesGeneratedCost;
+    } else if (model === "claude") {
+        const inputTokensCost = (inputTokens * 0.25) / 1000000;
+        const outputTokensCost = (outputTokens * 1.25) / 1000000;
+        moneyConsumed = inputTokensCost + outputTokensCost;
+    }
+
+    User.findByIdAndUpdate(
+        userId,
+        {
+            $inc: {
+                [`usageStats.${model}.inputTokens`]: inputTokens,
+                [`usageStats.${model}.outputTokens`]: outputTokens,
+                [`usageStats.${model}.inputCharacters`]: inputCharacters,
+                [`usageStats.${model}.outputCharacters`]: outputCharacters,
+                [`usageStats.${model}.imagesGenerated`]: imagesGenerated,
+                [`usageStats.${model}.moneyConsumed`]: moneyConsumed,
+            },
+        },
+        { new: true }
+    )
+        .then((updatedUser) => console.log(`Usage stats updated for user ${updatedUser._id}`))
+        .catch((err) => console.error("Error updating usage stats:", err));
+}