From fadd524c7d25a4db0fe5ab8ce4c437d352f67502 Mon Sep 17 00:00:00 2001
From: Javier Toledo <javier@theagilemonkeys.com>
Date: Tue, 29 Aug 2023 18:13:36 +0100
Subject: [PATCH] Core abstractions review

* Reviewed that documentation matches the code and fixed docs discrepancies
* Refactored the `Embedding` object to be a record and include metadata and an ID that will be retrieved afterwards from the vector store
* Implemented constructor that accepts explicit credentials in addition to supporting environment variables
* The OpenAI embeddings generator now injects the original text, the model name and a timestamp in the embeddings metadata automatically
---
 .idea/workspace.xml                           |  83 +++++++----
 .../ellmental/core/schema/Embedding.java      |  20 +--
 .../03_components/01_core_abstractions.md     |  35 ++++-
 .../docs/03_components/02_embeddings_space.md |  20 ---
 .../ellmental/helloworld/Main.java            |   5 +-
 .../EmbeddingsGenerationModel.java            |   3 +
 .../openai/OpenAIEmbeddingsModel.java         |  75 +++++++---
 .../openai/OpenAIEmbeddingsModelTest.java     |   4 +-
 .../embeddingsstore/EmbeddingsStore.java      |   5 +-
 .../pinecone/PineconeEmbeddingsStore.java     | 136 +++++++++++-------
 .../pinecone/QueryVectorResponseSchema.java   |   3 +-
 .../pinecone/UpsertVectorSchema.java          |  10 +-
 .../pinecone/PineconeEmbeddingsStoreTest.java |  10 +-
 13 files changed, 258 insertions(+), 151 deletions(-)

diff --git a/.idea/workspace.xml b/.idea/workspace.xml
index 9744266..0a1674e 100644
--- a/.idea/workspace.xml
+++ b/.idea/workspace.xml
@@ -5,15 +5,19 @@
   </component>
   <component name="ChangeListManager">
     <list default="true" id="9cad3848-1525-4d25-a87c-214202721c5d" name="Changes" comment="">
-      <change afterPath="$PROJECT_DIR$/.envrc" afterDir="false" />
-      <change afterPath="$PROJECT_DIR$/.idea/aws.xml" afterDir="false" />
-      <change afterPath="$PROJECT_DIR$/.idea/inspectionProfiles/Project_Default.xml" afterDir="false" />
-      <change afterPath="$PROJECT_DIR$/flake.lock" afterDir="false" />
-      <change afterPath="$PROJECT_DIR$/flake.nix" afterDir="false" />
-      <change beforePath="$PROJECT_DIR$/.idea/gradle.xml" beforeDir="false" afterPath="$PROJECT_DIR$/.idea/gradle.xml" afterDir="false" />
       <change beforePath="$PROJECT_DIR$/.idea/workspace.xml" beforeDir="false" afterPath="$PROJECT_DIR$/.idea/workspace.xml" afterDir="false" />
