Browse Source

fix: 分页查询向量库数据

zhongwei
haungxing 3 months ago
parent
commit
cc903588ab
  1. 7
      hzims-service/hzims-big-model/src/main/java/com/hnac/hzims/bigmodel/database/controller/WeaviateController.java
  2. 33
      hzims-service/hzims-big-model/src/main/java/com/hnac/hzims/bigmodel/database/service/WeaviateService.java

7
hzims-service/hzims-big-model/src/main/java/com/hnac/hzims/bigmodel/database/controller/WeaviateController.java

@ -27,8 +27,11 @@ public class WeaviateController {
}
@GetMapping("/list")
public R<List<WeaviateObject>> list(@RequestParam(value = "id",required = false) String id, @RequestParam("className") String className) {
return R.data(weaviateService.list(id,className));
public R<List<WeaviateObject>> list(@RequestParam(value = "id",required = false) String id,
@RequestParam("className") String className,
@RequestParam("current") Integer current,
@RequestParam("pageSize") Integer pageSize) {
return R.data(weaviateService.list(id,className,current,pageSize));
}
@DeleteMapping("/removeById")

33
hzims-service/hzims-big-model/src/main/java/com/hnac/hzims/bigmodel/database/service/WeaviateService.java

@ -8,27 +8,34 @@ import com.alibaba.fastjson.JSONObject;
import com.google.common.collect.Lists;
import com.hnac.hzims.bigmodel.configuration.BigModelInvokeApi;
import com.hnac.hzinfo.exception.HzServiceException;
import io.weaviate.client.Config;
import io.weaviate.client.WeaviateAuthClient;
import io.weaviate.client.WeaviateClient;
import io.weaviate.client.base.Result;
import io.weaviate.client.v1.auth.exception.AuthException;
import io.weaviate.client.v1.data.api.ObjectCreator;
import io.weaviate.client.v1.data.api.ObjectDeleter;
import io.weaviate.client.v1.data.api.ObjectUpdater;
import io.weaviate.client.v1.data.api.ObjectsGetter;
import io.weaviate.client.v1.data.model.WeaviateObject;
import io.weaviate.client.v1.graphql.model.GraphQLResponse;
import io.weaviate.client.v1.graphql.query.Get;
import io.weaviate.client.v1.graphql.query.fields.Field;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springblade.core.tool.api.ResultCode;
import org.springblade.core.tool.utils.BeanUtil;
import org.springblade.core.tool.utils.Func;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;
import java.lang.reflect.Field;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.util.*;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
import java.util.stream.Stream;
/**
* @Author: huangxing
@ -122,16 +129,22 @@ public class WeaviateService {
* @param className 表名
* @return 删除结果
*/
public Boolean delete(String id,String className) {
public Boolean delete(String ids,String className) {
ObjectDeleter deleter = weaviateClient.data().deleter();
if(Func.isNotEmpty(id)) {
deleter.withID(id);
}
if(Func.isNotEmpty(className)) {
deleter.withClassName(className);
if(Func.isEmpty(ids)) {
Result<List<WeaviateObject>> allObject = weaviateClient.data().objectsGetter().withClassName(className).withLimit(10000).run();
if(!allObject.hasErrors()) {
ids = allObject.getResult().stream().map(WeaviateObject::getId).collect(Collectors.joining(","));
}
Result<Boolean> result = deleter.run();
return !result.hasErrors();
}
Func.toStrList(",",ids).forEach(id -> {
Result<Boolean> result = deleter.withID(id).run();
if(result.hasErrors()) {
throw new HzServiceException(ResultCode.FAILURE,id + "记录删除失败!");
}
});
return true;
}
/**
@ -157,7 +170,7 @@ public class WeaviateService {
return !result.hasErrors();
}
public List<WeaviateObject> list(String id,String className) {
public List<WeaviateObject> list(String id,String className,Integer current,Integer pageSize) {
ObjectsGetter objectsGetter = weaviateClient.data().objectsGetter();
if(Func.isNotEmpty(id)) {
objectsGetter.withID(id);
@ -165,7 +178,7 @@ public class WeaviateService {
if(Func.isNotEmpty(className)) {
objectsGetter.withClassName(className);
}
Result<List<WeaviateObject>> result = objectsGetter.run();
Result<List<WeaviateObject>> result = objectsGetter.withLimit(pageSize).withOffset((current-1) * pageSize).run();
if(result.hasErrors()) {
throw new HzServiceException("查询失败!");
}

Loading…
Cancel
Save