int-layer/src/main/java/com/example/service/base/NotificationContextResolver.java

/*
* Copyright 2025 Google. This software is provided as-is, without warranty or representation for any use or purpose.
* Your use of it is subject to your agreement with Google.
*/
package com.example.service.base;

import com.example.service.notification.MemoryStoreNotificationService;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.ClassPathResource;
import org.springframework.stereotype.Service;
import org.springframework.util.FileCopyUtils;
import jakarta.annotation.PostConstruct;
import java.io.IOException;
import java.io.Reader;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.util.Objects;

@Service
public class NotificationContextResolver {

    private static final Logger logger = LoggerFactory.getLogger(NotificationContextResolver.class);
    private final GeminiClientService geminiService;

    @Value("${notificationcontext.geminimodel:gemini-2.0-flash-001}")
    private String geminiModelNameResolver;

    @Value("${notificationcontext.temperature:0.1f}")
    private Float resolverTemperature;

    @Value("${notificationcontext.maxOutputTokens:1024}")
    private Integer resolverMaxOutputTokens;

    @Value("${notificationcontext.topP:0.1f}")
    private Float resolverTopP;

    @Value("${notificationcontext.prompt:prompts/notification_context_resolver.txt}")
    private String promptFilePath;

    public static final String CATEGORY_DIALOGFLOW = "DIALOGFLOW";

    private String promptTemplate;

    public NotificationContextResolver(GeminiClientService geminiService, MemoryStoreNotificationService memoryStoreNotificationService) {
        this.geminiService = Objects.requireNonNull(geminiService,
                "GeminiClientService cannot be null for NotificationContextResolver.");
    }

    @PostConstruct
    public void loadPromptTemplate() {
        try {
            ClassPathResource resource = new ClassPathResource(promptFilePath);
            try (Reader reader = new InputStreamReader(resource.getInputStream(), StandardCharsets.UTF_8)) {
                this.promptTemplate = FileCopyUtils.copyToString(reader);
            }
            logger.info("Successfully loaded prompt template from '{}'.", promptFilePath);
        } catch (IOException e) {
            logger.error("Failed to load prompt template from '{}'. Please ensure the file exists.", promptFilePath, e);
            throw new IllegalStateException("Could not load prompt template.", e);
        }
    }

    public String resolveContext(String queryInputText, String notificationsJson, String conversationJson,
            String metadata, String userId, String sessionId, String userPhoneNumber) {
        logger.debug("resolveContext -> queryInputText: {}, notificationsJson: {}, conversationJson: {}, metadata: {}", queryInputText, notificationsJson, conversationJson, metadata);
        if (queryInputText == null || queryInputText.isBlank()) {
            logger.warn("Query input text for context resolution is null or blank. Returning {}.", CATEGORY_DIALOGFLOW);
            return CATEGORY_DIALOGFLOW;
        }

        String notificationContent = (notificationsJson != null && !notificationsJson.isBlank()) ? notificationsJson
                : "No metadata in notification.";

        String conversationHistory = (conversationJson != null && !conversationJson.isBlank()) ? conversationJson
                : "No conversation history.";

        String contextPrompt = String.format(
                this.promptTemplate,
                conversationHistory,
                notificationContent,
                metadata,
                queryInputText);

        logger.debug("Sending context resolution request to Gemini for input (first 100 chars): '{}'...",
                queryInputText.substring(0, Math.min(queryInputText.length(), 100)));

        try {
            String geminiResponse = geminiService.generateContent(
                    contextPrompt,
                    resolverTemperature,
                    resolverMaxOutputTokens,
                    geminiModelNameResolver,
                    resolverTopP);

            if (geminiResponse != null && !geminiResponse.isBlank()) {
                if (geminiResponse.trim().equalsIgnoreCase(CATEGORY_DIALOGFLOW)) {
                    logger.info("Resolved to {}. Input: '{}'", CATEGORY_DIALOGFLOW, queryInputText);
                    return CATEGORY_DIALOGFLOW;
                } else {
                    logger.info("Resolved to a specific response. Input: '{}'", queryInputText);
                    return geminiResponse;
                }
            } else {
                logger.warn("Gemini returned a null or blank response. Input: '{}'. Returning {}.",
                        queryInputText, CATEGORY_DIALOGFLOW);
                return CATEGORY_DIALOGFLOW;
            }
        } catch (Exception e) {
            logger.error("An error occurred during Gemini content generation for context resolution.", e);
            return CATEGORY_DIALOGFLOW;
        }
    }
}