package org.agoncal.fascicle.langchain4j.vectordb.pgvector;

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 java.util.List;

// tag::adocSkip[]

/**
 * @author Antonio Goncalves
 * http://www.antoniogoncalves.org
 * --
 */
// end::adocSkip[]
public class MusicianService {

  public static void main(String[] args) {
    MusicianService musicianService = new MusicianService();

    musicianService.usePGVectorToStoreEmbeddings();
  }

  public void usePGVectorToStoreEmbeddings() {
    System.out.println("### usePGVectorToStoreEmbeddings");

    // tag::adocSnippet[]
    EmbeddingStore<TextSegment> embeddingStore =
      PgVectorEmbeddingStore.builder()
        .host("localhost")
        .port(5432)
        .createTable(true)
        .dropTableFirst(true)
        .dimension(384)
        .table("langchain4j_collection")
        .user("agoncal")
        .password("agoncal")
        .database("agoncal")
        .build();
    // end::adocSnippet[]

    EmbeddingModel embeddingModel = new AllMiniLmL6V2EmbeddingModel();

    TextSegment segment1 = TextSegment.from("I've been to France twice.");
    Embedding embedding1 = embeddingModel.embed(segment1).content();
    embeddingStore.add(embedding1, segment1);

    TextSegment segment2 = TextSegment.from("New Delhi is the capital of India.");
    Embedding embedding2 = embeddingModel.embed(segment2).content();
    embeddingStore.add(embedding2, segment2);

    Embedding queryEmbedding = embeddingModel.embed("Did you ever travel abroad?").content();
    List<EmbeddingMatch<TextSegment>> relevant = embeddingStore.findRelevant(queryEmbedding, 1);
    EmbeddingMatch<TextSegment> embeddingMatch = relevant.get(0);

    System.out.println(embeddingMatch.score());
    System.out.println(embeddingMatch.embedded().text());
  }
}

dev.langchain4j.store.embedding.pgvector.PgVectorEmbeddingStore.builder
[((989, 1290), 'dev.langchain4j.store.embedding.pgvector.PgVectorEmbeddingStore.builder'), ((989, 1273), 'dev.langchain4j.store.embedding.pgvector.PgVectorEmbeddingStore.builder'), ((989, 1244), 'dev.langchain4j.store.embedding.pgvector.PgVectorEmbeddingStore.builder'), ((989, 1215), 'dev.langchain4j.store.embedding.pgvector.PgVectorEmbeddingStore.builder'), ((989, 1190), 'dev.langchain4j.store.embedding.pgvector.PgVectorEmbeddingStore.builder'), ((989, 1149), 'dev.langchain4j.store.embedding.pgvector.PgVectorEmbeddingStore.builder'), ((989, 1125), 'dev.langchain4j.store.embedding.pgvector.PgVectorEmbeddingStore.builder'), ((989, 1095), 'dev.langchain4j.store.embedding.pgvector.PgVectorEmbeddingStore.builder'), ((989, 1068), 'dev.langchain4j.store.embedding.pgvector.PgVectorEmbeddingStore.builder'), ((989, 1048), 'dev.langchain4j.store.embedding.pgvector.PgVectorEmbeddingStore.builder')]