-
Notifications
You must be signed in to change notification settings - Fork 24
Add SyncService client from EGW for embedded calling from Data API #2405
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.
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
Merged
Merged
Changes from all commits
Commits
Show all changes
30 commits
Select commit
Hold shift + click to select a range
5384f34
Add SyncService client from EGW
tatu-at-datastax f545b4f
...
tatu-at-datastax 51cede3
Merge branch 'main' into tatu/2387-add-egw-sync-service-client
tatu-at-datastax 7974351
Merge branch 'main' into tatu/2387-add-egw-sync-service-client
tatu-at-datastax 36f330d
Minor comment addition
tatu-at-datastax b3d9689
Merge branch 'main' into tatu/2387-add-egw-sync-service-client
tatu-at-datastax 15a44fe
Merge branch 'main' into tatu/2387-add-egw-sync-service-client
tatu-at-datastax 99d7663
Add wrapper to get shared secret from secret service in "embedded EGW…
tatu-at-datastax 3b3e843
refactoring: egw->syncservice
tatu-at-datastax 657913d
Add warning for missing/mismatching credentials
tatu-at-datastax 32dc72d
Fix resolution logic
tatu-at-datastax 6f71b26
One more fix
tatu-at-datastax 9cc09f1
Merge branch 'main' into tatu/2387-add-egw-sync-service-client
tatu-at-datastax 562ce0b
Merge branch 'main' into tatu/2387-add-egw-sync-service-client
tatu-at-datastax 1bdcb05
Merge branch 'main' into tatu/2387-add-egw-sync-service-client
tatu-at-datastax 6559e27
Merge branch 'main' into tatu/2387-add-egw-sync-service-client
tatu-at-datastax 7084913
Merge branch 'main' into tatu/2387-add-egw-sync-service-client
tatu-at-datastax 781e13a
Merge branch 'main' into tatu/2387-add-egw-sync-service-client
tatu-at-datastax 9553eb1
Merge branch 'main' into tatu/2387-add-egw-sync-service-client
tatu-at-datastax f10166e
Merge branch 'main' into tatu/2387-add-egw-sync-service-client
tatu-at-datastax 3c319ab
Merge branch 'main' into tatu/2387-add-egw-sync-service-client
tatu-at-datastax ec0a8ea
Fix minor regression
tatu-at-datastax df809e3
Fix issues from code review
tatu-at-datastax 0b54a64
./mvnw fmt:format
tatu-at-datastax aecb79f
Add some testing
tatu-at-datastax ac1a67a
Try disabling failing IT
tatu-at-datastax 46b7be6
Remove test that fails due to needing Sync Service
tatu-at-datastax f029481
Put failing test back, but now disabled
tatu-at-datastax c26148d
Merge branch 'main' into tatu/2387-add-egw-sync-service-client
tatu-at-datastax e5e4743
Merge branch 'main' into tatu/2387-add-egw-sync-service-client
tatu-at-datastax File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
162 changes: 162 additions & 0 deletions
162
...gate/sgv2/jsonapi/service/embedding/operation/SyncServiceCredentialResolvingProvider.java
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,162 @@ | ||
| package io.stargate.sgv2.jsonapi.service.embedding.operation; | ||
|
|
||
| import io.smallrye.mutiny.Uni; | ||
| import io.stargate.sgv2.jsonapi.api.request.EmbeddingCredentials; | ||
| import io.stargate.sgv2.jsonapi.api.request.tenant.Tenant; | ||
| import io.stargate.sgv2.jsonapi.syncservice.SyncServiceClient; | ||
| import java.util.*; | ||
| import org.slf4j.Logger; | ||
| import org.slf4j.LoggerFactory; | ||
|
|
||
| /** | ||
| * A decorator that wraps a direct {@link EmbeddingProvider} and resolves shared-secret credentials | ||
| * via {@link SyncServiceClient} before each {@link #vectorize} call. | ||
| * | ||
| * <p>When the Embedding Gateway (EGW) is disabled (standalone mode), collections configured with | ||
| * {@code SHARED_SECRET} authentication need their credential names resolved into actual secrets. | ||
| * This provider handles that resolution by calling {@link SyncServiceClient#getCredential} for each | ||
| * entry in the authentication map, then passing the resolved credentials to the delegate provider. | ||
| */ | ||
| public class SyncServiceCredentialResolvingProvider extends EmbeddingProvider { | ||
|
|
||
| private static final Logger LOGGER = | ||
| LoggerFactory.getLogger(SyncServiceCredentialResolvingProvider.class); | ||
|
|
||
| private static final String PROVIDER_KEY = "providerKey"; | ||
| private static final String ACCESS_ID = "accessId"; | ||
| private static final String SECRET_KEY = "secretKey"; | ||
| private static final String SECRET_ID = "secretId"; | ||
|
|
||
| private final EmbeddingProvider delegate; | ||
| private final SyncServiceClient syncServiceClient; | ||
| private final Map<String, String> authentication; | ||
| private final Tenant tenant; | ||
| private final String authToken; | ||
|
|
||
| public SyncServiceCredentialResolvingProvider( | ||
| EmbeddingProvider delegate, | ||
| SyncServiceClient syncServiceClient, | ||
| Map<String, String> authentication, | ||
| Tenant tenant, | ||
| String authToken) { | ||
| super( | ||
| delegate.modelProvider(), | ||
| delegate.providerConfig, | ||
| delegate.modelConfig, | ||
| delegate.serviceConfig, | ||
| delegate.dimension, | ||
| delegate.vectorizeServiceParameters); | ||
|
|
||
| this.delegate = delegate; | ||
| this.syncServiceClient = syncServiceClient; | ||
| this.authentication = authentication; | ||
| this.tenant = tenant; | ||
| this.authToken = authToken; | ||
| } | ||
|
|
||
| @Override | ||
| protected String errorMessageJsonPtr() { | ||
| // Not used directly — this wrapper never makes HTTP calls itself | ||
| return ""; | ||
| } | ||
|
|
||
| @Override | ||
| public Uni<BatchedEmbeddingResponse> vectorize( | ||
| int batchId, | ||
| List<String> texts, | ||
| EmbeddingCredentials embeddingCredentials, | ||
| EmbeddingRequestType embeddingRequestType) { | ||
|
|
||
| // Match EGW behavior: if caller already provided credentials via headers, use those | ||
| // directly and skip SyncService resolution | ||
| if (hasHeaderCredentials(embeddingCredentials)) { | ||
| return delegate.vectorize(batchId, texts, embeddingCredentials, embeddingRequestType); | ||
| } | ||
|
|
||
| return resolveCredentials() | ||
| .flatMap(resolved -> delegate.vectorize(batchId, texts, resolved, embeddingRequestType)); | ||
| } | ||
|
|
||
| private static boolean hasHeaderCredentials(EmbeddingCredentials creds) { | ||
| return creds.apiKey().isPresent() | ||
| || creds.accessId().isPresent() | ||
| || creds.secretId().isPresent(); | ||
| } | ||
|
|
||
| @Override | ||
| public int maxBatchSize() { | ||
| return delegate.maxBatchSize(); | ||
| } | ||
|
|
||
| /** | ||
| * Resolves credentials by calling SyncService for each entry in the authentication map. Each | ||
| * entry's key is the accepted token name (e.g. "providerKey"), and the value is the credential | ||
| * reference name stored in SyncService (e.g. "my-openai-cred"). | ||
| * | ||
| * <p>The SyncService response map is keyed by the credential reference name, so we extract the | ||
| * resolved secret using the credential name as key (matching EGW's EmbeddingServiceImpl | ||
| * behavior). | ||
| */ | ||
| private Uni<EmbeddingCredentials> resolveCredentials() { | ||
| String providerName = modelProvider().apiName(); | ||
|
|
||
| // Build parallel SyncService calls and track which accepted token name each one is for | ||
| List<String> acceptedNames = new ArrayList<>(); | ||
| List<String> credNames = new ArrayList<>(); | ||
| List<Uni<Map<String, String>>> resolveUnis = new ArrayList<>(); | ||
|
|
||
| for (Map.Entry<String, String> entry : authentication.entrySet()) { | ||
| acceptedNames.add(entry.getKey()); | ||
| credNames.add(entry.getValue()); | ||
| resolveUnis.add( | ||
| syncServiceClient.getCredential(authToken, tenant, providerName, entry.getValue())); | ||
| } | ||
|
|
||
| return Uni.join() | ||
| .all(resolveUnis) | ||
| .andFailFast() | ||
| .map( | ||
| results -> { | ||
| // For each resolved result, extract the secret using the credential name as key | ||
| // (SyncService response is keyed by credential reference name, not accepted name) | ||
| Map<String, String> resolvedByAcceptedName = new HashMap<>(); | ||
| for (int i = 0; i < results.size(); i++) { | ||
| Map<String, String> credMap = results.get(i); | ||
| if (credMap != null) { | ||
| String credName = credNames.get(i); | ||
| String acceptedName = acceptedNames.get(i); | ||
| String resolvedValue = credMap.get(credName); | ||
| if (resolvedValue != null) { | ||
| resolvedByAcceptedName.put(acceptedName, resolvedValue); | ||
| } else { | ||
| LOGGER.warn( | ||
| "SyncService response for credential '{}' (provider '{}') did not contain" | ||
| + " expected key '{}'; available keys: {}", | ||
| credName, | ||
| providerName, | ||
| credName, | ||
| credMap.keySet()); | ||
| } | ||
| } | ||
| } | ||
| return buildEmbeddingCredentials(resolvedByAcceptedName); | ||
| }); | ||
| } | ||
|
|
||
| /** | ||
| * Maps resolved credential values (keyed by accepted token name) to {@link EmbeddingCredentials}. | ||
| * The accepted token names come from the provider's SHARED_SECRET config (e.g. "providerKey", | ||
| * "accessId", "secretKey"). | ||
| */ | ||
| private EmbeddingCredentials buildEmbeddingCredentials( | ||
| Map<String, String> resolvedByAcceptedName) { | ||
| var apiKey = Optional.ofNullable(resolvedByAcceptedName.get(PROVIDER_KEY)); | ||
| var accessId = Optional.ofNullable(resolvedByAcceptedName.get(ACCESS_ID)); | ||
| var secretId = | ||
| Optional.ofNullable( | ||
| resolvedByAcceptedName.containsKey(SECRET_KEY) | ||
| ? resolvedByAcceptedName.get(SECRET_KEY) | ||
| : resolvedByAcceptedName.get(SECRET_ID)); | ||
| return new EmbeddingCredentials(tenant, apiKey, accessId, secretId, Optional.of(authToken)); | ||
| } | ||
| } | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Great improvement! Original code only resolves one cred