Adds options to EmbeddingsRequest

Additionally refactors the Embedding Models and Tests
This commit is contained in:
Markus Klenke
2024-02-25 20:53:45 +00:00
parent a10692e2f1
commit 63d4de4e24
12 changed files with 203 additions and 121 deletions

View File

@@ -6,10 +6,11 @@ import io.github.amithkoujalgi.ollama4j.core.models.chat.OllamaChatMessage;
import io.github.amithkoujalgi.ollama4j.core.models.chat.OllamaChatRequestBuilder;
import io.github.amithkoujalgi.ollama4j.core.models.chat.OllamaChatRequestModel;
import io.github.amithkoujalgi.ollama4j.core.models.chat.OllamaChatResult;
import io.github.amithkoujalgi.ollama4j.core.models.embeddings.OllamaEmbeddingResponseModel;
import io.github.amithkoujalgi.ollama4j.core.models.embeddings.OllamaEmbeddingsRequestModel;
import io.github.amithkoujalgi.ollama4j.core.models.generate.OllamaGenerateRequestModel;
import io.github.amithkoujalgi.ollama4j.core.models.request.CustomModelFileContentsRequest;
import io.github.amithkoujalgi.ollama4j.core.models.request.CustomModelFilePathRequest;
import io.github.amithkoujalgi.ollama4j.core.models.request.ModelEmbeddingsRequest;
import io.github.amithkoujalgi.ollama4j.core.models.request.ModelRequest;
import io.github.amithkoujalgi.ollama4j.core.models.request.OllamaChatEndpointCaller;
import io.github.amithkoujalgi.ollama4j.core.models.request.OllamaGenerateEndpointCaller;
@@ -313,8 +314,18 @@ public class OllamaAPI {
*/
public List<Double> generateEmbeddings(String model, String prompt)
throws IOException, InterruptedException, OllamaBaseException {
return generateEmbeddings(new OllamaEmbeddingsRequestModel(model, prompt));
}
/**
* Generate embeddings using a {@link OllamaEmbeddingsRequestModel}.
*
* @param modelRequest request for '/api/embeddings' endpoint
* @return embeddings
*/
public List<Double> generateEmbeddings(OllamaEmbeddingsRequestModel modelRequest) throws IOException, InterruptedException, OllamaBaseException{
URI uri = URI.create(this.host + "/api/embeddings");
String jsonData = new ModelEmbeddingsRequest(model, prompt).toString();
String jsonData = modelRequest.toString();
HttpClient httpClient = HttpClient.newHttpClient();
HttpRequest.Builder requestBuilder =
getRequestBuilderDefault(uri)
@@ -325,8 +336,8 @@ public class OllamaAPI {
int statusCode = response.statusCode();
String responseBody = response.body();
if (statusCode == 200) {
EmbeddingResponse embeddingResponse =
Utils.getObjectMapper().readValue(responseBody, EmbeddingResponse.class);
OllamaEmbeddingResponseModel embeddingResponse =
Utils.getObjectMapper().readValue(responseBody, OllamaEmbeddingResponseModel.class);
return embeddingResponse.getEmbedding();
} else {
throw new OllamaBaseException(statusCode + " - " + responseBody);

View File

@@ -1,4 +1,4 @@
package io.github.amithkoujalgi.ollama4j.core.models;
package io.github.amithkoujalgi.ollama4j.core.models.embeddings;
import com.fasterxml.jackson.annotation.JsonProperty;
@@ -7,7 +7,7 @@ import lombok.Data;
@SuppressWarnings("unused")
@Data
public class EmbeddingResponse {
public class OllamaEmbeddingResponseModel {
@JsonProperty("embedding")
private List<Double> embedding;
}

View File

@@ -0,0 +1,31 @@
package io.github.amithkoujalgi.ollama4j.core.models.embeddings;
import io.github.amithkoujalgi.ollama4j.core.utils.Options;
public class OllamaEmbeddingsRequestBuilder {
private OllamaEmbeddingsRequestBuilder(String model, String prompt){
request = new OllamaEmbeddingsRequestModel(model, prompt);
}
private OllamaEmbeddingsRequestModel request;
public static OllamaEmbeddingsRequestBuilder getInstance(String model, String prompt){
return new OllamaEmbeddingsRequestBuilder(model, prompt);
}
public OllamaEmbeddingsRequestModel build(){
return request;
}
public OllamaEmbeddingsRequestBuilder withOptions(Options options){
this.request.setOptions(options.getOptionsMap());
return this;
}
public OllamaEmbeddingsRequestBuilder withKeepAlive(String keepAlive){
this.request.setKeepAlive(keepAlive);
return this;
}
}

View File

@@ -0,0 +1,33 @@
package io.github.amithkoujalgi.ollama4j.core.models.embeddings;
import static io.github.amithkoujalgi.ollama4j.core.utils.Utils.getObjectMapper;
import java.util.Map;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.core.JsonProcessingException;
import lombok.Data;
import lombok.NoArgsConstructor;
import lombok.NonNull;
import lombok.RequiredArgsConstructor;
@Data
@RequiredArgsConstructor
@NoArgsConstructor
public class OllamaEmbeddingsRequestModel {
@NonNull
private String model;
@NonNull
private String prompt;
protected Map<String, Object> options;
@JsonProperty(value = "keep_alive")
private String keepAlive;
@Override
public String toString() {
try {
return getObjectMapper().writerWithDefaultPrettyPrinter().writeValueAsString(this);
} catch (JsonProcessingException e) {
throw new RuntimeException(e);
}
}
}

View File

@@ -1,23 +0,0 @@
package io.github.amithkoujalgi.ollama4j.core.models.request;
import static io.github.amithkoujalgi.ollama4j.core.utils.Utils.getObjectMapper;
import com.fasterxml.jackson.core.JsonProcessingException;
import lombok.AllArgsConstructor;
import lombok.Data;
@Data
@AllArgsConstructor
public class ModelEmbeddingsRequest {
private String model;
private String prompt;
@Override
public String toString() {
try {
return getObjectMapper().writerWithDefaultPrettyPrinter().writeValueAsString(this);
} catch (JsonProcessingException e) {
throw new RuntimeException(e);
}
}
}