lee/puter
1
0
Fork 0
mirror of https://github.com/HeyPuter/puter.git synced 2026-08-21 13:38:14 +00:00
Code Issues Actions 9 Packages Projects Releases Wiki Activity

feat: create and export UsageLimitedChatService for when user exceeds usage limit

Browse Source
This commit is contained in:
Andrew Shiroma
2025-03-14 17:56:40 -07:00
parent 3cdbcd83b3
commit 190326a29d
3 changed files with 297 additions and 74 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/backend/src/modules/puterai/AIChatService.js
+124 -74
View File
@@ -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;
}
src/backend/src/modules/puterai/PuterAIModule.js
+3
View File
@@ -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);
}
}
src/backend/src/modules/puterai/UsageLimitedChatService.js
+170
View File
@@ -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 <https://www.gnu.org/licenses/>.
*/
// 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<string[]>} Array containing the single model identifier
*/
async list () {
return ['usage-limited'];
},
/**
* Returns model details for the usage-limited model
* @returns {Promise<Object[]>} 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,
};