Clean up trace

Author: jaybcee

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/elastic/ElasticInferenceServiceSparseEmbeddingsRequest.java

@@ -16,15 +16,16 @@
 import org.elasticsearch.xpack.inference.common.Truncator;
 import org.elasticsearch.xpack.inference.external.request.HttpRequest;
 import org.elasticsearch.xpack.inference.external.request.Request;
-import org.elasticsearch.xpack.inference.external.request.TraceContextPropagator;
 import org.elasticsearch.xpack.inference.services.elastic.ElasticInferenceServiceSparseEmbeddingsModel;
 import org.elasticsearch.xpack.inference.telemetry.TraceContext;
+import org.elasticsearch.xpack.inference.telemetry.TraceContextAware;
 
 import java.net.URI;
 import java.nio.charset.StandardCharsets;
 import java.util.Objects;
 
-public class ElasticInferenceServiceSparseEmbeddingsRequest implements ElasticInferenceServiceRequest {
+public class ElasticInferenceServiceSparseEmbeddingsRequest
+    implements ElasticInferenceServiceRequest, TraceContextAware {
 
     private final URI uri;
 
@@ -56,7 +57,7 @@ public HttpRequest createHttpRequest() {
         ByteArrayEntity byteEntity = new ByteArrayEntity(requestEntity.getBytes(StandardCharsets.UTF_8));
         httpPost.setEntity(byteEntity);
 
-        TraceContextPropagator.propagateTraceContext(httpPost, traceContext);
+        propagateTraceContext(httpPost);
 
         httpPost.setHeader(new BasicHeader(HttpHeaders.CONTENT_TYPE, XContentType.JSON.mediaType()));
 
@@ -73,6 +74,7 @@ public URI getURI() {
         return this.uri;
     }
 
+    @Override
     public TraceContext getTraceContext() {
         return traceContext;
     }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/elastic/ElasticInferenceServiceUnifiedChatCompletionRequest.java

@@ -12,20 +12,21 @@
 import org.apache.http.entity.ByteArrayEntity;
 import org.apache.http.message.BasicHeader;
 import org.elasticsearch.common.Strings;
-import org.elasticsearch.tasks.Task;
 import org.elasticsearch.xcontent.XContentType;
 import org.elasticsearch.xpack.inference.external.http.sender.UnifiedChatInput;
 import org.elasticsearch.xpack.inference.external.request.HttpRequest;
 import org.elasticsearch.xpack.inference.external.request.Request;
 import org.elasticsearch.xpack.inference.external.request.openai.OpenAiRequest;
 import org.elasticsearch.xpack.inference.services.elastic.completion.ElasticInferenceServiceCompletionModel;
 import org.elasticsearch.xpack.inference.telemetry.TraceContext;
+import org.elasticsearch.xpack.inference.telemetry.TraceContextAware;
 
 import java.net.URI;
 import java.nio.charset.StandardCharsets;
 import java.util.Objects;
 
-public class ElasticInferenceServiceUnifiedChatCompletionRequest implements OpenAiRequest {
+public class ElasticInferenceServiceUnifiedChatCompletionRequest
+    implements OpenAiRequest, TraceContextAware {
 
     // Implementing OpenAiRequest to ensure compatibility with the OpenAI API interface
     // This allows the ElasticInferenceService to handle requests in a standardized manner
@@ -94,20 +95,8 @@ public boolean isStreaming() {
         return true;
     }
 
+    @Override
     public TraceContext getTraceContext() {
         return traceContext;
     }
-
-    private void propagateTraceContext(HttpPost httpPost) {
-        var traceParent = traceContext.traceParent();
-        var traceState = traceContext.traceState();
-
-        if (traceParent != null) {
-            httpPost.setHeader(Task.TRACE_PARENT_HTTP_HEADER, traceParent);
-        }
-
-        if (traceState != null) {
-            httpPost.setHeader(Task.TRACE_STATE, traceState);
-        }
-    }
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/TraceContextPropagator.java renamed to x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/telemetry/TraceContextAware.java

@@ -4,17 +4,16 @@
  * 2.0; you may not use this file except in compliance with the Elastic License
  * 2.0.
  */
-
-package org.elasticsearch.xpack.inference.external.request;
+package org.elasticsearch.xpack.inference.telemetry;
 
 import org.apache.http.client.methods.HttpPost;
 import org.elasticsearch.tasks.Task;
-import org.elasticsearch.xpack.inference.telemetry.TraceContext;
 
-public final class TraceContextPropagator {
-    private TraceContextPropagator() {} // Utility class
+public interface TraceContextAware {
+    TraceContext getTraceContext();
 
-    public static void propagateTraceContext(HttpPost httpPost, TraceContext traceContext) {
+    default void propagateTraceContext(HttpPost httpPost) {
+        TraceContext traceContext = this.getTraceContext();
         if (traceContext == null) {
             return;
         }