chore: Add getHeaders method to the ApiResponse

PiperOrigin-RevId: 755967387

Author: IvanLH

src/main/java/com/google/genai/ApiResponse.java

@@ -16,13 +16,17 @@
 
 package com.google.genai;
 
+import org.apache.http.Header;
 import org.apache.http.HttpEntity;
 
+
 /** The API response contains a response to a call to the GenAI APIs. */
 abstract class ApiResponse implements AutoCloseable {
   /** Gets the HttpEntity. */
   public abstract HttpEntity getEntity();
 
+  public abstract Header[] getHeaders();
+
   @Override
   public abstract void close();
 }

src/main/java/com/google/genai/HttpApiResponse.java

@@ -19,6 +19,7 @@
 import com.google.genai.errors.ApiException;
 import com.google.genai.errors.GenAiIOException;
 import java.io.IOException;
+import org.apache.http.Header;
 import org.apache.http.HttpEntity;
 import org.apache.http.client.methods.CloseableHttpResponse;
 
@@ -39,6 +40,12 @@ public HttpEntity getEntity() {
     return response.getEntity();
   }
 
+  /** Returns all of the headers from the response. */
+  @Override
+  public Header[] getHeaders() {
+    return response.getAllHeaders();
+  }
+
   /** Closes the Http response. */
   @Override
   public void close() {

src/main/java/com/google/genai/ReplayApiClient.java

@@ -36,9 +36,11 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.stream.Stream;
+import org.apache.http.Header;
 import org.apache.http.ProtocolVersion;
 import org.apache.http.StatusLine;
 import org.apache.http.entity.BasicHttpEntity;
+import org.apache.http.message.BasicHeader;
 import org.apache.http.message.BasicStatusLine;
 
 // TODO(b/369384123): Currently the ReplayApiClient mirrors the HttpApiClient. We will refactor the
@@ -131,10 +133,17 @@ public ApiResponse request(String httpMethod, String path, String requestJson) {
       Map<String, Object> responseMap = (Map<String, Object>) currentMember.get("response");
       Integer statusCode = (Integer) responseMap.get("status_code");
       List<Object> bodySegments = (List<Object>) responseMap.get("body_segments");
+      Map<String, String> headerMap = (Map<String, String>) responseMap.get("headers");
       StringBuilder responseBody = new StringBuilder();
       for (Object bodySegment : bodySegments) {
         responseBody.append(bodySegment.toString());
       }
+
+      Header[] headers = headerMap.entrySet()
+              .stream()
+              .map(entry -> new BasicHeader(entry.getKey(), entry.getValue()))
+              .toArray(Header[]::new);
+
       String responseString = responseBody.toString();
 
       BasicHttpEntity entity = new BasicHttpEntity();
@@ -144,7 +153,7 @@ public ApiResponse request(String httpMethod, String path, String requestJson) {
       StatusLine statusLine =
           new BasicStatusLine(new ProtocolVersion("HTTP", 1, 1), statusCode, "OK");
 
-      return new ReplayApiResponse(entity, statusLine);
+      return new ReplayApiResponse(entity, statusLine, headers);
     } else {
       // Note that if the client mode is "api", then the ReplayApiClient will not be used.
       throw new IllegalArgumentException("Invalid client mode: " + this.clientMode);

src/main/java/com/google/genai/ReplayApiResponse.java

@@ -16,6 +16,7 @@
 
 package com.google.genai;
 
+import org.apache.http.Header;
 import org.apache.http.HttpEntity;
 import org.apache.http.StatusLine;
 
@@ -25,17 +26,22 @@ final class ReplayApiResponse extends ApiResponse {
 
   private final HttpEntity entity;
   private final StatusLine statusLine;
+  private final Header[] headers;
 
-  public ReplayApiResponse(HttpEntity entity, StatusLine statusLine) {
+  public ReplayApiResponse(HttpEntity entity, StatusLine statusLine, Header[] headers) {
     this.entity = entity;
     this.statusLine = statusLine;
+    this.headers = headers;
   }
 
   @Override
   public HttpEntity getEntity() {
     return this.entity;
   }
 
+  @Override
+  public Header[] getHeaders() { return this.headers;}
+
   @Override
   public void close() {}
 }