dev: add support for tools

HeyPuter · Jan 23, 2025 · 1256614 · 1256614
1 parent cf7aa27
commit 1256614
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 5 deletions.
diff --git a/src/backend/src/modules/puterai/AIInterfaceService.js b/src/backend/src/modules/puterai/AIInterfaceService.js
@@ -75,6 +75,7 @@ class AIInterfaceService extends BaseService {
                     description: 'Get completions for a chat log.',
                     parameters: {
                         messages: { type: 'json' },
+                        tools: { type: 'json' },
                         vision: { type: 'flag' },
                         stream: { type: 'flag' },
                         model: { type: 'string' },

diff --git a/src/backend/src/modules/puterai/OpenAICompletionService.js b/src/backend/src/modules/puterai/OpenAICompletionService.js
@@ -152,7 +152,7 @@ class OpenAICompletionService extends BaseService {
              * AI Chat completion method.
              * See AIChatService for more details.
              */
-            async complete ({ messages, test_mode, stream, model }) {
+            async complete ({ messages, test_mode, stream, model, tools }) {
 
                 // for now this code (also in AIChatService.js) needs to be
                 // duplicated because this hasn't been moved to be under
@@ -195,6 +195,7 @@ class OpenAICompletionService extends BaseService {
 
                 return await this.complete(messages, {
                     model: model,
+                    tools,
                     moderation: true,
                     stream,
                 });
@@ -242,7 +243,7 @@ class OpenAICompletionService extends BaseService {
     * @returns {Promise<Object>} The completion response containing message and usage info
     * @throws {Error} If messages are invalid or content is flagged by moderation
     */
-    async complete (messages, { stream, moderation, model }) {
+    async complete (messages, { stream, moderation, model, tools }) {
         // Validate messages
         if ( ! Array.isArray(messages) ) {
             throw new Error('`messages` must be an array');
@@ -360,6 +361,7 @@ class OpenAICompletionService extends BaseService {
             user: user_private_uid,
             messages: messages,
             model: model,
+            ...(tools ? { tools } : {}),
             // max_tokens,
             stream,
             ...(stream ? {
@@ -449,9 +451,9 @@ class OpenAICompletionService extends BaseService {
         }
 
         // We need to moderate the completion too
-        if ( moderation ) {
-            const text = completion.choices[0].message.content;
-            const moderation_result = await this.check_moderation(text);
+        const mod_text = completion.choices[0].message.content;
+        if ( moderation && mod_text !== null ) {
+            const moderation_result = await this.check_moderation(mod_text);
             if ( moderation_result.flagged ) {
                 throw new Error('message is not allowed');
             }

diff --git a/src/puter-js/src/modules/AI.js b/src/puter-js/src/modules/AI.js
@@ -276,6 +276,10 @@ class AI{
             options.stream = settings.stream;
         }
 
+        if ( settings.tools ) {
+            options.tools = settings.tools;
+        }
+
         // Call the original chat.complete method
         return await utils.make_driver_method(['messages'], 'puter-chat-completion', driver, 'complete', {
             test_mode: testMode ?? false,