spring-projects / spring-ai

An Application Framework for AI Engineering
https://docs.spring.io/spring-ai/reference/1.0-SNAPSHOT/index.html
Apache License 2.0
2.78k stars 692 forks source link

Add support for configuring an endpoint URL when connecting with Bedrock and Spring AI #1018

Open bruno-oliveira opened 1 month ago

bruno-oliveira commented 1 month ago

Expected Behavior

Essentially, it should be possible to configure a dedicated "endpoint URL" for connecting to a Bedrock hosted endpoint from the side of Spring AI. An example that comes to mind is configuring the connection to Bedrock to use a VPC, for example, to avoid it going over the public internet.

Current Behavior

Currently, there is no way to configure the endpoint URL used to connect to Bedrock from Spring AI, while that seems to be exclusively exposed from the side of the Bedrock API from the Java side.

Context

The most simple context is as defined above: we want to customize the exising endpoint URL to connect to Bedrock in some way, such that we can do it with minimal code while still leveraging all of the abstractions in place for Spring AI.

A current, potential workaround is to create a custom implementation of the AbstractBedrockApi that looks like:

public abstract class CustomAbstractBedrockApi<I, O, SO> {
    private static final Logger logger = LoggerFactory.getLogger(CustomAbstractBedrockApi.class);
    private final String modelId;
    private final ObjectMapper objectMapper;
    private final Region region;
    private BedrockRuntimeClient client;
    private BedrockRuntimeAsyncClient clientStreaming;

...
public CustomAbstractBedrockApi(
            String modelId,
            AwsCredentialsProvider credentialsProvider,
            String region,
            ObjectMapper objectMapper,
            Duration timeout,
            BedrockRuntimeClient bedrockRuntimeClient,
            BedrockRuntimeAsyncClient bedrockRuntimeAsyncClient) {
        this(modelId, credentialsProvider, Region.of(region), objectMapper, timeout);
        client = bedrockRuntimeClient;
        clientStreaming = bedrockRuntimeAsyncClient;
    }

in order to enable configuration via DI of the runtime clients.

With this in place, all that would be needed to do would be, from a configuration class:

@Bean
    @ConditionalOnProperty("spring.ai.bedrock.endpoint-url")
    public BedrockRuntimeClient bedrockRuntimeClient() {
        log.info("Configured runtime client to use VPC");

        try {
            return BedrockRuntimeClient.builder()
                    .region(Region.of(getRegion()))
                    .credentialsProvider(DefaultCredentialsProvider.create())
                    .endpointOverride(new URI(endpointUrl))
                    .build();
        } catch (URISyntaxException e) {
            log.info("Failure when configuring the BedrockRuntimeClient: {}", e.getMessage());
            throw new RuntimeException(e);
        }
    }

While, in the current implementation, quite a lot of custom code is needed.

Extra details

1/ People might need or want to exclude some of the auto-configuration classes because they may need more fine-grained control over how to instantiate and create those Beans at runtime in their own code;

2/ The “interaction surface” between Spring AI and Bedrock API happens at the level of these BedrockRuntime classes that are embedded deep in the code of Spring AI, not exposed externally, deep in the AbstractBedrockApi class;

3/ Defining beans at the application level alone doesn’t help for auto-configuration if they can’t be “Dependency Injected” via a constructor;

PS: Do let me know if I am missing something.

bruno-oliveira commented 1 month ago

This seems to be related: https://github.com/spring-projects/spring-ai/issues/809#issuecomment-2147086881

So I am wondering if someone else has some better workarounds for this, has stumbled upon it, and if it's in the roadmap.

szarpul commented 1 month ago

I've created a PR where an extra constructor with endpoint override is defined. It's also defined for all the Bedrock API implementations.