@smythos/sre 1.5.37 → 1.5.40

package/src/subsystems/LLMManager/LLM.service/connectors/VertexAI.class.ts

@@ -3,16 +3,23 @@ import EventEmitter from 'events';
 
 import { JSON_RESPONSE_INSTRUCTION, BUILT_IN_MODEL_PREFIX } from '@sre/constants';
 import {
-    TLLMParams,
     TCustomLLMModel,
     APIKeySource,
     TVertexAISettings,
     ILLMRequestFuncParams,
     TGoogleAIRequestBody,
-    TLLMConnectorParams,
     ILLMRequestContext,
+    TLLMPreparedParams,
+    TLLMMessageBlock,
+    ToolData,
+    TLLMToolResultMessageBlock,
+    TLLMMessageRole,
+    TLLMChatResponse,
+    TLLMEvent,
 } from '@sre/types/LLM.types';
 import { LLMHelper } from '@sre/LLMManager/LLM.helper';
+import { BinaryInput } from '@sre/helpers/BinaryInput.helper';
+import { AccessCandidate } from '@sre/Security/AccessControl/AccessCandidate.class';
 
 import { LLMConnector } from '../LLMConnector';
 import { SystemEvents } from '@sre/Core/SystemEvents';
@@ -32,38 +39,66 @@ export class VertexAIConnector extends LLMConnector {
         return new VertexAI({
             project: projectId,
             location: region,
+            apiEndpoint: (modelInfo?.settings as TVertexAISettings)?.apiEndpoint,
             googleAuthOptions: {
                 credentials: credentials as any,
             },
         });
     }
 
