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") @GetMapping("/list")
public R<List<WeaviateObject>> list(@RequestParam(value = "id",required = false) String id, @RequestParam("className") String className) { public R<List<WeaviateObject>> list(@RequestParam(value = "id",required = false) String id,
return R.data(weaviateService.list(id,className)); @RequestParam("className") String className,
@RequestParam("current") Integer current,
@RequestParam("pageSize") Integer pageSize) {
return R.data(weaviateService.list(id,className,current,pageSize));
} }
@DeleteMapping("/removeById") @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.google.common.collect.Lists;
import com.hnac.hzims.bigmodel.configuration.BigModelInvokeApi; import com.hnac.hzims.bigmodel.configuration.BigModelInvokeApi;
import com.hnac.hzinfo.exception.HzServiceException; 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.WeaviateClient;
import io.weaviate.client.base.Result; 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.ObjectCreator;
import io.weaviate.client.v1.data.api.ObjectDeleter; import io.weaviate.client.v1.data.api.ObjectDeleter;
import io.weaviate.client.v1.data.api.ObjectUpdater; import io.weaviate.client.v1.data.api.ObjectUpdater;
import io.weaviate.client.v1.data.api.ObjectsGetter; import io.weaviate.client.v1.data.api.ObjectsGetter;
import io.weaviate.client.v1.data.model.WeaviateObject; 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.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j; 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.BeanUtil;
import org.springblade.core.tool.utils.Func; import org.springblade.core.tool.utils.Func;
import org.springframework.beans.factory.annotation.Value; import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service; import org.springframework.stereotype.Service;
import java.lang.reflect.Field;
import java.nio.ByteBuffer; import java.nio.ByteBuffer;
import java.nio.ByteOrder; import java.nio.ByteOrder;
import java.util.*; import java.util.*;
import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors; import java.util.stream.Collectors;
import java.util.stream.IntStream; import java.util.stream.IntStream;
import java.util.stream.Stream;
/** /**
* @Author: huangxing * @Author: huangxing
@ -122,16 +129,22 @@ public class WeaviateService {
* @param className 表名 * @param className 表名
* @return 删除结果 * @return 删除结果
*/ */
public Boolean delete(String id,String className) { public Boolean delete(String ids,String className) {
ObjectDeleter deleter = weaviateClient.data().deleter(); ObjectDeleter deleter = weaviateClient.data().deleter();
if(Func.isNotEmpty(id)) {
deleter.withID(id);
}
if(Func.isNotEmpty(className)) {
deleter.withClassName(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(); 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(); ObjectsGetter objectsGetter = weaviateClient.data().objectsGetter();
if(Func.isNotEmpty(id)) { if(Func.isNotEmpty(id)) {
objectsGetter.withID(id); objectsGetter.withID(id);
@ -165,7 +178,7 @@ public class WeaviateService {
if(Func.isNotEmpty(className)) { if(Func.isNotEmpty(className)) {
objectsGetter.withClassName(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()) { if(result.hasErrors()) {
throw new HzServiceException("查询失败!"); throw new HzServiceException("查询失败!");
} }

Loading…
Cancel
Save