content: /*
   * Copyright 2024-2025 the original author or authors.
   *
   * Licensed under the Apache License, Version 2.0 (the "License");
   * you may not use this file except in compliance with the License.
   * You may obtain a copy of the License at
   *
   *      https://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  package com.alibaba.cloud.ai.graph.node;
  
  import com.alibaba.cloud.ai.dashscope.api.DashScopeApi;
  import com.alibaba.cloud.ai.dashscope.embedding.DashScopeEmbeddingModel;
  import com.alibaba.cloud.ai.dashscope.embedding.DashScopeEmbeddingOptions;
  import com.alibaba.cloud.ai.dashscope.rerank.DashScopeRerankModel;
  import com.alibaba.cloud.ai.dashscope.rerank.DashScopeRerankOptions;
  import com.alibaba.cloud.ai.graph.OverAllState;
  import com.alibaba.cloud.ai.model.RerankModel;
  
  import org.springframework.ai.document.Document;
  import org.springframework.ai.document.MetadataMode;
  import org.springframework.ai.embedding.EmbeddingModel;
  import org.springframework.ai.vectorstore.SimpleVectorStore;
  import org.springframework.ai.vectorstore.filter.Filter;
  import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;
  
  import java.util.HashMap;
  import java.util.List;
  import java.util.Map;
  
  import org.junit.jupiter.api.Test;
  import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import static org.junit.jupiter.api.Assertions.assertEquals;
  
  public class KnowledgeRetrievalNodeTest {
  
  	private static final Logger logger = LoggerFactory.getLogger(KnowledgeRetrievalNode.class);
  
  	List<Document> documents = List.of(new Document(
  			"产品说明书:产品名称：智能机器人\n" + "产品描述：智能机器人是一个智能设备，能够自动完成各种任务。\n" + "功能：\n" + "1. 自动导航：机器人能够自动导航到指定位置。\n"
  					+ "2. 自动抓取：机器人能够自动抓取物品。\n" + "3. 自动放置：机器人能够自动放置物品。\n",
  			Map.of("type", "instruction", // 文档类型
  					"year", "2023", // 年份
  					"month", "06" // 月份
  			)),
  			new Document(
  					"产品说明书:产品名称：智能家居控制器\n" + "产品描述：智能家居控制器是一款集成化设备，可远程控制多种智能家电。\n" + "功能：\n"
  							+ "1. 远程控制：通过手机APP远程控制家电开关和调节。\n" + "2. 定时任务：设置家电定时开启或关闭。\n" + "3. 场景模式：支持多种场景模式一键切换。\n"
  							+ "4. 能耗统计：实时监控并统计家电能耗数据。\n",
  
  					Map.of("type", "instruction", // 文档类型
  							"year", "2024", // 年份
  							"month", "02" // 月份
  
  					)));
  
  	String apiKey = System.getenv().getOrDefault("AI_DASHSCOPE_API_KEY", "test-api-key");
  
  	DashScopeApi dashScopeApi = DashScopeApi.builder().apiKey(apiKey).build();
  
  	;
  
  	EmbeddingModel embeddingModel = new DashScopeEmbeddingModel(dashScopeApi, MetadataMode.EMBED,
  			DashScopeEmbeddingOptions.builder().withModel("text-embedding-v2").build());
  
  	SimpleVectorStore simpleVectorStore = SimpleVectorStore.builder(embeddingModel).build();
  
  	RerankModel rerankModel = new DashScopeRerankModel(dashScopeApi);
  
  	Filter.Expression filterExpression = new FilterExpressionBuilder().eq("type", "instruction").build();
  
  	DashScopeRerankOptions rerankOptions = new DashScopeRerankOptions();
  
  	Map<String, Object> initStateMap() {
  		Map<String, Object> modifiableMap = new HashMap<>();
  		modifiableMap.put("user_prompt", "你将作为一名机器人产品的专家，对于用户的使用需求作出解答");
  		modifiableMap.put("top_k", 5);
  		modifiableMap.put("similarity_threshold", 0.1);
  		modifiableMap.put("filter_expression", filterExpression);
  		modifiableMap.put("enable_ranker", true);
  		modifiableMap.put("rerank_model", rerankModel);
  		modifiableMap.put("rerank_options", rerankOptions);
  		modifiableMap.put("vector_store", simpleVectorStore);
  		return modifiableMap;
  	}
  
  	KnowledgeRetrievalNode.Builder initNodeBuilder() {
  		return KnowledgeRetrievalNode.builder()
  			.userPromptKey("user_prompt")
  			.topKKey("top_k")
  			.similarityThresholdKey("similarity_threshold")
  			.filterExpressionKey("filter_expression")
  			.enableRankerKey("enable_ranker")
  			.rerankModelKey("rerank_model")
  			.rerankOptionsKey("rerank_options")
  			.vectorStoreKey("vector_store");
  	}
  
  	@Test
  	@EnabledIfEnvironmentVariable(named = "AI_DASHSCOPE_API_KEY", matches = ".+")
  	void testValueFirst() throws Exception {
  		simpleVectorStore.add(documents);
  
  		KnowledgeRetrievalNode node = initNodeBuilder().topK(5).isKeyFirst(false).build();
  		Map<String, Object> stateMap = initStateMap();
  		// 修改topk
  		stateMap.put("top_k", 1);
  		Map<String, Object> newState = node.apply(new OverAllState(stateMap));
  		logger.info("文档检索结果加入prompt为{}", newState.get("user_prompt"));
  		assertEquals(2, node.documents.size());
  	}
  
  	@Test
  	@EnabledIfEnvironmentVariable(named = "AI_DASHSCOPE_API_KEY", matches = ".+")
  	void testTopK() throws Exception {
  
  		simpleVectorStore.add(documents);
  
  		KnowledgeRetrievalNode node = initNodeBuilder().build();
  		Map<String, Object> stateMap = initStateMap();
  		// 原本topk为5
  		Map<String, Object> newState = node.apply(new OverAllState(stateMap));
  		logger.info("文档检索结果加入prompt为{}", newState.get("user_prompt"));
  		assertEquals(2, node.documents.size());
  		// 修改topk
  		stateMap.put("top_k", 1);
  		newState = node.apply(new OverAllState(stateMap));
  		logger.info("文档检索结果加入prompt为{}", newState.get("user_prompt"));
  		assertEquals(1, node.documents.size());
  
  	}
  
  	@Test
  	@EnabledIfEnvironmentVariable(named = "AI_DASHSCOPE_API_KEY", matches = ".+")
  	void testSimilarityThreshold() throws Exception {
  
  		simpleVectorStore.add(documents);
  
  		KnowledgeRetrievalNode node = initNodeBuilder().build();
  		Map<String, Object> stateMap = initStateMap();
  		// 原本similarity_threshold为0，1
  		Map<String, Object> newState = node.apply(new OverAllState(stateMap));
  		logger.info("文档检索结果加入prompt为{}", newState.get("user_prompt"));
  		assertEquals(2, node.documents.size());
  		// 修改为0.5
  		stateMap.put("similarity_threshold", 0.5);
  		newState = node.apply(new OverAllState(stateMap));
  		logger.info("文档检索结果加入prompt为{}", newState.get("user_prompt"));
  		assertEquals(1, node.documents.size());
  
  	}
  
  	@Test
  	@EnabledIfEnvironmentVariable(named = "AI_DASHSCOPE_API_KEY", matches = ".+")
  	void testFilterExpression() throws Exception {
  
  		simpleVectorStore.add(documents);
  
  		KnowledgeRetrievalNode node = initNodeBuilder().build();
  		Map<String, Object> stateMap = initStateMap();
  		// 原本筛选条件是eq("type", "instruction")
  		Map<String, Object> newState = node.apply(new OverAllState(stateMap));
  		logger.info("文档检索结果加入prompt为{}", newState.get("user_prompt"));
  		assertEquals(2, node.documents.size());
  		// 现修改为eq("type", "book")
  		stateMap.put("filter_expression", new FilterExpressionBuilder().eq("type", "book").build());
  		newState = node.apply(new OverAllState(stateMap));
  		logger.info("文档检索结果加入prompt为{}", newState.get("user_prompt"));
  		assertEquals(0, node.documents.size());
  
  	}
  
  	@Test
  	@EnabledIfEnvironmentVariable(named = "AI_DASHSCOPE_API_KEY", matches = ".+")
  	void testRerank() throws Exception {
  
  		simpleVectorStore.add(documents);
  
  		KnowledgeRetrievalNode node = initNodeBuilder().build();
  		Map<String, Object> stateMap = initStateMap();
  
  		// rerankOptions原本默认topN为3
  		Map<String, Object> newState = node.apply(new OverAllState(stateMap));
  		logger.info("文档检索结果加入prompt为{}", newState.get("user_prompt"));
  		assertEquals(2, node.documents.size());
  		// 重设topN数量
  		rerankOptions.setTopN(1);
  		stateMap.put("rerank_options", rerankOptions);
  		newState = node.apply(new OverAllState(stateMap));
  		logger.info("文档检索结果加入prompt为{}", newState.get("user_prompt"));
  		assertEquals(1, node.documents.size());
  		// 使重排序失效
  		stateMap.put("enable_ranker", false);
  		newState = node.apply(new OverAllState(stateMap));
  		logger.info("文档检索结果加入prompt为{}", newState.get("user_prompt"));
  		assertEquals(2, node.documents.size());
  
  	}
  
  }
  