seeseemelk/ollama4j
1
0
Fork 0
forked from Mirror/ollama4j
Code Issues Pull Requests Actions Packages Releases 1 Activity

basic auth implemented for generate request

Browse Source
This commit is contained in:
Michael Wechner 2023-12-26 20:49:54 +01:00
parent 2027171cb9
commit 9d336e257c
1 changed files with 44 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

46
src/main/java/io/github/amithkoujalgi/ollama4j/core/OllamaAPI.java
View File

@ -37,6 +37,8 @@ public class OllamaAPI {
private final String host; private final String host;
private long requestTimeoutSeconds = 3; private long requestTimeoutSeconds = 3;
private boolean verbose = true; private boolean verbose = true;
private String username;
private String password;
/** /**
* Instantiates the Ollama API. * Instantiates the Ollama API.
@ -64,6 +66,14 @@ public class OllamaAPI {
this.verbose = verbose; this.verbose = verbose;
} }
/**
*
*/
public void setBasicAuth(String username, String password) {
this.username = username;
this.password = password;
}
/** /**
* API to check the reachability of Ollama server. * API to check the reachability of Ollama server.
* *
@ -426,14 +436,18 @@ public class OllamaAPI {
long startTime = System.currentTimeMillis(); long startTime = System.currentTimeMillis();
HttpClient httpClient = HttpClient.newHttpClient(); HttpClient httpClient = HttpClient.newHttpClient();
URI uri = URI.create(this.host + "/api/generate"); URI uri = URI.create(this.host + "/api/generate");
HttpRequest request = HttpRequest.Builder requestBuilder =
HttpRequest.newBuilder(uri) HttpRequest.newBuilder(uri)
.POST( .POST(
HttpRequest.BodyPublishers.ofString( HttpRequest.BodyPublishers.ofString(
Utils.getObjectMapper().writeValueAsString(ollamaRequestModel))) Utils.getObjectMapper().writeValueAsString(ollamaRequestModel)))
.header("Content-Type", "application/json") .header("Content-Type", "application/json")
.timeout(Duration.ofSeconds(requestTimeoutSeconds)) .timeout(Duration.ofSeconds(requestTimeoutSeconds));
.build(); if (basicAuthCredentialsSet()) {
requestBuilder.header("Authorization", getBasicAuthHeaderValue());
}
HttpRequest request = requestBuilder.build();
logger.debug("Ask model '" + ollamaRequestModel + "' ...");
HttpResponse<InputStream> response = HttpResponse<InputStream> response =
httpClient.send(request, HttpResponse.BodyHandlers.ofInputStream()); httpClient.send(request, HttpResponse.BodyHandlers.ofInputStream());
int statusCode = response.statusCode(); int statusCode = response.statusCode();
@ -444,9 +458,15 @@ public class OllamaAPI {
String line; String line;
while ((line = reader.readLine()) != null) { while ((line = reader.readLine()) != null) {
if (statusCode == 404) { if (statusCode == 404) {
logger.warn("Status code: 404 (Not Found)");
OllamaErrorResponseModel ollamaResponseModel = OllamaErrorResponseModel ollamaResponseModel =
Utils.getObjectMapper().readValue(line, OllamaErrorResponseModel.class); Utils.getObjectMapper().readValue(line, OllamaErrorResponseModel.class);
responseBuffer.append(ollamaResponseModel.getError()); responseBuffer.append(ollamaResponseModel.getError());
} else if (statusCode == 401) {
logger.warn("Status code: 401 (Unauthorized)");
OllamaErrorResponseModel ollamaResponseModel =
Utils.getObjectMapper().readValue("{\"error\":\"Unauthorized\"}", OllamaErrorResponseModel.class);
responseBuffer.append(ollamaResponseModel.getError());
}else { }else {
OllamaResponseModel ollamaResponseModel = OllamaResponseModel ollamaResponseModel =
Utils.getObjectMapper().readValue(line, OllamaResponseModel.class); Utils.getObjectMapper().readValue(line, OllamaResponseModel.class);
@ -457,10 +477,30 @@ public class OllamaAPI {
} }
} }
if (statusCode != 200) { if (statusCode != 200) {
logger.error("Status code " + statusCode + " instead 200");
throw new OllamaBaseException(responseBuffer.toString()); throw new OllamaBaseException(responseBuffer.toString());
} else { } else {
long endTime = System.currentTimeMillis(); long endTime = System.currentTimeMillis();
return new OllamaResult(responseBuffer.toString().trim(), endTime - startTime, statusCode); return new OllamaResult(responseBuffer.toString().trim(), endTime - startTime, statusCode);
} }
} }
/**
* @return basic authentication header value (encoded credentials)
*/
private String getBasicAuthHeaderValue() {
String credentialsToEncode = username + ":" + password;
return "Basic " + Base64.getEncoder().encodeToString(credentialsToEncode.getBytes());
}
/**
* @return true when Basic Auth credentials set
*/
private boolean basicAuthCredentialsSet() {
if (username != null && password != null) {
return true;
} else {
return false;
}
}
} }