Refactor Llama embedding and chat completion tests for consistency and clarity

Author: Jan-Kazlouski-elastic

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/llama/action/LlamaActionCreatorTests.java

@@ -30,7 +30,7 @@
 import org.elasticsearch.xpack.inference.logging.ThrottlerManager;
 import org.elasticsearch.xpack.inference.services.ServiceComponents;
 import org.elasticsearch.xpack.inference.services.llama.completion.LlamaChatCompletionModelTests;
-import org.elasticsearch.xpack.inference.services.llama.embeddings.LlamaEmbeddingModelTests;
+import org.elasticsearch.xpack.inference.services.llama.embeddings.LlamaEmbeddingsModelTests;
 import org.junit.After;
 import org.junit.Before;
 
@@ -226,7 +226,7 @@ public void testExecute_FailsFromInvalidResponseFormat_ForCompletionAction() thr
     }
 
     private PlainActionFuture<InferenceServiceResults> createEmbeddingsFuture(Sender sender, ServiceComponents threadPool) {
-        var model = LlamaEmbeddingModelTests.createEmbeddingsModel("model", getUrl(webServer), "secret");
+        var model = LlamaEmbeddingsModelTests.createEmbeddingsModel("model", getUrl(webServer), "secret");
         var actionCreator = new LlamaActionCreator(sender, threadPool);
         var action = actionCreator.create(model);
 

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/llama/completion/LlamaChatCompletionModelTests.java

@@ -8,6 +8,7 @@
 package org.elasticsearch.xpack.inference.services.llama.completion;
 
 import org.elasticsearch.common.settings.SecureString;
+import org.elasticsearch.inference.EmptySecretSettings;
 import org.elasticsearch.inference.TaskType;
 import org.elasticsearch.inference.UnifiedCompletionRequest;
 import org.elasticsearch.test.ESTestCase;
@@ -39,6 +40,16 @@ public static LlamaChatCompletionModel createChatCompletionModel(String modelId,
         );
     }
 
+    public static LlamaChatCompletionModel createChatCompletionModelNoAuth(String modelId, String url) {
+        return new LlamaChatCompletionModel(
+            "id",
+            TaskType.CHAT_COMPLETION,
+            "llama",
+            new LlamaChatCompletionServiceSettings(modelId, url, null),
+            EmptySecretSettings.INSTANCE
+        );
+    }
+
     public void testOverrideWith_UnifiedCompletionRequest_KeepsSameModelId() {
         var model = createCompletionModel("model_name", "url", "api_key");
         var request = new UnifiedCompletionRequest(

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/llama/embeddings/LlamaEmbeddingModelTests.java renamed to x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/llama/embeddings/LlamaEmbeddingsModelTests.java

@@ -13,7 +13,7 @@
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.xpack.inference.services.settings.DefaultSecretSettings;
 
-public class LlamaEmbeddingModelTests extends ESTestCase {
+public class LlamaEmbeddingsModelTests extends ESTestCase {
     public static LlamaEmbeddingsModel createEmbeddingsModel(String modelId, String url, String apiKey) {
         return new LlamaEmbeddingsModel(
             "id",

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/llama/request/completion/LlamaChatCompletionRequestTests.java

@@ -0,0 +1,81 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.inference.services.llama.request.completion;
+
+import org.apache.http.client.methods.HttpPost;
+import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.xpack.inference.external.http.sender.UnifiedChatInput;
+import org.elasticsearch.xpack.inference.services.llama.completion.LlamaChatCompletionModelTests;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Map;
+
+import static org.elasticsearch.xpack.inference.external.http.Utils.entityAsMap;
+import static org.hamcrest.Matchers.instanceOf;
+import static org.hamcrest.Matchers.is;
+
+public class LlamaChatCompletionRequestTests extends ESTestCase {
+
+    public void testCreateRequest_WithStreaming() throws IOException {
+        String input = randomAlphaOfLength(15);
+        var request = createRequest("model", "url", "secret", input, true);
+        var httpRequest = request.createHttpRequest();
+
+        assertThat(httpRequest.httpRequestBase(), instanceOf(HttpPost.class));
+        var httpPost = (HttpPost) httpRequest.httpRequestBase();
+
+        var requestMap = entityAsMap(httpPost.getEntity().getContent());
+        assertThat(request.getURI().toString(), is("url"));
+        assertThat(requestMap.get("stream"), is(true));
+        assertThat(requestMap.get("model"), is("model"));
+        assertThat(requestMap.get("n"), is(1));
+        assertThat(requestMap.get("stream_options"), is(Map.of("include_usage", true)));
+        assertThat(requestMap.get("messages"), is(List.of(Map.of("role", "user", "content", input))));
+        assertNotNull(httpPost.getHeaders("Authorization"));
+    }
+
+    public void testCreateRequest_NoStreaming_NoAuthorization() throws IOException {
+        String input = randomAlphaOfLength(15);
+        var request = createRequestWithNoAuth("model", "url", input, false);
+        var httpRequest = request.createHttpRequest();
+
+        assertThat(httpRequest.httpRequestBase(), instanceOf(HttpPost.class));
+        var httpPost = (HttpPost) httpRequest.httpRequestBase();
+
+        var requestMap = entityAsMap(httpPost.getEntity().getContent());
+        assertThat(request.getURI().toString(), is("url"));
+        assertThat(requestMap.get("stream"), is(false));
+        assertThat(requestMap.get("model"), is("model"));
+        assertThat(requestMap.get("n"), is(1));
+        assertNull(requestMap.get("stream_options"));
+        assertThat(requestMap.get("messages"), is(List.of(Map.of("role", "user", "content", input))));
+        assertNull(httpPost.getFirstHeader("Authorization"));
+    }
+
+    public void testTruncate_DoesNotReduceInputTextSize() {
+        String input = randomAlphaOfLength(5);
+        var request = createRequest("model", "url", "secret", input, true);
+        assertThat(request.truncate(), is(request));
+    }
+
+    public void testTruncationInfo_ReturnsNull() {
+        var request = createRequest("model", "url", "secret", randomAlphaOfLength(5), true);
+        assertNull(request.getTruncationInfo());
+    }
+
+    public static LlamaChatCompletionRequest createRequest(String modelId, String url, String apiKey, String input, boolean stream) {
+        var chatCompletionModel = LlamaChatCompletionModelTests.createChatCompletionModel(modelId, url, apiKey);
+        return new LlamaChatCompletionRequest(new UnifiedChatInput(List.of(input), "user", stream), chatCompletionModel);
+    }
+
+    public static LlamaChatCompletionRequest createRequestWithNoAuth(String modelId, String url, String input, boolean stream) {
+        var chatCompletionModel = LlamaChatCompletionModelTests.createChatCompletionModelNoAuth(modelId, url);
+        return new LlamaChatCompletionRequest(new UnifiedChatInput(List.of(input), "user", stream), chatCompletionModel);
+    }
+}

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/mistral/request/completion/MistralChatCompletionRequestTests.java

@@ -49,7 +49,7 @@ public void testTruncate_DoesNotReduceInputTextSize() throws IOException {
         var requestMap = entityAsMap(httpPost.getEntity().getContent());
         assertThat(requestMap, aMapWithSize(4));
 
-        // We do not truncate for Hugging Face chat completions
+        // We do not truncate for Mistral chat completions
         assertThat(requestMap.get("messages"), is(List.of(Map.of("role", "user", "content", input))));
         assertThat(requestMap.get("model"), is("model"));
         assertThat(requestMap.get("n"), is(1));