From 21bd427cef5c453697d706511fba4bf611ae75e3 Mon Sep 17 00:00:00 2001 From: KernelDeimos Date: Tue, 1 Apr 2025 17:42:37 -0400 Subject: [PATCH] refactor: generalize handling of costly services Usage/cost tracking per user was over-fitted to LLM/AI services. This commit adds CostService and updates AIChatService to use that instead of updating the `ai_usage` table. The `ai_usage` table is now considered deprecated. A service called PermissiveCreditService is enabled by default in SelfHostedModule. This is so users who host their own Puter instances don't have to worry about usage limits, since they're providing their own keys anyway. --- src/backend/src/CoreModule.js | 3 ++ .../src/modules/puterai/AIChatService.js | 19 +++---- .../selfhosted/PermissiveCreditService.js | 22 ++++++++ .../modules/selfhosted/SelfHostedModule.js | 3 ++ .../src/services/drivers/CostService.js | 53 +++++++++++++++++++ .../src/services/drivers/DriverService.js | 15 +++++- 6 files changed, 101 insertions(+), 14 deletions(-) create mode 100644 src/backend/src/modules/selfhosted/PermissiveCreditService.js create mode 100644 src/backend/src/services/drivers/CostService.js diff --git a/src/backend/src/CoreModule.js b/src/backend/src/CoreModule.js index 10cf8cc34..b4d37e8de 100644 --- a/src/backend/src/CoreModule.js +++ b/src/backend/src/CoreModule.js @@ -293,6 +293,9 @@ const install = async ({ services, app, useapi, modapi }) => { const { DevTODService } = require('./services/DevTODService'); services.registerService('__dev-tod', DevTODService); + const { CostService } = require("./services/drivers/CostService"); + services.registerService('cost', CostService); + const { DriverService } = require("./services/drivers/DriverService"); services.registerService('driver', DriverService); diff --git a/src/backend/src/modules/puterai/AIChatService.js b/src/backend/src/modules/puterai/AIChatService.js index 0d39a42e6..5cd6035cd 100644 --- a/src/backend/src/modules/puterai/AIChatService.js +++ b/src/backend/src/modules/puterai/AIChatService.js @@ -128,8 +128,9 @@ class AIChatService extends BaseService { this.log.noticeme('COST INFO', values); - await this.db.insert('ai_usage', values); - + const svc_cost = this.services.get('cost'); + svc_cost.record_cost({ cost: values.cost }); + // USD cost from microcents const cost_usc = values.cost / 1000000; const cost_usd = cost_usc / 100; @@ -404,11 +405,8 @@ class AIChatService extends BaseService { }); // Updated: Check usage and get a boolean result instead of throwing error - const usageAllowed = await this.check_usage_({ - actor: Context.get('actor'), - service: service_used, - model: model_used, - }); + const svc_cost = this.services.get('cost'); + const usageAllowed = await svc_cost.get_funding_allowed(); // Handle usage limits reached case if ( !usageAllowed ) { @@ -496,11 +494,8 @@ class AIChatService extends BaseService { }); // Check usage for fallback model too (with updated method) - const fallbackUsageAllowed = await this.check_usage_({ - actor: Context.get('actor'), - service: fallback_service_name, - model: fallback_model_name, - }); + const svc_cost = this.services.get('cost'); + const fallbackUsageAllowed = await svc_cost.get_funding_allowed(); // If usage not allowed for fallback, use usage-limited-chat instead if (!fallbackUsageAllowed) { diff --git a/src/backend/src/modules/selfhosted/PermissiveCreditService.js b/src/backend/src/modules/selfhosted/PermissiveCreditService.js new file mode 100644 index 000000000..ffba55fd5 --- /dev/null +++ b/src/backend/src/modules/selfhosted/PermissiveCreditService.js @@ -0,0 +1,22 @@ +const BaseService = require("../../services/BaseService"); + +/** + * PermissiveCreditService listens to the event where DriverService asks + * for a credit context, and always provides one that allows use of + * cost-incurring services for no charge. This grants free use to + * everyone to services that incur a cost, as long as the user has + * permission to call the respective service. + */ +class PermissiveCreditService extends BaseService { + static MODULES = { + uuidv4: require('uuid').v4, + } + _init () { + const svc_event = this.services.get('event'); + svc_event.on(`credit.check-available`, (_, event) => { + event.available = Number.MAX_SAFE_INTEGER; + }); + } +} + +module.exports = PermissiveCreditService; diff --git a/src/backend/src/modules/selfhosted/SelfHostedModule.js b/src/backend/src/modules/selfhosted/SelfHostedModule.js index e9ad85423..02c867711 100644 --- a/src/backend/src/modules/selfhosted/SelfHostedModule.js +++ b/src/backend/src/modules/selfhosted/SelfHostedModule.js @@ -35,6 +35,9 @@ class SelfHostedModule extends AdvancedBase { const DevWatcherService = require('./DevWatcherService'); const path_ = require('path'); + const PermissiveCreditService = require("./PermissiveCreditService"); + services.registerService('__permissive-credit', PermissiveCreditService); + const { DBKVService } = require("../../services/DBKVService"); services.registerService('puter-kvstore', DBKVService); diff --git a/src/backend/src/services/drivers/CostService.js b/src/backend/src/services/drivers/CostService.js new file mode 100644 index 000000000..e2837161b --- /dev/null +++ b/src/backend/src/services/drivers/CostService.js @@ -0,0 +1,53 @@ +const APIError = require("../../api/APIError"); +const { Context } = require("../../util/context"); +const BaseService = require("../BaseService"); + +class CreditContext { + constructor (svc, o) { + for ( const k in o ) this[k] = o[k]; + this.svc = svc; + } +} + +class CostService extends BaseService { + static MODULES = { + uuidv4: require('uuid').v4, + } + + _init () { + const svc_cost = this; + const svc_event = this.services.get('event'); + svc_event.on('driver.create-call-context', async (_, event) => { + event.context = event.context.sub({ + // Future Use + }); + }); + } + async get_credit_context (params) { + return new CreditContext(this, params); + } + + async get_funding_allowed () { + const cost_uuid = this.modules.uuidv4(); + const svc_event = this.services.get('event'); + const event = { + actor: Context.get('actor'), + available: 0, + cost_uuid, + }; + await svc_event.emit('credit.check-available', event); + return event.available > 0; + } + async record_cost ({ cost }) { + const svc_event = this.services.get('event'); + const event = { + actor: Context.get('actor'), + cost, + }; + await svc_event.emit('credit.record-cost', event); + } +} + +module.exports = { + CostService, +}; diff --git a/src/backend/src/services/drivers/DriverService.js b/src/backend/src/services/drivers/DriverService.js index 58d41a693..2a15bcfbc 100644 --- a/src/backend/src/services/drivers/DriverService.js +++ b/src/backend/src/services/drivers/DriverService.js @@ -268,9 +268,20 @@ class DriverService extends BaseService { skip_usage = true; } - return await Context.sub({ + const svc_event = this.services.get('event'); + const event = {}; + event.context = Context.sub({ client_driver_call, - }).arun(async () => { + }); + event.call_details = { + service: driver, + iface, method, args, + skip_usage, + }; + + svc_event.emit('driver.create-call-context', event); + + return event.context.arun(async () => { const result = await this.call_new_({ actor, service,