Skip to content

Commit 0f71bad

Browse files
OPIK-547 Store and retrieve LLM provider api key (#845)
1 parent 9547626 commit 0f71bad

16 files changed

+709
-1
lines changed

apps/opik-backend/config.yml

+3
Original file line numberDiff line numberDiff line change
@@ -88,3 +88,6 @@ metadata:
8888

8989
cors:
9090
enabled: ${CORS:-false}
91+
92+
encryption:
93+
key: ${OPIK_ENCRYPTION_KEY:-'GiTHubiLoVeYouAA'}

apps/opik-backend/src/main/java/com/comet/opik/OpikApplication.java

+2-1
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import com.comet.opik.api.error.JsonInvalidFormatExceptionMapper;
44
import com.comet.opik.infrastructure.ConfigurationModule;
5+
import com.comet.opik.infrastructure.EncryptionUtilsModule;
56
import com.comet.opik.infrastructure.OpikConfiguration;
67
import com.comet.opik.infrastructure.auth.AuthModule;
78
import com.comet.opik.infrastructure.bi.BiModule;
@@ -69,7 +70,7 @@ public void initialize(Bootstrap<OpikConfiguration> bootstrap) {
6970
.withPlugins(new SqlObjectPlugin(), new Jackson2Plugin()))
7071
.modules(new DatabaseAnalyticsModule(), new IdGeneratorModule(), new AuthModule(), new RedisModule(),
7172
new RateLimitModule(), new NameGeneratorModule(), new HttpModule(), new EventModule(),
72-
new ConfigurationModule(), new BiModule())
73+
new ConfigurationModule(), new BiModule(), new EncryptionUtilsModule())
7374
.installers(JobGuiceyInstaller.class)
7475
.listen(new OpikGuiceyLifecycleEventListener())
7576
.enableAutoConfig()
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
package com.comet.opik.api;
2+
3+
import com.fasterxml.jackson.annotation.JsonProperty;
4+
import lombok.Getter;
5+
import lombok.RequiredArgsConstructor;
6+
7+
@Getter
8+
@RequiredArgsConstructor
9+
public enum LlmProvider {
10+
11+
@JsonProperty("openai")
12+
OPEN_AI;
13+
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
package com.comet.opik.api;
2+
3+
import com.comet.opik.utils.ProviderApiKeyDeserializer;
4+
import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
5+
import com.fasterxml.jackson.annotation.JsonView;
6+
import com.fasterxml.jackson.databind.PropertyNamingStrategies;
7+
import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
8+
import com.fasterxml.jackson.databind.annotation.JsonNaming;
9+
import io.swagger.v3.oas.annotations.media.Schema;
10+
import jakarta.validation.constraints.NotBlank;
11+
import lombok.Builder;
12+
import lombok.NonNull;
13+
14+
import java.time.Instant;
15+
import java.util.UUID;
16+
17+
@Builder(toBuilder = true)
18+
@JsonIgnoreProperties(ignoreUnknown = true)
19+
@JsonNaming(PropertyNamingStrategies.SnakeCaseStrategy.class)
20+
public record ProviderApiKey(
21+
@JsonView( {
22+
View.Public.class}) @Schema(accessMode = Schema.AccessMode.READ_ONLY) UUID id,
23+
@JsonView({View.Public.class, View.Write.class}) @NonNull LlmProvider provider,
24+
@JsonView({
25+
View.Write.class}) @NotBlank @JsonDeserialize(using = ProviderApiKeyDeserializer.class) String apiKey,
26+
@JsonView({View.Public.class}) @Schema(accessMode = Schema.AccessMode.READ_ONLY) Instant createdAt,
27+
@JsonView({View.Public.class}) @Schema(accessMode = Schema.AccessMode.READ_ONLY) String createdBy,
28+
@JsonView({View.Public.class}) @Schema(accessMode = Schema.AccessMode.READ_ONLY) Instant lastUpdatedAt,
29+
@JsonView({View.Public.class}) @Schema(accessMode = Schema.AccessMode.READ_ONLY) String lastUpdatedBy){
30+
@Override
31+
public String toString() {
32+
return "ProviderApiKey{" +
33+
"id=" + id +
34+
", provider='" + provider + '\'' +
35+
", createdAt=" + createdAt +
36+
", createdBy='" + createdBy + '\'' +
37+
", lastUpdatedAt=" + lastUpdatedAt +
38+
", lastUpdatedBy='" + lastUpdatedBy + '\'' +
39+
'}';
40+
}
41+
42+
public static class View {
43+
public static class Write {
44+
}
45+
46+
public static class Public {
47+
}
48+
}
49+
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
package com.comet.opik.api;
2+
3+
import com.comet.opik.utils.ProviderApiKeyDeserializer;
4+
import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
5+
import com.fasterxml.jackson.databind.PropertyNamingStrategies;
6+
import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
7+
import com.fasterxml.jackson.databind.annotation.JsonNaming;
8+
import jakarta.validation.constraints.NotBlank;
9+
import lombok.Builder;
10+
import lombok.Getter;
11+
import lombok.ToString;
12+
13+
@Builder(toBuilder = true)
14+
@JsonIgnoreProperties(ignoreUnknown = true)
15+
@JsonNaming(PropertyNamingStrategies.SnakeCaseStrategy.class)
16+
@Getter
17+
public class ProviderApiKeyUpdate {
18+
@ToString.Exclude
19+
@NotBlank
20+
@JsonDeserialize(using = ProviderApiKeyDeserializer.class)
21+
String apiKey;
22+
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
1+
package com.comet.opik.api.resources.v1.priv;
2+
3+
import com.codahale.metrics.annotation.Timed;
4+
import com.comet.opik.api.ProviderApiKey;
5+
import com.comet.opik.api.ProviderApiKeyUpdate;
6+
import com.comet.opik.api.error.ErrorMessage;
7+
import com.comet.opik.domain.LlmProviderApiKeyService;
8+
import com.comet.opik.infrastructure.auth.RequestContext;
9+
import com.fasterxml.jackson.annotation.JsonView;
10+
import io.swagger.v3.oas.annotations.Operation;
11+
import io.swagger.v3.oas.annotations.headers.Header;
12+
import io.swagger.v3.oas.annotations.media.Content;
13+
import io.swagger.v3.oas.annotations.media.Schema;
14+
import io.swagger.v3.oas.annotations.parameters.RequestBody;
15+
import io.swagger.v3.oas.annotations.responses.ApiResponse;
16+
import io.swagger.v3.oas.annotations.tags.Tag;
17+
import jakarta.inject.Inject;
18+
import jakarta.inject.Provider;
19+
import jakarta.validation.Valid;
20+
import jakarta.ws.rs.Consumes;
21+
import jakarta.ws.rs.GET;
22+
import jakarta.ws.rs.PATCH;
23+
import jakarta.ws.rs.POST;
24+
import jakarta.ws.rs.Path;
25+
import jakarta.ws.rs.PathParam;
26+
import jakarta.ws.rs.Produces;
27+
import jakarta.ws.rs.core.Context;
28+
import jakarta.ws.rs.core.MediaType;
29+
import jakarta.ws.rs.core.Response;
30+
import jakarta.ws.rs.core.UriInfo;
31+
import lombok.NonNull;
32+
import lombok.RequiredArgsConstructor;
33+
import lombok.extern.slf4j.Slf4j;
34+
35+
import java.util.UUID;
36+
37+
@Path("/v1/private/llm-provider-key")
38+
@Produces(MediaType.APPLICATION_JSON)
39+
@Consumes(MediaType.APPLICATION_JSON)
40+
@Timed
41+
@Slf4j
42+
@RequiredArgsConstructor(onConstructor_ = @Inject)
43+
@Tag(name = "LlmProviderKey", description = "LLM Provider Key")
44+
public class LlmProviderApiKeyResource {
45+
46+
private final @NonNull LlmProviderApiKeyService llmProviderApiKeyService;
47+
private final @NonNull Provider<RequestContext> requestContext;
48+
49+
@GET
50+
@Path("{id}")
51+
@Operation(operationId = "getLlmProviderApiKeyById", summary = "Get LLM Provider's ApiKey by id", description = "Get LLM Provider's ApiKey by id", responses = {
52+
@ApiResponse(responseCode = "200", description = "ProviderApiKey resource", content = @Content(schema = @Schema(implementation = ProviderApiKey.class))),
53+
@ApiResponse(responseCode = "404", description = "Not found", content = @Content(schema = @Schema(implementation = ErrorMessage.class)))})
54+
@JsonView({ProviderApiKey.View.Public.class})
55+
public Response getById(@PathParam("id") UUID id) {
56+
57+
String workspaceId = requestContext.get().getWorkspaceId();
58+
59+
log.info("Getting Provider's ApiKey by id '{}' on workspace_id '{}'", id, workspaceId);
60+
61+
ProviderApiKey providerApiKey = llmProviderApiKeyService.get(id, workspaceId);
62+
63+
log.info("Got Provider's ApiKey by id '{}' on workspace_id '{}'", id, workspaceId);
64+
65+
return Response.ok().entity(providerApiKey).build();
66+
}
67+
68+
@POST
69+
@Operation(operationId = "storeLlmProviderApiKey", summary = "Store LLM Provider's ApiKey", description = "Store LLM Provider's ApiKey", responses = {
70+
@ApiResponse(responseCode = "201", description = "Created", headers = {
71+
@Header(name = "Location", required = true, example = "${basePath}/v1/private/proxy/api_key/{apiKeyId}", schema = @Schema(implementation = String.class))}),
72+
@ApiResponse(responseCode = "401", description = "Bad Request", content = @Content(schema = @Schema(implementation = ErrorMessage.class))),
73+
@ApiResponse(responseCode = "403", description = "Access forbidden", content = @Content(schema = @Schema(implementation = ErrorMessage.class)))
74+
})
75+
public Response saveApiKey(
76+
@RequestBody(content = @Content(schema = @Schema(implementation = ProviderApiKey.class))) @JsonView(ProviderApiKey.View.Write.class) @Valid ProviderApiKey providerApiKey,
77+
@Context UriInfo uriInfo) {
78+
String workspaceId = requestContext.get().getWorkspaceId();
79+
String userName = requestContext.get().getUserName();
80+
log.info("Save api key for provider '{}', on workspace_id '{}'", providerApiKey.provider(), workspaceId);
81+
var providerApiKeyId = llmProviderApiKeyService.saveApiKey(providerApiKey, userName, workspaceId).id();
82+
log.info("Saved api key for provider '{}', on workspace_id '{}'", providerApiKey.provider(), workspaceId);
83+
84+
var uri = uriInfo.getAbsolutePathBuilder().path("/%s".formatted(providerApiKeyId)).build();
85+
86+
return Response.created(uri).build();
87+
}
88+
89+
@PATCH
90+
@Path("{id}")
91+
@Operation(operationId = "updateLlmProviderApiKey", summary = "Update LLM Provider's ApiKey", description = "Update LLM Provider's ApiKey", responses = {
92+
@ApiResponse(responseCode = "204", description = "No Content"),
93+
@ApiResponse(responseCode = "401", description = "Bad Request", content = @Content(schema = @Schema(implementation = ErrorMessage.class))),
94+
@ApiResponse(responseCode = "403", description = "Access forbidden", content = @Content(schema = @Schema(implementation = ErrorMessage.class))),
95+
@ApiResponse(responseCode = "404", description = "Not found", content = @Content(schema = @Schema(implementation = ErrorMessage.class)))
96+
})
97+
public Response updateApiKey(@PathParam("id") UUID id,
98+
@RequestBody(content = @Content(schema = @Schema(implementation = ProviderApiKeyUpdate.class))) @Valid ProviderApiKeyUpdate providerApiKeyUpdate) {
99+
String workspaceId = requestContext.get().getWorkspaceId();
100+
String userName = requestContext.get().getUserName();
101+
102+
log.info("Updating api key for provider with id '{}' on workspaceId '{}'", id, workspaceId);
103+
llmProviderApiKeyService.updateApiKey(id, providerApiKeyUpdate, userName, workspaceId);
104+
log.info("Updated api key for provider with id '{}' on workspaceId '{}'", id, workspaceId);
105+
106+
return Response.noContent().build();
107+
}
108+
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
package com.comet.opik.domain;
2+
3+
import com.comet.opik.api.ProviderApiKey;
4+
import com.comet.opik.infrastructure.db.UUIDArgumentFactory;
5+
import org.jdbi.v3.sqlobject.config.RegisterArgumentFactory;
6+
import org.jdbi.v3.sqlobject.config.RegisterConstructorMapper;
7+
import org.jdbi.v3.sqlobject.customizer.Bind;
8+
import org.jdbi.v3.sqlobject.customizer.BindMethods;
9+
import org.jdbi.v3.sqlobject.statement.SqlQuery;
10+
import org.jdbi.v3.sqlobject.statement.SqlUpdate;
11+
12+
import java.util.Optional;
13+
import java.util.UUID;
14+
15+
@RegisterConstructorMapper(ProviderApiKey.class)
16+
@RegisterArgumentFactory(UUIDArgumentFactory.class)
17+
public interface LlmProviderApiKeyDAO {
18+
19+
@SqlUpdate("INSERT INTO llm_provider_api_key (id, provider, workspace_id, api_key, created_by, last_updated_by) VALUES (:bean.id, :bean.provider, :workspaceId, :bean.apiKey, :bean.createdBy, :bean.lastUpdatedBy)")
20+
void save(@Bind("workspaceId") String workspaceId,
21+
@BindMethods("bean") ProviderApiKey providerApiKey);
22+
23+
@SqlUpdate("UPDATE llm_provider_api_key SET " +
24+
"api_key = :apiKey, " +
25+
"last_updated_by = :lastUpdatedBy " +
26+
"WHERE id = :id AND workspace_id = :workspaceId")
27+
void update(@Bind("id") UUID id,
28+
@Bind("workspaceId") String workspaceId,
29+
@Bind("apiKey") String encryptedApiKey,
30+
@Bind("lastUpdatedBy") String lastUpdatedBy);
31+
32+
@SqlQuery("SELECT * FROM llm_provider_api_key WHERE id = :id AND workspace_id = :workspaceId")
33+
ProviderApiKey findById(@Bind("id") UUID id, @Bind("workspaceId") String workspaceId);
34+
35+
default Optional<ProviderApiKey> fetch(UUID id, String workspaceId) {
36+
return Optional.ofNullable(findById(id, workspaceId));
37+
}
38+
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,118 @@
1+
package com.comet.opik.domain;
2+
3+
import com.comet.opik.api.ProviderApiKey;
4+
import com.comet.opik.api.ProviderApiKeyUpdate;
5+
import com.comet.opik.api.error.EntityAlreadyExistsException;
6+
import com.comet.opik.api.error.ErrorMessage;
7+
import com.google.inject.ImplementedBy;
8+
import jakarta.inject.Inject;
9+
import jakarta.inject.Singleton;
10+
import jakarta.ws.rs.NotFoundException;
11+
import jakarta.ws.rs.core.Response;
12+
import lombok.NonNull;
13+
import lombok.RequiredArgsConstructor;
14+
import lombok.extern.slf4j.Slf4j;
15+
import org.jdbi.v3.core.statement.UnableToExecuteStatementException;
16+
import ru.vyarus.guicey.jdbi3.tx.TransactionTemplate;
17+
18+
import java.sql.SQLIntegrityConstraintViolationException;
19+
import java.util.List;
20+
import java.util.UUID;
21+
22+
import static com.comet.opik.infrastructure.db.TransactionTemplateAsync.READ_ONLY;
23+
import static com.comet.opik.infrastructure.db.TransactionTemplateAsync.WRITE;
24+
25+
@ImplementedBy(LlmProviderApiKeyServiceImpl.class)
26+
public interface LlmProviderApiKeyService {
27+
28+
ProviderApiKey get(UUID id, String workspaceId);
29+
ProviderApiKey saveApiKey(ProviderApiKey providerApiKey, String userName, String workspaceId);
30+
void updateApiKey(UUID id, ProviderApiKeyUpdate providerApiKeyUpdate, String userName, String workspaceId);
31+
}
32+
33+
@Slf4j
34+
@Singleton
35+
@RequiredArgsConstructor(onConstructor_ = @Inject)
36+
class LlmProviderApiKeyServiceImpl implements LlmProviderApiKeyService {
37+
38+
private static final String PROVIDER_API_KEY_ALREADY_EXISTS = "Api key for this provider already exists";
39+
private final @NonNull IdGenerator idGenerator;
40+
private final @NonNull TransactionTemplate template;
41+
42+
@Override
43+
public ProviderApiKey get(UUID id, String workspaceId) {
44+
log.info("Getting provider api key with id '{}', workspaceId '{}'", id, workspaceId);
45+
46+
ProviderApiKey providerApiKey = template.inTransaction(READ_ONLY, handle -> {
47+
48+
var repository = handle.attach(LlmProviderApiKeyDAO.class);
49+
50+
return repository.fetch(id, workspaceId).orElseThrow(this::createNotFoundError);
51+
});
52+
log.info("Got provider api key with id '{}', workspaceId '{}'", id, workspaceId);
53+
54+
return providerApiKey.toBuilder()
55+
.build();
56+
}
57+
58+
@Override
59+
public ProviderApiKey saveApiKey(@NonNull ProviderApiKey providerApiKey, String userName, String workspaceId) {
60+
UUID apiKeyId = idGenerator.generateId();
61+
62+
var newProviderApiKey = providerApiKey.toBuilder()
63+
.id(apiKeyId)
64+
.createdBy(userName)
65+
.lastUpdatedBy(userName)
66+
.build();
67+
68+
try {
69+
template.inTransaction(WRITE, handle -> {
70+
71+
var repository = handle.attach(LlmProviderApiKeyDAO.class);
72+
repository.save(workspaceId, newProviderApiKey);
73+
74+
return newProviderApiKey;
75+
});
76+
77+
return get(apiKeyId, workspaceId);
78+
} catch (UnableToExecuteStatementException e) {
79+
if (e.getCause() instanceof SQLIntegrityConstraintViolationException) {
80+
throw newConflict();
81+
} else {
82+
throw e;
83+
}
84+
}
85+
}
86+
87+
@Override
88+
public void updateApiKey(@NonNull UUID id, @NonNull ProviderApiKeyUpdate providerApiKeyUpdate, String userName,
89+
String workspaceId) {
90+
91+
template.inTransaction(WRITE, handle -> {
92+
93+
var repository = handle.attach(LlmProviderApiKeyDAO.class);
94+
95+
ProviderApiKey providerApiKey = repository.fetch(id, workspaceId)
96+
.orElseThrow(this::createNotFoundError);
97+
98+
repository.update(providerApiKey.id(),
99+
workspaceId,
100+
providerApiKeyUpdate.getApiKey(),
101+
userName);
102+
103+
return null;
104+
});
105+
}
106+
107+
private EntityAlreadyExistsException newConflict() {
108+
log.info(PROVIDER_API_KEY_ALREADY_EXISTS);
109+
return new EntityAlreadyExistsException(new ErrorMessage(List.of(PROVIDER_API_KEY_ALREADY_EXISTS)));
110+
}
111+
112+
private NotFoundException createNotFoundError() {
113+
String message = "Provider api key not found";
114+
log.info(message);
115+
return new NotFoundException(message,
116+
Response.status(Response.Status.NOT_FOUND).entity(new ErrorMessage(List.of(message))).build());
117+
}
118+
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
package com.comet.opik.infrastructure;
2+
3+
import com.fasterxml.jackson.annotation.JsonProperty;
4+
import jakarta.validation.Valid;
5+
import jakarta.validation.constraints.NotNull;
6+
import lombok.Data;
7+
import lombok.ToString;
8+
9+
@Data
10+
public class EncryptionConfig {
11+
12+
@Valid
13+
@JsonProperty
14+
@NotNull @ToString.Exclude
15+
private String key;
16+
}

0 commit comments

Comments
 (0)