sf-service-models-api

public with sharing class CaseSummarizer {

    private static final String MODEL_NAME = 'sfdc_ai__DefaultGPT4Omni';

    @TestVisible
    private static String mockResponse;

    @AuraEnabled
    public static String summarize(String caseDescription) {
        if (String.isBlank(caseDescription)) return null;

        String prompt = buildPrompt(caseDescription);
        String raw = callModel(prompt);
        if (String.isBlank(raw)) return null;

        return raw.trim();
    }

    private static String buildPrompt(String caseDescription) {
        return 'You are a Service Cloud agent assistant.\n\n'
            + 'CASE DESCRIPTION:\n' + caseDescription + '\n\n'
            + 'INSTRUCTIONS:\n'
            + '- Summarize in 2 sentences max.\n'
            + '- Plain text only. No markdown, no preamble.';
    }

    private static String callModel(String prompt) {
        if (Test.isRunningTest() && mockResponse != null) {
            return mockResponse;
        }

        aiplatform.ModelsAPI modelsAPI = new aiplatform.ModelsAPI();

        aiplatform.ModelsAPI_ChatMessageRequest msg = new aiplatform.ModelsAPI_ChatMessageRequest();
        msg.role = 'user';
        msg.content = prompt;

        aiplatform.ModelsAPI_ChatGenerationsRequest body = new aiplatform.ModelsAPI_ChatGenerationsRequest();
        body.messages = new List<aiplatform.ModelsAPI_ChatMessageRequest>{ msg };

        aiplatform.ModelsAPI.createChatGenerations_Request req =
            new aiplatform.ModelsAPI.createChatGenerations_Request();
        req.modelName = MODEL_NAME;
        req.body = body;

        aiplatform.ModelsAPI.createChatGenerations_Response resp =
            modelsAPI.createChatGenerations(req);

        if (resp.Code200 == null
            || resp.Code200.generationDetails == null
            || resp.Code200.generationDetails.generations == null
            || resp.Code200.generationDetails.generations.isEmpty()) {
            return null;
        }

        return resp.Code200.generationDetails.generations[0].content;
    }
}

private static String buildExtractionPrompt(String text) {
    return 'You are extracting fields from a customer message.\n\n'
        + 'TEXT:\n' + text + '\n\n'
        + 'INSTRUCTIONS:\n'
        + '- Respond ONLY with a valid JSON object with these keys:\n'
        + '  customer_name (string), severity (one of: Low, Medium, High), wants_callback (true/false).\n'
        + '- Omit a key when no clear answer exists. Do not invent values.\n'
        + '- No markdown, no code fences, no explanation.\n'
        + 'Example: {"customer_name":"Jane Doe","severity":"High"}';
}

private static Map<String, Object> parseJsonResponse(String raw) {
    if (String.isBlank(raw)) return new Map<String, Object>();
    String cleaned = raw.trim();

    if (cleaned.startsWith('```')) {
        cleaned = cleaned.replaceAll('```[a-zA-Z]*\\n?', '')
                         .replaceAll('```', '')
                         .trim();
    }

    try {
        return (Map<String, Object>) JSON.deserializeUntyped(cleaned);
    } catch (Exception e) {
        return new Map<String, Object>();
    }
}

private static final Set<String> VALID_SEVERITIES = new Set<String>{ 'Low', 'Medium', 'High' };

String severity = (String) parsed.get('severity');
if (severity != null && !VALID_SEVERITIES.contains(severity)) {
    severity = null; // drop invalid
}

@IsTest
static void summarize_returnsModelOutput() {
    CaseSummarizer.mockResponse = 'Customer reports a flat tire on I-95. Roadside dispatched.';

    Test.startTest();
    String result = CaseSummarizer.summarize('Caller says they have a flat tire on I-95...');
    Test.stopTest();

    System.assertEquals(
        'Customer reports a flat tire on I-95. Roadside dispatched.',
        result
    );
}