111 lines
4.0 KiB
Java
111 lines
4.0 KiB
Java
package com.gxwebsoft.ai.util;
|
||
|
||
import com.aliyun.bailian20231229.models.RetrieveRequest;
|
||
import com.aliyun.bailian20231229.models.RetrieveResponse;
|
||
import com.aliyun.teautil.models.RuntimeOptions;
|
||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||
|
||
import java.util.*;
|
||
|
||
/**
|
||
* 检索知识库
|
||
* @author GIIT-YC
|
||
*
|
||
*/
|
||
public class KnowledgeBaseRetrieve {
|
||
|
||
static String ALIBABA_CLOUD_ACCESS_KEY_ID = "LTAI5tD5YRKuxWz6Eg7qrM4P";
|
||
static String ALIBABA_CLOUD_ACCESS_KEY_SECRET = "bO8TBDXflOwbtSKimPpG8XrJnyzgTk";
|
||
static String WORKSPACE_ID = "llm-4pf5auwewoz34zqu";
|
||
|
||
|
||
/**
|
||
* 检查并提示设置必要的环境变量。
|
||
*
|
||
* @return true 如果所有必需的环境变量都已设置，否则 false
|
||
*/
|
||
public static boolean checkEnvironmentVariables() {
|
||
Map<String, String> requiredVars = new HashMap<>();
|
||
requiredVars.put("ALIBABA_CLOUD_ACCESS_KEY_ID", "阿里云访问密钥ID");
|
||
requiredVars.put("ALIBABA_CLOUD_ACCESS_KEY_SECRET", "阿里云访问密钥密码");
|
||
requiredVars.put("WORKSPACE_ID", "阿里云百炼业务空间ID");
|
||
|
||
List<String> missingVars = new ArrayList<>();
|
||
for (Map.Entry<String, String> entry : requiredVars.entrySet()) {
|
||
String value = System.getenv(entry.getKey());
|
||
if (value == null || value.isEmpty()) {
|
||
missingVars.add(entry.getKey());
|
||
System.out.println("错误：请设置 " + entry.getKey() + " 环境变量 (" + entry.getValue() + ")");
|
||
}
|
||
}
|
||
|
||
return missingVars.isEmpty();
|
||
}
|
||
|
||
/**
|
||
* 初始化客户端（Client）。
|
||
*
|
||
* @return 配置好的客户端对象
|
||
*/
|
||
public static com.aliyun.bailian20231229.Client createClient() throws Exception {
|
||
com.aliyun.teaopenapi.models.Config config = new com.aliyun.teaopenapi.models.Config()
|
||
.setAccessKeyId(ALIBABA_CLOUD_ACCESS_KEY_ID)
|
||
.setAccessKeySecret(ALIBABA_CLOUD_ACCESS_KEY_SECRET);
|
||
// 下方接入地址以公有云的公网接入地址为例，可按需更换接入地址。
|
||
config.endpoint = "bailian.cn-beijing.aliyuncs.com";
|
||
return new com.aliyun.bailian20231229.Client(config);
|
||
}
|
||
|
||
/**
|
||
* 在指定的知识库中检索信息。
|
||
*
|
||
* @param client 客户端对象（bailian20231229Client）
|
||
* @param workspaceId 业务空间ID
|
||
* @param indexId 知识库ID
|
||
* @param query 检索查询语句
|
||
* @return 阿里云百炼服务的响应
|
||
*/
|
||
public static RetrieveResponse retrieveIndex(com.aliyun.bailian20231229.Client client, String workspaceId,
|
||
String indexId, String query) throws Exception {
|
||
RetrieveRequest retrieveRequest = new RetrieveRequest();
|
||
retrieveRequest.setIndexId(indexId);
|
||
retrieveRequest.setQuery(query);
|
||
retrieveRequest.setDenseSimilarityTopK(null);
|
||
RuntimeOptions runtime = new RuntimeOptions();
|
||
return client.retrieveWithOptions(workspaceId, retrieveRequest, null, runtime);
|
||
}
|
||
|
||
/**
|
||
* 使用阿里云百炼服务检索知识库。
|
||
*/
|
||
public static void main(String[] args) {
|
||
// if (!checkEnvironmentVariables()) {
|
||
// System.out.println("环境变量校验未通过。");
|
||
// return;
|
||
// }
|
||
|
||
try {
|
||
// 步骤1：初始化客户端（Client）
|
||
System.out.println("步骤1：创建Client");
|
||
com.aliyun.bailian20231229.Client client = createClient();
|
||
|
||
// 步骤2：检索知识库
|
||
System.out.println("步骤2：检索知识库");
|
||
Scanner scanner = new Scanner(System.in);
|
||
System.out.print("请输入知识库ID："); // 即 CreateIndex 接口返回的 Data.Id，您也可以在阿里云百炼控制台的知识库页面获取。
|
||
String indexId = scanner.nextLine();
|
||
System.out.print("请输入检索query：");
|
||
String query = scanner.nextLine();
|
||
String workspaceId = WORKSPACE_ID;
|
||
RetrieveResponse resp = retrieveIndex(client, workspaceId, indexId, query);
|
||
|
||
// 请自行安装jackson-databind。将响应体responsebody转换为 JSON 字符串
|
||
ObjectMapper mapper = new ObjectMapper();
|
||
String result = mapper.writeValueAsString(resp.getBody());
|
||
System.out.println(result);
|
||
} catch (Exception e) {
|
||
System.out.println("发生错误：" + e.getMessage());
|
||
}
|
||
}
|
||
}
|