DIP/src/main/java/at/procon/dip/embedding/service/DefaultQueryEmbeddingService.java

package at.procon.dip.embedding.service;

import at.procon.dip.embedding.model.EmbeddingProviderResult;
import at.procon.dip.embedding.model.EmbeddingUseCase;
import at.procon.dip.embedding.registry.EmbeddingModelRegistry;
import java.util.List;
import lombok.RequiredArgsConstructor;
import org.springframework.stereotype.Service;

@Service
@RequiredArgsConstructor
public class DefaultQueryEmbeddingService implements QueryEmbeddingService {

    private final EmbeddingExecutionService executionService;
    private final EmbeddingModelRegistry modelRegistry;

    @Override
    public float[] embedQuery(String queryText) {
        return embedQuery(queryText, modelRegistry.getRequiredDefaultQueryModelKey());
    }

    @Override
    public float[] embedQuery(String queryText, String modelKey) {
        EmbeddingProviderResult result = executionService.embedTexts(
                modelKey,
                EmbeddingUseCase.QUERY,
                List.of(queryText)
        );

        if (result.vectors() == null || result.vectors().isEmpty()) {
            throw new IllegalStateException("Embedding provider returned no query vector for model " + modelKey);
        }

        return result.vectors().getFirst();
    }
}