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.85k stars 711 forks source link

Add support for custom WebClientBuilder #739

Open fmunch opened 3 months ago

fmunch commented 3 months ago

This might be related to #512 and #609.

Expected Behavior

The *Api client classes should offer a way to customize the WebClient used. For example to add an Authorization header for custom Ollama setups, or to add some logging.

Current Behavior

Most *Api clients offer a constructor with a RestClient.Builder parameter, but also create a hardcoded WebClient.Builder:

public OllamaApi(String baseUrl, RestClient.Builder restClientBuilder) {
    // ...
    this.restClient = restClientBuilder.baseUrl(baseUrl).defaultHeaders(defaultHeaders).build();
    this.webClient = WebClient.builder().baseUrl(baseUrl).defaultHeaders(defaultHeaders).build();
}

public MistralAiApi(String baseUrl, String mistralAiApiKey, RestClient.Builder restClientBuilder,
        ResponseErrorHandler responseErrorHandler) {
    // ...
    this.restClient = restClientBuilder.baseUrl(baseUrl)
        .defaultHeaders(jsonContentHeaders)
        .defaultStatusHandler(responseErrorHandler)
        .build();
    this.webClient = WebClient.builder().baseUrl(baseUrl).defaultHeaders(jsonContentHeaders).build();
}

public OpenAiApi(String baseUrl, String openAiToken, RestClient.Builder restClientBuilder,
        ResponseErrorHandler responseErrorHandler) {
    this.restClient = restClientBuilder
            .baseUrl(baseUrl)
            .defaultHeaders(ApiUtils.getJsonContentHeaders(openAiToken))
            .defaultStatusHandler(responseErrorHandler)
            .build();
    this.webClient = WebClient.builder()
            .baseUrl(baseUrl)
            .defaultHeaders(ApiUtils.getJsonContentHeaders(openAiToken))
            .build();
}

This does not allow to customize the WebClient used in streaming methods.

Context

We want to deploy a Spring AI cluster that calls an Ollama instance, ideally with some sort of authentication.

markpollack commented 1 month ago

Sounds super reasonable. Will investigate.