Adds an embedding store for Azure Cosmos DB for NoSQL (langchain4j#1115)

zouanchao · May 23, 2024 · 240d34d · 240d34d
1 parent 9e382d4
commit 240d34d
Show file tree

Hide file tree

Showing 9 changed files with 26 additions and 62 deletions.
diff --git a/langchain4j-azure-cosmos-no-sql/pom.xml → langchain4j-azure-cosmos-nosql/pom.xml b/langchain4j-azure-cosmos-no-sql/pom.xml → langchain4j-azure-cosmos-nosql/pom.xml
@@ -3,28 +3,30 @@
          xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
          xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
     <modelVersion>4.0.0</modelVersion>
+
     <parent>
         <groupId>dev.langchain4j</groupId>
         <artifactId>langchain4j-parent</artifactId>
         <version>0.31.0-SNAPSHOT</version>
         <relativePath>../langchain4j-parent/pom.xml</relativePath>
     </parent>
 
-    <artifactId>langchain4j-azure-cosmos-no-sql</artifactId>
-    <packaging>jar</packaging>
-
-    <name>LangChain4j :: Integration :: Azure CosmosDB No SQL</name>
+    <artifactId>langchain4j-azure-cosmos-nosql</artifactId>
+    <name>LangChain4j :: Integration :: Azure CosmosDB NoSQL</name>
 
     <dependencies>
+
         <dependency>
             <groupId>dev.langchain4j</groupId>
             <artifactId>langchain4j-core</artifactId>
         </dependency>
+
         <dependency>
             <groupId>com.azure</groupId>
             <artifactId>azure-cosmos</artifactId>
             <version>4.60.0</version>
         </dependency>
+
         <dependency>
             <groupId>org.projectlombok</groupId>
             <artifactId>lombok</artifactId>
@@ -67,19 +69,7 @@
             <artifactId>junit-jupiter</artifactId>
             <scope>test</scope>
         </dependency>
-    </dependencies>
 
-    <build>
-        <plugins>
-            <plugin>
-                <groupId>org.apache.maven.plugins</groupId>
-                <artifactId>maven-compiler-plugin</artifactId>
-                <configuration>
-                    <source>${maven.compiler.source}</source>
-                    <target>${maven.compiler.target}</target>
-                </configuration>
-            </plugin>
-        </plugins>
-    </build>
+    </dependencies>
 
 </project>
diff --git a/...os/no/sql/AzureCosmosDbNoSqlDocument.java → ...mos/nosql/AzureCosmosDbNoSqlDocument.java b/...os/no/sql/AzureCosmosDbNoSqlDocument.java → ...mos/nosql/AzureCosmosDbNoSqlDocument.java
@@ -1,4 +1,4 @@
-package dev.langchain4j.store.embedding.azure.cosmos.no.sql;
+package dev.langchain4j.store.embedding.azure.cosmos.nosql;
 
 import lombok.AllArgsConstructor;
 import lombok.Builder;

diff --git a/...sql/AzureCosmosDbNoSqlEmbeddingStore.java → ...sql/AzureCosmosDbNoSqlEmbeddingStore.java b/...sql/AzureCosmosDbNoSqlEmbeddingStore.java → ...sql/AzureCosmosDbNoSqlEmbeddingStore.java
@@ -1,15 +1,9 @@
-package dev.langchain4j.store.embedding.azure.cosmos.no.sql;
+package dev.langchain4j.store.embedding.azure.cosmos.nosql;
 
 import com.azure.cosmos.CosmosClient;
 import com.azure.cosmos.CosmosContainer;
 import com.azure.cosmos.CosmosDatabase;
-import com.azure.cosmos.models.CosmosBulkOperations;
-import com.azure.cosmos.models.CosmosContainerProperties;
-import com.azure.cosmos.models.CosmosItemOperation;
-import com.azure.cosmos.models.CosmosQueryRequestOptions;
-import com.azure.cosmos.models.CosmosVectorIndexSpec;
-import com.azure.cosmos.models.CosmosVectorEmbeddingPolicy;
-import com.azure.cosmos.models.PartitionKey;
+import com.azure.cosmos.models.*;
 import com.azure.cosmos.util.CosmosPagedIterable;
 import dev.langchain4j.data.embedding.Embedding;
 import dev.langchain4j.data.segment.TextSegment;
@@ -23,13 +17,15 @@
 import java.util.List;
 import java.util.stream.Collectors;
 
-import static dev.langchain4j.internal.Utils.isNullOrBlank;
-import static dev.langchain4j.internal.Utils.isNullOrEmpty;
-import static dev.langchain4j.internal.Utils.randomUUID;
+import static dev.langchain4j.internal.Utils.*;
 import static dev.langchain4j.internal.ValidationUtils.ensureTrue;
-import static dev.langchain4j.store.embedding.azure.cosmos.no.sql.MappingUtils.toNoSqlDbDocument;
+import static dev.langchain4j.store.embedding.azure.cosmos.nosql.MappingUtils.toNoSqlDbDocument;
 import static java.util.Collections.singletonList;
 
+/**
+ * You can read more about vector search using Azure Cosmos DB NoSQL
+ * <a href="https://aka.ms/CosmosVectorSearch">here</a>.
+ */
 public class AzureCosmosDbNoSqlEmbeddingStore implements EmbeddingStore<TextSegment> {
 
     private static final Logger log = LoggerFactory.getLogger(AzureCosmosDbNoSqlEmbeddingStore.class);
@@ -44,8 +40,6 @@ public class AzureCosmosDbNoSqlEmbeddingStore implements EmbeddingStore<TextSegm
     private final CosmosDatabase database;
     private final CosmosContainer container;
 
-    // You can read more about vector search using AzureCosmosDBNoSQL here
-    // https://aka.ms/CosmosVectorSearch
     @Builder
     public AzureCosmosDbNoSqlEmbeddingStore(CosmosClient cosmosClient,
                                             String databaseName,

diff --git a/...ql/AzureCosmosDbNoSqlMatchedDocument.java → ...ql/AzureCosmosDbNoSqlMatchedDocument.java b/...ql/AzureCosmosDbNoSqlMatchedDocument.java → ...ql/AzureCosmosDbNoSqlMatchedDocument.java
@@ -1,4 +1,4 @@
-package dev.langchain4j.store.embedding.azure.cosmos.no.sql;
+package dev.langchain4j.store.embedding.azure.cosmos.nosql;
 
 import lombok.AllArgsConstructor;
 import lombok.Data;

diff --git a/...ing/azure/cosmos/no/sql/MappingUtils.java → ...ding/azure/cosmos/nosql/MappingUtils.java b/...ing/azure/cosmos/no/sql/MappingUtils.java → ...ding/azure/cosmos/nosql/MappingUtils.java
@@ -1,4 +1,4 @@
-package dev.langchain4j.store.embedding.azure.cosmos.no.sql;
+package dev.langchain4j.store.embedding.azure.cosmos.nosql;
 
 import dev.langchain4j.data.document.Metadata;
 import dev.langchain4j.data.embedding.Embedding;

diff --git a/...l/AzureCosmosDbNoSqlEmbeddingStoreIT.java → ...l/AzureCosmosDbNoSqlEmbeddingStoreIT.java b/...l/AzureCosmosDbNoSqlEmbeddingStoreIT.java → ...l/AzureCosmosDbNoSqlEmbeddingStoreIT.java
@@ -1,20 +1,10 @@
-package dev.langchain4j.store.embedding.azure.cosmos.no.sql;
+package dev.langchain4j.store.embedding.azure.cosmos.nosql;
 
 import com.azure.cosmos.ConsistencyLevel;
 import com.azure.cosmos.CosmosClient;
 import com.azure.cosmos.CosmosClientBuilder;
 import com.azure.cosmos.CosmosDatabase;
-import com.azure.cosmos.models.CosmosContainerProperties;
-import com.azure.cosmos.models.CosmosVectorDataType;
-import com.azure.cosmos.models.CosmosVectorDistanceFunction;
-import com.azure.cosmos.models.CosmosVectorEmbedding;
-import com.azure.cosmos.models.CosmosVectorEmbeddingPolicy;
-import com.azure.cosmos.models.CosmosVectorIndexSpec;
-import com.azure.cosmos.models.CosmosVectorIndexType;
-import com.azure.cosmos.models.IncludedPath;
-import com.azure.cosmos.models.IndexingMode;
-import com.azure.cosmos.models.IndexingPolicy;
-import com.azure.cosmos.models.PartitionKeyDefinition;
+import com.azure.cosmos.models.*;
 import dev.langchain4j.data.embedding.Embedding;
 import dev.langchain4j.data.segment.TextSegment;
 import dev.langchain4j.model.embedding.AllMiniLmL6V2QuantizedEmbeddingModel;
@@ -36,7 +26,7 @@
 
 @EnabledIfEnvironmentVariable(named = "AZURE_COSMOS_HOST", matches = ".+")
 @EnabledIfEnvironmentVariable(named = "AZURE_COSMOS_MASTER_KEY", matches = ".+")
-public class AzureCosmosDbNoSqlEmbeddingStoreIT extends EmbeddingStoreIT {
+class AzureCosmosDbNoSqlEmbeddingStoreIT extends EmbeddingStoreIT {
 
     protected static Logger logger = LoggerFactory.getLogger(AzureCosmosDbNoSqlEmbeddingStoreIT.class);
 

diff --git a/...AzureCosmosDbNoSqlEmbeddingStoreTest.java → ...AzureCosmosDbNoSqlEmbeddingStoreTest.java b/...AzureCosmosDbNoSqlEmbeddingStoreTest.java → ...AzureCosmosDbNoSqlEmbeddingStoreTest.java
@@ -1,35 +1,25 @@
-package dev.langchain4j.store.embedding.azure.cosmos.no.sql;
+package dev.langchain4j.store.embedding.azure.cosmos.nosql;
 
 import com.azure.cosmos.ConsistencyLevel;
 import com.azure.cosmos.CosmosClient;
 import com.azure.cosmos.CosmosClientBuilder;
-import com.azure.cosmos.models.CosmosVectorDataType;
-import com.azure.cosmos.models.CosmosVectorDistanceFunction;
-import com.azure.cosmos.models.CosmosVectorEmbedding;
-import com.azure.cosmos.models.CosmosVectorEmbeddingPolicy;
-import com.azure.cosmos.models.CosmosVectorIndexSpec;
-
+import com.azure.cosmos.models.*;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
-
 @EnabledIfEnvironmentVariable(named = "AZURE_COSMOS_HOST", matches = ".+")
 @EnabledIfEnvironmentVariable(named = "AZURE_COSMOS_MASTER_KEY", matches = ".+")
-public class AzureCosmosDbNoSqlEmbeddingStoreTest {
+class AzureCosmosDbNoSqlEmbeddingStoreTest {
 
     private static final String DATABASE_NAME = "test_db";
     private static final String CONTAINER_NAME = "test_container";
 
-    private final static Logger logger = LoggerFactory.getLogger(AzureCosmosDbNoSqlEmbeddingStoreTest.class);
-
     @Test
     void should_fail_if_cosmosClient_missing() {
         assertThrows(IllegalArgumentException.class, () -> {

diff --git a/langchain4j-bom/pom.xml b/langchain4j-bom/pom.xml
@@ -247,7 +247,7 @@
 
             <dependency>
                 <groupId>dev.langchain4j</groupId>
-                <artifactId>langchain4j-azure-cosmos-no-sql</artifactId>
+                <artifactId>langchain4j-azure-cosmos-nosql</artifactId>
                 <version>${project.version}</version>
             </dependency>
 

diff --git a/pom.xml b/pom.xml
@@ -42,6 +42,7 @@
         <!-- embedding stores -->
         <module>langchain4j-azure-ai-search</module>
         <module>langchain4j-azure-cosmos-mongo-vcore</module>
+        <module>langchain4j-azure-cosmos-nosql</module>
         <module>langchain4j-cassandra</module>
         <module>langchain4j-chroma</module>
         <module>langchain4j-elasticsearch</module>
@@ -57,7 +58,6 @@
         <module>langchain4j-vearch</module>
         <module>langchain4j-vespa</module>
         <module>langchain4j-weaviate</module>
-        <module>langchain4j-azure-cosmos-no-sql</module>
 
         <!-- document loaders -->
         <module>document-loaders/langchain4j-document-loader-amazon-s3</module>