Spring/spring-ai
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add observability support for Ollama

Browse Source 
- Improve ITs to reuse a single container across tests
This commit is contained in:
Thomas Vitale
2024-08-16 15:57:51 +02:00
committed by Mark Pollack
parent 3bef5c12bf
commit 10e1e13fa2
27 changed files with 691 additions and 216 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
models/spring-ai-ollama/pom.xml
View File

@@ -74,6 +74,12 @@
<scope>test</scope>
</dependency>
<dependency>
<groupId>io.micrometer</groupId>
<artifactId>micrometer-observation-test</artifactId>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.testcontainers</groupId>
<artifactId>junit-jupiter</artifactId>

215
models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/OllamaChatModel.java
View File

@@ -21,17 +21,22 @@ import java.util.List;
import java.util.Map;
import java.util.Set;
import io.micrometer.observation.Observation;
import io.micrometer.observation.ObservationRegistry;
import io.micrometer.observation.contextpropagation.ObservationThreadLocalAccessor;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.ToolResponseMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.metadata.ChatGenerationMetadata;
import org.springframework.ai.chat.metadata.ChatResponseMetadata;
import org.springframework.ai.chat.model.AbstractToolCallSupport;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.Generation;
import org.springframework.ai.chat.model.*;
import org.springframework.ai.chat.observation.ChatModelObservationContext;
import org.springframework.ai.chat.observation.ChatModelObservationConvention;
import org.springframework.ai.chat.observation.ChatModelObservationDocumentation;
import org.springframework.ai.chat.observation.DefaultChatModelObservationConvention;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.ChatOptionsBuilder;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.ModelOptionsUtils;
import org.springframework.ai.model.function.FunctionCallback;
@@ -64,6 +69,8 @@ import reactor.core.publisher.Flux;
*/
public class OllamaChatModel extends AbstractToolCallSupport implements ChatModel {
private static final ChatModelObservationConvention DEFAULT_OBSERVATION_CONVENTION = new DefaultChatModelObservationConvention();
/**
* Low-level Ollama API library.
*/
@@ -72,61 +79,97 @@ public class OllamaChatModel extends AbstractToolCallSupport implements ChatMode
/**
* Default options to be used for all chat requests.
*/
private OllamaOptions defaultOptions;
private final OllamaOptions defaultOptions;
public OllamaChatModel(OllamaApi chatApi) {
this(chatApi, OllamaOptions.create().withModel(OllamaOptions.DEFAULT_MODEL));
/**
* Observation registry used for instrumentation.
*/
private final ObservationRegistry observationRegistry;
/**
* Conventions to use for generating observations.
*/
private ChatModelObservationConvention observationConvention = DEFAULT_OBSERVATION_CONVENTION;
public OllamaChatModel(OllamaApi ollamaApi) {
this(ollamaApi, OllamaOptions.create().withModel(OllamaOptions.DEFAULT_MODEL));
}
public OllamaChatModel(OllamaApi chatApi, OllamaOptions defaultOptions) {
this(chatApi, defaultOptions, null);
public OllamaChatModel(OllamaApi ollamaApi, OllamaOptions defaultOptions) {
this(ollamaApi, defaultOptions, null);
}
public OllamaChatModel(OllamaApi chatApi, OllamaOptions defaultOptions,
public OllamaChatModel(OllamaApi ollamaApi, OllamaOptions defaultOptions,
FunctionCallbackContext functionCallbackContext) {
this(chatApi, defaultOptions, functionCallbackContext, List.of());
this(ollamaApi, defaultOptions, functionCallbackContext, List.of());
}
public OllamaChatModel(OllamaApi ollamaApi, OllamaOptions defaultOptions,
FunctionCallbackContext functionCallbackContext, List<FunctionCallback> toolFunctionCallbacks) {
this(ollamaApi, defaultOptions, functionCallbackContext, toolFunctionCallbacks, ObservationRegistry.NOOP);
}
public OllamaChatModel(OllamaApi chatApi, OllamaOptions defaultOptions,
FunctionCallbackContext functionCallbackContext, List<FunctionCallback> toolFunctionCallbacks) {
FunctionCallbackContext functionCallbackContext, List<FunctionCallback> toolFunctionCallbacks,
ObservationRegistry observationRegistry) {
super(functionCallbackContext, defaultOptions, toolFunctionCallbacks);
Assert.notNull(chatApi, "OllamaApi must not be null");
Assert.notNull(defaultOptions, "DefaultOptions must not be null");
Assert.notNull(chatApi, "ollamaApi must not be null");
Assert.notNull(defaultOptions, "defaultOptions must not be null");
Assert.notNull(observationRegistry, "ObservationRegistry must not be null");
this.chatApi = chatApi;
this.defaultOptions = defaultOptions;
this.observationRegistry = observationRegistry;
}
@Override
public ChatResponse call(Prompt prompt) {
OllamaApi.ChatRequest request = ollamaChatRequest(prompt, false);
OllamaApi.ChatResponse response = this.chatApi.chat(ollamaChatRequest(prompt, false));
ChatModelObservationContext observationContext = ChatModelObservationContext.builder()
.prompt(prompt)
.provider(OllamaApi.PROVIDER_NAME)
.requestOptions(buildRequestOptions(request))
.build();
List<AssistantMessage.ToolCall> toolCalls = response.message().toolCalls() == null ? List.of()
: response.message()
.toolCalls()
.stream()
.map(toolCall -> new AssistantMessage.ToolCall("", "function", toolCall.function().name(),
ModelOptionsUtils.toJsonString(toolCall.function().arguments())))
.toList();
ChatResponse response = ChatModelObservationDocumentation.CHAT_MODEL_OPERATION
.observation(this.observationConvention, DEFAULT_OBSERVATION_CONVENTION, () -> observationContext,
this.observationRegistry)
.observe(() -> {
var assistantMessage = new AssistantMessage(response.message().content(), Map.of(), toolCalls);
OllamaApi.ChatResponse ollamaResponse = this.chatApi.chat(request);
ChatGenerationMetadata generationMetadata = ChatGenerationMetadata.NULL;
if (response.promptEvalCount() != null && response.evalCount() != null) {
generationMetadata = ChatGenerationMetadata.from(response.doneReason(), null);
}
List<AssistantMessage.ToolCall> toolCalls = ollamaResponse.message().toolCalls() == null ? List.of()
: ollamaResponse.message()
.toolCalls()
.stream()
.map(toolCall -> new AssistantMessage.ToolCall("", "function", toolCall.function().name(),
ModelOptionsUtils.toJsonString(toolCall.function().arguments())))
.toList();
var generator = new Generation(assistantMessage, generationMetadata);
var chatResponse = new ChatResponse(List.of(generator), from(response));
var assistantMessage = new AssistantMessage(ollamaResponse.message().content(), Map.of(), toolCalls);
if (isToolCall(chatResponse, Set.of("stop"))) {
var toolCallConversation = handleToolCalls(prompt, chatResponse);
ChatGenerationMetadata generationMetadata = ChatGenerationMetadata.NULL;
if (ollamaResponse.promptEvalCount() != null && ollamaResponse.evalCount() != null) {
generationMetadata = ChatGenerationMetadata.from(ollamaResponse.doneReason(), null);
}
var generator = new Generation(assistantMessage, generationMetadata);
ChatResponse chatResponse = new ChatResponse(List.of(generator), from(ollamaResponse));
observationContext.setResponse(chatResponse);
return chatResponse;
});
if (response != null && isToolCall(response, Set.of("stop"))) {
var toolCallConversation = handleToolCalls(prompt, response);
// Recursively call the call method with the tool call message
// conversation that contains the call responses.
return this.call(new Prompt(toolCallConversation, prompt.getOptions()));
}
return chatResponse;
return response;
}
public static ChatResponseMetadata from(OllamaApi.ChatResponse response) {
@@ -147,40 +190,64 @@ public class OllamaChatModel extends AbstractToolCallSupport implements ChatMode
@Override
public Flux<ChatResponse> stream(Prompt prompt) {
return Flux.deferContextual(contextView -> {
OllamaApi.ChatRequest request = ollamaChatRequest(prompt, true);
Flux<OllamaApi.ChatResponse> ollamaResponse = this.chatApi.streamingChat(ollamaChatRequest(prompt, true));
final ChatModelObservationContext observationContext = ChatModelObservationContext.builder()
.prompt(prompt)
.provider(OllamaApi.PROVIDER_NAME)
.requestOptions(buildRequestOptions(request))
.build();
Flux<ChatResponse> chatResponse = ollamaResponse.map(chunk -> {
String content = (chunk.message() != null) ? chunk.message().content() : "";
List<AssistantMessage.ToolCall> toolCalls = chunk.message().toolCalls() == null ? List.of()
: chunk.message()
.toolCalls()
.stream()
.map(toolCall -> new AssistantMessage.ToolCall("", "function", toolCall.function().name(),
ModelOptionsUtils.toJsonString(toolCall.function().arguments())))
.toList();
Observation observation = ChatModelObservationDocumentation.CHAT_MODEL_OPERATION.observation(
this.observationConvention, DEFAULT_OBSERVATION_CONVENTION, () -> observationContext,
this.observationRegistry);
var assistantMessage = new AssistantMessage(content, Map.of(), toolCalls);
observation.parentObservation(contextView.getOrDefault(ObservationThreadLocalAccessor.KEY, null)).start();
ChatGenerationMetadata generationMetadata = ChatGenerationMetadata.NULL;
if (chunk.promptEvalCount() != null && chunk.evalCount() != null) {
generationMetadata = ChatGenerationMetadata.from(chunk.doneReason(), null);
}
Flux<OllamaApi.ChatResponse> ollamaResponse = this.chatApi.streamingChat(request);
var generator = new Generation(assistantMessage, generationMetadata);
return new ChatResponse(List.of(generator), from(chunk));
});
Flux<ChatResponse> chatResponse = ollamaResponse.map(chunk -> {
String content = (chunk.message() != null) ? chunk.message().content() : "";
List<AssistantMessage.ToolCall> toolCalls = chunk.message().toolCalls() == null ? List.of()
: chunk.message()
.toolCalls()
.stream()
.map(toolCall -> new AssistantMessage.ToolCall("", "function", toolCall.function().name(),
ModelOptionsUtils.toJsonString(toolCall.function().arguments())))
.toList();
return chatResponse.flatMap(response -> {
if (isToolCall(response, Set.of("stop"))) {
var toolCallConversation = handleToolCalls(prompt, response);
// Recursively call the stream method with the tool call message
// conversation that contains the call responses.
return this.stream(new Prompt(toolCallConversation, prompt.getOptions()));
}
else {
return Flux.just(response);
}
var assistantMessage = new AssistantMessage(content, Map.of(), toolCalls);
ChatGenerationMetadata generationMetadata = ChatGenerationMetadata.NULL;
if (chunk.promptEvalCount() != null && chunk.evalCount() != null) {
generationMetadata = ChatGenerationMetadata.from(chunk.doneReason(), null);
}
var generator = new Generation(assistantMessage, generationMetadata);
return new ChatResponse(List.of(generator), from(chunk));
});
// @formatter:off
Flux<ChatResponse> chatResponseFlux = chatResponse.flatMap(response -> {
if (isToolCall(response, Set.of("stop"))) {
var toolCallConversation = handleToolCalls(prompt, response);
// Recursively call the stream method with the tool call message
// conversation that contains the call responses.
return this.stream(new Prompt(toolCallConversation, prompt.getOptions()));
}
else {
return Flux.just(response);
}
})
.doOnError(observation::error)
.doFinally(s -> {
observation.stop();
})
.contextWrite(ctx -> ctx.put(ObservationThreadLocalAccessor.KEY, observation));
// @formatter:on
return new MessageAggregator().aggregate(chatResponseFlux, observationContext::setResponse);
});
}
@@ -216,13 +283,10 @@ public class OllamaChatModel extends AbstractToolCallSupport implements ChatMode
.build());
}
else if (message instanceof ToolResponseMessage toolMessage) {
List<OllamaApi.Message> responseMessages = toolMessage.getResponses()
return toolMessage.getResponses()
.stream()
.map(tr -> OllamaApi.Message.builder(Role.TOOL).withContent(tr.responseData()).build())
.toList();
return responseMessages;
}
throw new IllegalArgumentException("Unsupported message type: " + message.getMessageType());
}).flatMap(List::stream).toList();
@@ -290,9 +354,32 @@ public class OllamaChatModel extends AbstractToolCallSupport implements ChatMode
}).toList();
}
private ChatOptions buildRequestOptions(OllamaApi.ChatRequest request) {
var options = ModelOptionsUtils.mapToClass(request.options(), OllamaOptions.class);
return ChatOptionsBuilder.builder()
.withModel(request.model())
.withFrequencyPenalty(options.getFrequencyPenalty())
.withMaxTokens(options.getMaxTokens())
.withPresencePenalty(options.getPresencePenalty())
.withStopSequences(options.getStopSequences())
.withTemperature(options.getTemperature())
.withTopK(options.getTopK())
.withTopP(options.getTopP())
.build();
}
@Override
public ChatOptions getDefaultOptions() {
return OllamaOptions.fromOptions(this.defaultOptions);
}
/**
* Use the provided convention for reporting observation data
* @param observationConvention The provided convention
*/
public void setObservationConvention(ChatModelObservationConvention observationConvention) {
Assert.notNull(observationConvention, "observationConvention cannot be null");
this.observationConvention = observationConvention;
}
}

