GH-1826 Fix EmbeddingModel's usage on Document#embedding

- Since the Document object's reference to the `embedding` is deprecated and will be removed, the VectorStore implementations require a way to store the embedding of the corresponding Document objects

     - One way to fix this is, to have the EmbeddingModel#embed to return the embeddings in the same order as that of the Documents passed to it.

       - Since both the Document and embedding collections use the List object, their iteration operation will make sure to keep them in line with the same order.

       - A fix is required to preserve the order when batching strategy is applied.
	  - Updated the Javadoc for BatchingStrategy
          - Fixed the Document List order in TokenCountBatchingStrategy

    - Refactored the vector store implementations to update this change

Resolves #GH-1826
This commit is contained in:
Ilayaperumal Gopinathan
2024-11-25 22:21:11 +00:00
committed by Christian Tzolov
parent 6cfe5e79e8
commit ebd29e0959
24 changed files with 118 additions and 94 deletions

View File

@@ -66,12 +66,12 @@ class EmbeddingIT extends AbstractIT {
@Test
void embeddingBatchDocuments() throws Exception {
assertThat(this.embeddingModel).isNotNull();
List<float[]> embedded = this.embeddingModel.embed(
List<float[]> embeddings = this.embeddingModel.embed(
List.of(new Document("Hello world"), new Document("Hello Spring"), new Document("Hello Spring AI!")),
OpenAiEmbeddingOptions.builder().withModel(OpenAiApi.DEFAULT_EMBEDDING_MODEL).build(),
new TokenCountBatchingStrategy());
assertThat(embedded.size()).isEqualTo(3);
embedded.forEach(embedding -> assertThat(embedding.length).isEqualTo(this.embeddingModel.dimensions()));
assertThat(embeddings.size()).isEqualTo(3);
embeddings.forEach(embedding -> assertThat(embedding.length).isEqualTo(this.embeddingModel.dimensions()));
}
@Test