Skip to content

Commit

Permalink
Support Ollama JSON Structured Output
Browse files Browse the repository at this point in the history
Ollama has recently introduced native support for JSON structured output, as described in https://ollama.com/blog/structured-outputs.
This PR introduces support for it, both for directly passing a JSON schema and when using the Spring AI output conversion APIs.

Signed-off-by: Thomas Vitale <[email protected]>
  • Loading branch information
ThomasVitale authored and tzolov committed Dec 18, 2024
1 parent be9d6a1 commit 6ab7e20
Show file tree
Hide file tree
Showing 8 changed files with 133 additions and 12 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -25,11 +25,9 @@
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Consumer;

import com.fasterxml.jackson.annotation.JsonFormat;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonInclude.Include;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.annotation.JsonFormat.Feature;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import reactor.core.publisher.Flux;
Expand Down Expand Up @@ -418,7 +416,7 @@ public Message build() {
* @param model The model to use for completion. It should be a name familiar to Ollama from the <a href="https://ollama.com/library">Library</a>.
* @param messages The list of messages in the chat. This can be used to keep a chat memory.
* @param stream Whether to stream the response. If false, the response will be returned as a single response object rather than a stream of objects.
* @param format The format to return the response in. Currently, the only accepted value is "json".
* @param format The format to return the response in. It can either be the String "json" or a Map containing a JSON Schema definition.
* @param keepAlive Controls how long the model will stay loaded into memory following this request (default: 5m).
* @param tools List of tools the model has access to.
* @param options Model-specific options. For example, "temperature" can be set through this field, if the model supports it.
Expand All @@ -435,7 +433,7 @@ public record ChatRequest(
@JsonProperty("model") String model,
@JsonProperty("messages") List<Message> messages,
@JsonProperty("stream") Boolean stream,
@JsonProperty("format") String format,
@JsonProperty("format") Object format,
@JsonProperty("keep_alive") String keepAlive,
@JsonProperty("tools") List<Tool> tools,
@JsonProperty("options") Map<String, Object> options
Expand Down Expand Up @@ -507,7 +505,7 @@ public static class Builder {
private final String model;
private List<Message> messages = List.of();
private boolean stream = false;
private String format;
private Object format;
private String keepAlive;
private List<Tool> tools = List.of();
private Map<String, Object> options = Map.of();
Expand All @@ -527,7 +525,7 @@ public Builder withStream(boolean stream) {
return this;
}

public Builder withFormat(String format) {
public Builder withFormat(Object format) {
this.format = format;
return this;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -287,7 +287,7 @@ public class OllamaOptions implements FunctionCallingOptions, EmbeddingOptions {
* Part of Chat completion <a href="https://github.com/ollama/ollama/blob/main/docs/api.md#parameters-1">advanced parameters</a>.
*/
@JsonProperty("format")
private String format;
private Object format;

/**
* Sets the length of time for Ollama to keep the model loaded. Valid values for this
Expand Down Expand Up @@ -411,7 +411,7 @@ public OllamaOptions withModel(OllamaModel model) {
return this;
}

public OllamaOptions withFormat(String format) {
public OllamaOptions withFormat(Object format) {
this.format = format;
return this;
}
Expand Down Expand Up @@ -614,11 +614,11 @@ public void setModel(String model) {
this.model = model;
}

public String getFormat() {
public Object getFormat() {
return this.format;
}

public void setFormat(String format) {
public void setFormat(Object format) {
this.format = format;
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
import java.util.Map;
import java.util.stream.Collectors;

import com.fasterxml.jackson.annotation.JsonProperty;
import org.junit.jupiter.api.Test;

import org.springframework.ai.chat.client.ChatClient;
Expand Down Expand Up @@ -228,6 +229,31 @@ void beanStreamOutputConverterRecords() {
assertThat(actorsFilms.movies()).hasSize(5);
}

// Example inspired by https://ollama.com/blog/structured-outputs
@Test
void jsonSchemaFormatStructuredOutput() {
var outputConverter = new BeanOutputConverter<>(CountryInfo.class);
var userPromptTemplate = new PromptTemplate("""
Tell me about {country}.
""");
Map<String, Object> model = Map.of("country", "denmark");
var prompt = userPromptTemplate.create(model,
OllamaOptions.builder()
.withModel(OllamaModel.LLAMA3_2.getName())
.withFormat(outputConverter.getJsonSchemaMap())
.build());

var chatResponse = this.chatModel.call(prompt);

var countryInfo = outputConverter.convert(chatResponse.getResult().getOutput().getText());
assertThat(countryInfo).isNotNull();
assertThat(countryInfo.capital()).isEqualToIgnoringCase("Copenhagen");
}

record CountryInfo(@JsonProperty(required = true) String name, @JsonProperty(required = true) String capital,
@JsonProperty(required = true) List<String> languages) {
}

record ActorsFilmsRecord(String actor, List<String> movies) {

}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,7 @@ void multiModalityTest() {
var response = this.chatModel.call(new Prompt(List.of(userMessage)));

logger.info(response.getResult().getOutput().getText());
assertThat(response.getResult().getOutput().getText()).contains("bananas", "apple");
assertThat(response.getResult().getOutput().getText()).containsAnyOf("bananas", "apple");
}

@SpringBootConfiguration
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@
*/
public final class OllamaImage {

public static final DockerImageName DEFAULT_IMAGE = DockerImageName.parse("ollama/ollama:0.5.1");
public static final DockerImageName DEFAULT_IMAGE = DockerImageName.parse("ollama/ollama:0.5.2");

private OllamaImage() {

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -99,6 +99,7 @@ public void pullAndDeleteModelFromOllama() {
assertThat(isModelWithLatestVersionAvailable).isFalse();
}

@Disabled
@Test
public void pullAndDeleteModelFromHuggingFace() {
// Pull model with explicit version.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
package org.springframework.ai.converter;

import java.lang.reflect.Type;
import java.util.Map;
import java.util.Objects;

import com.fasterxml.jackson.core.JsonProcessingException;
Expand Down Expand Up @@ -54,6 +55,7 @@
* @author Josh Long
* @author Sebastien Deleuze
* @author Soby Chacko
* @author Thomas Vitale
*/
public class BeanOutputConverter<T> implements StructuredOutputConverter<T> {

Expand Down Expand Up @@ -220,4 +222,14 @@ public String getJsonSchema() {
return this.jsonSchema;
}

public Map<String, Object> getJsonSchemaMap() {
try {
return this.objectMapper.readValue(this.jsonSchema, Map.class);
}
catch (JsonProcessingException ex) {
logger.error("Could not parse the JSON Schema to a Map object", ex);
throw new IllegalStateException(ex);
}
}

}
Original file line number Diff line number Diff line change
Expand Up @@ -269,6 +269,90 @@ photo was taken in an area with metallic decorations or fixtures. The overall se
where fruits are being displayed, possibly for convenience or aesthetic purposes.
----

== Structured Outputs

Ollama provides custom https://ollama.com/blog/structured-outputs[Structured Outputs] APIs that ensure your model generates responses conforming strictly to your provided `JSON Schema`.
In addition to the existing Spring AI model-agnostic xref::api/structured-output-converter.adoc[Structured Output Converter], these APIs offer enhanced control and precision.

=== Configuration

Spring AI allows you to configure your response format programmatically using the `OllamaOptions` builder.

==== Using the Chat Options Builder

You can set the response format programmatically with the `OllamaOptions` builder as shown below:

[source,java]
----
String jsonSchema = """
{
"type": "object",
"properties": {
"steps": {
"type": "array",
"items": {
"type": "object",
"properties": {
"explanation": { "type": "string" },
"output": { "type": "string" }
},
"required": ["explanation", "output"],
"additionalProperties": false
}
},
"final_answer": { "type": "string" }
},
"required": ["steps", "final_answer"],
"additionalProperties": false
}
""";
Prompt prompt = new Prompt("how can I solve 8x + 7 = -23",
OllamaOptions.builder()
.withModel(OllamaModel.LLAMA3_2.getName())
.withFormat(new ObjectMapper().readValue(jsonSchema, Map.class))
.build());
ChatResponse response = this.ollamaChatModel.call(this.prompt);
----

==== Integrating with BeanOutputConverter Utilities

You can leverage existing xref::api/structured-output-converter.adoc#_bean_output_converter[BeanOutputConverter] utilities to automatically generate the JSON Schema from your domain objects and later convert the structured response into domain-specific instances:

[source,java]
----
record MathReasoning(
@JsonProperty(required = true, value = "steps") Steps steps,
@JsonProperty(required = true, value = "final_answer") String finalAnswer) {
record Steps(
@JsonProperty(required = true, value = "items") Items[] items) {
record Items(
@JsonProperty(required = true, value = "explanation") String explanation,
@JsonProperty(required = true, value = "output") String output) {
}
}
}
var outputConverter = new BeanOutputConverter<>(MathReasoning.class);
Prompt prompt = new Prompt("how can I solve 8x + 7 = -23",
OllamaOptions.builder()
.withModel(OllamaModel.LLAMA3_2.getName())
.withFormat(outputConverter.getJsonSchemaMap())
.build());
ChatResponse response = this.ollamaChatModel.call(this.prompt);
String content = this.response.getResult().getOutput().getText();
MathReasoning mathReasoning = this.outputConverter.convert(this.content);
----

NOTE: Ensure you use the `@JsonProperty(required = true,...)` annotation for generating a schema that accurately marks fields as `required`.
Although this is optional for JSON Schema, it's recommended for the structured response to function correctly.

== OpenAI API Compatibility

Ollama is OpenAI API-compatible and you can use the xref:api/chat/openai-chat.adoc[Spring AI OpenAI] client to talk to Ollama and use tools.
Expand Down

0 comments on commit 6ab7e20

Please sign in to comment.