-    protected async request({ acRequest, body, context }: ILLMRequestFuncParams): Promise<any> {
-        const messages = body.messages;
-        delete body.messages;
-
+    protected async request({ acRequest, body, context }: ILLMRequestFuncParams): Promise<TLLMChatResponse> {
         try {
             const vertexAI = await this.getClient(context);
-            const generativeModel = vertexAI.getGenerativeModel(body);
 
-            const result = await generativeModel.generateContent({ contents: messages });
-            const content = result?.response?.candidates?.[0]?.content?.parts?.[0]?.text;
-            const usage = result?.response?.usageMetadata;
+            // Separate contents from model configuration
+            const contents = body.contents;
+            delete body.contents;
 
-            this.reportUsage(usage, {
-                modelEntryName: context.modelEntryName,
-                keySource: context.isUserKey ? APIKeySource.User : APIKeySource.Smyth,
-                agentId: context.agentId,
-                teamId: context.teamId,
-            });
+            // VertexAI expects contents in a specific format: {contents: [...]}
+            const requestParam = { contents };
+
+            const model = vertexAI.getGenerativeModel(body);
+
+            const result = await model.generateContent(requestParam);
+            const response = await result.response;
+
+            const content = response.candidates?.[0]?.content?.parts?.[0]?.text || '';
+            const finishReason = response.candidates?.[0]?.finishReason || 'stop';
+            const usage = response.usageMetadata;
+
+            let toolsData: ToolData[] = [];
+            let useTool = false;
+
+            // Check for function calls in the response
+            const functionCalls = response.candidates?.[0]?.content?.parts?.filter((part) => part.functionCall);
+            if (functionCalls && functionCalls.length > 0) {
+                functionCalls.forEach((call, index) => {
+                    toolsData.push({
+                        index,
+                        id: call.functionCall?.name + '_' + index, // VertexAI doesn't provide IDs like Anthropic
+                        type: 'function',
+                        name: call.functionCall?.name,
+                        arguments: call.functionCall?.args,
+                        role: TLLMMessageRole.Assistant,
+                    });
+                });
+                useTool = true;
+            }
+
+            if (usage) {
+                this.reportUsage(usage, {
+                    modelEntryName: context.modelEntryName,
+                    keySource: context.isUserKey ? APIKeySource.User : APIKeySource.Smyth,
+                    agentId: context.agentId,
+                    teamId: context.teamId,
+                });
+            }
 
             return {
                 content,
-                finishReason: 'stop',
-                useTool: false,
-                toolsData: [],
-                message: { content, role: 'assistant' },
-                usage,
+                finishReason,
+                toolsData,
+                useTool,
             };
         } catch (error) {
             throw error;
@@ -71,85 +106,125 @@ export class VertexAIConnector extends LLMConnector {
     }
 
     protected async streamRequest({ acRequest, body, context }: ILLMRequestFuncParams): Promise<EventEmitter> {
-        // Simulate streaming similar to Perplexity's approach - fallback to regular request
         const emitter = new EventEmitter();
 
-        setTimeout(() => {
+        setTimeout(async () => {
             try {
-                this.request({ acRequest, body, context })
-                    .then((response) => {
-                        const finishReason = response.finishReason;
-                        const usage = response.usage;
-
-                        // Emit the content as a stream-like response
-                        emitter.emit('interrupted', finishReason);
-                        emitter.emit('content', response.content);
-                        emitter.emit('end', undefined, usage, finishReason);
-                    })
-                    .catch((error) => {
-                        emitter.emit('error', error.message || error.toString());
+                const vertexAI = await this.getClient(context);
+
+                // Separate contents from model configuration
+                const contents = body.contents;
+                delete body.contents;
+
+                const vertexModel = vertexAI.getGenerativeModel(body);
+
+                // VertexAI expects contents in a specific format: {contents: [...]}
+                const requestParam = { contents };
+
+                const streamResult = await vertexModel.generateContentStream(requestParam);
+
+                let toolsData: ToolData[] = [];
+                let usageData: any[] = [];
+
+                for await (const chunk of streamResult.stream) {
+                    const chunkText = chunk.candidates?.[0]?.content?.parts?.[0]?.text || '';
+                    if (chunkText) {
+                        emitter.emit('content', chunkText);
+                    }
+                }
+
+                const aggregatedResponse = await streamResult.response;
+
+                // Check for function calls in the final response (like Anthropic does)
+                const functionCalls = aggregatedResponse.candidates?.[0]?.content?.parts?.filter((part) => part.functionCall);
+                if (functionCalls && functionCalls.length > 0) {
+                    functionCalls.forEach((call, index) => {
+                        toolsData.push({
+                            index,
+                            id: call.functionCall?.name + '_' + index,
+                            type: 'function',
+                            name: call.functionCall?.name,
+                            arguments: call.functionCall?.args,
+                            role: TLLMMessageRole.Assistant,
+                        });
+                    });
+
+                    emitter.emit(TLLMEvent.ToolInfo, toolsData);
+                }
+
+                const usage = aggregatedResponse.usageMetadata;
+
+                if (usage) {
+                    const reportedUsage = this.reportUsage(usage, {
+                        modelEntryName: context.modelEntryName,
+                        keySource: context.isUserKey ? APIKeySource.User : APIKeySource.Smyth,
+                        agentId: context.agentId,
+                        teamId: context.teamId,
                     });
+                    usageData.push(reportedUsage);
+                }
+
+                const finishReason = (aggregatedResponse.candidates?.[0]?.finishReason || 'stop').toLowerCase();
+
+                if (finishReason !== 'stop') {
+                    emitter.emit('interrupted', finishReason);
+                }
+
+                setTimeout(() => {
+                    emitter.emit('end', toolsData, usageData, finishReason);
+                }, 100);
             } catch (error) {
-                emitter.emit('error', error.message || error.toString());
+                emitter.emit('error', error);
             }
         }, 100);
 
         return emitter;
     }
 
-    protected async webSearchRequest({ acRequest, body, context }: ILLMRequestFuncParams): Promise<EventEmitter> {
-        throw new Error('Web search is not supported for Vertex AI');
-    }
-
-    protected async reqBodyAdapter(params: TLLMParams): Promise<TGoogleAIRequestBody> {
-        let messages = params?.messages || [];
-
-        //#region Separate system message and add JSON response instruction if needed
-        let systemInstruction;
-        const { systemMessage, otherMessages } = LLMHelper.separateSystemMessages(messages);
+    protected async reqBodyAdapter(params: TLLMPreparedParams): Promise<TGoogleAIRequestBody> {
+        const model = params?.model;
+        const { messages, systemMessage } = await this.prepareMessages(params);
 
-        if ('content' in systemMessage) {
-            systemInstruction = systemMessage.content;
-        }
-
-        messages = otherMessages;
+        let body: any = {
+            model: model as string,
+            contents: messages, // This will be separated in the request methods
+        };
 
         const responseFormat = params?.responseFormat || '';
+        let systemInstruction = systemMessage || '';
+
         if (responseFormat === 'json') {
-            systemInstruction = JSON_RESPONSE_INSTRUCTION;
+            systemInstruction += (systemInstruction ? '\n\n' : '') + JSON_RESPONSE_INSTRUCTION;
         }
-        //#endregion Separate system message and add JSON response instruction if needed
-
-        const modelInfo = params.modelInfo as TCustomLLMModel;
-
-        let body: TGoogleAIRequestBody = {
-            model: modelInfo?.settings?.customModel || modelInfo?.settings?.foundationModel,
-            messages,
-        };
 
         const config: GenerationConfig = {};
 
-        if (params?.maxTokens !== undefined) config.maxOutputTokens = params.maxTokens;
-        if (params?.temperature !== undefined) config.temperature = params.temperature;
-        if (params?.topP !== undefined) config.topP = params.topP;
-        if (params?.topK !== undefined) config.topK = params.topK;
-        if (params?.stopSequences?.length) config.stopSequences = params.stopSequences;
+        if (params.maxTokens !== undefined) config.maxOutputTokens = params.maxTokens;
+        if (params.temperature !== undefined) config.temperature = params.temperature;
+        if (params.topP !== undefined) config.topP = params.topP;
+        if (params.topK !== undefined) config.topK = params.topK;
+        if (params.stopSequences?.length) config.stopSequences = params.stopSequences;
 
         if (systemInstruction) {
-            body.systemInstruction = systemInstruction;
+            body.systemInstruction = {
+                role: 'system',
+                parts: [{ text: systemInstruction }],
+            };
         }
 
         if (Object.keys(config).length > 0) {
-            body.generationConfig = config as any;
+            body.generationConfig = config;
+        }
+
+        // Handle tools configuration
+        if (params?.toolsConfig?.tools && params?.toolsConfig?.tools.length > 0) {
+            body.tools = this.formatToolsForVertexAI(params.toolsConfig.tools);
         }
 
         return body;
     }
 
-    protected reportUsage(
-        usage: UsageMetadata & { cachedContentTokenCount?: number },
-        metadata: { modelEntryName: string; keySource: APIKeySource; agentId: string; teamId: string }
-    ) {
+    protected reportUsage(usage: UsageMetadata, metadata: { modelEntryName: string; keySource: APIKeySource; agentId: string; teamId: string }) {
         // SmythOS (built-in) models have a prefix, so we need to remove it to get the model name
         const modelName = metadata.modelEntryName.replace(BUILT_IN_MODEL_PREFIX, '');
 
@@ -168,26 +243,181 @@ export class VertexAIConnector extends LLMConnector {
         return usageData;
     }
 
-    public formatToolsConfig({ type = 'function', toolDefinitions, toolChoice = 'auto' }) {
-        throw new Error('Tool configuration is not currently implemented for Vertex AI');
-    }
+    private async prepareMessages(params: TLLMPreparedParams) {
+        const messages = params?.messages || [];
+        const files: BinaryInput[] = params?.files || [];
 
-    public getConsistentMessages(messages) {
-        const _messages = LLMHelper.removeDuplicateUserMessages(messages);
+        let processedMessages = [...messages];
 
-        return _messages.map((message) => {
-            let textBlock = [];
+        // Handle system messages - VertexAI uses systemInstruction separately
+        const { systemMessage, otherMessages } = LLMHelper.separateSystemMessages(processedMessages);
+        processedMessages = otherMessages;
 
-            if (message?.parts) {
-                textBlock = message.parts;
-            } else if (message?.content) {
-                textBlock = Array.isArray(message.content) ? message.content : [{ text: message.content as string }];
+        // Handle files if present
+        if (files?.length > 0) {
+            const fileData = await this.processFiles(files, params.agentId);
+
+            // Add file data to the last user message
+            const userMessage = processedMessages.pop();
+            if (userMessage) {
+                const content = [{ text: userMessage.content as string }, ...fileData];
+                processedMessages.push({
+                    role: userMessage.role,
+                    parts: content,
+                });
             }
+        }
+
+        // Convert messages to VertexAI format
+        let vertexAIMessages = this.convertMessagesToVertexAIFormat(processedMessages);
+
+        // Ensure we have at least one message with content
+        if (!vertexAIMessages || vertexAIMessages.length === 0) {
+            vertexAIMessages = [
+                {
+                    role: 'user',
+                    parts: [{ text: 'Hello' }],
+                },
+            ];
+        }
+
+        return {
+            messages: vertexAIMessages,
+            systemMessage: (systemMessage as any)?.content || '',
+        };
+    }
+
+    private async processFiles(files: BinaryInput[], agentId: string) {
+        const fileData = [];
+
+        for (const file of files) {
+            const bufferData = await file.readData(AccessCandidate.agent(agentId));
+            const base64Data = bufferData.toString('base64');
+
+            fileData.push({
+                inlineData: {
+                    data: base64Data,
+                    mimeType: file.mimetype,
+                },
+            });
+        }
+
+        return fileData;
+    }
+
+    private convertMessagesToVertexAIFormat(messages: TLLMMessageBlock[]) {
+        return messages
+            .filter((message) => message && (message.content || message.parts))
+            .map((message) => {
+                let parts;
+
+                if (typeof message.content === 'string') {
+                    parts = message.content.trim() ? [{ text: message.content.trim() }] : [{ text: 'Continue' }];
+                } else if (message.parts && Array.isArray(message.parts)) {
+                    parts = message.parts;
+                } else if (message.content) {
+                    parts = [{ text: String(message.content) || 'Continue' }];
+                } else {
+                    parts = [{ text: 'Continue' }];
+                }
+
+                return {
+                    role: message.role === TLLMMessageRole.Assistant ? 'model' : 'user',
+                    parts,
+                };
+            });
+    }
+
+    private formatToolsForVertexAI(tools: any[]) {
+        return [
+            {
+                functionDeclarations: tools.map((tool) => ({
+                    name: tool.name,
+                    description: tool.description || '',
+                    parameters: {
+                        type: 'object',
+                        properties: tool.properties || {},
+                        required: tool.requiredFields || [],
+                    },
+                })),
+            },
+        ];
+    }
+
+    public formatToolsConfig({ toolDefinitions, toolChoice = 'auto' }) {
+        const tools = toolDefinitions.map((tool) => {
+            const { name, description, properties, requiredFields } = tool;
 
             return {
-                role: message.role,
-                parts: textBlock,
+                name,
+                description,
+                properties,
+                requiredFields,
             };
         });
+
+        return {
+            tools,
+            toolChoice: {
+                type: toolChoice,
+            },
+        };
+    }
+
+    public transformToolMessageBlocks({
+        messageBlock,
+        toolsData,
+    }: {
+        messageBlock: TLLMMessageBlock;
+        toolsData: ToolData[];
+    }): TLLMToolResultMessageBlock[] {
+        const messageBlocks: TLLMToolResultMessageBlock[] = [];
+
+        if (messageBlock) {
+            const parts = [];
+
+            if (typeof messageBlock.content === 'string') {
+                parts.push({ text: messageBlock.content });
+            } else if (Array.isArray(messageBlock.content)) {
+                parts.push(...messageBlock.content);
+            }
+
+            if (messageBlock.tool_calls) {
+                const functionCalls = messageBlock.tool_calls.map((toolCall: any) => ({
+                    functionCall: {
+                        name: toolCall?.function?.name,
+                        args:
+                            typeof toolCall?.function?.arguments === 'string'
+                                ? JSON.parse(toolCall.function.arguments)
+                                : toolCall?.function?.arguments || {},
+                    },
+                }));
+                parts.push(...functionCalls);
+            }
+
+            messageBlocks.push({
+                role: messageBlock.role,
+                parts,
+            });
+        }
+
+        // Transform tool results
+        const toolResults = toolsData.map((toolData) => ({
+            role: TLLMMessageRole.User,
+            parts: [
+                {
+                    functionResponse: {
+                        name: toolData.name,
+                        response: {
+                            name: toolData.name,
+                            content: toolData.result,
+                        },
+                    },
+                },
+            ],
+        }));
+
+        messageBlocks.push(...toolResults);
+        return messageBlocks;
     }
 }