Merge pull request #55 from MEITREX/ollama-addition

Ollama Change to DGX

Author: ZiyaSelim

src/main/java/de/unistuttgart/iste/meitrex/common/config/OllamaConfig.java

@@ -20,4 +20,7 @@ public class OllamaConfig {
     // Should be the name of the folder in which the prompts reside in the resource directory
     @Value("${ollama.promptFolder:prompt_templates}")
     private String promptFolder;
+
+    @Value("${ollama.apiKey:}")
+    private String apiKey;
 }

src/main/java/de/unistuttgart/iste/meitrex/common/ollama/OllamaClient.java

@@ -9,12 +9,17 @@
 import org.springframework.stereotype.Service;
 
 import javax.annotation.Nullable;
-import java.io.*;
+import java.io.BufferedReader;
+import java.io.FileNotFoundException;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.InputStreamReader;
 import java.net.URI;
 import java.net.http.HttpClient;
 import java.net.http.HttpRequest;
 import java.net.http.HttpResponse;
 import java.nio.charset.StandardCharsets;
+import java.time.Duration;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Optional;
@@ -98,7 +103,7 @@ public String fillTemplate(final String promptTemplate, final Map<String, String
         while (matcher.find()) {
             final String fullPlaceholder = matcher.group();
 
-            final String key = fullPlaceholder.substring(2, fullPlaceholder.length() - 2);
+            final String key = fullPlaceholder.substring(2, fullPlaceholder.length() - 2).trim();
             final String value = args.get(key);
 
             if (value != null) {
@@ -131,13 +136,16 @@ public <ResponseType> ResponseType startQuery(
             @Nullable final String modelOverride) {
         try {
             final String promptTemplate = getTemplate(templateFileName);
-
             final String filledPrompt = fillTemplate(promptTemplate, argMap);
 
             final TypeReference<HashMap<String, Object>> typeRef = new TypeReference<>() {};
             final String jsonSchema = jsonSchemaService.getJsonSchema(responseType);
             final Map<String, Object> schemaObject = jsonMapper.readValue(jsonSchema, typeRef);
 
+            if (schemaObject.get("properties") instanceof Map<?, ?> properties) {
+                schemaObject.put("required", properties.keySet());
+            }
+
             // Determine which model to use: override or default config
             final String modelToUse = (modelOverride != null && !modelOverride.isBlank())
                     ? modelOverride
@@ -192,18 +200,24 @@ public <ResponseType> ResponseType startQuery(
      * @throws IOException if the request or response handling fails
      * @throws InterruptedException if the HTTP call is interrupted
      */
-    private OllamaResponse queryLLM(OllamaRequest request) throws IOException, InterruptedException {
+    private OllamaResponse queryLLM(final OllamaRequest request) throws IOException, InterruptedException {
         final String json = jsonMapper.writeValueAsString(request);
 
-        HttpRequest req = HttpRequest.newBuilder()
+        final HttpRequest.Builder reqBuilder = HttpRequest.newBuilder()
                 .uri(URI.create(this.config.getUrl() + "/" + this.config.getEndpoint()))
                 .header("Content-Type", "application/json")
-                .POST(HttpRequest.BodyPublishers.ofString(json))
-                .build();
+                .timeout(Duration.ofMinutes(10))
+                .POST(HttpRequest.BodyPublishers.ofString(json));
+
+        if (this.config.getApiKey() != null && !this.config.getApiKey().isBlank()) {
+            reqBuilder.header("Authorization", "Bearer " + this.config.getApiKey());
+        }
+
+        final HttpResponse<String> response = client.send(reqBuilder.build(), HttpResponse.BodyHandlers.ofString());
 
-        HttpResponse<String> response = client.send(req, HttpResponse.BodyHandlers.ofString());
+        log.info("RAW OLLAMA HTTP RESPONSE BODY: {}", response.body());
 
-        OllamaResponse result = jsonMapper.readValue(response.body(), OllamaResponse.class);
+        final OllamaResponse result = jsonMapper.readValue(response.body(), OllamaResponse.class);
 
         if (result.getError() != null) {
             throw new RuntimeException("Ollama returned error: " + result.getError());
@@ -221,9 +235,11 @@ private OllamaResponse queryLLM(OllamaRequest request) throws IOException, Inter
      * @return an optional of the parsed response
      * @param <ResponseType> the type to cast to
      */
-    public <ResponseType> Optional<ResponseType> parseResponse(OllamaResponse ollamaResponse, Class<ResponseType> responseType) {
+    public <ResponseType> Optional<ResponseType> parseResponse(
+            final OllamaResponse ollamaResponse,
+            final Class<ResponseType> responseType) {
         final String response = ollamaResponse.getResponse();
-        if(responseType == null || response == null) {
+        if (responseType == null || response == null) {
             log.info("Response is null or empty: {}", response);
             return Optional.empty();
         }

src/main/java/de/unistuttgart/iste/meitrex/common/ollama/OllamaRequest.java

@@ -1,18 +1,12 @@
 package de.unistuttgart.iste.meitrex.common.ollama;
 
 import com.fasterxml.jackson.annotation.JsonProperty;
-import lombok.Getter;
 import java.util.Map;
 
-public record OllamaRequest(@JsonProperty("model") String model, @JsonProperty("prompt") String prompt,
-                            @JsonProperty("stream") boolean stream,
-                            @JsonProperty("format") Map<String, Object> format) {
-
-    public OllamaRequest(String model, String prompt, boolean stream, Map<String, Object> format) {
-        this.model = model;
-        this.prompt = prompt;
-        this.stream = stream;
-        this.format = format;
-    }
-}
+public record OllamaRequest(
+        @JsonProperty("model") String model,
+        @JsonProperty("prompt") String prompt,
+        @JsonProperty("stream") boolean stream,
+        @JsonProperty("format") Map<String, Object> format
+) {}
 

src/main/java/de/unistuttgart/iste/meitrex/common/ollama/OllamaResponse.java

@@ -10,35 +10,33 @@
 public class OllamaResponse {
 
     @JsonProperty("total_duration")
-    long totalDuration;
+    private long totalDuration;
     @JsonProperty("load_duration")
-    long loadDuration;
+    private long loadDuration;
     @JsonProperty("prompt_eval_count")
-    long promptEvalCount;
+    private long promptEvalCount;
     @JsonProperty("prompt_eval_duration")
-    long promptEvalDuration;
+    private long promptEvalDuration;
     @JsonProperty("eval_count")
-    long evalCount;
+    private long evalCount;
     @JsonProperty("eval_duration")
-    long evalDuration;
+    private long evalDuration;
     @JsonProperty("model")
-    String model;
+    private String model;
     @JsonProperty("created_at")
-    String createdAt;
+    private String createdAt;
     @JsonProperty("response")
-    String response;
+    private String response;
     @JsonProperty("done")
-    boolean done;
+    private boolean done;
     @JsonProperty("done_reason")
-    String doneReason;
+    private String doneReason;
     @JsonProperty("context")
-    long[] context;
+    private long[] context;
     @JsonProperty("error")
-    String error;
-
+    private String error;
 
     /**
-     *
      * @return The total duration of the request in milliseconds.
      */
     @JsonIgnore
@@ -47,7 +45,6 @@ public long getTotalDuration() {
     }
 
     /**
-     *
      * @return The duration of the model loading in milliseconds.
      */
     @JsonIgnore
@@ -56,7 +53,6 @@ public long getLoadDuration() {
     }
 
     /**
-     *
      * @return The number of prompt evaluations.
      */
     @JsonIgnore
@@ -65,7 +61,6 @@ public long getPromptEvalCount() {
     }
 
     /**
-     *
      * @return The duration of the prompt evaluation in milliseconds.
      */
     @JsonIgnore
@@ -74,7 +69,6 @@ public long getPromptEvalDuration() {
     }
 
     /**
-     *
      * @return The number of evaluations.
      */
     @JsonIgnore
@@ -83,7 +77,6 @@ public long getEvalCount() {
     }
 
     /**
-     *
      * @return The duration of the evaluation in milliseconds.
      */
     @JsonIgnore
@@ -92,7 +85,6 @@ public long getEvalDuration() {
     }
 
     /**
-     *
      * @return The model used for the request.
      */
     @JsonIgnore
@@ -101,7 +93,6 @@ public String getModel() {
     }
 
     /**
-     *
      * @return The creation time of the request in ISO 8601 format.
      */
     @JsonIgnore
@@ -110,7 +101,6 @@ public String getCreatedAt() {
     }
 
     /**
-     *
      * @return The response from the model.
      */
     @JsonIgnore
@@ -119,17 +109,14 @@ public String getResponse() {
     }
 
     /**
-     *
      * @return Whether the request is done or not.
      */
     @JsonIgnore
     public boolean isDone() {
         return done;
     }
 
-
     /**
-     *
      * @return The reason why the request is done.
      */
     @JsonIgnore
@@ -138,7 +125,6 @@ public String getDoneReason() {
     }
 
     /**
-     *
      * @return The context of the request.
      */
     @JsonIgnore
@@ -147,5 +133,7 @@ public long[] getContext() {
     }
 
     @JsonIgnore
-    public String getError() {return error;}
+    public String getError() {
+        return error;
+    }
 }