diff --git a/app-builder/plugins/data-mate-knowledge/pom.xml b/app-builder/plugins/data-mate-knowledge/pom.xml index 1a637780c..ce7358c12 100644 --- a/app-builder/plugins/data-mate-knowledge/pom.xml +++ b/app-builder/plugins/data-mate-knowledge/pom.xml @@ -139,4 +139,4 @@ - \ No newline at end of file + diff --git a/app-builder/plugins/data-mate-knowledge/src/main/java/modelengine/fit/jade/datamate/knowledge/dto/DataMateKnowledgeListQueryParam.java b/app-builder/plugins/data-mate-knowledge/src/main/java/modelengine/fit/jade/datamate/knowledge/dto/DataMateKnowledgeListQueryParam.java index 4574efd33..f5e53d1aa 100644 --- a/app-builder/plugins/data-mate-knowledge/src/main/java/modelengine/fit/jade/datamate/knowledge/dto/DataMateKnowledgeListQueryParam.java +++ b/app-builder/plugins/data-mate-knowledge/src/main/java/modelengine/fit/jade/datamate/knowledge/dto/DataMateKnowledgeListQueryParam.java @@ -21,7 +21,7 @@ @SerializeStrategy(include = SerializeStrategy.Include.NON_NULL) public class DataMateKnowledgeListQueryParam { /** - * 页码，从0开始。 + * 页码，从1开始。 */ private Integer page; @@ -40,4 +40,3 @@ public class DataMateKnowledgeListQueryParam { */ private String description; } - diff --git a/app-builder/plugins/data-mate-knowledge/src/main/java/modelengine/fit/jade/datamate/knowledge/external/DataMateKnowledgeBaseManager.java b/app-builder/plugins/data-mate-knowledge/src/main/java/modelengine/fit/jade/datamate/knowledge/external/DataMateKnowledgeBaseManager.java index acb7b2f8d..1f287004f 100644 --- a/app-builder/plugins/data-mate-knowledge/src/main/java/modelengine/fit/jade/datamate/knowledge/external/DataMateKnowledgeBaseManager.java +++ b/app-builder/plugins/data-mate-knowledge/src/main/java/modelengine/fit/jade/datamate/knowledge/external/DataMateKnowledgeBaseManager.java @@ -35,6 +35,8 @@ import modelengine.fitframework.util.MapBuilder; import modelengine.fitframework.util.ObjectUtils; import modelengine.fitframework.util.StringUtils; +import modelengine.jade.authentication.context.UserContext; +import modelengine.jade.authentication.context.UserContextHolder; import modelengine.jade.knowledge.code.KnowledgeManagerRetCode; import modelengine.jade.knowledge.exception.KnowledgeException; @@ -52,6 +54,8 @@ public class DataMateKnowledgeBaseManager { private static final Logger log = Logger.get(DataMateKnowledgeBaseManager.class); private static final String BEARER = "Bearer "; private static final String CONTENT_TYPE_JSON = "application/json"; + private static final String USER_HEADER = "User"; + private static final String RETRIEVE_USER = "admin"; /** 默认访问超时时间（秒）。 */ private static final int DEFAULT_TIMEOUT_SECONDS = 30; @@ -89,6 +93,7 @@ public DataMateKnowledgeListEntity listRepos(String apiKey, DataMateKnowledgeLis if (StringUtils.isNotEmpty(apiKey)) { request.headers().set(AUTHORIZATION, BEARER + apiKey); } + this.addCurrentUserHeader(request); try { Object object = this.httpClient.get().exchangeForEntity(request, Object.class); Map response = @@ -118,6 +123,7 @@ public DataMateRetrievalResult retrieve(String apiKey, DataMateRetrievalParam pa if (StringUtils.isNotEmpty(apiKey)) { request.headers().set(AUTHORIZATION, BEARER + apiKey); } + this.addUserHeader(request, RETRIEVE_USER); request.headers().set(CONTENT_TYPE, CONTENT_TYPE_JSON); try { Object object = this.httpClient.get().exchangeForEntity(request, Object.class); @@ -141,6 +147,20 @@ private KnowledgeException handleException(HttpClientResponseException ex) { return new KnowledgeException(retCode, ex, ex.getSimpleMessage()); } + private void addCurrentUserHeader(HttpClassicClientRequest request) { + UserContext context = UserContextHolder.get(); + String userName = context == null ? StringUtils.EMPTY : context.getName(); + this.addUserHeader(request, userName); + } + + private void addUserHeader(HttpClassicClientRequest request, String userName) { + log.info("DataMate request user: {}", StringUtils.isEmpty(userName) ? "" : userName); + if (StringUtils.isEmpty(userName)) { + return; + } + request.headers().set(USER_HEADER, userName); + } + private HttpClassicClient getHttpClient() { int timeoutMs = this.timeoutSeconds * 1000; Map custom = MapBuilder.get() @@ -155,4 +175,3 @@ private HttpClassicClient getHttpClient() { .build()); } } - diff --git a/app-builder/plugins/data-mate-knowledge/src/main/java/modelengine/fit/jade/datamate/knowledge/service/DataMateKnowledgeRepoServiceImpl.java b/app-builder/plugins/data-mate-knowledge/src/main/java/modelengine/fit/jade/datamate/knowledge/service/DataMateKnowledgeRepoServiceImpl.java index bfbe02d53..01fb5760b 100644 --- a/app-builder/plugins/data-mate-knowledge/src/main/java/modelengine/fit/jade/datamate/knowledge/service/DataMateKnowledgeRepoServiceImpl.java +++ b/app-builder/plugins/data-mate-knowledge/src/main/java/modelengine/fit/jade/datamate/knowledge/service/DataMateKnowledgeRepoServiceImpl.java @@ -132,7 +132,7 @@ public List retrieve(String apiKey, FlatKnowledgeOption optio } private PageVoKnowledgeList queryKnowledgeList(String apiKey, ListRepoQueryParam param) { - int page = Math.max(param.getPageIndex() - 1, 0); + int page = Math.max(param.getPageIndex(), 1); int size = param.getPageSize(); DataMateKnowledgeListEntity listEntity = this.executeQuery(apiKey, param.getRepoName(), page, size); List content = listEntity.getContent() == null diff --git a/app-builder/plugins/data-mate-knowledge/src/main/resources/application.yml b/app-builder/plugins/data-mate-knowledge/src/main/resources/application.yml index 420680ee6..e79d4b3a0 100644 --- a/app-builder/plugins/data-mate-knowledge/src/main/resources/application.yml +++ b/app-builder/plugins/data-mate-knowledge/src/main/resources/application.yml @@ -6,7 +6,6 @@ fit: # timeout: 访问超时时间（秒），默认 30，可按场景自定义 datamate: url: - list: 'http://datamate-gateway:8080/api/knowledge-base/list' - retrieve: 'http://datamate-gateway:8080/api/knowledge-base/retrieve' + list: 'http://datamate-backend-python:18000/api/knowledge-base/list' + retrieve: 'http://datamate-backend-python:18000/api/knowledge-base/retrieve' # timeout: 30 - diff --git a/app-builder/plugins/data-mate-knowledge/src/test/java/modelengine/fit/jade/datamate/knowledge/DataMateKnowledgeBaseManagerTest.java b/app-builder/plugins/data-mate-knowledge/src/test/java/modelengine/fit/jade/datamate/knowledge/DataMateKnowledgeBaseManagerTest.java index d87352438..360c46550 100644 --- a/app-builder/plugins/data-mate-knowledge/src/test/java/modelengine/fit/jade/datamate/knowledge/DataMateKnowledgeBaseManagerTest.java +++ b/app-builder/plugins/data-mate-knowledge/src/test/java/modelengine/fit/jade/datamate/knowledge/DataMateKnowledgeBaseManagerTest.java @@ -17,6 +17,8 @@ import modelengine.fit.jade.datamate.knowledge.entity.DataMateRetrievalChunksEntity; import modelengine.fit.jade.datamate.knowledge.entity.DataMateRetrievalResult; import modelengine.fit.jade.datamate.knowledge.external.DataMateKnowledgeBaseManager; +import modelengine.jade.authentication.context.UserContext; +import modelengine.jade.authentication.context.UserContextHolder; import modelengine.jade.knowledge.exception.KnowledgeException; import org.junit.jupiter.api.AfterEach; @@ -31,6 +33,7 @@ import java.lang.reflect.Field; import java.util.HashMap; import java.util.Map; +import java.util.concurrent.atomic.AtomicReference; /** * 表示 {@link DataMateKnowledgeBaseManager} 的测试集。 @@ -81,6 +84,28 @@ public void shouldOkWhenListRepo() { DataMateKnowledgeEntity::getDescription).containsExactly("2", "test2", "test2知识库"); } + @Test + @DisplayName("查询知识库列表时携带 User 请求头") + public void shouldAddUserHeaderWhenListRepo() { + DataMateKnowledgeListQueryParam param = DataMateKnowledgeListQueryParam.builder().name("user").build(); + DataMateKnowledgeListEntity entity = this.executeListWithContext(param, + new UserContext("test-user", "", "")); + + assertThat(entity.getContent()).hasSize(2); + assertThat(UserContextHolder.get()).isNull(); + } + + @Test + @DisplayName("用户名缺失时不添加 User 请求头") + public void shouldOmitUserHeaderWhenUserNameMissing() { + DataMateKnowledgeListQueryParam param = + DataMateKnowledgeListQueryParam.builder().name("missing-user").build(); + DataMateKnowledgeListEntity entity = this.executeListWithContext(param, + new UserContext("", "", "")); + + assertThat(entity.getContent()).hasSize(2); + } + @Test @DisplayName("查询知识库列表失败，抛出异常") public void shouldFailWhenListRepoThrowException() { @@ -100,6 +125,19 @@ public void shouldOkWhenRetrieve() { assertThat(result.getData().get(0).content()).isEqualTo("content1"); } + @Test + @DisplayName("检索知识库时固定携带 admin 用户") + public void shouldAddAdminUserHeaderWhenRetrieve() { + DataMateRetrievalParam param = DataMateRetrievalParam.builder().query("admin").build(); + AtomicReference result = new AtomicReference<>(); + UserContext context = new UserContext("test-user", "", ""); + + UserContextHolder.apply(context, () -> result.set(this.manager.retrieve(this.apiKey, param))); + + assertThat(result.get().getData()).hasSize(3); + assertThat(UserContextHolder.get()).isNull(); + } + @Test @DisplayName("检索知识库失败，抛出异常") public void shouldFailWhenRetrieveThrowException() { @@ -108,4 +146,11 @@ public void shouldFailWhenRetrieveThrowException() { .extracting("code") .isEqualTo(130703005); } + + private DataMateKnowledgeListEntity executeListWithContext(DataMateKnowledgeListQueryParam param, + UserContext context) { + AtomicReference entity = new AtomicReference<>(); + UserContextHolder.apply(context, () -> entity.set(this.manager.listRepos(this.apiKey, param))); + return entity.get(); + } } diff --git a/app-builder/plugins/data-mate-knowledge/src/test/java/modelengine/fit/jade/datamate/knowledge/MockedDataMateKnowledgeBaseInnerController.java b/app-builder/plugins/data-mate-knowledge/src/test/java/modelengine/fit/jade/datamate/knowledge/MockedDataMateKnowledgeBaseInnerController.java index ec9a75caf..01f6db844 100644 --- a/app-builder/plugins/data-mate-knowledge/src/test/java/modelengine/fit/jade/datamate/knowledge/MockedDataMateKnowledgeBaseInnerController.java +++ b/app-builder/plugins/data-mate-knowledge/src/test/java/modelengine/fit/jade/datamate/knowledge/MockedDataMateKnowledgeBaseInnerController.java @@ -6,12 +6,14 @@ package modelengine.fit.jade.datamate.knowledge; +import static modelengine.fit.http.protocol.MessageHeaderNames.AUTHORIZATION; import static modelengine.fitframework.util.IoUtils.content; import modelengine.fit.http.annotation.PostMapping; import modelengine.fit.http.annotation.RequestBody; import modelengine.fit.http.annotation.RequestMapping; import modelengine.fit.http.client.HttpClientException; +import modelengine.fit.http.server.HttpClassicServerRequest; import modelengine.fitframework.annotation.Component; import modelengine.fitframework.serialization.ObjectSerializer; import modelengine.fit.jade.datamate.knowledge.entity.DataMateResponse; @@ -28,6 +30,11 @@ @Component @RequestMapping(path = "/v2", group = "DataMate知识库内部接口打桩") public class MockedDataMateKnowledgeBaseInnerController { + private static final String USER_HEADER = "User"; + private static final String EXPECTED_AUTHORIZATION = "Bearer 123"; + private static final String EXPECTED_USER = "test-user"; + private static final String EXPECTED_RETRIEVE_USER = "admin"; + private final ObjectSerializer serializer; public MockedDataMateKnowledgeBaseInnerController(ObjectSerializer serializer) { @@ -35,22 +42,43 @@ public MockedDataMateKnowledgeBaseInnerController(ObjectSerializer serializer) { } @PostMapping(path = "/knowledgeBase") - public Map listRepos(@RequestBody MockedDataMateKnowledgeListQueryParam param) throws IOException { - if (param.getName().equals("error")) { + public Map listRepos(HttpClassicServerRequest request, + @RequestBody MockedDataMateKnowledgeListQueryParam param) throws IOException { + if ("error".equals(param.getName())) { throw new HttpClientException("error"); } + this.validateRequestContext(request, "user".equals(param.getName()) ? EXPECTED_USER : null); String resourceName = "/listRepoResult.json"; String jsonContent = content(DataMateResponse.class, resourceName); return serializer.deserialize(jsonContent, Map.class); } @PostMapping(path = "/knowledgebases/query") - public Map retrieve(@RequestBody MockedDataMateRetrievalParam param) throws IOException { - if (param.getQuery().equals("error")) { + public Map retrieve(HttpClassicServerRequest request, + @RequestBody MockedDataMateRetrievalParam param) throws IOException { + if ("error".equals(param.getQuery())) { throw new HttpClientException("error"); } + this.validateRequestContext(request, EXPECTED_RETRIEVE_USER); String resourceName = "/retrieveResult.json"; String jsonContent = content(DataMateResponse.class, resourceName); return serializer.deserialize(jsonContent, Map.class); } + + private void validateRequestContext(HttpClassicServerRequest request, String expectedUser) { + this.validate(EXPECTED_AUTHORIZATION.equals(request.headers().first(AUTHORIZATION).orElse(null)), + "The authorization header is incorrect."); + if (expectedUser != null) { + this.validate(expectedUser.equals(request.headers().first(USER_HEADER).orElse(null)), + "The User header is incorrect."); + return; + } + this.validate(request.headers().first(USER_HEADER).isEmpty(), "The User header should be absent."); + } + + private void validate(boolean expression, String message) { + if (!expression) { + throw new HttpClientException(message); + } + } } diff --git a/app-builder/plugins/data-mate-knowledge/src/test/java/modelengine/fit/jade/datamate/knowledge/MockedDataMateKnowledgeListQueryParam.java b/app-builder/plugins/data-mate-knowledge/src/test/java/modelengine/fit/jade/datamate/knowledge/MockedDataMateKnowledgeListQueryParam.java index 6ef86a3c5..820306536 100644 --- a/app-builder/plugins/data-mate-knowledge/src/test/java/modelengine/fit/jade/datamate/knowledge/MockedDataMateKnowledgeListQueryParam.java +++ b/app-builder/plugins/data-mate-knowledge/src/test/java/modelengine/fit/jade/datamate/knowledge/MockedDataMateKnowledgeListQueryParam.java @@ -19,7 +19,7 @@ @SerializeStrategy(include = SerializeStrategy.Include.NON_NULL) public class MockedDataMateKnowledgeListQueryParam { /** - * 页码，从0开始。 + * 页码，从1开始。 */ private Integer page; diff --git a/app-builder/plugins/data-mate-knowledge/src/test/java/modelengine/fit/jade/datamate/knowledge/service/DataMateKnowledgeRepoServiceImplTest.java b/app-builder/plugins/data-mate-knowledge/src/test/java/modelengine/fit/jade/datamate/knowledge/service/DataMateKnowledgeRepoServiceImplTest.java index ac6a83a7b..db9a27b7a 100644 --- a/app-builder/plugins/data-mate-knowledge/src/test/java/modelengine/fit/jade/datamate/knowledge/service/DataMateKnowledgeRepoServiceImplTest.java +++ b/app-builder/plugins/data-mate-knowledge/src/test/java/modelengine/fit/jade/datamate/knowledge/service/DataMateKnowledgeRepoServiceImplTest.java @@ -10,9 +10,12 @@ import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.clearInvocations; import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import modelengine.fit.jade.datamate.knowledge.dto.DataMateKnowledgeListQueryParam; import modelengine.fit.jade.datamate.knowledge.entity.DataMateKnowledgeEntity; import modelengine.fit.jade.datamate.knowledge.entity.DataMateKnowledgeListEntity; import modelengine.fit.jade.datamate.knowledge.entity.DataMateRetrievalChunksEntity; @@ -36,6 +39,7 @@ import org.junit.jupiter.api.DisplayName; import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; import java.util.Collections; import java.util.List; @@ -70,7 +74,7 @@ void shouldReturnPageWhenListRepos() { .build(); DataMateKnowledgeListEntity listEntity = DataMateKnowledgeListEntity.builder() - .page(0) + .page(1) .size(10) .totalElements(1) .totalPages(1) @@ -83,6 +87,7 @@ void shouldReturnPageWhenListRepos() { param.setRepoName("test"); when(knowledgeBaseManager.listRepos(anyString(), any())).thenReturn(listEntity); + clearInvocations(knowledgeBaseManager); PageVo result = knowledgeRepoService.listRepos(API_KEY, param); @@ -91,13 +96,17 @@ void shouldReturnPageWhenListRepos() { assertThat(result.getItems().get(0)) .extracting(KnowledgeRepo::id, KnowledgeRepo::name, KnowledgeRepo::description, KnowledgeRepo::type) .containsExactly("kb-1", "测试知识库", "描述", "embedding-v1"); + ArgumentCaptor queryParamCaptor = + ArgumentCaptor.forClass(DataMateKnowledgeListQueryParam.class); + verify(knowledgeBaseManager).listRepos(anyString(), queryParamCaptor.capture()); + assertThat(queryParamCaptor.getValue().getPage()).isEqualTo(1); } @Test @DisplayName("listRepos - 列表为空时返回空分页") void shouldReturnEmptyPageWhenContentEmpty() { DataMateKnowledgeListEntity listEntity = DataMateKnowledgeListEntity.builder() - .page(0) + .page(1) .size(10) .totalElements(0) .totalPages(0) @@ -120,7 +129,7 @@ void shouldReturnEmptyPageWhenContentEmpty() { @DisplayName("listRepos - content 为 null 时按空列表处理") void shouldReturnEmptyPageWhenContentNull() { DataMateKnowledgeListEntity listEntity = DataMateKnowledgeListEntity.builder() - .page(0) + .page(1) .size(10) .totalElements(0) .totalPages(0) @@ -143,7 +152,7 @@ void shouldReturnEmptyPageWhenContentNull() { @DisplayName("listRepos - totalElements 为 null 时 total 为 0") void shouldUseZeroTotalWhenTotalElementsNull() { DataMateKnowledgeListEntity listEntity = DataMateKnowledgeListEntity.builder() - .page(0) + .page(1) .size(10) .totalElements(null) .totalPages(null)