Merge pull request #239 from jmartisk/pinecone-improvements

Various Pinecone improvements

Author: geoand

pinecone/runtime/src/main/java/io/quarkiverse/langchain4j/pinecone/PineconeEmbeddingStore.java

@@ -25,7 +25,9 @@
 import dev.langchain4j.store.embedding.EmbeddingMatch;
 import dev.langchain4j.store.embedding.EmbeddingStore;
 import dev.langchain4j.store.embedding.RelevanceScore;
+import io.quarkiverse.langchain4j.pinecone.runtime.CreateIndexPodSpec;
 import io.quarkiverse.langchain4j.pinecone.runtime.CreateIndexRequest;
+import io.quarkiverse.langchain4j.pinecone.runtime.CreateIndexSpec;
 import io.quarkiverse.langchain4j.pinecone.runtime.DistanceMetric;
 import io.quarkiverse.langchain4j.pinecone.runtime.PineconeIndexOperationsApi;
 import io.quarkiverse.langchain4j.pinecone.runtime.PineconeVectorOperationsApi;
@@ -55,11 +57,13 @@ public PineconeEmbeddingStore(String apiKey,
             String namespace,
             String textFieldName,
             Duration timeout,
-            Integer dimension) {
+            Integer dimension,
+            String podType,
+            Duration indexReadinessTimeout) {
         this.indexName = indexName;
         this.dimension = dimension;
         String baseUrl = "https://" + indexName + "-" + projectId + ".svc." + environment + ".pinecone.io";
-        String baseUrlIndexOperations = "https://controller." + environment + ".pinecone.io";
+        String baseUrlIndexOperations = "https://api.pinecone.io";
         try {
             ClientHeadersFactory clientHeadersFactory = new ClientHeadersFactory() {
                 @Override
@@ -91,15 +95,18 @@ public MultivaluedMap<String, String> update(MultivaluedMap<String, String> inco
         this.indexExists = new LazyValue<>(new Supplier<Object>() {
             @Override
             public Object get() {
-                if (indexOperations.listIndexes().contains(indexName)) {
+                if (indexOperations.listIndexes().getIndexes().stream().anyMatch(i -> i.getName().equals(indexName))) {
                     Log.info("Pinecone index " + indexName + " already exists");
                 } else {
                     if (dimension == null) {
                         throw new IllegalArgumentException(
                                 "quarkus.langchain4j.pinecone.dimension must be specified when creating a new index");
                     }
-                    indexOperations.createIndex(new CreateIndexRequest(indexName, dimension, DistanceMetric.COSINE));
-                    Log.info("Created Pinecone index " + indexName + " with dimension = " + dimension);
+                    CreateIndexSpec spec = new CreateIndexSpec(new CreateIndexPodSpec(environment, podType));
+                    indexOperations.createIndex(new CreateIndexRequest(indexName, dimension, DistanceMetric.COSINE, spec));
+                    Log.info("Created Pinecone index " + indexName + " with dimension = " + dimension + ", " +
+                            "now waiting for it to be become ready...");
+                    waitForIndexToBecomeReady(indexName, indexReadinessTimeout);
                 }
                 return new Object();
             }
@@ -199,4 +206,20 @@ private void addAllInternal(List<String> ids, List<Embedding> embeddings, List<T
         Log.debug("Added embeddings: " + response.getUpsertedCount());
     }
 
+    private void waitForIndexToBecomeReady(String indexName, Duration timeout) {
+        long start = System.currentTimeMillis();
+        while (System.currentTimeMillis() - start < timeout.toMillis()) {
+            try {
+                Thread.sleep(1000);
+            } catch (InterruptedException e) {
+                throw new RuntimeException(e);
+            }
+            if (indexOperations.describeIndex(indexName).getStatus().isReady()) {
+                Log.info("Pinecone index " + indexName + " is now ready");
+                return;
+            }
+        }
+        throw new RuntimeException("Index " + indexName + " did not become ready within " + timeout);
+    }
+
 }

pinecone/runtime/src/main/java/io/quarkiverse/langchain4j/pinecone/runtime/CreateIndexPodSpec.java

@@ -0,0 +1,27 @@
+package io.quarkiverse.langchain4j.pinecone.runtime;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+import io.quarkus.runtime.annotations.RegisterForReflection;
+
+@RegisterForReflection
+public class CreateIndexPodSpec {
+
+    private final String environment;
+
+    @JsonProperty("pod_type")
+    private final String podType;
+
+    public CreateIndexPodSpec(String environment, String podType) {
+        this.environment = environment;
+        this.podType = podType;
+    }
+
+    public String getEnvironment() {
+        return environment;
+    }
+
+    public String getPodType() {
+        return podType;
+    }
+}

pinecone/runtime/src/main/java/io/quarkiverse/langchain4j/pinecone/runtime/CreateIndexRequest.java

@@ -14,11 +14,13 @@ public class CreateIndexRequest {
     private final String name;
     private final Integer dimension;
     private final DistanceMetric metric;
+    private final CreateIndexSpec spec;
 
-    public CreateIndexRequest(String name, Integer dimension, DistanceMetric metric) {
+    public CreateIndexRequest(String name, Integer dimension, DistanceMetric metric, CreateIndexSpec spec) {
         this.name = name;
         this.dimension = dimension;
         this.metric = metric;
+        this.spec = spec;
     }
 
     public String getName() {
@@ -32,4 +34,8 @@ public Integer getDimension() {
     public DistanceMetric getMetric() {
         return metric;
     }
+
+    public CreateIndexSpec getSpec() {
+        return spec;
+    }
 }

pinecone/runtime/src/main/java/io/quarkiverse/langchain4j/pinecone/runtime/CreateIndexSpec.java

@@ -0,0 +1,17 @@
+package io.quarkiverse.langchain4j.pinecone.runtime;
+
+import io.quarkus.runtime.annotations.RegisterForReflection;
+
+@RegisterForReflection
+public class CreateIndexSpec {
+
+    private final CreateIndexPodSpec pod;
+
+    public CreateIndexSpec(CreateIndexPodSpec pod) {
+        this.pod = pod;
+    }
+
+    public CreateIndexPodSpec getPod() {
+        return pod;
+    }
+}

pinecone/runtime/src/main/java/io/quarkiverse/langchain4j/pinecone/runtime/DescribeIndexResponse.java

@@ -0,0 +1,45 @@
+package io.quarkiverse.langchain4j.pinecone.runtime;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+
+public class DescribeIndexResponse {
+
+    private final String name;
+
+    private final DistanceMetric metric;
+
+    private final int dimension;
+
+    private final String host;
+
+    private final IndexStatus status;
+
+    @JsonCreator
+    public DescribeIndexResponse(String name, DistanceMetric metric, int dimension, String host, IndexStatus status) {
+        this.name = name;
+        this.metric = metric;
+        this.dimension = dimension;
+        this.host = host;
+        this.status = status;
+    }
+
+    public String getName() {
+        return name;
+    }
+
+    public DistanceMetric getMetric() {
+        return metric;
+    }
+
+    public int getDimension() {
+        return dimension;
+    }
+
+    public String getHost() {
+        return host;
+    }
+
+    public IndexStatus getStatus() {
+        return status;
+    }
+}

pinecone/runtime/src/main/java/io/quarkiverse/langchain4j/pinecone/runtime/DistanceMetric.java

@@ -1,9 +1,14 @@
 package io.quarkiverse.langchain4j.pinecone.runtime;
 
+import com.fasterxml.jackson.annotation.JsonProperty;
+
 public enum DistanceMetric {
 
+    @JsonProperty("euclidean")
     EUCLIDEAN,
+    @JsonProperty("cosine")
     COSINE,
+    @JsonProperty("dotproduct")
     DOTPRODUCT
 
 }

pinecone/runtime/src/main/java/io/quarkiverse/langchain4j/pinecone/runtime/IndexStatus.java

@@ -0,0 +1,27 @@
+package io.quarkiverse.langchain4j.pinecone.runtime;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+
+import io.quarkus.runtime.annotations.RegisterForReflection;
+
+@RegisterForReflection
+public class IndexStatus {
+
+    private final boolean ready;
+
+    private final String state;
+
+    @JsonCreator
+    public IndexStatus(boolean ready, String state) {
+        this.ready = ready;
+        this.state = state;
+    }
+
+    public boolean isReady() {
+        return ready;
+    }
+
+    public String getState() {
+        return state;
+    }
+}

pinecone/runtime/src/main/java/io/quarkiverse/langchain4j/pinecone/runtime/ListIndexesResponse.java

@@ -0,0 +1,22 @@
+package io.quarkiverse.langchain4j.pinecone.runtime;
+
+import java.util.List;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+
+import io.quarkus.runtime.annotations.RegisterForReflection;
+
+@RegisterForReflection
+public class ListIndexesResponse {
+
+    private final List<DescribeIndexResponse> indexes;
+
+    @JsonCreator
+    public ListIndexesResponse(List<DescribeIndexResponse> indexes) {
+        this.indexes = indexes;
+    }
+
+    public List<DescribeIndexResponse> getIndexes() {
+        return indexes;
+    }
+}

pinecone/runtime/src/main/java/io/quarkiverse/langchain4j/pinecone/runtime/PineconeConfig.java

@@ -39,6 +39,19 @@ public interface PineconeConfig {
      */
     Optional<Integer> dimension();
 
+    /**
+     * The type of the pod to use. This is only used if the index doesn't exist yet and needs to be created.
+     * The format: One of `s1`, `p1`, or `p2` appended with `.` and one of `x1`, `x2`, `x4`, or `x8`.
+     */
+    @WithDefault("s1.x1")
+    String podType();
+
+    /**
+     * The timeout duration for the index to become ready. Only relevant if the index doesn't exist yet and needs to be
+     * created. If not specified, 1 minute will be used.
+     */
+    Optional<Duration> indexReadinessTimeout();
+
     /**
      * The namespace.
      */

pinecone/runtime/src/main/java/io/quarkiverse/langchain4j/pinecone/runtime/PineconeExceptionMapper.java

@@ -0,0 +1,16 @@
+package io.quarkiverse.langchain4j.pinecone.runtime;
+
+import jakarta.ws.rs.core.Response;
+
+import org.eclipse.microprofile.rest.client.ext.ResponseExceptionMapper;
+
+public class PineconeExceptionMapper implements ResponseExceptionMapper<RuntimeException> {
+    @Override
+    public RuntimeException toThrowable(Response response) {
+        // FIXME: do this for some more status codes?
+        if (response.getStatus() == 400) {
+            return new RuntimeException("Pinecone returned 400, error: " + response.readEntity(String.class));
+        }
+        return null;
+    }
+}