adding support for messages.

Author: Vaijanath Rao

pom.xml

@@ -5,7 +5,7 @@
 
 	<groupId>de.kherud</groupId>
 	<artifactId>llama</artifactId>
-	<version>4.0.1</version>
+	<version>4.0.0</version>
 	<packaging>jar</packaging>
 
 	<name>${project.groupId}:${project.artifactId}</name>
@@ -65,7 +65,11 @@
 			<version>24.1.0</version>
 			<scope>compile</scope>
 		</dependency>
-		
+		<dependency>
+    		<groupId>com.fasterxml.jackson.core</groupId>
+    		<artifactId>jackson-databind</artifactId>
+    		<version>2.16.0</version> <!-- Use the latest version -->
+		</dependency>
 	</dependencies>
 
 	<build>

src/main/cpp/jllama.cpp

@@ -786,6 +786,20 @@ JNIEXPORT jobject JNICALL Java_de_kherud_llama_LlamaModel_rerank(JNIEnv *env, jo
 
 }
 
+JNIEXPORT jstring JNICALL Java_de_kherud_llama_LlamaModel_applyTemplate(JNIEnv *env, jobject obj, jstring jparams){
+ jlong server_handle = env->GetLongField(obj, f_model_pointer);
+    auto *ctx_server = reinterpret_cast<server_context *>(server_handle); // NOLINT(*-no-int-to-ptr)
+
+    std::string c_params = parse_jstring(env, jparams);
+    json data = json::parse(c_params);
+    
+    json templateData = oaicompat_completion_params_parse(data, ctx_server->params_base.use_jinja, ctx_server->params_base.reasoning_format, ctx_server->chat_templates.get());
+    std::string tok_str = templateData.at("prompt");
+	jstring jtok_str = env->NewStringUTF(tok_str.c_str());
+	
+	return jtok_str;
+}
+
 JNIEXPORT jintArray JNICALL Java_de_kherud_llama_LlamaModel_encode(JNIEnv *env, jobject obj, jstring jprompt) {
     jlong server_handle = env->GetLongField(obj, f_model_pointer);
     auto *ctx_server = reinterpret_cast<server_context *>(server_handle); // NOLINT(*-no-int-to-ptr)

src/main/cpp/jllama.h

@@ -1,8 +1,13 @@
 package de.kherud.llama;
 
 import java.util.Collection;
+import java.util.List;
 import java.util.Map;
 
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.node.ArrayNode;
+import com.fasterxml.jackson.databind.node.ObjectNode;
+
 import de.kherud.llama.args.MiroStat;
 import de.kherud.llama.args.Sampler;
 
@@ -12,6 +17,9 @@
  * {@link LlamaModel#complete(InferenceParameters)}.
  */
 public final class InferenceParameters extends JsonParameters {
+	
+    private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); // Reusable ObjectMapper
+
 
 	private static final String PARAM_PROMPT = "prompt";
 	private static final String PARAM_INPUT_PREFIX = "input_prefix";
@@ -47,6 +55,7 @@ public final class InferenceParameters extends JsonParameters {
 	private static final String PARAM_STREAM = "stream";
 	private static final String PARAM_USE_CHAT_TEMPLATE = "use_chat_template";
 	private static final String PARAM_USE_JINJA = "use_jinja";
+	private static final String PARAM_MESSAGES = "messages";
 
 	public InferenceParameters(String prompt) {
 		// we always need a prompt
@@ -493,7 +502,41 @@ public InferenceParameters setUseChatTemplate(boolean useChatTemplate) {
 		return this;
 	}
 
-	
+	/**
+     * Set the messages for chat-based inference.
+     * - Allows **only one** system message.
+     * - Allows **one or more** user/assistant messages.
+     */
+    public InferenceParameters setMessages(String systemMessage, List<Pair<String, String>> messages) {
+        ArrayNode messagesArray = OBJECT_MAPPER.createArrayNode();
+
+        // Add system message (if provided)
+        if (systemMessage != null && !systemMessage.isEmpty()) {
+            ObjectNode systemObj = OBJECT_MAPPER.createObjectNode();
+            systemObj.put("role", "system");
+            systemObj.put("content", systemMessage);
+            messagesArray.add(systemObj);
+        }
+
+        // Add user/assistant messages
+        for (Pair<String, String> message : messages) {
+            String role = message.getKey();
+            String content = message.getValue();
+
+            if (!role.equals("user") && !role.equals("assistant")) {
+                throw new IllegalArgumentException("Invalid role: " + role + ". Role must be 'user' or 'assistant'.");
+            }
+
+            ObjectNode messageObj = OBJECT_MAPPER.createObjectNode();
+            messageObj.put("role", role);
+            messageObj.put("content", content);
+            messagesArray.add(messageObj);
+        }
+
+        // Convert ArrayNode to a JSON string and store it in parameters
+        parameters.put(PARAM_MESSAGES, messagesArray.toString());
+        return this;
+    }
 
 
 

src/main/java/de/kherud/llama/InferenceParameters.java

@@ -163,4 +163,9 @@ public List<Pair<String, Float>> rerank(boolean reRank, String query, String ...
 	}
 
 	public native LlamaOutput rerank(String query, String... documents);
+	
+	public  String applyTemplate(InferenceParameters parameters) {
+		return applyTemplate(parameters.toString());
+	}
+	public native String applyTemplate(String parametersJson);
 }

src/main/java/de/kherud/llama/LlamaModel.java

@@ -316,4 +316,20 @@ public void testJsonSchemaToGrammar() {
 		String actualGrammar = LlamaModel.jsonSchemaToGrammar(schema);
 		Assert.assertEquals(expectedGrammar, actualGrammar);
 	}
+	
+	@Test
+	public void testTemplate() {
+		
+		List<Pair<String, String>> userMessages = new ArrayList<>();
+        userMessages.add(new Pair<>("user", "What is the best book?"));
+        userMessages.add(new Pair<>("assistant", "It depends on your interests. Do you like fiction or non-fiction?"));
+        
+		InferenceParameters params = new InferenceParameters("A book recommendation system.")
+				.setMessages("Book", userMessages)
+				.setTemperature(0.95f)
+				.setStopStrings("\"\"\"")
+				.setNPredict(nPredict)
+				.setSeed(42);
+		Assert.assertEquals(model.applyTemplate(params), "<|im_start|>system\nBook<|im_end|>\n<|im_start|>user\nWhat is the best book?<|im_end|>\n<|im_start|>assistant\nIt depends on your interests. Do you like fiction or non-fiction?<|im_end|>\n<|im_start|>assistant\n");
+	}
 }