RagEngine/src/main/java/com/bw/search/service/impl/RagSearchServiceImpl.java

package com.bw.search.service.impl;

import java.util.Map;

import org.springframework.stereotype.Service;

import com.alibaba.fastjson.JSONObject;
import com.bw.search.cache.ConfigCache;
import com.bw.search.common.Res;
import com.bw.search.entity.Constants;
import com.bw.search.entity.SearchResponse;
import com.bw.search.service.RagSearchService;
import com.bw.search.utils.DateUtil;

import lombok.extern.slf4j.Slf4j;

@Service
@Slf4j
public class RagSearchServiceImpl implements RagSearchService {
	
	@Override
	public Res<?> search(String dataJson) {
		log.info("向量检索参数：{}",dataJson);
		//转换对象
		JSONObject parseObject = JSONObject.parseObject(dataJson);
		String id = parseObject.getString(Constants.ID);			
		Map<String, Object> knowResult = getKnowledge(id);
		if(knowResult == null) {
			log.error("向量检索失败！");
			return Res.fail("知识库获取失败！");
		}
		log.info("知识库结果已获取：{}",JSONObject.toJSONString(knowResult));
		//响应体数据
		SearchResponse searchResponse = new SearchResponse();
		searchResponse.setIds(JSONObject.parseArray((String)knowResult.get(Constants.IDS), String.class));
		return Res.ok(searchResponse);
	}

	
	
	/**
	 * 获取知识库结果
	 * 同步
	 * @param chatId
	 * @return
	 */
	private Map<String, Object> getKnowledge(String id) {
		Map<String, Object> knowResult = null;
		int retryTime = 30;
		while(ConfigCache.isStart) {
			if (retryTime > 0) {
				if(ConfigCache.searchResult.containsKey(id)) {
					knowResult = (Map<String, Object>) ConfigCache.searchResult.get(id);
					break;
				}else {
					log.info("知识获取中请稍后...{}",retryTime);
					DateUtil.sleep(1000);
				}
				retryTime --;
			}else {
				break;
			}
		}
		//获取之后删除
		ConfigCache.searchResult.remove(id);
		return knowResult;
	}
}