From 190326a29dc988e4cd918c10ed4d0c55bd612b44 Mon Sep 17 00:00:00 2001 From: Andrew Shiroma Date: Fri, 14 Mar 2025 17:56:40 -0700 Subject: [PATCH] feat: create and export UsageLimitedChatService for when user exceeds usage limit --- .../src/modules/puterai/AIChatService.js | 198 +++++++++++------- .../src/modules/puterai/PuterAIModule.js | 3 + .../puterai/UsageLimitedChatService.js | 170 +++++++++++++++ 3 files changed, 297 insertions(+), 74 deletions(-) create mode 100644 src/backend/src/modules/puterai/UsageLimitedChatService.js diff --git a/src/backend/src/modules/puterai/AIChatService.js b/src/backend/src/modules/puterai/AIChatService.js index 45d6967cc..67f08deb4 100644 --- a/src/backend/src/modules/puterai/AIChatService.js +++ b/src/backend/src/modules/puterai/AIChatService.js @@ -351,7 +351,7 @@ class AIChatService extends BaseService { * @param {string} options.model - The name of a model to use * @returns {TypedValue|Object} Returns either a TypedValue with streaming response or a completion object */ - async complete (parameters) { + async complete(parameters) { const client_driver_call = Context.get('client_driver_call'); let { test_mode, intended_service, response_metadata } = client_driver_call; @@ -363,49 +363,61 @@ class AIChatService extends BaseService { parameters }; await svc_event.emit('ai.prompt.validate', event); - if ( ! event.allow ) { + if (!event.allow) { test_mode = true; } - if ( parameters.messages ) { - parameters.messages = + if (parameters.messages) { + parameters.messages = Messages.normalize_messages(parameters.messages); } - - if ( ! test_mode && ! await this.moderate(parameters) ) { + + if (!test_mode && !await this.moderate(parameters)) { test_mode = true; } - - if ( ! test_mode ) { + + if (!test_mode) { Context.set('moderated', true); } - - if ( test_mode ) { + + if (test_mode) { intended_service = 'fake-chat'; - if ( event.abuse ) { + if (event.abuse) { parameters.model = 'abuse'; } } - - if ( parameters.tools ) { + + if (parameters.tools) { FunctionCalling.normalize_tools_object(parameters.tools); } - - if ( intended_service === this.service_name ) { + + if (intended_service === this.service_name) { throw new Error('Calling ai-chat directly is not yet supported'); } - + const svc_driver = this.services.get('driver'); let ret, error; let service_used = intended_service; let model_used = this.get_model_from_request(parameters, { intended_service }); - await this.check_usage_({ + + // 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, }); + + // Handle usage limits reached case + if (!usageAllowed) { + // The check_usage_ method has already updated the intended_service to 'usage-limited-chat' + service_used = 'usage-limited-chat'; + model_used = 'usage-limited'; + // Update intended_service to match service_used + intended_service = service_used; + } + try { ret = await svc_driver.call_new_({ actor: Context.get('actor'), @@ -418,10 +430,10 @@ class AIChatService extends BaseService { } catch (e) { const tried = []; let model = model_used; - + // TODO: if conflict models exist, add service name tried.push(model); - + error = e; // Distinguishing between user errors and service errors @@ -429,100 +441,128 @@ class AIChatService extends BaseService { // services. This is a best-effort attempt to catch user // errors and throw them as 400s. const is_request_error = (() => { - if ( e instanceof APIError ) { + if (e instanceof APIError) { return true; } - if ( e.type === 'invalid_request_error' ) { + if (e.type === 'invalid_request_error') { return true; } let some_error = e; - while ( some_error ) { - if ( some_error.type === 'invalid_request_error' ) { + while (some_error) { + if (some_error.type === 'invalid_request_error') { return true; } some_error = some_error.error ?? some_error.cause; } return false; })(); - - if ( is_request_error ) { + + if (is_request_error) { throw APIError.create('error_400_from_delegate', null, { delegate: intended_service, message: e.message, - }) + }); } console.error(e); - - if ( config.disable_fallback_mechanisms ) { + + if (config.disable_fallback_mechanisms) { throw e; } - + this.log.error('error calling service', { intended_service, model, error: e, }); - while ( !! error ) { + while (!!error) { const fallback = this.get_fallback_model({ model, tried, }); - - if ( ! fallback ) { + + if (!fallback) { throw new Error('no fallback model available'); } - + const { fallback_service_name, fallback_model_name, } = fallback; - + this.log.warn('model fallback', { intended_service, fallback_service_name, fallback_model_name }); - - await this.check_usage_({ + + // 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, }); - try { + + // If usage not allowed for fallback, use usage-limited-chat instead + if (!fallbackUsageAllowed) { + // The check_usage_ method has already updated intended_service + service_used = 'usage-limited-chat'; + model_used = 'usage-limited'; + // Clear the error to exit the fallback loop + error = null; + + // Call the usage-limited service ret = await svc_driver.call_new_({ actor: Context.get('actor'), - service_name: fallback_service_name, + service_name: 'usage-limited-chat', skip_usage: true, iface: 'puter-chat-completion', method: 'complete', - args: { - ...parameters, + args: parameters, + }); + } else { + // Normal fallback flow continues + try { + ret = await svc_driver.call_new_({ + actor: Context.get('actor'), + service_name: fallback_service_name, + skip_usage: true, + iface: 'puter-chat-completion', + method: 'complete', + args: { + ...parameters, + model: fallback_model_name, + }, + }); + error = null; + service_used = fallback_service_name; + model_used = fallback_model_name; + response_metadata.fallback = { + service: fallback_service_name, model: fallback_model_name, - }, - }); - error = null; - service_used = fallback_service_name; - model_used = fallback_model_name; - response_metadata.fallback = { - service: fallback_service_name, - model: fallback_model_name, - tried: tried, - }; - } catch (e) { - error = e; - tried.push(fallback_model_name); - this.log.error('error calling fallback', { - intended_service, - model, - error: e, - }); + tried: tried, + }; + } catch (e) { + error = e; + tried.push(fallback_model_name); + this.log.error('error calling fallback', { + intended_service, + model, + error: e, + }); + } } } } + ret.result.via_ai_chat_service = true; response_metadata.service_used = service_used; - + + // Add flag if we're using the usage-limited service + if (service_used === 'usage-limited-chat') { + response_metadata.usage_limited = true; + } + const username = Context.get('actor').type?.user?.username; - + if ( // Check if we have 'ai-chat-intermediate' response type; // this means we're streaming and usage comes from a promise. @@ -540,19 +580,19 @@ class AIChatService extends BaseService { usage, }); })(); - - if ( ret.result.value.init_chat_stream ) { + + if (ret.result.value.init_chat_stream) { const stream = new PassThrough(); const retval = new TypedValue({ $: 'stream', content_type: 'application/x-ndjson', chunked: true, }, stream); - + const chatStream = new Streaming.AIChatStream({ stream, }); - + (async () => { try { await ret.result.value.init_chat_stream({ chatStream }); @@ -564,10 +604,10 @@ class AIChatService extends BaseService { stream.end(); } })(); - + return retval; } - + return ret.result.value.response; } else { await svc_event.emit('ai.prompt.report-usage', { @@ -588,9 +628,9 @@ class AIChatService extends BaseService { model_used, service_used, }); - - - if ( parameters.response?.normalize ) { + + + if (parameters.response?.normalize) { ret.result.message = Messages.normalize_single_message(ret.result.message); ret.result = { @@ -599,7 +639,7 @@ class AIChatService extends BaseService { normalized: true, }; } - + return ret.result; } } @@ -621,7 +661,7 @@ class AIChatService extends BaseService { const svc_event = this.services.get('event'); const reading = await svc_permission.scan(actor, `paid-services:ai-chat`); const options = PermissionUtil.reading_to_options(reading); - + // Query current ai usage in terms of cost const [row] = await this.db.read( 'SELECT SUM(`cost`) AS sum FROM `ai_usage` ' + @@ -639,10 +679,20 @@ class AIChatService extends BaseService { permission_options: options, }; await svc_event.emit('ai.prompt.check-usage', event); - if ( event.error ) throw event.error; - if ( ! event.allowed ) { - throw new APIError('forbidden'); + + // Check if usage is not allowed or there's an error + if (event.error || !event.allowed) { + // Instead of throwing an error, modify the intended_service + const client_driver_call = Context.get('client_driver_call'); + client_driver_call.intended_service = 'usage-limited-chat'; + client_driver_call.response_metadata.usage_limited = true; + + // Return false to indicate that usage is not allowed and service has been changed + return false; } + + // Usage is allowed + return true; } diff --git a/src/backend/src/modules/puterai/PuterAIModule.js b/src/backend/src/modules/puterai/PuterAIModule.js index b6fa45a14..5348fa61f 100644 --- a/src/backend/src/modules/puterai/PuterAIModule.js +++ b/src/backend/src/modules/puterai/PuterAIModule.js @@ -116,6 +116,9 @@ class PuterAIModule extends AdvancedBase { const{ AITestModeService } = require('./AITestModeService'); services.registerService('ai-test-mode', AITestModeService); + + const { UsageLimitedChatService } = require('./UsageLimitedChatService'); + services.registerService('usage-limited-chat', UsageLimitedChatService); } } diff --git a/src/backend/src/modules/puterai/UsageLimitedChatService.js b/src/backend/src/modules/puterai/UsageLimitedChatService.js new file mode 100644 index 000000000..904703a99 --- /dev/null +++ b/src/backend/src/modules/puterai/UsageLimitedChatService.js @@ -0,0 +1,170 @@ +/* + * Copyright (C) 2024-present Puter Technologies Inc. + * + * This file is part of Puter. + * + * Puter is free software: you can redistribute it and/or modify + * it under the terms of the GNU Affero General Public License as published + * by the Free Software Foundation, either version 3 of the License, or + * (at your option) any later version. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * GNU Affero General Public License for more details. + * + * You should have received a copy of the GNU Affero General Public License + * along with this program. If not, see . + */ + +// METADATA // {"ai-commented":{"service":"claude"}} +const { default: dedent } = require("dedent"); +const BaseService = require("../../services/BaseService"); +const { PassThrough } = require("stream"); +const { TypedValue } = require("../../services/drivers/meta/Runtime"); +const Streaming = require("./lib/Streaming"); + +/** +* UsageLimitedChatService - A specialized chat service that returns resource exhaustion messages. +* Extends BaseService to provide responses indicating the user has exceeded their usage limits. +* Follows the same response format as real AI providers but with a custom message about upgrading. +* Can handle both streaming and non-streaming requests consistently. +*/ +class UsageLimitedChatService extends BaseService { + get_default_model () { + return 'usage-limited'; + } + + static IMPLEMENTS = { + ['puter-chat-completion']: { + /** + * Returns a list of available model names + * @returns {Promise} Array containing the single model identifier + */ + async list () { + return ['usage-limited']; + }, + + /** + * Returns model details for the usage-limited model + * @returns {Promise} Array containing the model details + */ + async models () { + return [{ + id: 'usage-limited', + name: 'Usage Limited', + context: 16384, + cost: { + currency: 'usd-cents', + tokens: 1_000_000, + input: 0, + output: 0, + }, + }]; + }, + + /** + * Simulates a chat completion request with a usage limit message + * @param {Object} params - The completion parameters + * @param {Array} params.messages - Array of chat messages (unused) + * @param {boolean} params.stream - Whether to stream the response + * @param {string} params.model - The model to use (unused) + * @returns {Object|TypedValue} A chat completion response or streamed response + */ + async complete ({ messages, stream, model, customLimitMessage }) { + const limitMessage = customLimitMessage || dedent(` + You have reached your AI usage limit for this account. + `); + + // If streaming is requested, return a streaming response + if (stream) { + const streamObj = new PassThrough(); + const retval = new TypedValue({ + $: 'stream', + content_type: 'application/x-ndjson', + chunked: true, + }, streamObj); + + const chatStream = new Streaming.AIChatStream({ + stream: streamObj, + }); + + // Schedule the streaming response + setTimeout(() => { + chatStream.write({ + type: 'content_block_start', + index: 0, + }); + + chatStream.write({ + type: 'content_block_delta', + index: 0, + delta: { + type: 'text', + text: limitMessage, + }, + }); + + chatStream.write({ + type: 'content_block_stop', + index: 0, + }); + + chatStream.write({ + type: 'message_stop', + stop_reason: 'end_turn', + }); + + chatStream.end(); + }, 10); + + // Return a TypedValue with usage_promise for proper integration + return new TypedValue({ $: 'ai-chat-intermediate' }, { + stream: true, + init_chat_stream: async ({ chatStream: cs }) => { + // Copy contents from our stream to the provided one + chatStream.pipe(cs.stream); + }, + usage_promise: Promise.resolve({ + input_tokens: 0, + output_tokens: 1, + }), + }); + } + + // Non-streaming response + return { + "index": 0, + message: { + "id": "00000000-0000-0000-0000-000000000000", + "type": "message", + "role": "assistant", + "model": "usage-limited", + "content": [ + { + "type": "text", + "text": limitMessage + } + ], + "stop_reason": "end_turn", + "stop_sequence": null, + "usage": { + "input_tokens": 0, + "output_tokens": 1 + } + }, + "usage": { + "input_tokens": 0, + "output_tokens": 1 + }, + "logprobs": null, + "finish_reason": "stop" + }; + } + } + } +} + +module.exports = { + UsageLimitedChatService, +}; \ No newline at end of file