Skip to content

Commit

Permalink
added aiUtils file for ai multiplier management
Browse files Browse the repository at this point in the history
added incrbyfloat function in TdCache
added quotes field in profile model
added multiplier support in incrementTokenCount
added check chatbots limit in chatbot create
updated QuoteManager class
updated tests
  • Loading branch information
Giovanni Troisi committed Mar 25, 2024
1 parent bb36955 commit 5fc81b7
Show file tree
Hide file tree
Showing 10 changed files with 228 additions and 37 deletions.
3 changes: 3 additions & 0 deletions models/profile.js
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,9 @@ var ProfileSchema = new Schema({
type: String,
default: 'free',
},
quotes: {
type: Object
},
subStart: {
type: Date,
},
Expand Down
14 changes: 13 additions & 1 deletion routes/faq_kb.js
Original file line number Diff line number Diff line change
Expand Up @@ -19,10 +19,22 @@ const trainingService = require('../services/trainingService');

let chatbot_templates_api_url = process.env.CHATBOT_TEMPLATES_API_URL

router.post('/', function (req, res) {
router.post('/', async function (req, res) {
winston.debug('create BOT ', req.body);
//create(name, url, projectid, user_id, type, description, webhook_url, webhook_enabled, language, template)
//faqService.create(req.body.name, req.body.url, req.projectid, req.user.id, req.body.type, req.body.description, undefined, undefined, req.body.language, req.body.template, req.body.mainCategory, req.body.intentsEngine).then(function (savedFaq_kb) {
let quoteManager = req.app.get('quote_manager');
let limits = await quoteManager.getPlanLimits(req.project);
let chatbots_limit = limits.chatbots;
winston.debug("chatbots_limit for project " + req.projectid + ": " + chatbots_limit);

let chatbots_count = await Faq_kb.countDocuments({ id_project: req.projectid, type: { $ne: "identity" } }).exec();
winston.debug("chatbots_count for project " + req.projectid + ": " + chatbots_count);

if (chatbots_count >= chatbots_limit) {
return res.status(403).send({ success: false, error: "Maximum number of chatbots reached for the current plan", plan_limit: chatbots_limit })
}

faqService.create(req.body.name, req.body.url, req.projectid, req.user.id, req.body.type, req.body.description, req.body.webhook_url, req.body.webhook_enabled, req.body.language, req.body.template, req.body.mainCategory, req.body.intentsEngine, req.body.attributes).then(function (savedFaq_kb) {
res.json(savedFaq_kb);
});
Expand Down
9 changes: 8 additions & 1 deletion routes/openai.js
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ var { KBSettings } = require('../models/kb_setting');
var openaiService = require('../services/openaiService');
var winston = require('../config/winston');
const { QuoteManager } = require('../services/QuoteManager');
const { MODEL_MULTIPLIER } = require('../utils/aiUtils');

router.post('/', async (req, res) => {

Expand Down Expand Up @@ -60,8 +61,14 @@ router.post('/', async (req, res) => {
json.messages.unshift(message);
}

let multiplier = MODEL_MULTIPLIER[json.model];
if (!multiplier) {
multiplier = 1;
winston.info("No multiplier found for AI model")
}

openaiService.completions(json, gptkey).then(async (response) => {
let data = { createdAt: new Date(), tokens: response.data.usage.total_tokens }
let data = { createdAt: new Date(), tokens: response.data.usage.total_tokens, multiplier: multiplier }
if (usePublicKey === true) {
let incremented_key = await quoteManager.incrementTokenCount(req.project, data);
winston.verbose("Tokens quota incremented for key " + incremented_key);
Expand Down
31 changes: 21 additions & 10 deletions services/QuoteManager.js
Original file line number Diff line number Diff line change
Expand Up @@ -5,14 +5,14 @@ const messageEvent = require('../event/messageEvent');
const emailEvent = require('../event/emailEvent');

const PLANS_LIST = {
FREE_TRIAL: { requests: 3000, messages: 0, tokens: 250000, email: 200, kbs: 50 }, // same as PREMIUM
SANDBOX: { requests: 200, messages: 0, tokens: 10000, email: 200, kbs: 50 },
BASIC: { requests: 800, messages: 0, tokens: 50000, email: 200, kbs: 150},
PREMIUM: { requests: 3000, messages: 0, tokens: 250000, email: 200, kbs: 300},
CUSTOM: { requests: 3000, messages: 0, tokens: 250000, email: 200, kbs: 300}
FREE_TRIAL: { requests: 3000, messages: 0, tokens: 250000, email: 200, chatbots: 20, kbs: 50 }, // same as PREMIUM
SANDBOX: { requests: 200, messages: 0, tokens: 10000, email: 200, chatbots: 2, kbs: 50 },
BASIC: { requests: 800, messages: 0, tokens: 50000, email: 200, chatbots: 5, kbs: 150},
PREMIUM: { requests: 3000, messages: 0, tokens: 250000, email: 200, chatbots: 20, kbs: 300},
CUSTOM: { requests: 3000, messages: 0, tokens: 250000, email: 200, chatbots: 20, kbs: 300}
}

const typesList = ['requests', 'messages', 'email', 'tokens']
const typesList = ['requests', 'messages', 'email', 'tokens', 'chatbots', 'kbs']

let quotes_enabled = true;

Expand Down Expand Up @@ -74,8 +74,11 @@ class QuoteManager {
winston.debug("QUOTES DISABLED - incrementTokenCount")
return key;
}

await this.tdCache.incrby(key, data.tokens);

let tokens = data.tokens * data.multiplier;
await this.tdCache.incrbyfloat(key, tokens);
// await this.tdCache.incrby(key, tokens);

return key;
}
// INCREMENT KEY SECTION - END
Expand Down Expand Up @@ -156,6 +159,7 @@ class QuoteManager {
this.project = project;
let limits = await this.getPlanLimits();
winston.verbose("limits for current plan: ", limits)

let quote = await this.getCurrentQuote(project, object, type);
winston.verbose("getCurrentQuote resp: ", quote)

Expand Down Expand Up @@ -212,7 +216,14 @@ class QuoteManager {
} else {
limits = PLANS_LIST.FREE_TRIAL;
}
return limits;

if (this.project?.profile?.quotes) {
let profile_quotes = this.project?.profile?.quotes;
const merged_quotes = Object.assign({}, limits, profile_quotes);
return merged_quotes;
} else {
return limits;
}
}


Expand All @@ -226,7 +237,7 @@ class QuoteManager {
}
}

winston.info("QUOTES ENABLED ? ", quotes_enabled);
winston.info("QUOTES ENABLED: " + quotes_enabled);

// TODO - Try to generalize to avoid repetition
let incrementEventHandler = (object) => { }
Expand Down
68 changes: 68 additions & 0 deletions test/faqkbRoute.js
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,74 @@ describe('FaqKBRoute', () => {

}).timeout(20000);

/**
* This test will be no longer available after merge with master because
* the profile section can no longer be modified via api.
*/
// it('createMaximumNumberExceeded', (done) => {

// var email = "test-signup-" + Date.now() + "@email.com";
// var pwd = "pwd";

// userService.signup(email, pwd, "Test Firstname", "Test lastname").then(function (savedUser) {
// projectService.create("test-faqkb-create", savedUser._id).then(function (savedProject) {

// chai.request(server)
// .put('/projects/' + savedProject._id)
// .auth(email, pwd)
// .send({ profile: { quotes: { chatbots: 2 } } })
// .end((err, res) => {

// if (log) { console.log("res.body", res.body); }

// chai.request(server)
// .post('/' + savedProject._id + '/faq_kb')
// .auth(email, pwd)
// .send({ "name": "testbot1", type: "external", language: 'en' })
// .end((err, res) => {
// if (log) { console.log("res.body", res.body); }
// res.should.have.status(200);
// res.body.should.be.a('object');
// expect(res.body.name).to.equal("testbot1");
// expect(res.body.language).to.equal("en");

// chai.request(server)
// .post('/' + savedProject._id + '/faq_kb')
// .auth(email, pwd)
// .send({ "name": "testbot2", type: "external", language: 'en' })
// .end((err, res) => {
// if (log) { console.log("res.body", res.body); }
// res.should.have.status(200);
// res.body.should.be.a('object');
// expect(res.body.name).to.equal("testbot2");
// expect(res.body.language).to.equal("en");

// chai.request(server)
// .post('/' + savedProject._id + '/faq_kb')
// .auth(email, pwd)
// .send({ "name": "testbot3", type: "external", language: 'en' })
// .end((err, res) => {

// if (log) { console.log("res.body", res.body); }

// res.should.have.status(403);
// res.body.should.be.a('object');
// expect(res.body.success).to.equal(false);
// expect(res.body.error).to.equal("Maximum number of chatbots reached for the current plan");
// expect(res.body.plan_limit).to.equal(2);

// done()

// });
// });
// });
// })
// });
// });

// }).timeout(20000);


// it('train with tiledesk-ai', (done) => {
// var email = "test-signup-" + Date.now() + "@email.com";
// var pwd = "pwd";
Expand Down
7 changes: 3 additions & 4 deletions test/messageService.js
Original file line number Diff line number Diff line change
Expand Up @@ -69,9 +69,8 @@ describe('messageService', function () {
let obj = { createdAt: new Date() }

let quotes = await qm.getAllQuotes(savedProject, obj);
console.log("quotes: ", quotes);
quotes.messages.quote.should.be.a('string');
expect(quotes.messages.quote).to.equal('1');
quotes.messages.quote.should.be.a('number');
expect(quotes.messages.quote).to.equal(1);

done();

Expand All @@ -84,7 +83,7 @@ describe('messageService', function () {
});

});
});
}).timeout(10000);

it('createMessage', function (done) {
// this.timeout(10000);
Expand Down
104 changes: 85 additions & 19 deletions test/openaiRoute.js
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
//During the test the env variable is set to test
process.env.NODE_ENV = 'test';
// process.env.GPTKEY = 'customgptkey';

let log = false;
var projectService = require('../services/projectService');
Expand Down Expand Up @@ -67,7 +68,43 @@ describe('OpenaiRoute', () => {
// })
// }).timeout(20000)

it('completions missing gptkey', (done) => {
// it('completions missing gptkey', (done) => {

// var email = "test-signup-" + Date.now() + "@email.com";
// var pwd = "pwd";

// userService.signup(email, pwd, "Test Firstname", "Test Lastname").then((savedUser) => {
// projectService.create("test-openai-create", savedUser._id).then((savedProject) => {

// chai.request(server)
// .post('/' + savedProject._id + '/kbsettings')
// .auth(email, pwd)
// .send({}) // can be empty
// .end((err, res) => {
// if (log) { console.log("create kbsettings res.body: ", res.body); }
// res.should.have.status(200);
// res.body.should.be.a('object');

// chai.request(server)
// .post('/' + savedProject._id + '/openai')
// .auth(email, pwd)
// .send({ question: "Provide 3 names for a dog", context: "you are an awesome assistant", max_tokens: 100, temperature: 0, model: "gpt-3.5-turbo" })
// .end((err, res) => {
// if (log) { console.log("res.body: ", res.body); }
// console.log("res.body: ", res.body)
// // res.should.have.status(400);
// // res.body.should.be.a('object');
// // expect(res.body.success).to.equal(false);
// // expect(res.body.message).to.equal("Missing gptkey parameter");

// done();
// })
// })
// })
// })
// }).timeout(20000)

it('newCompletionsMissingGptkey', (done) => {

var email = "test-signup-" + Date.now() + "@email.com";
var pwd = "pwd";
Expand All @@ -76,32 +113,61 @@ describe('OpenaiRoute', () => {
projectService.create("test-openai-create", savedUser._id).then((savedProject) => {

chai.request(server)
.post('/' + savedProject._id + '/kbsettings')
.post('/' + savedProject._id + '/openai')
.auth(email, pwd)
.send({}) // can be empty
.send({ question: "Provide 3 names for a dog", context: "you are an awesome assistant", max_tokens: 100, temperature: 0, model: "gpt-3.5-turbo" })
.end((err, res) => {
if (log) { console.log("create kbsettings res.body: ", res.body); }
res.should.have.status(200);
if (log) { console.log("res.body: ", res.body); }
res.should.have.status(400);
res.body.should.be.a('object');
expect(res.body.success).to.equal(false);
expect(res.body.message).to.equal("Missing gptkey parameter");

chai.request(server)
.post('/' + savedProject._id + '/openai')
.auth(email, pwd)
.send({ question: "Provide 3 names for a dog", context: "you are an awesome assistant", max_tokens: 100, temperature: 0, model: "gpt-3.5-turbo" })
.end((err, res) => {
if (log) {console.log("res.body: ", res.body);}
res.should.have.status(400);
res.body.should.be.a('object');
expect(res.body.success).to.equal(false);
expect(res.body.message).to.equal("Missing gptkey parameter");

done();
})
done();
})


})
})
}).timeout(20000)
}).timeout(10000)


/**
* This test will be no longer available after merge with master because
* the profile section can no longer be modified via api.
*/
// it('completionsWithProfileModified', (done) => {

// var email = "test-signup-" + Date.now() + "@email.com";
// var pwd = "pwd";

// userService.signup(email, pwd, "Test Firstname", "Test Lastname").then((savedUser) => {
// projectService.create("test-openai-create", savedUser._id).then((savedProject) => {

// chai.request(server)
// .put('/projects/' + savedProject._id)
// .auth(email, pwd)
// .send({ profile: { quotes: { tokens: 400000, kbs: 100 } } })
// .end((err, res) => {
// if (log) { console.log("res.body: ", res.body); };
// console.log("res.body: ", res.body)

// chai.request(server)
// .post('/' + savedProject._id + '/openai')
// .auth(email, pwd)
// .send({ question: "Provide 3 names for a dog", context: "you are an awesome assistant", max_tokens: 100, temperature: 0, model: "gpt-3.5-turbo" })
// .end((err, res) => {

// if (log) { console.log("res.body: ", res.body); }
// done();

// })


// })
// })
// })
// }).timeout(10000)

});

Expand Down
8 changes: 6 additions & 2 deletions test/quoteManager.js
Original file line number Diff line number Diff line change
@@ -1,3 +1,6 @@
process.env.NODE_ENV = 'test';
process.env.QUOTES_ENABLED = 'true';

const { QuoteManager } = require('../services/QuoteManager');
const pubModulesManager = require('../pubmodules/pubModulesManager'); // on constructor init is undefined beacusae pub module is loaded after
var projectService = require('../services/projectService');
Expand Down Expand Up @@ -198,7 +201,7 @@ describe('QuoteManager', function () {
chai.request(server)
.post('/' + savedProject._id + "/openai/quotes")
.auth(email, pwd)
.send({ createdAt: createdAt , tokens: 128 })
.send({ createdAt: createdAt , tokens: 128, multiplier: 25 })
.end((err, res) => {
if (log) { console.log("res.body", res.body )};
res.should.have.status(200);
Expand All @@ -210,7 +213,8 @@ describe('QuoteManager', function () {
let message_resp = "value incremented for key " + key;
expect(res.body.message).to.equal(message_resp);
expect(res.body.key).to.equal(key);
expect(res.body.currentQuote).to.equal(128);
let expected_quote = 128 * 25;
expect(res.body.currentQuote).to.equal(expected_quote);


done();
Expand Down
Loading

0 comments on commit 5fc81b7

Please sign in to comment.