Skip to content

Commit 577a507

Browse files
committed
Allow OllamaApiAutoConfiguration conditional instantiation based on chat model and embedding model properties
- Introduce OllamaChatOrEmbeddingCondition for this logic - Add unit tests covering the different scenarios Signed-off-by: Nicolas Krier <[email protected]>
1 parent bf5ebce commit 577a507

File tree

2 files changed

+76
-2
lines changed

2 files changed

+76
-2
lines changed

auto-configurations/models/spring-ai-autoconfigure-model-ollama/src/main/java/org/springframework/ai/model/ollama/autoconfigure/OllamaApiAutoConfiguration.java

Lines changed: 39 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,16 +16,26 @@
1616

1717
package org.springframework.ai.model.ollama.autoconfigure;
1818

19+
import org.jspecify.annotations.NonNull;
20+
21+
import org.springframework.ai.model.SpringAIModelProperties;
22+
import org.springframework.ai.model.SpringAIModels;
1923
import org.springframework.ai.ollama.api.OllamaApi;
2024
import org.springframework.ai.retry.autoconfigure.SpringAiRetryAutoConfiguration;
2125
import org.springframework.beans.factory.ObjectProvider;
2226
import org.springframework.boot.autoconfigure.AutoConfiguration;
27+
import org.springframework.boot.autoconfigure.condition.ConditionMessage;
28+
import org.springframework.boot.autoconfigure.condition.ConditionOutcome;
2329
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
2430
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
31+
import org.springframework.boot.autoconfigure.condition.SpringBootCondition;
2532
import org.springframework.boot.context.properties.EnableConfigurationProperties;
2633
import org.springframework.boot.restclient.autoconfigure.RestClientAutoConfiguration;
2734
import org.springframework.boot.webclient.autoconfigure.WebClientAutoConfiguration;
2835
import org.springframework.context.annotation.Bean;
36+
import org.springframework.context.annotation.ConditionContext;
37+
import org.springframework.context.annotation.Conditional;
38+
import org.springframework.core.type.AnnotatedTypeMetadata;
2939
import org.springframework.web.client.ResponseErrorHandler;
3040
import org.springframework.web.client.RestClient;
3141
import org.springframework.web.reactive.function.client.WebClient;
@@ -37,11 +47,13 @@
3747
* @author Eddú Meléndez
3848
* @author Thomas Vitale
3949
* @author Ilayaperumal Gopinathan
50+
* @author Nicolas Krier
4051
* @since 0.8.0
4152
*/
4253
@AutoConfiguration(after = { RestClientAutoConfiguration.class, WebClientAutoConfiguration.class,
4354
SpringAiRetryAutoConfiguration.class })
4455
@ConditionalOnClass(OllamaApi.class)
56+
@Conditional(OllamaApiAutoConfiguration.OllamaChatOrEmbeddingCondition.class)
4557
@EnableConfigurationProperties(OllamaConnectionProperties.class)
4658
public class OllamaApiAutoConfiguration {
4759

@@ -53,7 +65,7 @@ PropertiesOllamaConnectionDetails ollamaConnectionDetails(OllamaConnectionProper
5365

5466
@Bean
5567
@ConditionalOnMissingBean
56-
public OllamaApi ollamaApi(OllamaConnectionDetails connectionDetails,
68+
OllamaApi ollamaApi(OllamaConnectionDetails connectionDetails,
5769
ObjectProvider<RestClient.Builder> restClientBuilderProvider,
5870
ObjectProvider<WebClient.Builder> webClientBuilderProvider, ResponseErrorHandler responseErrorHandler) {
5971
return OllamaApi.builder()
@@ -79,4 +91,30 @@ public String getBaseUrl() {
7991

8092
}
8193

94+
static class OllamaChatOrEmbeddingCondition extends SpringBootCondition {
95+
96+
@Override
97+
public @NonNull ConditionOutcome getMatchOutcome(@NonNull ConditionContext context,
98+
@NonNull AnnotatedTypeMetadata metadata) {
99+
var messageBuilder = ConditionMessage.forCondition("OllamaChatOrEmbeddingCondition");
100+
var environment = context.getEnvironment();
101+
var chatModel = environment.getProperty(SpringAIModelProperties.CHAT_MODEL, SpringAIModels.OLLAMA);
102+
103+
if (SpringAIModels.OLLAMA.equals(chatModel)) {
104+
return ConditionOutcome.match(messageBuilder.because("Chat model corresponds to Ollama."));
105+
}
106+
107+
var embeddingModel = environment.getProperty(SpringAIModelProperties.EMBEDDING_MODEL,
108+
SpringAIModels.OLLAMA);
109+
110+
if (SpringAIModels.OLLAMA.equals(embeddingModel)) {
111+
return ConditionOutcome.match(messageBuilder.because("Embedding model corresponds to Ollama."));
112+
}
113+
114+
return ConditionOutcome
115+
.noMatch(messageBuilder.because("Neither chat model nor embedding model correspond to Ollama."));
116+
}
117+
118+
}
119+
82120
}

auto-configurations/models/spring-ai-autoconfigure-model-ollama/src/test/java/org/springframework/ai/model/ollama/autoconfigure/OllamaModelConfigurationTests.java

Lines changed: 37 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,10 +16,14 @@
1616

1717
package org.springframework.ai.model.ollama.autoconfigure;
1818

19+
import java.util.List;
20+
1921
import org.junit.jupiter.api.Test;
2022

2123
import org.springframework.ai.ollama.OllamaChatModel;
2224
import org.springframework.ai.ollama.OllamaEmbeddingModel;
25+
import org.springframework.ai.ollama.api.OllamaApi;
26+
import org.springframework.ai.utils.SpringAiTestAutoConfigurations;
2327
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
2428

2529
import static org.assertj.core.api.Assertions.assertThat;
@@ -28,11 +32,43 @@
2832
* Unit Tests for Ollama auto-configurations conditional enabling of models.
2933
*
3034
* @author Ilayaperumal Gopinathan
35+
* @author Nicolas Krier
3136
*/
32-
public class OllamaModelConfigurationTests {
37+
class OllamaModelConfigurationTests {
3338

3439
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner();
3540

41+
@Test
42+
void apiActivation() {
43+
this.contextRunner.withConfiguration(SpringAiTestAutoConfigurations.of(OllamaApiAutoConfiguration.class))
44+
.withPropertyValues("spring.ai.model.chat=none", "spring.ai.model.embedding=none")
45+
.run(context -> {
46+
assertThat(context).hasNotFailed();
47+
assertThat(context).doesNotHaveBean(OllamaApiAutoConfiguration.PropertiesOllamaConnectionDetails.class);
48+
assertThat(context).doesNotHaveBean(OllamaApi.class);
49+
});
50+
51+
List.of(
52+
// @formatter:off
53+
this.contextRunner
54+
.withConfiguration(SpringAiTestAutoConfigurations.of(OllamaApiAutoConfiguration.class)),
55+
this.contextRunner
56+
.withConfiguration(SpringAiTestAutoConfigurations.of(OllamaApiAutoConfiguration.class))
57+
.withPropertyValues("spring.ai.model.chat=ollama", "spring.ai.model.embedding=ollama"),
58+
this.contextRunner
59+
.withConfiguration(SpringAiTestAutoConfigurations.of(OllamaApiAutoConfiguration.class))
60+
.withPropertyValues("spring.ai.model.chat=ollama", "spring.ai.model.embedding=none"),
61+
this.contextRunner
62+
.withConfiguration(SpringAiTestAutoConfigurations.of(OllamaApiAutoConfiguration.class))
63+
.withPropertyValues("spring.ai.model.chat=none", "spring.ai.model.embedding=ollama")
64+
// @formatter:on
65+
).forEach(runner -> runner.run(context -> {
66+
assertThat(context).hasNotFailed();
67+
assertThat(context).hasSingleBean(OllamaApiAutoConfiguration.PropertiesOllamaConnectionDetails.class);
68+
assertThat(context).hasSingleBean(OllamaApi.class);
69+
}));
70+
}
71+
3672
@Test
3773
void chatModelActivation() {
3874
this.contextRunner.withConfiguration(BaseOllamaIT.ollamaAutoConfig(OllamaChatAutoConfiguration.class))

0 commit comments

Comments
 (0)