From 5746c7156afef8403b121d5093badeabe2aedf42 Mon Sep 17 00:00:00 2001
From: KernelDeimos <eric.alex.dube@gmail.com>
Date: Mon, 25 Nov 2024 11:56:07 -0500
Subject: [PATCH] dev: migrate Mistral to new controller

---
 .../src/modules/puterai/AIChatService.js      |  18 ++-
 .../src/modules/puterai/MistralAIService.js   | 123 +++++++++++++++++-
 2 files changed, 136 insertions(+), 5 deletions(-)

diff --git a/src/backend/src/modules/puterai/AIChatService.js b/src/backend/src/modules/puterai/AIChatService.js
index 823bb6e3..88c93d46 100644
--- a/src/backend/src/modules/puterai/AIChatService.js
+++ b/src/backend/src/modules/puterai/AIChatService.js
@@ -25,13 +25,27 @@ class AIChatService extends BaseService {
             
             // Populate simple model list
             {
-                const models = await delegate.list();
+                const models = await (async () => {
+                    try {
+                        return await delegate.list() ?? [];
+                    } catch (e) {
+                        this.log.error(e);
+                        return [];
+                    }
+                })();
                 this.simple_model_list.push(...models);
             }
 
             // Populate detail model list and map
             {
-                const models = await delegate.models();
+                const models = await (async () => {
+                    try {
+                        return await delegate.models() ?? [];
+                    } catch (e) {
+                        this.log.error(e);
+                        return [];
+                    }
+                })();
                 const annotated_models = [];
                 for ( const model of models ) {
                     annotated_models.push({
diff --git a/src/backend/src/modules/puterai/MistralAIService.js b/src/backend/src/modules/puterai/MistralAIService.js
index b57219bb..195bd8d8 100644
--- a/src/backend/src/modules/puterai/MistralAIService.js
+++ b/src/backend/src/modules/puterai/MistralAIService.js
@@ -3,23 +3,140 @@ const BaseService = require("../../services/BaseService");
 const { TypedValue } = require("../../services/drivers/meta/Runtime");
 const { nou } = require("../../util/langutil");
 
+const axios = require('axios');
+
 class MistralAIService extends BaseService {
     static MODULES = {
         '@mistralai/mistralai': require('@mistralai/mistralai'),
     }
+    _construct () {
+        this.costs_ = {
+            'mistral-large-latest': {
+                currency: 'usd-cents',
+                tokens: 1_000_000,
+                input: 200,
+                output: 600,
+            },
+            'pixtral-large-latest': {
+                currency: 'usd-cents',
+                tokens: 1_000_000,
+                input: 200,
+                output: 600,
+            },
+            'mistral-small-latest': {
+                currency: 'usd-cents',
+                tokens: 1_000_000,
+                input: 20,
+                output: 60,
+            },
+            'codestral-latest': {
+                currency: 'usd-cents',
+                tokens: 1_000_000,
+                input: 20,
+                output: 60,
+            },
+            'ministral-8b-latest': {
+                currency: 'usd-cents',
+                tokens: 1_000_000,
+                input: 10,
+                output: 10,
+            },
+            'ministral-3b-latest': {
+                currency: 'usd-cents',
+                tokens: 1_000_000,
+                input: 4,
+                output: 4,
+            },
+            'pixtral-12b': {
+                currency: 'usd-cents',
+                tokens: 1_000_000,
+                input: 15,
+                output: 15,
+            },
+            'mistral-nemo': {
+                currency: 'usd-cents',
+                tokens: 1_000_000,
+                input: 15,
+                output: 15,
+            },
+            'open-mistral-7b': {
+                currency: 'usd-cents',
+                tokens: 1_000_000,
+                input: 25,
+                output: 25,
+            },
+            'open-mixtral-8x7b': {
+                currency: 'usd-cents',
+                tokens: 1_000_000,
+                input: 7,
+                output: 7,
+            },
+            'open-mixtral-8x22b': {
+                currency: 'usd-cents',
+                tokens: 1_000_000,
+                input: 2,
+                output: 6,
+            },
+        };
+    }
     async _init () {
         const require = this.require;
         const { Mistral } = require('@mistralai/mistralai');
+        this.api_base_url = 'https://api.mistral.ai/v1';
         this.client = new Mistral({
             apiKey: this.config.apiKey,
         });
+
+        const svc_aiChat = this.services.get('ai-chat');
+        svc_aiChat.register_provider({
+            service_name: this.service_name,
+            alias: true,
+        });
+
+        // TODO: make this event-driven so it doesn't hold up boot
+        await this.populate_models_();
+    }
+    async populate_models_ () {
+        const resp = await axios({
+            method: 'get',
+            url: this.api_base_url + '/models',
+            headers: {
+                Authorization: `Bearer ${this.config.apiKey}`
+            }
+        })
+
+        const response_json = resp.data;
+        const models = response_json.data;
+        this.models_array_ = [];
+        for ( const api_model of models ) {
+            
+            let cost = this.costs_[api_model.id];
+            if ( ! cost ) for ( const alias of api_model.aliases ) {
+                cost = this.costs_[alias];
+                if ( cost ) break;
+            }
+            if ( ! cost ) continue;
+            const model = {
+                id: api_model.id,
+                name: api_model.description,
+                aliases: api_model.aliases,
+                context: api_model.max_context_length,
+                capabilities: api_model.capabilities,
+                vision: api_model.capabilities.vision,
+                cost,
+            };
+
+            this.models_array_.push(model);
+        }
+        // return resp.data;
     }
     static IMPLEMENTS = {
         'puter-chat-completion': {
+            async models () {
+                return this.models_array_;
+            },
             async list () {
-                // They send: { "object": "list", data }
-                const funny_wrapper = await this.client.models.list();
-                return funny_wrapper.data;
+                return this.models_array_.map(m => m.id);
             },
             async complete ({ messages, stream, model }) {