91
models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/OllamaEmbeddingModel.java
View File

@@ -21,17 +21,14 @@ import java.util.concurrent.atomic.AtomicInteger;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import io.micrometer.observation.ObservationRegistry;
import org.springframework.ai.chat.metadata.EmptyUsage;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.AbstractEmbeddingModel;
import org.springframework.ai.embedding.Embedding;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.EmbeddingOptions;
import org.springframework.ai.embedding.EmbeddingRequest;
import org.springframework.ai.embedding.EmbeddingResponse;
import org.springframework.ai.embedding.EmbeddingResponseMetadata;
import org.springframework.ai.embedding.*;
import org.springframework.ai.embedding.observation.DefaultEmbeddingModelObservationConvention;
import org.springframework.ai.embedding.observation.EmbeddingModelObservationContext;
import org.springframework.ai.embedding.observation.EmbeddingModelObservationConvention;
import org.springframework.ai.embedding.observation.EmbeddingModelObservationDocumentation;
import org.springframework.ai.model.ModelOptionsUtils;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaApi.EmbeddingsResponse;
@@ -53,26 +50,47 @@ import org.springframework.util.StringUtils;
* most up-to-date information on available models.
*
* @author Christian Tzolov
* @author Thomas Vitale
* @since 0.8.0
*/
public class OllamaEmbeddingModel extends AbstractEmbeddingModel {
private final Logger logger = LoggerFactory.getLogger(getClass());
private static final EmbeddingModelObservationConvention DEFAULT_OBSERVATION_CONVENTION = new DefaultEmbeddingModelObservationConvention();
private final OllamaApi ollamaApi;
/**
* Default options to be used for all chat requests.
*/
private OllamaOptions defaultOptions = OllamaOptions.create().withModel(OllamaOptions.DEFAULT_MODEL);
private final OllamaOptions defaultOptions;
/**
* Observation registry used for instrumentation.
*/
private final ObservationRegistry observationRegistry;
/**
* Conventions to use for generating observations.
*/
private EmbeddingModelObservationConvention observationConvention = DEFAULT_OBSERVATION_CONVENTION;
public OllamaEmbeddingModel(OllamaApi ollamaApi) {
this.ollamaApi = ollamaApi;
this(ollamaApi, OllamaOptions.create().withModel(OllamaOptions.DEFAULT_MODEL));
}
public OllamaEmbeddingModel(OllamaApi ollamaApi, OllamaOptions defaultOptions) {
this(ollamaApi, defaultOptions, ObservationRegistry.NOOP);
}
public OllamaEmbeddingModel(OllamaApi ollamaApi, OllamaOptions defaultOptions,
ObservationRegistry observationRegistry) {
Assert.notNull(ollamaApi, "openAiApi must not be null");
Assert.notNull(defaultOptions, "options must not be null");
Assert.notNull(observationRegistry, "observationRegistry must not be null");
this.ollamaApi = ollamaApi;
this.defaultOptions = defaultOptions;
this.observationRegistry = observationRegistry;
}
@Override
@@ -82,25 +100,39 @@ public class OllamaEmbeddingModel extends AbstractEmbeddingModel {
@Override
public EmbeddingResponse call(EmbeddingRequest request) {
Assert.notEmpty(request.getInstructions(), "At least one text is required!");
OllamaApi.EmbeddingsRequest ollamaEmbeddingRequest = ollamaEmbeddingRequest(request.getInstructions(),
request.getOptions());
EmbeddingsResponse response = this.ollamaApi.embed(ollamaEmbeddingRequest);
var observationContext = EmbeddingModelObservationContext.builder()
.embeddingRequest(request)
.provider(OllamaApi.PROVIDER_NAME)
.requestOptions(buildRequestOptions(ollamaEmbeddingRequest))
.build();
AtomicInteger indexCounter = new AtomicInteger(0);
return EmbeddingModelObservationDocumentation.EMBEDDING_MODEL_OPERATION
.observation(this.observationConvention, DEFAULT_OBSERVATION_CONVENTION, () -> observationContext,
this.observationRegistry)
.observe(() -> {
EmbeddingsResponse response = this.ollamaApi.embed(ollamaEmbeddingRequest);
List<Embedding> embeddings = response.embeddings()
.stream()
.map(e -> new Embedding(e, indexCounter.getAndIncrement()))
.toList();
AtomicInteger indexCounter = new AtomicInteger(0);
EmbeddingResponseMetadata embeddingResponseMetadata = new EmbeddingResponseMetadata(response.model(),
new EmptyUsage());
List<Embedding> embeddings = response.embeddings()
.stream()
.map(e -> new Embedding(e, indexCounter.getAndIncrement()))
.toList();
return new EmbeddingResponse(embeddings, embeddingResponseMetadata);
EmbeddingResponseMetadata embeddingResponseMetadata = new EmbeddingResponseMetadata(response.model(),
new EmptyUsage());
EmbeddingResponse embeddingResponse = new EmbeddingResponse(embeddings, embeddingResponseMetadata);
observationContext.setResponse(embeddingResponse);
return embeddingResponse;
});
}
/**
@@ -126,9 +158,22 @@ public class OllamaEmbeddingModel extends AbstractEmbeddingModel {
OllamaOptions.filterNonSupportedFields(mergedOptions.toMap()), mergedOptions.getTruncate());
}
private EmbeddingOptions buildRequestOptions(OllamaApi.EmbeddingsRequest request) {
return EmbeddingOptionsBuilder.builder().withModel(request.model()).build();
}
/**
* Use the provided convention for reporting observation data
* @param observationConvention The provided convention
*/
public void setObservationConvention(EmbeddingModelObservationConvention observationConvention) {
Assert.notNull(observationConvention, "observationConvention cannot be null");
this.observationConvention = observationConvention;
}
public static class DurationParser {
private static Pattern PATTERN = Pattern.compile("(\\d+)(ms|s|m|h)");
private static final Pattern PATTERN = Pattern.compile("(\\d+)(ms|s|m|h)");
public static Duration parse(String input) {

5
models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/api/OllamaApi.java
View File

@@ -26,6 +26,7 @@ import java.util.function.Consumer;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.springframework.ai.model.ModelOptionsUtils;
import org.springframework.ai.observation.conventions.AiProvider;
import org.springframework.boot.context.properties.bind.ConstructorBinding;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
@@ -55,7 +56,9 @@ public class OllamaApi {
private static final Log logger = LogFactory.getLog(OllamaApi.class);
private final static String DEFAULT_BASE_URL = "http://localhost:11434";
private static final String DEFAULT_BASE_URL = "http://localhost:11434";
public static final String PROVIDER_NAME = AiProvider.OLLAMA.value();
public static final String REQUEST_BODY_NULL_ERROR = "The request body can not be null.";

9
models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/api/OllamaModel.java
View File

@@ -22,7 +22,7 @@ import org.springframework.ai.model.ChatModelDescription;
*
* @author Siarhei Blashuk
* @author Thomas Vitale
* @since 0.8.1
* @since 1.0.0
*/
public enum OllamaModel implements ChatModelDescription {
@@ -106,7 +106,12 @@ public enum OllamaModel implements ChatModelDescription {
/**
* Uncensored Llama 2 model
*/
LLAMA2_UNCENSORED("llama2-uncensored");
LLAMA2_UNCENSORED("llama2-uncensored"),
/**
* A high-performing open embedding model with a large token context window.
*/
NOMIC_EMBED_TEXT("nomic-embed-text");
private final String id;

26
models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/BaseOllamaIT.java Normal file
View File

@@ -0,0 +1,26 @@
package org.springframework.ai.ollama;
import org.testcontainers.ollama.OllamaContainer;
public class BaseOllamaIT {
public static final OllamaContainer ollamaContainer;
static {
ollamaContainer = new OllamaContainer(OllamaImage.DEFAULT_IMAGE).withReuse(true);
ollamaContainer.start();
}
/**
* Change the value to false in order to run multiple Ollama IT tests locally reusing
* the same container image Also add the entry
*
* testcontainers.reuse.enable=true
*
* to the file .testcontainers.properties located in your home directory
*/
public static boolean isDisabled() {
return true;
}
}

25
models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/OllamaChatModelFunctionCallingIT.java
View File

@@ -15,16 +15,10 @@
*/
package org.springframework.ai.ollama;
import static org.assertj.core.api.Assertions.assertThat;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Disabled;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.DisabledIf;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.messages.AssistantMessage;
@@ -43,24 +37,25 @@ import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.context.annotation.Bean;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.testcontainers.ollama.OllamaContainer;
import reactor.core.publisher.Flux;
@Disabled("For manual smoke testing only.")
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;
import static org.assertj.core.api.Assertions.assertThat;
@Testcontainers
@SpringBootTest(classes = OllamaChatModelFunctionCallingIT.Config.class)
class OllamaChatModelFunctionCallingIT {
@DisabledIf("isDisabled")
class OllamaChatModelFunctionCallingIT extends BaseOllamaIT {
private static final Logger logger = LoggerFactory.getLogger(OllamaChatModelFunctionCallingIT.class);
private static final String MODEL = OllamaModel.MISTRAL.getName();
@Container
static OllamaContainer ollamaContainer = new OllamaContainer(OllamaImage.DEFAULT_IMAGE);
static String baseUrl = "http://localhost:11434";
@BeforeAll

27
models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/OllamaChatModelIT.java
View File

@@ -15,19 +15,11 @@
*/
package org.springframework.ai.ollama;
import static org.assertj.core.api.Assertions.assertThat;
import java.io.IOException;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Disabled;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.DisabledIf;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
@@ -49,22 +41,25 @@ import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.context.annotation.Bean;
import org.springframework.core.convert.support.DefaultConversionService;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.testcontainers.ollama.OllamaContainer;
import java.io.IOException;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import static org.assertj.core.api.Assertions.assertThat;
@SpringBootTest
@Testcontainers
@Disabled("For manual smoke testing only.")
class OllamaChatModelIT {
@DisabledIf("isDisabled")
class OllamaChatModelIT extends BaseOllamaIT {
private static final String MODEL = OllamaModel.MISTRAL.getName();
private static final Log logger = LogFactory.getLog(OllamaChatModelIT.class);
@Container
static OllamaContainer ollamaContainer = new OllamaContainer(OllamaImage.DEFAULT_IMAGE);
static String baseUrl = "http://localhost:11434";
@BeforeAll

23
models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/OllamaChatModelMultimodalIT.java
View File

@@ -15,23 +15,16 @@
*/
package org.springframework.ai.ollama;
import java.io.IOException;
import java.util.List;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Disabled;
import org.junit.jupiter.api.Test;
import org.springframework.ai.ollama.api.OllamaModel;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.testcontainers.ollama.OllamaContainer;
import org.springframework.ai.model.Media;
import org.junit.jupiter.api.condition.DisabledIf;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.Media;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaModel;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.SpringBootConfiguration;
@@ -39,13 +32,19 @@ import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.context.annotation.Bean;
import org.springframework.core.io.ClassPathResource;
import org.springframework.util.MimeTypeUtils;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.testcontainers.ollama.OllamaContainer;
import java.io.IOException;
import java.util.List;
import static org.assertj.core.api.Assertions.assertThat;
@SpringBootTest
@Testcontainers
@Disabled("For manual smoke testing only.")
class OllamaChatModelMultimodalIT {
@DisabledIf("isDisabled")
class OllamaChatModelMultimodalIT extends BaseOllamaIT {
private static final String MODEL = OllamaModel.MOONDREAM.getName();

177
models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/OllamaChatModelObservationIT.java Normal file
View File

@@ -0,0 +1,177 @@
/*
* Copyright 2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.ollama;
import io.micrometer.observation.tck.TestObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistryAssert;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.DisabledIf;
import org.springframework.ai.chat.metadata.ChatResponseMetadata;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.observation.DefaultChatModelObservationConvention;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.function.FunctionCallbackContext;
import org.springframework.ai.observation.conventions.AiOperationType;
import org.springframework.ai.observation.conventions.AiProvider;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaModel;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.context.annotation.Bean;
import reactor.core.publisher.Flux;
import java.util.List;
import java.util.stream.Collectors;
import static org.assertj.core.api.Assertions.assertThat;
import static org.springframework.ai.chat.observation.ChatModelObservationDocumentation.HighCardinalityKeyNames;
import static org.springframework.ai.chat.observation.ChatModelObservationDocumentation.LowCardinalityKeyNames;
/**
* Integration tests for observation instrumentation in {@link OllamaChatModel}.
*
* @author Thomas Vitale
*/
@SpringBootTest(classes = OllamaChatModelObservationIT.Config.class)
@DisabledIf("isDisabled")
public class OllamaChatModelObservationIT extends BaseOllamaIT {
@Autowired
TestObservationRegistry observationRegistry;
@Autowired
OllamaChatModel chatModel;
@BeforeEach
void beforeEach() {
observationRegistry.clear();
}
@Test
void observationForChatOperation() {
var options = OllamaOptions.builder()
.withModel(OllamaModel.MISTRAL.getName())
.withFrequencyPenalty(0f)
.withNumPredict(2048)
.withPresencePenalty(0f)
.withStop(List.of("this-is-the-end"))
.withTemperature(0.7f)
.withTopK(1)
.withTopP(1f)
.build();
Prompt prompt = new Prompt("Why does a raven look like a desk?", options);
ChatResponse chatResponse = chatModel.call(prompt);
assertThat(chatResponse.getResult().getOutput().getContent()).isNotEmpty();
ChatResponseMetadata responseMetadata = chatResponse.getMetadata();
assertThat(responseMetadata).isNotNull();
validate(responseMetadata);
}
@Test
void observationForStreamingChatOperation() {
var options = OllamaOptions.builder()
.withModel(OllamaModel.MISTRAL.getName())
.withFrequencyPenalty(0f)
.withNumPredict(2048)
.withPresencePenalty(0f)
.withStop(List.of("this-is-the-end"))
.withTemperature(0.7f)
.withTopK(1)
.withTopP(1f)
.build();
Prompt prompt = new Prompt("Why does a raven look like a desk?", options);
Flux<ChatResponse> chatResponseFlux = chatModel.stream(prompt);
List<ChatResponse> responses = chatResponseFlux.collectList().block();
assertThat(responses).isNotEmpty();
assertThat(responses).hasSizeGreaterThan(10);
String aggregatedResponse = responses.subList(0, responses.size() - 1)
.stream()
.map(r -> r.getResult().getOutput().getContent())
.collect(Collectors.joining());
assertThat(aggregatedResponse).isNotEmpty();
ChatResponse lastChatResponse = responses.get(responses.size() - 1);
ChatResponseMetadata responseMetadata = lastChatResponse.getMetadata();
assertThat(responseMetadata).isNotNull();
validate(responseMetadata);
}
private void validate(ChatResponseMetadata responseMetadata) {
TestObservationRegistryAssert.assertThat(observationRegistry)
.doesNotHaveAnyRemainingCurrentObservation()
.hasObservationWithNameEqualTo(DefaultChatModelObservationConvention.DEFAULT_NAME)
.that()
.hasContextualNameEqualTo("chat " + OllamaModel.MISTRAL.getName())
.hasLowCardinalityKeyValue(LowCardinalityKeyNames.AI_OPERATION_TYPE.asString(),
AiOperationType.CHAT.value())
.hasLowCardinalityKeyValue(LowCardinalityKeyNames.AI_PROVIDER.asString(), AiProvider.OLLAMA.value())
.hasLowCardinalityKeyValue(LowCardinalityKeyNames.REQUEST_MODEL.asString(), OllamaModel.MISTRAL.getName())
.hasLowCardinalityKeyValue(LowCardinalityKeyNames.RESPONSE_MODEL.asString(), responseMetadata.getModel())
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.REQUEST_FREQUENCY_PENALTY.asString(), "0.0")
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.REQUEST_MAX_TOKENS.asString(), "2048")
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.REQUEST_PRESENCE_PENALTY.asString(), "0.0")
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.REQUEST_STOP_SEQUENCES.asString(),
"[\"this-is-the-end\"]")
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.REQUEST_TEMPERATURE.asString(), "0.7")
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.REQUEST_TOP_K.asString(), "1")
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.REQUEST_TOP_P.asString(), "1.0")
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.RESPONSE_ID.asString(), responseMetadata.getId())
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.RESPONSE_FINISH_REASONS.asString(), "[\"stop\"]")
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.USAGE_INPUT_TOKENS.asString(),
String.valueOf(responseMetadata.getUsage().getPromptTokens()))
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.USAGE_OUTPUT_TOKENS.asString(),
String.valueOf(responseMetadata.getUsage().getGenerationTokens()))
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.USAGE_TOTAL_TOKENS.asString(),
String.valueOf(responseMetadata.getUsage().getTotalTokens()))
.hasBeenStarted()
.hasBeenStopped();
}
@SpringBootConfiguration
static class Config {
@Bean
public TestObservationRegistry observationRegistry() {
return TestObservationRegistry.create();
}
@Bean
public OllamaApi openAiApi() {
return new OllamaApi();
}
@Bean
public OllamaChatModel openAiChatModel(OllamaApi openAiApi, TestObservationRegistry observationRegistry) {
return new OllamaChatModel(openAiApi, OllamaOptions.create(), new FunctionCallbackContext(), List.of(),
observationRegistry);
}
}
}

1
models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/OllamaChatRequestTests.java
View File

@@ -16,7 +16,6 @@
package org.springframework.ai.ollama;
import org.junit.jupiter.api.Test;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.ChatOptionsBuilder;
import org.springframework.ai.chat.prompt.Prompt;

23
models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/OllamaEmbeddingModelIT.java
View File

@@ -15,41 +15,40 @@
*/
package org.springframework.ai.ollama;
import java.io.IOException;
import java.util.List;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Disabled;
import org.junit.jupiter.api.Test;
import org.springframework.ai.ollama.api.OllamaModel;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.testcontainers.ollama.OllamaContainer;
import org.junit.jupiter.api.condition.DisabledIf;
import org.springframework.ai.embedding.EmbeddingRequest;
import org.springframework.ai.embedding.EmbeddingResponse;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaApiIT;
import org.springframework.ai.ollama.api.OllamaModel;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.context.annotation.Bean;
import org.testcontainers.junit.jupiter.Testcontainers;
import java.io.IOException;
import java.util.List;
import static org.assertj.core.api.Assertions.assertThat;
@SpringBootTest
@Disabled("For manual smoke testing only.")
@DisabledIf("isDisabled")
@Testcontainers
class OllamaEmbeddingModelIT {
class OllamaEmbeddingModelIT extends BaseOllamaIT {
private static final String MODEL = OllamaModel.MISTRAL.getName();
private static final Log logger = LogFactory.getLog(OllamaApiIT.class);
@Container
static OllamaContainer ollamaContainer = new OllamaContainer(OllamaImage.DEFAULT_IMAGE);
// @Container
// static OllamaContainer ollamaContainer = new
// OllamaContainer(OllamaImage.DEFAULT_IMAGE);
static String baseUrl = "http://localhost:11434";

112
models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/OllamaEmbeddingModelObservationIT.java Normal file
View File

@@ -0,0 +1,112 @@
/*
* Copyright 2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.ollama;
import io.micrometer.common.KeyValue;
import io.micrometer.observation.tck.TestObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistryAssert;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.DisabledIf;
import org.springframework.ai.embedding.EmbeddingRequest;
import org.springframework.ai.embedding.EmbeddingResponse;
import org.springframework.ai.embedding.EmbeddingResponseMetadata;
import org.springframework.ai.embedding.observation.DefaultEmbeddingModelObservationConvention;
import org.springframework.ai.observation.conventions.AiOperationType;
import org.springframework.ai.observation.conventions.AiProvider;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaModel;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.context.annotation.Bean;
import java.util.List;
import static org.assertj.core.api.Assertions.assertThat;
import static org.springframework.ai.embedding.observation.EmbeddingModelObservationDocumentation.HighCardinalityKeyNames;
import static org.springframework.ai.embedding.observation.EmbeddingModelObservationDocumentation.LowCardinalityKeyNames;
/**
* Integration tests for observation instrumentation in {@link OllamaEmbeddingModel}.
*
* @author Thomas Vitale
*/
@SpringBootTest(classes = OllamaEmbeddingModelObservationIT.Config.class)
@DisabledIf("isDisabled")
public class OllamaEmbeddingModelObservationIT extends BaseOllamaIT {
@Autowired
TestObservationRegistry observationRegistry;
@Autowired
OllamaEmbeddingModel embeddingModel;
@Test
void observationForEmbeddingOperation() {
var options = OllamaOptions.builder().withModel(OllamaModel.NOMIC_EMBED_TEXT.getName()).build();
EmbeddingRequest embeddingRequest = new EmbeddingRequest(List.of("Here comes the sun"), options);
EmbeddingResponse embeddingResponse = embeddingModel.call(embeddingRequest);
assertThat(embeddingResponse.getResults()).isNotEmpty();
EmbeddingResponseMetadata responseMetadata = embeddingResponse.getMetadata();
assertThat(responseMetadata).isNotNull();
TestObservationRegistryAssert.assertThat(observationRegistry)
.doesNotHaveAnyRemainingCurrentObservation()
.hasObservationWithNameEqualTo(DefaultEmbeddingModelObservationConvention.DEFAULT_NAME)
.that()
.hasContextualNameEqualTo("embedding " + OllamaModel.NOMIC_EMBED_TEXT.getName())
.hasLowCardinalityKeyValue(LowCardinalityKeyNames.AI_OPERATION_TYPE.asString(),
AiOperationType.EMBEDDING.value())
.hasLowCardinalityKeyValue(LowCardinalityKeyNames.AI_PROVIDER.asString(), AiProvider.OLLAMA.value())
.hasLowCardinalityKeyValue(LowCardinalityKeyNames.REQUEST_MODEL.asString(),
OllamaModel.NOMIC_EMBED_TEXT.getName())
.hasLowCardinalityKeyValue(LowCardinalityKeyNames.RESPONSE_MODEL.asString(), responseMetadata.getModel())
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.REQUEST_EMBEDDING_DIMENSIONS.asString(),
KeyValue.NONE_VALUE)
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.USAGE_INPUT_TOKENS.asString(),
String.valueOf(responseMetadata.getUsage().getPromptTokens()))
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.USAGE_TOTAL_TOKENS.asString(),
String.valueOf(responseMetadata.getUsage().getTotalTokens()))
.hasBeenStarted()
.hasBeenStopped();
}
@SpringBootConfiguration
static class Config {
@Bean
public TestObservationRegistry observationRegistry() {
return TestObservationRegistry.create();
}
@Bean
public OllamaApi openAiApi() {
return new OllamaApi();
}
@Bean
public OllamaEmbeddingModel openAiEmbeddingModel(OllamaApi openAiApi,
TestObservationRegistry observationRegistry) {
return new OllamaEmbeddingModel(openAiApi, OllamaOptions.builder().build(), observationRegistry);
}
}
}

