Bläddra i källkod

feat: add test_mode flag

KernelDeimos 9 månader sedan
förälder
incheckning
9a9bd5eaf0

+ 18 - 0
src/backend/src/modules/puterai/AWSTextractService.js

@@ -1,6 +1,7 @@
 const { TextractClient, AnalyzeDocumentCommand, InvalidS3ObjectException } = require("@aws-sdk/client-textract");
 
 const BaseService = require("../../services/BaseService");
+const APIError = require("../../api/APIError");
 
 class AWSTextractService extends BaseService {
     _construct () {
@@ -10,6 +11,23 @@ class AWSTextractService extends BaseService {
     static IMPLEMENTS = {
         ['puter-ocr']: {
             async recognize ({ source, test_mode }) {
+                if ( test_mode ) {
+                    return {
+                        blocks: [
+                            {
+                                type: 'text/textract:WORD',
+                                confidence: 0.9999998807907104,
+                                text: 'Hello',
+                            },
+                            {
+                                type: 'text/puter:sample-output',
+                                confidence: 1,
+                                text: 'The test_mode flag is set to true. This is a sample output.',
+                            },
+                        ]
+                    };
+                }
+
                 const resp = await this.analyze_document(source);
 
                 // Simplify the response for common interface

+ 26 - 1
src/backend/src/modules/puterai/OpenAICompletionService.js

@@ -20,7 +20,32 @@ class OpenAICompletionService extends BaseService {
 
     static IMPLEMENTS = {
         ['puter-chat-completion']: {
-            async complete ({ messages, vision }) {
+            async complete ({ messages, test_mode }) {
+                if ( test_mode ) {
+                    const { LoremIpsum } = require('lorem-ipsum');
+                    const li = new LoremIpsum({
+                        sentencesPerParagraph: {
+                            max: 8,
+                            min: 4
+                        },
+                        wordsPerSentence: {
+                            max: 20,
+                            min: 12
+                        },
+                    });
+                    return {
+                        "index": 0,
+                        "message": {
+                            "role": "assistant",
+                            "content": li.generateParagraphs(
+                                Math.floor(Math.random() * 3) + 1
+                            ),
+                        },
+                        "logprobs": null,
+                        "finish_reason": "stop"
+                    }
+                }
+
                 const model = 'gpt-4o';
                 return await this.complete(messages, {
                     model,

+ 4 - 4
src/backend/src/modules/puterai/OpenAIImageGenerationService.js

@@ -19,10 +19,6 @@ class OpenAIImageGenerationService extends BaseService {
     static IMPLEMENTS = {
         ['puter-image-generation']: {
             async generate ({ prompt, test_mode }) {
-                const url = await this.generate(prompt, {
-                    ratio: this.constructor.RATIO_SQUARE,
-                });
-
                 if ( test_mode ) {
                     return new TypedValue({
                         $: 'string:url:web',
@@ -30,6 +26,10 @@ class OpenAIImageGenerationService extends BaseService {
                     }, 'https://puter-sample-data.puter.site/image_example.png');
                 }
 
+                const url = await this.generate(prompt, {
+                    ratio: this.constructor.RATIO_SQUARE,
+                });
+
                 const image = new TypedValue({
                     $: 'string:url:web',
                     content_type: 'image'