fix(ai): add context possibiloty

frischHWC · Jun 12, 2024 · aeb95b2 · aeb95b2
1 parent 530caf1
commit aeb95b2
Show file tree

Hide file tree

Showing 7 changed files with 71 additions and 25 deletions.
diff --git a/src/main/java/com/datagen/model/type/BedrockField.java b/src/main/java/com/datagen/model/type/BedrockField.java
@@ -55,9 +55,10 @@ public class BedrockField extends Field<String> {
   private final String modelId;
   private final BedrockModelType bedrockmodeltype;
   private JSONObject preparedRequest = null;
+  private final String context;
 
   public BedrockField(String name, String url, String user, String password,
-                      String request, String modelType, Float temperature, String region, Integer maxTokens) {
+                      String request, String modelType, Float temperature, String region, Integer maxTokens, String context) {
     this.name = name;
     this.url = url;
     this.user = user;
@@ -76,6 +77,10 @@ public BedrockField(String name, String url, String user, String password,
         .region(this.region)
         .build();
 
+    var contextAsMessage = context!=null?"Use the following information to answer the question:"+System.lineSeparator()+context:"";
+    this.context = "Generate only the answer."+System.lineSeparator()+contextAsMessage;
+    log.debug("Will provide following System information to the model: {}", context);
+
     // See model Ids available at: https://docs.aws.amazon.com/bedrock/latest/userguide/model-ids.html
     this.modelId = modelType == null ? "amazon.titan-text-lite-v1" : modelType;
     /*
@@ -107,6 +112,7 @@ public BedrockField(String name, String url, String user, String password,
           yield new JSONObject()
               .put("temperature", this.temperature)
               .put("stop_sequences", List.of("\n\nHuman:"))
+              .put("system", context)
               .put("max_tokens_to_sample", this.maxTokens);
         case MISTRAL:
           yield new JSONObject()
@@ -134,11 +140,14 @@ public String generateComputedValue(Row row) {
       case ANTHROPIC -> preparedRequest.put("prompt",
           "Human: " + stringToEvaluate + "\\n\\nAssistant:");
       case MISTRAL -> preparedRequest.put("prompt",
-          "<s>[INST] " + stringToEvaluate + "[/INST]");
+          "<s>[INST] " + stringToEvaluate + "[/INST]" + context);
       case TITAN -> preparedRequest.put("inputText", stringToEvaluate);
+      case LLAMA -> preparedRequest.put("prompt", stringToEvaluate);
       default -> preparedRequest.put("prompt", stringToEvaluate);
       }
 
+      log.debug("Request to Bedrock is: {}", preparedRequest.toString());
+
       // Encode and send the request.
       var response = bedrockRuntimeClient.invokeModel(req -> req
           .accept("application/json")
@@ -169,7 +178,7 @@ public String generateComputedValue(Row row) {
           e);
     }
 
-    return responseText;
+    return responseText.trim().replaceAll("\\n[ \\t]*\\n","");
   }
 
   @Override

diff --git a/src/main/java/com/datagen/model/type/Field.java b/src/main/java/com/datagen/model/type/Field.java
@@ -145,7 +145,8 @@ public static Field instantiateField(
       Float frequencyPenalty,
       Float presencePenalty,
       Integer maxTokens,
-      Float topP
+      Float topP,
+      String context
       ) {
     if (name == null || name.isEmpty()) {
       throw new IllegalStateException(
@@ -296,7 +297,8 @@ yield new OllamaField(name,
             presencePenalty == null ? Float.valueOf(properties.get(
                 ApplicationConfigs.OLLAMA_PRESENCE_PENALTY_DEFAULT)) : presencePenalty,
             topP == null ? Float.valueOf(properties.get(
-                ApplicationConfigs.OLLAMA_TOP_P_DEFAULT)) : topP
+                ApplicationConfigs.OLLAMA_TOP_P_DEFAULT)) : topP,
+            context
         );
 
       case "BEDROCK":
@@ -314,7 +316,8 @@ yield new BedrockField(name, url,
                 ApplicationConfigs.BEDROCK_TEMPERATURE_DEFAULT)) : temperature,
             properties.get(ApplicationConfigs.BEDROCK_REGION),
             maxTokens == null ? Integer.valueOf(properties.get(
-                ApplicationConfigs.BEDROCK_MAX_TOKENS_DEFAULT)) : maxTokens
+                ApplicationConfigs.BEDROCK_MAX_TOKENS_DEFAULT)) : maxTokens,
+            context
             );
 
       case "OPENAI":
@@ -336,7 +339,8 @@ yield new OpenAIField(name, url,
             maxTokens == null ? Integer.valueOf(properties.get(
                 ApplicationConfigs.OPENAI_MAX_TOKENS_DEFAULT)) : maxTokens,
             topP == null ? Float.valueOf(properties.get(
-                ApplicationConfigs.OPENAI_TOP_P_DEFAULT)) : topP
+                ApplicationConfigs.OPENAI_TOP_P_DEFAULT)) : topP,
+            context
         );
 
       default:

diff --git a/src/main/java/com/datagen/model/type/OllamaField.java b/src/main/java/com/datagen/model/type/OllamaField.java
@@ -28,13 +28,16 @@
 import org.apache.kudu.Type;
 import org.apache.kudu.client.PartialRow;
 import org.apache.orc.TypeDescription;
+import org.springframework.ai.chat.messages.SystemMessage;
+import org.springframework.ai.chat.messages.UserMessage;
 import org.springframework.ai.chat.prompt.Prompt;
 import org.springframework.ai.ollama.OllamaChatClient;
 import org.springframework.ai.ollama.api.OllamaApi;
 import org.springframework.ai.ollama.api.OllamaOptions;
 
 import java.sql.SQLException;
 import java.util.LinkedList;
+import java.util.List;
 
 @Slf4j
 public class OllamaField extends Field<String> {
@@ -46,11 +49,12 @@ public class OllamaField extends Field<String> {
   private final OllamaApi ollamaApi;
   private final OllamaChatClient ollamaChatClient;
   private final OllamaOptions ollamaOptions;
+  private final SystemMessage systemMessage;
 
 
   public OllamaField(String name, String url, String user, String password, String request,
                      String modelType, Float temperature, Float frequencyPenalty,
-                     Float presencePenalty, Float topP) {
+                     Float presencePenalty, Float topP, String context) {
     this.name = name;
     this.url = url;
     this.user = user;
@@ -64,18 +68,24 @@ public OllamaField(String name, String url, String user, String password, String
         .withFrequencyPenalty(frequencyPenalty == null ? 1.0f : frequencyPenalty)
         .withPresencePenalty(presencePenalty == null ? 1.0f : presencePenalty)
         .withTopP(topP == null ? 1.0f : topP);
+
+    var contextAsMessage = context!=null?"Use the following information to answer the question:"+System.lineSeparator()+context:"";
+    this.systemMessage = new SystemMessage("Generate only the answer."+System.lineSeparator()+contextAsMessage);
+    log.debug("Will provide following System information to the model: {}", systemMessage.getContent());
   }
 
   @Override
   public String generateComputedValue(Row row) {
     String stringToEvaluate = ParsingUtils.injectRowValuesToAString(row, requestToInject);
     log.debug("Asking to Ollama: {}", stringToEvaluate);
+    UserMessage userMessage = new UserMessage(stringToEvaluate);
 
     return this.ollamaChatClient.call(
         new Prompt(
-            stringToEvaluate,
+            List.of(userMessage, this.systemMessage),
             this.ollamaOptions
-        )).getResult().getOutput().getContent();
+        )).getResult().getOutput().getContent()
+        .trim().replaceAll("\\n[ \\t]*\\n","");
   }
 
   @Override

diff --git a/src/main/java/com/datagen/model/type/OpenAIField.java b/src/main/java/com/datagen/model/type/OpenAIField.java
@@ -28,12 +28,16 @@
 import org.apache.kudu.Type;
 import org.apache.kudu.client.PartialRow;
 import org.apache.orc.TypeDescription;
+import org.springframework.ai.chat.messages.SystemMessage;
+import org.springframework.ai.chat.messages.UserMessage;
+import org.springframework.ai.chat.prompt.Prompt;
 import org.springframework.ai.openai.OpenAiChatClient;
 import org.springframework.ai.openai.OpenAiChatOptions;
 import org.springframework.ai.openai.api.OpenAiApi;
 
 import java.sql.SQLException;
 import java.util.LinkedList;
+import java.util.List;
 
 @Slf4j
 public class OpenAIField extends Field<String> {
@@ -47,10 +51,11 @@ public class OpenAIField extends Field<String> {
   private final OpenAiChatClient openAiChatClient;
   private final OpenAiChatOptions openAiChatOptions;
   private final String modelId;
+  private final SystemMessage systemMessage;
 
   public OpenAIField(String name, String url, String user, String password,
                      String request, String modelType, Float temperature, Float frequencyPenalty,
-                     Float presencePenalty, Integer maxTokens, Float topP) {
+                     Float presencePenalty, Integer maxTokens, Float topP, String context) {
     this.name = name;
     this.url = url;
     this.user = user;
@@ -71,14 +76,24 @@ public OpenAIField(String name, String url, String user, String password,
         .build();
     this.openAiChatClient = new OpenAiChatClient(openAiApi, openAiChatOptions);
 
+    var contextAsMessage = context!=null?"Use the following information to answer the question:"+System.lineSeparator()+context:"";
+    this.systemMessage = new SystemMessage("Generate only the answer."+System.lineSeparator()+contextAsMessage);
+    log.debug("Will provide following System information to the model: {}", systemMessage.getContent());
+
   }
 
   @Override
   public String generateComputedValue(Row row) {
     String stringToEvaluate =
         ParsingUtils.injectRowValuesToAString(row, requestToInject);
     log.debug("Asking to OpenAI: {}", stringToEvaluate);
-    return openAiChatClient.call(stringToEvaluate);
+    UserMessage userMessage = new UserMessage(stringToEvaluate);
+
+    return openAiChatClient.call(new Prompt(
+        List.of(userMessage, this.systemMessage),
+        this.openAiChatOptions
+    )).getResult().getOutput().getContent()
+        .trim().replaceAll("\\n[ \\t]*\\n","");
   }
 
   @Override

diff --git a/src/main/java/com/datagen/parsers/JsonParser.java b/src/main/java/com/datagen/parsers/JsonParser.java
@@ -312,6 +312,13 @@ private T getOneField(JsonNode jsonField, Map<ApplicationConfigs, String> proper
       topP = null;
     }
 
+    String context;
+    try {
+      context = jsonField.get("context").asText();
+    } catch (NullPointerException e) {
+      context = null;
+    }
+
     JsonNode filtersArray = jsonField.get("filters");
     List<JsonNode> filters = new ArrayList<>();
     try {
@@ -391,7 +398,8 @@ private T getOneField(JsonNode jsonField, Map<ApplicationConfigs, String> proper
         frequencyPenalty,
         presencePenalty,
         maxTokens,
-        topP);
+        topP,
+        context);
   }
 
   private Map<String, String> mapColNameToColQual(String mapping) {

diff --git a/src/main/resources/models/example-model-ai.json b/src/main/resources/models/example-model-ai.json
@@ -18,7 +18,8 @@
       "user": "<YOUR_ACCESS_KEY>",
       "password": "<YOUR_SECRET_KEY>",
       "temperature": 1.0,
-      "max_tokens": 256
+      "max_tokens": 256,
+      "context": "For their birthday, kids wants a toy, young wants a new car, adults wants a prosper family and friends, elder wants a strong health"
     },
     {
       "name": "birthday_wish_ollama",
@@ -28,7 +29,8 @@
       "temperature": 1.0,
       "frequency_penalty": 1.5,
       "presence_penalty": 1.3,
-      "top_p": 1.0
+      "top_p": 1.0,
+      "context": "For their birthday, kids wants a toy, young wants a new car, adults wants a prosper family and friends, elder wants a strong health"
     },
     {
       "name": "birthday_wish_openai",
@@ -40,7 +42,8 @@
       "frequency_penalty": 1.5,
       "presence_penalty": 1.3,
       "max_tokens": 256,
-      "top_p": 1.0
+      "top_p": 1.0,
+      "context": "For their birthday, kids wants a toy, young wants a new car, adults wants a prosper family and friends, elder wants a strong health"
     }
   ],
   "Table_Names": {

diff --git a/src/main/resources/models/example-model.json b/src/main/resources/models/example-model.json
@@ -7,20 +7,17 @@
     {
       "name": "age",
       "type": "LONG",
-      "min": 18,
-      "max": 99
+      "min": 5,
+      "max": 85
     },
     {
-      "name": "birthday_wish_openai",
-      "type": "OPENAI",
+      "name": "birthday_wish_bedrock",
+      "type": "BEDROCK",
       "request": "generate a one line birthday wish to ${name} who is ${age} years old today",
-      "model_type": "gpt-4o",
-      "password": "",
+      "model_type": "mistral.mistral-small-2402-v1:0",
       "temperature": 1.0,
-      "frequency_penalty": 1.5,
-      "presence_penalty": 1.3,
       "max_tokens": 256,
-      "top_p": 1.0
+      "context": "For their birthday, kids wants a toy, young wants a new car, adults wants a prosper family and friends, elder wants a strong health"
     }
   ],
   "Table_Names": {