feat: store http headers through an interceptors invocation chain (#550)

epam · Oct 29, 2024 · f88c451 · f88c451
1 parent f0a708b
commit f88c451
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 3 deletions.
diff --git a/server/src/main/java/com/epam/aidial/core/server/ProxyContext.java b/server/src/main/java/com/epam/aidial/core/server/ProxyContext.java
@@ -26,6 +26,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.Set;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -217,4 +218,15 @@ public ProxyContext exposeHeaders() {
 
         return this;
     }
-}
+
+    public String getRequestHeader(String name) {
+        String value = request.getHeader(name);
+        if (value != null) {
+            return value;
+        }
+        return Optional.ofNullable(apiKeyData)
+                .map(ApiKeyData::getHttpHeaders)
+                .map(h -> h.get(name))
+                .orElse(null);
+    }
+}
diff --git a/server/src/main/java/com/epam/aidial/core/server/data/ApiKeyData.java b/server/src/main/java/com/epam/aidial/core/server/data/ApiKeyData.java
@@ -1,6 +1,7 @@
 package com.epam.aidial.core.server.data;
 
 import com.epam.aidial.core.config.Key;
+import com.epam.aidial.core.server.Proxy;
 import com.epam.aidial.core.server.ProxyContext;
 import com.epam.aidial.core.server.security.ExtractedClaims;
 import com.fasterxml.jackson.annotation.JsonIgnore;
@@ -10,6 +11,9 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 /**
  * The container keeps data associated with API key.
@@ -23,6 +27,10 @@
  */
 @Data
 public class ApiKeyData {
+    private static final List<String> HTTP_HEADERS_TO_STORE = List.of(
+            Proxy.HEADER_JOB_TITLE,
+            Proxy.HEADER_CONVERSATION_ID
+    );
     // per request key is available with during the request lifetime. It's generated in runtime
     private String perRequestKey;
     // the key of root request initiator
@@ -47,6 +55,8 @@ public class ApiKeyData {
     // deployment triggers interceptors
     private String initialDeployment;
     private String initialDeploymentApi;
+    // Original HTTP headers to be stored during an interceptor invocation chain
+    private Map<String, String> httpHeaders = new HashMap<>();
 
     public ApiKeyData() {
     }
@@ -57,6 +67,7 @@ public static void initFromContext(ApiKeyData proxyApiKeyData, ProxyContext cont
         proxyApiKeyData.setInterceptors(context.getInterceptors());
         proxyApiKeyData.setInitialDeployment(context.getInitialDeployment());
         proxyApiKeyData.setInitialDeploymentApi(context.getInitialDeploymentApi());
+        proxyApiKeyData.setHttpHeaders(collectHttpHeaders(context));
 
         if (apiKeyData.getPerRequestKey() == null) {
             proxyApiKeyData.setOriginalKey(context.getKey());
@@ -80,4 +91,13 @@ public static void initFromContext(ApiKeyData proxyApiKeyData, ProxyContext cont
     public boolean isInterceptor() {
         return perRequestKey != null && interceptors != null && interceptorIndex >= 0 && interceptorIndex < interceptors.size();
     }
+
+    private static Map<String, String> collectHttpHeaders(ProxyContext context) {
+        if (!context.getApiKeyData().getHttpHeaders().isEmpty()) {
+            return context.getApiKeyData().getHttpHeaders();
+        }
+        return context.getRequest().headers().entries().stream()
+                .filter(h -> HTTP_HEADERS_TO_STORE.contains(h.getKey()))
+                .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+    }
 }
diff --git a/server/src/main/java/com/epam/aidial/core/server/log/GfLogStore.java b/server/src/main/java/com/epam/aidial/core/server/log/GfLogStore.java
@@ -69,7 +69,7 @@ private void append(ProxyContext context, LogEntry entry) throws JsonProcessingE
         HttpServerResponse response = context.getResponse();
 
         append(entry, "{\"apiType\":\"DialOpenAI\",\"chat\":{\"id\":\"", false);
-        append(entry, request.getHeader(Proxy.HEADER_CONVERSATION_ID), true);
+        append(entry, context.getRequestHeader(Proxy.HEADER_CONVERSATION_ID), true);
 
         append(entry, "\"},\"project\":{\"id\":\"", false);
         append(entry, context.getProject(), true);
@@ -78,7 +78,7 @@ private void append(ProxyContext context, LogEntry entry) throws JsonProcessingE
         append(entry, context.getUserHash(), true);
 
         append(entry, "\",\"title\":\"", false);
-        append(entry, request.getHeader(Proxy.HEADER_JOB_TITLE), true);
+        append(entry, context.getRequestHeader(Proxy.HEADER_JOB_TITLE), true);
         append(entry, "\"}", false);
 
         TokenUsage tokenUsage = context.getTokenUsage();