@tiledesk/tiledesk-server 2.19.3 → 2.19.5

package/services/aiManager.js

@@ -21,6 +21,7 @@ const default_engine_hybrid = require('../config/kb/engine.hybrid');
 const default_embedding = require('../config/kb/embedding');
 const integrationService = require('./integrationService');
 const situatedContext = require('../config/kb/situatedContext');
+const { MODELS_MULTIPLIER } = require('../utils/aiUtils');
 
 // Job managers
 let jobManager = new JobManager(AMQP_MANAGER_URL, {
@@ -94,7 +95,7 @@ class AiManager {
       this.saveBulk(operations, kbs, namespace.id_project, namespace.id).then( async (result) => {
         let hybrid = namespace.hybrid;
         let engine = namespace.engine || default_engine;
-        let embedding = namespace.embedding || default_embedding;
+        let embedding = namespace.embedding || Object.assign({}, default_embedding);
         embedding.api_key = process.env.EMBEDDING_API_KEY || process.env.GPTKEY;
 
         let situated_context;
@@ -514,25 +515,6 @@ class AiManager {
     return true;
   }
 
-  // from webhook
-  // async scheduleScrape(resources) {
-
-  //   let scheduler = new Scheduler({ jobManager: jobManager });
-  
-  //   resources.forEach(r => {
-  //     winston.debug("(Webhook) Schedule job with following data: ", r);
-  //     scheduler.trainSchedule(r, async (err, result) => {
-  //       if (err) {
-  //         winston.error("Scheduling error: ", err);
-  //       } else {
-  //         winston.info("Scheduling result: ", result);
-  //       }
-  //     });
-  //   })
-  
-  //   return true;
-  // }
-
   async startScrape(data) {
 
     if (!data.gptkey) {
@@ -625,6 +607,33 @@ class AiManager {
       : undefined;
   }
 
+  async getTrackingTokens({ qaResult, publicKey, model, obj, quoteManager, project }) {
+    if (publicKey !== true) {
+      return qaResult.prompt_token_size;
+    }
+    let modelKey;
+    if (typeof model === 'string') {
+      modelKey = model;
+    } else if (model && typeof model.name === 'string') {
+      modelKey = model.name;
+    }
+    let multiplier = MODELS_MULTIPLIER[modelKey];
+    if (!multiplier) {
+      multiplier = 1;
+      winston.info("No multiplier found for AI model (qa) " + modelKey);
+    }
+    obj.multiplier = multiplier;
+    obj.tokens = qaResult.prompt_token_size;
+    const incremented_key = await quoteManager.incrementTokenCount(project, obj);
+    winston.verbose("incremented_key: ", incremented_key);
+    const trackingTokens = obj.tokens * obj.multiplier;
+    if (qaResult != null && typeof qaResult.prompt_token_size === 'number') {
+      qaResult.legacy_prompt_token_size = qaResult.prompt_token_size;
+      qaResult.prompt_token_size = trackingTokens;
+    }
+    return trackingTokens;
+  }
+
 }
 
 const aiManager = new AiManager();

package/services/kbQuestionService.js

@@ -0,0 +1,98 @@
+const { Namespace, AnsweredQuestion, UnansweredQuestion } = require('../models/kb_setting');
+const winston = require('../config/winston');
+
+class KbQuestionService {
+  
+  constructor() {}
+
+  async createAnsweredQuestion(id_project, data) {
+    try {
+      const { namespace, question, answer, tokens, request_id } = data;
+
+      if (!namespace || !question || !answer) {
+        throw new Error({ status: 422, error: 'Missing required parameters: namespace, question and answer' });
+      }
+
+      const isValidNamespace = await this.validateNamespace(id_project, namespace);
+      if (!isValidNamespace) {
+        throw new Error({ status: 403, error: 'Not allowed. The namespace does not belong to the current project.' });
+      }
+
+      const answeredQuestion = new AnsweredQuestion({
+        id_project,
+        namespace,
+        question,
+        answer,
+        tokens,
+        request_id,
+      });
+
+      return await answeredQuestion.save();
+
+    } catch (error) {
+      throw error;
+    }
+  }
+
+  async createUnansweredQuestion(id_project, data) {
+    try {
+      const { namespace, question, request_id, sender } = data;
+
+      if (!namespace || !question) {
+        throw new Error({ status: 422, error: 'Missing required parameters: namespace and question' });
+      }
+
+      const isValidNamespace = await this.validateNamespace(id_project, namespace);
+      if (!isValidNamespace) {
+        throw new Error({ status: 403, error: 'Not allowed. The namespace does not belong to the current project.' });
+      }
+
+      const unansweredQuestion = new UnansweredQuestion({
+        id_project,
+        namespace,
+        question,
+        request_id,
+        sender,
+      });
+      return await unansweredQuestion.save();
+
+    } catch (error) {
+      throw error;
+    }
+  }
+
+  async validateNamespace(id_project, namespace_id) {
+    try {
+      const namespace = await Namespace.findOne({
+        id_project: id_project,
+        id: namespace_id
+      });
+      return !!namespace;
+    } catch (error) {
+      throw error;
+    }
+  }
+
+  async trackKbQaResult(id_project, { namespace, question }, answer, tokens) {
+    try {
+      if (answer.success === true) {
+        await this.createAnsweredQuestion(id_project, {
+          namespace,
+          question,
+          answer: answer.answer,
+          tokens,
+        });
+      } else if (answer.success === false) {
+        await this.createUnansweredQuestion(id_project, { namespace, question });
+      } else {
+        winston.warn('qa answer.success is not true or false: ', answer.success);
+      }
+    } catch (error) {
+      winston.error('qa err: ', error);
+    }
+  }
+
+}
+
+const kbQuestionService = new KbQuestionService();
+module.exports = kbQuestionService;