Newer
Older
gcprag / src / main / java / net / curtlewis / gcprag / chunk / ChunkService.java
package net.curtlewis.gcprag.chunk;

import java.util.ArrayList;
import java.util.List;
import java.util.UUID;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.EmbeddingRequest;
import org.springframework.ai.embedding.EmbeddingResponse;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.stereotype.Service;

@Service
public class ChunkService {

    private final ChunkRepository chunkRepository;
    private final EmbeddingModel embeddingModel;

    public ChunkService(ChunkRepository chunkRepository, @Qualifier("ollamaEmbeddingModel") EmbeddingModel embeddingModel) {
        this.chunkRepository = chunkRepository;
        this.embeddingModel = embeddingModel;
    }

    /**
     * Get all documents
     */
    public List<ChunkEntity> getAllDocuments() {
        return chunkRepository.findAll();
    }

    /**
     * Regenerate embeddings for all documents (useful for migration or model
     * changes)
     */
    public void regenerateAllEmbeddings() {
        List<ChunkEntity> documents = chunkRepository.findAll();

        for (ChunkEntity document : documents) {
            List<Double> embedding = generateEmbedding(document.getContent());

            float[] embeddingFloats = new float[embedding.size()];
            int index = 0;
            for (Double d : embedding) {
                embeddingFloats[index++] = d.floatValue();
            }

            float[] embeddingArray = new float[embeddingFloats.length];
            for (int i = 0; i < embeddingFloats.length; i++) {
                embeddingArray[i] = embeddingFloats[i];
            }

            document.setEmbedding(embeddingArray);
            // document.setDateLastModified(LocalDateTime.now());

        }

        chunkRepository.saveAll(documents);
    }

    /**
     * Get document by ID
     */
    public ChunkEntity getDocument(UUID id) {
        return chunkRepository.findById(id)
                .orElseThrow(() -> new RuntimeException("Document not found"));
    }

    /**
     * Generate embedding for text
     */
    private List<Double> generateEmbedding(String text) {
        EmbeddingRequest request = new EmbeddingRequest(List.of(text), null);
        EmbeddingResponse response = embeddingModel.call(request);
        float[] floatArray = response.getResults().get(0).getOutput();
        List<Double> doublesList = new ArrayList<>();
        for (float f : floatArray) {
            doublesList.add((double) f);
        }

        return doublesList;
    }

}