diff --git a/.gitignore b/.gitignore
index ab697c25d..5abeff98d 100644
--- a/.gitignore
+++ b/.gitignore
@@ -29,4 +29,4 @@ dist/
# Local Netlify folder
.netlify
-src/emulator/release/
+src/emulator/release/
\ No newline at end of file
diff --git a/src/backend/.gitignore b/src/backend/.gitignore
index 0bd7d2ef3..ed0ef7365 100644
--- a/src/backend/.gitignore
+++ b/src/backend/.gitignore
@@ -148,4 +148,4 @@ creds*
thumbnail-service
# init sql generated from ./run.sh
-init.sql
+init.sql
\ No newline at end of file
diff --git a/src/backend/src/modules/puterai/AIChatService.js b/src/backend/src/modules/puterai/AIChatService.js
index 45d6967cc..f2ff4e7b4 100644
--- a/src/backend/src/modules/puterai/AIChatService.js
+++ b/src/backend/src/modules/puterai/AIChatService.js
@@ -366,7 +366,7 @@ class AIChatService extends BaseService {
if ( ! event.allow ) {
test_mode = true;
}
-
+
if ( parameters.messages ) {
parameters.messages =
Messages.normalize_messages(parameters.messages);
@@ -401,11 +401,23 @@ class AIChatService extends BaseService {
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'),
@@ -423,7 +435,7 @@ class AIChatService extends BaseService {
tried.push(model);
error = e;
-
+
// Distinguishing between user errors and service errors
// is very messy because of different conventions between
// services. This is a best-effort attempt to catch user
@@ -482,45 +494,73 @@ class AIChatService extends BaseService {
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 (
@@ -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');
+
+ // If the user has exceeded their usage limit, apply usage-limited-chat which lets them know
+ 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 the user has gone over their limit and service has been changed
+ return false;
}
+
+ // Return true if the user has tokens to spend
+ 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..777470941
--- /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