fix TestEmbeddingExtractor_Extract

pkg/extractors/embedder_test.go

@@ -76,11 +76,25 @@ func TestEmbeddingExtractor_Extract(t *testing.T) {
 
 	assert.Equal(t, len(expectedEmbeddingRecords), len(embeddedMessages))
 
-	// Test if the length of embeddedMessages is equal to the length of messageEvent.Messages
-	for i, r := range embeddedMessages {
-		assert.Equal(t, expectedEmbeddingRecords[i].TextUUID, r.TextUUID)
-		assert.Equal(t, expectedEmbeddingRecords[i].Text, r.Text)
-		compareFloat32Vectors(t, expectedEmbeddingRecords[i].Embedding, r.Embedding, 0.001)
+	expectedEmbeddingRecordsMap := make(map[string]models.MessageEmbedding)
+	for _, r := range expectedEmbeddingRecords {
+		expectedEmbeddingRecordsMap[r.TextUUID.String()] = r
+	}
+
+	embeddedMessagesMap := make(map[string]models.MessageEmbedding)
+	for _, r := range embeddedMessages {
+		embeddedMessagesMap[r.TextUUID.String()] = r
+	}
+
+	assert.Equal(t, len(expectedEmbeddingRecordsMap), len(embeddedMessagesMap))
+
+	for uuid, expectedRecord := range expectedEmbeddingRecordsMap {
+		actualRecord, ok := embeddedMessagesMap[uuid]
+		assert.True(t, ok)
+
+		assert.Equal(t, expectedRecord.TextUUID, actualRecord.TextUUID)
+		assert.Equal(t, expectedRecord.Text, actualRecord.Text)
+		compareFloat32Vectors(t, expectedRecord.Embedding, actualRecord.Embedding, 0.001)
 	}
 }
 

pkg/store/postgres/documents.go

@@ -268,6 +268,7 @@ func (dc *DocumentCollectionDAO) CreateDocuments(
 
 // UpdateDocuments updates the document_id, metadata, and embedding columns of the
 // given documents in the given collection. The documents must have non-nil uuids.
+//
 // **IMPORTANT:** We determine which columns to update based on the fields that are
 // non-zero in the given documents. This means that all documents must have data
 // for the same fields. If a document is missing data for a field, there could be