package com.honvay.flychat.langchain.llama.embedding;

import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.openai.OpenAiEmbeddingModel;
import dev.langchain4j.model.openai.OpenAiModelName;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;

import java.time.Duration;
import java.util.List;
import java.util.stream.Collectors;

@Service
public class OpenAiEmbeddingService implements EmbeddingService {

    private final String apiKey;

    public OpenAiEmbeddingService(@Value("${openai.apiKey:}") String apiKey) {
        this.apiKey = apiKey;
    }

    @Override
   public List<float[]> embed(List<String> texts){

       List<TextSegment> segments = texts.stream()
               .map(TextSegment::from)
               .collect(Collectors.toList());

       EmbeddingModel embeddingModel = OpenAiEmbeddingModel.builder()
               .apiKey(apiKey) // https://platform.openai.com/account/api-keys
               .modelName(OpenAiModelName.TEXT_EMBEDDING_ADA_002)
               .timeout(Duration.ofSeconds(15))
               .build();

        List<Embedding> embeddings = embeddingModel.embedAll(segments);
       return embeddings
               .stream()
               .map(Embedding::vector)
               .collect(Collectors.toList());
   }

    @Override
    public float[] embed(String text) {

        EmbeddingModel embeddingModel = OpenAiEmbeddingModel.builder()
                .apiKey(apiKey) // https://platform.openai.com/account/api-keys
                .modelName(OpenAiModelName.TEXT_EMBEDDING_ADA_002)
                .timeout(Duration.ofSeconds(15))
                .build();

        Embedding embedding = embeddingModel.embed(text);
        return embedding.vector();
    }

}

dev.langchain4j.model.openai.OpenAiEmbeddingModel.builder
[((981, 1228), 'dev.langchain4j.model.openai.OpenAiEmbeddingModel.builder'), ((981, 1204), 'dev.langchain4j.model.openai.OpenAiEmbeddingModel.builder'), ((981, 1156), 'dev.langchain4j.model.openai.OpenAiEmbeddingModel.builder'), ((981, 1042), 'dev.langchain4j.model.openai.OpenAiEmbeddingModel.builder'), ((1539, 1790), 'dev.langchain4j.model.openai.OpenAiEmbeddingModel.builder'), ((1539, 1765), 'dev.langchain4j.model.openai.OpenAiEmbeddingModel.builder'), ((1539, 1716), 'dev.langchain4j.model.openai.OpenAiEmbeddingModel.builder'), ((1539, 1601), 'dev.langchain4j.model.openai.OpenAiEmbeddingModel.builder')]