18
models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/OllamaEmbeddingModelTests.java
View File

@@ -15,20 +15,13 @@
*/
package org.springframework.ai.ollama;
import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.Mockito.when;
import java.time.Duration;
import java.util.List;
import java.util.Map;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.ArgumentCaptor;
import org.mockito.Captor;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;
import org.springframework.ai.embedding.EmbeddingOptions;
import org.springframework.ai.embedding.EmbeddingOptionsBuilder;
import org.springframework.ai.embedding.EmbeddingRequest;
import org.springframework.ai.embedding.EmbeddingResponse;
import org.springframework.ai.embedding.EmbeddingResultMetadata;
@@ -37,6 +30,13 @@ import org.springframework.ai.ollama.api.OllamaApi.EmbeddingsRequest;
import org.springframework.ai.ollama.api.OllamaApi.EmbeddingsResponse;
import org.springframework.ai.ollama.api.OllamaOptions;
import java.time.Duration;
import java.util.List;
import java.util.Map;
import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.Mockito.when;
/**
* @author Christian Tzolov
* @since 1.0.0
@@ -65,7 +65,7 @@ public class OllamaEmbeddingModelTests {
var embeddingModel = new OllamaEmbeddingModel(ollamaApi, defaultOptions);
EmbeddingResponse response = embeddingModel
.call(new EmbeddingRequest(List.of("Input1", "Input2", "Input3"), EmbeddingOptions.EMPTY));
.call(new EmbeddingRequest(List.of("Input1", "Input2", "Input3"), EmbeddingOptionsBuilder.builder().build()));
assertThat(response.getResults()).hasSize(2);
assertThat(response.getResults().get(0).getIndex()).isEqualTo(0);

11
models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/OllamaEmbeddingRequestTests.java
View File

@@ -19,22 +19,23 @@ import org.junit.jupiter.api.Test;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaOptions;
import static org.assertj.core.api.Assertions.assertThat;
import java.util.List;
import static org.assertj.core.api.Assertions.assertThat;
/**
* @author Christian Tzolov
* @author Thomas Vitale
*/
public class OllamaEmbeddingRequestTests {
OllamaEmbeddingModel chatModel = new OllamaEmbeddingModel(new OllamaApi(),
OllamaEmbeddingModel embeddingModel = new OllamaEmbeddingModel(new OllamaApi(),
new OllamaOptions().withModel("DEFAULT_MODEL").withMainGPU(11).withUseMMap(true).withNumGPU(1));
@Test
public void ollamaEmbeddingRequestDefaultOptions() {
var request = chatModel.ollamaEmbeddingRequest(List.of("Hello"), null);
var request = embeddingModel.ollamaEmbeddingRequest(List.of("Hello"), null);
assertThat(request.model()).isEqualTo("DEFAULT_MODEL");
assertThat(request.options().get("num_gpu")).isEqualTo(1);
@@ -52,7 +53,7 @@ public class OllamaEmbeddingRequestTests {
.withUseMMap(true)//
.withNumGPU(2);
var request = chatModel.ollamaEmbeddingRequest(List.of("Hello"), promptOptions);
var request = embeddingModel.ollamaEmbeddingRequest(List.of("Hello"), promptOptions);
assertThat(request.model()).isEqualTo("PROMPT_MODEL");
assertThat(request.options().get("num_gpu")).isEqualTo(2);

2
models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/OllamaImage.java
View File

@@ -22,6 +22,6 @@ import org.testcontainers.utility.DockerImageName;
*/
public class OllamaImage {
public static final DockerImageName DEFAULT_IMAGE = DockerImageName.parse("ollama/ollama:0.2.8");
public static final DockerImageName DEFAULT_IMAGE = DockerImageName.parse("ollama/ollama:0.3.6");
}

34
models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/api/OllamaApiIT.java
View File

@@ -15,18 +15,12 @@
*/
package org.springframework.ai.ollama.api;
import static org.assertj.core.api.Assertions.assertThat;
import java.io.IOException;
import java.util.List;
import java.util.stream.Collectors;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Disabled;
import org.junit.jupiter.api.Test;
import org.springframework.ai.ollama.OllamaImage;
import org.junit.jupiter.api.condition.DisabledIf;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.ollama.BaseOllamaIT;
import org.springframework.ai.ollama.api.OllamaApi.ChatRequest;
import org.springframework.ai.ollama.api.OllamaApi.ChatResponse;
import org.springframework.ai.ollama.api.OllamaApi.EmbeddingsRequest;
@@ -35,26 +29,28 @@ import org.springframework.ai.ollama.api.OllamaApi.GenerateRequest;
import org.springframework.ai.ollama.api.OllamaApi.GenerateResponse;
import org.springframework.ai.ollama.api.OllamaApi.Message;
import org.springframework.ai.ollama.api.OllamaApi.Message.Role;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.testcontainers.ollama.OllamaContainer;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Flux;;
import java.io.IOException;
import java.util.List;
import java.util.stream.Collectors;
import static org.assertj.core.api.Assertions.assertThat;
;
/**
* @author Christian Tzolov
* @author Thomas Vitale
*/
@Disabled("For manual smoke testing only.")
@Testcontainers
public class OllamaApiIT {
@DisabledIf("isDisabled")
public class OllamaApiIT extends BaseOllamaIT {
private static final String MODEL = OllamaModel.ORCA_MINI.getName();
private static final Log logger = LogFactory.getLog(OllamaApiIT.class);
@Container
static OllamaContainer ollamaContainer = new OllamaContainer(OllamaImage.DEFAULT_IMAGE);
private static final Logger logger = LoggerFactory.getLogger(OllamaApiIT.class);
static OllamaApi ollamaApi;

4
models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/api/OllamaModelOptionsTests.java
View File

@@ -15,10 +15,10 @@
*/
package org.springframework.ai.ollama.api;
import java.util.List;
import org.junit.jupiter.api.Test;
import java.util.List;
import static org.assertj.core.api.Assertions.assertThat;
/**

4
models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/api/tool/MockWeatherService.java
View File

@@ -15,14 +15,14 @@
*/
package org.springframework.ai.ollama.api.tool;
import java.util.function.Function;
import com.fasterxml.jackson.annotation.JsonClassDescription;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonInclude.Include;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.annotation.JsonPropertyDescription;
import java.util.function.Function;
/**
* @author Christian Tzolov
*/

27
models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/api/tool/OllamaApiToolFunctionCallIT.java
View File

@@ -16,36 +16,34 @@
package org.springframework.ai.ollama.api.tool;
import static org.assertj.core.api.Assertions.assertThat;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Disabled;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.DisabledIf;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.model.ModelOptionsUtils;
import org.springframework.ai.ollama.OllamaImage;
import org.springframework.ai.ollama.BaseOllamaIT;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaApi.ChatResponse;
import org.springframework.ai.ollama.api.OllamaApi.Message;
import org.springframework.ai.ollama.api.OllamaApi.Message.Role;
import org.springframework.ai.ollama.api.OllamaApi.Message.ToolCall;
import org.springframework.ai.ollama.api.OllamaModel;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.testcontainers.ollama.OllamaContainer;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import static org.assertj.core.api.Assertions.assertThat;
/**
* @author Christian Tzolov
*/
@Disabled("For manual smoke testing only.")
@Testcontainers
public class OllamaApiToolFunctionCallIT {
@DisabledIf("isDisabled")
public class OllamaApiToolFunctionCallIT extends BaseOllamaIT {
private static final String MODEL = OllamaModel.MISTRAL.getName();
@@ -53,9 +51,6 @@ public class OllamaApiToolFunctionCallIT {
MockWeatherService weatherService = new MockWeatherService();
@Container
static OllamaContainer ollamaContainer = new OllamaContainer(OllamaImage.DEFAULT_IMAGE);
static String baseUrl = "http://localhost:11434";
@BeforeAll

2
models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/OllamaWithOpenAiChatModelIT.java
View File

@@ -75,7 +75,7 @@ class OllamaWithOpenAiChatModelIT {
private static final String DEFAULT_OLLAMA_MODEL = "mistral";
@Container
static OllamaContainer ollamaContainer = new OllamaContainer("ollama/ollama:0.2.8");
static OllamaContainer ollamaContainer = new OllamaContainer("ollama/ollama:0.3.6");
static String baseUrl = "http://localhost:11434";

21
spring-ai-core/src/main/java/org/springframework/ai/converter/BeanOutputConverter.java
View File

@@ -168,8 +168,25 @@ public class BeanOutputConverter<T> implements StructuredOutputConverter<T> {
*/
public T convert(@NonNull String text) {
try {
if (text.startsWith("```json") && text.endsWith("```")) {
text = text.substring(7, text.length() - 3);
// Remove leading and trailing whitespace
text = text.trim();
// Check for and remove triple backticks and "json" identifier
if (text.startsWith("```") && text.endsWith("```")) {
// Remove the first line if it contains "```json"
String[] lines = text.split("\n", 2);
if (lines[0].trim().equalsIgnoreCase("```json")) {
text = lines.length > 1 ? lines[1] : "";
}
else {
text = text.substring(3); // Remove leading ```
}
// Remove trailing ```
text = text.substring(0, text.length() - 3);
// Trim again to remove any potential whitespace
text = text.trim();
}
return (T) this.objectMapper.readValue(text, this.typeRef);
}

25
spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/ollama/OllamaAutoConfiguration.java
View File

@@ -17,11 +17,15 @@ package org.springframework.ai.autoconfigure.ollama;
import java.util.List;
import io.micrometer.observation.ObservationRegistry;
import org.springframework.ai.chat.observation.ChatModelObservationConvention;
import org.springframework.ai.embedding.observation.EmbeddingModelObservationConvention;
import org.springframework.ai.model.function.FunctionCallback;
import org.springframework.ai.model.function.FunctionCallbackContext;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.ollama.OllamaEmbeddingModel;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.boot.autoconfigure.AutoConfiguration;
import org.springframework.boot.autoconfigure.ImportAutoConfiguration;
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
@@ -39,6 +43,7 @@ import org.springframework.web.client.RestClient;
*
* @author Christian Tzolov
* @author Eddú Meléndez
* @author Thomas Vitale
* @since 0.8.0
*/
@AutoConfiguration(after = RestClientAutoConfiguration.class)
@@ -65,18 +70,30 @@ public class OllamaAutoConfiguration {
@ConditionalOnProperty(prefix = OllamaChatProperties.CONFIG_PREFIX, name = "enabled", havingValue = "true",
matchIfMissing = true)
public OllamaChatModel ollamaChatModel(OllamaApi ollamaApi, OllamaChatProperties properties,
List<FunctionCallback> toolFunctionCallbacks, FunctionCallbackContext functionCallbackContext) {
List<FunctionCallback> toolFunctionCallbacks, FunctionCallbackContext functionCallbackContext,
ObjectProvider<ObservationRegistry> observationRegistry,
ObjectProvider<ChatModelObservationConvention> observationConvention) {
var chatModel = new OllamaChatModel(ollamaApi, properties.getOptions(), functionCallbackContext,
toolFunctionCallbacks, observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP));
return new OllamaChatModel(ollamaApi, properties.getOptions(), functionCallbackContext, toolFunctionCallbacks);
observationConvention.ifAvailable(chatModel::setObservationConvention);
return chatModel;
}
@Bean
@ConditionalOnMissingBean
@ConditionalOnProperty(prefix = OllamaEmbeddingProperties.CONFIG_PREFIX, name = "enabled", havingValue = "true",
matchIfMissing = true)
public OllamaEmbeddingModel ollamaEmbeddingModel(OllamaApi ollamaApi, OllamaEmbeddingProperties properties) {
public OllamaEmbeddingModel ollamaEmbeddingModel(OllamaApi ollamaApi, OllamaEmbeddingProperties properties,
ObjectProvider<ObservationRegistry> observationRegistry,
ObjectProvider<EmbeddingModelObservationConvention> observationConvention) {
var embeddingModel = new OllamaEmbeddingModel(ollamaApi, properties.getOptions(),
observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP));
return new OllamaEmbeddingModel(ollamaApi, properties.getOptions());
observationConvention.ifAvailable(embeddingModel::setObservationConvention);
return embeddingModel;
}
static class PropertiesOllamaConnectionDetails implements OllamaConnectionDetails {

4
spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/ollama/OllamaChatAutoConfigurationIT.java
View File

@@ -60,11 +60,11 @@ public class OllamaChatAutoConfigurationIT {
private static final Log logger = LogFactory.getLog(OllamaChatAutoConfigurationIT.class);
private static String MODEL_NAME = "mistral";
private static final String MODEL_NAME = "mistral";
private static final String OLLAMA_WITH_MODEL = "%s-%s".formatted(MODEL_NAME, OllamaImage.IMAGE);
private static final OllamaContainer ollamaContainer;
private static OllamaContainer ollamaContainer;
static {
ollamaContainer = new OllamaContainer(OllamaDockerImageName.image());

11
spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/ollama/OllamaEmbeddingAutoConfigurationIT.java
View File

@@ -18,11 +18,11 @@ package org.springframework.ai.autoconfigure.ollama;
import java.io.IOException;
import java.util.List;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Disabled;
import org.junit.jupiter.api.Test;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
@@ -37,20 +37,21 @@ import static org.assertj.core.api.Assertions.assertThat;
/**
* @author Christian Tzolov
* @author Thomas Vitale
* @since 0.8.0
*/
@Disabled("For manual smoke testing only.")
@Testcontainers
public class OllamaEmbeddingAutoConfigurationIT {
private static final Log logger = LogFactory.getLog(OllamaEmbeddingAutoConfigurationIT.class);
private static final Logger logger = LoggerFactory.getLogger(OllamaEmbeddingAutoConfigurationIT.class);
private static String MODEL_NAME = "orca-mini";
private static final String MODEL_NAME = "orca-mini";
@Container
static OllamaContainer ollamaContainer = new OllamaContainer(OllamaImage.IMAGE);
static String baseUrl;
static String baseUrl = "http://localhost:11434";
@BeforeAll
public static void beforeAll() throws IOException, InterruptedException {

2
spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/ollama/OllamaImage.java
View File

@@ -17,6 +17,6 @@ package org.springframework.ai.autoconfigure.ollama;
public class OllamaImage {
public static final String IMAGE = "ollama/ollama:0.2.8";
public static final String IMAGE = "ollama/ollama:0.3.6";
}

2
spring-ai-spring-boot-testcontainers/src/test/java/org/springframework/ai/testcontainers/service/connection/ollama/OllamaContainerConnectionDetailsFactoryTest.java
View File

@@ -53,7 +53,7 @@ class OllamaContainerConnectionDetailsFactoryTest {
@Container
@ServiceConnection
static OllamaContainer ollama = new OllamaContainer("ollama/ollama:0.3.2");
static OllamaContainer ollama = new OllamaContainer("ollama/ollama:0.3.6");
@Autowired
private OllamaEmbeddingModel embeddingModel;