upgrade test

Author: davidkyle

x-pack/plugin/inference/qa/rolling-upgrade/src/javaRestTest/java/org/elasticsearch/xpack/application/CohereServiceUpgradeIT.java

@@ -39,10 +39,12 @@ public class CohereServiceUpgradeIT extends InferenceUpgradeTestCase {
     // TODO: replace with proper test features
     private static final String COHERE_EMBEDDINGS_ADDED_TEST_FEATURE = "gte_v8.13.0";
     private static final String COHERE_RERANK_ADDED_TEST_FEATURE = "gte_v8.14.0";
+    private static final String COHERE_COMPLETIONS_ADDED_TEST_FEATURE = "gte_v8.15.0";
     private static final String COHERE_V2_API_ADDED_TEST_FEATURE = "inference.cohere.v2";
 
     private static MockWebServer cohereEmbeddingsServer;
     private static MockWebServer cohereRerankServer;
+    private static MockWebServer cohereCompletionsServer;
 
     private enum ApiVersion {
         V1,
@@ -60,12 +62,16 @@ public static void startWebServer() throws IOException {
 
         cohereRerankServer = new MockWebServer();
         cohereRerankServer.start();
+
+        cohereCompletionsServer = new MockWebServer();
+        cohereCompletionsServer.start();
     }
 
     @AfterClass
     public static void shutdown() {
         cohereEmbeddingsServer.close();
         cohereRerankServer.close();
+        cohereCompletionsServer.close();
     }
 
     @SuppressWarnings("unchecked")
@@ -326,6 +332,80 @@ private void assertRerank(String inferenceId) throws IOException {
         assertThat(inferenceMap.entrySet(), not(empty()));
     }
 
+    @SuppressWarnings("unchecked")
+    public void testCohereCompletions() throws IOException {
+        var completionsSupported = oldClusterHasFeature(COHERE_COMPLETIONS_ADDED_TEST_FEATURE);
+        assumeTrue("Cohere completions not supported", completionsSupported);
+
+        ApiVersion oldClusterApiVersion = oldClusterHasFeature(COHERE_V2_API_ADDED_TEST_FEATURE) ? ApiVersion.V2 : ApiVersion.V1;
+
+        final String oldClusterId = "old-cluster-completions";
+
+        if (isOldCluster()) {
+            // queue a response as PUT will call the service
+            cohereCompletionsServer.enqueue(new MockResponse().setResponseCode(200).setBody(completionsResponse(oldClusterApiVersion)));
+            put(oldClusterId, completionsConfig(getUrl(cohereCompletionsServer)), TaskType.COMPLETION);
+
+            var configs = (List<Map<String, Object>>) get(TaskType.COMPLETION, oldClusterId).get("endpoints");
+            assertThat(configs, hasSize(1));
+            assertEquals("cohere", configs.get(0).get("service"));
+            var serviceSettings = (Map<String, Object>) configs.get(0).get("service_settings");
+            assertThat(serviceSettings, hasEntry("model_id", "command"));
+        } else if (isMixedCluster()) {
+            var configs = (List<Map<String, Object>>) get(TaskType.COMPLETION, oldClusterId).get("endpoints");
+            assertThat(configs, hasSize(1));
+            assertEquals("cohere", configs.get(0).get("service"));
+            var serviceSettings = (Map<String, Object>) configs.get(0).get("service_settings");
+            assertThat(serviceSettings, hasEntry("model_id", "command"));
+        } else if (isUpgradedCluster()) {
+            // check old cluster model
+            var configs = (List<Map<String, Object>>) get(TaskType.COMPLETION, oldClusterId).get("endpoints");
+            var serviceSettings = (Map<String, Object>) configs.get(0).get("service_settings");
+            assertThat(serviceSettings, hasEntry("model_id", "command"));
+
+            final String newClusterId = "new-cluster-completions";
+            {
+                cohereCompletionsServer.enqueue(new MockResponse().setResponseCode(200).setBody(completionsResponse(oldClusterApiVersion)));
+                var inferenceMap = inference(oldClusterId, TaskType.COMPLETION, "some text");
+                assertThat(inferenceMap.entrySet(), not(empty()));
+                assertVersionInPath(cohereCompletionsServer.requests().getLast(), "chat", oldClusterApiVersion);
+            }
+            {
+                // new cluster uses the V2 API
+                cohereCompletionsServer.enqueue(new MockResponse().setResponseCode(200).setBody(completionsResponse(ApiVersion.V2)));
+                put(newClusterId, completionsConfig(getUrl(cohereCompletionsServer)), TaskType.COMPLETION);
+
+                cohereCompletionsServer.enqueue(new MockResponse().setResponseCode(200).setBody(completionsResponse(ApiVersion.V2)));
+                var inferenceMap = inference(newClusterId, TaskType.COMPLETION, "some text");
+                assertThat(inferenceMap.entrySet(), not(empty()));
+                assertVersionInPath(cohereCompletionsServer.requests().getLast(), "chat", ApiVersion.V2);
+            }
+
+            {
+                // new endpoints use the V2 API which require the model to be set
+                final String upgradedClusterNoModel = "upgraded-cluster-missing-model-id";
+                var jsonBody = Strings.format("""
+                    {
+                        "service": "cohere",
+                        "service_settings": {
+                            "url": "%s",
+                            "api_key": "XXXX"
+                        }
+                    }
+                    """, getUrl(cohereEmbeddingsServer));
+
+                var e = expectThrows(ResponseException.class, () -> put(upgradedClusterNoModel, jsonBody, TaskType.COMPLETION));
+                assertThat(
+                    e.getMessage(),
+                    containsString("Validation Failed: 1: The [service_settings.model_id] field is required for the Cohere V2 API.")
+                );
+            }
+
+            delete(oldClusterId);
+            delete(newClusterId);
+        }
+    }
+
     private String embeddingConfigByte(String url) {
         return embeddingConfigTemplate(url, "byte");
     }
@@ -451,4 +531,86 @@ private String rerankResponse() {
             """;
     }
 
+    private String completionsConfig(String url) {
+        return Strings.format("""
+            {
+                "service": "cohere",
+                "service_settings": {
+                    "api_key": "XXXX",
+                    "model_id": "command",
+                    "url": "%s"
+                }
+            }
+            """, url);
+    }
+
+    private String completionsResponse(ApiVersion version) {
+        return switch (version) {
+            case V1 -> v1CompletionsResponse();
+            case V2 -> v2CompletionsResponse();
+        };
+    }
+
+    private String v1CompletionsResponse() {
+        return """
+            {
+                "response_id": "some id",
+                "text": "result",
+                "generation_id": "some id",
+                "chat_history": [
+                    {
+                        "role": "USER",
+                        "message": "some input"
+                    },
+                    {
+                        "role": "CHATBOT",
+                        "message": "v1 response from the llm"
+                    }
+                ],
+                "finish_reason": "COMPLETE",
+                "meta": {
+                    "api_version": {
+                        "version": "1"
+                    },
+                    "billed_units": {
+                        "input_tokens": 4,
+                        "output_tokens": 191
+                    },
+                    "tokens": {
+                        "input_tokens": 70,
+                        "output_tokens": 191
+                    }
+                }
+            }
+            """;
+    }
+
+    private String v2CompletionsResponse() {
+        return """
+            {
+              "id": "c14c80c3-18eb-4519-9460-6c92edd8cfb4",
+              "finish_reason": "COMPLETE",
+              "message": {
+                "role": "assistant",
+                "content": [
+                  {
+                    "type": "text",
+                    "text": "v2 response from the LLM"
+                  }
+                ]
+              },
+              "usage": {
+                "billed_units": {
+                  "input_tokens": 1,
+                  "output_tokens": 2
+                },
+                "tokens": {
+                  "input_tokens": 3,
+                  "output_tokens": 4
+                }
+              }
+            }
+            """;
+    }
+
 }