
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.AllMiniLmL6V2EmbeddingModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.pgvector.PgVectorEmbeddingStore;
import org.testcontainers.containers.PostgreSQLContainer;
import org.testcontainers.utility.DockerImageName;

import java.util.List;

public class PgVectorEmbeddingStoreExample {

    public static void main(String[] args) {

        DockerImageName dockerImageName = DockerImageName.parse("pgvector/pgvector:pg16");
        try (PostgreSQLContainer<?> postgreSQLContainer = new PostgreSQLContainer<>(dockerImageName)) {
            postgreSQLContainer.start();

            EmbeddingStore<TextSegment> embeddingStore = PgVectorEmbeddingStore.builder()
                    .host(postgreSQLContainer.getHost())
                    .port(postgreSQLContainer.getFirstMappedPort())
                    .database(postgreSQLContainer.getDatabaseName())
                    .user(postgreSQLContainer.getUsername())
                    .password(postgreSQLContainer.getPassword())
                    .table("test")
                    .dimension(384)
                    .build();

            EmbeddingModel embeddingModel = new AllMiniLmL6V2EmbeddingModel();

            TextSegment segment1 = TextSegment.from("I like football.");
            Embedding embedding1 = embeddingModel.embed(segment1).content();
            embeddingStore.add(embedding1, segment1);

            TextSegment segment2 = TextSegment.from("The weather is good today.");
            Embedding embedding2 = embeddingModel.embed(segment2).content();
            embeddingStore.add(embedding2, segment2);

            Embedding queryEmbedding = embeddingModel.embed("What is your favourite sport?").content();
            List<EmbeddingMatch<TextSegment>> relevant = embeddingStore.findRelevant(queryEmbedding, 1);
            EmbeddingMatch<TextSegment> embeddingMatch = relevant.get(0);

            System.out.println(embeddingMatch.score()); // 0.8144288608390052
            System.out.println(embeddingMatch.embedded().text()); // I like football.

            postgreSQLContainer.stop();
        }
    }
}

[((924, 1376), 'dev.langchain4j.store.embedding.pgvector.PgVectorEmbeddingStore.builder'), ((924, 1347), 'dev.langchain4j.store.embedding.pgvector.PgVectorEmbeddingStore.builder'), ((924, 1311), 'dev.langchain4j.store.embedding.pgvector.PgVectorEmbeddingStore.builder'), ((924, 1276), 'dev.langchain4j.store.embedding.pgvector.PgVectorEmbeddingStore.builder'), ((924, 1211), 'dev.langchain4j.store.embedding.pgvector.PgVectorEmbeddingStore.builder'), ((924, 1150), 'dev.langchain4j.store.embedding.pgvector.PgVectorEmbeddingStore.builder'), ((924, 1081), 'dev.langchain4j.store.embedding.pgvector.PgVectorEmbeddingStore.builder'), ((924, 1013), 'dev.langchain4j.store.embedding.pgvector.PgVectorEmbeddingStore.builder')]