Skip to content

feat(rag): Milvus support configure database name #978

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
guanxuc wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat(rag): Milvus support configure database name #978

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
124 changes: 115 additions & 9 deletions ...ntscope-extensions-rag-simple/src/main/java/io/agentscope/core/rag/store/MilvusStore.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,8 @@
import io.milvus.v2.service.collection.request.AddFieldReq;
import io.milvus.v2.service.collection.request.CreateCollectionReq;
import io.milvus.v2.service.collection.request.HasCollectionReq;
import io.milvus.v2.service.database.request.CreateDatabaseReq;
import io.milvus.v2.service.database.response.ListDatabasesResp;
import io.milvus.v2.service.vector.request.DeleteReq;
import io.milvus.v2.service.vector.request.InsertReq;
import io.milvus.v2.service.vector.request.SearchReq;
Expand Down Expand Up @@ -75,6 +77,7 @@
* // Using builder with authentication and timeout options (recommended)
* try (MilvusStore store = MilvusStore.builder()
* .uri("http://localhost:19530")
* .databaseName("my_database")
* .collectionName("my_collection")
* .dimensions(1024)
* .token("root:Milvus")
Expand Down Expand Up @@ -120,6 +123,7 @@ public class MilvusStore implements VDBStoreBase, AutoCloseable {

private final String uri;
private final String collectionName;
private final String databaseName;
private final int dimensions;
private final IndexParam.MetricType metricType;
private final MilvusClientV2 milvusClient;
Expand All @@ -137,6 +141,8 @@ private MilvusStore(Builder builder) throws VectorStoreException {
this.collectionName = builder.collectionName;
this.dimensions = builder.dimensions;
this.metricType = builder.metricType;
this.databaseName = builder.databaseName;
Map<String, String> databaseProperties = builder.databaseProperties;
String token = builder.token;
String username = builder.username;
String password = builder.password;
Expand Down Expand Up @@ -166,6 +172,12 @@ private MilvusStore(Builder builder) throws VectorStoreException {

log.debug("Initialized Milvus client: uri={}, collection={}", uri, collectionName);

// Ensure database exists
if (databaseName != null && !databaseName.trim().isEmpty()) {
ensureDataBase(tempClient, databaseProperties);
tempClient.useDatabase(databaseName);
}

// Ensure collection exists
ensureCollection(tempClient);

Expand All @@ -191,15 +203,22 @@ private MilvusStore(Builder builder) throws VectorStoreException {
/**
* Creates a new MilvusStore with minimal configuration.
*
* @param uri the Milvus server URI (e.g., "http://localhost:19530")
* @param uri the Milvus server URI (e.g., "http://localhost:19530")
* @param databaseName the name of the database to use
* @param collectionName the name of the collection to use
* @param dimensions the dimension of vectors that will be stored
* @param dimensions the dimension of vectors that will be stored
* @return a new MilvusStore instance
* @throws VectorStoreException if initialization fails
*/
public static MilvusStore create(String uri, String collectionName, int dimensions)
public static MilvusStore create(
String uri, String databaseName, String collectionName, int dimensions)
throws VectorStoreException {
return builder().uri(uri).collectionName(collectionName).dimensions(dimensions).build();
return builder()
.uri(uri)
.databaseName(databaseName)
.collectionName(collectionName)
.dimensions(dimensions)
.build();
}

@Override
Expand Down Expand Up @@ -354,7 +373,10 @@ private void ensureCollection(MilvusClientV2 client) throws VectorStoreException
try {
// Check if collection exists
HasCollectionReq hasCollectionReq =
HasCollectionReq.builder().collectionName(collectionName).build();
HasCollectionReq.builder()
.databaseName(databaseName)
.collectionName(collectionName)
.build();
boolean exists = client.hasCollection(hasCollectionReq);

if (exists) {
Expand All @@ -371,6 +393,32 @@ private void ensureCollection(MilvusClientV2 client) throws VectorStoreException
}
}

/**
* Ensures the database exists, creating it if necessary.
*
* @param client the MilvusClientV2 to use
* @throws VectorStoreException if database creation fails
*/
private void ensureDataBase(MilvusClientV2 client, Map<String, String> databaseProperties)
throws VectorStoreException {
try {
// Check if database exists
ListDatabasesResp listDatabasesResp = client.listDatabases();
boolean exists = listDatabasesResp.getDatabaseNames().contains(databaseName);
if (exists) {
log.debug("DatabaseName '{}' already exists", databaseName);
return;
}

// Create database
createDataBase(client, databaseProperties);
log.debug("Created database '{}'", databaseName);
} catch (Exception e) {
throw new VectorStoreException(
"Failed to ensure databaseName exists: " + databaseName, e);
}
}

/**
* Creates a new collection with the specified dimensions and schema.
*
Expand Down Expand Up @@ -435,6 +483,7 @@ private void createCollection(MilvusClientV2 client) {
// Create the collection with schema and index
CreateCollectionReq createCollectionReq =
CreateCollectionReq.builder()
.databaseName(databaseName)
.collectionName(collectionName)
.collectionSchema(schema)
.indexParams(indexParams)
Expand All @@ -443,6 +492,21 @@ private void createCollection(MilvusClientV2 client) {
client.createCollection(createCollectionReq);
}

/**
* Creates a new database with the specified properties.
*
* @param client the MilvusClientV2 to use
* @param databaseProperties the properties for the database
*/
private void createDataBase(MilvusClientV2 client, Map<String, String> databaseProperties) {
CreateDatabaseReq createDatabaseReq =
CreateDatabaseReq.builder()
.databaseName(databaseName)
.properties(databaseProperties)
.build();
client.createDatabase(createDatabaseReq);
}

/**
* Adds multiple documents to Milvus.
*
Expand Down Expand Up @@ -488,13 +552,20 @@ private void addDocumentsToMilvus(List<Document> documents) {
if (customPayload != null && !customPayload.isEmpty()) {
JsonObject payloadObject = GSON.toJsonTree(customPayload).getAsJsonObject();
row.add(FIELD_PAYLOAD, payloadObject);
} else {
row.add(FIELD_PAYLOAD, new JsonObject());
}

rows.add(row);
}

// Insert data
InsertReq insertReq = InsertReq.builder().collectionName(collectionName).data(rows).build();
InsertReq insertReq =
InsertReq.builder()
.databaseName(databaseName)
.collectionName(collectionName)
.data(rows)
.build();

InsertResp insertResp = milvusClient.insert(insertReq);
log.debug(
Expand All @@ -507,7 +578,7 @@ private void addDocumentsToMilvus(List<Document> documents) {
* Searches for similar documents in Milvus.
*
* @param queryEmbedding the query embedding vector
* @param limit the maximum number of results
* @param limit the maximum number of results
* @param scoreThreshold optional minimum score threshold
* @return a list of documents with scores set
*/
Expand All @@ -524,6 +595,7 @@ private List<Document> searchDocumentsInMilvus(
// Build search request
SearchReq.SearchReqBuilder searchBuilder =
SearchReq.builder()
.databaseName(databaseName)
.collectionName(collectionName)
.data(Collections.singletonList(queryVector))
.limit(limit)
Expand Down Expand Up @@ -573,7 +645,7 @@ private List<Document> searchDocumentsInMilvus(
* Reconstructs a Document from Milvus search result.
*
* @param result the search result from Milvus
* @param score the similarity score
* @param score the similarity score
* @return the reconstructed Document, or null if reconstruction fails
*/
private Document reconstructDocumentFromResult(SearchResp.SearchResult result, double score) {
Expand Down Expand Up @@ -631,6 +703,7 @@ private Document reconstructDocumentFromResult(SearchResp.SearchResult result, d
private boolean deleteDocumentFromMilvus(String docId) {
DeleteReq deleteReq =
DeleteReq.builder()
.databaseName(databaseName)
.collectionName(collectionName)
.filter(FIELD_DOC_ID + " in [\"" + docId + "\"]")
.build();
Expand Down Expand Up @@ -678,6 +751,15 @@ public IndexParam.MetricType getMetricType() {
return metricType;
}

/**
* Gets the database name.
*
* @return the database name
*/
public String getDatabaseName() {
return databaseName;
}

/**
* Closes the Milvus client and releases all resources.
*
Expand Down Expand Up @@ -757,6 +839,8 @@ public static class Builder {
private String password;
private long connectTimeoutMs = DEFAULT_CONNECT_TIMEOUT_MS;
private IndexParam.MetricType metricType = IndexParam.MetricType.COSINE;
private String databaseName;
private Map<String, String> databaseProperties;

private Builder() {}

Expand Down Expand Up @@ -858,6 +942,28 @@ public Builder metricType(IndexParam.MetricType metricType) {
return this;
}

/**
* Sets the database name.
*
* @param databaseName the database name
* @return this builder
*/
public Builder databaseName(String databaseName) {
this.databaseName = databaseName;
return this;
}

/**
* Sets the database properties.
*
* @param databaseProperties the database properties
* @return this builder
*/
public Builder databaseProperties(Map<String, String> databaseProperties) {
this.databaseProperties = databaseProperties;
return this;
}

/**
* Builds a new MilvusStore instance.
*
Expand All @@ -866,7 +972,7 @@ public Builder metricType(IndexParam.MetricType metricType) {
*
* @return a new MilvusStore instance
* @throws IllegalArgumentException if required parameters are invalid
* @throws VectorStoreException if client initialization or collection creation fails
* @throws VectorStoreException if client initialization or collection creation fails
*/
public MilvusStore build() throws VectorStoreException {
if (uri == null || uri.trim().isEmpty()) {
Expand Down
Loading
Loading