An example Apex action using the Salesforce models API to send prompts to a custom model.

custom-llm-action-example.cls

public class CustomLLMAction {
    @InvocableMethod(
        label='Your label here'
        description='Using a custom LLM to ...'
        category='AI/ML'
    )
    public static List<GenerateResponseResult> generateResponse(List<GenerateResponseRequest> requests) {
        List<GenerateResponseResult> results = new List<GenerateResponseResult>();
        
        for(GenerateResponseRequest request : requests) {
            GenerateResponseResult result = new GenerateResponseResult();
            
            try {
                // Create generate chat request
                aiplatform.ModelsAPI.createChatGenerations_Request apiRequest = 
                    new aiplatform.ModelsAPI.createChatGenerations_Request();
                
                // Specify model
                apiRequest.modelName = 'your_model_api_name_here';
                
                // Create request body
                aiplatform.ModelsAPI_ChatGenerationsRequest body = 
                    new aiplatform.ModelsAPI_ChatGenerationsRequest();
                apiRequest.body = body;
                
                // Add chat messages to body
                List<aiplatform.ModelsAPI_ChatMessageRequest> messagesList = 
                    new List<aiplatform.ModelsAPI_ChatMessageRequest>();
                
                // Add user message
                aiplatform.ModelsAPI_ChatMessageRequest modelPrompt = 
                    new aiplatform.ModelsAPI_ChatMessageRequest();
                modelPrompt.content = request.modelPrompt;
                modelPrompt.role = 'user';
                messagesList.add(modelPrompt);
                
                body.messages = messagesList;
                
                // Make request
                aiplatform.ModelsAPI modelsAPI = new aiplatform.ModelsAPI();
                aiplatform.ModelsAPI.createChatGenerations_Response response = 
                    modelsAPI.createChatGenerations(apiRequest);
                
                // Extract generated response from response
                if (response.Code200?.generationDetails?.generations != null && 
                    !response.Code200.generationDetails.generations.isEmpty()) {
                    result.generatedResponse = response.Code200.generationDetails.generations[0].content;
                    result.isSuccess = true;
                } else {
                    result.isSuccess = false;
                    result.errorMessage = 'No generated response in response';
                }
                
            } catch(aiplatform.ModelsAPI.createChatGenerations_ResponseException e) {
                result.isSuccess = false;
                result.errorMessage = 'Error: ' + e.getMessage() + 
                    ' (Response Code: ' + e.responseCode + ')';
            } catch(Exception e) {
                result.isSuccess = false;
                result.errorMessage = 'Unexpected error: ' + e.getMessage();
            }
            
            results.add(result);
        }
        
        return results;
    }
    
    public class GenerateResponseRequest {    
        @InvocableVariable(
            label='Model prompt'
            description='The prompt to send to the code instruction model'
            required=true
        )
            public String modelPrompt;
    }
    
    public class GenerateResponseResult {
        @InvocableVariable(
            label='Generated Response'
            description='The response generated by the code instruction model'
        )
        public String generatedResponse;
        
        @InvocableVariable(
            label='Success'
            description='Indicates if the response generation was successful'
        )
        public Boolean isSuccess;
        
        @InvocableVariable(
            label='Error Message'
            description='Error message if the generation failed'
        )
        public String errorMessage;
    }
}