-      <change beforePath="$PROJECT_DIR$/build.gradle" beforeDir="false" afterPath="$PROJECT_DIR$/build.gradle" afterDir="false" />
-      <change beforePath="$PROJECT_DIR$/gradle/wrapper/gradle-wrapper.properties" beforeDir="false" afterPath="$PROJECT_DIR$/gradle/wrapper/gradle-wrapper.properties" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/core/src/main/java/com/theagilemonkeys/ellmental/core/schema/Embedding.java" beforeDir="false" afterPath="$PROJECT_DIR$/core/src/main/java/com/theagilemonkeys/ellmental/core/schema/Embedding.java" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/docs-site/docs/03_components/01_core_abstractions.md" beforeDir="false" afterPath="$PROJECT_DIR$/docs-site/docs/03_components/01_core_abstractions.md" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/docs-site/docs/03_components/02_embeddings_space.md" beforeDir="false" afterPath="$PROJECT_DIR$/docs-site/docs/03_components/02_embeddings_space.md" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/examples/simplejava/src/main/java/com/theagilemonkeys/ellmental/helloworld/Main.java" beforeDir="false" afterPath="$PROJECT_DIR$/examples/simplejava/src/main/java/com/theagilemonkeys/ellmental/helloworld/Main.java" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/modules/embeddingsgeneration/src/main/java/com/theagilemonkeys/ellmental/embeddingsgeneration/EmbeddingsGenerationModel.java" beforeDir="false" afterPath="$PROJECT_DIR$/modules/embeddingsgeneration/src/main/java/com/theagilemonkeys/ellmental/embeddingsgeneration/EmbeddingsGenerationModel.java" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/modules/embeddingsgeneration/src/main/java/com/theagilemonkeys/ellmental/embeddingsgeneration/openai/OpenAIEmbeddingsModel.java" beforeDir="false" afterPath="$PROJECT_DIR$/modules/embeddingsgeneration/src/main/java/com/theagilemonkeys/ellmental/embeddingsgeneration/openai/OpenAIEmbeddingsModel.java" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/modules/embeddingsgeneration/src/test/java/com/theagilemonkeys/ellmental/embeddingsgeneration/openai/OpenAIEmbeddingsModelTest.java" beforeDir="false" afterPath="$PROJECT_DIR$/modules/embeddingsgeneration/src/test/java/com/theagilemonkeys/ellmental/embeddingsgeneration/openai/OpenAIEmbeddingsModelTest.java" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/modules/embeddingsstore/src/main/java/com/theagilemonkeys/ellmental/embeddingsstore/EmbeddingsStore.java" beforeDir="false" afterPath="$PROJECT_DIR$/modules/embeddingsstore/src/main/java/com/theagilemonkeys/ellmental/embeddingsstore/EmbeddingsStore.java" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/modules/embeddingsstore/src/main/java/com/theagilemonkeys/ellmental/embeddingsstore/pinecone/PineconeEmbeddingsStore.java" beforeDir="false" afterPath="$PROJECT_DIR$/modules/embeddingsstore/src/main/java/com/theagilemonkeys/ellmental/embeddingsstore/pinecone/PineconeEmbeddingsStore.java" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/modules/embeddingsstore/src/main/java/com/theagilemonkeys/ellmental/embeddingsstore/pinecone/QueryVectorResponseSchema.java" beforeDir="false" afterPath="$PROJECT_DIR$/modules/embeddingsstore/src/main/java/com/theagilemonkeys/ellmental/embeddingsstore/pinecone/QueryVectorResponseSchema.java" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/modules/embeddingsstore/src/main/java/com/theagilemonkeys/ellmental/embeddingsstore/pinecone/UpsertVectorSchema.java" beforeDir="false" afterPath="$PROJECT_DIR$/modules/embeddingsstore/src/main/java/com/theagilemonkeys/ellmental/embeddingsstore/pinecone/UpsertVectorSchema.java" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/modules/embeddingsstore/src/test/java/com/theagilemonkeys/ellmental/embeddingsstore/pinecone/PineconeEmbeddingsStoreTest.java" beforeDir="false" afterPath="$PROJECT_DIR$/modules/embeddingsstore/src/test/java/com/theagilemonkeys/ellmental/embeddingsstore/pinecone/PineconeEmbeddingsStoreTest.java" afterDir="false" />
     </list>
     <option name="SHOW_DIALOG" value="false" />
     <option name="HIGHLIGHT_CONFLICTS" value="true" />
@@ -25,37 +29,60 @@
       <ProjectState />
     </projectState>
   </component>
+  <component name="ExternalProjectsManager">
+    <system id="GRADLE">
+      <state>
+        <task path="$PROJECT_DIR$">
+          <activation />
+        </task>
+        <projects_view>
+          <tree_state>
+            <expand>
+              <path>
+                <item name="" type="6a2764b6:ExternalProjectsStructure$RootNode" />
+                <item name="eLLMental" type="f1a62948:ProjectNode" />
+              </path>
+            </expand>
+            <select />
+          </tree_state>
+        </projects_view>
+      </state>
+    </system>
+  </component>
   <component name="Git.Settings">
     <option name="RECENT_GIT_ROOT_PATH" value="$PROJECT_DIR$" />
   </component>
   <component name="MarkdownSettingsMigration">
     <option name="stateVersion" value="1" />
   </component>
+  <component name="ProblemsViewState">
+    <option name="selectedTabId" value="ProjectErrors" />
+  </component>
   <component name="ProjectId" id="2Thv794LljpJB6whwEafjMfg13O" />
   <component name="ProjectLevelVcsManager" settingsEditedManually="true" />
   <component name="ProjectViewState">
     <option name="hideEmptyMiddlePackages" value="true" />
     <option name="showLibraryContents" value="true" />
   </component>
