From a3711974109eece202f05206514a9f342e4b1a48 Mon Sep 17 00:00:00 2001
From: Jan Martiska <jmartisk@redhat.com>
Date: Wed, 6 Dec 2023 10:12:33 +0100
Subject: [PATCH] Adjust to recent Pinecone query response format change

---
 .../PineconeEmbeddingStoreTest.java           |  2 +-
 .../pinecone/PineconeEmbeddingStore.java      | 20 ++++++++++++-------
 .../pinecone/runtime/QueryResponse.java       | 11 +++++-----
 3 files changed, 19 insertions(+), 14 deletions(-)

diff --git a/pinecone/deployment/src/test/java/io/quarkiverse/langchain4j/pinecone/deployment/PineconeEmbeddingStoreTest.java b/pinecone/deployment/src/test/java/io/quarkiverse/langchain4j/pinecone/deployment/PineconeEmbeddingStoreTest.java
index 4ef60756b..fd7f02d24 100644
--- a/pinecone/deployment/src/test/java/io/quarkiverse/langchain4j/pinecone/deployment/PineconeEmbeddingStoreTest.java
+++ b/pinecone/deployment/src/test/java/io/quarkiverse/langchain4j/pinecone/deployment/PineconeEmbeddingStoreTest.java
@@ -98,7 +98,7 @@ public void cleanup() {
      */
     private static void delay() {
         try {
-            int timeout = 40;
+            int timeout = 30;
             Log.info("Waiting " + timeout + " seconds to allow Pinecone time to process deletions");
             TimeUnit.SECONDS.sleep(timeout);
         } catch (InterruptedException e) {
diff --git a/pinecone/runtime/src/main/java/io/quarkiverse/langchain4j/pinecone/PineconeEmbeddingStore.java b/pinecone/runtime/src/main/java/io/quarkiverse/langchain4j/pinecone/PineconeEmbeddingStore.java
index 092f6d11b..a581489fa 100644
--- a/pinecone/runtime/src/main/java/io/quarkiverse/langchain4j/pinecone/PineconeEmbeddingStore.java
+++ b/pinecone/runtime/src/main/java/io/quarkiverse/langchain4j/pinecone/PineconeEmbeddingStore.java
@@ -149,13 +149,19 @@ public List<EmbeddingMatch<TextSegment>> findRelevant(Embedding embedding, int m
         QueryRequest request = new QueryRequest(namespace, (long) maxResults, true, true, embedding.vector());
         QueryResponse response = vectorOperations.query(request);
         return response
-                .getMatches().stream().map(match -> new EmbeddingMatch<>(
-                        RelevanceScore.fromCosineSimilarity(match.getScore()),
-                        match.getId(),
-                        new Embedding(match.getValues()),
-                        match.getMetadata().get(textFieldName) != null ? new TextSegment(
-                                match.getMetadata().get(textFieldName),
-                                new Metadata(mapWithoutKey(match.getMetadata(), textFieldName))) : null))
+                .getMatches().stream().map(match -> {
+                    String text = match.getMetadata() != null &&
+                            match.getMetadata().get(textFieldName) != null
+                                    ? match.getMetadata().get(textFieldName)
+                                    : null;
+                    return new EmbeddingMatch<>(
+                            RelevanceScore.fromCosineSimilarity(match.getScore()),
+                            match.getId(),
+                            new Embedding(match.getValues()),
+                            text != null ? new TextSegment(
+                                    text,
+                                    new Metadata(mapWithoutKey(match.getMetadata(), textFieldName))) : null);
+                })
                 .filter(match -> match.score() >= minScore)
                 .collect(toList());
     }
diff --git a/pinecone/runtime/src/main/java/io/quarkiverse/langchain4j/pinecone/runtime/QueryResponse.java b/pinecone/runtime/src/main/java/io/quarkiverse/langchain4j/pinecone/runtime/QueryResponse.java
index fc2e6260a..2cfbda5b4 100644
--- a/pinecone/runtime/src/main/java/io/quarkiverse/langchain4j/pinecone/runtime/QueryResponse.java
+++ b/pinecone/runtime/src/main/java/io/quarkiverse/langchain4j/pinecone/runtime/QueryResponse.java
@@ -15,20 +15,19 @@ public class QueryResponse {
 
     private final List<VectorMatch> matches;
 
-    private String metadata;
+    private String namespace;
 
     @JsonCreator
-    public QueryResponse(List<VectorMatch> matches, String metadata) {
+    public QueryResponse(List<VectorMatch> matches, String metadata, String namespace) {
         this.matches = matches;
-        this.metadata = metadata;
+        this.namespace = namespace;
     }
 
     public List<VectorMatch> getMatches() {
         return matches;
     }
 
-    public String getMetadata() {
-        return metadata;
+    public String getNamespace() {
+        return namespace;
     }
-
 }