fix: implement proper token tracking and cost calculations - Add Prompt model import and update token tracking in generate.js - Update savePrompt to handle token data - Fix token calculations for both streaming and non-streaming responses - Add proper error handling for token updates
This commit is contained in:
@ -3,16 +3,27 @@ const ErrorLog = require('../../models/Error');
|
|||||||
|
|
||||||
/**
|
/**
|
||||||
* Save a prompt to the database
|
* Save a prompt to the database
|
||||||
|
* @param {Object} params - The prompt data
|
||||||
|
* @param {string} params.model - The model used
|
||||||
|
* @param {string} [params.prompt] - The prompt text (for non-chat requests)
|
||||||
|
* @param {Array} [params.messages] - Array of chat messages (for chat requests)
|
||||||
|
* @param {Object} [params.request_data] - Complete request data
|
||||||
|
* @param {Object} [params.token_data] - Token counts and cost calculations
|
||||||
|
* @returns {Promise<Prompt>} The created prompt record
|
||||||
*/
|
*/
|
||||||
async function savePrompt({ model, prompt, messages, request_data }) {
|
async function savePrompt({ model, prompt, messages, request_data, token_data }) {
|
||||||
return Prompt.createPrompt({ model, prompt, messages, request_data });
|
return Prompt.createPrompt({ model, prompt, messages, request_data, token_data });
|
||||||
}
|
}
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* Save an error to the database
|
* Save an error to the database
|
||||||
|
* @param {Object} params - The error data
|
||||||
|
* @param {string} params.error_message - The error message
|
||||||
|
* @param {Object} [params.details] - Additional error details
|
||||||
|
* @returns {Promise<ErrorLog>} The created error record
|
||||||
*/
|
*/
|
||||||
async function saveError({ error_message, details }) {
|
async function saveError({ error_message, details }) {
|
||||||
return ErrorLog.createError({ error_message, details });
|
return ErrorLog.create({ error_message, details });
|
||||||
}
|
}
|
||||||
|
|
||||||
module.exports = { savePrompt, saveError };
|
module.exports = { savePrompt, saveError };
|
||||||
|
@ -3,6 +3,7 @@ const fs = require("fs");
|
|||||||
const path = require("path");
|
const path = require("path");
|
||||||
const { savePrompt, saveError } = require("../controller/generate");
|
const { savePrompt, saveError } = require("../controller/generate");
|
||||||
const { processRequestTokens } = require("../../utils/tokenCounter");
|
const { processRequestTokens } = require("../../utils/tokenCounter");
|
||||||
|
const Prompt = require("../../models/Prompt");
|
||||||
|
|
||||||
// Constants for file handling
|
// Constants for file handling
|
||||||
const MAX_FILE_SIZE = 30000; // ~30KB per file
|
const MAX_FILE_SIZE = 30000; // ~30KB per file
|
||||||
@ -455,6 +456,11 @@ async function handleGenerate(req, res) {
|
|||||||
prompt: isChatRequest ? null : cleanedRequest.prompt,
|
prompt: isChatRequest ? null : cleanedRequest.prompt,
|
||||||
messages: isChatRequest ? cleanedRequest.messages : null,
|
messages: isChatRequest ? cleanedRequest.messages : null,
|
||||||
request_data: requestData,
|
request_data: requestData,
|
||||||
|
token_data: processRequestTokens(
|
||||||
|
cleanedRequest,
|
||||||
|
'', // Empty response initially
|
||||||
|
parseFloat(process.env.YOUR_MODEL_COST) || 0
|
||||||
|
)
|
||||||
});
|
});
|
||||||
promptId = prompt.id;
|
promptId = prompt.id;
|
||||||
|
|
||||||
@ -512,7 +518,7 @@ async function handleGenerate(req, res) {
|
|||||||
responseContent,
|
responseContent,
|
||||||
parseFloat(process.env.YOUR_MODEL_COST) || 0
|
parseFloat(process.env.YOUR_MODEL_COST) || 0
|
||||||
);
|
);
|
||||||
await updatePromptTokens(promptId, tokenData);
|
await Prompt.updatePromptTokens(promptId, tokenData);
|
||||||
}
|
}
|
||||||
res.end();
|
res.end();
|
||||||
}
|
}
|
||||||
@ -527,7 +533,7 @@ async function handleGenerate(req, res) {
|
|||||||
responseContent,
|
responseContent,
|
||||||
parseFloat(process.env.YOUR_MODEL_COST) || 0
|
parseFloat(process.env.YOUR_MODEL_COST) || 0
|
||||||
);
|
);
|
||||||
await updatePromptTokens(promptId, tokenData);
|
await Prompt.updatePromptTokens(promptId, tokenData);
|
||||||
|
|
||||||
res.status(ollamaResponse.status).json(ollamaResponse.data);
|
res.status(ollamaResponse.status).json(ollamaResponse.data);
|
||||||
}
|
}
|
||||||
|
Reference in New Issue
Block a user