From 363b14f8db91c1c4cd87b68a95de59b4ea2a333f Mon Sep 17 00:00:00 2001 From: BorisTkachenko <35521895+BorisTkachenko@users.noreply.github.com> Date: Thu, 12 Dec 2024 09:27:29 +0100 Subject: [PATCH] [OPIK-548] List LLM Provider api keys (#864) --- .../java/com/comet/opik/OpikApplication.java | 5 +- .../com/comet/opik/api/ProviderApiKey.java | 17 +++ .../v1/priv/LlmProviderApiKeyResource.java | 34 ++++-- .../opik/domain/LlmProviderApiKeyDAO.java | 5 + .../opik/domain/LlmProviderApiKeyService.java | 33 ++++-- .../opik/infrastructure/EncryptionUtils.java | 25 +--- .../infrastructure/EncryptionUtilsModule.java | 11 -- .../LlmProviderApiKeyResourceClient.java | 101 ++++++++++++++++ .../priv/LlmProviderApiKeyResourceTest.java | 112 ++++++++---------- 9 files changed, 232 insertions(+), 111 deletions(-) delete mode 100644 apps/opik-backend/src/main/java/com/comet/opik/infrastructure/EncryptionUtilsModule.java create mode 100644 apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/resources/LlmProviderApiKeyResourceClient.java diff --git a/apps/opik-backend/src/main/java/com/comet/opik/OpikApplication.java b/apps/opik-backend/src/main/java/com/comet/opik/OpikApplication.java index 1ad8300518..9d2da4d1fd 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/OpikApplication.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/OpikApplication.java @@ -2,7 +2,7 @@ import com.comet.opik.api.error.JsonInvalidFormatExceptionMapper; import com.comet.opik.infrastructure.ConfigurationModule; -import com.comet.opik.infrastructure.EncryptionUtilsModule; +import com.comet.opik.infrastructure.EncryptionUtils; import com.comet.opik.infrastructure.OpikConfiguration; import com.comet.opik.infrastructure.auth.AuthModule; import com.comet.opik.infrastructure.bi.BiModule; @@ -70,7 +70,7 @@ public void initialize(Bootstrap bootstrap) { .withPlugins(new SqlObjectPlugin(), new Jackson2Plugin())) .modules(new DatabaseAnalyticsModule(), new IdGeneratorModule(), new AuthModule(), new RedisModule(), new RateLimitModule(), new NameGeneratorModule(), new HttpModule(), new EventModule(), - new ConfigurationModule(), new BiModule(), new EncryptionUtilsModule()) + new ConfigurationModule(), new BiModule()) .installers(JobGuiceyInstaller.class) .listen(new OpikGuiceyLifecycleEventListener()) .enableAutoConfig() @@ -79,6 +79,7 @@ public void initialize(Bootstrap bootstrap) { @Override public void run(OpikConfiguration configuration, Environment environment) { + EncryptionUtils.setConfig(configuration); // Resources var jersey = environment.jersey(); diff --git a/apps/opik-backend/src/main/java/com/comet/opik/api/ProviderApiKey.java b/apps/opik-backend/src/main/java/com/comet/opik/api/ProviderApiKey.java index 187ea94d25..0cac56cd6e 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/api/ProviderApiKey.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/api/ProviderApiKey.java @@ -12,6 +12,7 @@ import lombok.NonNull; import java.time.Instant; +import java.util.List; import java.util.UUID; @Builder(toBuilder = true) @@ -27,6 +28,7 @@ public record ProviderApiKey( @JsonView({View.Public.class}) @Schema(accessMode = Schema.AccessMode.READ_ONLY) String createdBy, @JsonView({View.Public.class}) @Schema(accessMode = Schema.AccessMode.READ_ONLY) Instant lastUpdatedAt, @JsonView({View.Public.class}) @Schema(accessMode = Schema.AccessMode.READ_ONLY) String lastUpdatedBy){ + @Override public String toString() { return "ProviderApiKey{" + @@ -46,4 +48,19 @@ public static class Write { public static class Public { } } + + public record ProviderApiKeyPage( + @JsonView( { + Project.View.Public.class}) int page, + @JsonView({View.Public.class}) int size, + @JsonView({View.Public.class}) long total, + @JsonView({View.Public.class}) List content, + @JsonView({View.Public.class}) List sortableBy) + implements + com.comet.opik.api.Page{ + + public static ProviderApiKeyPage empty(int page) { + return new ProviderApiKeyPage(page, 0, 0, List.of(), List.of()); + } + } } diff --git a/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/LlmProviderApiKeyResource.java b/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/LlmProviderApiKeyResource.java index 0f88f18481..73e69d6717 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/LlmProviderApiKeyResource.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/LlmProviderApiKeyResource.java @@ -1,6 +1,8 @@ package com.comet.opik.api.resources.v1.priv; import com.codahale.metrics.annotation.Timed; +import com.comet.opik.api.Page; +import com.comet.opik.api.Project; import com.comet.opik.api.ProviderApiKey; import com.comet.opik.api.ProviderApiKeyUpdate; import com.comet.opik.api.error.ErrorMessage; @@ -46,21 +48,37 @@ public class LlmProviderApiKeyResource { private final @NonNull LlmProviderApiKeyService llmProviderApiKeyService; private final @NonNull Provider requestContext; + @GET + @Operation(operationId = "findLlmProviderKeys", summary = "Find LLM Provider's ApiKeys", description = "Find LLM Provider's ApiKeys", responses = { + @ApiResponse(responseCode = "200", description = "LLMProviderApiKey resource", content = @Content(schema = @Schema(implementation = Project.ProjectPage.class))) + }) + @JsonView({ProviderApiKey.View.Public.class}) + public Response find() { + + String workspaceId = requestContext.get().getWorkspaceId(); + + log.info("Find LLM Provider's ApiKeys for workspaceId '{}'", workspaceId); + Page providerApiKeyPage = llmProviderApiKeyService.find(workspaceId); + log.info("Found LLM Provider's ApiKeys for workspaceId '{}'", workspaceId); + + return Response.ok().entity(providerApiKeyPage).build(); + } + @GET @Path("{id}") @Operation(operationId = "getLlmProviderApiKeyById", summary = "Get LLM Provider's ApiKey by id", description = "Get LLM Provider's ApiKey by id", responses = { - @ApiResponse(responseCode = "200", description = "ProviderApiKey resource", content = @Content(schema = @Schema(implementation = ProviderApiKey.class))), + @ApiResponse(responseCode = "200", description = "LLMProviderApiKey resource", content = @Content(schema = @Schema(implementation = ProviderApiKey.class))), @ApiResponse(responseCode = "404", description = "Not found", content = @Content(schema = @Schema(implementation = ErrorMessage.class)))}) @JsonView({ProviderApiKey.View.Public.class}) public Response getById(@PathParam("id") UUID id) { String workspaceId = requestContext.get().getWorkspaceId(); - log.info("Getting Provider's ApiKey by id '{}' on workspace_id '{}'", id, workspaceId); + log.info("Getting LLM Provider's ApiKey by id '{}' on workspace_id '{}'", id, workspaceId); - ProviderApiKey providerApiKey = llmProviderApiKeyService.get(id, workspaceId); + ProviderApiKey providerApiKey = llmProviderApiKeyService.find(id, workspaceId); - log.info("Got Provider's ApiKey by id '{}' on workspace_id '{}'", id, workspaceId); + log.info("Got LLM Provider's ApiKey by id '{}' on workspace_id '{}'", id, workspaceId); return Response.ok().entity(providerApiKey).build(); } @@ -77,9 +95,9 @@ public Response saveApiKey( @Context UriInfo uriInfo) { String workspaceId = requestContext.get().getWorkspaceId(); String userName = requestContext.get().getUserName(); - log.info("Save api key for provider '{}', on workspace_id '{}'", providerApiKey.provider(), workspaceId); + log.info("Save api key for LLM provider '{}', on workspace_id '{}'", providerApiKey.provider(), workspaceId); var providerApiKeyId = llmProviderApiKeyService.saveApiKey(providerApiKey, userName, workspaceId).id(); - log.info("Saved api key for provider '{}', on workspace_id '{}'", providerApiKey.provider(), workspaceId); + log.info("Saved api key for LLM provider '{}', on workspace_id '{}'", providerApiKey.provider(), workspaceId); var uri = uriInfo.getAbsolutePathBuilder().path("/%s".formatted(providerApiKeyId)).build(); @@ -99,9 +117,9 @@ public Response updateApiKey(@PathParam("id") UUID id, String workspaceId = requestContext.get().getWorkspaceId(); String userName = requestContext.get().getUserName(); - log.info("Updating api key for provider with id '{}' on workspaceId '{}'", id, workspaceId); + log.info("Updating api key for LLM provider with id '{}' on workspaceId '{}'", id, workspaceId); llmProviderApiKeyService.updateApiKey(id, providerApiKeyUpdate, userName, workspaceId); - log.info("Updated api key for provider with id '{}' on workspaceId '{}'", id, workspaceId); + log.info("Updated api key for LLM provider with id '{}' on workspaceId '{}'", id, workspaceId); return Response.noContent().build(); } diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/LlmProviderApiKeyDAO.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/LlmProviderApiKeyDAO.java index eb03d0967a..b8de197a44 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/domain/LlmProviderApiKeyDAO.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/domain/LlmProviderApiKeyDAO.java @@ -9,6 +9,7 @@ import org.jdbi.v3.sqlobject.statement.SqlQuery; import org.jdbi.v3.sqlobject.statement.SqlUpdate; +import java.util.List; import java.util.Optional; import java.util.UUID; @@ -32,6 +33,10 @@ void update(@Bind("id") UUID id, @SqlQuery("SELECT * FROM llm_provider_api_key WHERE id = :id AND workspace_id = :workspaceId") ProviderApiKey findById(@Bind("id") UUID id, @Bind("workspaceId") String workspaceId); + @SqlQuery("SELECT * FROM llm_provider_api_key " + + " WHERE workspace_id = :workspaceId ") + List find(@Bind("workspaceId") String workspaceId); + default Optional fetch(UUID id, String workspaceId) { return Optional.ofNullable(findById(id, workspaceId)); } diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/LlmProviderApiKeyService.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/LlmProviderApiKeyService.java index 9996c49122..5b8bcac421 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/domain/LlmProviderApiKeyService.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/domain/LlmProviderApiKeyService.java @@ -1,5 +1,6 @@ package com.comet.opik.domain; +import com.comet.opik.api.Page; import com.comet.opik.api.ProviderApiKey; import com.comet.opik.api.ProviderApiKeyUpdate; import com.comet.opik.api.error.EntityAlreadyExistsException; @@ -12,6 +13,7 @@ import lombok.NonNull; import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; +import org.apache.commons.collections4.CollectionUtils; import org.jdbi.v3.core.statement.UnableToExecuteStatementException; import ru.vyarus.guicey.jdbi3.tx.TransactionTemplate; @@ -25,7 +27,8 @@ @ImplementedBy(LlmProviderApiKeyServiceImpl.class) public interface LlmProviderApiKeyService { - ProviderApiKey get(UUID id, String workspaceId); + ProviderApiKey find(UUID id, String workspaceId); + Page find(String workspaceId); ProviderApiKey saveApiKey(ProviderApiKey providerApiKey, String userName, String workspaceId); void updateApiKey(UUID id, ProviderApiKeyUpdate providerApiKeyUpdate, String userName, String workspaceId); } @@ -40,8 +43,7 @@ class LlmProviderApiKeyServiceImpl implements LlmProviderApiKeyService { private final @NonNull TransactionTemplate template; @Override - public ProviderApiKey get(UUID id, String workspaceId) { - log.info("Getting provider api key with id '{}', workspaceId '{}'", id, workspaceId); + public ProviderApiKey find(@NonNull UUID id, @NonNull String workspaceId) { ProviderApiKey providerApiKey = template.inTransaction(READ_ONLY, handle -> { @@ -49,14 +51,29 @@ public ProviderApiKey get(UUID id, String workspaceId) { return repository.fetch(id, workspaceId).orElseThrow(this::createNotFoundError); }); - log.info("Got provider api key with id '{}', workspaceId '{}'", id, workspaceId); return providerApiKey.toBuilder() .build(); } @Override - public ProviderApiKey saveApiKey(@NonNull ProviderApiKey providerApiKey, String userName, String workspaceId) { + public Page find(@NonNull String workspaceId) { + List providerApiKeys = template.inTransaction(READ_ONLY, handle -> { + var repository = handle.attach(LlmProviderApiKeyDAO.class); + return repository.find(workspaceId); + }); + + if (CollectionUtils.isEmpty(providerApiKeys)) { + return ProviderApiKey.ProviderApiKeyPage.empty(0); + } + + return new ProviderApiKey.ProviderApiKeyPage( + 0, providerApiKeys.size(), providerApiKeys.size(), + providerApiKeys, List.of()); + } + + @Override + public ProviderApiKey saveApiKey(@NonNull ProviderApiKey providerApiKey, @NonNull String userName, @NonNull String workspaceId) { UUID apiKeyId = idGenerator.generateId(); var newProviderApiKey = providerApiKey.toBuilder() @@ -74,7 +91,7 @@ public ProviderApiKey saveApiKey(@NonNull ProviderApiKey providerApiKey, String return newProviderApiKey; }); - return get(apiKeyId, workspaceId); + return find(apiKeyId, workspaceId); } catch (UnableToExecuteStatementException e) { if (e.getCause() instanceof SQLIntegrityConstraintViolationException) { throw newConflict(); @@ -85,8 +102,8 @@ public ProviderApiKey saveApiKey(@NonNull ProviderApiKey providerApiKey, String } @Override - public void updateApiKey(@NonNull UUID id, @NonNull ProviderApiKeyUpdate providerApiKeyUpdate, String userName, - String workspaceId) { + public void updateApiKey(@NonNull UUID id, @NonNull ProviderApiKeyUpdate providerApiKeyUpdate, @NonNull String userName, + @NonNull String workspaceId) { template.inTransaction(WRITE, handle -> { diff --git a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/EncryptionUtils.java b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/EncryptionUtils.java index 839ad16127..7bbb9910de 100644 --- a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/EncryptionUtils.java +++ b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/EncryptionUtils.java @@ -1,6 +1,6 @@ package com.comet.opik.infrastructure; -import jakarta.inject.Inject; +import lombok.NonNull; import javax.crypto.BadPaddingException; import javax.crypto.Cipher; @@ -16,29 +16,17 @@ public class EncryptionUtils { - private static void init() { - if (key != null) return; - synchronized (EncryptionUtils.class) { - if (key == null) { - byte[] keyBytes = config.getEncryption().getKey().getBytes(StandardCharsets.UTF_8); - key = new SecretKeySpec(keyBytes, ALGO); - } - } - } - private static final String ALGO = "AES"; private static final Base64.Encoder mimeEncoder = Base64.getMimeEncoder(); private static final Base64.Decoder mimeDecoder = Base64.getMimeDecoder(); - private static OpikConfiguration config; private static Key key; - @Inject - static void setConfig(OpikConfiguration config) { - EncryptionUtils.config = config; + public static void setConfig(@NonNull OpikConfiguration config) { + byte[] keyBytes = config.getEncryption().getKey().getBytes(StandardCharsets.UTF_8); + key = new SecretKeySpec(keyBytes, ALGO); } - public static String encrypt(String data) { - init(); + public static String encrypt(@NonNull String data) { try { Cipher c = Cipher.getInstance(ALGO); c.init(Cipher.ENCRYPT_MODE, key); @@ -50,8 +38,7 @@ public static String encrypt(String data) { } } - public static String decrypt(String encryptedData) { - init(); + public static String decrypt(@NonNull String encryptedData) { try { Cipher c = Cipher.getInstance(ALGO); c.init(Cipher.DECRYPT_MODE, key); diff --git a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/EncryptionUtilsModule.java b/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/EncryptionUtilsModule.java deleted file mode 100644 index e1861f031a..0000000000 --- a/apps/opik-backend/src/main/java/com/comet/opik/infrastructure/EncryptionUtilsModule.java +++ /dev/null @@ -1,11 +0,0 @@ -package com.comet.opik.infrastructure; - -import ru.vyarus.dropwizard.guice.module.support.DropwizardAwareModule; - -public class EncryptionUtilsModule extends DropwizardAwareModule { - - @Override - protected void configure() { - requestStaticInjection(EncryptionUtils.class); - } -} diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/resources/LlmProviderApiKeyResourceClient.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/resources/LlmProviderApiKeyResourceClient.java new file mode 100644 index 0000000000..ad3514057d --- /dev/null +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/utils/resources/LlmProviderApiKeyResourceClient.java @@ -0,0 +1,101 @@ +package com.comet.opik.api.resources.utils.resources; + +import com.comet.opik.api.LlmProvider; +import com.comet.opik.api.Page; +import com.comet.opik.api.ProviderApiKey; +import com.comet.opik.api.ProviderApiKeyUpdate; +import com.comet.opik.api.resources.utils.TestUtils; +import jakarta.ws.rs.HttpMethod; +import jakarta.ws.rs.client.Entity; +import jakarta.ws.rs.core.HttpHeaders; +import jakarta.ws.rs.core.MediaType; +import lombok.RequiredArgsConstructor; +import ru.vyarus.dropwizard.guice.test.ClientSupport; +import uk.co.jemos.podam.api.PodamUtils; + +import java.util.UUID; + +import static com.comet.opik.infrastructure.auth.RequestContext.WORKSPACE_HEADER; +import static org.assertj.core.api.Assertions.assertThat; + +@RequiredArgsConstructor +public class LlmProviderApiKeyResourceClient { + private static final String RESOURCE_PATH = "%s/v1/private/llm-provider-key"; + + private final ClientSupport client; + private final String baseURI; + + public ProviderApiKey createProviderApiKey(String providerApiKey, String apiKey, String workspaceName, + int expectedStatus) { + ProviderApiKey body = ProviderApiKey.builder().provider(randomLlmProvider()).apiKey(providerApiKey).build(); + try (var actualResponse = client.target(RESOURCE_PATH.formatted(baseURI)) + .request() + .accept(MediaType.APPLICATION_JSON_TYPE) + .header(HttpHeaders.AUTHORIZATION, apiKey) + .header(WORKSPACE_HEADER, workspaceName) + .post(Entity.json(body))) { + + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(expectedStatus); + if (expectedStatus == 201) { + return body.toBuilder() + .id(TestUtils.getIdFromLocation(actualResponse.getLocation())) + .build(); + } + + return null; + } + } + + public void updateProviderApiKey(UUID id, String providerApiKey, String apiKey, String workspaceName, + int expectedStatus) { + try (var actualResponse = client.target(RESOURCE_PATH.formatted(baseURI)) + .path(id.toString()) + .request() + .accept(MediaType.APPLICATION_JSON_TYPE) + .header(HttpHeaders.AUTHORIZATION, apiKey) + .header(WORKSPACE_HEADER, workspaceName) + .method(HttpMethod.PATCH, Entity.json(ProviderApiKeyUpdate.builder().apiKey(providerApiKey).build()))) { + + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(expectedStatus); + } + } + + public ProviderApiKey getById(UUID id, String workspaceName, String apiKey) { + try (var actualResponse = client.target(RESOURCE_PATH.formatted(baseURI)) + .path(id.toString()) + .request() + .header(HttpHeaders.AUTHORIZATION, apiKey) + .header(WORKSPACE_HEADER, workspaceName) + .get()) { + + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(200); + assertThat(actualResponse.hasEntity()).isTrue(); + + var actualEntity = actualResponse.readEntity(ProviderApiKey.class); + assertThat(actualEntity.apiKey()).isBlank(); + + return actualEntity; + } + } + + public Page getAll(String workspaceName, String apiKey) { + try (var actualResponse = client.target(RESOURCE_PATH.formatted(baseURI)) + .request() + .header(HttpHeaders.AUTHORIZATION, apiKey) + .header(WORKSPACE_HEADER, workspaceName) + .get()) { + + assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(200); + assertThat(actualResponse.hasEntity()).isTrue(); + + var actualEntity = actualResponse.readEntity(ProviderApiKey.ProviderApiKeyPage.class); + actualEntity.content().forEach(providerApiKey -> assertThat(providerApiKey.apiKey()).isBlank()); + + return actualEntity; + } + } + + public LlmProvider randomLlmProvider() { + return LlmProvider.values()[PodamUtils.getIntegerInRange(0, LlmProvider.values().length - 1)]; + } +} diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/LlmProviderApiKeyResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/LlmProviderApiKeyResourceTest.java index dd10524f36..5f3b617971 100644 --- a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/LlmProviderApiKeyResourceTest.java +++ b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/LlmProviderApiKeyResourceTest.java @@ -1,8 +1,7 @@ package com.comet.opik.api.resources.v1.priv; -import com.comet.opik.api.LlmProvider; +import com.comet.opik.api.Page; import com.comet.opik.api.ProviderApiKey; -import com.comet.opik.api.ProviderApiKeyUpdate; import com.comet.opik.api.resources.utils.AuthTestUtils; import com.comet.opik.api.resources.utils.ClickHouseContainerUtils; import com.comet.opik.api.resources.utils.ClientSupportUtils; @@ -10,17 +9,13 @@ import com.comet.opik.api.resources.utils.MySQLContainerUtils; import com.comet.opik.api.resources.utils.RedisContainerUtils; import com.comet.opik.api.resources.utils.TestDropwizardAppExtensionUtils; -import com.comet.opik.api.resources.utils.TestUtils; import com.comet.opik.api.resources.utils.WireMockUtils; +import com.comet.opik.api.resources.utils.resources.LlmProviderApiKeyResourceClient; import com.comet.opik.domain.LlmProviderApiKeyDAO; import com.comet.opik.infrastructure.DatabaseAnalyticsFactory; import com.comet.opik.infrastructure.EncryptionUtils; import com.comet.opik.podam.PodamFactoryUtils; import com.redis.testcontainers.RedisContainer; -import jakarta.ws.rs.HttpMethod; -import jakarta.ws.rs.client.Entity; -import jakarta.ws.rs.core.HttpHeaders; -import jakarta.ws.rs.core.MediaType; import org.jdbi.v3.core.Jdbi; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.BeforeAll; @@ -37,12 +32,11 @@ import uk.co.jemos.podam.api.PodamFactory; import java.sql.SQLException; +import java.util.List; import java.util.UUID; -import static com.comet.opik.api.LlmProvider.OPEN_AI; import static com.comet.opik.api.resources.utils.ClickHouseContainerUtils.DATABASE_NAME; import static com.comet.opik.api.resources.utils.MigrationUtils.CLICKHOUSE_CHANGELOG_FILE; -import static com.comet.opik.infrastructure.auth.RequestContext.WORKSPACE_HEADER; import static com.comet.opik.infrastructure.db.TransactionTemplateAsync.READ_ONLY; import static org.assertj.core.api.Assertions.assertThat; @@ -81,6 +75,7 @@ class LlmProviderApiKeyResourceTest { private String baseURI; private ClientSupport client; private TransactionTemplate mySqlTemplate; + private LlmProviderApiKeyResourceClient llmProviderApiKeyResourceClient; @BeforeAll void setUpAll(ClientSupport client, Jdbi jdbi, @@ -96,6 +91,7 @@ void setUpAll(ClientSupport client, Jdbi jdbi, this.baseURI = "http://localhost:%d".formatted(client.getPort()); this.client = client; this.mySqlTemplate = mySqlTemplate; + this.llmProviderApiKeyResourceClient = new LlmProviderApiKeyResourceClient(this.client, this.baseURI); ClientSupportUtils.config(client); } @@ -116,19 +112,19 @@ void createAndUpdateProviderApiKey() { String workspaceName = UUID.randomUUID().toString(); String apiKey = UUID.randomUUID().toString(); String workspaceId = UUID.randomUUID().toString(); - var provider = OPEN_AI; String providerApiKey = factory.manufacturePojo(String.class); mockTargetWorkspace(apiKey, workspaceName, workspaceId); - var id = createProviderApiKey(provider, providerApiKey, apiKey, workspaceName, 201); - var expectedProviderApiKey = ProviderApiKey.builder().id(id).provider(provider).build(); + var expectedProviderApiKey = llmProviderApiKeyResourceClient.createProviderApiKey(providerApiKey, apiKey, + workspaceName, 201); getAndAssertProviderApiKey(expectedProviderApiKey, apiKey, workspaceName); - checkEncryption(id, workspaceId, providerApiKey); + checkEncryption(expectedProviderApiKey.id(), workspaceId, providerApiKey); String newProviderApiKey = factory.manufacturePojo(String.class); - updateProviderApiKey(id, newProviderApiKey, apiKey, workspaceName, 204); - checkEncryption(id, workspaceId, newProviderApiKey); + llmProviderApiKeyResourceClient.updateProviderApiKey(expectedProviderApiKey.id(), newProviderApiKey, apiKey, + workspaceName, 204); + checkEncryption(expectedProviderApiKey.id(), workspaceId, newProviderApiKey); } @Test @@ -138,13 +134,12 @@ void createProviderApiKeyForExistingProviderShouldFail() { String workspaceName = UUID.randomUUID().toString(); String apiKey = UUID.randomUUID().toString(); String workspaceId = UUID.randomUUID().toString(); - var provider = OPEN_AI; String providerApiKey = factory.manufacturePojo(String.class); mockTargetWorkspace(apiKey, workspaceName, workspaceId); - createProviderApiKey(provider, providerApiKey, apiKey, workspaceName, 201); - createProviderApiKey(provider, providerApiKey, apiKey, workspaceName, 409); + llmProviderApiKeyResourceClient.createProviderApiKey(providerApiKey, apiKey, workspaceName, 201); + llmProviderApiKeyResourceClient.createProviderApiKey(providerApiKey, apiKey, workspaceName, 409); } @Test @@ -159,56 +154,37 @@ void updateProviderFail() { mockTargetWorkspace(apiKey, workspaceName, workspaceId); // for non-existing id - updateProviderApiKey(UUID.randomUUID(), providerApiKey, apiKey, workspaceName, 404); + llmProviderApiKeyResourceClient.updateProviderApiKey(UUID.randomUUID(), providerApiKey, apiKey, workspaceName, + 404); } - private UUID createProviderApiKey(LlmProvider provider, String providerApiKey, String apiKey, String workspaceName, - int expectedStatus) { - try (var actualResponse = client.target(URL_TEMPLATE.formatted(baseURI)) - .request() - .accept(MediaType.APPLICATION_JSON_TYPE) - .header(HttpHeaders.AUTHORIZATION, apiKey) - .header(WORKSPACE_HEADER, workspaceName) - .post(Entity.json(ProviderApiKey.builder().provider(provider).apiKey(providerApiKey).build()))) { - - assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(expectedStatus); - if (expectedStatus == 201) { - return TestUtils.getIdFromLocation(actualResponse.getLocation()); - } - - return null; - } - } + @Test + @DisplayName("Create and get provider Api Keys List") + void createAndGetProviderApiKeyList() { + + String workspaceName = UUID.randomUUID().toString(); + String apiKey = UUID.randomUUID().toString(); + String workspaceId = UUID.randomUUID().toString(); + String providerApiKey = factory.manufacturePojo(String.class); + + mockTargetWorkspace(apiKey, workspaceName, workspaceId); + + // No LLM Provider api keys, expect empty response + var actualProviderApiKeyPage = llmProviderApiKeyResourceClient.getAll(workspaceName, apiKey); + assertPage(actualProviderApiKeyPage, List.of()); + + // Create LLM Provider api key + var expectedProviderApiKey = llmProviderApiKeyResourceClient.createProviderApiKey(providerApiKey, apiKey, + workspaceName, 201); + actualProviderApiKeyPage = llmProviderApiKeyResourceClient.getAll(workspaceName, apiKey); + assertPage(actualProviderApiKeyPage, List.of(expectedProviderApiKey)); - private void updateProviderApiKey(UUID id, String providerApiKey, String apiKey, String workspaceName, - int expectedStatus) { - try (var actualResponse = client.target(URL_TEMPLATE.formatted(baseURI)) - .path(id.toString()) - .request() - .accept(MediaType.APPLICATION_JSON_TYPE) - .header(HttpHeaders.AUTHORIZATION, apiKey) - .header(WORKSPACE_HEADER, workspaceName) - .method(HttpMethod.PATCH, Entity.json(ProviderApiKeyUpdate.builder().apiKey(providerApiKey).build()))) { - - assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(expectedStatus); - } } private void getAndAssertProviderApiKey(ProviderApiKey expected, String apiKey, String workspaceName) { - try (var actualResponse = client.target(URL_TEMPLATE.formatted(baseURI)) - .path(expected.id().toString()) - .request() - .header(HttpHeaders.AUTHORIZATION, apiKey) - .header(WORKSPACE_HEADER, workspaceName) - .get()) { - - assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(200); - assertThat(actualResponse.hasEntity()).isTrue(); - - var actualEntity = actualResponse.readEntity(ProviderApiKey.class); - assertThat(actualEntity.provider()).isEqualTo(expected.provider()); - assertThat(actualEntity.apiKey()).isBlank(); - } + var actualEntity = llmProviderApiKeyResourceClient.getById(expected.id(), workspaceName, apiKey); + assertThat(actualEntity.provider()).isEqualTo(expected.provider()); + assertThat(actualEntity.apiKey()).isBlank(); } private void checkEncryption(UUID id, String workspaceId, String expectedApiKey) { @@ -218,4 +194,14 @@ private void checkEncryption(UUID id, String workspaceId, String expectedApiKey) }); assertThat(EncryptionUtils.decrypt(actualEncryptedApiKey)).isEqualTo(expectedApiKey); } -} \ No newline at end of file + + private void assertPage(Page actual, List expected) { + assertThat(actual.content()).hasSize(expected.size()); + assertThat(actual.page()).isEqualTo(0); + assertThat(actual.total()).isEqualTo(expected.size()); + assertThat(actual.size()).isEqualTo(expected.size()); + + assertThat(actual.content().stream().map(ProviderApiKey::provider).toList()) + .isEqualTo(expected.stream().map(ProviderApiKey::provider).toList()); + } +}