forked from TencentBlueKing/bk-repo
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
fix: 增加Milvus Http客户端 TencentBlueKing#2573
- Loading branch information
Showing
19 changed files
with
295 additions
and
34 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
59 changes: 59 additions & 0 deletions
59
...src/main/kotlin/com/tencent/bkrepo/job/batch/task/cache/preload/ai/milvus/MilvusClient.kt
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,59 @@ | ||
| package com.tencent.bkrepo.job.batch.task.cache.preload.ai.milvus | ||
|
|
||
| import com.tencent.bkrepo.common.api.constant.HttpHeaders | ||
| import com.tencent.bkrepo.common.api.constant.MediaTypes | ||
| import com.tencent.bkrepo.common.api.util.toJsonString | ||
| import com.tencent.bkrepo.common.storage.innercos.http.toRequestBody | ||
| import com.tencent.bkrepo.job.batch.task.cache.preload.ai.milvus.request.CreateCollectionReq | ||
| import okhttp3.MediaType.Companion.toMediaType | ||
| import okhttp3.OkHttpClient | ||
| import okhttp3.Request | ||
|
|
||
| class MilvusClient( | ||
| private val clientProperties: MilvusClientProperties | ||
| ) { | ||
| val client = OkHttpClient.Builder().addInterceptor { | ||
| val newReq = it.request().newBuilder() | ||
|
|
||
| // add token to header | ||
| with(clientProperties) { | ||
| val token = if (username.isNotEmpty() && password.isNotEmpty()) { | ||
| "Bearer ${username}:${password}" | ||
| } else if (!clientProperties.token.isNullOrEmpty()) { | ||
| "Bearer ${clientProperties.token}" | ||
| } else { | ||
| null | ||
| } | ||
| token?.let { newReq.header(HttpHeaders.AUTHORIZATION, it) } | ||
| } | ||
|
|
||
| it.proceed(newReq.build()) | ||
| }.build() | ||
|
|
||
| fun createCollection(req: CreateCollectionReq) { | ||
| val reqBody = req.toJsonString().toRequestBody(MediaTypes.APPLICATION_JSON.toMediaType()) | ||
| val request = Request.Builder() | ||
| .post(reqBody) | ||
| .url("${clientProperties.uri}/v2/vectordb/collections/create") | ||
| .build() | ||
| client.newCall(request).execute().use { response -> | ||
| if (!response.isSuccessful) { | ||
| throw RuntimeException("create collection[${req.collectionName}] failed") | ||
| } | ||
| } | ||
| } | ||
|
|
||
| fun collectionExists(collectionName: String): Boolean { | ||
| TODO() | ||
| } | ||
|
|
||
| fun dropCollection(collectionName: String) { | ||
| TODO() | ||
| } | ||
|
|
||
| fun search() {} | ||
|
|
||
| fun describeIndex(collectionName: String) {} | ||
|
|
||
| fun createIndex(collectionName: String) {} | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
40 changes: 40 additions & 0 deletions
40
...lin/com/tencent/bkrepo/job/batch/task/cache/preload/ai/milvus/request/CollectionSchema.kt
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,40 @@ | ||
| package com.tencent.bkrepo.job.batch.task.cache.preload.ai.milvus.request | ||
|
|
||
| import com.fasterxml.jackson.annotation.JsonProperty | ||
|
|
||
| data class CollectionSchema( | ||
| private val enableDynamicField: Boolean = java.lang.Boolean.TRUE, | ||
| private val fields: MutableList<FieldSchema> = ArrayList(), | ||
| @JsonProperty("autoID") | ||
| private val autoId: Boolean = false, | ||
| ) { | ||
| fun addField(fieldSchema: FieldSchema): CollectionSchema { | ||
| if (fieldSchema.dataType == DataType.Array.name) { | ||
| if (fieldSchema.elementDataType.isNullOrEmpty() || fieldSchema.elementTypeParams?.maxCapacity == null) { | ||
| throw IllegalArgumentException("Element type, maxCapacity are required for array field") | ||
| } | ||
| } | ||
|
|
||
| if (fieldSchema.dataType == DataType.FloatVector.name || | ||
| fieldSchema.dataType == DataType.BinaryVector.name || | ||
| fieldSchema.dataType == DataType.Float16Vector.name || | ||
| fieldSchema.dataType == DataType.BFloat16Vector.name | ||
| ) { | ||
| if (fieldSchema.elementTypeParams?.dim == null) { | ||
| throw IllegalArgumentException("Dimension is required for vector field") | ||
| } | ||
| } | ||
|
|
||
| fields.add(fieldSchema) | ||
| return this | ||
| } | ||
|
|
||
| fun getField(fieldName: String): FieldSchema? { | ||
| for (field in fields) { | ||
| if (field.fieldName == fieldName) { | ||
| return field | ||
| } | ||
| } | ||
| return null | ||
| } | ||
| } |
7 changes: 7 additions & 0 deletions
7
...lin/com/tencent/bkrepo/job/batch/task/cache/preload/ai/milvus/request/ConsistencyLevel.kt
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,7 @@ | ||
| package com.tencent.bkrepo.job.batch.task.cache.preload.ai.milvus.request | ||
|
|
||
| enum class ConsistencyLevel(private val code: Int) { | ||
| Strong(0), | ||
| Bounded(2), | ||
| Eventually(3), | ||
| } |
29 changes: 29 additions & 0 deletions
29
.../com/tencent/bkrepo/job/batch/task/cache/preload/ai/milvus/request/CreateCollectionReq.kt
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,29 @@ | ||
| package com.tencent.bkrepo.job.batch.task.cache.preload.ai.milvus.request | ||
|
|
||
| import com.fasterxml.jackson.annotation.JsonProperty | ||
|
|
||
|
|
||
| data class CreateCollectionReq( | ||
| val dbName: String, | ||
| val collectionName: String, | ||
| val dimension: Int, | ||
| val metricType: String = MetricType.COSINE.name, | ||
| val idType: DataType = DataType.Int64, | ||
| @JsonProperty("autoID") | ||
| val autoId: Boolean = false, | ||
| val primaryFieldName: String = "id", | ||
| val vectorFieldName: String = "vector", | ||
| val schema: CollectionSchema? = null, | ||
| val indexParams: List<IndexParam> = ArrayList(), | ||
| val params: Params? = null, | ||
| ) { | ||
| data class Params( | ||
| @JsonProperty("max_length") | ||
| val maxLength: Int? = null, | ||
| val enableDynamicField: Boolean? = null, | ||
| val shardsNum: Int? = null, | ||
| val consistencyLevel: String? = ConsistencyLevel.Bounded.name, | ||
| val partitionsNum: Int? = null, | ||
| val ttlSeconds: Int? = null, | ||
| ) | ||
| } |
24 changes: 24 additions & 0 deletions
24
...main/kotlin/com/tencent/bkrepo/job/batch/task/cache/preload/ai/milvus/request/DataType.kt
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,24 @@ | ||
| package com.tencent.bkrepo.job.batch.task.cache.preload.ai.milvus.request | ||
|
|
||
| enum class DataType(private val code: Int) { | ||
| None(0), | ||
| Bool(1), | ||
| Int8(2), | ||
| Int16(3), | ||
| Int32(4), | ||
| Int64(5), | ||
|
|
||
| Float(10), | ||
| Double(11), | ||
|
|
||
| String(20), | ||
| VarChar(21), // variable-length strings with a specified maximum length | ||
| Array(22), | ||
| JSON(23), | ||
|
|
||
| BinaryVector(100), | ||
| FloatVector(101), | ||
| Float16Vector(102), | ||
| BFloat16Vector(103), | ||
| SparseFloatVector(104) | ||
| } |
33 changes: 33 additions & 0 deletions
33
...n/kotlin/com/tencent/bkrepo/job/batch/task/cache/preload/ai/milvus/request/FieldSchema.kt
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,33 @@ | ||
| package com.tencent.bkrepo.job.batch.task.cache.preload.ai.milvus.request | ||
|
|
||
| import com.fasterxml.jackson.annotation.JsonProperty | ||
|
|
||
| data class FieldSchema ( | ||
| val fieldName: String, | ||
| val dataType: String, | ||
| val isPrimary: Boolean = false, | ||
| val isPartitionKey: Boolean = false, | ||
| /** | ||
| * An optional parameter for Array field values | ||
| */ | ||
| val elementDataType: String? = null, | ||
| val elementTypeParams: ElementTypeParams? = null | ||
| ) | ||
|
|
||
| data class ElementTypeParams( | ||
| /** | ||
| * An optional parameter for VarChar values that determines the maximum length of the value in the current field. | ||
| */ | ||
| @JsonProperty("max_length") | ||
| val maxLength: Int? = 65535, | ||
| /** | ||
| * An optional parameter for FloatVector or BinaryVector fields that determines the vector dimension. | ||
| */ | ||
| val dim: Int? = null, | ||
| /** | ||
| * An optional parameter for Array field values that | ||
| * determines the maximum number of elements in the current array field. | ||
| */ | ||
| @JsonProperty("max_capacity") | ||
| val maxCapacity: Int? = null, | ||
| ) |
28 changes: 28 additions & 0 deletions
28
...in/kotlin/com/tencent/bkrepo/job/batch/task/cache/preload/ai/milvus/request/IndexParam.kt
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,28 @@ | ||
| package com.tencent.bkrepo.job.batch.task.cache.preload.ai.milvus.request | ||
|
|
||
| import com.fasterxml.jackson.annotation.JsonProperty | ||
|
|
||
| data class IndexParam( | ||
| var fieldName: String, | ||
| val metricType: MetricType = MetricType.COSINE, | ||
| val indexName: String? = null, | ||
| val params: Map<String, Any>? = null, | ||
| ) | ||
|
|
||
| data class Params( | ||
| @JsonProperty("index_type") | ||
| val indexType: IndexType = IndexType.AUTOINDEX, | ||
| /** | ||
| * The maximum degree of the node and applies only when index_type is set to HNSW. | ||
| */ | ||
| @JsonProperty("M") | ||
| val m: Int, | ||
| /** | ||
| * The search scope. This applies only when index_type is set to HNSW | ||
| */ | ||
| val efConstruction: Int, | ||
| /** | ||
| * The number of cluster units. This applies to IVF-related index types. | ||
| */ | ||
| val nlist: Int, | ||
| ) |
Oops, something went wrong.