package com.lanyuanxiaoyao.service.ai.knowledge;

import io.qdrant.client.QdrantClient;
import io.qdrant.client.QdrantGrpcClient;
import io.qdrant.client.grpc.Collections;
import java.net.http.HttpClient;
import java.util.List;
import java.util.concurrent.ExecutionException;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.MetadataMode;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.openai.OpenAiEmbeddingModel;
import org.springframework.ai.openai.OpenAiEmbeddingOptions;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.qdrant.QdrantVectorStore;
import org.springframework.http.client.JdkClientHttpRequestFactory;
import org.springframework.http.client.reactive.JdkClientHttpConnector;
import org.springframework.web.client.RestClient;
import org.springframework.web.reactive.function.client.WebClient;

/**
 * @author lanyuanxiaoyao
 * @version 20250527
 */
public class TestEmbedding {
  public static void main(String[] args) throws ExecutionException, InterruptedException {
    HttpClient httpClient = HttpClient.newBuilder()
        .version(HttpClient.Version.HTTP_1_1)
        .build();
    EmbeddingModel model = new OpenAiEmbeddingModel(
        OpenAiApi.builder()
            .baseUrl("http://132.121.206.65:10086")
            .apiKey("*XMySqV%>hR&v>>g*NwCs3tpQ5FVMFEF2VHVTj<MYQd$&@$sY7CgqNyea4giJi4")
            .webClientBuilder(WebClient.builder().clientConnector(new JdkClientHttpConnector(httpClient)))
            .restClientBuilder(RestClient.builder().requestFactory(new JdkClientHttpRequestFactory(httpClient)))
            .build(),
        MetadataMode.ALL,
        OpenAiEmbeddingOptions.builder()
            .model("Bge-m3-vllm")
            .build()
    );
    QdrantClient client = new QdrantClient(
        QdrantGrpcClient.newBuilder("localhost", 6334, false).build()
    );
    client.createCollectionAsync(
        "test",
        Collections.VectorParams.newBuilder()
            .setDistance(Collections.Distance.Cosine)
            .setSize(1024)
            .build()
    ).get();
    VectorStore store = QdrantVectorStore.builder(client, model)
        .initializeSchema(true)
        .collectionName("test")
        .build();
    store.add(List.of(new Document("hello world")));
  }
}
