meals-made-easy-api/src/main/java/app/mealsmadeeasy/api/recipe/job/RecipeInferJobHandler.java

package app.mealsmadeeasy.api.recipe.job;

import app.mealsmadeeasy.api.ai.InferenceService;
import app.mealsmadeeasy.api.ai.OcrService;
import app.mealsmadeeasy.api.file.File;
import app.mealsmadeeasy.api.file.FileService;
import app.mealsmadeeasy.api.job.Job;
import app.mealsmadeeasy.api.job.JobHandler;
import app.mealsmadeeasy.api.recipe.RecipeDraft;
import app.mealsmadeeasy.api.recipe.RecipeService;
import com.fasterxml.jackson.annotation.JsonProperty;
import lombok.RequiredArgsConstructor;
import org.jetbrains.annotations.Nullable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.content.Media;
import org.springframework.ai.converter.BeanOutputConverter;
import org.springframework.core.io.InputStreamResource;
import org.springframework.stereotype.Component;
import org.springframework.util.MimeType;

import java.io.IOException;
import java.io.InputStream;
import java.time.OffsetDateTime;
import java.util.ArrayList;
import java.util.List;

@Component
@RequiredArgsConstructor
public class RecipeInferJobHandler implements JobHandler<RecipeInferJobPayload> {

    private record RecipeExtraction(
            @JsonProperty(required = true) String title,
            @JsonProperty(required = true) List<IngredientExtraction> ingredients
    ) {}

    private record IngredientExtraction(
            String amount,
            @JsonProperty(required = true) String name,
            String notes
    ) {}

    public static final String JOB_KEY = "RECIPE_INFER_JOB";

    private static final Logger logger = LoggerFactory.getLogger(RecipeInferJobHandler.class);

    private final FileService fileService;
    private final RecipeService recipeService;
    private final OcrService ocrService;
    private final InferenceService inferenceService;
    private final BeanOutputConverter<RecipeExtraction> extractionConverter =
            new BeanOutputConverter<>(RecipeExtraction.class);

    @Override
    public Class<RecipeInferJobPayload> getPayloadType() {
        return RecipeInferJobPayload.class;
    }

    @Override
    public String getJobKey() {
        return "RECIPE_INFER_JOB";
    }

    @Override
    public void handle(Job job, RecipeInferJobPayload payload) {
        logger.debug("Starting recipe inference job {}", job.getId());
        final File sourceFile = this.fileService.getById(payload.fileId());
        final InputStream sourceFileContent;
        try {
            sourceFileContent = this.fileService.getFileContentById(payload.fileId());
        } catch (IOException e) {
            throw new RuntimeException(e);
        }

        // 1. OCR of recipe image
        final Media sourceFileMedia = Media.builder()
                .data(new InputStreamResource(sourceFileContent))
                .mimeType(MimeType.valueOf(sourceFile.getMimeType()))
                .build();
        final @Nullable String fullMarkdownText = this.ocrService.readMarkdown(sourceFileMedia);
        if (fullMarkdownText == null) {
            throw new RuntimeException("fullMarkdownText from ocr came back null");
        }

        logger.debug("ocr returned:\n{}", fullMarkdownText);

        // 2. Extract title and ingredients
        final Message extractSystemMessage = SystemMessage.builder()
                .text("""
                    Extract from the given Markdown food recipe the title and ingredients.
                    For each ingredient, you must provide the name; amount and notes are optional.
                    
                    Here is an example of a full extraction:
                    
                    ```json
                    {
                        "title": "Delicious Recipe",
                        "ingredients": [
                            {
                                "amount": "1 tablespoon",
                                "name": "garlic powder"
                            },
                            {
                                "amount": "1 lb",
                                "name": "tofu",
                                "notes": "pressed and dried"
                            },
                            {
                                "name": "salt and pepper"
                            }
                        ]
                    }
                    ```
                    """.stripIndent().trim()
                )
                .build();
        final Message extractUserMessage = UserMessage.builder()
                .text(fullMarkdownText)
                .build();

        final @Nullable RecipeExtraction extraction = this.inferenceService.extract(
                extractSystemMessage,
                extractUserMessage,
                this.extractionConverter
        );
        if (extraction == null) {
            throw new RuntimeException("extraction returned null");
        }

        logger.debug("extract returned:\n{}", extraction);

        // 3. Get draft and set props from extraction
        final @Nullable RecipeDraft recipeDraft = this.recipeService.getDraftById(payload.recipeDraftId());
        if (recipeDraft == null) {
            throw new RuntimeException("Recipe draft not found");
        }

        recipeDraft.setTitle(extraction.title());

        // todo: calculate slug

        // ingredients
        if (recipeDraft.getIngredients() == null) {
            recipeDraft.setIngredients(new ArrayList<>());
        }
        recipeDraft.getIngredients().addAll(extraction.ingredients().stream()
                .map(extractedIngredient -> {
                    final RecipeDraft.IngredientDraft ingredientDraft = new RecipeDraft.IngredientDraft();
                    ingredientDraft.setAmount(extractedIngredient.amount());
                    ingredientDraft.setName(extractedIngredient.name());
                    ingredientDraft.setNotes(extractedIngredient.notes());
                    return ingredientDraft;
                })
                .toList()
        );

        // other props
        recipeDraft.setRawText(fullMarkdownText);
        recipeDraft.setState(RecipeDraft.State.ENTER_DATA);

        // save inference for later
        if (recipeDraft.getInferences() == null) {
            recipeDraft.setInferences(new ArrayList<>());
        }
        final RecipeDraft.RecipeDraftInference inference = new RecipeDraft.RecipeDraftInference();
        inference.setTitle(extraction.title());
        inference.setRawText(fullMarkdownText);
        inference.setInferredAt(OffsetDateTime.now());
        recipeDraft.getInferences().add(inference);

        logger.debug("Recipe infer job completed, saving to db...");

        this.recipeService.saveDraft(recipeDraft);
    }

}