dev: fix test_mode flag

This commit is contained in:
KernelDeimos
2024-11-20 11:57:41 -05:00
parent 5ab87979bc
commit d911754e2b
7 changed files with 60 additions and 1 deletions
@@ -0,0 +1,12 @@
const BaseService = require("../../services/BaseService");
class AITestModeService extends BaseService {
async _init () {
const svc_driver = this.services.get('driver');
svc_driver.register_test_service('puter-chat-completion', 'openai-completion');
}
}
module.exports = {
AITestModeService,
};
@@ -12,6 +12,11 @@ class AWSPollyService extends BaseService {
}
static IMPLEMENTS = {
['driver-capabilities']: {
supports_test_mode (iface, method_name) {
return iface === 'puter-tts' && method_name === 'synthesize';
}
},
['puter-tts']: {
async list_voices () {
const polly_voices = await this.describe_voices();
@@ -9,6 +9,11 @@ class AWSTextractService extends BaseService {
}
static IMPLEMENTS = {
['driver-capabilities']: {
supports_test_mode (iface, method_name) {
return iface === 'puter-ocr' && method_name === 'recognize';
}
},
['puter-ocr']: {
async recognize ({ source, test_mode }) {
if ( test_mode ) {
@@ -22,6 +22,12 @@ class OpenAICompletionService extends BaseService {
}
static IMPLEMENTS = {
['driver-capabilities']: {
supports_test_mode (iface, method_name) {
return iface === 'puter-chat-completion' &&
method_name === 'complete';
}
},
['puter-chat-completion']: {
async list () {
return [
@@ -56,6 +56,9 @@ class PuterAIModule extends AdvancedBase {
const { FakeChatService } = require('./FakeChatService');
services.registerService('fake-chat', FakeChatService);
const{ AITestModeService } = require('./AITestModeService');
services.registerService('ai-test-mode', AITestModeService);
}
}
@@ -226,6 +226,10 @@ class PermissionService extends BaseService {
}
async scan (actor, permission_options, _reserved, state) {
if ( ! state ) this.log.info('scan', {
actor: actor.uid,
permission_options,
});
const reading = [];
if ( ! state ) {
@@ -37,6 +37,7 @@ class DriverService extends BaseService {
_construct () {
this.drivers = {};
this.interface_to_implementation = {};
this.interface_to_test_service = {};
}
async ['__on_registry.collections'] () {
@@ -77,6 +78,10 @@ class DriverService extends BaseService {
register_driver (interface_name, implementation) {
this.interface_to_implementation[interface_name] = implementation;
}
register_test_service (interface_name, service_name) {
this.interface_to_test_service[interface_name] = service_name;
}
get_interface (interface_name) {
const o = {};
@@ -113,7 +118,8 @@ class DriverService extends BaseService {
async _call ({ driver, iface, method, args }) {
console.log('??', driver, iface, method, args);
const processed_args = await this._process_args(iface, method, args);
if ( Context.get('test_mode') ) {
const test_mode = Context.get('test_mode');
if ( test_mode ) {
processed_args.test_mode = true;
}
@@ -141,18 +147,31 @@ class DriverService extends BaseService {
driver = driver ?? iface_to_driver[iface] ?? iface;
let skip_usage = false;
if ( test_mode && this.interface_to_test_service[iface] ) {
driver = this.interface_to_test_service[iface];
}
const driver_service_exists = (() => {
console.log('CHECKING FOR THIS', driver, iface);
return this.services.has(driver) &&
this.services.get(driver).list_traits()
.includes(iface);
})();
if ( driver_service_exists ) {
const service = this.services.get(driver);
const caps = service.as('driver-capabilities');
if ( test_mode && caps && caps.supports_test_mode(iface, method) ) {
skip_usage = true;
}
return await this.call_new_({
actor,
service,
service_name: driver,
iface, method, args: processed_args,
skip_usage,
});
}
@@ -240,6 +259,7 @@ class DriverService extends BaseService {
service,
service_name,
iface, method, args,
skip_usage,
}) {
const svc_permission = this.services.get('permission');
const reading = await svc_permission.scan(
@@ -324,6 +344,8 @@ class DriverService extends BaseService {
{
name: 'enforce monthly usage limit',
on_call: async args => {
if ( skip_usage ) return args;
// Typo-Tolerance
if ( effective_policy?.['monthy-limit'] ) {
effective_policy['monthly-limit'] = effective_policy['monthy-limit'];
@@ -343,6 +365,8 @@ class DriverService extends BaseService {
return args;
},
on_return: async result => {
if ( skip_usage ) return result;
console.log('monthly usage is returning');
const svc_monthlyUsage = services.get('monthly-usage');
const extra = {