withHttpResponse clean up

Author: ygree

dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/main/java/datadog/trace/instrumentation/openai_java/HttpResponseWrappers.java

@@ -55,7 +55,7 @@ public void close() {
 
   public static <T> HttpResponseFor<T> wrapHttpResponse(
       HttpResponseFor<T> response, AgentSpan span, BiConsumer<AgentSpan, T> afterParse) {
-    DECORATE.withHttpResponse(span, response);
+    DECORATE.withHttpResponse(span, response.headers());
     return new DDHttpResponseFor<T>(response) {
       @Override
       public T afterParse(T t) {
@@ -82,7 +82,7 @@ public static <T> HttpResponseFor<StreamResponse<T>> wrapHttpResponseStream(
       HttpResponseFor<StreamResponse<T>> response,
       final AgentSpan span,
       BiConsumer<AgentSpan, List<T>> decorate) {
-    DECORATE.withHttpResponse(span, response);
+    DECORATE.withHttpResponse(span, response.headers());
     return new DDHttpResponseFor<StreamResponse<T>>(response) {
       @Override
       public StreamResponse<T> afterParse(StreamResponse<T> streamResponse) {

dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/main/java/datadog/trace/instrumentation/openai_java/OpenAiDecorator.java

@@ -2,7 +2,6 @@
 
 import com.openai.core.ClientOptions;
 import com.openai.core.http.Headers;
-import com.openai.core.http.HttpResponse;
 import datadog.trace.api.DDSpanId;
 import datadog.trace.api.llmobs.LLMObsContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
@@ -70,9 +69,11 @@ private String currentLlmParentSpanId() {
     return Long.toString(parentLlmSpanId);
   }
 
-  public void withHttpResponse(AgentSpan span, HttpResponse response) {
-    Headers headers = response.headers();
-    setTagFromHeader(span, OPENAI_ORGANIZATION_NAME, headers, "openai-organization");
+  public void withHttpResponse(AgentSpan span, Headers headers) {
+    List<String> values = headers.values("openai-organization");
+    if (!values.isEmpty()) {
+      span.setTag(OPENAI_ORGANIZATION_NAME, values.get(0));
+    }
 
     setMetricFromHeader(
         span,
@@ -93,14 +94,6 @@ public void withHttpResponse(AgentSpan span, HttpResponse response) {
         "x-ratelimit-remaining-tokens");
   }
 
-  private static void setTagFromHeader(AgentSpan span, String tag, Headers headers, String header) {
-    List<String> values = headers.values(header);
-    if (values.isEmpty()) {
-      return;
-    }
-    span.setTag(tag, values.get(0));
-  }
-
   private static void setMetricFromHeader(
       AgentSpan span, String metric, Headers headers, String header) {
     List<String> values = headers.values(header);