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