Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[OPIK-548] List LLM Provider api keys #864

Merged
merged 5 commits into from
Dec 12, 2024
Merged

[OPIK-548] List LLM Provider api keys #864

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
579113a
Follow up for OPIK-547
Dec 11, 2024
f09b3c9
OPIK-548 List LLM Provider api keys
Dec 11, 2024
885b30c
Fix comment
Dec 11, 2024
4aa28c9
Add annotations
Dec 11, 2024
e08c5a2
Update annotation to NotNull
Dec 11, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions apps/opik-backend/src/main/java/com/comet/opik/OpikApplication.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -70,7 +70,7 @@ public void initialize(Bootstrap<OpikConfiguration> 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()
Expand All @@ -79,6 +79,7 @@ public void initialize(Bootstrap<OpikConfiguration> bootstrap) {

@Override
public void run(OpikConfiguration configuration, Environment environment) {
EncryptionUtils.setConfig(configuration);
// Resources
var jersey = environment.jersey();

Expand Down
17 changes: 17 additions & 0 deletions apps/opik-backend/src/main/java/com/comet/opik/api/ProviderApiKey.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
import lombok.NonNull;

import java.time.Instant;
import java.util.List;
import java.util.UUID;

@Builder(toBuilder = true)
Expand All @@ -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{" +
Expand All @@ -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<ProviderApiKey> content,
@JsonView({View.Public.class}) List<String> sortableBy)
implements
com.comet.opik.api.Page<ProviderApiKey>{

public static ProviderApiKeyPage empty(int page) {
return new ProviderApiKeyPage(page, 0, 0, List.of(), List.of());
}
}
}
34 changes: 26 additions & 8 deletions ...backend/src/main/java/com/comet/opik/api/resources/v1/priv/LlmProviderApiKeyResource.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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;
Expand Down Expand Up @@ -46,21 +48,37 @@ public class LlmProviderApiKeyResource {
private final @NonNull LlmProviderApiKeyService llmProviderApiKeyService;
private final @NonNull Provider<RequestContext> 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<ProviderApiKey> 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();
}
Expand All @@ -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();

Expand All @@ -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();
}
Expand Down
5 changes: 5 additions & 0 deletions apps/opik-backend/src/main/java/com/comet/opik/domain/LlmProviderApiKeyDAO.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand All @@ -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<ProviderApiKey> find(@Bind("workspaceId") String workspaceId);

default Optional<ProviderApiKey> fetch(UUID id, String workspaceId) {
return Optional.ofNullable(findById(id, workspaceId));
}
Expand Down
33 changes: 25 additions & 8 deletions apps/opik-backend/src/main/java/com/comet/opik/domain/LlmProviderApiKeyService.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -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;

Expand All @@ -25,7 +27,8 @@
@ImplementedBy(LlmProviderApiKeyServiceImpl.class)
public interface LlmProviderApiKeyService {

ProviderApiKey get(UUID id, String workspaceId);
ProviderApiKey find(UUID id, String workspaceId);
Page<ProviderApiKey> find(String workspaceId);
ProviderApiKey saveApiKey(ProviderApiKey providerApiKey, String userName, String workspaceId);
void updateApiKey(UUID id, ProviderApiKeyUpdate providerApiKeyUpdate, String userName, String workspaceId);
}
Expand All @@ -40,23 +43,37 @@ 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 -> {

var repository = handle.attach(LlmProviderApiKeyDAO.class);

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<ProviderApiKey> find(@NonNull String workspaceId) {
List<ProviderApiKey> 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());
thiagohora marked this conversation as resolved.
Show resolved Hide resolved
}

@Override
public ProviderApiKey saveApiKey(@NonNull ProviderApiKey providerApiKey, @NonNull String userName, @NonNull String workspaceId) {
UUID apiKeyId = idGenerator.generateId();

var newProviderApiKey = providerApiKey.toBuilder()
Expand All @@ -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();
Expand All @@ -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 -> {

Expand Down
25 changes: 6 additions & 19 deletions apps/opik-backend/src/main/java/com/comet/opik/infrastructure/EncryptionUtils.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
package com.comet.opik.infrastructure;

import jakarta.inject.Inject;
import lombok.NonNull;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
Expand All @@ -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 {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As key is now nullable, I suggest validating it's not null to avoid NPE

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We set key at the application boot-up, so it will never be null.
I would rather not add additional checks here, since this scenario is not feasible. If you insist I could add it, but seems unnecassary.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

the annotations also serve as documentation and to guarantee new usage will respect the constraints. Basically all public methods should have it

Copy link
Contributor Author

@BorisTkachenko BorisTkachenko Dec 11, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@thiagohora I added annotations to public methods. But it's not related to what Ido is concerned. key is a private field and is not a part of public api.

Cipher c = Cipher.getInstance(ALGO);
c.init(Cipher.ENCRYPT_MODE, key);
Expand All @@ -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 {
BorisTkachenko marked this conversation as resolved.
Show resolved Hide resolved
Cipher c = Cipher.getInstance(ALGO);
c.init(Cipher.DECRYPT_MODE, key);
Expand Down
11 changes: 0 additions & 11 deletions apps/opik-backend/src/main/java/com/comet/opik/infrastructure/EncryptionUtilsModule.java
View file
Open in desktop

This file was deleted.

1 change: 0 additions & 1 deletion ...end/src/test/java/com/comet/opik/api/resources/utils/resources/ChatCompletionsClient.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,6 @@
import jakarta.ws.rs.core.HttpHeaders;
import jakarta.ws.rs.core.MediaType;
import jakarta.ws.rs.core.Response;
import lombok.Builder;
import org.apache.http.HttpStatus;
import org.glassfish.jersey.client.ChunkedInput;
import ru.vyarus.dropwizard.guice.test.ClientSupport;
Expand Down
Loading
Loading