Get and delete embeddings space

docs-site/docs/03_components/01_core_abstractions.md

@@ -55,12 +55,23 @@ public abstract class EmbeddingsStore {
 
 ### `PineconeEmbeddingsStore`
 
-eLLMental provides a concrete implementation for Pinecone, which requires defining an URL, an API Key and a space.
+eLLMental provides a concrete implementation for Pinecone, which requires defining a URL, an API Key, and a namespace.
 
 ```java
 EmbeddingsStore pineconeStore = new PineconeEmbeddingsStore("YOUR_PINECONE_URL", "YOUR_PINECONE_API_KEY", "YOUR_PINECONE_NAMESPACE");
 
-// You can now insert or perform similarity searches using the pineconeStore instance:
+// You can now insert, fetch, delete or perform similarity searches using the pineconeStore instance:
 pineconeStore.store(someEmbedding);
 List<Embedding> similarEmbeddings = pineconeStore.similaritySearch(referenceEmbedding, 5);
+
+// Get back an embedding by id without a namespace
+UUID someEmbeddingId = someEmbedding.id()
+Embedding embedding = pineconeStore.get(someEmbeddingId);
+
+// Get a namespaced embedding by id
+UUID anotherEmbeddingId = UUID.fromString("01870603-f211-7b9a-a7ea-4a98f5320ff8")
+Embedding anotherEmbedding = pineconeStore.get(anotherEmbeddingId, "my-namespace");
+
+// Delete an embedding by id
+pineconeStore.delete(anotherEmbeddingId);
 ```

modules/embeddingsspace/build.gradle

@@ -1,5 +1,6 @@
 plugins {
     id 'java'
+    id "io.freefair.lombok" version "8.3"
 }
 
 group = "com.theagilemonkeys.ellmental"
@@ -17,6 +18,8 @@ dependencies {
 
     testImplementation platform('org.junit:junit-bom:5.9.1')
     testImplementation 'org.junit.jupiter:junit-jupiter'
+    testImplementation 'org.mockito:mockito-core:3.+'
+    testImplementation 'org.mockito:mockito-junit-jupiter:5.5.0'
 }
 
 test {

modules/embeddingsspace/src/main/java/com/theagilemonkeys/ellmental/embeddingsspace/EmbeddingsSpaceComponent.java

@@ -3,19 +3,17 @@
 import com.theagilemonkeys.ellmental.core.schema.Embedding;
 import com.theagilemonkeys.ellmental.embeddingsgeneration.EmbeddingsGenerationModel;
 import com.theagilemonkeys.ellmental.embeddingsstore.EmbeddingsStore;
+import lombok.RequiredArgsConstructor;
 
 import java.util.List;
+import java.util.UUID;
 
+@RequiredArgsConstructor
 public class EmbeddingsSpaceComponent {
 
     private final EmbeddingsGenerationModel embeddingsGenerationModel;
     private final EmbeddingsStore embeddingsStore;
 
-    public EmbeddingsSpaceComponent(EmbeddingsGenerationModel embeddingsGenerationModel, EmbeddingsStore embeddingsStore) {
-        this.embeddingsGenerationModel = embeddingsGenerationModel;
-        this.embeddingsStore = embeddingsStore;
-    }
-
     private Embedding generate(String text) {
         return embeddingsGenerationModel.generateEmbedding(text);
     }
@@ -36,4 +34,11 @@ public List<Embedding> mostSimilarEmbeddings(String referenceText, int limit) {
         return mostSimilarEmbeddings(embedding, limit);
     }
 
+    public Embedding get(UUID uuid) {
+        return embeddingsStore.get(uuid);
+    }
+
+    public void delete(UUID uuid) {
+        embeddingsStore.delete(uuid);
+    }
 }

modules/embeddingsspace/src/test/java/com/theagilemonkeys/ellmental/embeddingsspace/EmbeddingsSpaceComponentTest.java

@@ -1,21 +1,100 @@
 package com.theagilemonkeys.ellmental.embeddingsspace;
 
+import com.theagilemonkeys.ellmental.core.schema.Embedding;
+import com.theagilemonkeys.ellmental.embeddingsgeneration.EmbeddingsGenerationModel;
+import com.theagilemonkeys.ellmental.embeddingsstore.EmbeddingsStore;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.*;
+import org.mockito.junit.jupiter.MockitoExtension;
 
-import static org.junit.jupiter.api.Assertions.assertTrue;
+import java.util.List;
+import java.util.UUID;
 
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.mockito.Mockito.*;
+
+@ExtendWith(MockitoExtension.class)
 public class EmbeddingsSpaceComponentTest {
 
-    // TODO: Use mockito to Mock these
-    // @Spy
-    // private final EmbeddingsGenerationModel model;
-    // @Spy
-    // private final EmbeddingsStore store;
-    // @Inject
-    // private final EmbeddingsSpaceComponent embeddingsSpaceComponent;
+    @Mock
+    private EmbeddingsGenerationModel embeddingsGenerationModel;
+    @Mock
+    private EmbeddingsStore embeddingsStore;
+    @InjectMocks
+    private EmbeddingsSpaceComponent embeddingsSpaceComponent;
+    @Captor
+    private ArgumentCaptor<String> sampleTextCaptor;
+
+    private final String sampleText = "sample text";
+    private final UUID embeddingId = UUID.randomUUID();
+    private final Embedding embeddingMock = new Embedding(embeddingId, List.of(1.0, 2.0, 3.0), null);
+
+    @Test
+    public void testSaveEmbedding() {
+        when(embeddingsGenerationModel.generateEmbedding(sampleText)).thenReturn(embeddingMock);
+
+        embeddingsSpaceComponent.save(sampleText);
+
+        verify(embeddingsGenerationModel).generateEmbedding(sampleTextCaptor.capture());
+        String capturedSampleText = sampleTextCaptor.getValue();
+        assertEquals(sampleText, capturedSampleText);
+
+        verify(embeddingsStore).store(embeddingMock);
+
+        verifyNoMoreInteractions(embeddingsGenerationModel, embeddingsStore);
+    }
+
+    @Test
+    public void testMostSimilarEmbeddings() {
+        int limit = 1;
+        Embedding similarEmbeddingMock = new Embedding(UUID.randomUUID(), List.of(4.0, 5.0, 6.0), null);
+
+        when(embeddingsStore.similaritySearch(embeddingMock, limit)).thenReturn(List.of(similarEmbeddingMock));
+
+        List<Embedding> embeddings = embeddingsSpaceComponent.mostSimilarEmbeddings(embeddingMock, limit);
+
+        assertEquals(embeddings.get(0), similarEmbeddingMock);
+
+        verifyNoMoreInteractions(embeddingsGenerationModel, embeddingsStore);
+    }
+
+    @Test
+    public void testMostSimilarEmbeddingsWhenUsingString() {
+        int limit = 1;
+        Embedding similarEmbeddingMock = new Embedding(UUID.randomUUID(), List.of(4.0, 5.0, 6.0), null);
+
+        when(embeddingsGenerationModel.generateEmbedding(sampleText)).thenReturn(embeddingMock);
+        when(embeddingsStore.similaritySearch(embeddingMock, limit)).thenReturn(List.of(similarEmbeddingMock));
+
+        List<Embedding> embeddings = embeddingsSpaceComponent.mostSimilarEmbeddings(sampleText, limit);
+
+        assertEquals(embeddings.get(0), similarEmbeddingMock);
+
+        verify(embeddingsGenerationModel).generateEmbedding(sampleTextCaptor.capture());
+        String capturedSampleText = sampleTextCaptor.getValue();
+        assertEquals(sampleText, capturedSampleText);
+
+        verifyNoMoreInteractions(embeddingsGenerationModel, embeddingsStore);
+    }
+
+    @Test
+    public void testGetEmbedding() {
+        when(embeddingsStore.get(embeddingId)).thenReturn(embeddingMock);
+
+        Embedding embedding = embeddingsSpaceComponent.get(embeddingId);
+
+        assertEquals(embeddingMock, embedding);
+
+        verifyNoMoreInteractions(embeddingsGenerationModel, embeddingsStore);
+    }
 
     @Test
-    public void todo() {
-        assertTrue(true);
+    public void testDeleteEmbedding() {
+        embeddingsSpaceComponent.delete(embeddingId);
+
+        verify(embeddingsStore).delete(embeddingId);
+
+        verifyNoMoreInteractions(embeddingsGenerationModel, embeddingsStore);
     }
-}
+}

modules/embeddingsstore/build.gradle

@@ -1,5 +1,6 @@
 plugins {
     id 'java'
+    id "io.freefair.lombok" version "8.3"
 }
 
 group = 'com.theagilemonkeys.ellemental'
@@ -16,6 +17,8 @@ dependencies {
 
     testImplementation platform('org.junit:junit-bom:5.9.1')
     testImplementation 'org.junit.jupiter:junit-jupiter'
+    testImplementation 'org.mockito:mockito-core:3.+'
+    testImplementation 'org.mockito:mockito-junit-jupiter:5.5.0'
 }
 
 test {

modules/embeddingsstore/src/main/java/com/theagilemonkeys/ellmental/embeddingsstore/EmbeddingsStore.java

@@ -2,6 +2,7 @@
 
 import java.util.List;
 import java.util.Map;
+import java.util.UUID;
 
 import com.theagilemonkeys.ellmental.core.schema.Embedding;
 
@@ -11,5 +12,7 @@
 public abstract class EmbeddingsStore {
     //TODO: check if there is an issue using a map instead of a metadata class
     public abstract void store(Embedding embedding);
+    public abstract Embedding get(UUID uuid);
+    public abstract void delete(UUID uuid);
     public abstract List<Embedding> similaritySearch(Embedding reference, int limit);
 }

modules/embeddingsstore/src/main/java/com/theagilemonkeys/ellmental/embeddingsstore/pinecone/DeleteVectorSchema.java

@@ -0,0 +1,15 @@
+package com.theagilemonkeys.ellmental.embeddingsstore.pinecone;
+
+import lombok.Data;
+
+import java.util.List;
+import java.util.Map;
+import java.util.UUID;
+
+@Data
+public class DeleteVectorSchema {
+    private final List<UUID> ids;
+    private boolean deleteAll;
+    private String namespace;
+    private Map<String, String> filterMetadata;
+}

modules/embeddingsstore/src/main/java/com/theagilemonkeys/ellmental/embeddingsstore/pinecone/FetchVectorResponseSchema.java

@@ -0,0 +1,11 @@
+package com.theagilemonkeys.ellmental.embeddingsstore.pinecone;
+
+import lombok.Data;
+
+import java.util.List;
+
+@Data
+public class FetchVectorResponseSchema {
+    private final List<Vector> vectors;
+    private final String namespace;
+}

modules/embeddingsstore/src/main/java/com/theagilemonkeys/ellmental/embeddingsstore/pinecone/FetchVectorSchema.java

@@ -0,0 +1,14 @@
+package com.theagilemonkeys.ellmental.embeddingsstore.pinecone;
+
+import lombok.AllArgsConstructor;
+import lombok.Data;
+
+import java.util.List;
+import java.util.UUID;
+
+@Data
+@AllArgsConstructor
+public class FetchVectorSchema {
+    private final List<UUID> ids;
+    private String namespace;
+}

modules/embeddingsstore/src/main/java/com/theagilemonkeys/ellmental/embeddingsstore/pinecone/PineconeEmbeddingsStore.java

@@ -3,10 +3,12 @@
 import com.google.gson.Gson;
 import com.theagilemonkeys.ellmental.embeddingsstore.EmbeddingsStore;
 import java.io.IOException;
-import java.util.ArrayList;
-import java.util.List;
+import java.util.*;
+
 import com.theagilemonkeys.ellmental.core.schema.Embedding;
 import com.theagilemonkeys.ellmental.core.errors.MissingRequiredCredentialException;
+import lombok.Data;
+import lombok.RequiredArgsConstructor;
 import okhttp3.*;
 import static java.net.HttpURLConnection.HTTP_BAD_REQUEST;
 
@@ -57,6 +59,51 @@ public void store(Embedding embedding) {
 
     }
 
+    /**
+     * Retrieves and embedding in the embeddings store based on its UUID.
+     * @param uuid UUID to retrieve.
+     */
+    public Embedding get(UUID uuid) {
+        Embedding embedding = null;
+        try {
+            embedding = this.fetch(List.of(uuid), null);
+        } catch (IOException e) {
+            System.out.println("VectorStore error on fetch: " + e.getMessage());
+        }
+        return embedding;
+    }
+
+    /**
+     * Retrieves and embedding in the embeddings store based on its UUID.
+     * @param uuid UUID to retrieve.
+     * @param namespace String to look into.
+     */
+    public Embedding get(UUID uuid, String namespace) {
+        Embedding embedding = null;
+        try {
+            embedding = this.fetch(List.of(uuid), namespace);
+        } catch (IOException e) {
+            System.out.println("VectorStore error on fetch: " + e.getMessage());
+        }
+        return embedding;
+    }
+
+    /**
+     * Deletes an embedding in the embeddings store based on its UUID.
+     * @param uuid UUID to delete.
+     */
+    public void delete(UUID uuid) {
+        DeleteVectorSchema request = new DeleteVectorSchema(List.of(uuid));
+        String requestBodyJson = new Gson().toJson(request);
+
+        try {
+            this.post("/vectors/delete", requestBodyJson);
+        } catch (IOException e) {
+            System.out.println("VectorStore error on delete: " + e.getMessage());
+        }
+
+    }
+
     public List<Embedding> similaritySearch(Embedding reference, int limit) {
         QueryVectorRequestSchema body = new QueryVectorRequestSchema(
                 limit,
@@ -92,6 +139,45 @@ private boolean validateEnvironment() {
         return true;
     }
 
+    private Embedding fetch(List<UUID> ids, String namespace) throws IOException {
+        if (!validateEnvironment()) {
+            return null;
+        }
+
+        HttpUrl.Builder urlBuilder = Objects.requireNonNull(HttpUrl.parse(this.url + "/vectors/fetch")).newBuilder();
+        urlBuilder.addQueryParameter("ids", ids.toString());
+        if (namespace != null) {
+            urlBuilder.addQueryParameter("namespace", namespace);
+        }
+
+        HttpUrl url = urlBuilder.build();
+
+        Request request = new Request.Builder()
+                .url(url)
+                .header("accept", "application/json")
+                .header("content-type", "application/json")
+                .header("Api-Key", apiKey)
+                .get()
+                .build();
+
+        try (Response response = new OkHttpClient().newCall(request).execute()) {
+            if (response.code() >= HTTP_BAD_REQUEST) {
+                throw new IOException(this.url);
+            }
+
+            ResponseBody responseBody = response.body();
+
+            if (responseBody != null) {
+                String json = responseBody.string();
+                FetchVectorResponseSchema responseSchema = new Gson().fromJson(json, FetchVectorResponseSchema.class);
+                Vector vector = responseSchema.getVectors().get(0);
+                return new Embedding(vector.getId(), vector.getValues(), vector.getMetadata());
+            } else {
+                return null;
+            }
+        }
+    }
+
     private String post(String path, String bodyString) throws IOException {
         if (!validateEnvironment()) {
             return null;

modules/embeddingsstore/src/main/java/com/theagilemonkeys/ellmental/embeddingsstore/pinecone/SparseValues.java

@@ -0,0 +1,8 @@
+package com.theagilemonkeys.ellmental.embeddingsstore.pinecone;
+
+import java.util.List;
+
+public class SparseValues {
+    public List<Integer> indices;
+    public List<Integer> values;
+}