-  <component name="PropertiesComponent"><![CDATA[{
-  "keyToString": {
-    "RunOnceActivity.OpenProjectViewOnStart": "true",
-    "RunOnceActivity.ShowReadmeOnStart": "true",
-    "SHARE_PROJECT_CONFIGURATION_FILES": "true",
-    "TODO_SCOPE": "All Places",
-    "WebServerToolWindowFactoryState": "false",
-    "git-widget-placeholder": "main",
-    "jdk.selected.JAVA_MODULE": "openjdk-20",
-    "last_opened_file_path": "/Users/javier/projects/eLLMental",
-    "node.js.detected.package.eslint": "true",
-    "node.js.detected.package.tslint": "true",
-    "node.js.selected.package.eslint": "(autodetect)",
-    "node.js.selected.package.tslint": "(autodetect)",
-    "nodejs_package_manager_path": "npm",
-    "settings.editor.selected.configurable": "preferences.pluginManager",
-    "vue.rearranger.settings.migration": "true"
+  <component name="PropertiesComponent">{
+  &quot;keyToString&quot;: {
+    &quot;RunOnceActivity.OpenProjectViewOnStart&quot;: &quot;true&quot;,
+    &quot;RunOnceActivity.ShowReadmeOnStart&quot;: &quot;true&quot;,
+    &quot;SHARE_PROJECT_CONFIGURATION_FILES&quot;: &quot;true&quot;,
+    &quot;TODO_SCOPE&quot;: &quot;All Places&quot;,
+    &quot;WebServerToolWindowFactoryState&quot;: &quot;false&quot;,
+    &quot;git-widget-placeholder&quot;: &quot;main&quot;,
+    &quot;jdk.selected.JAVA_MODULE&quot;: &quot;openjdk-20&quot;,
+    &quot;last_opened_file_path&quot;: &quot;/Users/javier/projects/eLLMental&quot;,
+    &quot;node.js.detected.package.eslint&quot;: &quot;true&quot;,
+    &quot;node.js.detected.package.tslint&quot;: &quot;true&quot;,
+    &quot;node.js.selected.package.eslint&quot;: &quot;(autodetect)&quot;,
+    &quot;node.js.selected.package.tslint&quot;: &quot;(autodetect)&quot;,
+    &quot;nodejs_package_manager_path&quot;: &quot;npm&quot;,
+    &quot;settings.editor.selected.configurable&quot;: &quot;preferences.pluginManager&quot;,
+    &quot;vue.rearranger.settings.migration&quot;: &quot;true&quot;
   }
-}]]></component>
+}</component>
   <component name="RecentsManager">
     <key name="CopyFile.RECENT_KEYS">
       <recent name="$PROJECT_DIR$" />
@@ -70,7 +97,7 @@
         <option name="executionName" />
         <option name="externalProjectPath" value="$PROJECT_DIR$" />
         <option name="externalSystemIdString" value="GRADLE" />
-        <option name="scriptParameters" value="" />
+        <option name="scriptParameters" value="--scan" />
         <option name="taskDescriptions">
           <list />
         </option>
@@ -111,6 +138,8 @@
       <workItem from="1691507960273" duration="17280000" />
       <workItem from="1691687224885" duration="6374000" />
       <workItem from="1692634591758" duration="494000" />
+      <workItem from="1693322245615" duration="4859000" />
+      <workItem from="1693327144223" duration="1788000" />
     </task>
     <servers />
   </component>
diff --git a/core/src/main/java/com/theagilemonkeys/ellmental/core/schema/Embedding.java b/core/src/main/java/com/theagilemonkeys/ellmental/core/schema/Embedding.java
index e845f13..c4a638f 100644
--- a/core/src/main/java/com/theagilemonkeys/ellmental/core/schema/Embedding.java
+++ b/core/src/main/java/com/theagilemonkeys/ellmental/core/schema/Embedding.java
@@ -1,12 +1,14 @@
 package com.theagilemonkeys.ellmental.core.schema;
 
 import java.util.List;
-
-public class Embedding {
-
-    public Embedding(List<Double> vector) {
-        this.vector = vector;
-    }
-    public List<Double> vector;
-
-}
+import java.util.Map;
+import java.util.UUID;
+
+/**
+ * Embeddings represent a point in the embeddings space, representing the semantics of a given text.
+ */
+public record Embedding(
+        UUID id,
+        List<Double> vector,
+        Map<String, String> metadata
+) {}
diff --git a/docs-site/docs/03_components/01_core_abstractions.md b/docs-site/docs/03_components/01_core_abstractions.md
index 9327987..732d4c7 100644
--- a/docs-site/docs/03_components/01_core_abstractions.md
+++ b/docs-site/docs/03_components/01_core_abstractions.md
@@ -1,6 +1,23 @@
 # Core Abstractions
 
-eLLMental uses different 3rd party components and APIs and provides a unified interface. To ensure extensibility and avoid tight coupling with any specific API, the library provides a series of abstract classes that define the expected interface for these components to work with eLLMental. To use eLLMental, you can provide your own implementation or use one of the built-in concrete implementations.  
+eLLMental uses different 3rd party components and APIs and provides a unified interface. To ensure extensibility and avoid tight coupling with any specific API, the library provides a series of abstract classes that define the expected interface for these components to work with eLLMental. To use eLLMental, you can provide your own implementation or use one of the built-in concrete implementations.
+
+## `Embedding` object
+
+In eLLMental embeddings are represented by the `Embedding` record, which has the following attributes:
+
+- `id`: An unique identifier of the embedding.
+- `vector`: A numeric vector that represents the semantic location of the text.
+- `metadata`: Additional information associated with the embedding. It can be used to store the original text, the model used to generate the embedding, or any other information you may find useful.
+
+```java
+public record Embedding(
+        UUID id,
+        List<Double> vector,
+        Map<String, String> metadata
+) {}
+
+```
 
 ## `EmbeddingsGenerationModel`
 
@@ -17,19 +34,22 @@ public abstract class EmbeddingsGenerationModel {
 eLLMental provides an implementation to use [OpenAI's embeddings model](https://platform.openai.com/docs/guides/embeddings). This model is only accessible via API, so you'll need to initialize it with a valid OpenAI API key.
 
 ```java
+// You can explicitly initialize the model with your API key, or use the constructor without parameters and set the `OPENAI_API_KEY` environment variable.
 EmbeddingsGenerationModel openAIModel = new OpenAIEmbeddingsGenerationModel("YOUR_OPENAI_API_KEY");
 
-// You'll rarely need to interact directly with the `openAIModel`, but you can use it to generate an embedding:
-openAIModel.generateEmbedding("Sample string");
+// You'll rarely need to interact directly with the `openAIModel`, but you can use it to generate an embedding object:
+Embedding embedding = openAIModel.generateEmbedding("Sample string");
 ```
 
+The OpenAI embeddings generator will automatically include the original text, the timestamp and the model used to generate the embedding in the metadata. 
+
 ## `EmbeddingsStore`
 
 This abstract class defines the expected interface for a persistence mechanism capable of storing and querying embeddings:
 
 ```java
 public abstract class EmbeddingsStore {
-    public abstract void store(Embedding embedding, Metadata metadata);
+    public abstract void store(Embedding embedding);
     public abstract List<Embedding> similaritySearch(Embedding reference, int limit);
 }
 ```
@@ -39,9 +59,10 @@ public abstract class EmbeddingsStore {
 eLLMental provides a concrete implementation for Pinecone, which requires defining an URL, an API Key and a space.
 
 ```java
-EmbeddingsStore pineconeStore = new PineconeEmbeddingsStore("YOUR_PINECONE_URL", "YOUR_PINECONE_API_KEY", "YOUR_PINECONE_SPACE");
+// You can explicitly initialize the store with your credentials, or use the constructor with no parameters and set the `PINECONE_URL`, `PINECONE_API_KEY` and `PINECONE_NAMESPACE` environment variables.
+EmbeddingsStore pineconeStore = new PineconeEmbeddingsStore("YOUR_PINECONE_URL", "YOUR_PINECONE_API_KEY", "YOUR_PINECONE_NAMESPACE");
 
-// You can insert or perform similarity searches using this object. Metadata is optional.
-pineconeStore.store(someEmbedding, someMetadata);
+// You can now insert or perform similarity searches using the pineconeStore instance:
+pineconeStore.store(someEmbedding);
 List<Embedding> similarEmbeddings = pineconeStore.similaritySearch(referenceEmbedding, 5);
 ```
\ No newline at end of file
diff --git a/docs-site/docs/03_components/02_embeddings_space.md b/docs-site/docs/03_components/02_embeddings_space.md
index cc267f4..285b6e2 100644
--- a/docs-site/docs/03_components/02_embeddings_space.md
+++ b/docs-site/docs/03_components/02_embeddings_space.md
@@ -11,26 +11,6 @@ Leveraging the power of embeddings models, this component allows you to represen
 [//]: # (TODO: Highlight this as a warning message in Docusaurus)
 > **Warning**: The current version supports the [OpenAI embeddings model](https://platform.openai.com/docs/guides/embeddings) and [Pinecone](https://www.pinecone.io) for storage. Please provide the necessary credentials for these services. Note: these services may involve costs, always review their pricing details before use.
 
-## `Embedding` object
-
-In eLLMental embeddings are represented by the `Embedding` class, which has the following attributes:
-
-- `id`: The unique identifier of the embedding.
-- `embedding`: The numeric representation of the text.
-- `metadata`: Additional information associated with the embedding.
-- `createdAt`: The timestamp of the embedding creation.
-- `modelId`: The identifier of the embeddings model used to generate the embedding.
-
-```java
-public class Embedding {
-    private final String id;
-    private final float[] embedding;
-    private final Map<String, String> metadata;
-    private final Instant createdAt;
-    private final String modelId;
-}
-```
-
 The `EmbeddingsSpaceComponent` interface defines the following methods:
 
 ## Constructor
diff --git a/examples/simplejava/src/main/java/com/theagilemonkeys/ellmental/helloworld/Main.java b/examples/simplejava/src/main/java/com/theagilemonkeys/ellmental/helloworld/Main.java
index af457b8..2dc7e4f 100644
--- a/examples/simplejava/src/main/java/com/theagilemonkeys/ellmental/helloworld/Main.java
+++ b/examples/simplejava/src/main/java/com/theagilemonkeys/ellmental/helloworld/Main.java
@@ -16,10 +16,7 @@ public static void main(String[] args) {
 
         // Step 2: save the generated embeddings to a store (Pinecone in this case)
         EmbeddingsStore embeddingStore = new PineconeEmbeddingsStore();
-        Map<String, String> metadata = new HashMap<>();
-        metadata.put("key1", "value1");
-        metadata.put("key2", "value2");
-        embeddingStore.store(embedding, metadata);
+        embeddingStore.store(embedding);
 
         // Step 3: search for the embedding in the store
         List<Embedding> searchEmbeddings = embeddingStore.similaritySearch(embedding, 5);
diff --git a/modules/embeddingsgeneration/src/main/java/com/theagilemonkeys/ellmental/embeddingsgeneration/EmbeddingsGenerationModel.java b/modules/embeddingsgeneration/src/main/java/com/theagilemonkeys/ellmental/embeddingsgeneration/EmbeddingsGenerationModel.java
index 0eddceb..c3c47a9 100644
--- a/modules/embeddingsgeneration/src/main/java/com/theagilemonkeys/ellmental/embeddingsgeneration/EmbeddingsGenerationModel.java
+++ b/modules/embeddingsgeneration/src/main/java/com/theagilemonkeys/ellmental/embeddingsgeneration/EmbeddingsGenerationModel.java
@@ -3,6 +3,9 @@
 
 import com.theagilemonkeys.ellmental.core.schema.Embedding;
 
+/**
+ * Abstract class that defines the interface expected by eLLMental for a valid embeddings generation model.
+ */
 public abstract class EmbeddingsGenerationModel {
     public abstract Embedding generateEmbedding(String text);
 }
\ No newline at end of file
diff --git a/modules/embeddingsgeneration/src/main/java/com/theagilemonkeys/ellmental/embeddingsgeneration/openai/OpenAIEmbeddingsModel.java b/modules/embeddingsgeneration/src/main/java/com/theagilemonkeys/ellmental/embeddingsgeneration/openai/OpenAIEmbeddingsModel.java
index 54a9e2c..7ddeac3 100644
--- a/modules/embeddingsgeneration/src/main/java/com/theagilemonkeys/ellmental/embeddingsgeneration/openai/OpenAIEmbeddingsModel.java
+++ b/modules/embeddingsgeneration/src/main/java/com/theagilemonkeys/ellmental/embeddingsgeneration/openai/OpenAIEmbeddingsModel.java
@@ -11,30 +11,52 @@
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
+import java.util.UUID;
 
+/**
+ * OpenAI `EmbeddingsGenerationModel` implementation.
+ */
 public class OpenAIEmbeddingsModel extends EmbeddingsGenerationModel {
-    private static OpenAiService service;
+    private static String openAIKey;
+    private static OpenAiService cachedService;
     public static String embeddingOpenAiModel = "text-embedding-ada-002";
 
+    /**
+     * Constructor that initializes the OpenAI embeddings model.
+     * It will try to load the API key from the environment variable `OPEN_AI_API_KEY`.
+     */
     public OpenAIEmbeddingsModel() {
-        var dotenv = Dotenv
-                .configure()
-                .ignoreIfMissing()
-                .ignoreIfMalformed()
-                .load();
-        String openAIKey = dotenv.get("OPEN_AI_API_KEY");
-        if (openAIKey == null) {
-            throw new EnvironmentVariableNotDeclaredException("Environment variable OPEN_AI_API_KEY is not declared.");
-        }
-        service = new OpenAiService(openAIKey);
+        this(null);
+    }
 
+    /**
+     * Constructor that initializes the OpenAI embeddings model with an explicit API Key.
+     *
+     * @param openAIKey OpenAI API key.
+     */
+    public OpenAIEmbeddingsModel(String APIKey) {
+        if (APIKey != null) {
+            openAIKey = APIKey;
+        } else { // Tries to load it from the environment variable
+            var dotenv = Dotenv
+                    .configure()
+                    .ignoreIfMissing()
+                    .ignoreIfMalformed()
+                    .load();
+            openAIKey = dotenv.get("OPEN_AI_API_KEY");
+        }
     }
 
+    /**
+     * Generates an embedding for the given input string.
+     *
+     * @param inputString Input string to generate the embedding for.
+     * @return An embedding object for the given text.
+     */
     public Embedding generateEmbedding(String inputString) {
-
-        // TODO: the embeddings function from the library uses an array as input. We are
-        // only using a length 1 array.
-        // Check if we should implement an array option.
+        /* TODO: the embeddings function from the library uses an array as input.
+            We are only using a length 1 array. Should implement an array option? */
         List<String> embeddingsInput = new ArrayList<>();
         embeddingsInput.add(inputString);
 
@@ -43,8 +65,27 @@ public Embedding generateEmbedding(String inputString) {
                 .input(embeddingsInput)
                 .build();
 
-        List<Double> embedding = service.createEmbeddings(embeddingRequest).getData().get(0).getEmbedding();
+        List<Double> vector = getService()
+                .createEmbeddings(embeddingRequest)
+                .getData()
+                .get(0)
+                .getEmbedding();
 
-        return new Embedding(embedding);
+        Map<String, String> metadata = new java.util.HashMap<>();
+        metadata.put("input", inputString);
+        metadata.put("source", "OpenAI");
+        metadata.put("model", embeddingOpenAiModel);
+        metadata.put("createdAt", java.time.LocalDateTime.now().toString());
+        return new Embedding(UUID.randomUUID(), vector, metadata);
+    }
+
+    private OpenAiService getService() {
+        if (cachedService == null) {
+            if (openAIKey == null) {
+                throw new EnvironmentVariableNotDeclaredException("Environment variable OPEN_AI_API_KEY is required.");
+            }
+            cachedService = new OpenAiService(openAIKey);
+        }
+        return cachedService;
     }
 }
diff --git a/modules/embeddingsgeneration/src/test/java/com/theagilemonkeys/ellmental/embeddingsgeneration/openai/OpenAIEmbeddingsModelTest.java b/modules/embeddingsgeneration/src/test/java/com/theagilemonkeys/ellmental/embeddingsgeneration/openai/OpenAIEmbeddingsModelTest.java
index eb0be69..affae64 100644
--- a/modules/embeddingsgeneration/src/test/java/com/theagilemonkeys/ellmental/embeddingsgeneration/openai/OpenAIEmbeddingsModelTest.java
+++ b/modules/embeddingsgeneration/src/test/java/com/theagilemonkeys/ellmental/embeddingsgeneration/openai/OpenAIEmbeddingsModelTest.java
@@ -15,8 +15,8 @@ public void testGenerateEmbedding(){
         Embedding embedding =  openAI.generateEmbedding("The Agile Monkeys rule!");
         TestValues testValues = new TestValues();
 
-        assertEquals(embedding.vector.size(), testValues.testGenerateEmbeddingExpectedValue.size());
-        assertArrayEquals(embedding.vector.toArray(), testValues.testGenerateEmbeddingExpectedValue.toArray());
+        assertEquals(embedding.vector().size(), testValues.testGenerateEmbeddingExpectedValue.size());
+        assertArrayEquals(embedding.vector().toArray(), testValues.testGenerateEmbeddingExpectedValue.toArray());
     }
 }
 
diff --git a/modules/embeddingsstore/src/main/java/com/theagilemonkeys/ellmental/embeddingsstore/EmbeddingsStore.java b/modules/embeddingsstore/src/main/java/com/theagilemonkeys/ellmental/embeddingsstore/EmbeddingsStore.java
index 7bb7977..66f7f49 100644
--- a/modules/embeddingsstore/src/main/java/com/theagilemonkeys/ellmental/embeddingsstore/EmbeddingsStore.java
+++ b/modules/embeddingsstore/src/main/java/com/theagilemonkeys/ellmental/embeddingsstore/EmbeddingsStore.java
@@ -5,8 +5,11 @@
 
 import com.theagilemonkeys.ellmental.core.schema.Embedding;
 
+/**
+ * Abstract class that defines the expected interface for a persistence mechanism capable of storing and querying embeddings.
+ */
 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, Map<String,String> metadata);
+    public abstract void store(Embedding embedding);
     public abstract List<Embedding> similaritySearch(Embedding reference, int limit);
 }
\ No newline at end of file
diff --git a/modules/embeddingsstore/src/main/java/com/theagilemonkeys/ellmental/embeddingsstore/pinecone/PineconeEmbeddingsStore.java b/modules/embeddingsstore/src/main/java/com/theagilemonkeys/ellmental/embeddingsstore/pinecone/PineconeEmbeddingsStore.java
index 75ea335..5af1b82 100644
--- a/modules/embeddingsstore/src/main/java/com/theagilemonkeys/ellmental/embeddingsstore/pinecone/PineconeEmbeddingsStore.java
+++ b/modules/embeddingsstore/src/main/java/com/theagilemonkeys/ellmental/embeddingsstore/pinecone/PineconeEmbeddingsStore.java
@@ -12,67 +12,53 @@
 import static java.net.HttpURLConnection.HTTP_BAD_REQUEST;
 import io.github.cdimascio.dotenv.Dotenv;
 
+/**
+ * Implementation for Pinecone EmbeddingsStore
+ */
 public class PineconeEmbeddingsStore extends EmbeddingsStore {
-
     private final String url;
     private final String apiKey;
     private final String namespace;
-    private static OkHttpClient client;
 
+    /**
+     * Constructor that initializes the Pinecone embeddings store loading the configuration from environment variables:
+     *  It will load the API key from the environment variable `PINECONE_API_KEY`.
+     *  It will load the URL from the environment variable `PINECONE_URL`.
+     *  It will load the namespace from the environment variable `PINECONE_NAMESPACE`.
+     */
     public PineconeEmbeddingsStore() {
-        var dotenv = Dotenv
-                .configure()
-                .ignoreIfMissing()
-                .ignoreIfMalformed()
-                .load();
-        url = dotenv.get("PINECONE_URL");
-        apiKey = dotenv.get("PINECONE_API_KEY");
-        namespace = dotenv.get("PINECONE_NAMESPACE");
-
-        if (url == null) {
-            throw new EnvironmentVariableNotDeclaredException("Environement variable PINECONE_URL is not declared.");
-        } else if (apiKey == null) {
-            throw new EnvironmentVariableNotDeclaredException(
-                    "Environement variable PINECONE_API_KEY is not declared.");
-        } else if (namespace == null) {
-            // TODO: need to test code using pinecone with namespace
-            throw new EnvironmentVariableNotDeclaredException(
-                    "Environement variable PINECONE_NAMESPACE is not declared.");
-        }
-
-        client = new OkHttpClient();
+        this(null, null, null);
     }
 
-    private String post(String path, String bodyString) throws IOException {
-        MediaType JSON = MediaType.parse("application/json; charset=utf-8");
-
-        RequestBody body = RequestBody.create(bodyString, JSON);
-
-        Request request = new Request.Builder()
-                .url(url + path)
-                .header("accept", "application/json")
-                .header("content-type", "application/json")
-                .header("Api-Key", apiKey)
-                .post(body)
-                .build();
-
-        try (Response response = client.newCall(request).execute()) {
-            if (response.code() >= HTTP_BAD_REQUEST) {
-                throw new IOException(url);
-            }
-
-            ResponseBody responseBody = response.body();
-
-            if (responseBody != null) {
-                return responseBody.string();
-            } else {
-                return null;
-            }
+    /**
+     * Constructor that initializes the Pinecone embeddings store with an explicit url, apiKey and namespace.
+     *
+     * @param url Pinecone URL.
+     * @param apiKey Pinecone API key.
+     * @param namespace Pinecone namespace.
+     */
+    public PineconeEmbeddingsStore(String url, String apiKey, String namespace) {
+        if (url == null || apiKey == null || namespace == null) {
+            var dotenv = Dotenv
+                    .configure()
+                    .ignoreIfMissing()
+                    .ignoreIfMalformed()
+                    .load();
+            if (url == null) url = dotenv.get("PINECONE_URL");
+            if (apiKey == null) apiKey = dotenv.get("PINECONE_API_KEY");
+            if (namespace == null) namespace = dotenv.get("PINECONE_NAMESPACE");
         }
+        this.url = url;
+        this.apiKey = apiKey;
+        this.namespace = namespace;
     }
 
-    public void store(Embedding embedding, Map<String, String> metadata) {
-        UpsertVectorSchema schema = new UpsertVectorSchema(embedding, metadata);
+    /**
+     * Stores an embedding in the embeddings store.
+     * @param embedding Embedding object to store.
+     */
+    public void store(Embedding embedding) {
+        UpsertVectorSchema schema = new UpsertVectorSchema(embedding);
         String requestBodyJson = new Gson().toJson(schema);
 
         try {
@@ -88,8 +74,8 @@ public List<Embedding> similaritySearch(Embedding reference, int limit) {
                 limit,
                 true,
                 true,
-                reference.vector,
-                null);
+                reference.vector(),
+                namespace);
         String requestBodyJson = new Gson().toJson(body);
 
         try {
@@ -97,8 +83,8 @@ public List<Embedding> similaritySearch(Embedding reference, int limit) {
             QueryVectorResponseSchema response = new Gson().fromJson(responseString, QueryVectorResponseSchema.class);
             ArrayList<Embedding> matchEmbeddings = new ArrayList<>();
 
-            for (Match m : response.matches) {
-                matchEmbeddings.add(new Embedding(m.values));
+            for (Match match : response.matches) { // TODO: Make sure this array is sorted by similarity using the score field
+                matchEmbeddings.add(new Embedding(match.id, match.values, match.metadata));
             }
 
             return matchEmbeddings;
@@ -107,4 +93,46 @@ public List<Embedding> similaritySearch(Embedding reference, int limit) {
             return null;
         }
     }
+
+    private boolean validateEnvironment() {
+        if (url == null) {
+            throw new EnvironmentVariableNotDeclaredException("Environment variable PINECONE_URL is not declared.");
+        } else if (apiKey == null) {
+            throw new EnvironmentVariableNotDeclaredException(
+                    "Environment variable PINECONE_API_KEY is not declared.");
+        }
+        // TODO: check if namespace is required. From Pinecone documentation this doesn't seem to be the case, so I removed the check here.
+        return true;
+    }
+
+    private String post(String path, String bodyString) throws IOException {
+        if (!validateEnvironment()) {
+            return null;
+        }
+        MediaType JSON = MediaType.parse("application/json; charset=utf-8");
+
+        RequestBody body = RequestBody.create(bodyString, JSON);
+
+        Request request = new Request.Builder()
+                .url(url + path)
+                .header("accept", "application/json")
+                .header("content-type", "application/json")
+                .header("Api-Key", apiKey)
+                .post(body)
+                .build();
+
+        try (Response response = new OkHttpClient().newCall(request).execute()) {
+            if (response.code() >= HTTP_BAD_REQUEST) {
+                throw new IOException(url);
+            }
+
+            ResponseBody responseBody = response.body();
+
+            if (responseBody != null) {
+                return responseBody.string();
+            } else {
+                return null;
+            }
+        }
+    }
 }
diff --git a/modules/embeddingsstore/src/main/java/com/theagilemonkeys/ellmental/embeddingsstore/pinecone/QueryVectorResponseSchema.java b/modules/embeddingsstore/src/main/java/com/theagilemonkeys/ellmental/embeddingsstore/pinecone/QueryVectorResponseSchema.java
index 684a45f..47b5ef5 100644
--- a/modules/embeddingsstore/src/main/java/com/theagilemonkeys/ellmental/embeddingsstore/pinecone/QueryVectorResponseSchema.java
+++ b/modules/embeddingsstore/src/main/java/com/theagilemonkeys/ellmental/embeddingsstore/pinecone/QueryVectorResponseSchema.java
@@ -2,6 +2,7 @@
 
 import java.util.List;
 import java.util.Map;
+import java.util.UUID;
 
 public class QueryVectorResponseSchema {
     String namespace;
@@ -9,7 +10,7 @@ public class QueryVectorResponseSchema {
 }
 
 class Match {
-    String id;
+    UUID id;
     Double score;
     List<Double> values;
     Map<String,String> metadata;
diff --git a/modules/embeddingsstore/src/main/java/com/theagilemonkeys/ellmental/embeddingsstore/pinecone/UpsertVectorSchema.java b/modules/embeddingsstore/src/main/java/com/theagilemonkeys/ellmental/embeddingsstore/pinecone/UpsertVectorSchema.java
index f0ddfae..5559adc 100644
--- a/modules/embeddingsstore/src/main/java/com/theagilemonkeys/ellmental/embeddingsstore/pinecone/UpsertVectorSchema.java
+++ b/modules/embeddingsstore/src/main/java/com/theagilemonkeys/ellmental/embeddingsstore/pinecone/UpsertVectorSchema.java
@@ -4,16 +4,16 @@
 import com.theagilemonkeys.ellmental.core.schema.Embedding;
 import java.util.List;
 import java.util.Map;
-import java.util.UUID;
 import java.util.ArrayList;
+import java.util.UUID;
 
 
 public class UpsertVectorSchema {
     List<Vector> vectors;
 
-    UpsertVectorSchema(Embedding embedding, Map<String,String> metadata ) {
+    UpsertVectorSchema(Embedding embedding) {
         vectors = new ArrayList<>();
-        vectors.add(new Vector(UUID.randomUUID().toString(), embedding.vector, metadata));
+        vectors.add(new Vector(embedding.id(), embedding.vector(), embedding.metadata()));
     }
 
 }
@@ -23,7 +23,7 @@ class SparseValues {
 }
 class  Vector {
 
-    Vector (String id, List<Double> values, Map<String,String> metadata) {
+    Vector (UUID id, List<Double> values, Map<String,String> metadata) {
         this.id = id;
         this.values = values;
         this.metadata = metadata;
@@ -32,5 +32,5 @@ class  Vector {
     public SparseValues sparseValues;
     public Map<String,String> metadata;
     public List<Double> values;
-    public String id;
+    public UUID id;
 }
diff --git a/modules/embeddingsstore/src/test/java/com/theagilemonkeys/ellmental/embeddingsstore/pinecone/PineconeEmbeddingsStoreTest.java b/modules/embeddingsstore/src/test/java/com/theagilemonkeys/ellmental/embeddingsstore/pinecone/PineconeEmbeddingsStoreTest.java
index c484fe4..d176196 100644
--- a/modules/embeddingsstore/src/test/java/com/theagilemonkeys/ellmental/embeddingsstore/pinecone/PineconeEmbeddingsStoreTest.java
+++ b/modules/embeddingsstore/src/test/java/com/theagilemonkeys/ellmental/embeddingsstore/pinecone/PineconeEmbeddingsStoreTest.java
@@ -9,6 +9,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.UUID;
 
 public class PineconeEmbeddingsStoreTest {
     //TODO: add mock reponse to improve store test
@@ -16,16 +17,17 @@ public class PineconeEmbeddingsStoreTest {
     public void testStore(){
         EmbeddingsStore embeddingStore = new PineconeEmbeddingsStore();
         TestValues testValues = new TestValues();
-        Embedding embedding = new Embedding(testValues.testGenerateEmbeddingExpectedValue);
 
         Map<String, String> metadata = new HashMap<>();
         metadata.put("key1", "value1");
         metadata.put("key2", "value2");
 
-        embeddingStore.store(embedding, metadata);
+        Embedding embedding = new Embedding(UUID.randomUUID(), testValues.testGenerateEmbeddingExpectedValue, metadata);
 
-        assertEquals(embedding.vector.size(), testValues.testGenerateEmbeddingExpectedValue.size());
-        assertArrayEquals(embedding.vector.toArray(), testValues.testGenerateEmbeddingExpectedValue.toArray());
+        embeddingStore.store(embedding);
+
+        assertEquals(embedding.vector().size(), testValues.testGenerateEmbeddingExpectedValue.size());
+        assertArrayEquals(embedding.vector().toArray(), testValues.testGenerateEmbeddingExpectedValue.toArray());
     }
 }